feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'

src/main.py

@@ -1,41 +1,67 @@
+#!/usr/bin/env python3
 """
-Command‑line interface for the FAQ bot.
+FAQ Bot using ChromaDB and LangChain
 """
 
-import argparse
 import os
+import sys
+import argparse
+from pathlib import Path
 
-from .bot import FAQBot
+from chromadb import Client
+from chromadb.config import Settings
+
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.llms.openai import OpenAIChat
+from langchain.chains import RetrievalQA
+from langchain.vectorstores import Chroma
+
+from ingest import ingest_faq
+from retriever import get_answer
+
+def init_chroma(collection_name: str) -> Client:
+    """
+    Initialize a ChromaDB client and create a collection if it does not exist.
+    """
+    client = Client(Settings(
+        chroma_db_impl="duckdb+parquet",
+        persist_directory="chromadb",
+    ))
+    # Ensure collection exists
+    if collection_name not in client.list_collections():
+        client.create_collection(name=collection_name)
+    return client
 
 def main():
     parser = argparse.ArgumentParser(description="FAQ Bot CLI")
-    parser.add_argument(
-        "--persist-dir",
-        type=str,
-        default="chromadb_persist",
-        help="Directory to persist ChromaDB data",
-    )
-    parser.add_argument(
-        "--openai-key",
-        type=str,
-        default=os.getenv("OPENAI_API_KEY"),
-        help="OpenAI API key (optional)",
-    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    ingest_parser = subparsers.add_parser("ingest", help="Ingest FAQ file into ChromaDB")
+    ingest_parser.add_argument("faq_file", type=Path, help="Path to FAQ text file")
+    ingest_parser.add_argument("--collection", type=str, default="faq_collection", help="Chroma collection name")
+
+    query_parser = subparsers.add_parser("ask", help="Ask a question to the FAQ bot")
+    query_parser.add_argument("question", type=str, help="Your question")
+    query_parser.add_argument("--collection", type=str, default="faq_collection", help="Chroma collection name")
+
     args = parser.parse_args()
 
-    bot = FAQBot(persist_dir=args.persist_dir, openai_api_key=args.openai_key)
+    # Ensure OpenAI API key is set
+    if "OPENAI_API_KEY" not in os.environ:
+        print("Error: OPENAI_API_KEY environment variable not set.", file=sys.stderr)
+        sys.exit(1)
 
-    print("FAQ Bot is ready. Type your question (Ctrl+C to exit).")
-    while True:
-        try:
-            question = input("\n> ")
-            if not question.strip():
-                continue
-            answer = bot.ask(question)
-            print(f"\nAnswer: {answer}")
-        except (KeyboardInterrupt, EOFError):
-            print("\nGoodbye!")
-            break
+    client = init_chroma(args.collection)
+
+    if args.command == "ingest":
+        ingest_faq(args.faq_file, client, args.collection)
+        print(f"Ingestion completed. Collection '{args.collection}' updated.")
+    elif args.command == "ask":
+        answer = get_answer(args.question, client, args.collection)
+        print("\nAnswer:\n")
+        print(answer)
+    else:
+        parser.print_help()
 
 if __name__ == "__main__":
     main()