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

src/main.py

@@ -1,68 +1,64 @@
+"""
+Main entry point for the FAQ bot.
+"""
+
+import argparse
 import os
-from fastapi import FastAPI, HTTPException, Query
-from pydantic import BaseModel
+import json
+from src.embedding import embed_text
+from src.database import ChromaDB
+from src.mcp_tools import generate_answer
 
-from .vector_store import QdrantVectorStore
-from .mcp_tool import MCPTool
-
-app = FastAPI(
-    title="FAQ Bot API",
-    description="A simple FAQ bot using Qdrant as the vector store.",
-    version="1.0.0",
-)
-
-# Initialize the vector store and MCP tool
-vector_store = QdrantVectorStore(
-    host=os.getenv("QDRANT_HOST"),
-    port=int(os.getenv("QDRANT_PORT", "6333")),
-    collection_name=os.getenv("QDRANT_COLLECTION", "faq_collection"),
-)
-mcp_tool = MCPTool(vector_store)
-
-
-class Document(BaseModel):
-    id: str
-    text: str
-    metadata: dict | None = None
-
-
-@app.post("/documents", status_code=201)
-def add_document(doc: Document):
+def load_sample_faq() -> list:
     """
-    Add a new document to the vector store.
+    Load a small sample FAQ dataset.
     """
-    try:
-        vector_store.add_document(doc.id, doc.text, doc.metadata)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-    return {"status": "added", "id": doc.id}
+    return [
+        {"id": "0", "text": "What is the return policy? Our return policy allows returns within 30 days of purchase."},
+        {"id": "1", "text": "How do I track my order? You can track your order using the tracking link sent to your email."},
+        {"id": "2", "text": "What payment methods are accepted? We accept Visa, MasterCard, and PayPal."},
+        {"id": "3", "text": "Do you ship internationally? Yes, we ship to most countries worldwide."},
+        {"id": "4", "text": "How can I contact customer support? You can contact us via email at support@example.com."},
+    ]
 
+def main():
+    parser = argparse.ArgumentParser(description="FAQ Bot using ChromaDB and MCP-tools")
+    parser.add_argument("--question", type=str, help="Your question to ask the bot")
+    args = parser.parse_args()
 
-@app.get("/search")
-def search(
-    query: str = Query(..., description="Search query"),
-    top_k: int = Query(5, ge=1, le=20, description="Number of results to return"),
-):
-    """
-    Search for documents similar to the query.
-    """
-    try:
-        results = vector_store.search(query, top_k=top_k)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-    return {"query": query, "results": results}
+    if not args.question:
+        print("Please provide a question using --question")
+        return
 
+    # Initialize database
+    db = ChromaDB()
 
-@app.get("/answer")
-def answer(
-    query: str = Query(..., description="Question to answer"),
-    top_k: int = Query(3, ge=1, le=10, description="Number of answers to return"),
-):
-    """
-    Get the best answer(s) for the query using MCPTool.
-    """
-    try:
-        answers = mcp_tool.answer(query, top_k=top_k)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-    return {"query": query, "answers": answers}
+    # If the collection is empty, load sample data
+    if db.is_empty():
+        print("Database empty. Loading sample FAQ data...")
+        sample_data = load_sample_faq()
+        for item in sample_data:
+            text = item["text"]
+            doc_id = item.get("id")
+            embedding = embed_text(text)
+            db.add_document(text, embedding, doc_id=doc_id)
+        print("Sample data loaded.")
+
+    # Generate embedding for the question
+    question_embedding = embed_text(args.question)
+
+    # Query the database for relevant documents
+    results = db.query(question_embedding, k=5)
+
+    # Extract context documents
+    context = [res["document"] for res in results]
+
+    # Generate answer using MCP-tools
+    answer = generate_answer(context, args.question)
+
+    # Output the answer
+    print("\nAnswer:")
+    print(answer)
+
+if __name__ == "__main__":
+    main()