Update schedule_agent.py

schedule_agent.py

@@ -1,247 +1,202 @@
-from langchain_google_genai import ChatGoogleGenerativeAI
-
-
-from langchain.tools import tool
-
+from langchain.prompts import PromptTemplate
 from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import add_messages
-from langgraph.prebuilt import ToolNode, tools_condition,InjectedState
+
 from langchain_core.messages import (
-    SystemMessage,
+
     HumanMessage,
-    AIMessage,
-    ToolMessage,
+
 )
-from langgraph.types import Command, interrupt
+
 from langgraph.checkpoint.memory import MemorySaver
-from langchain_core.tools.base import InjectedToolCallId
+
 
 #structuring
 import ast
-from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
 #error handling with output parser
 from langchain.output_parsers import RetryOutputParser
 
-
-from dataclasses import dataclass
 from typing_extensions import TypedDict
-from typing import Annotated, Literal
-from pydantic import BaseModel, Field
-
-
-
-import os
-import requests
-import json
-from dotenv import load_dotenv 
-from os import listdir
-from os.path import isfile, join
-
-
-load_dotenv()
-
-
-# loading the necessary api keys
-GOOGLE_API_KEY=os.getenv('google_api_key')
 
+from pydantic import BaseModel, Field
+#get graph visuals
+from IPython.display import Image, display
+from langchain_core.runnables.graph import MermaidDrawMethod
 
-    
 
-GEMINI_MODEL='gemini-2.0-flash'
 
-llm = ChatGoogleGenerativeAI(google_api_key=GOOGLE_API_KEY, model=GEMINI_MODEL, temperature=0.3)
-    
 # state
 class State(TypedDict):
   """
   A dictionnary representing the state of the agent.
   """
-  messages: Annotated[list, add_messages]
+  node_message: str
   trip_data: dict
-
-# defining the tools for the agent to use
-
-@tool
-def local_files_browser(tool_call_id: Annotated[str, InjectedToolCallId]) -> str:
-  """
-  tool to list the local schedule files.
-  args:none
-  """
-  mypath=f'schedules/'
-  onlyfiles = [f for f in listdir(mypath) if isfile(join(mypath, f))]
-  if not onlyfiles:
-    return Command(update={'messages':[ToolMessage(f'No files are available, try to upload one',tool_call_id=tool_call_id)]})
-  else:
-    return Command(update={'messages':[ToolMessage(f'Here are the available schedules: {onlyfiles}',tool_call_id=tool_call_id)]})
-  
+  query: str
+  route:str
+
+
+class llm_nodes:
+
+  def __init__(self, llm:any):
+     self.model=llm
+  def schedule_creator_node(self,state:State):
+    llm=self.model
+    parser = JsonOutputParser()
+    prompt = PromptTemplate(
+    template="Answer the user query.\n{format_instructions}\n{query}\n",
+    input_variables=["query"],
+    partial_variables={"format_instructions": parser.get_format_instructions()},
+    )
+
+    chain = prompt | llm
+    result=chain.invoke({"query": f'from this query: {state.get('query')} turn the data into a schedule into a json format in the output, do not include ```json```, do not include comments either'})
+    try:
+
+          result=parser.parse(result.content)
+          return {'trip_data':result,
+                  'node_message':result}
+    except:
+        try: 
+          retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
+          result=retry_parser.parse_with_prompt(result.content, prompt)
+          return {'trip_data':result,
+                  'node_message':result}
+        except:
+          return {'trip_data':result.content,
+                        'node_message': f'created the schedule:{result.content}, but formating failed '}
+        
 
 
-@tool
-def schedule_loader(tool_call_id: Annotated[str, InjectedToolCallId],filename: str) -> Command:
-  """
-  Use this tool to load the schedule from local directory, which is a text file.
-  args: filename - the name of the file, include the extention.
-  return: schedule in a json format
-  """
-  try:
-    with open(f'schedules/{filename}', 'rb') as f:
-      schedule=f.read()
-
-      try:
-        parser = JsonOutputParser()
-        prompt = PromptTemplate(
-        template="Answer the user query.\n{format_instructions}\n{query}\n",
-        input_variables=["query"],
-        partial_variables={"format_instructions": parser.get_format_instructions()},
-        )
-
-        chain = prompt | llm
-        result=chain.invoke({"query": f'format this schedule: {str(schedule)} into a json format in the output, do not include ```json```, do not include comments either'})
-        result=parser.parse(result.content)
-        return Command(update={'trip_data':result,
-                'messages': [ToolMessage('Succesfully uploaded schedule',tool_call_id=tool_call_id)]})
-      except:
+  def schedule_editor_node(self,state:State):
+    """
+    Tool to make modifications to the schedule such as add, delete or modify.
+    Pass the query to the llm to edit the schedule.
+    args: query - the query to edit the schedule.
+    return: modified schedule in a json format
+    """
+    llm=self.model
+    file=state['trip_data']
+    # result=llm.invoke(f'Edit this schedule: {str(file)} following the instructions in the query: {query}, and include the changes in the schedule, but do not mention them specifically, only include the updated schedule json format in the output, do not include ```json```, do not include comments either')
+    parser = JsonOutputParser()
+    prompt = PromptTemplate(
+    template="Answer the user query.\n{format_instructions}\n{query}\n",
+    input_variables=["query"],
+    partial_variables={"format_instructions": parser.get_format_instructions()},
+    )
+
+    chain = prompt | llm
+    result=chain.invoke({"query":  f'Edit this schedule: {str(file)} following the instructions in the query: {state.get('query')}, and include the changes in the schedule, but do not mention them specifically, only include the updated schedule json format in the output, do not include ```json```, do not include comments either'})
+    try:
+          
+          result=parser.parse(result.content)
+          return {'trip_data':result,
+                  'node_message': f'edited the schedule with these changes:{result}'}
+    except:
         try: 
           retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
           result=retry_parser.parse_with_prompt(result.content, prompt)
-          return Command(update={'trip_data':result,
-                'messages': [ToolMessage('Succesfully uploaded schedule',tool_call_id=tool_call_id)]})
+          return {'trip_data':result,
+                  'node_message': f'edited the schedule with these changes:{result}'}
         except:
-          return Command(update={'trip_data':result.content,
-                        'messages': [ToolMessage(f'loaded the schedule:{result.content}, but formating failed ',tool_call_id=tool_call_id)]})
-  except:
-      return Command(update={'messages':[ToolMessage('No Schedule please try a different filename, or include the extention eg. filename.txt',tool_call_id=tool_call_id)]})
-  
-
-@tool
-def schedule_creator(tool_call_id: Annotated[str, InjectedToolCallId], schedule:str)->str:
-  """Tool to create a schedule from the chat with the agent
-  and then uses an llm to structure it.
-  args: schedule - the schedule from the chat
-  """
+          return {'trip_data':result.content,
+                        'node_message': f'edited the schedule with these changes:{result}, but formating failed '}
+        
+        
+  def agent_node(self,state:State):
+      llm=self.model
+      class Form(BaseModel):
+          route: str = Field(description= 'Return one of: schedule_creator, schedule_editor, show_schedule')
+          
+      parser=JsonOutputParser(pydantic_object=Form)
+      instruction=parser.get_format_instructions()
+      response=llm.invoke([HumanMessage(content=f"Based on this query: {state['query']}, select the appropriate route. Options are: schedule_creator, schedule_editor, show_schedule\n\n{instruction}")])
+      response=parser.parse(response.content)
+      route=response.get('route')
+
+      return {'route':route}
   
-  try:
-        parser = JsonOutputParser()
-        prompt = PromptTemplate(
-        template="Answer the user query.\n{format_instructions}\n{query}\n",
-        input_variables=["query"],
-        partial_variables={"format_instructions": parser.get_format_instructions()},
-        )
-
-        chain = prompt | llm
-        result=chain.invoke({"query": f'format this schedule: {str(schedule)} into a json format in the output, do not include ```json```, do not include comments either'})
-        result=parser.parse(result.content)
-        return Command(update={'trip_data':result,
-                'messages': [ToolMessage('Succesfully created schedule',tool_call_id=tool_call_id)]})
-  except:
-      try: 
-        retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
-        result=retry_parser.parse_with_prompt(result.content, prompt)
-        return Command(update={'trip_data':result,
-              'messages': [ToolMessage('Succesfully created schedule',tool_call_id=tool_call_id)]})
-      except:
-        return Command(update={'trip_data':result.content,
-                      'messages': [ToolMessage(f'created the schedule:{result.content}, but formating failed ',tool_call_id=tool_call_id)]})
-
-
-@tool
-def get_schedule(state: Annotated[dict, InjectedState])-> str:
-  """
-  Use this tool to get the information about the schedule once it has been loaded.
-  args: none
-  return: schedule
-  """
-  return state['trip_data']
-
-@tool
-def schedule_editor(query:str,state: Annotated[dict, InjectedState],tool_call_id: Annotated[str, InjectedToolCallId])-> str:
-  """
-  Tool to make modifications to the schedule such as add, delete or modify.
-  Pass the query to the llm to edit the schedule.
-  args: query - the query to edit the schedule.
-  return: modified schedule in a json format
-  """
-  file=state['trip_data']
-  try:
-        parser = JsonOutputParser()
-        prompt = PromptTemplate(
-        template="Answer the user query.\n{format_instructions}\n{query}\n",
-        input_variables=["query"],
-        partial_variables={"format_instructions": parser.get_format_instructions()},
-        )
-
-        chain = prompt | llm
-        result=chain.invoke({"query":  f'Edit this schedule: {str(file)} following the instructions in the query: {query}, and include the changes in the schedule, but do not mention them specifically, only include the updated schedule json format in the output, do not include ```json```, do not include comments either'})
-        result=parser.parse(result.content)
-        return Command(update={'trip_data':result,
-                'messages': [ToolMessage(f'edited the schedule with these changes:{result} ',tool_call_id=tool_call_id)]})
-  except:
-      try: 
-        retry_parser = RetryOutputParser.from_llm(parser=parser, llm=llm)
-        result=retry_parser.parse_with_prompt(result.content, prompt)
-        return Command(update={'trip_data':result,
-              'messages': [ToolMessage(f'edited the schedule with these changes:{result} ',tool_call_id=tool_call_id)]})
-      except:
-        return Command(update={'trip_data':result.content,
-                      'messages': [ToolMessage(f'edited the schedule with these changes:{result}, but formating failed ',tool_call_id=tool_call_id)]})
-
-@tool
-def save_schedule(state: Annotated[dict, InjectedState],tool_call_id: Annotated[str, InjectedToolCallId], filename: str) -> str:
-    """ 
-    Tool to save the schedule with a specified filename.
-    agrs: filename the name of the file, no need to include the extentions of the file
+  def show_schedule_node(self,state: State):
     """
-    file= state['trip_data']
-    return f'{filename} saved (however, since this is a demo, it is not really saved:))'
-
-
-
-
+    Use this tool to get the information about the schedule once it has been loaded.
+    args: none
+    return: schedule
+    """
+    schedule=state.get('trip_data')
+    if schedule:
+      return {"node_message":schedule}
+    else:
+       return{"node_message":"no schedule found, please upload one or add it in the chat"}
+
+def route(state:State):
+    route=state.get('route')
+    routing_map={
+    'schedule_creator': 'to_schedule_creator',
+    'schedule_editor': 'to_schedule_editor',
+    'show_schedule': 'to_show_schedule'
+    }
+    return routing_map.get(route)
+
+# langgraph
+#loading tools
 class Schedule_agent:
-    def __init__(self,llm:any):
+    def __init__(self, llm:any):
         self.agent=self._setup(llm)
     def _setup(self,llm):
-        
-        langgraph_tools=[get_schedule,schedule_creator,local_files_browser, save_schedule, schedule_editor,schedule_loader]
-
+        nodes=llm_nodes(llm)
+   
 
         graph_builder = StateGraph(State)
 
-        # Modification: tell the LLM which tools it can call
-        llm_with_tools = llm.bind_tools(langgraph_tools)
-        tool_node = ToolNode(tools=langgraph_tools)
-        def chatbot(state: State):
-            """ travel assistant that answers user questions about their trip.
-            Depending on the request, leverage which tools to use if necessary."""
-            return {"messages": [llm_with_tools.invoke(state['messages'])]}
 
-        graph_builder.add_node("chatbot", chatbot)
+        graph_builder.add_node("agent",nodes.agent_node)
 
 
-        graph_builder.add_node("tools", tool_node)
+        graph_builder.add_node('schedule_creator', nodes.schedule_creator_node)
+        graph_builder.add_node('schedule_editor', nodes.schedule_editor_node)
+        graph_builder.add_node('show_schedule',nodes.show_schedule_node)
         # Any time a tool is called, we return to the chatbot to decide the next step
-        graph_builder.set_entry_point("chatbot")
-        graph_builder.add_edge("tools", "chatbot")
+        graph_builder.set_entry_point("agent")
         graph_builder.add_conditional_edges(
-            "chatbot",
-            tools_condition,
+            "agent",
+            route,{
+                'to_schedule_creator': 'schedule_creator',
+                'to_schedule_editor': 'schedule_editor',
+                'to_show_schedule': 'show_schedule'
+            }
         )
+        graph_builder.add_edge('schedule_creator',END)
+        graph_builder.add_edge('schedule_editor',END)
+        graph_builder.add_edge('show_schedule',END)
         memory=MemorySaver()
         graph=graph_builder.compile(checkpointer=memory)
         return graph
         
-    def stream(self,input:str):
-        config = {"configurable": {"thread_id": "1"}}
-        input_message = HumanMessage(content=input)
-        for event in self.agent.stream({"messages": [input_message]}, config, stream_mode="values"):
-            event["messages"][-1].pretty_print()
 
+    def display_graph(self):
+        return display(
+                        Image(
+                                self.agent.get_graph().draw_mermaid_png(
+                                    draw_method=MermaidDrawMethod.API,
+                                )
+                            )
+                        )
     def chat(self,input:str):
         config = {"configurable": {"thread_id": "1"}}
-        response=self.agent.invoke({'messages':HumanMessage(content=str(input))},config)
-        return response['messages'][-1].content
+        response=self.agent.invoke({'query':input
+                                    },config)
+        return response
+
+    def stream(self,input:str):
+        config = {"configurable": {"thread_id": "1"}}
+        for event in self.agent.stream({'query':input
+                                        }, config, stream_mode="updates"):
+            print(event)
     
     def get_state(self, state_val:str):
         config = {"configurable": {"thread_id": "1"}}
-        return self.agent.get_state(config).values[state_val]
+        return self.agent.get_state(config).values[state_val]
+    
+    def update_state(self, data: dict):
+      config = {"configurable": {"thread_id": "1"}}
+      return self.agent.update_state(config, data)