Spaces:
Sleeping
Sleeping
File size: 1,743 Bytes
ed65c68 5d970ca ed65c68 5d970ca ed65c68 5d970ca ed65c68 |
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 58 59 |
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"
|