Правки solution.py: Внеси изменения согласно инструкции: удаляем OpenAI и FastMC

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG-памятью/solution.py

@@ -1,80 +1,54 @@
 import os
-from typing import List
-
-from fastmcp import App, Route  # FastMCP – lightweight web framework
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores.qdrant import Qdrant
-from langchain.llms.openai import OpenAI
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
 from rich.console import Console
 
-# Конфигурация
+# Импортируем необходимые модули, которые реализованы в отдельных файлах
+from vector_store import get_vector_store
+from agent import RAGAgent
+from init_knowledge_base import init_kb
+
+console = Console()
+
+# Конфигурация окружения
 QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
 QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
     raise RuntimeError("Не задана переменная окружения OPENAI_API_KEY")
 
-# Инициализация консоли rich
-console = Console()
-
-# Векторный хранилище в Qdrant
-embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
-qdrant_store = Qdrant(
-    client=None,
-    collection_name="rag_memory",
-    embeddings=embeddings,
-    url=f"http://{QDRANT_HOST}:{QDRANT_PORT}",
+# Инициализация векторного хранилища
+vector_store = get_vector_store(
+    host=QDRANT_HOST,
+    port=QDRANT_PORT,
+    api_key=OPENAI_API_KEY,
 )
 
-# LLM и цепочка RAG
-llm = OpenAI(openai_api_key=OPENAI_API_KEY, temperature=0.7)
-prompt_template = PromptTemplate(
-    input_variables=["context", "question"],
-    template="Ниже приведена информация:\n{context}\n\nВопрос: {question}\nОтвет:",
-)
-rag_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    chain_type="stuff",
-    retriever=qdrant_store.as_retriever(search_kwargs={"k": 5}),
-    return_source_documents=True,
-    chain_type_kwargs={"prompt": prompt_template},
+# При необходимости загружаем начальные документы
+init_kb(vector_store)
+
+# Создаём агента RAG, передавая ему векторное хранилище и ключ OpenAI
+agent = RAGAgent(
+    vector_store=vector_store,
+    openai_api_key=OPENAI_API_KEY,
 )
 
-app = App()
-
-
-@app.post("/add")
-def add_document(text: str):
+def main() -> None:
     """
-    Добавляет новый документ в память агента.
+    Интерактивный цикл: пользователь вводит вопрос, агент отвечает.
+    Ввод 'exit' или 'quit' завершает программу.
     """
+    console.print("[bold green]Запуск интерактивного режима. Введите 'exit' для выхода.[/]")
+    while True:
+        question = input("Вопрос: ")
+        if question.lower() in ("exit", "quit"):
+            break
         try:
-        qdrant_store.add_texts([text])
-        console.log(f"[green]Документ добавлен[/green]")
-        return {"status": "ok"}
+            answer, sources = agent.ask(question)
+            console.print(f"[blue]Ответ:[/]\n{answer}")
+            console.print("[yellow]Источники:[/]")
+            for src in sources:
+                console.print(f"- {src}")
         except Exception as e:
             console.print_exception()
-        return {"status": "error", "detail": str(e)}
-
-
-@app.post("/ask")
-def ask(question: str):
-    """
-    Делает запрос к агенту с RAG‑памятью.
-    """
-    try:
-        result = rag_chain({"question": question})
-        answer = result["answer"]
-        sources = [doc.metadata.get("source", "unknown") for doc in result["source_documents"]]
-        console.log(f"[blue]Ответ[/blue]: {answer}")
-        return {"answer": answer, "sources": sources}
-    except Exception as e:
-        console.print_exception()
-        return {"status": "error", "detail": str(e)}
-
 
 if __name__ == "__main__":
-    # Запуск FastMCP сервера
-    app.run(host="0.0.0.0", port=8000)
+    main()