fix tool call in multiple call with json structured

example/AutoAgent.ipynb

@@ -58,15 +58,14 @@
     "                              system_prompt= \"You are that one who write the time in Persian. when you wrote time, then in new line write [#finish#]\")\n",
     "\n",
     "manager = AutoAgentManager(\n",
-    "    init_message=\"what time is it?\",\n",
     "    agents= [agent1,agent2,agent3],\n",
     "    first_agent=agent1,\n",
     "    max_round=5,\n",
     "    termination_fn=simple_termination,\n",
     "    termination_word=\"[#finish#]\"\n",
     ")\n",
     "\n",
-    "res = manager.start()\n",
+    "res = manager.start(\"what time is it ?\")\n",
     "res.content"
    ]
   },

example/BaseMemory.ipynb

@@ -72,7 +72,6 @@
     "                              system_prompt= \"You are that one who write the time in Persian. when you wrote time, then in new line write [#finish#]\", memory=BaseMemory)\n",
     "\n",
     "manager = AutoAgentManager(\n",
-    "    init_message=\"what time is it?\",\n",
     "    agents= [agent1,agent2,agent3],\n",
     "    first_agent=agent1,\n",
     "    max_round=5,\n",
@@ -81,7 +80,7 @@
     ")\n",
     "\n",
     "\n",
-    "res = manager.start()\n",
+    "res = manager.start(\"what time is it ?\")\n",
     "res.content"
    ]
   },

example/SmartPrompt.ipynb

@@ -102,7 +102,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "1026a3b8",
    "metadata": {},
    "outputs": [

iragent/agent.py

@@ -4,7 +4,6 @@
 import warnings
 from typing import Any, Callable, Dict, List, get_type_hints
 
-import requests
 from openai import OpenAI
 
 from .message import Message
@@ -107,85 +106,6 @@ def call_message(self, message: Message, **kwargs) -> str:
             self.memory.add_message(res)
         return res
 
-    def _call_ollama(self, msgs: List[Dict], message: Message) -> Message:
-        """!
-        This function use http call for ollama provider.
-        @param msgs:
-            this is a list of dictionary
-        """
-        payload = {"model": self.model, "messages": msgs, "stream": False}
-        function_payload = {
-            "tools": [{"type": "function", "function": f} for f in self.fn]
-        }
-        payload.update(function_payload)
-
-        try:
-            response = requests.post(
-                f"{self.base_url.removesuffix('/v1')}/api/chat", json=payload
-            )
-        except Exception as e:
-            raise ValueError(f"Error calling Ollama: {str(e)}")
-
-        response.raise_for_status()
-        result = response.json()
-        msg = result["message"]
-        tool_calls = msg.get("tool_calls", [])
-        if not tool_calls:
-            reply = msg.get("content", "")
-            return Message(
-                sender=self.name,
-                reciever=self.next_agent or message.sender,
-                content=reply.strip(),
-                metadata={"reply_to": message.metadata.get("message_id")},
-            )
-
-        # Handle tool call (assume one for now)
-        tool = tool_calls[0]
-        fn_name = tool["function"]["name"]
-        arguments = tool["function"]["arguments"]
-
-        if fn_name in self.function_map:
-            result_str = str(self.function_map[fn_name](**arguments))
-
-            # Add tool call + tool response to messages for a second round
-            followup_msgs = msgs + [
-                {"role": "assistant", "tool_calls": tool_calls, "content": ""},
-                {
-                    "role": "tool",
-                    "tool_call_id": tool.get("id", fn_name),
-                    "name": fn_name,
-                    "content": result_str,
-                },
-            ]
-
-            followup_payload = {
-                "model": self.model,
-                "messages": followup_msgs,
-                "stream": False,
-            }
-
-            followup_response = requests.post(
-                f"{self.base_url.removesuffix('/v1')}/api/chat", json=followup_payload
-            )
-            followup_response.raise_for_status()
-            followup = followup_response.json()
-
-            final_reply = followup["message"]["content"]
-            return Message(
-                sender=self.name,
-                reciever=self.next_agent or message.sender,
-                content=final_reply.strip(),
-                metadata={"reply_to": message.metadata.get("message_id")},
-            )
-
-        # fallback if function is not found
-        return Message(
-            sender=self.name,
-            reciever=self.next_agent or message.sender,
-            content=f"Function `{fn_name}` is not defined.",
-            metadata={"reply_to": message.metadata.get("message_id")},
-        )
-
     def _call_ollama_v2(self, msgs: List[Dict], message: Message) -> Message:
         """
         There is some different when you want to use ollama or openai call. this function work with "role":"tool".
@@ -204,42 +124,63 @@ def _call_ollama_v2(self, msgs: List[Dict], message: Message) -> Message:
         response = self.client.chat.completions.create(**kwargs)
 
         msg = response.choices[0].message
+
         # for function call
         if msg.tool_calls:
-            fn_name = msg.tool_calls[0].function.name
-            arguments = json.loads(msg.tool_calls[0].function.arguments)
-            if fn_name in self.function_map:
-                result = self.function_map[fn_name](**arguments)
-                followup = self.client.chat.completions.create(
-                    model=self.model,
-                    messages=msgs
-                    + [msg, {"role": "tool","tool_call_id": msg.tool_calls[0].id, "name": fn_name, "content": str(result)}],
-                    max_tokens=self.max_token,
-                    temperature=self.temprature,
-                )
-                return Message(
-                    sender=self.name,
-                    reciever=self.next_agent or message.sender,
-                    content=followup.choices[0].message.content.strip(),
-                    metadata={"reply_to": message.metadata.get("message_id")},
+            msgs.append(msg.model_dump())
+
+            for tool_call in msg.tool_calls:
+                fn_name = tool_call.function.name
+                arguments = json.loads(tool_call.function.arguments)
+                fn = self.function_map.get(
+                    fn_name, lambda **_: f"Function `{fn_name}` not found."
                 )
-        # Handle response format
-        if self.response_format:
-            try:
-                parsed_content = json.loads(msg.content)
-            except json.JSONDecodeError:
-                parsed_content = {"error": "Invalid JSON response", "raw": msg.content}
+                try:
+                    result = fn(**arguments)
+                except Exception as e:
+                    result = f"Error executing {fn_name}: {str(e)}"
+                msgs.append({
+                    "role": "tool",
+                    "tool_call_id": tool_call.id,
+                    "name": fn_name,
+                    "content": str(result)
+                })
+            followup = self.client.chat.completions.create(
+                model=self.model,
+                messages=msgs,
+                max_tokens=self.max_token,
+                temperature=self.temprature
+            )
+            follow_msg = followup.choices[0].message
+            content = follow_msg.content.strip() if follow_msg.content else ""
+            # Structured output handling
+            if self.response_format:
+                try:
+                    content = json.loads(content)
+                except json.JSONDecodeError:
+                    content = {"error": "Invalid JSON response", "raw": content}  
+
             return Message(
                 sender=self.name,
                 reciever=self.next_agent or message.sender,
-                content=parsed_content,
+                content=followup.choices[0].message.content.strip(),
                 metadata={"reply_to": message.metadata.get("message_id")},
             )
+        # --- Normal assistant reply (no tools) ---
+        msgs.append(msg.model_dump())
+        content = msg.content.strip() if msg.content else ""
+        # Structured output handling
+        if self.response_format:
+            try:
+                content = json.loads(content)
+            except json.JSONDecodeError:
+                content = {"error": "Invalid JSON response", "raw": content}
+
 
         return Message(
             sender=self.name,
             reciever=self.next_agent or message.sender,
-            content=response.choices[0].message.content.strip(),
+            content=content,
             metadata={"reply_to": message.metadata.get("message_id")},
         )
 

iragent/models.py

@@ -199,29 +199,21 @@ def __init__(
         max_round: int = 3,
         termination_fn: Callable = None,
         termination_word: str = None,
-        init_message: str = None,
     ) -> None:
         self.auto_agent = Agent(
             "agent_manager",
-            system_prompt="You are the agent manager.",
+            system_prompt="You are the agent manager. Who route the message between agents and user.",
             model=first_agent.model,
             base_url=first_agent.base_url,
             api_key=first_agent.api_key,
             temprature=0.1,
-            max_token=1024,
+            max_token=4096,
             memory=BaseMemory
         )
         self.first_agent = first_agent
         self.termination_fn = termination_fn
         self.max_round = max_round
         self.agents = {agent.name: agent for agent in agents}
-        self.init_msg = Message(
-            sender="user",
-            reciever=first_agent.name,
-            content=init_message,
-            intent="User request",
-            metadata={},
-        )
         self.termination_word = termination_word
 
     def start(self, message = None) -> Message:
@@ -246,13 +238,15 @@ def start(self, message = None) -> Message:
             for _ in range(self.max_round):
                 next_agent = (self.auto_agent.call_message(
                     Message(
-                        sender="auto_router",
+                        sender=self.first_agent.name,
                         reciever="agent_manager",
                         content=AUTO_AGENT_PROMPT.format(
-                            list_agents_info, last_msg.sender, last_msg.content
+                            list_agents_info, message ,last_msg.sender, last_msg.content
                         ),
                     )
-                )).content
+                )).content.strip()
+                if next_agent.lower() == "finish":
+                    return last_msg
                 if next_agent in self.agents.keys():
                     break
             last_msg.reciever = next_agent
@@ -544,7 +538,6 @@ def __init__(self, agent_factory: AgentFactory) -> None:
             max_token = 512
         )
         self.manager = AutoAgentManager(
-            init_message="",
             agents= [self.writer,self.reader],
             first_agent=self.writer,
             max_round=5,
@@ -561,8 +554,7 @@ def generate(self, input: str, output: str) -> str:
         {}
 
         """
-        self.manager.init_msg.content = msg.format(input, output)
-        return self.manager.start()
+        return self.manager.start(msg.format(input, output))
 
 
 class SmartAgentBuilder:
@@ -635,7 +627,6 @@ def create_agent(self, task: str, init_message: str = None) -> list[Agent]:
         )
         print(f"Agents are created : {' | '.join([a.name for a in agents_object])}")
         manager = AutoAgentManager(
-            init_message= init_message if init_message is not None else None,
             agents= agents_object,
             first_agent=agents_object[0],
             max_round=2 * len(agents_object),
@@ -811,7 +802,6 @@ def __init__(self, kg: KnowledgeGraphBuilder, agent_factory: AgentFactory) -> No
             max_token = 1024
         )
         self.manager = AutoAgentManager(
-            init_message=None,
             agents=[self.retriever_agent, self.generator_agent],
             first_agent=self.retriever_agent,
             max_round=10,

iragent/prompts.py

@@ -1,17 +1,24 @@
 AUTO_AGENT_PROMPT = """
 You are the Auto Agent Manager in a multi-agent AI system.
 
-Your job is to decide which agent should handle the next step based on the output of the previous agent.
+Your job: decide which agent should handle the next step based on the output of the previous agent.
 
-You will be given:
-1. A list of agents with their names and descriptions (system prompts)
-2. The output message from the last agent
+You are given:
+1. A list of available agents, with their names and role descriptions
+2. The last message produced by the current agent
 
-Respond with only the name of the next agent to route the message to.
+Your rules:
+- Always respond with only the name of the next agent, nothing else.
+- If the last agent’s response looks like a direct answer to the user or user's original request(e.g., it contains the final explanation, numbers, or insights requested by the user), respond with "finish".
+- Otherwise, select the most appropriate agent from the list.
 
 agents: {}
 
-{} message: {}
+User request:
+{}
+
+Last message (from {}):
+{}
 """
 
 SUMMARIZER_PROMPT = """

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "iragent"
-version = "0.1.8"
+version = "0.1.9"
 description = "A simple multi-agent framework"
 authors = [{ name = "Parsa Bakhtiari", email = "spacenavard1@gmail.com" }]
 readme = "README.md"

tests/test_auto_agent_manager.py

@@ -77,14 +77,13 @@ def test_auto_agent_routing(self, mock_auto_router):
         ]
 
         manager = AutoAgentManager(
-            init_message="What time is it now?",
             agents=[self.agent_a, self.agent_b, self.agent_c],
             first_agent=self.agent_a,
             max_round=5,
             termination_fn=simple_termination,
             termination_word="[#finish#]",
         )
 
-        result = manager.start()
+        result = manager.start("What time is it now?")
         self.assertTrue(result.content.strip())
         self.assertIn("15", result.content)