Arxip222
8bdacb4f7a
All checks were successful
continuous-integration/drone/push Build is passing
* Delete legacy from bot * Clear old models * Единый http клиент * РАГ полечен
131 lines
6.0 KiB
Python
131 lines
6.0 KiB
Python
"""
|
||
RAG сервис для бота - вызывает API бэкенда
|
||
"""
|
||
from tg_bot.config.settings import settings
|
||
from tg_bot.infrastructure.http_client import create_http_session
|
||
|
||
|
||
class RAGService:
|
||
"""Сервис для работы с RAG через API бэкенда"""
|
||
|
||
async def get_or_create_conversation(
|
||
self,
|
||
user_telegram_id: str,
|
||
collection_id: str = None
|
||
) -> str | None:
|
||
"""Получить или создать беседу для пользователя"""
|
||
try:
|
||
async with create_http_session() as session:
|
||
async with session.get(
|
||
f"{settings.BACKEND_URL}/collections/",
|
||
headers={"X-Telegram-ID": user_telegram_id}
|
||
) as collections_response:
|
||
if collections_response.status != 200:
|
||
return None
|
||
collections = await collections_response.json()
|
||
|
||
if not collections:
|
||
if not collection_id:
|
||
async with session.post(
|
||
f"{settings.BACKEND_URL}/collections",
|
||
json={
|
||
"name": "Основная коллекция",
|
||
"description": "Коллекция по умолчанию",
|
||
"is_public": False
|
||
},
|
||
headers={"X-Telegram-ID": user_telegram_id}
|
||
) as create_collection_response:
|
||
if create_collection_response.status in [200, 201]:
|
||
collection_data = await create_collection_response.json()
|
||
collection_id = collection_data.get("collection_id")
|
||
else:
|
||
collection_id = collection_id
|
||
else:
|
||
collection_id = collections[0].get("collection_id")
|
||
|
||
if not collection_id:
|
||
return None
|
||
|
||
async with session.get(
|
||
f"{settings.BACKEND_URL}/conversations",
|
||
headers={"X-Telegram-ID": user_telegram_id}
|
||
) as conversations_response:
|
||
if conversations_response.status == 200:
|
||
conversations = await conversations_response.json()
|
||
for conv in conversations:
|
||
if conv.get("collection_id") == str(collection_id):
|
||
return conv.get("conversation_id")
|
||
|
||
async with session.post(
|
||
f"{settings.BACKEND_URL}/conversations",
|
||
json={"collection_id": str(collection_id)},
|
||
headers={"X-Telegram-ID": user_telegram_id}
|
||
) as create_conversation_response:
|
||
if create_conversation_response.status in [200, 201]:
|
||
conversation_data = await create_conversation_response.json()
|
||
return conversation_data.get("conversation_id")
|
||
|
||
return None
|
||
except Exception as e:
|
||
print(f"Error getting/creating conversation: {e}")
|
||
return None
|
||
|
||
async def generate_answer_with_rag(
|
||
self,
|
||
question: str,
|
||
user_telegram_id: str
|
||
) -> dict:
|
||
"""Генерирует ответ используя RAG через API бэкенда"""
|
||
try:
|
||
conversation_id = await self.get_or_create_conversation(user_telegram_id)
|
||
if not conversation_id:
|
||
return {
|
||
"answer": "Не удалось создать беседу. Попробуйте позже.",
|
||
"sources": [],
|
||
"usage": {}
|
||
}
|
||
|
||
async with create_http_session() as session:
|
||
async with session.post(
|
||
f"{settings.BACKEND_URL}/rag/question",
|
||
json={
|
||
"conversation_id": str(conversation_id),
|
||
"question": question,
|
||
"top_k": 20,
|
||
"rerank_top_n": 5
|
||
},
|
||
headers={"X-Telegram-ID": user_telegram_id}
|
||
) as response:
|
||
if response.status == 200:
|
||
result = await response.json()
|
||
sources = []
|
||
for source in result.get("sources", []):
|
||
sources.append({
|
||
"title": source.get("title", "Без названия"),
|
||
"document_id": source.get("document_id"),
|
||
"chunk_id": source.get("chunk_id"),
|
||
"index": source.get("index", 0)
|
||
})
|
||
|
||
return {
|
||
"answer": result.get("answer", "Не удалось сгенерировать ответ."),
|
||
"sources": sources,
|
||
"usage": result.get("usage", {}),
|
||
"conversation_id": str(conversation_id)
|
||
}
|
||
else:
|
||
error_text = await response.text()
|
||
print(f"RAG API error: {response.status} - {error_text}")
|
||
return {
|
||
"answer": "Ошибка при генерации ответа. Попробуйте позже.",
|
||
"sources": [],
|
||
"usage": {}
|
||
}
|
||
except Exception as e:
|
||
print(f"Error generating answer with RAG: {e}")
|
||
return {
|
||
"answer": "Произошла ошибка при генерации ответа. Попробуйте позже.",
|
||
"sources": [],
|
||
"usage": {}
|
||
}
|