src/blogGenerator/nodes/route_deciding_node.py

from pydantic import BaseModel, Field
from typing_extensions import Literal
from blogGenerator.state.state import State
from langchain_core.messages import SystemMessage, HumanMessage


class Route(BaseModel):
    step: Literal["youtube", "topic"] = Field(
        None, description="The next step in the routing process"
    )


class RouteDecidingNode:
    """
    Decides where to route the query based on the user input.
    """

    def __init__(self, model):
        self.llm = model

    def process(self, state: State):

        # print(f"state['user_message'] : {state['user_message']}")
        # print(f"Node Called : RouteDecidingNode")

        route = self.llm.with_structured_output(Route)

        decision = route.invoke(
            [
                SystemMessage(content="Route the user message to youtube or topic."),
                HumanMessage(content=state["user_message"]),
            ]
        )

        # print(f"decision : {decision}")

        if decision.step == "youtube":
            extract_url = self.llm.invoke(
                [
                    SystemMessage(
                        content="Extract youtube url from user message. Only extract youtube link. Don't add any message."
                    ),
                    HumanMessage(content=state["user_message"]),
                ]
            )

            return {"yt_url": extract_url.content, "decision": "youtube"}

        return {"decision": decision.step}

    def route_decision(self, state: State):

        # print(f"state : {state}")
        # Return the node name you want to visit next
        if state["decision"] == "youtube":
            return "youtube"
        elif state["decision"] == "topic":
            return "topic"