80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
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) |