massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/v1/agent.py

@@ -1,19 +1,18 @@
-import os
-import sys
-import time
+# -*- coding: utf-8 -*-
 import json
-from typing import Callable, Union, Optional, List, Dict
-from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
+import time
+from abc import ABC
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from typing import Any, Callable, Dict, List, Optional
 
 from dotenv import load_dotenv
 
-load_dotenv()
+from .backends import gemini, grok, oai
+from .types import AgentResponse, AgentState, ModelConfig, TaskInput
+from .utils import function_to_json, get_agent_type_from_model
 
-from .types import TaskInput, AgentState, AgentResponse, ModelConfig
-from .utils import get_agent_type_from_model, function_to_json
-from abc import ABC, abstractmethod
-from typing import Any, Callable, Union, Optional, List, Dict
-from .backends import oai, gemini, grok
+load_dotenv()
 
 # TASK_INSTRUCTION = """
 # Please use your expertise and tools (if available) to fully verify if the best CURRENT ANSWER addresses the ORIGINAL MESSAGE.
@@ -24,7 +23,7 @@ from .backends import oai, gemini, grok
 # """
 
 SYSTEM_INSTRUCTION = f"""
-You are evaluating answers from multiple agents for final response to a message. 
+You are evaluating answers from multiple agents for final response to a message.
 
 For every aspect, claim, reasoning steps in the CURRENT ANSWERS, verify correctness, factual accuracy, and completeness using your expertise, reasoning, and available tools.
 
@@ -111,9 +110,7 @@ class MassAgent(ABC):
             "grok": grok.process_message,
         }
         if self.agent_type not in process_message_impl_map:
-            raise ValueError(
-                f"Unknown agent type: {self.agent_type}. Available types: {list(process_message_impl_map.keys())}"
-            )
+            raise ValueError(f"Unknown agent type: {self.agent_type}. Available types: {list(process_message_impl_map.keys())}")
 
         # Get the appropriate process_message implementation based on the agent type
         self.process_message_impl = process_message_impl_map[self.agent_type]
@@ -130,9 +127,7 @@ class MassAgent(ABC):
         self.stream_callback = stream_callback
         self.kwargs = kwargs
 
-    def process_message(
-        self, messages: List[Dict[str, str]], tools: List[str] = None
-    ) -> AgentResponse:
+    def process_message(self, messages: List[Dict[str, str]], tools: List[str] = None) -> AgentResponse:
         """
         Core LLM inference function for task processing.
 
@@ -164,9 +159,7 @@ class MassAgent(ABC):
         try:
             # Use ThreadPoolExecutor to implement timeout
             with ThreadPoolExecutor(max_workers=1) as executor:
-                future = executor.submit(
-                    self.process_message_impl, messages=messages, tools=tools, **config
-                )
+                future = executor.submit(self.process_message_impl, messages=messages, tools=tools, **config)
 
                 try:
                     # Wait for result with timeout
@@ -202,11 +195,9 @@ class MassAgent(ABC):
         """
         # Use the orchestrator to update the answer and notify other agents to restart
         self.orchestrator.notify_answer_update(self.agent_id, new_answer)
-        return f"The new answer has been added."
+        return "The new answer has been added."
 
-    def vote(
-        self, agent_id: int, reason: str = "", invalid_vote_options: List[int] = []
-    ):
+    def vote(self, agent_id: int, reason: str = "", invalid_vote_options: List[int] = []):
         """
         Vote for the representative agent, who you believe has found the correct solution.
 
@@ -255,9 +246,7 @@ class MassAgent(ABC):
                 deduplicated_function_calls.append(func_call)
         return deduplicated_function_calls
 
-    def _execute_function_calls(
-        self, function_calls: List[Dict], invalid_vote_options: List[int] = []
-    ):
+    def _execute_function_calls(self, function_calls: List[Dict], invalid_vote_options: List[int] = []):
         """Execute function calls and return function outputs."""
         from .tools import register_tool
 
@@ -311,9 +300,7 @@ class MassAgent(ABC):
 
                 # DEBUGGING
                 with open("function_calls.txt", "a") as f:
-                    f.write(
-                        f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] Agent {self.agent_id} ({self.model}):\n"
-                    )
+                    f.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] Agent {self.agent_id} ({self.model}):\n")
                     f.write(f"{json.dumps(error_output, indent=2)}\n")
                     f.write(f"Successful called: {False}\n")
 
@@ -335,7 +322,7 @@ class MassAgent(ABC):
                     "new_answer": {
                         "type": "string",
                         "description": "Your new answer, which should be self-contained, complete, and ready to serve as the definitive final response.",
-                    }
+                    },
                 },
                 "required": ["new_answer"],
             },
@@ -360,16 +347,8 @@ class MassAgent(ABC):
             },
         }
         # Check if there are any available options to vote for. If not, only return the add_answer schema.
-        available_options = [
-            agent_id
-            for agent_id, agent_state in self.orchestrator.agent_states.items()
-            if agent_state.curr_answer
-        ]
-        return (
-            [add_answer_schema, vote_schema]
-            if available_options
-            else [add_answer_schema]
-        )
+        available_options = [agent_id for agent_id, agent_state in self.orchestrator.agent_states.items() if agent_state.curr_answer]
+        return [add_answer_schema, vote_schema] if available_options else [add_answer_schema]
 
     def _get_registered_tools(self) -> List[Dict[str, Any]]:
         """Return the tool schema for the tools that are available to this agent."""
@@ -420,9 +399,7 @@ class MassAgent(ABC):
         agent_votes = []
         for agent_id, agent_state in self.orchestrator.agent_states.items():
             if agent_state.curr_vote:
-                agent_votes.append(
-                    f"**Vote for Agent {agent_state.curr_vote.target_id}**: {agent_state.curr_vote.reason}"
-                )
+                agent_votes.append(f"**Vote for Agent {agent_state.curr_vote.target_id}**: {agent_state.curr_vote.reason}")
         return agent_votes
 
     def _get_task_input(self, task: TaskInput) -> str:
@@ -440,11 +417,7 @@ class MassAgent(ABC):
         all_agent_answers = self._get_all_answers()
         all_agent_answers_str = "\n\n".join(all_agent_answers)
         # Check if in debate mode or not
-        voted_agents = [
-            agent_id
-            for agent_id, agent_state in self.orchestrator.agent_states.items()
-            if agent_state.curr_vote is not None
-        ]
+        voted_agents = [agent_id for agent_id, agent_state in self.orchestrator.agent_states.items() if agent_state.curr_vote is not None]
         if len(voted_agents) == len(self.orchestrator.agent_states):
             # Case 2: All agents have voted and are debating. Can not use agent status to check as they have been updated to 'working/debate'
             all_agent_votes = self._get_all_votes()
@@ -458,9 +431,7 @@ class MassAgent(ABC):
         else:
             # Case 3: All agents are working and not in debating
             status = "working"
-            task_input = AGENT_ANSWER_MESSAGE.format(
-                task=task.question, agent_answers=all_agent_answers_str
-            )
+            task_input = AGENT_ANSWER_MESSAGE.format(task=task.question, agent_answers=all_agent_answers_str)
 
         return status, task_input
 
@@ -500,17 +471,13 @@ class MassAgent(ABC):
 
         # Initialize working messages
         curr_round = 0
-        working_status, working_messages, all_tools = self._get_curr_messages_and_tools(
-            task
-        )
+        working_status, working_messages, all_tools = self._get_curr_messages_and_tools(task)
 
         # Start the task solving loop
         while curr_round < self.max_rounds and self.state.status == "working":
             try:
                 # Call LLM with current conversation
-                result = self.process_message(
-                    messages=working_messages, tools=all_tools
-                )
+                result = self.process_message(messages=working_messages, tools=all_tools)
 
                 # Before Making the new result into effect, check if there is any update from other agents that are unseen by this agent
                 agents_with_update = self.check_update()
@@ -521,30 +488,19 @@ class MassAgent(ABC):
 
                 # Add assistant response
                 if result.text:
-                    working_messages.append(
-                        {"role": "assistant", "content": result.text}
-                    )
+                    working_messages.append({"role": "assistant", "content": result.text})
 
                 # Execute function calls if any
                 if result.function_calls:
                     # Deduplicate function calls by their name
-                    result.function_calls = self.deduplicate_function_calls(
-                        result.function_calls
-                    )
+                    result.function_calls = self.deduplicate_function_calls(result.function_calls)
                     # Not voting if there is any update
-                    function_outputs, successful_called = self._execute_function_calls(
-                        result.function_calls, invalid_vote_options=agents_with_update
-                    )
+                    function_outputs, successful_called = self._execute_function_calls(result.function_calls, invalid_vote_options=agents_with_update)
 
                     renew_conversation = False
-                    for function_call, function_output, successful_called in zip(
-                        result.function_calls, function_outputs, successful_called
-                    ):
+                    for function_call, function_output, successful_called in zip(result.function_calls, function_outputs, successful_called):
                         # If call `add_answer`, we need to rebuild the conversation history with new answers
-                        if (
-                            function_call.get("name") == "add_answer"
-                            and successful_called
-                        ):
+                        if function_call.get("name") == "add_answer" and successful_called:
                             renew_conversation = True
                             break
 
@@ -553,17 +509,15 @@ class MassAgent(ABC):
                             renew_conversation = True
                             break
 
-                    if (
-                        not renew_conversation
-                    ):  # Add all function call results to the current conversation and continue the loop
-                        for function_call, function_output in zip(
-                            result.function_calls, function_outputs
-                        ):
+                    if not renew_conversation:  # Add all function call results to the current conversation and continue the loop
+                        for function_call, function_output in zip(result.function_calls, function_outputs):
                             working_messages.extend([function_call, function_output])
                     else:  # Renew the conversation
-                        working_status, working_messages, all_tools = (
-                            self._get_curr_messages_and_tools(task)
-                        )
+                        (
+                            working_status,
+                            working_messages,
+                            all_tools,
+                        ) = self._get_curr_messages_and_tools(task)
                 else:
                     # No function calls - check if we should continue or stop
                     if self.state.status == "voted":
@@ -573,15 +527,17 @@ class MassAgent(ABC):
                         # Check if there is any update from other agents that are unseen by this agent
                         if has_update and working_status != "initial":
                             # The vote option has changed, thus we need to renew the conversation within the loop
-                            working_status, working_messages, all_tools = (
-                                self._get_curr_messages_and_tools(task)
-                            )
+                            (
+                                working_status,
+                                working_messages,
+                                all_tools,
+                            ) = self._get_curr_messages_and_tools(task)
                         else:  # Continue the current conversation and prompting checkin
                             working_messages.append(
                                 {
                                     "role": "user",
                                     "content": "Finish your work above by making a tool call of `vote` or `add_answer`. Make sure you actually call the tool.",
-                                }
+                                },
                             )
 
                 curr_round += 1
@@ -592,9 +548,7 @@ class MassAgent(ABC):
                     break
 
             except Exception as e:
-                print(
-                    f"❌ Agent {self.agent_id} error in round {self.state.chat_round}: {e}"
-                )
+                print(f"❌ Agent {self.agent_id} error in round {self.state.chat_round}: {e}")
                 if self.orchestrator:
                     self.orchestrator.mark_agent_failed(self.agent_id, str(e))
 

massgen/v1/agents.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 """
 MassAgent implementations that wrap the existing agent backends.
 
@@ -5,20 +6,14 @@ This module provides MassAgent-compatible wrappers for the existing
 OpenAI, Gemini, and Grok agent implementations.
 """
 
-import os
-import sys
-import copy
-import time
-import traceback
-from typing import Callable, Union, Optional, List, Dict, Any
+from typing import Callable, Dict, List, Optional
 
 from dotenv import load_dotenv
 
-load_dotenv()
-
 from .agent import MassAgent
 from .types import ModelConfig, TaskInput
-from .tools import register_tool
+
+load_dotenv()
 
 
 class OpenAIMassAgent(MassAgent):
@@ -32,7 +27,6 @@ class OpenAIMassAgent(MassAgent):
         stream_callback: Optional[Callable] = None,
         **kwargs,
     ):
-
         # Pass all configuration to parent, including agent_type
         super().__init__(
             agent_id=agent_id,
@@ -54,7 +48,6 @@ class GrokMassAgent(OpenAIMassAgent):
         stream_callback: Optional[Callable] = None,
         **kwargs,
     ):
-
         # Pass all configuration to parent, including agent_type
         super().__init__(
             agent_id=agent_id,
@@ -76,7 +69,6 @@ class GeminiMassAgent(OpenAIMassAgent):
         stream_callback: Optional[Callable] = None,
         **kwargs,
     ):
-
         # Pass all configuration to parent, including agent_type
         super().__init__(
             agent_id=agent_id,
@@ -157,19 +149,13 @@ class GeminiMassAgent(OpenAIMassAgent):
                 if working_messages[-1].get("role", "") == "user":
                     if not function_call_enabled:
                         working_messages[-1]["content"] += (
-                            "\n\n"
-                            + "Note that the `add_answer` and `vote` tools are not enabled now. Please prioritize using the built-in tools to analyze the task first."
+                            "\n\n" + "Note that the `add_answer` and `vote` tools are not enabled now. Please prioritize using the built-in tools to analyze the task first."
                         )
                     else:
-                        working_messages[-1]["content"] += (
-                            "\n\n"
-                            + "Note that the `add_answer` and `vote` tools are enabled now."
-                        )
+                        working_messages[-1]["content"] += "\n\n" + "Note that the `add_answer` and `vote` tools are enabled now."
 
                 # Call LLM with current conversation
-                result = self.process_message(
-                    messages=working_messages, tools=available_tools
-                )
+                result = self.process_message(messages=working_messages, tools=available_tools)
 
                 # Before Making the new result into effect, check if there is any update from other agents that are unseen by this agent
                 agents_with_update = self.check_update()
@@ -180,29 +166,18 @@ class GeminiMassAgent(OpenAIMassAgent):
 
                 # Add assistant response
                 if result.text:
-                    working_messages.append(
-                        {"role": "assistant", "content": result.text}
-                    )
+                    working_messages.append({"role": "assistant", "content": result.text})
 
                 # Execute function calls if any
                 if result.function_calls:
                     # Deduplicate function calls by their name
-                    result.function_calls = self.deduplicate_function_calls(
-                        result.function_calls
-                    )
-                    function_outputs, successful_called = self._execute_function_calls(
-                        result.function_calls, invalid_vote_options=agents_with_update
-                    )
+                    result.function_calls = self.deduplicate_function_calls(result.function_calls)
+                    function_outputs, successful_called = self._execute_function_calls(result.function_calls, invalid_vote_options=agents_with_update)
 
                     renew_conversation = False
-                    for function_call, function_output, successful_called in zip(
-                        result.function_calls, function_outputs, successful_called
-                    ):
+                    for function_call, function_output, successful_called in zip(result.function_calls, function_outputs, successful_called):
                         # If call `add_answer`, we need to rebuild the conversation history with new answers
-                        if (
-                            function_call.get("name") == "add_answer"
-                            and successful_called
-                        ):
+                        if function_call.get("name") == "add_answer" and successful_called:
                             renew_conversation = True
                             break
 
@@ -213,17 +188,13 @@ class GeminiMassAgent(OpenAIMassAgent):
 
                     if not renew_conversation:
                         # Add all function call results to the current conversation
-                        for function_call, function_output in zip(
-                            result.function_calls, function_outputs
-                        ):
+                        for function_call, function_output in zip(result.function_calls, function_outputs):
                             working_messages.extend([function_call, function_output])
                         # If we have used custom tools, switch to built-in tools in the next round
                         if tool_switch:
                             available_tools = built_in_tools
                             function_call_enabled = False
-                            print(
-                                f"🔄 Agent {self.agent_id} (Gemini) switching to built-in tools in the next round"
-                            )
+                            print(f"🔄 Agent {self.agent_id} (Gemini) switching to built-in tools in the next round")
                     else:  # Renew the conversation
                         (
                             working_status,
@@ -259,16 +230,14 @@ class GeminiMassAgent(OpenAIMassAgent):
                                 {
                                     "role": "user",
                                     "content": "Finish your work above by making a tool call of `vote` or `add_answer`. Make sure you actually call the tool.",
-                                }
+                                },
                             )
 
                     # Switch to custom tools in the next round
                     if tool_switch:
                         available_tools = system_tools + custom_tools
                         function_call_enabled = True
-                        print(
-                            f"🔄 Agent {self.agent_id} (Gemini) switching to custom tools in the next round"
-                        )
+                        print(f"🔄 Agent {self.agent_id} (Gemini) switching to custom tools in the next round")
 
                 curr_round += 1
                 self.state.chat_round += 1
@@ -278,9 +247,7 @@ class GeminiMassAgent(OpenAIMassAgent):
                     break
 
             except Exception as e:
-                print(
-                    f"❌ Agent {self.agent_id} error in round {self.state.chat_round}: {e}"
-                )
+                print(f"❌ Agent {self.agent_id} error in round {self.state.chat_round}: {e}")
                 if self.orchestrator:
                     self.orchestrator.mark_agent_failed(self.agent_id, str(e))
 
@@ -318,9 +285,7 @@ def create_agent(
     }
 
     if agent_type not in agent_classes:
-        raise ValueError(
-            f"Unknown agent type: {agent_type}. Available types: {list(agent_classes.keys())}"
-        )
+        raise ValueError(f"Unknown agent type: {agent_type}. Available types: {list(agent_classes.keys())}")
 
     return agent_classes[agent_type](
         agent_id=agent_id,