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.
- massgen/__init__.py +1 -1
- massgen/agent_config.py +17 -0
- massgen/api_params_handler/_api_params_handler_base.py +1 -0
- massgen/api_params_handler/_chat_completions_api_params_handler.py +8 -1
- massgen/api_params_handler/_claude_api_params_handler.py +8 -1
- massgen/api_params_handler/_gemini_api_params_handler.py +73 -0
- massgen/api_params_handler/_response_api_params_handler.py +8 -1
- massgen/backend/base.py +31 -0
- massgen/backend/{base_with_mcp.py → base_with_custom_tool_and_mcp.py} +282 -11
- massgen/backend/chat_completions.py +182 -92
- massgen/backend/claude.py +115 -18
- massgen/backend/claude_code.py +378 -14
- massgen/backend/docs/CLAUDE_API_RESEARCH.md +3 -3
- massgen/backend/gemini.py +1275 -1607
- massgen/backend/gemini_mcp_manager.py +545 -0
- massgen/backend/gemini_trackers.py +344 -0
- massgen/backend/gemini_utils.py +43 -0
- massgen/backend/response.py +129 -70
- massgen/cli.py +577 -110
- massgen/config_builder.py +376 -27
- massgen/configs/README.md +111 -80
- massgen/configs/basic/multi/three_agents_default.yaml +1 -1
- massgen/configs/basic/single/single_agent.yaml +1 -1
- massgen/configs/providers/openai/gpt5_nano.yaml +3 -3
- massgen/configs/tools/custom_tools/claude_code_custom_tool_example.yaml +32 -0
- massgen/configs/tools/custom_tools/claude_code_custom_tool_example_no_path.yaml +28 -0
- massgen/configs/tools/custom_tools/claude_code_custom_tool_with_mcp_example.yaml +40 -0
- massgen/configs/tools/custom_tools/claude_code_custom_tool_with_wrong_mcp_example.yaml +38 -0
- massgen/configs/tools/custom_tools/claude_code_wrong_custom_tool_with_mcp_example.yaml +38 -0
- massgen/configs/tools/custom_tools/claude_custom_tool_example.yaml +24 -0
- massgen/configs/tools/custom_tools/claude_custom_tool_example_no_path.yaml +22 -0
- massgen/configs/tools/custom_tools/claude_custom_tool_with_mcp_example.yaml +35 -0
- massgen/configs/tools/custom_tools/claude_custom_tool_with_wrong_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/claude_wrong_custom_tool_with_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/gemini_custom_tool_example.yaml +24 -0
- massgen/configs/tools/custom_tools/gemini_custom_tool_example_no_path.yaml +22 -0
- massgen/configs/tools/custom_tools/gemini_custom_tool_with_mcp_example.yaml +35 -0
- massgen/configs/tools/custom_tools/gemini_custom_tool_with_wrong_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/gemini_wrong_custom_tool_with_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/github_issue_market_analysis.yaml +94 -0
- massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_example.yaml +24 -0
- massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_example_no_path.yaml +22 -0
- massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_with_mcp_example.yaml +35 -0
- massgen/configs/tools/custom_tools/gpt5_nano_custom_tool_with_wrong_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/gpt5_nano_wrong_custom_tool_with_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/gpt_oss_custom_tool_example.yaml +25 -0
- massgen/configs/tools/custom_tools/gpt_oss_custom_tool_example_no_path.yaml +23 -0
- massgen/configs/tools/custom_tools/gpt_oss_custom_tool_with_mcp_example.yaml +34 -0
- massgen/configs/tools/custom_tools/gpt_oss_custom_tool_with_wrong_mcp_example.yaml +34 -0
- massgen/configs/tools/custom_tools/gpt_oss_wrong_custom_tool_with_mcp_example.yaml +34 -0
- massgen/configs/tools/custom_tools/grok3_mini_custom_tool_example.yaml +24 -0
- massgen/configs/tools/custom_tools/grok3_mini_custom_tool_example_no_path.yaml +22 -0
- massgen/configs/tools/custom_tools/grok3_mini_custom_tool_with_mcp_example.yaml +35 -0
- massgen/configs/tools/custom_tools/grok3_mini_custom_tool_with_wrong_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/grok3_mini_wrong_custom_tool_with_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/qwen_api_custom_tool_example.yaml +25 -0
- massgen/configs/tools/custom_tools/qwen_api_custom_tool_example_no_path.yaml +23 -0
- massgen/configs/tools/custom_tools/qwen_api_custom_tool_with_mcp_example.yaml +36 -0
- massgen/configs/tools/custom_tools/qwen_api_custom_tool_with_wrong_mcp_example.yaml +34 -0
- massgen/configs/tools/custom_tools/qwen_api_wrong_custom_tool_with_mcp_example.yaml +34 -0
- massgen/configs/tools/custom_tools/qwen_local_custom_tool_example.yaml +24 -0
- massgen/configs/tools/custom_tools/qwen_local_custom_tool_example_no_path.yaml +22 -0
- massgen/configs/tools/custom_tools/qwen_local_custom_tool_with_mcp_example.yaml +35 -0
- massgen/configs/tools/custom_tools/qwen_local_custom_tool_with_wrong_mcp_example.yaml +33 -0
- massgen/configs/tools/custom_tools/qwen_local_wrong_custom_tool_with_mcp_example.yaml +33 -0
- massgen/configs/tools/filesystem/claude_code_context_sharing.yaml +1 -1
- massgen/configs/voting/gemini_gpt_voting_sensitivity.yaml +67 -0
- massgen/formatter/_chat_completions_formatter.py +104 -0
- massgen/formatter/_claude_formatter.py +120 -0
- massgen/formatter/_gemini_formatter.py +448 -0
- massgen/formatter/_response_formatter.py +88 -0
- massgen/frontend/coordination_ui.py +4 -2
- massgen/logger_config.py +35 -3
- massgen/message_templates.py +56 -6
- massgen/orchestrator.py +179 -10
- massgen/stream_chunk/base.py +3 -0
- massgen/tests/custom_tools_example.py +392 -0
- massgen/tests/mcp_test_server.py +17 -7
- massgen/tests/test_config_builder.py +423 -0
- massgen/tests/test_custom_tools.py +401 -0
- massgen/tests/test_tools.py +127 -0
- massgen/tool/README.md +935 -0
- massgen/tool/__init__.py +39 -0
- massgen/tool/_async_helpers.py +70 -0
- massgen/tool/_basic/__init__.py +8 -0
- massgen/tool/_basic/_two_num_tool.py +24 -0
- massgen/tool/_code_executors/__init__.py +10 -0
- massgen/tool/_code_executors/_python_executor.py +74 -0
- massgen/tool/_code_executors/_shell_executor.py +61 -0
- massgen/tool/_exceptions.py +39 -0
- massgen/tool/_file_handlers/__init__.py +10 -0
- massgen/tool/_file_handlers/_file_operations.py +218 -0
- massgen/tool/_manager.py +634 -0
- massgen/tool/_registered_tool.py +88 -0
- massgen/tool/_result.py +66 -0
- massgen/tool/_self_evolution/_github_issue_analyzer.py +369 -0
- massgen/tool/docs/builtin_tools.md +681 -0
- massgen/tool/docs/exceptions.md +794 -0
- massgen/tool/docs/execution_results.md +691 -0
- massgen/tool/docs/manager.md +887 -0
- massgen/tool/docs/workflow_toolkits.md +529 -0
- massgen/tool/workflow_toolkits/__init__.py +57 -0
- massgen/tool/workflow_toolkits/base.py +55 -0
- massgen/tool/workflow_toolkits/new_answer.py +126 -0
- massgen/tool/workflow_toolkits/vote.py +167 -0
- {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/METADATA +89 -131
- {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/RECORD +111 -36
- {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/WHEEL +0 -0
- {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/entry_points.txt +0 -0
- {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/licenses/LICENSE +0 -0
- {massgen-0.1.0a3.dist-info → massgen-0.1.1.dist-info}/top_level.txt +0 -0
massgen/__init__.py
CHANGED
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
|
)
|
|
@@ -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
|
-
#
|
|
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
|
-
#
|
|
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
|
|
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
|
|
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
|
-
|
|
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.
|
|
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.
|
|
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.
|
|
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.
|
|
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.
|
|
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
|
|
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
|
|
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.
|
|
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) -> "
|
|
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:
|