flows/base.py

from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Generic, TypeVar

from langgraph.graph import StateGraph
from pydantic import BaseModel

InputT = TypeVar("InputT", bound=BaseModel)
OutputT = TypeVar("OutputT", bound=BaseModel)


class FlowState(BaseModel, Generic[InputT, OutputT]):
    flow_input: InputT
    flow_output: OutputT | None = None


class Flow(ABC, Generic[InputT, OutputT]):
    def __init__(self, output_keys: list[str] = ["flow_output"]):
        self.output_keys = output_keys

    @abstractmethod
    def get_state_graph(self) -> StateGraph: ...

    def run(self, inpt: InputT) -> OutputT:
        chain = self.get_state_graph().compile()
        result = chain.invoke(FlowState(flow_input=inpt), output_keys=self.output_keys)
        return self.post_run(result)

    def post_run(self, result: dict[str, Any]):
        """
        Implement this method in the subclass to post-process the result.
        """
        return result

    def export_graph_diagram(self, output_png_path: str | Path) -> None:
        if isinstance(output_png_path, Path):
            output_png_path = output_png_path.as_posix()
        self.get_state_graph().compile().get_graph().draw_png(output_png_path)