good, next tool exec

Author: HomunMage

07 Agent Choose Tool/main.py

@@ -1,12 +1,63 @@
 from langgraph.graph import StateGraph, END
-from typing import TypedDict, Literal
-import random
+from typing import TypedDict, Literal, Callable, Dict, List, Any
+import inspect
 import json
+import random
 from langchain_community.chat_models import ChatOllama
 from langchain_core.prompts import PromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from abc import ABC, abstractmethod
 
+# Tool registry to hold information about tools
+tool_registry: List[Dict[str, Any]] = []
+
+# Decorator to register tools
+def register_tool(func: Callable) -> Callable:
+    signature = inspect.signature(func)
+    docstring = func.__doc__ or ""
+    params = [
+        {"name": param.name, "type": param.annotation}
+        for param in signature.parameters.values()
+    ]
+    tool_info = {
+        "name": func.__name__,
+        "description": docstring,
+        "parameters": params
+    }
+    tool_registry.append(tool_info)
+    return func
+
+# Define the tools with detailed parameter descriptions in the docstrings
+@register_tool
+def add(a: int, b: int) -> int:
+    """
+    :function: add   
+    :param int a: First number to add
+    :param int b: Second number to add
+    :return: Sum of a and b
+    """
+    return a + b
+
+@register_tool
+def ls() -> List[str]:
+    """
+    :function: ls
+    :return: List of filenames in the current directory
+    """
+    # Fake implementation
+    return ["file1.txt", "file2.txt", "file3.txt"]
+
+@register_tool
+def filewrite(name: str, content: str) -> None:
+    """
+    :function: filewrite
+    :param str name: Name of the file
+    :param str content: Content to write to the file
+    :return: None
+    """
+    # Fake implementation
+    print(f"Writing to {name}: {content}")
+
 # Specify the local language model
 local_llm = "mistral"
 llm = ChatOllama(model=local_llm, format="json", temperature=0)
@@ -22,7 +73,7 @@ class ToolState(TypedDict):
     history: str
     use_tool: bool
     tool_exec: str
-
+    tools_list: str
 
 # Define the base class for tasks
 class AgentBase(ABC):
@@ -41,10 +92,14 @@ def execute(self) -> ToolState:
         template = self.get_prompt_template()
         prompt = PromptTemplate.from_template(template)        
         llm_chain = prompt | llm | StrOutputParser()
-        generation = llm_chain.invoke({"history": self.state["history"], "use_tool": self.state["use_tool"]})
+        generation = llm_chain.invoke({
+            "history": self.state["history"], 
+            "use_tool": self.state["use_tool"],
+            "tools_list": self.state["tools_list"]
+        })
         data = json.loads(generation)
-        self.state["use_tool"] = data.get("use_tool", "")        
-        self.state["tool_exec"] = data
+        self.state["use_tool"] = data.get("use_tool", False)        
+        self.state["tool_exec"] = data.get("tool_exec", "")
 
         self.state["history"] += "\n" + generation
         self.state["history"] = clip_history(self.state["history"])
@@ -56,8 +111,9 @@ class ChatAgent(AgentBase):
     def get_prompt_template(self) -> str:
         return """
             {history}
-            As ChatAgent, decide we need use tool/py or not
-            if we don't need tool, just reply, otherwirse, let tool agent to handle
+            Available tools: {tools_list}
+            As ChatAgent, decide if we need to use a tool or not.
+            If we don't need a tool, just reply; otherwise, let the ToolAgent handle it.
             Output the JSON in the format: {{"scenario": "your reply", "use_tool": True/False}}
         """
 
@@ -71,12 +127,16 @@ def get_prompt_template(self) -> str:
         """
 
 def ToolExecute(state: ToolState) -> ToolState:
-    choice = self.llm_output(self.state["tool_exec"])
-    tool_name = choice["use_tool"]
-    args = self.convert_args(tool_name, choice["args"])
+    choice = json.loads(state["tool_exec"])
+    tool_name = choice["function"]
+    args = choice["args"]
     result = globals()[tool_name](*args)
+    state["history"] += f"\nExecuted {tool_name} with result: {result}"
+    state["history"] = clip_history(state["history"])
+    state["use_tool"] = False
+    return state
 
-# for conditional edges
+# For conditional edges
 def check_use_tool(state: ToolState) -> Literal["use tool", "not use tool"]:
     if state.get("use_tool") == True:
         return "use tool"
@@ -86,7 +146,7 @@ def check_use_tool(state: ToolState) -> Literal["use tool", "not use tool"]:
 # Define the state machine
 workflow = StateGraph(ToolState)
 
-# Initialize tasks for DM and Player
+# Initialize tasks for ChatAgent and ToolAgent
 def chat_agent(state: ToolState) -> ToolState:
     return ChatAgent(state).execute()
 
@@ -100,7 +160,6 @@ def tool_agent(state: ToolState) -> ToolState:
 workflow.set_entry_point("chat_agent")
 
 # Define edges between nodes
-
 workflow.add_conditional_edges(
     "chat_agent",
     check_use_tool,
@@ -111,19 +170,28 @@ def tool_agent(state: ToolState) -> ToolState:
 )
 
 workflow.add_edge('tool_agent', 'tool')
+workflow.add_edge('tool', 'chat_agent')
 
-
+# Generate the tools list
+tools_list = json.dumps([
+    {
+        "name": tool["name"],
+        "description": tool["description"]
+    }
+    for tool in tool_registry
+])
 
 # Compile the workflow into a runnable app
 app = workflow.compile()
 
 # Initialize the state
 initial_state = ToolState(
     history="help me ls files in current folder",
-    use_tool=False, 
-    )
+    use_tool=False,
+    tool_exec="",
+    tools_list=tools_list
+)
 
 for s in app.stream(initial_state):
     # Print the current state
-    print("for s in app.stream(initial_state):")
     print(s)