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

src/main.py

@@ -1,67 +1,78 @@
-#!/usr/bin/env python3
-"""
-FAQ Bot using ChromaDB and LangChain
-"""
-
 import os
 import sys
-import argparse
 from pathlib import Path
 
-from chromadb import Client
-from chromadb.config import Settings
+from dotenv import load_dotenv
+from langchain_community.document_loaders import TextLoader
+from langchain_community.embeddings import OllamaEmbeddings
+from langchain_community.vectorstores.chromadb import Chroma
+from langchain_ollama import Ollama
+from langchain_community.tools.mcp_tool import MCPTool
 
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.llms.openai import OpenAIChat
-from langchain.chains import RetrievalQA
-from langchain.vectorstores import Chroma
+# Load environment variables (if any)
+load_dotenv()
 
-from ingest import ingest_faq
-from retriever import get_answer
+# Directory containing FAQ documents (plain text files)
+DATA_DIR = Path("data")
+# Directory where ChromaDB will persist its data
+CHROMA_DIR = Path("chroma_db")
 
-def init_chroma(collection_name: str) -> Client:
+# Ensure directories exist
+DATA_DIR.mkdir(parents=True, exist_ok=True)
+CHROMA_DIR.mkdir(parents=True, exist_ok=True)
+
+# Load all text files from the data directory
+documents = []
+for txt_file in DATA_DIR.glob("*.txt"):
+    loader = TextLoader(str(txt_file))
+    documents.extend(loader.load())
+
+# Create embeddings using Ollama's nomic-embed-text model
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+# Create or load the ChromaDB vector store
+vectorstore = Chroma.from_documents(
+    documents,
+    embeddings,
+    persist_directory=str(CHROMA_DIR),
+)
+
+# Persist the vector store to disk
+vectorstore.persist()
+
+# Initialize the Ollama LLM for generation (e.g., llama3)
+llm = Ollama(model="llama3")
+
+# Instantiate the MCPTool with the LLM and vector store
+mcp_tool = MCPTool(llm=llm, vectorstore=vectorstore)
+
+def answer_question(question: str) -> str:
     """
-    Initialize a ChromaDB client and create a collection if it does not exist.
+    Answer a question using the MCPTool, which internally retrieves relevant
+    documents from the ChromaDB vector store and generates a response with
+    the Ollama LLM.
     """
-    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
+    return mcp_tool.run(question)
 
-def main():
-    parser = argparse.ArgumentParser(description="FAQ Bot CLI")
-    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()
-
-    # 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)
-
-    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()
+def main() -> None:
+    """
+    Simple command-line interface for the FAQ bot.
+    """
+    print("FAQ Bot powered by ChromaDB and Ollama.")
+    print("Type 'exit' to quit.")
+    while True:
+        try:
+            user_input = input("\nYour question: ").strip()
+        except (KeyboardInterrupt, EOFError):
+            print("\nExiting.")
+            break
+        if user_input.lower() in {"exit", "quit"}:
+            print("Goodbye!")
+            break
+        if not user_input:
+            continue
+        answer = answer_question(user_input)
+        print(f"\nAnswer: {answer}")
 
 if __name__ == "__main__":
     main()