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

massgen/backend/grok.py

@@ -1,26 +1,30 @@
-from __future__ import annotations
-
+# -*- coding: utf-8 -*-
 """
-Grok/xAI backend implementation using OpenAI-compatible API.
-Clean implementation with only Grok-specific features.
+Grok/xAI backend is using the chat_completions backend for streaming.
+It overrides methods for Grok-specific features (Grok Live Search).
 
 ✅ TESTED: Backend works correctly with architecture
-- ✅ Grok API integration working
-- ✅ Tool message conversion compatible with Chat Completions format
-- ✅ Streaming functionality working correctly  
+- ✅ Grok API integration working (through chat_completions)
+- ✅ Streaming functionality working correctly
 - ✅ SingleAgent integration working
 - ✅ Error handling and pricing calculations implemented
+- ✅ Web search is working through Grok Live Search
+- ✅ MCP is working
 
 TODO for future releases:
 - Test multi-agent orchestrator integration
-- Test web search capabilities with tools
 - Validate advanced Grok-specific features
 """
+# -*- coding: utf-8 -*-
+from __future__ import annotations
 
 import os
-from typing import Dict, List, Any, AsyncGenerator, Optional
+from typing import Any, Dict, List, Optional
+
+from openai import AsyncOpenAI
+
+from ..logger_config import log_stream_chunk
 from .chat_completions import ChatCompletionsBackend
-from .base import StreamChunk
 
 
 class GrokBackend(ChatCompletionsBackend):
@@ -31,72 +35,32 @@ class GrokBackend(ChatCompletionsBackend):
         self.api_key = api_key or os.getenv("XAI_API_KEY")
         self.base_url = "https://api.x.ai/v1"
 
-    async def stream_with_tools(
-        self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs
-    ) -> AsyncGenerator[StreamChunk, None]:
-        """Stream response using xAI's OpenAI-compatible API."""
-
-        # Convert messages for Grok API compatibility
-        grok_messages = self._convert_messages_for_grok(messages)
-
-        try:
-            import openai
-
-            # Use OpenAI client with xAI base URL
-            client = openai.AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
-
-            # Extract parameters
-            model = kwargs.get("model", "grok-3-mini")
-            max_tokens = kwargs.get("max_tokens", None)
-            temperature = kwargs.get("temperature", None)
-            enable_web_search = kwargs.get("enable_web_search", False)
-
-            # Convert tools to Chat Completions format
-            converted_tools = (
-                self.convert_tools_to_chat_completions_format(tools) if tools else None
-            )
-
-            # Chat Completions API parameters
-            api_params = {
-                "model": model,
-                "messages": grok_messages,
-                "tools": converted_tools,
-                "max_tokens": max_tokens,
-                "temperature": temperature,
-                "stream": True,
-            }
-
-            # Add Live Search parameters if enabled (Grok-specific)
-            if enable_web_search:
-                search_params_kwargs = {"mode": "auto", "return_citations": True}
-
-                # Allow override of search parameters from backend params
-                max_results = kwargs.get("max_search_results")
-                if max_results is not None:
-                    search_params_kwargs["max_search_results"] = max_results
-
-                search_mode = kwargs.get("search_mode")
-                if search_mode is not None:
-                    search_params_kwargs["mode"] = search_mode
-
-                return_citations = kwargs.get("return_citations")
-                if return_citations is not None:
-                    search_params_kwargs["return_citations"] = return_citations
-
-                # Use extra_body to pass search_parameters to xAI API
-                api_params["extra_body"] = {"search_parameters": search_params_kwargs}
-
-            # Create stream
-            stream = await client.chat.completions.create(**api_params)
-
-            # Use base class streaming handler
-            async for chunk in self.handle_chat_completions_stream(
-                stream, enable_web_search
-            ):
-                yield chunk
-
-        except Exception as e:
-            yield StreamChunk(type="error", error=f"Grok API error: {e}")
+    def _create_client(self, **kwargs) -> AsyncOpenAI:
+        """Create OpenAI client configured for xAI's Grok API."""
+        import openai
+
+        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)
+        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", {})
+            if isinstance(existing_extra, dict) and "search_parameters" in existing_extra:
+                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["search_parameters"] = search_params
+            api_params["extra_body"] = merged_extra
+
+        return api_params
 
     def get_provider_name(self) -> str:
         """Get the name of this provider."""
@@ -105,83 +69,3 @@ class GrokBackend(ChatCompletionsBackend):
     def get_supported_builtin_tools(self) -> List[str]:
         """Get list of builtin tools supported by Grok."""
         return ["web_search"]
-
-    def estimate_tokens(self, text: str) -> int:
-        """Estimate token count for text (rough approximation)."""
-        return int(len(text.split()) * 1.3)
-
-    def calculate_cost(
-        self, input_tokens: int, output_tokens: int, model: str
-    ) -> float:
-        """Calculate cost for token usage."""
-        model_lower = model.lower()
-
-        # Handle -mini models with lower costs
-        if "grok-2" in model_lower:
-            if "mini" in model_lower:
-                input_cost = (input_tokens / 1_000_000) * 1.0  # Lower cost for mini
-                output_cost = (output_tokens / 1_000_000) * 5.0
-            else:
-                input_cost = (input_tokens / 1_000_000) * 2.0
-                output_cost = (output_tokens / 1_000_000) * 10.0
-        elif "grok-3" in model_lower:
-            if "mini" in model_lower:
-                input_cost = (input_tokens / 1_000_000) * 2.5  # Lower cost for mini
-                output_cost = (output_tokens / 1_000_000) * 7.5
-            else:
-                input_cost = (input_tokens / 1_000_000) * 5.0
-                output_cost = (output_tokens / 1_000_000) * 15.0
-        elif "grok-4" in model_lower:
-            if "mini" in model_lower:
-                input_cost = (input_tokens / 1_000_000) * 4.0  # Lower cost for mini
-                output_cost = (output_tokens / 1_000_000) * 10.0
-            else:
-                input_cost = (input_tokens / 1_000_000) * 8.0
-                output_cost = (output_tokens / 1_000_000) * 20.0
-        else:
-            # Default fallback (assume grok-3 pricing)
-            input_cost = (input_tokens / 1_000_000) * 5.0
-            output_cost = (output_tokens / 1_000_000) * 15.0
-
-        return input_cost + output_cost
-
-    def _convert_messages_for_grok(
-        self, messages: List[Dict[str, Any]]
-    ) -> List[Dict[str, Any]]:
-        """
-        Convert messages for Grok API compatibility.
-
-        Grok expects tool call arguments as JSON strings in conversation history,
-        but returns them as objects in responses.
-        """
-        import json
-
-        converted_messages = []
-
-        for message in messages:
-            # Create a copy to avoid modifying the original
-            converted_msg = dict(message)
-
-            # Convert tool_calls arguments from objects to JSON strings
-            if message.get("role") == "assistant" and "tool_calls" in message:
-                converted_tool_calls = []
-                for tool_call in message["tool_calls"]:
-                    converted_call = dict(tool_call)
-                    if "function" in converted_call:
-                        converted_function = dict(converted_call["function"])
-                        arguments = converted_function.get("arguments")
-
-                        # Convert arguments to JSON string if it's an object
-                        if isinstance(arguments, dict):
-                            converted_function["arguments"] = json.dumps(arguments)
-                        elif arguments is None:
-                            converted_function["arguments"] = "{}"
-                        # If it's already a string, keep it as-is
-
-                        converted_call["function"] = converted_function
-                    converted_tool_calls.append(converted_call)
-                converted_msg["tool_calls"] = converted_tool_calls
-
-            converted_messages.append(converted_msg)
-
-        return converted_messages

massgen/backend/inference.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+"""
+Inference backend supporting both vLLM and SGLang servers using OpenAI-compatible Chat Completions API.
+
+Defaults are tailored for local inference servers:
+- vLLM: base_url: http://localhost:8000/v1, api_key: "EMPTY"
+- SGLang: base_url: http://localhost:30000/v1, api_key: "EMPTY"
+
+This backend delegates most behavior to ChatCompletionsBackend, only
+customizing provider naming, API key resolution, and backend-specific extra_body parameters.
+"""
+from __future__ import annotations
+
+import os
+from typing import Any, AsyncGenerator, Dict, List, Optional, Set
+
+from ..api_params_handler._chat_completions_api_params_handler import (
+    ChatCompletionsAPIParamsHandler,
+)
+from .base import StreamChunk
+from .chat_completions import ChatCompletionsBackend
+
+
+class InferenceAPIParamsHandler(ChatCompletionsAPIParamsHandler):
+    """API params handler for InferenceBackend that excludes backend-specific parameters."""
+
+    def get_excluded_params(self) -> Set[str]:
+        """Get parameters to exclude from Chat Completions API calls, including backend-specific ones."""
+        return (
+            super()
+            .get_excluded_params()
+            .union(
+                {
+                    "chat_template_kwargs",
+                    "top_k",
+                    "repetition_penalty",
+                    "separate_reasoning",  # SGLang-specific parameter
+                },
+            )
+        )
+
+
+class InferenceBackend(ChatCompletionsBackend):
+    """Backend for local inference servers (vLLM and SGLang).
+
+    This backend connects to inference servers running with OpenAI-compatible API.
+    It supports both vLLM and SGLang specific parameters like guided generation,
+    thinking mode, and separate reasoning.
+    """
+
+    def __init__(self, backend_type: str = "vllm", api_key: Optional[str] = None, **kwargs):
+        """Initialize inference backend.
+
+        Args:
+            backend_type: Type of backend ("vllm" or "sglang")
+            api_key: API key (usually "EMPTY" for local servers)
+            **kwargs: Additional arguments passed to parent
+        """
+        self._backend_type = backend_type.lower()
+
+        # Set default base URLs based on backend type
+        if "base_url" not in kwargs:
+            if self._backend_type == "sglang":
+                kwargs["base_url"] = "http://localhost:30000/v1"
+            else:  # vllm
+                kwargs["base_url"] = "http://localhost:8000/v1"
+
+        # Determine API key based on backend type before calling parent
+        if api_key is None:
+            if self._backend_type == "sglang":
+                api_key = os.getenv("SGLANG_API_KEY") or "EMPTY"
+            else:  # vllm
+                api_key = os.getenv("VLLM_API_KEY") or "EMPTY"
+
+        # Initialize parent with the correct API key
+        super().__init__(api_key, **kwargs)
+
+        # Override the API params handler to exclude backend-specific parameters
+        self.api_params_handler = InferenceAPIParamsHandler(self)
+
+    def get_provider_name(self) -> str:
+        """Get the provider name for this backend."""
+        if self._backend_type == "sglang":
+            return "SGLang"
+        return "vLLM"
+
+    def _build_extra_body(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Build backend-specific extra_body parameters and strip them from kwargs.
+
+        Args:
+            kwargs: Keyword arguments that may contain backend parameters
+
+        Returns:
+            Dictionary of backend-specific parameters for extra_body
+        """
+        extra_body: Dict[str, Any] = {}
+
+        # Add vLLM and sglang specific parameters from kwargs while preventing them from reaching parent payload
+        top_k = kwargs.pop("top_k", None)
+        if top_k is not None:
+            extra_body["top_k"] = top_k
+
+        repetition_penalty = kwargs.pop("repetition_penalty", None)
+        if repetition_penalty is not None:
+            extra_body["repetition_penalty"] = repetition_penalty
+
+        # Unified chat template handling for both vLLM and SGLang , Some models different way to add it
+        chat_template_kwargs = kwargs.pop("chat_template_kwargs", None)
+        if chat_template_kwargs is not None:
+            extra_body["chat_template_kwargs"] = chat_template_kwargs
+
+        # SGLang-specific parameters handling for separate reasoning
+        if self._backend_type == "sglang":
+            separate_reasoning = kwargs.pop("separate_reasoning", None)
+            if separate_reasoning is not None:
+                extra_body["separate_reasoning"] = separate_reasoning
+
+        return extra_body
+
+    async def stream_with_tools(
+        self,
+        messages: List[Dict[str, Any]],
+        tools: List[Dict[str, Any]],
+        **kwargs,
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """Stream response using OpenAI-compatible Chat Completions API with backend-specific parameters.
+
+        Args:
+            messages: List of messages
+            tools: List of tool definitions
+            **kwargs: Additional parameters including backend-specific ones
+
+        Yields:
+            StreamChunk objects
+        """
+        # Build backend-specific extra_body parameters
+        extra_body = self._build_extra_body(kwargs)
+
+        # Add extra_body to kwargs if we have backend-specific parameters
+        if extra_body:
+            # Add to existing extra_body if present
+            if "extra_body" in kwargs:
+                kwargs["extra_body"].update(extra_body)
+            else:
+                kwargs["extra_body"] = extra_body
+
+        # Delegate to parent with backend-specific parameters in extra_body
+        async for chunk in super().stream_with_tools(messages, tools, **kwargs):
+            yield chunk
+
+    def get_supported_builtin_tools(self) -> List[str]:
+        """Return list of supported builtin tools.
+
+        Local inference servers (vLLM/SGLang) do not provide provider-specific builtin tools.
+        """
+        return []

massgen/backend/lmstudio.py

@@ -0,0 +1,171 @@
+# -*- coding: utf-8 -*-
+"""
+LM Studio backend using an OpenAI-compatible Chat Completions API.
+
+Defaults are tailored for a local LM Studio server:
+- base_url: http://localhost:1234/v1
+- api_key:  "lm-studio" (LM Studio accepts any non-empty key)
+
+This backend delegates most behavior to ChatCompletionsBackend, only
+customizing provider naming, API key resolution, and cost calculation.
+"""
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+import platform
+import shutil
+import subprocess
+import time
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+import lmstudio as lms
+
+from .base import StreamChunk
+from .chat_completions import ChatCompletionsBackend
+
+
+class LMStudioBackend(ChatCompletionsBackend):
+    """LM Studio backend (OpenAI-compatible, local server)."""
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        super().__init__(api_key="lm-studio", **kwargs)  # Override to avoid environment-variable enforcement; LM Studio accepts any key
+        self._models_attempted = set()  # Track models this instance has attempted to load
+        self.start_lmstudio_server(**kwargs)
+
+    async def stream_with_tools(self, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]], **kwargs) -> AsyncGenerator[StreamChunk, None]:
+        """Stream response using OpenAI-compatible Chat Completions API.
+
+        LM Studio does not require special message conversions; this delegates to
+        the generic ChatCompletions implementation while preserving our defaults.
+        """
+
+        # Ensure LM Studio defaults
+        base_url = kwargs.get("base_url", "http://localhost:1234/v1")
+        kwargs["base_url"] = base_url
+
+        async for chunk in super().stream_with_tools(messages, tools, **kwargs):
+            yield chunk
+
+        # self.end_lmstudio_server()
+
+    def get_supported_builtin_tools(self) -> List[str]:  # type: ignore[override]
+        # LM Studio (local OpenAI-compatible) does not provide provider-builtins
+        return []
+
+    def start_lmstudio_server(self, **kwargs):
+        """Start LM Studio server after checking CLI and model availability."""
+        self._ensure_cli_installed()
+        self._start_server()
+        model_name = kwargs.get("model", "")
+        if model_name:
+            self._handle_model(model_name)
+
+    def _ensure_cli_installed(self):
+        """Ensure LM Studio CLI is installed."""
+        if shutil.which("lms"):
+            return
+        print("LM Studio CLI not found. Installing...")
+        try:
+            system = platform.system().lower()
+            install_commands = {
+                "darwin": (["brew", "install", "lmstudio"], False),
+                "linux": (["curl", "-sSL", "https://lmstudio.ai/install.sh", "|", "sh"], True),
+                "windows": (["powershell", "-Command", "iwr -useb https://lmstudio.ai/install.ps1 | iex"], False),
+            }
+            if system not in install_commands:
+                raise RuntimeError(f"Unsupported platform: {system}")
+            cmd, use_shell = install_commands[system]
+            subprocess.run(cmd, shell=use_shell, check=True)
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(f"Failed to install LM Studio CLI: {e}") from e
+
+    def _start_server(self):
+        """Start the LM Studio server in background mode."""
+        try:
+            with subprocess.Popen(
+                ["lms", "server", "start"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            ) as process:
+                time.sleep(3)
+                if process.poll() is None:
+                    print("LM Studio server started successfully (running in background).")
+                else:
+                    self._handle_server_output(process)
+        except Exception as e:
+            raise RuntimeError(f"Failed to start LM Studio server: {e}") from e
+
+    def _handle_server_output(self, process):
+        """Handle server process output."""
+        stdout, stderr = process.communicate(timeout=1)
+        if stdout:
+            print(f"Server output: {stdout}")
+        if stderr:
+            self._process_stderr(stderr)
+        print("LM Studio server started successfully.")
+
+    def _process_stderr(self, stderr):
+        """Process server stderr output."""
+        stderr_lower = stderr.lower()
+        if "success" in stderr_lower or "running on port" in stderr_lower:
+            print(f"Server info: {stderr.strip()}")
+        elif "warning" in stderr_lower or "warn" in stderr_lower:
+            print(f"Server warning: {stderr.strip()}")
+        else:
+            print(f"Server error: {stderr.strip()}")
+
+    def _handle_model(self, model_name):
+        """Handle model downloading and loading."""
+        self._ensure_model_downloaded(model_name)
+        self._load_model_if_needed(model_name)
+
+    def _ensure_model_downloaded(self, model_name):
+        """Ensure model is downloaded locally."""
+        try:
+            downloaded = lms.list_downloaded_models()
+            model_keys = [m.model_key for m in downloaded]
+            if model_name not in model_keys:
+                print(f"Model '{model_name}' not found locally. Downloading...")
+                subprocess.run(["lms", "get", model_name], check=True)
+                print(f"Model '{model_name}' downloaded successfully.")
+        except Exception as e:
+            print(f"Warning: Could not check/download model: {e}")
+
+    def _load_model_if_needed(self, model_name):
+        """Load model if not already loaded."""
+        try:
+            if model_name in self._models_attempted:
+                print(f"Model '{model_name}' load already attempted by this instance.")
+                return
+
+            time.sleep(5)
+            loaded = lms.list_loaded_models()
+            loaded_identifiers = [m.identifier for m in loaded]
+            if model_name not in loaded_identifiers:
+                print(f"Model '{model_name}' not loaded. Loading...")
+                self._models_attempted.add(model_name)
+                subprocess.run(["lms", "load", model_name], check=True)
+                print(f"Model '{model_name}' loaded successfully.")
+            else:
+                print(f"Model '{model_name}' is already loaded.")
+                self._models_attempted.add(model_name)
+        except subprocess.CalledProcessError as e:
+            print(f"Warning: Failed to load model '{model_name}': {e}")
+        except Exception as e:
+            print(f"Warning: Could not check loaded models: {e}")
+
+    def end_lmstudio_server(self):
+        """Stop the LM Studio server after receiving all chunks."""
+        try:
+            # Use lms server end command as specified in requirement
+            result = subprocess.run(["lms", "server", "end"], capture_output=True, text=True, check=False)
+
+            if result.returncode == 0:
+                print("LM Studio server ended successfully.")
+            else:
+                # Fallback to stop command if end doesn't work
+                subprocess.run(["lms", "server", "stop"], check=True)
+                print("LM Studio server stopped successfully.")
+        except Exception as e:
+            print(f"Warning: Failed to end LM Studio server: {e}")