Update LG_agent.py

LG_agent.py

@@ -1,48 +1,52 @@
-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 langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode, tools_condition
 from langgraph.graph import START, StateGraph
-from langgraph.prebuilt import ToolNode
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from tools import all_tools
+import os
+
+HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+if not HUGGINGFACEHUB_API_TOKEN:
+    raise ValueError("Missing Hugging Face token. Set HUGGINGFACEHUB_API_TOKEN environment variable.")
+
+# 1. Setup once
+llm = HuggingFaceEndpoint(
+    repo_id="tiiuae/falcon-7b-instruct",  # Change model here easily
+    huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
+    temperature=0,
+    max_new_tokens=512,
+)
 
-def build_llm():
-    return ChatHuggingFace(
-        llm=HuggingFaceEndpoint(
-            # repo_id="google/gemma-7b-it",
-            repo_id="tiiuae/falcon-7b-instruct",
-            temperature=0,
-            max_new_tokens=512,
-        ),
-        verbose=True,
-    )
+chat = ChatHuggingFace(llm=llm, verbose=True)
+chat_with_tools = chat.bind_tools(all_tools)
 
+# 2. Define the agent state
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
 
+# 3. Assistant node
+
 def assistant(state: AgentState):
-    chat_with_tools = build_llm().bind_tools(all_tools)  # Bind tools dynamically
     return {
         "messages": [chat_with_tools.invoke(state["messages"])],
     }
 
+# 4. Build the agent graph
 def build_graph(max_steps: int = 5):
     builder = StateGraph(AgentState)
-    
+
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(all_tools))
-    
+
     builder.add_edge(START, "assistant")
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-    )
+    builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
-    
+
     graph = builder.compile()
 
-    # Wrap the graph with step limiter
     def limited_invoke(input_state):
         steps = 0
         state = input_state
@@ -50,12 +54,14 @@ def build_graph(max_steps: int = 5):
             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
+                break
             steps += 1
         return state
 
     return limited_invoke
 
+# 5. BasicAgent class
+
 class BasicAgent:
     def __init__(self, max_steps: int = 5):
         self.graph = build_graph(max_steps)
@@ -66,4 +72,4 @@ class BasicAgent:
             final_message = response["messages"][-1]
             return final_message.content if hasattr(final_message, "content") else "No final message."
         else:
-            return "No response."
+            return "No response."