diff --git a/example/AutoAgent.ipynb b/example/AutoAgent.ipynb index 291700e..0779de3 100644 --- a/example/AutoAgent.ipynb +++ b/example/AutoAgent.ipynb @@ -58,7 +58,6 @@ " 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", @@ -66,7 +65,7 @@ " termination_word=\"[#finish#]\"\n", ")\n", "\n", - "res = manager.start()\n", + "res = manager.start(\"what time is it ?\")\n", "res.content" ] }, diff --git a/example/BaseMemory.ipynb b/example/BaseMemory.ipynb index 59bed6f..115f9a3 100644 --- a/example/BaseMemory.ipynb +++ b/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" ] }, diff --git a/example/SmartPrompt.ipynb b/example/SmartPrompt.ipynb index 44f7bc1..538eebe 100644 --- a/example/SmartPrompt.ipynb +++ b/example/SmartPrompt.ipynb @@ -102,7 +102,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "1026a3b8", "metadata": {}, "outputs": [ diff --git a/iragent/__pycache__/agent.cpython-312.pyc b/iragent/__pycache__/agent.cpython-312.pyc index b7ca1fa..0cedbf2 100644 Binary files a/iragent/__pycache__/agent.cpython-312.pyc and b/iragent/__pycache__/agent.cpython-312.pyc differ diff --git a/iragent/__pycache__/models.cpython-312.pyc b/iragent/__pycache__/models.cpython-312.pyc index f8117ba..7e07aaa 100644 Binary files a/iragent/__pycache__/models.cpython-312.pyc and b/iragent/__pycache__/models.cpython-312.pyc differ diff --git a/iragent/__pycache__/prompts.cpython-312.pyc b/iragent/__pycache__/prompts.cpython-312.pyc index ebcf1fe..149da58 100644 Binary files a/iragent/__pycache__/prompts.cpython-312.pyc and b/iragent/__pycache__/prompts.cpython-312.pyc differ diff --git a/iragent/agent.py b/iragent/agent.py index 6c97c46..0042558 100644 --- a/iragent/agent.py +++ b/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")}, ) diff --git a/iragent/models.py b/iragent/models.py index be08b91..9869fc5 100644 --- a/iragent/models.py +++ b/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, diff --git a/iragent/prompts.py b/iragent/prompts.py index 9946cca..6d6668a 100644 --- a/iragent/prompts.py +++ b/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 = """ diff --git a/pyproject.toml b/pyproject.toml index d47f566..630d229 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_auto_agent_manager.py b/tests/test_auto_agent_manager.py index d7e0b50..1f93015 100644 --- a/tests/test_auto_agent_manager.py +++ b/tests/test_auto_agent_manager.py @@ -77,7 +77,6 @@ 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, @@ -85,6 +84,6 @@ def test_auto_agent_routing(self, mock_auto_router): termination_word="[#finish#]", ) - result = manager.start() + result = manager.start("What time is it now?") self.assertTrue(result.content.strip()) self.assertIn("15", result.content)