Удалить solution.py
This commit is contained in:
-71
@@ -1,71 +0,0 @@
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from langchain.embeddings.openai import OpenAIEmbeddings
|
||||
from langchain.vectorstores.qdrant import Qdrant
|
||||
from langchain.chains.question_answering import load_qa_chain
|
||||
from langchain.llms.openai import ChatOpenAI
|
||||
from rich.console import Console
|
||||
|
||||
# Конфигурация
|
||||
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
|
||||
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_memory")
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("Не задана переменная окружения OPENAI_API_KEY")
|
||||
|
||||
console = Console()
|
||||
|
||||
# Инициализация компонентов
|
||||
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
|
||||
llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
|
||||
|
||||
# Создание или подключение к коллекции Qdrant
|
||||
vectorstore = Qdrant(
|
||||
client=None,
|
||||
collection_name=QDRANT_COLLECTION,
|
||||
embeddings=embeddings,
|
||||
url=QDRANT_URL,
|
||||
)
|
||||
|
||||
app = FastAPI(title="RAG Agent")
|
||||
|
||||
class Document(BaseModel):
|
||||
content: str
|
||||
|
||||
class QueryRequest(BaseModel):
|
||||
question: str
|
||||
top_k: int = 5
|
||||
|
||||
@app.post("/add_document")
|
||||
def add_document(doc: Document):
|
||||
"""
|
||||
Добавляет документ в память агента.
|
||||
"""
|
||||
try:
|
||||
vectorstore.add_texts([doc.content])
|
||||
console.log(f"[green]Документ добавлен:[/green] {doc.content[:50]}...")
|
||||
return {"status": "ok"}
|
||||
except Exception as e:
|
||||
console.print_exception()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/ask")
|
||||
def ask(request: QueryRequest):
|
||||
"""
|
||||
Отвечает на вопрос, используя RAG.
|
||||
"""
|
||||
try:
|
||||
# Получаем похожие документы
|
||||
docs = vectorstore.similarity_search_with_score(request.question, k=request.top_k)
|
||||
contexts = [doc.page_content for doc, _ in docs]
|
||||
console.log(f"[blue]Найдено контекстов:[/blue] {len(contexts)}")
|
||||
|
||||
chain = load_qa_chain(llm=llm, chain_type="stuff")
|
||||
answer = chain.run(input_documents=[{"content": c} for c in contexts], question=request.question)
|
||||
return {"answer": answer}
|
||||
except Exception as e:
|
||||
console.print_exception()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
Reference in New Issue
Block a user