Rename LG_agent_disabled.py to LG_agent.py

LG_agent_disabled.py → LG_agent.py

@@ -2,24 +2,21 @@ from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from typing import Annotated, TypedDict
 from langgraph.graph.message import add_messages
 from langchain_core.messages import HumanMessage, AIMessage, AnyMessage
-from tools import search_duckduckgo
 from langgraph.prebuilt import tools_condition
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import ToolNode
+from tools import all_tools
 
+def build_llm():
+    return ChatHuggingFace(
+        llm=HuggingFaceEndpoint(
+            repo_id="google/flan-t5-small",  # Change easily here
+            temperature=0,
+            max_new_tokens=512,
+        ),
+        verbose=True,
+    )
 
-llm_endpoint = HuggingFaceEndpoint(
-    repo_id="google/flan-t5-small",
-    temperature=0,
-    max_new_tokens=512
-)
-
-chat = ChatHuggingFace(llm=llm_endpoint, verbose=True)
-
-tools = [search_duckduckgo]
-chat_with_tools = chat.bind_tools(tools)
-
-# Generate the AgentState and Agent graph
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
 
@@ -28,28 +25,43 @@ def assistant(state: AgentState):
         "messages": [chat_with_tools.invoke(state["messages"])],
     }
 
-def build_graph():
+def build_graph(max_steps: int = 5):
     builder = StateGraph(AgentState)
     
     builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
+    builder.add_node("tools", ToolNode(all_tools))
     
     builder.add_edge(START, "assistant")
     builder.add_conditional_edges(
         "assistant",
-        # If the latest message requires a tool, route to tools
-        # Otherwise, provide a direct response
         tools_condition,
     )
     builder.add_edge("tools", "assistant")
-    agent = builder.compile()
-    return agent
+    
+    graph = builder.compile()
+
+    # Wrap the graph with step limiter
+    def limited_invoke(input_state):
+        steps = 0
+        state = input_state
+        while steps < max_steps:
+            state = graph.invoke(state)
+            latest_message = state["messages"][-1] if state["messages"] else None
+            if latest_message and isinstance(latest_message, AIMessage):
+                break  # Agent finished with direct answer
+            steps += 1
+        return state
+
+    return limited_invoke
 
 class BasicAgent:
-    def __init__(self):
-        self.graph = build_graph()
+    def __init__(self, max_steps: int = 5):
+        self.graph = build_graph(max_steps)
 
     def __call__(self, question: str) -> str:
-        response = self.graph.invoke({"messages": [HumanMessage(content=question)]})
-        final_message = response['messages'][-1]
-        return final_message.content if final_message else "No response."
+        response = self.graph({"messages": [HumanMessage(content=question)]})
+        if response.get("messages"):
+            final_message = response["messages"][-1]
+            return final_message.content if hasattr(final_message, "content") else "No final message."
+        else:
+            return "No response."