sigma-terminal 3.4.1__py3-none-any.whl → 3.5.0__py3-none-any.whl

sigma/llm.py

@@ -1,786 +0,0 @@
-"""LLM client implementations for all providers."""
-
-import asyncio
-import json
-import re
-import time
-from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, Callable, Optional
-
-from sigma.config import LLMProvider, get_settings, ErrorCode, SigmaError
-
-
-# Rate limiting configuration
-class RateLimiter:
-    """Simple rate limiter to prevent API flooding."""
-    
-    def __init__(self, requests_per_minute: int = 10, min_interval: float = 1.0):
-        self.requests_per_minute = requests_per_minute
-        self.min_interval = min_interval
-        self.last_request_time = 0
-        self.request_count = 0
-        self.window_start = time.time()
-    
-    async def wait(self):
-        """Wait if necessary to respect rate limits."""
-        current_time = time.time()
-        
-        # Reset window if a minute has passed
-        if current_time - self.window_start >= 60:
-            self.window_start = current_time
-            self.request_count = 0
-        
-        # Check if we've hit the rate limit
-        if self.request_count >= self.requests_per_minute:
-            wait_time = 60 - (current_time - self.window_start)
-            if wait_time > 0:
-                await asyncio.sleep(wait_time)
-                self.window_start = time.time()
-                self.request_count = 0
-        
-        # Ensure minimum interval between requests
-        time_since_last = current_time - self.last_request_time
-        if time_since_last < self.min_interval:
-            await asyncio.sleep(self.min_interval - time_since_last)
-        
-        self.last_request_time = time.time()
-        self.request_count += 1
-
-
-# Global rate limiters per provider (generous limits to avoid rate limiting)
-_rate_limiters = {
-    "google": RateLimiter(requests_per_minute=60, min_interval=0.2),      # Gemini free tier is generous
-    "openai": RateLimiter(requests_per_minute=60, min_interval=0.2),      # GPT-5 tier 1+
-    "anthropic": RateLimiter(requests_per_minute=40, min_interval=0.3),   # Claude standard
-    "groq": RateLimiter(requests_per_minute=30, min_interval=0.2),        # Groq free tier
-    "xai": RateLimiter(requests_per_minute=30, min_interval=0.3),         # Grok standard
-    "ollama": RateLimiter(requests_per_minute=120, min_interval=0.05),    # Local, no limits
-}
-
-
-class BaseLLM(ABC):
-    """Base class for LLM clients."""
-    
-    provider_name: str = "base"
-    
-    @abstractmethod
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        """Generate a response."""
-        pass
-    
-    async def _rate_limit(self):
-        """Apply rate limiting."""
-        limiter = _rate_limiters.get(self.provider_name)
-        if limiter:
-            await limiter.wait()
-
-
-class GoogleLLM(BaseLLM):
-    """Google Gemini client."""
-    
-    provider_name = "google"
-    
-    def __init__(self, api_key: str, model: str):
-        from google import genai
-        self.client = genai.Client(api_key=api_key)
-        self.model_name = model
-    
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        await self._rate_limit()
-        from google.genai import types
-        
-        # Extract system prompt and build contents
-        system_prompt = None
-        contents = []
-        
-        for msg in messages:
-            role = msg["role"]
-            content = msg["content"]
-            
-            if role == "system":
-                system_prompt = content
-            elif role == "user":
-                contents.append(types.Content(
-                    role="user",
-                    parts=[types.Part(text=content)]
-                ))
-            elif role == "assistant":
-                contents.append(types.Content(
-                    role="model",
-                    parts=[types.Part(text=content)]
-                ))
-        
-        # Build config
-        config = types.GenerateContentConfig(
-            system_instruction=system_prompt,
-        )
-        
-        # Add tools if provided
-        if tools:
-            function_declarations = []
-            for tool in tools:
-                if tool.get("type") == "function":
-                    func = tool["function"]
-                    function_declarations.append(types.FunctionDeclaration(
-                        name=func["name"],
-                        description=func.get("description", ""),
-                        parameters=func.get("parameters", {}),
-                    ))
-            if function_declarations:
-                config.tools = [types.Tool(function_declarations=function_declarations)]
-        
-        # Generate
-        response = self.client.models.generate_content(
-            model=self.model_name,
-            contents=contents,
-            config=config,
-        )
-        
-        # Handle function calls
-        if response.candidates:
-            candidate = response.candidates[0]
-            if candidate.content and candidate.content.parts:
-                # Collect all function calls first
-                function_calls = []
-                for part in candidate.content.parts:
-                    if hasattr(part, 'function_call') and part.function_call:
-                        function_calls.append(part.function_call)
-                
-                # If there are function calls, process all of them
-                if function_calls and on_tool_call:
-                    # Add the model's response with function calls
-                    contents.append(candidate.content)
-                    
-                    # Execute all function calls and build responses
-                    function_responses = []
-                    for fc in function_calls:
-                        args = dict(fc.args) if fc.args else {}
-                        result = await on_tool_call(fc.name, args)
-                        function_responses.append(types.Part(
-                            function_response=types.FunctionResponse(
-                                name=fc.name,
-                                response={"result": str(result)}
-                            )
-                        ))
-                    
-                    # Add all function responses in a single user message
-                    contents.append(types.Content(
-                        role="user",
-                        parts=function_responses
-                    ))
-                    
-                    # Get final response
-                    final_response = self.client.models.generate_content(
-                        model=self.model_name,
-                        contents=contents,
-                        config=config,
-                    )
-                    
-                    # Check if there are more function calls in the response
-                    if final_response.candidates:
-                        final_candidate = final_response.candidates[0]
-                        if final_candidate.content and final_candidate.content.parts:
-                            for part in final_candidate.content.parts:
-                                if hasattr(part, 'function_call') and part.function_call:
-                                    # Recursive call to handle chained tool calls
-                                    new_contents = contents + [final_candidate.content]
-                                    fc = part.function_call
-                                    args = dict(fc.args) if fc.args else {}
-                                    result = await on_tool_call(fc.name, args)
-                                    new_contents.append(types.Content(
-                                        role="user",
-                                        parts=[types.Part(
-                                            function_response=types.FunctionResponse(
-                                                name=fc.name,
-                                                response={"result": str(result)}
-                                            )
-                                        )]
-                                    ))
-                                    final_final = self.client.models.generate_content(
-                                        model=self.model_name,
-                                        contents=new_contents,
-                                        config=config,
-                                    )
-                                    return final_final.text or ""
-                    
-                    return final_response.text or ""
-        
-        return response.text or ""
-
-
-class OpenAILLM(BaseLLM):
-    """OpenAI client."""
-    
-    provider_name = "openai"
-    
-    def __init__(self, api_key: str, model: str):
-        from openai import AsyncOpenAI
-        self.client = AsyncOpenAI(api_key=api_key)
-        self.model = model
-    
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        await self._rate_limit()
-        kwargs = {
-            "model": self.model,
-            "messages": messages,
-        }
-        
-        if tools:
-            kwargs["tools"] = tools
-            kwargs["tool_choice"] = "auto"
-        
-        response = await self.client.chat.completions.create(**kwargs)
-        message = response.choices[0].message
-        
-        # Handle tool calls
-        if message.tool_calls and on_tool_call:
-            tool_results = []
-            for tc in message.tool_calls:
-                args = json.loads(tc.function.arguments)
-                result = await on_tool_call(tc.function.name, args)
-                tool_results.append({
-                    "tool_call_id": tc.id,
-                    "role": "tool",
-                    "content": json.dumps(result)
-                })
-            
-            # Continue with tool results
-            messages = messages + [message.model_dump()] + tool_results
-            return await self.generate(messages, tools, on_tool_call)
-        
-        return message.content or ""
-
-
-class AnthropicLLM(BaseLLM):
-    """Anthropic Claude client."""
-    
-    provider_name = "anthropic"
-    
-    def __init__(self, api_key: str, model: str):
-        from anthropic import AsyncAnthropic
-        self.client = AsyncAnthropic(api_key=api_key)
-        self.model = model
-    
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        await self._rate_limit()
-        # Extract system message
-        system = ""
-        filtered_messages = []
-        for msg in messages:
-            if msg["role"] == "system":
-                system = msg["content"]
-            else:
-                filtered_messages.append(msg)
-        
-        kwargs = {
-            "model": self.model,
-            "max_tokens": 4096,
-            "messages": filtered_messages,
-        }
-        
-        if system:
-            kwargs["system"] = system
-        
-        if tools:
-            # Convert to Anthropic format
-            kwargs["tools"] = [
-                {
-                    "name": t["function"]["name"],
-                    "description": t["function"].get("description", ""),
-                    "input_schema": t["function"].get("parameters", {})
-                }
-                for t in tools if t.get("type") == "function"
-            ]
-        
-        response = await self.client.messages.create(**kwargs)
-        
-        # Handle tool use
-        result_text = ""
-        for block in response.content:
-            if block.type == "text":
-                result_text += block.text
-            elif block.type == "tool_use" and on_tool_call:
-                result = await on_tool_call(block.name, block.input)
-                # Continue conversation
-                filtered_messages.append({
-                    "role": "assistant",
-                    "content": response.content
-                })
-                filtered_messages.append({
-                    "role": "user",
-                    "content": [{
-                        "type": "tool_result",
-                        "tool_use_id": block.id,
-                        "content": json.dumps(result)
-                    }]
-                })
-                return await self.generate(
-                    [{"role": "system", "content": system}] + filtered_messages,
-                    tools, on_tool_call
-                )
-        
-        return result_text
-
-
-class GroqLLM(BaseLLM):
-    """Groq client."""
-    
-    provider_name = "groq"
-    
-    def __init__(self, api_key: str, model: str):
-        from groq import AsyncGroq
-        self.client = AsyncGroq(api_key=api_key)
-        self.model = model
-    
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        await self._rate_limit()
-        kwargs = {
-            "model": self.model,
-            "messages": messages,
-        }
-        
-        if tools:
-            kwargs["tools"] = tools
-            kwargs["tool_choice"] = "auto"
-        
-        response = await self.client.chat.completions.create(**kwargs)
-        message = response.choices[0].message
-        
-        # Handle tool calls (similar to OpenAI)
-        if message.tool_calls and on_tool_call:
-            tool_results = []
-            for tc in message.tool_calls:
-                args = json.loads(tc.function.arguments)
-                result = await on_tool_call(tc.function.name, args)
-                tool_results.append({
-                    "tool_call_id": tc.id,
-                    "role": "tool",
-                    "content": json.dumps(result)
-                })
-            
-            messages = messages + [{"role": "assistant", "tool_calls": message.tool_calls}] + tool_results
-            return await self.generate(messages, tools, on_tool_call)
-        
-        return message.content or ""
-
-
-class OllamaLLM(BaseLLM):
-    """Ollama local client with native tool call support."""
-    
-    provider_name = "ollama"
-    
-    def __init__(self, host: str, model: str):
-        self.host = host.rstrip("/")
-        self.model = model
-    
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        await self._rate_limit()
-        import aiohttp
-        
-        # Convert tools to Ollama format
-        ollama_tools = None
-        if tools:
-            ollama_tools = []
-            for tool in tools:
-                if tool.get("type") == "function":
-                    f = tool["function"]
-                    ollama_tools.append({
-                        "type": "function",
-                        "function": {
-                            "name": f["name"],
-                            "description": f.get("description", ""),
-                            "parameters": f.get("parameters", {})
-                        }
-                    })
-        
-        request_body = {
-            "model": self.model,
-            "messages": messages,
-            "stream": False
-        }
-        
-        if ollama_tools:
-            request_body["tools"] = ollama_tools
-        
-        async with aiohttp.ClientSession() as session:
-            try:
-                async with session.post(
-                    f"{self.host}/api/chat",
-                    json=request_body,
-                    timeout=aiohttp.ClientTimeout(total=120)
-                ) as resp:
-                    if resp.status != 200:
-                        error_text = await resp.text()
-                        return f"Ollama error: {error_text}"
-                    
-                    data = await resp.json()
-                    message = data.get("message", {})
-                    
-                    # Check for tool calls in response
-                    tool_calls = message.get("tool_calls", [])
-                    
-                    if tool_calls and on_tool_call:
-                        # Process tool calls
-                        updated_messages = messages.copy()
-                        updated_messages.append(message)
-                        
-                        for tc in tool_calls:
-                            func = tc.get("function", {})
-                            tool_name = func.get("name", "")
-                            tool_args = func.get("arguments", {})
-                            
-                            # If arguments is a string, parse it
-                            if isinstance(tool_args, str):
-                                try:
-                                    tool_args = json.loads(tool_args)
-                                except json.JSONDecodeError:
-                                    tool_args = {}
-                            
-                            # Execute the tool
-                            result = await on_tool_call(tool_name, tool_args)
-                            
-                            # Add tool result to messages
-                            updated_messages.append({
-                                "role": "tool",
-                                "content": json.dumps(result) if not isinstance(result, str) else result
-                            })
-                        
-                        # Get final response with tool results
-                        return await self._continue_with_tools(session, updated_messages, ollama_tools, on_tool_call)
-                    
-                    # Check for text-based tool calls (fallback for older models)
-                    content = message.get("content", "")
-                    if "TOOL_CALL:" in content and on_tool_call:
-                        result = await self._parse_text_tool_call(content, on_tool_call)
-                        if result:
-                            return result
-                    
-                    return content
-                    
-            except aiohttp.ClientError as e:
-                return f"Connection error: {e}. Is Ollama running?"
-            except asyncio.TimeoutError:
-                return "Request timed out. Try a simpler query or check Ollama status."
-    
-    async def _continue_with_tools(
-        self,
-        session,
-        messages: list[dict],
-        tools: Optional[list[dict]],
-        on_tool_call: Optional[Callable],
-        depth: int = 0
-    ) -> str:
-        """Continue conversation after tool calls."""
-        import aiohttp
-        
-        if depth > 5:  # Prevent infinite loops
-            return "Maximum tool call depth reached."
-        
-        request_body = {
-            "model": self.model,
-            "messages": messages,
-            "stream": False
-        }
-        if tools:
-            request_body["tools"] = tools
-        
-        async with session.post(
-            f"{self.host}/api/chat",
-            json=request_body,
-            timeout=aiohttp.ClientTimeout(total=120)
-        ) as resp:
-            data = await resp.json()
-            message = data.get("message", {})
-            
-            # Check for more tool calls
-            tool_calls = message.get("tool_calls", [])
-            if tool_calls and on_tool_call:
-                updated_messages = messages.copy()
-                updated_messages.append(message)
-                
-                for tc in tool_calls:
-                    func = tc.get("function", {})
-                    tool_name = func.get("name", "")
-                    tool_args = func.get("arguments", {})
-                    
-                    if isinstance(tool_args, str):
-                        try:
-                            tool_args = json.loads(tool_args)
-                        except json.JSONDecodeError:
-                            tool_args = {}
-                    
-                    result = await on_tool_call(tool_name, tool_args)
-                    updated_messages.append({
-                        "role": "tool",
-                        "content": json.dumps(result) if not isinstance(result, str) else result
-                    })
-                
-                return await self._continue_with_tools(session, updated_messages, tools, on_tool_call, depth + 1)
-            
-            return message.get("content", "")
-    
-    async def _parse_text_tool_call(self, content: str, on_tool_call: Callable) -> Optional[str]:
-        """Parse text-based tool calls for older models."""
-        # Pattern: TOOL_CALL: tool_name({"arg": "value"}) or TOOL_CALL: tool_name(arg=value)
-        pattern = r'TOOL_CALL:\s*(\w+)\s*\(([^)]*)\)'
-        match = re.search(pattern, content)
-        
-        if not match:
-            return None
-        
-        tool_name = match.group(1)
-        args_str = match.group(2).strip()
-        
-        # Try to parse arguments
-        try:
-            if args_str.startswith("{"):
-                args = json.loads(args_str)
-            else:
-                # Parse key=value format
-                args = {}
-                for part in args_str.split(","):
-                    if "=" in part:
-                        k, v = part.split("=", 1)
-                        args[k.strip()] = v.strip().strip('"\'')
-        except:
-            args = {"symbol": args_str} if args_str else {}
-        
-        # Execute tool
-        result = await on_tool_call(tool_name, args)
-        
-        # Format result for response
-        if isinstance(result, dict):
-            result_str = json.dumps(result, indent=2)
-        else:
-            result_str = str(result)
-        
-        # Return combined response
-        return f"Tool result:\n```json\n{result_str}\n```"
-    
-    def _format_tools_for_prompt(self, tools: list[dict]) -> str:
-        """Format tools as text for prompt injection (legacy fallback)."""
-        lines = ["You have access to these tools:"]
-        for tool in tools:
-            if tool.get("type") == "function":
-                f = tool["function"]
-                params = f.get("parameters", {}).get("properties", {})
-                param_str = ", ".join(params.keys()) if params else ""
-                lines.append(f"- {f['name']}({param_str}): {f.get('description', '')}")
-        lines.append("\nTo use a tool, respond with: TOOL_CALL: tool_name(args)")
-        lines.append("Example: TOOL_CALL: get_stock_quote(symbol=\"AAPL\")")
-        return "\n".join(lines)
-
-
-class XaiLLM(BaseLLM):
-    """xAI Grok client (uses OpenAI-compatible API)."""
-    
-    provider_name = "xai"
-    
-    def __init__(self, api_key: str, model: str):
-        from openai import AsyncOpenAI
-        self.client = AsyncOpenAI(
-            api_key=api_key,
-            base_url="https://api.x.ai/v1"
-        )
-        self.model = model
-    
-    async def generate(
-        self,
-        messages: list[dict],
-        tools: Optional[list[dict]] = None,
-        on_tool_call: Optional[Callable] = None,
-    ) -> str:
-        await self._rate_limit()
-        
-        try:
-            kwargs = {
-                "model": self.model,
-                "messages": messages,
-            }
-            
-            if tools:
-                kwargs["tools"] = tools
-                kwargs["tool_choice"] = "auto"
-            
-            response = await self.client.chat.completions.create(**kwargs)
-            message = response.choices[0].message
-            
-            # Handle tool calls
-            if message.tool_calls and on_tool_call:
-                tool_results = []
-                for tc in message.tool_calls:
-                    args = json.loads(tc.function.arguments)
-                    result = await on_tool_call(tc.function.name, args)
-                    tool_results.append({
-                        "tool_call_id": tc.id,
-                        "role": "tool",
-                        "content": json.dumps(result)
-                    })
-                
-                # Continue with tool results
-                messages = messages + [message.model_dump()] + tool_results
-                return await self.generate(messages, tools, on_tool_call)
-            
-            return message.content or ""
-            
-        except Exception as e:
-            error_str = str(e)
-            if "401" in error_str or "invalid_api_key" in error_str:
-                raise SigmaError(
-                    ErrorCode.API_KEY_INVALID,
-                    "xAI API key is invalid",
-                    {"provider": "xai"}
-                )
-            elif "429" in error_str or "rate_limit" in error_str:
-                raise SigmaError(
-                    ErrorCode.API_KEY_RATE_LIMITED,
-                    "xAI rate limit exceeded. Please wait.",
-                    {"provider": "xai"}
-                )
-            raise
-
-
-def _parse_api_error(error: Exception, provider: str) -> SigmaError:
-    """Parse API errors into SigmaError with proper codes."""
-    error_str = str(error).lower()
-    
-    if "401" in error_str or "invalid_api_key" in error_str or "api key not valid" in error_str:
-        return SigmaError(
-            ErrorCode.API_KEY_INVALID,
-            f"{provider.title()} API key is invalid. Check your key at /keys",
-            {"provider": provider, "original_error": str(error)[:200]}
-        )
-    elif "403" in error_str or "forbidden" in error_str:
-        return SigmaError(
-            ErrorCode.API_KEY_INVALID,
-            f"{provider.title()} API key doesn't have access to this model",
-            {"provider": provider}
-        )
-    elif "429" in error_str or "rate_limit" in error_str or "quota" in error_str:
-        return SigmaError(
-            ErrorCode.API_KEY_RATE_LIMITED,
-            f"{provider.title()} rate limit exceeded. Wait a moment.",
-            {"provider": provider}
-        )
-    elif "404" in error_str or "not found" in error_str or "does not exist" in error_str:
-        return SigmaError(
-            ErrorCode.MODEL_NOT_FOUND,
-            f"Model not found. Try /models to see available options.",
-            {"provider": provider}
-        )
-    elif "timeout" in error_str:
-        return SigmaError(
-            ErrorCode.TIMEOUT,
-            "Request timed out. Try a simpler query.",
-            {"provider": provider}
-        )
-    elif "connection" in error_str:
-        return SigmaError(
-            ErrorCode.CONNECTION_ERROR,
-            f"Cannot connect to {provider.title()}. Check internet.",
-            {"provider": provider}
-        )
-    else:
-        return SigmaError(
-            ErrorCode.PROVIDER_ERROR,
-            f"{provider.title()} error: {str(error)[:150]}",
-            {"provider": provider, "original_error": str(error)[:300]}
-        )
-
-
-def get_llm(provider: LLMProvider, model: Optional[str] = None) -> BaseLLM:
-    """Get LLM client for a provider."""
-    settings = get_settings()
-    
-    if model is None:
-        model = settings.get_model(provider)
-    
-    if provider == LLMProvider.GOOGLE:
-        api_key = settings.google_api_key
-        if not api_key:
-            raise SigmaError(
-                ErrorCode.API_KEY_MISSING,
-                "Google API key not configured. Use /keys to set up.",
-                {"provider": "google", "hint": "Get key at: https://aistudio.google.com/apikey"}
-            )
-        return GoogleLLM(api_key, model)
-    
-    elif provider == LLMProvider.OPENAI:
-        api_key = settings.openai_api_key
-        if not api_key:
-            raise SigmaError(
-                ErrorCode.API_KEY_MISSING,
-                "OpenAI API key not configured. Use /keys to set up.",
-                {"provider": "openai", "hint": "Get key at: https://platform.openai.com/api-keys"}
-            )
-        return OpenAILLM(api_key, model)
-    
-    elif provider == LLMProvider.ANTHROPIC:
-        api_key = settings.anthropic_api_key
-        if not api_key:
-            raise SigmaError(
-                ErrorCode.API_KEY_MISSING,
-                "Anthropic API key not configured. Use /keys to set up.",
-                {"provider": "anthropic", "hint": "Get key at: https://console.anthropic.com/settings/keys"}
-            )
-        return AnthropicLLM(api_key, model)
-    
-    elif provider == LLMProvider.GROQ:
-        api_key = settings.groq_api_key
-        if not api_key:
-            raise SigmaError(
-                ErrorCode.API_KEY_MISSING,
-                "Groq API key not configured. Use /keys to set up.",
-                {"provider": "groq", "hint": "Get key at: https://console.groq.com/keys"}
-            )
-        return GroqLLM(api_key, model)
-    
-    elif provider == LLMProvider.XAI:
-        api_key = settings.xai_api_key
-        if not api_key:
-            raise SigmaError(
-                ErrorCode.API_KEY_MISSING,
-                "xAI API key not configured. Use /keys to set up.",
-                {"provider": "xai", "hint": "Get key at: https://console.x.ai"}
-            )
-        return XaiLLM(api_key, model)
-    
-    elif provider == LLMProvider.OLLAMA:
-        return OllamaLLM(settings.ollama_host, model)
-    
-    else:
-        raise SigmaError(
-            ErrorCode.PROVIDER_UNAVAILABLE,
-            f"Unsupported provider: {provider}",
-            {"provider": str(provider)}
-        )