Upload agents_nodes.py

agents/agents_nodes.py

@@ -1,76 +1,26 @@
 import json
 from langchain_core.messages import AIMessage, ToolMessage
 from langgraph.prebuilt import ToolNode
-from langchain_huggingface import HuggingFacePipeline
 from utils.state_utils import AgentState
 from tools.financial_tools import time_value_tool
-from transformers import pipeline
-import torch
+from langchain_ollama import ChatOllama
 
-time_value_schema = {
-    "name": "time_value_tool",
-    "description": "Computes time value of money factors using financial formulas",
-    "parameters": {
-        "type": "object",
-        "properties": {
-            "CF": {"type": "number"},
-            "F": {"type": "string"},
-            "i": {"type": "number"},
-            "n": {"type": "number"},
-            "g": {"type": "number", "nullable": True}
-        },
-        "required": ["CF", "F", "i", "n"]
-    }
-}
-
-
-
-text_generator = pipeline(
-    "text-generation",  # Task type
-    #model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
-    # model="google/gemma-3n-E2B-it",
-    model="google/gemma-3n-e4b-it",
-    #model="Qwen/Qwen3-Embedding-0.6B",
-    device= "cpu",
-    torch_dtype=torch.bfloat16,
-    max_new_tokens=200  # Limit output length
-)
 
-llm = HuggingFacePipeline(pipeline=text_generator)
+# LLL instantation
 
-llm = HuggingFacePipeline(pipeline=text_generator)
-# llm_instantiated = llm.bind_tools(
-#     [time_value_tool],
-#     tool_choice={"type": "function", "function": {"name": "time_value_tool"}}
-# )
-
-llm_instantiated = llm.bind(
-    tools=[time_value_schema],
-    tool_choice={
-        "type": "function",
-        "function": {"name": "time_value_tool"}
-    }
+llm = ChatOllama(model="qwen3:4b", temperature=0)
+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 not (hasattr(response, 'tool_calls') and response.tool_calls):
-#         error_message = AIMessage(content="Error: Model failed to generate tool call.")
-#         return {"messages": [error_message]}
-#     return {"messages": [response]}
-
 def agent_node(state: AgentState):
     response = llm_instantiated.invoke(state["messages"])
-    if not getattr(response, 'tool_calls', None):
-        return {
-            "messages": [response],
-            "output": {"error": "Failed to generate tool call"}
-        }
+    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]}
     return {"messages": [response]}
 
-
-
 # Tool node executes the tool
 tool_node = ToolNode([time_value_tool])