Update agents/agents_nodes.py

agents/agents_nodes.py

@@ -1,17 +1,3 @@
-import json
-from transformers import pipeline
-import torch
-from langchain_core.messages import AIMessage, ToolMessage
-from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace,HuggingFaceEndpoint
-from langchain_core.runnables import RunnableLambda
-from langgraph.prebuilt import ToolNode
-from utils.state_utils import AgentState
-from tools.financial_tools import time_value_tool
-
-
-
-# LLM instantation
-
 text_generator = pipeline(
     "text-generation",  # Task type
     #model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
@@ -25,41 +11,13 @@ text_generator = pipeline(
 
 llm = HuggingFacePipeline(pipeline=text_generator)
 
-# llm_instantiated = llm.bind(
-#     functions=[time_value_tool],
-#     function_call={"name": "time_value_tool"},
-# )
-
-# llm_instantiated = llm.bind(
-#     tools=[{"type": "function", "function": time_value_tool.get_schema()}],
-#     tool_choice="auto"
-# )
-
-llm_instantiated = llm.bind(
-    tools=[{
-        "type": "function",
-        "function": {
-            "name": time_value_tool.name,
-            "description": time_value_tool.description,
-            "parameters": time_value_tool.args
-        }
-    }],
-    tool_choice={
-        "type": "function",
-        "function": {"name": time_value_tool.name}
-    }
+llm = HuggingFacePipeline(pipeline=text_generator)
+llm_instantiated = llm.bind_tools(
+    [time_value_tool],
+    tool_choice={"type": "function", "function": {"name": "time_value_tool"}}
 )
-
-
 def agent_node(state: AgentState):
     response = llm_instantiated.invoke(state["messages"])
-    if response.tool_calls:
-        tool_name = response.tool_calls["name"]
-    if tool_name != "time_value_tool":
-        return {
-            "output": {"error": f"Invalid tool call: {tool_name}"},
-            "messages": [response]
-        }  
     if not (hasattr(response, 'tool_calls') and response.tool_calls):
         error_message = AIMessage(content="Error: Model failed to generate tool call.")
         return {"messages": [error_message]}
@@ -75,12 +33,8 @@ F_MAPPING = {
     "A/P": "Annual", "A/F": "Annual", "A/G": "Annual"
 }
 
-
 def format_output(state: AgentState):
     try:
-        if state.get("output"):
-            return state
-        
         # The last message should be the ToolMessage (from the tool node)
         if not state["messages"] or not isinstance(state["messages"][-1], ToolMessage):
             return {"output": {"error": "No tool result found in the last message"}}