Агент с RAG-памятью: solution.py

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

@@ -0,0 +1,80 @@
+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
+
+# Конфигурация
+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}",
+)
+
+# 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},
+)
+
+app = App()
+
+
+@app.post("/add")
+def add_document(text: str):
+    """
+    Добавляет новый документ в память агента.
+    """
+    try:
+        qdrant_store.add_texts([text])
+        console.log(f"[green]Документ добавлен[/green]")
+        return {"status": "ok"}
+    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)