diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG-памятью/solution.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG-памятью/solution.py new file mode 100644 index 0000000..79e80a4 --- /dev/null +++ b/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) \ No newline at end of file