Spaces:
Sleeping
Sleeping
| from utils.model_loader import ModelLoader | |
| from prompt_library.prompt import SYSTEM_PROMPT | |
| from langgraph.graph import StateGraph, MessagesState, END, START | |
| from langgraph.graph.message import add_messages # Is reducer function that just aggregrates the messages. | |
| from langgraph.prebuilt import ToolNode, tools_condition | |
| from tools.weather_info_tool import WeatherInfoTool | |
| from tools.place_search_tool import PlaceSearchTool | |
| from tools.expense_calculator_tool import CalculatorTool | |
| from tools.currency_conversion_tool import CurrencyConverterTool | |
| from typing import TypedDict, Sequence, Optional, Annotated | |
| from langchain_core.messages import BaseMessage, HumanMessage | |
| from datetime import datetime | |
| class AgentState(TypedDict): | |
| # messages is a list (or other sequence) of BaseMessage objects, and it has extra metadata attached to it: add_messages | |
| ## Sequence is a type hint that represents any ordered, iterable collection of items — like lists, tuples, or strings | |
| ## Annotated type lets you attach metadata to a type. It's not used by Python itself, but frameworks (like LangGraph, Pydantic, FastAPI) can use it. | |
| messages: Annotated[Sequence[BaseMessage], add_messages] | |
| class GraphBuilder(): | |
| def __init__(self,model_provider: str = "openai"): | |
| self.model_loader = ModelLoader(model_provider=model_provider) | |
| self.llm = self.model_loader.load_llm() | |
| self.tools = [] | |
| self.weather_tools = WeatherInfoTool() | |
| self.place_search_tools = PlaceSearchTool() | |
| self.calculator_tools = CalculatorTool() | |
| self.currency_converter_tools = CurrencyConverterTool() | |
| self.tools.extend([* self.weather_tools.weather_tool_list, | |
| * self.place_search_tools.place_search_tool_list, | |
| * self.calculator_tools.calculator_tool_list, | |
| * self.currency_converter_tools.currency_converter_tool_list]) | |
| self.llm_with_tools = self.llm.bind_tools(tools=self.tools) | |
| self.graph = None | |
| self.system_prompt = SYSTEM_PROMPT | |
| def agent_function(self,state: MessagesState): | |
| """Main agent function""" | |
| user_question = state["messages"] | |
| input_question = [self.system_prompt] + user_question | |
| # with mlflow.start_run(run_name="genai-trip-planner-llm-call", nested=True): | |
| response = self.llm_with_tools.invoke(input_question) | |
| return {"messages": [response]} | |
| def should_continue(self, state: AgentState) -> AgentState: | |
| messages = state['messages'] | |
| last_message = messages[-1] | |
| if not last_message.tool_calls: | |
| return "end" | |
| else: | |
| return "continue" | |
| def build_graph(self): | |
| graph_builder=StateGraph(MessagesState) | |
| graph_builder.add_node("agent", self.agent_function) | |
| graph_builder.add_node("tools", ToolNode(tools=self.tools)) | |
| graph_builder.add_edge(START,"agent") | |
| # graph_builder.add_conditional_edges("agent",tools_condition) | |
| graph_builder.add_conditional_edges( | |
| source="agent", | |
| path=self.should_continue, | |
| path_map={ | |
| # Edge: Node | |
| "end": END, | |
| "continue": "tools" | |
| } | |
| ) | |
| graph_builder.add_edge("tools","agent") | |
| # graph_builder.add_edge("agent",END) | |
| self.graph = graph_builder.compile() | |
| return self.graph | |
| def __call__(self): | |
| return self.build_graph() |