NeuroBio_Agent / nodes.py
arnavmishra4's picture
Update nodes.py
f1deed6 verified
Raw
History Blame Contribute Delete
1.17 kB
from langchain_core.messages import ToolMessage
from config import llm
from tools import all_tools
llm_with_tools = llm.bind_tools(all_tools)
tools_map = {t.name : t for t in all_tools}
max_iterations = 10
def call_llm(state):
response = llm_with_tools.invoke(state["messages"])
return {
"messages":[response],
"retry_count": state["retry_count"]+1,
}
def run_tools(state):
last = state["messages"][-1]
results = [] # must be a plain list
for tool_call in last.tool_calls:
fn = tools_map[tool_call["name"]]
try:
result = fn.invoke(tool_call["args"])
except Exception as e:
result = f"Tool '{tool_call['name']}' failed: {e}"
results.append(
ToolMessage(content=str(result), tool_call_id=tool_call["id"])
)
return {"messages": results} # return the list, operator.add appends it
def should_continue(state):
last = state["messages"][-1]
if state['retry_count'] >= max_iterations:
return "end"
if getattr(last,"tool_calls",None):
return "tools"
return "end"