rag/chain.py

from __future__ import annotations

from functools import lru_cache
from typing import Any

from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_google_genai import ChatGoogleGenerativeAI
from pydantic import BaseModel

from .config import settings
from .prompts import prompt
from .retriever import get_retriever


class StructuredAnswer(BaseModel):
    answer: str
    sources: list[str]


@lru_cache(maxsize=1)
def get_chain() -> Any:
    llm_kwargs: dict[str, Any] = {
        "model": settings.google_genai_model,
        "temperature": settings.google_genai_temperature,
    }

    if settings.google_genai_api_key:
        llm_kwargs["api_key"] = settings.google_genai_api_key

    llm = ChatGoogleGenerativeAI(**llm_kwargs)
    context_formatter = RunnableLambda(
        lambda docs: "\n\n".join(
            f"[{i + 1}] id={doc.metadata.get('id', '')}\n{doc.page_content}"
            for i, doc in enumerate(docs)
        )
    )
    output_parser = PydanticOutputParser(pydantic_object=StructuredAnswer)
    prompt_with_format = prompt.partial(
        format_instructions=output_parser.get_format_instructions()
    )

    return (
        {"context": get_retriever() | context_formatter, "question": RunnablePassthrough()}
        | prompt_with_format
        | llm
        | output_parser
    )


def answer(question: str) -> StructuredAnswer:
    return get_chain().invoke(question)


async def aanswer(question: str) -> StructuredAnswer:
    return await get_chain().ainvoke(question)