massgen 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

massgen/backend/grok.py

@@ -18,6 +18,7 @@ TODO for future releases:
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
+import logging
 import os
 from typing import Any, Dict, List, Optional
 
@@ -26,6 +27,8 @@ from openai import AsyncOpenAI
 from ..logger_config import log_stream_chunk
 from .chat_completions import ChatCompletionsBackend
 
+logger = logging.getLogger(__name__)
+
 
 class GrokBackend(ChatCompletionsBackend):
     """Grok backend using xAI's OpenAI-compatible API."""
@@ -41,12 +44,10 @@ class GrokBackend(ChatCompletionsBackend):
 
         return openai.AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
 
-    def _build_base_api_params(self, messages: List[Dict[str, Any]], all_params: Dict[str, Any]) -> Dict[str, Any]:
-        """Build base API params for xAI's Grok API."""
-        api_params = super()._build_base_api_params(messages, all_params)
-
-        # Add Live Search parameters if enabled (Grok-specific)
+    def _add_grok_search_params(self, api_params: Dict[str, Any], all_params: Dict[str, Any]) -> Dict[str, Any]:
+        """Add Grok Live Search parameters to API params if web search is enabled."""
         enable_web_search = all_params.get("enable_web_search", False)
+
         if enable_web_search:
             # Check for conflict with manually specified search_parameters
             existing_extra = api_params.get("extra_body", {})
@@ -54,14 +55,46 @@ class GrokBackend(ChatCompletionsBackend):
                 error_message = "Conflict: Cannot use both 'enable_web_search: true' and manual 'extra_body.search_parameters'. Use one or the other."
                 log_stream_chunk("backend.grok", "error", error_message, self.agent_id)
                 raise ValueError(error_message)
+
             # Merge search_parameters into existing extra_body
             search_params = {"mode": "auto", "return_citations": True}
-            merged_extra = existing_extra.copy()
+            merged_extra = existing_extra.copy() if existing_extra else {}
             merged_extra["search_parameters"] = search_params
             api_params["extra_body"] = merged_extra
 
         return api_params
 
+    async def _stream_with_custom_and_mcp_tools(
+        self,
+        current_messages: List[Dict[str, Any]],
+        tools: List[Dict[str, Any]],
+        client,
+        **kwargs,
+    ):
+        """Override to add Grok-specific search parameters before API call."""
+        # Build API params using parent method
+        all_params = {**self.config, **kwargs}
+        api_params = await self.api_params_handler.build_api_params(current_messages, tools, all_params)
+
+        # Add provider tools (web search, code interpreter) if enabled
+        # Note: For Grok, get_provider_tools() won't add web_search function tool
+        provider_tools = self.api_params_handler.get_provider_tools(all_params)
+
+        if provider_tools:
+            if "tools" not in api_params:
+                api_params["tools"] = []
+            api_params["tools"].extend(provider_tools)
+
+        # Add Grok-specific web search parameters via extra_body
+        api_params = self._add_grok_search_params(api_params, all_params)
+
+        # Start streaming
+        stream = await client.chat.completions.create(**api_params)
+
+        # Delegate to parent's stream processing
+        async for chunk in super()._process_stream(stream, all_params, self.agent_id):
+            yield chunk
+
     def get_provider_name(self) -> str:
         """Get the name of this provider."""
         return "Grok"

massgen/backend/response.py

@@ -356,18 +356,25 @@ class ResponseBackend(CustomToolAndMCPBackend):
             # Execute MCP function calls
             mcp_functions_executed = False
 
-            # Check if planning mode is enabled - block MCP tool execution during planning
+            # Check if planning mode is enabled - selectively block MCP tool execution during planning
             if self.is_planning_mode_enabled():
-                logger.info("[MCP] Planning mode enabled - blocking all MCP tool execution")
-                yield StreamChunk(
-                    type="mcp_status",
-                    status="planning_mode_blocked",
-                    content="🚫 [MCP] Planning mode active - MCP tools blocked during coordination",
-                    source="planning_mode",
-                )
-                # Skip all MCP tool execution but still continue with workflow
-                yield StreamChunk(type="done")
-                return
+                blocked_tools = self.get_planning_mode_blocked_tools()
+
+                if not blocked_tools:
+                    # Empty set means block ALL MCP tools (backward compatible)
+                    logger.info("[Response] Planning mode enabled - blocking ALL MCP tool execution")
+                    yield StreamChunk(
+                        type="mcp_status",
+                        status="planning_mode_blocked",
+                        content="🚫 [MCP] Planning mode active - all MCP tools blocked during coordination",
+                        source="planning_mode",
+                    )
+                    # Skip all MCP tool execution but still continue with workflow
+                    yield StreamChunk(type="done")
+                    return
+                else:
+                    # Selective blocking - log but continue to check each tool individually
+                    logger.info(f"[Response] Planning mode enabled - selective blocking of {len(blocked_tools)} tools")
 
             # Ensure every captured function call gets a result to prevent hanging
             for call in captured_function_calls:

massgen/chat_agent.py

@@ -365,10 +365,15 @@ class ConfigurableAgent(SingleAgent):
             backend: LLM backend
             session_id: Optional session identifier
         """
+        # Extract system message without triggering deprecation warning
+        system_message = None
+        if hasattr(config, "_custom_system_instruction"):
+            system_message = config._custom_system_instruction
+
         super().__init__(
             backend=backend,
             agent_id=config.agent_id,
-            system_message=config.custom_system_instruction,
+            system_message=system_message,
             session_id=session_id,
         )
         self.config = config
@@ -411,8 +416,9 @@ class ConfigurableAgent(SingleAgent):
                 return backend_params["append_system_prompt"]
 
         # Fall back to custom_system_instruction (deprecated but still supported)
-        if self.config and self.config.custom_system_instruction:
-            return self.config.custom_system_instruction
+        # Access private attribute directly to avoid deprecation warning
+        if self.config and hasattr(self.config, "_custom_system_instruction") and self.config._custom_system_instruction:
+            return self.config._custom_system_instruction
 
         # Finally fall back to parent class implementation
         return super().get_configurable_system_message()