kg-scr/kg_ocr/rag/query.py

from __future__ import annotations

from typing import Any

DEFAULT_CHAT_MODEL = "openrouter/minimax/minimax-m2.5:free"

SYSTEM_PROMPT = (
    "Answer ONLY using the provided context. Cite which parts you're drawing "
    "from. If the context doesn't cover something, say 'not in my documents'."
)


def retrieve(embeddings: Any, query: str, limit: int = 3) -> list[dict]:
    """Search embeddings and return results with scores."""
    return embeddings.search(query, limit)


def ask_wllm(
    embeddings: Any,
    question: str,
    model: str = DEFAULT_CHAT_MODEL,
    limit: int = 3,
) -> str:
    """RAG: retrieve context from embeddings, then answer with an LLM."""
    try:
        import litellm
    except ImportError as exc:
        raise ImportError("ask_wllm needs litellm: uv sync --extra kg") from exc
    from dotenv import load_dotenv

    load_dotenv()

    results = retrieve(embeddings, question, limit)
    context = "\n\n".join([r["text"] for r in results])

    response = litellm.completion(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"Context from my documents:\n{context}\n\nQuestion: {question}",
            },
        ],
    )

    return response.choices[0].message.content