massgen 0.1.0a3__py3-none-any.whl → 0.1.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of massgen might be problematic. Click here for more details.

Files changed (111) hide show
  1. massgen/__init__.py +1 -1
  2. massgen/agent_config.py +17 -0
  3. massgen/api_params_handler/_api_params_handler_base.py +1 -0
  4. massgen/api_params_handler/_chat_completions_api_params_handler.py +8 -1
  5. massgen/api_params_handler/_claude_api_params_handler.py +8 -1
  6. massgen/api_params_handler/_gemini_api_params_handler.py +73 -0
  7. massgen/api_params_handler/_response_api_params_handler.py +8 -1
  8. massgen/backend/base.py +31 -0
  9. massgen/backend/{base_with_mcp.py → base_with_custom_tool_and_mcp.py} +282 -11
  10. massgen/backend/chat_completions.py +182 -92
  11. massgen/backend/claude.py +115 -18
  12. massgen/backend/claude_code.py +378 -14
  13. massgen/backend/docs/CLAUDE_API_RESEARCH.md +3 -3
  14. massgen/backend/gemini.py +1275 -1607
  15. massgen/backend/gemini_mcp_manager.py +545 -0
  16. massgen/backend/gemini_trackers.py +344 -0
  17. massgen/backend/gemini_utils.py +43 -0
  18. massgen/backend/response.py +129 -70
  19. massgen/cli.py +577 -110
  20. massgen/config_builder.py +376 -27
  21. massgen/configs/README.md +111 -80
  22. massgen/configs/basic/multi/three_agents_default.yaml +1 -1
  23. massgen/configs/basic/single/single_agent.yaml +1 -1
  24. massgen/configs/providers/openai/gpt5_nano.yaml +3 -3
  25. massgen/configs/tools/custom_tools/claude_code_custom_tool_example.yaml +32 -0
  26. massgen/configs/tools/custom_tools/claude_code_custom_tool_example_no_path.yaml +28 -0
  27. massgen/configs/tools/custom_tools/claude_code_custom_tool_with_mcp_example.yaml +40 -0
  28. massgen/configs/tools/custom_tools/claude_code_custom_tool_with_wrong_mcp_example.yaml +38 -0
  29. massgen/configs/tools/custom_tools/claude_code_wrong_custom_tool_with_mcp_example.yaml +38 -0
  30. massgen/configs/tools/custom_tools/claude_custom_tool_example.yaml +24 -0
  31. massgen/configs/tools/custom_tools/claude_custom_tool_example_no_path.yaml +22 -0
  32. massgen/configs/tools/custom_tools/claude_custom_tool_with_mcp_example.yaml +35 -0
  33. massgen/configs/tools/custom_tools/claude_custom_tool_with_wrong_mcp_example.yaml +33 -0
  34. massgen/configs/tools/custom_tools/claude_wrong_custom_tool_with_mcp_example.yaml +33 -0
  35. massgen/configs/tools/custom_tools/gemini_custom_tool_example.yaml +24 -0
  36. massgen/configs/tools/custom_tools/gemini_custom_tool_example_no_path.yaml +22 -0
  37. massgen/configs/tools/custom_tools/gemini_custom_tool_with_mcp_example.yaml +35 -0
  38. massgen/configs/tools/custom_tools/gemini_custom_tool_with_wrong_mcp_example.yaml +33 -0
  39. massgen/configs/tools/custom_tools/gemini_wrong_custom_tool_with_mcp_example.yaml +33 -0
  40. massgen/configs/tools/custom_tools/github_issue_market_analysis.yaml +94 -0
  41. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_example.yaml +24 -0
  42. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_example_no_path.yaml +22 -0
  43. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_with_mcp_example.yaml +35 -0
  44. massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_with_wrong_mcp_example.yaml +33 -0
  45. massgen/configs/tools/custom_tools/gpt5_nano_wrong_custom_tool_with_mcp_example.yaml +33 -0
  46. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_example.yaml +25 -0
  47. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_example_no_path.yaml +23 -0
  48. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_with_mcp_example.yaml +34 -0
  49. massgen/configs/tools/custom_tools/gpt_oss_custom_tool_with_wrong_mcp_example.yaml +34 -0
  50. massgen/configs/tools/custom_tools/gpt_oss_wrong_custom_tool_with_mcp_example.yaml +34 -0
  51. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_example.yaml +24 -0
  52. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_example_no_path.yaml +22 -0
  53. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_with_mcp_example.yaml +35 -0
  54. massgen/configs/tools/custom_tools/grok3_mini_custom_tool_with_wrong_mcp_example.yaml +33 -0
  55. massgen/configs/tools/custom_tools/grok3_mini_wrong_custom_tool_with_mcp_example.yaml +33 -0
  56. massgen/configs/tools/custom_tools/qwen_api_custom_tool_example.yaml +25 -0
  57. massgen/configs/tools/custom_tools/qwen_api_custom_tool_example_no_path.yaml +23 -0
  58. massgen/configs/tools/custom_tools/qwen_api_custom_tool_with_mcp_example.yaml +36 -0
  59. massgen/configs/tools/custom_tools/qwen_api_custom_tool_with_wrong_mcp_example.yaml +34 -0
  60. massgen/configs/tools/custom_tools/qwen_api_wrong_custom_tool_with_mcp_example.yaml +34 -0
  61. massgen/configs/tools/custom_tools/qwen_local_custom_tool_example.yaml +24 -0
  62. massgen/configs/tools/custom_tools/qwen_local_custom_tool_example_no_path.yaml +22 -0
  63. massgen/configs/tools/custom_tools/qwen_local_custom_tool_with_mcp_example.yaml +35 -0
  64. massgen/configs/tools/custom_tools/qwen_local_custom_tool_with_wrong_mcp_example.yaml +33 -0
  65. massgen/configs/tools/custom_tools/qwen_local_wrong_custom_tool_with_mcp_example.yaml +33 -0
  66. massgen/configs/tools/filesystem/claude_code_context_sharing.yaml +1 -1
  67. massgen/configs/voting/gemini_gpt_voting_sensitivity.yaml +67 -0
  68. massgen/formatter/_chat_completions_formatter.py +104 -0
  69. massgen/formatter/_claude_formatter.py +120 -0
  70. massgen/formatter/_gemini_formatter.py +448 -0
  71. massgen/formatter/_response_formatter.py +88 -0
  72. massgen/frontend/coordination_ui.py +4 -2
  73. massgen/logger_config.py +35 -3
  74. massgen/message_templates.py +56 -6
  75. massgen/orchestrator.py +179 -10
  76. massgen/stream_chunk/base.py +3 -0
  77. massgen/tests/custom_tools_example.py +392 -0
  78. massgen/tests/mcp_test_server.py +17 -7
  79. massgen/tests/test_config_builder.py +423 -0
  80. massgen/tests/test_custom_tools.py +401 -0
  81. massgen/tests/test_tools.py +127 -0
  82. massgen/tool/README.md +935 -0
  83. massgen/tool/__init__.py +39 -0
  84. massgen/tool/_async_helpers.py +70 -0
  85. massgen/tool/_basic/__init__.py +8 -0
  86. massgen/tool/_basic/_two_num_tool.py +24 -0
  87. massgen/tool/_code_executors/__init__.py +10 -0
  88. massgen/tool/_code_executors/_python_executor.py +74 -0
  89. massgen/tool/_code_executors/_shell_executor.py +61 -0
  90. massgen/tool/_exceptions.py +39 -0
  91. massgen/tool/_file_handlers/__init__.py +10 -0
  92. massgen/tool/_file_handlers/_file_operations.py +218 -0
  93. massgen/tool/_manager.py +634 -0
  94. massgen/tool/_registered_tool.py +88 -0
  95. massgen/tool/_result.py +66 -0
  96. massgen/tool/_self_evolution/_github_issue_analyzer.py +369 -0
  97. massgen/tool/docs/builtin_tools.md +681 -0
  98. massgen/tool/docs/exceptions.md +794 -0
  99. massgen/tool/docs/execution_results.md +691 -0
  100. massgen/tool/docs/manager.md +887 -0
  101. massgen/tool/docs/workflow_toolkits.md +529 -0
  102. massgen/tool/workflow_toolkits/__init__.py +57 -0
  103. massgen/tool/workflow_toolkits/base.py +55 -0
  104. massgen/tool/workflow_toolkits/new_answer.py +126 -0
  105. massgen/tool/workflow_toolkits/vote.py +167 -0
  106. {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/METADATA +89 -131
  107. {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/RECORD +111 -36
  108. {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/WHEEL +0 -0
  109. {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/entry_points.txt +0 -0
  110. {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/licenses/LICENSE +0 -0
  111. {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/top_level.txt +0 -0
massgen/__init__.py CHANGED
@@ -68,7 +68,7 @@ from .chat_agent import (
68
68
  from .message_templates import MessageTemplates, get_templates
69
69
  from .orchestrator import Orchestrator, create_orchestrator
70
70
 
71
- __version__ = "0.1.0a3"
71
+ __version__ = "0.1.1"
72
72
  __author__ = "MassGen Contributors"
73
73
 
74
74
 
massgen/agent_config.py CHANGED
@@ -58,6 +58,9 @@ class AgentConfig:
58
58
  timeout_config: Timeout and resource limit configuration
59
59
  coordination_config: Coordination behavior configuration (e.g., planning mode)
60
60
  skip_coordination_rounds: Debug/test mode - skip voting rounds and go straight to final presentation (default: False)
61
+ voting_sensitivity: Controls how critical agents are when voting ("lenient", "balanced", "strict")
62
+ max_new_answers_per_agent: Maximum number of new answers each agent can provide (None = unlimited)
63
+ answer_novelty_requirement: How different new answers must be from existing ones ("lenient", "balanced", "strict")
61
64
  """
62
65
 
63
66
  # Core backend configuration (includes tool enablement)
@@ -66,6 +69,11 @@ class AgentConfig:
66
69
  # Framework configuration
67
70
  message_templates: Optional["MessageTemplates"] = None
68
71
 
72
+ # Voting behavior configuration
73
+ voting_sensitivity: str = "lenient"
74
+ max_new_answers_per_agent: Optional[int] = None
75
+ answer_novelty_requirement: str = "lenient"
76
+
69
77
  # Agent customization
70
78
  agent_id: Optional[str] = None
71
79
  _custom_system_instruction: Optional[str] = field(default=None, init=False)
@@ -696,6 +704,9 @@ class AgentConfig:
696
704
  "backend_params": self.backend_params,
697
705
  "agent_id": self.agent_id,
698
706
  "custom_system_instruction": self.custom_system_instruction,
707
+ "voting_sensitivity": self.voting_sensitivity,
708
+ "max_new_answers_per_agent": self.max_new_answers_per_agent,
709
+ "answer_novelty_requirement": self.answer_novelty_requirement,
699
710
  "timeout_config": {
700
711
  "orchestrator_timeout_seconds": self.timeout_config.orchestrator_timeout_seconds,
701
712
  },
@@ -730,6 +741,9 @@ class AgentConfig:
730
741
  backend_params = data.get("backend_params", {})
731
742
  agent_id = data.get("agent_id")
732
743
  custom_system_instruction = data.get("custom_system_instruction")
744
+ voting_sensitivity = data.get("voting_sensitivity", "lenient")
745
+ max_new_answers_per_agent = data.get("max_new_answers_per_agent")
746
+ answer_novelty_requirement = data.get("answer_novelty_requirement", "lenient")
733
747
 
734
748
  # Handle timeout_config
735
749
  timeout_config = TimeoutConfig()
@@ -756,6 +770,9 @@ class AgentConfig:
756
770
  message_templates=message_templates,
757
771
  agent_id=agent_id,
758
772
  custom_system_instruction=custom_system_instruction,
773
+ voting_sensitivity=voting_sensitivity,
774
+ max_new_answers_per_agent=max_new_answers_per_agent,
775
+ answer_novelty_requirement=answer_novelty_requirement,
759
776
  timeout_config=timeout_config,
760
777
  coordination_config=coordination_config,
761
778
  )
@@ -21,6 +21,7 @@ class APIParamsHandlerBase(ABC):
21
21
  """
22
22
  self.backend = backend_instance
23
23
  self.formatter = backend_instance.formatter
24
+ self.custom_tool_manager = backend_instance.custom_tool_manager
24
25
 
25
26
  @abstractmethod
26
27
  async def build_api_params(
@@ -23,6 +23,7 @@ class ChatCompletionsAPIParamsHandler(APIParamsHandlerBase):
23
23
  "enable_code_interpreter",
24
24
  "allowed_tools",
25
25
  "exclude_tools",
26
+ "custom_tools", # Custom tools configuration (processed separately)
26
27
  },
27
28
  )
28
29
 
@@ -114,11 +115,17 @@ class ChatCompletionsAPIParamsHandler(APIParamsHandlerBase):
114
115
  if provider_tools:
115
116
  combined_tools.extend(provider_tools)
116
117
 
117
- # User-defined tools
118
+ # Workflow tools
118
119
  if tools:
119
120
  converted_tools = self.formatter.format_tools(tools)
120
121
  combined_tools.extend(converted_tools)
121
122
 
123
+ # Add custom tools
124
+ custom_tools = self.custom_tool_manager.registered_tools
125
+ if custom_tools:
126
+ converted_custom_tools = self.formatter.format_custom_tools(custom_tools)
127
+ combined_tools.extend(converted_custom_tools)
128
+
122
129
  # MCP tools
123
130
  mcp_tools = self.get_mcp_tools()
124
131
  if mcp_tools:
@@ -22,6 +22,7 @@ class ClaudeAPIParamsHandler(APIParamsHandlerBase):
22
22
  "enable_code_execution",
23
23
  "allowed_tools",
24
24
  "exclude_tools",
25
+ "custom_tools", # Custom tools configuration (processed separately)
25
26
  "_has_files_api_files",
26
27
  },
27
28
  )
@@ -97,11 +98,17 @@ class ClaudeAPIParamsHandler(APIParamsHandlerBase):
97
98
  if provider_tools:
98
99
  combined_tools.extend(provider_tools)
99
100
 
100
- # User-defined tools
101
+ # Workflow tools
101
102
  if tools:
102
103
  converted_tools = self.formatter.format_tools(tools)
103
104
  combined_tools.extend(converted_tools)
104
105
 
106
+ # Add custom tools
107
+ custom_tools = self.custom_tool_manager.registered_tools
108
+ if custom_tools:
109
+ converted_custom_tools = self.formatter.format_custom_tools(custom_tools)
110
+ combined_tools.extend(converted_custom_tools)
111
+
105
112
  # MCP tools
106
113
  mcp_tools = self.get_mcp_tools()
107
114
  if mcp_tools:
@@ -0,0 +1,73 @@
1
+ # -*- coding: utf-8 -*-
2
+ """
3
+ Gemini API parameters handler building SDK config with parameter mapping and exclusions.
4
+ """
5
+
6
+ from typing import Any, Dict, List, Set
7
+
8
+ from ._api_params_handler_base import APIParamsHandlerBase
9
+
10
+
11
+ class GeminiAPIParamsHandler(APIParamsHandlerBase):
12
+ def get_excluded_params(self) -> Set[str]:
13
+ base = self.get_base_excluded_params()
14
+ extra = {
15
+ "enable_web_search",
16
+ "enable_code_execution",
17
+ "use_multi_mcp",
18
+ "mcp_sdk_auto",
19
+ "allowed_tools",
20
+ "exclude_tools",
21
+ "custom_tools",
22
+ }
23
+ return set(base) | extra
24
+
25
+ def get_provider_tools(self, all_params: Dict[str, Any]) -> List[Dict[str, Any]]:
26
+ """
27
+ These are SDK Tool objects (from google.genai.types), not JSON tool declarations.
28
+ """
29
+ tools: List[Any] = []
30
+
31
+ if all_params.get("enable_web_search", False):
32
+ try:
33
+ from google.genai.types import GoogleSearch, Tool
34
+
35
+ tools.append(Tool(google_search=GoogleSearch()))
36
+ except Exception:
37
+ # Gracefully ignore if SDK not available
38
+ pass
39
+
40
+ if all_params.get("enable_code_execution", False):
41
+ try:
42
+ from google.genai.types import Tool, ToolCodeExecution
43
+
44
+ tools.append(Tool(code_execution=ToolCodeExecution()))
45
+ except Exception:
46
+ # Gracefully ignore if SDK not available
47
+ pass
48
+
49
+ return tools
50
+
51
+ async def build_api_params(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], all_params: Dict[str, Any]) -> Dict[str, Any]:
52
+ """
53
+ Build a config dict for google-genai Client.generate_content_stream.
54
+ - Map max_tokens -> max_output_tokens
55
+ - Do not include 'model' here; caller extracts it
56
+ - Do not add builtin tools; stream logic handles them
57
+ - Do not handle MCP tools or coordination schema here
58
+ """
59
+ excluded = self.get_excluded_params()
60
+ config: Dict[str, Any] = {}
61
+
62
+ for key, value in all_params.items():
63
+ if key in excluded or value is None:
64
+ continue
65
+ if key == "max_tokens":
66
+ config["max_output_tokens"] = value
67
+ elif key == "model":
68
+ # Caller will extract model separately
69
+ continue
70
+ else:
71
+ config[key] = value
72
+
73
+ return config
@@ -22,6 +22,7 @@ class ResponseAPIParamsHandler(APIParamsHandlerBase):
22
22
  "enable_code_interpreter",
23
23
  "allowed_tools",
24
24
  "exclude_tools",
25
+ "custom_tools", # Custom tools configuration (processed separately)
25
26
  "_has_file_search_files", # Internal flag for file search tracking
26
27
  },
27
28
  )
@@ -88,11 +89,17 @@ class ResponseAPIParamsHandler(APIParamsHandlerBase):
88
89
  if provider_tools:
89
90
  combined_tools.extend(provider_tools)
90
91
 
91
- # Add framework tools
92
+ # Add workflow tools
92
93
  if tools:
93
94
  converted_tools = self.formatter.format_tools(tools)
94
95
  combined_tools.extend(converted_tools)
95
96
 
97
+ # Add custom tools
98
+ custom_tools = self.custom_tool_manager.registered_tools
99
+ if custom_tools:
100
+ converted_custom_tools = self.formatter.format_custom_tools(custom_tools)
101
+ combined_tools.extend(converted_custom_tools)
102
+
96
103
  # Add MCP tools (use OpenAI format)
97
104
  mcp_tools = self._convert_mcp_tools_to_openai_format()
98
105
  if mcp_tools:
massgen/backend/base.py CHANGED
@@ -59,6 +59,14 @@ class LLMBackend(ABC):
59
59
  # Initialize utility classes
60
60
  self.token_usage = TokenUsage()
61
61
 
62
+ # # Initialize tool manager
63
+ # self.custom_tool_manager = ToolManager()
64
+
65
+ # # Register custom tools if specified
66
+ # custom_tools = kwargs.get("custom_tools", [])
67
+ # if custom_tools:
68
+ # self._register_custom_tools(custom_tools)
69
+
62
70
  # Planning mode flag - when True, MCP tools should be blocked during coordination
63
71
  self._planning_mode_enabled: bool = False
64
72
 
@@ -125,6 +133,29 @@ class LLMBackend(ABC):
125
133
  self.api_params_handler = None
126
134
  self.coordination_stage = None
127
135
 
136
+ # def _register_custom_tools(self, tool_names: list[str]) -> None:
137
+ # """Register custom tool functions.
138
+
139
+ # Args:
140
+ # tool_names: List of tool names to register
141
+ # """
142
+ # import importlib
143
+
144
+ # for tool_name in tool_names:
145
+ # try:
146
+ # # Try to import from tool module
147
+ # module = importlib.import_module("massgen.tool")
148
+ # if hasattr(module, tool_name):
149
+ # tool_func = getattr(module, tool_name)
150
+ # self.custom_tool_manager.add_tool_function(tool_func)
151
+ # print(f"Successfully registered custom tool: {tool_name}")
152
+ # else:
153
+ # print(f"Warning: Tool '{tool_name}' not found in massgen.tool")
154
+ # except ImportError as e:
155
+ # print(f"Warning: Could not import tool module: {e}")
156
+ # except Exception as e:
157
+ # print(f"Error registering tool '{tool_name}': {e}")
158
+
128
159
  def _setup_permission_hooks(self):
129
160
  """Setup permission hooks for function-based backends (default behavior)."""
130
161
  # Create per-agent hook manager
@@ -17,6 +17,7 @@ from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
17
17
  import httpx
18
18
 
19
19
  from ..logger_config import log_backend_activity, logger
20
+ from ..tool import ToolManager
20
21
  from .base import LLMBackend, StreamChunk
21
22
 
22
23
 
@@ -107,13 +108,22 @@ SUPPORTED_AUDIO_MIME_TYPES = {
107
108
  }
108
109
 
109
110
 
110
- class MCPBackend(LLMBackend):
111
+ class CustomToolAndMCPBackend(LLMBackend):
111
112
  """Base backend class with MCP (Model Context Protocol) support."""
112
113
 
113
114
  def __init__(self, api_key: Optional[str] = None, **kwargs):
114
115
  """Initialize backend with MCP support."""
115
116
  super().__init__(api_key, **kwargs)
116
117
 
118
+ # Custom tools support - initialize before api_params_handler
119
+ self.custom_tool_manager = ToolManager()
120
+ self._custom_tool_names: set[str] = set()
121
+
122
+ # Register custom tools if provided
123
+ custom_tools = kwargs.get("custom_tools", [])
124
+ if custom_tools:
125
+ self._register_custom_tools(custom_tools)
126
+
117
127
  # MCP integration (filesystem MCP server may have been injected by base class)
118
128
  self.mcp_servers = self.config.get("mcp_servers", [])
119
129
  self.allowed_tools = kwargs.pop("allowed_tools", None)
@@ -169,6 +179,261 @@ class MCPBackend(LLMBackend):
169
179
  async def _process_stream(self, stream, all_params, agent_id: Optional[str] = None) -> AsyncGenerator[StreamChunk, None]:
170
180
  """Process stream."""
171
181
 
182
+ # Custom tools support
183
+ def _register_custom_tools(self, custom_tools: List[Dict[str, Any]]) -> None:
184
+ """Register custom tools with the tool manager.
185
+
186
+ Supports flexible configuration:
187
+ - function: str | List[str]
188
+ - description: str (shared) | List[str] (1-to-1 mapping)
189
+ - preset_args: dict (shared) | List[dict] (1-to-1 mapping)
190
+
191
+ Examples:
192
+ # Single function
193
+ function: "my_func"
194
+ description: "My description"
195
+
196
+ # Multiple functions with shared description
197
+ function: ["func1", "func2"]
198
+ description: "Shared description"
199
+
200
+ # Multiple functions with individual descriptions
201
+ function: ["func1", "func2"]
202
+ description: ["Description 1", "Description 2"]
203
+
204
+ # Multiple functions with mixed (shared desc, individual args)
205
+ function: ["func1", "func2"]
206
+ description: "Shared description"
207
+ preset_args: [{"arg1": "val1"}, {"arg1": "val2"}]
208
+
209
+ Args:
210
+ custom_tools: List of custom tool configurations
211
+ """
212
+ # Collect unique categories and create them if needed
213
+ categories = set()
214
+ for tool_config in custom_tools:
215
+ if isinstance(tool_config, dict):
216
+ category = tool_config.get("category", "default")
217
+ if category != "default":
218
+ categories.add(category)
219
+
220
+ # Create categories that don't exist
221
+ for category in categories:
222
+ if category not in self.custom_tool_manager.tool_categories:
223
+ self.custom_tool_manager.setup_category(
224
+ category_name=category,
225
+ description=f"Custom {category} tools",
226
+ enabled=True,
227
+ )
228
+
229
+ # Register each custom tool
230
+ for tool_config in custom_tools:
231
+ try:
232
+ if isinstance(tool_config, dict):
233
+ # Extract base configuration
234
+ path = tool_config.get("path")
235
+ category = tool_config.get("category", "default")
236
+
237
+ # Normalize function field to list
238
+ func_field = tool_config.get("function")
239
+ if isinstance(func_field, str):
240
+ functions = [func_field]
241
+ elif isinstance(func_field, list):
242
+ functions = func_field
243
+ else:
244
+ logger.error(
245
+ f"Invalid function field type: {type(func_field)}. " f"Must be str or List[str].",
246
+ )
247
+ continue
248
+
249
+ if not functions:
250
+ logger.error("Empty function list in tool config")
251
+ continue
252
+
253
+ num_functions = len(functions)
254
+
255
+ # Process name field (can be str or List[str])
256
+ name_field = tool_config.get("name")
257
+ names = self._process_field_for_functions(
258
+ name_field,
259
+ num_functions,
260
+ "name",
261
+ )
262
+ if names is None:
263
+ continue # Validation error, skip this tool
264
+
265
+ # Process description field (can be str or List[str])
266
+ desc_field = tool_config.get("description")
267
+ descriptions = self._process_field_for_functions(
268
+ desc_field,
269
+ num_functions,
270
+ "description",
271
+ )
272
+ if descriptions is None:
273
+ continue # Validation error, skip this tool
274
+
275
+ # Process preset_args field (can be dict or List[dict])
276
+ preset_field = tool_config.get("preset_args")
277
+ preset_args_list = self._process_field_for_functions(
278
+ preset_field,
279
+ num_functions,
280
+ "preset_args",
281
+ )
282
+ if preset_args_list is None:
283
+ continue # Validation error, skip this tool
284
+
285
+ # Register each function with its corresponding values
286
+ for i, func in enumerate(functions):
287
+ # Load the function first if custom name is needed
288
+ if names[i] and names[i] != func:
289
+ # Need to load function and apply custom name
290
+ if path:
291
+ loaded_func = self.custom_tool_manager._load_function_from_path(path, func)
292
+ else:
293
+ loaded_func = self.custom_tool_manager._load_builtin_function(func)
294
+
295
+ if loaded_func is None:
296
+ logger.error(f"Could not load function '{func}' from path: {path}")
297
+ continue
298
+
299
+ # Apply custom name by modifying __name__ attribute
300
+ loaded_func.__name__ = names[i]
301
+
302
+ # Register with loaded function (no path needed)
303
+ self.custom_tool_manager.add_tool_function(
304
+ path=None,
305
+ func=loaded_func,
306
+ category=category,
307
+ preset_args=preset_args_list[i],
308
+ description=descriptions[i],
309
+ )
310
+ else:
311
+ # No custom name or same as function name, use normal registration
312
+ self.custom_tool_manager.add_tool_function(
313
+ path=path,
314
+ func=func,
315
+ category=category,
316
+ preset_args=preset_args_list[i],
317
+ description=descriptions[i],
318
+ )
319
+
320
+ # Use custom name for logging and tracking if provided
321
+ registered_name = names[i] if names[i] else func
322
+
323
+ # Track tool name for categorization
324
+ if registered_name.startswith("custom_tool__"):
325
+ self._custom_tool_names.add(registered_name)
326
+ else:
327
+ self._custom_tool_names.add(f"custom_tool__{registered_name}")
328
+
329
+ logger.info(
330
+ f"Registered custom tool: {registered_name} from {path} " f"(category: {category}, " f"desc: '{descriptions[i][:50] if descriptions[i] else 'None'}...')",
331
+ )
332
+
333
+ except Exception as e:
334
+ func_name = tool_config.get("function", "unknown")
335
+ logger.error(
336
+ f"Failed to register custom tool {func_name}: {e}",
337
+ exc_info=True,
338
+ )
339
+
340
+ def _process_field_for_functions(
341
+ self,
342
+ field_value: Any,
343
+ num_functions: int,
344
+ field_name: str,
345
+ ) -> Optional[List[Any]]:
346
+ """Process a config field that can be a single value or list.
347
+
348
+ Conversion rules:
349
+ - None → [None, None, ...] (repeated num_functions times)
350
+ - Single value (not list) → [value, value, ...] (shared)
351
+ - List with matching length → use as-is (1-to-1 mapping)
352
+ - List with wrong length → ERROR (return None)
353
+
354
+ Args:
355
+ field_value: The field value from config
356
+ num_functions: Number of functions being registered
357
+ field_name: Name of the field (for error messages)
358
+
359
+ Returns:
360
+ List of values (one per function), or None if validation fails
361
+
362
+ Examples:
363
+ _process_field_for_functions(None, 3, "desc")
364
+ → [None, None, None]
365
+
366
+ _process_field_for_functions("shared", 3, "desc")
367
+ → ["shared", "shared", "shared"]
368
+
369
+ _process_field_for_functions(["a", "b", "c"], 3, "desc")
370
+ → ["a", "b", "c"]
371
+
372
+ _process_field_for_functions(["a", "b"], 3, "desc")
373
+ → None (error logged)
374
+ """
375
+ # Case 1: None or missing field → use None for all functions
376
+ if field_value is None:
377
+ return [None] * num_functions
378
+
379
+ # Case 2: Single value (not a list) → share across all functions
380
+ if not isinstance(field_value, list):
381
+ return [field_value] * num_functions
382
+
383
+ # Case 3: List value → must match function count exactly
384
+ if len(field_value) == num_functions:
385
+ return field_value
386
+ else:
387
+ # Length mismatch → validation error
388
+ logger.error(
389
+ f"Configuration error: {field_name} is a list with "
390
+ f"{len(field_value)} items, but there are {num_functions} functions. "
391
+ f"Either use a single value (shared) or a list with exactly "
392
+ f"{num_functions} items (1-to-1 mapping).",
393
+ )
394
+ return None
395
+
396
+ async def _execute_custom_tool(self, call: Dict[str, Any]) -> str:
397
+ """Execute a custom tool and return the result.
398
+
399
+ Args:
400
+ call: Function call dictionary with name and arguments
401
+
402
+ Returns:
403
+ The execution result as a string
404
+ """
405
+ import json
406
+
407
+ tool_request = {
408
+ "name": call["name"],
409
+ "input": json.loads(call["arguments"]) if isinstance(call["arguments"], str) else call["arguments"],
410
+ }
411
+
412
+ result_text = ""
413
+ try:
414
+ async for result in self.custom_tool_manager.execute_tool(tool_request):
415
+ # Accumulate results
416
+ if hasattr(result, "output_blocks"):
417
+ for block in result.output_blocks:
418
+ if hasattr(block, "data"):
419
+ result_text += str(block.data)
420
+ elif hasattr(block, "content"):
421
+ result_text += str(block.content)
422
+ elif hasattr(result, "content"):
423
+ result_text += str(result.content)
424
+ else:
425
+ result_text += str(result)
426
+ except Exception as e:
427
+ logger.error(f"Error in custom tool execution: {e}")
428
+ result_text = f"Error: {str(e)}"
429
+
430
+ return result_text or "Tool executed successfully"
431
+
432
+ def _get_custom_tools_schemas(self) -> List[Dict[str, Any]]:
433
+ """Get OpenAI-formatted schemas for all registered custom tools."""
434
+ return self.custom_tool_manager.fetch_tool_schemas()
435
+
436
+ # MCP support methods
172
437
  async def _setup_mcp_tools(self) -> None:
173
438
  """Initialize MCP client for mcp_tools-based servers (stdio + streamable-http)."""
174
439
  if not self.mcp_servers or self._mcp_initialized:
@@ -783,14 +1048,16 @@ class MCPBackend(LLMBackend):
783
1048
  async for chunk in self.yield_mcp_status_chunks(use_mcp):
784
1049
  yield chunk
785
1050
 
786
- if use_mcp:
1051
+ use_custom_tools = bool(self._custom_tool_names)
1052
+
1053
+ if use_mcp or use_custom_tools:
787
1054
  # MCP MODE: Recursive function call detection and execution
788
1055
  logger.info("Using recursive MCP execution mode")
789
1056
 
790
1057
  current_messages = self._trim_message_history(messages.copy())
791
1058
 
792
1059
  # Start recursive MCP streaming
793
- async for chunk in self._stream_with_mcp_tools(current_messages, tools, client, **kwargs):
1060
+ async for chunk in self._stream_with_custom_and_mcp_tools(current_messages, tools, client, **kwargs):
794
1061
  yield chunk
795
1062
 
796
1063
  else:
@@ -798,7 +1065,7 @@ class MCPBackend(LLMBackend):
798
1065
  logger.info("Using no-MCP mode")
799
1066
 
800
1067
  # Start non-MCP streaming
801
- async for chunk in self._stream_without_mcp_tools(messages, tools, client, **kwargs):
1068
+ async for chunk in self._stream_without_custom_and_mcp_tools(messages, tools, client, **kwargs):
802
1069
  yield chunk
803
1070
 
804
1071
  except Exception as e:
@@ -808,7 +1075,7 @@ class MCPBackend(LLMBackend):
808
1075
  await self._record_mcp_circuit_breaker_failure(e, agent_id)
809
1076
 
810
1077
  # Handle MCP exceptions with fallback
811
- async for chunk in self._stream_handle_mcp_exceptions(e, messages, tools, client, **kwargs):
1078
+ async for chunk in self._stream_handle_custom_and_mcp_exceptions(e, messages, tools, client, **kwargs):
812
1079
  yield chunk
813
1080
  else:
814
1081
  logger.error(f"Streaming error: {e}")
@@ -824,7 +1091,7 @@ class MCPBackend(LLMBackend):
824
1091
 
825
1092
  if isinstance(e, (MCPConnectionError, MCPTimeoutError, MCPServerError, MCPError)):
826
1093
  # Handle MCP exceptions with fallback
827
- async for chunk in self._stream_handle_mcp_exceptions(e, messages, tools, client, **kwargs):
1094
+ async for chunk in self._stream_handle_custom_and_mcp_exceptions(e, messages, tools, client, **kwargs):
828
1095
  yield chunk
829
1096
  else:
830
1097
  # Generic setup error: still notify if MCP was configured
@@ -837,7 +1104,7 @@ class MCPBackend(LLMBackend):
837
1104
  )
838
1105
 
839
1106
  # Proceed with non-MCP streaming
840
- async for chunk in self._stream_without_mcp_tools(messages, tools, client, **kwargs):
1107
+ async for chunk in self._stream_without_custom_and_mcp_tools(messages, tools, client, **kwargs):
841
1108
  yield chunk
842
1109
  except Exception as inner_e:
843
1110
  logger.error(f"Streaming error during MCP setup fallback: {inner_e}")
@@ -845,7 +1112,7 @@ class MCPBackend(LLMBackend):
845
1112
  finally:
846
1113
  await self._cleanup_client(client)
847
1114
 
848
- async def _stream_without_mcp_tools(
1115
+ async def _stream_without_custom_and_mcp_tools(
849
1116
  self,
850
1117
  messages: List[Dict[str, Any]],
851
1118
  tools: List[Dict[str, Any]],
@@ -885,7 +1152,7 @@ class MCPBackend(LLMBackend):
885
1152
  async for chunk in self._process_stream(stream, all_params, agent_id):
886
1153
  yield chunk
887
1154
 
888
- async def _stream_handle_mcp_exceptions(
1155
+ async def _stream_handle_custom_and_mcp_exceptions(
889
1156
  self,
890
1157
  error: Exception,
891
1158
  messages: List[Dict[str, Any]],
@@ -921,7 +1188,7 @@ class MCPBackend(LLMBackend):
921
1188
  content=f"\n⚠️ {user_message} ({error}); continuing without MCP tools\n",
922
1189
  )
923
1190
 
924
- async for chunk in self._stream_without_mcp_tools(messages, tools, client, **kwargs):
1191
+ async for chunk in self._stream_without_custom_and_mcp_tools(messages, tools, client, **kwargs):
925
1192
  yield chunk
926
1193
 
927
1194
  def _track_mcp_function_names(self, tools: List[Dict[str, Any]]) -> None:
@@ -1012,7 +1279,7 @@ class MCPBackend(LLMBackend):
1012
1279
  self._mcp_functions.clear()
1013
1280
  self._mcp_function_names.clear()
1014
1281
 
1015
- async def __aenter__(self) -> "MCPBackend":
1282
+ async def __aenter__(self) -> "CustomToolAndMCPBackend":
1016
1283
  """Async context manager entry."""
1017
1284
  # Initialize MCP tools if configured
1018
1285
  if MCPResourceManager:
@@ -1087,6 +1354,10 @@ class MCPBackend(LLMBackend):
1087
1354
  """Check if a tool call is an MCP function."""
1088
1355
  return tool_name in self._mcp_functions
1089
1356
 
1357
+ def is_custom_tool_call(self, tool_name: str) -> bool:
1358
+ """Check if a tool call is a custom tool function."""
1359
+ return tool_name in self._custom_tool_names
1360
+
1090
1361
  def get_mcp_tools_formatted(self) -> List[Dict[str, Any]]:
1091
1362
  """Get MCP tools formatted for specific API format."""
1092
1363
  if not self._mcp_functions: