fee9b9acb1
Separate ASR (2 workers) from summary/RAG post-processing, add LightRAG chat API, batch upload fixes, and local model mounts for Docker deployment. Co-authored-by: Cursor <cursoragent@cursor.com>
155 lines
4.8 KiB
Python
155 lines
4.8 KiB
Python
"""Запросы к RAG и генерация ответов чат-бота через DeepSeek."""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from lightrag import QueryParam
|
||
from openai import AsyncOpenAI
|
||
|
||
from src.rag.indexer import get_global_rag, get_project_rag
|
||
|
||
|
||
async def retrieve_context(
|
||
question: str,
|
||
working_dir_base: Path,
|
||
project_name: Optional[str] = None,
|
||
mode: str = "hybrid",
|
||
api_key: str = "",
|
||
base_url: str = "https://opencode.ai/zen/v1",
|
||
index_model: str = "mimo-v2.5-free",
|
||
) -> str:
|
||
"""Извлекает релевантный контекст из LightRAG.
|
||
|
||
Args:
|
||
question: вопрос пользователя.
|
||
working_dir_base: базовая директория индексов.
|
||
project_name: если None — ищет в глобальном индексе.
|
||
mode: режим поиска LightRAG (naive, local, global, hybrid).
|
||
|
||
Returns:
|
||
Строка с найденным контекстом.
|
||
"""
|
||
if project_name:
|
||
rag = await get_project_rag(
|
||
project_name,
|
||
working_dir_base,
|
||
model=index_model,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
)
|
||
else:
|
||
rag = await get_global_rag(
|
||
working_dir_base,
|
||
model=index_model,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
)
|
||
|
||
# only_need_context=True — возвращает только найденные фрагменты без генерации ответа
|
||
param = QueryParam(mode=mode, only_need_context=True)
|
||
context = await rag.aquery(question, param=param)
|
||
return context if context else ""
|
||
|
||
|
||
async def generate_chat_response(
|
||
question: str,
|
||
context: str,
|
||
history: List[Dict[str, str]],
|
||
api_key: str,
|
||
base_url: str = "https://opencode.ai/zen/v1",
|
||
model: str = "deepseek-v4-flash-free",
|
||
) -> str:
|
||
"""Генерирует ответ чат-бота через DeepSeek (или другую модель).
|
||
|
||
Args:
|
||
question: вопрос пользователя.
|
||
context: контекст из RAG.
|
||
history: список {"question": ..., "answer": ...} предыдущих сообщений.
|
||
api_key: API ключ.
|
||
base_url: base URL OpenCode.
|
||
model: модель для чата.
|
||
|
||
Returns:
|
||
Ответ ассистента.
|
||
"""
|
||
if not api_key:
|
||
raise ValueError(
|
||
"OPENCODE_API_KEY не задан. Укажите rag.opencode_api_key в config.yaml "
|
||
"или переменную окружения OPENCODE_API_KEY."
|
||
)
|
||
client = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
||
|
||
system_prompt = (
|
||
"Ты — ассистент по протоколам совещаний строительной компании. "
|
||
"Отвечай на основе предоставленного контекста из протоколов. "
|
||
"Если в контексте нет ответа — так и скажи, не выдумывай. "
|
||
"Отвечай кратко и по делу."
|
||
)
|
||
|
||
messages = [{"role": "system", "content": system_prompt}]
|
||
|
||
for h in history:
|
||
messages.append({"role": "user", "content": h["question"]})
|
||
messages.append({"role": "assistant", "content": h["answer"]})
|
||
|
||
user_prompt = f"""Контекст из протоколов совещаний:
|
||
---
|
||
{context}
|
||
---
|
||
|
||
Вопрос: {question}
|
||
"""
|
||
messages.append({"role": "user", "content": user_prompt})
|
||
|
||
response = await client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=0.4,
|
||
max_tokens=2048,
|
||
)
|
||
content = response.choices[0].message.content
|
||
return content if content is not None else ""
|
||
|
||
|
||
async def rag_chat(
|
||
question: str,
|
||
working_dir_base: Path,
|
||
history: List[Dict[str, str]],
|
||
api_key: str,
|
||
project_name: Optional[str] = None,
|
||
base_url: str = "https://opencode.ai/zen/v1",
|
||
chat_model: str = "deepseek-v4-flash-free",
|
||
mode: str = "hybrid",
|
||
index_model: str = "mimo-v2.5-free",
|
||
) -> Dict[str, Any]:
|
||
"""Полный цикл RAG-чата: retrieval + generation.
|
||
|
||
Returns:
|
||
{"answer": str, "context": str, "project": str | None}
|
||
"""
|
||
context = await retrieve_context(
|
||
question=question,
|
||
working_dir_base=working_dir_base,
|
||
project_name=project_name,
|
||
mode=mode,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
index_model=index_model,
|
||
)
|
||
|
||
answer = await generate_chat_response(
|
||
question=question,
|
||
context=context,
|
||
history=history,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
model=chat_model,
|
||
)
|
||
|
||
return {
|
||
"answer": answer,
|
||
"context": context,
|
||
"project": project_name,
|
||
}
|