transcription/src/rag/query.py

"""Запросы к 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,
    }