File size: 1,600 Bytes
7744be2 e92d49a d35bd88 e92d49a d35bd88 e92d49a d35bd88 e92d49a d35bd88 7744be2 d35bd88 e92d49a d35bd88 e92d49a d35bd88 e92d49a d35bd88 e92d49a d35bd88 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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)
|