base.py

from typing import Annotated, Literal, TypedDict

from langchain_core.embeddings import Embeddings
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.graph.message import add_messages

from tools.external_search.state import SearchState, WebSearchResult
from tools.external_search.workflow import create_search_workflow
from tools.preprocessing.state import PreprocessingState
from tools.preprocessing.workflow import create_preprocessing_workflow

SAFETY_SAFE: Literal["true"] = "true"
SAFETY_UNSAFE: Literal["false"] = "false"
DEFAULT_LANGUAGE: str = "English"


class AdditionalContent(TypedDict, total=False):
    search_results: list[WebSearchResult]


class InputState(TypedDict, total=False):
    user_input: str  # User input text


class OutputState(TypedDict, total=False):
    answer: str  #  LLM response streamed to the user
    additional_content: AdditionalContent  # sends on graph completion


class BaseState(InputState, OutputState, total=False):
    rephrased_input: (
        str  # contextualized, LLM-generated standalone query from user input
    )
    chat_history: Annotated[list[BaseMessage], add_messages]
    safety: str  # LLM-assessed safety level of the user input
    reason_unsafe: str


class BaseGraphBuilder:

    def __init__(
        self,
        llm: BaseChatModel,
        embedding: Embeddings,
    ) -> None:
        self.preprocessing_workflow: Runnable = create_preprocessing_workflow(llm)
        self.search_workflow: Runnable = create_search_workflow(llm)

    async def preprocess(self, state: BaseState, config: RunnableConfig) -> BaseState:
        result: PreprocessingState = await self.preprocessing_workflow.ainvoke(
            PreprocessingState(
                user_input=state["user_input"],
                chat_history=state.get("chat_history", []),
            ),
            config,
        )
        mapped_state = BaseState(
            rephrased_input=result.get("rephrased_input", ""),
            safety=result.get("safety", SAFETY_SAFE),
            reason_unsafe=result.get("reason_unsafe", ""),
        )
        return BaseState(**state, **mapped_state)

    async def postprocess(self, state: BaseState, config: RunnableConfig) -> BaseState:
        search_results: list[WebSearchResult] = []
        if state.get("safety") == SAFETY_SAFE and config.get("configurable", {}).get(
            "enable_postprocess", False
        ):
            result: SearchState = await self.search_workflow.ainvoke(
                SearchState(
                    input=state["rephrased_input"],
                    generation=state["answer"],
                ),
                config=RunnableConfig(callbacks=config["callbacks"]),
            )
            search_results = result["search_results"]
        return BaseState(
            **state,
            additional_content=AdditionalContent(search_results=search_results),
        )