kolega-code 0.1.0__py3-none-any.whl

kolega_code/llm/providers/openai.py

@@ -0,0 +1,363 @@
+from typing import Any, AsyncContextManager, Dict, List, Optional
+import json
+import logging
+import math
+
+import tiktoken
+from openai import AsyncOpenAI, OpenAI
+
+from ..models import ImageBlock, Message, MessageChunk, MessageHistory, ToolCall, ToolDefinition, ToolResult
+from ..tool_execution_ids import ToolExecutionIdRegistry
+from .base import BaseLLMProvider
+from .models import GenerationParams, TokenCount
+
+
+class OpenAIStreamWrapper:
+    def __init__(self, openai_stream, requested_include_usage: bool = False):
+        self.openai_stream = openai_stream
+        self.final_content = ""
+        self.final_tool_calls = {}
+        self.stop_reason = None
+        self.usage_data = None
+        self.tool_execution_ids = ToolExecutionIdRegistry()
+
+        self._closed = False
+        self._requested_include_usage = requested_include_usage
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if hasattr(self.openai_stream, "aclose"):
+            await self.openai_stream.aclose()
+
+        self._closed = True
+        return False
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._closed:
+            raise StopAsyncIteration
+
+        try:
+            chunk = await self.openai_stream.__anext__()
+
+            # Some providers emit usage-only events with no choices; guard accesses
+            if hasattr(chunk, "choices") and chunk.choices:
+                choice0 = chunk.choices[0]
+                delta = getattr(choice0, "delta", None)
+                if delta is not None:
+                    content = getattr(delta, "content", None) or ""
+                    if content:
+                        self.final_content += content
+
+                    for tool_call in getattr(delta, "tool_calls", []) or []:
+                        index = tool_call.index
+
+                        if index not in self.final_tool_calls:
+                            self.final_tool_calls[index] = tool_call
+
+                        if self.final_tool_calls[index].function.arguments != tool_call.function.arguments:
+                            if self.final_tool_calls[index].function.arguments is None:
+                                self.final_tool_calls[index].function.arguments = ""
+
+                            self.final_tool_calls[index].function.arguments += tool_call.function.arguments
+
+                # Capture stop reason if present
+                self.stop_reason = getattr(choice0, "finish_reason", self.stop_reason)
+
+            # Capture usage data from final chunk
+            if hasattr(chunk, "usage") and chunk.usage:
+                self.usage_data = {
+                    "prompt_tokens": chunk.usage.prompt_tokens,
+                    "completion_tokens": chunk.usage.completion_tokens,
+                    "total_tokens": chunk.usage.total_tokens,
+                }
+                # Capture cached prompt tokens if available (e.g., DashScope/Qwen)
+                details = getattr(chunk.usage, "prompt_tokens_details", None)
+                cached = None
+                if details is not None:
+                    cached = getattr(details, "cached_tokens", None)
+                    if cached is None and isinstance(details, dict):
+                        cached = details.get("cached_tokens")
+                if cached is not None:
+                    self.usage_data["cache_read_input_tokens"] = cached
+
+            # Return a safe chunk representation; ignore events with no choices
+            if hasattr(chunk, "choices") and chunk.choices:
+                return MessageChunk.from_openai(chunk)
+            else:
+                return MessageChunk(type="ignore", text="")
+
+        except StopAsyncIteration:
+            raise
+
+    async def get_final_message(self):
+        message = Message.from_openai_stream(
+            role="assistant",
+            content=self.final_content,
+            tool_calls=self.final_tool_calls,
+            stop_reason=self.stop_reason,
+            tool_execution_ids=self.tool_execution_ids,
+        )
+
+        # Add usage data if available
+        if self.usage_data:
+            message.usage_metadata.update(self.usage_data)
+        else:
+            logger = logging.getLogger(__name__)
+            if self._requested_include_usage:
+                logger.warning(
+                    "OpenAIStreamWrapper: include_usage requested but provider emitted no usage; billing may be skipped"
+                )
+            else:
+                logger.warning(
+                    "OpenAIStreamWrapper: no usage metadata captured from streaming response; billing may be skipped"
+                )
+
+        return message
+
+
+class OpenAIProvider(BaseLLMProvider):
+
+    models_max_completion_tokens = ["o3-mini", "o3", "o3-2025-04-16", "o4-mini"]
+
+    def __init__(
+        self,
+        api_key: str,
+        max_retries: int = 3,
+        requests_per_minute: Optional[int] = None,
+        tokens_per_minute: Optional[int] = None,
+        base_url: Optional[str] = None,
+    ):
+        super().__init__(api_key, max_retries, requests_per_minute, tokens_per_minute, base_url)
+        self.async_client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+        self.sync_client = OpenAI(api_key=api_key, base_url=base_url)
+
+    @property
+    def retry_decorator(self):
+        """Get retry decorator with configured max retries"""
+        return self.get_retry_decorator()
+
+    def _prepare_generation_params(self, params: Optional[GenerationParams] = None) -> Dict[str, Any]:
+        """Convert common parameters to provider-specific format"""
+        generation_params = {
+            "model": "gpt-4o",  # Default model
+        }
+
+        if params:
+            if params.temperature is not None:
+                generation_params["temperature"] = params.temperature
+            if params.max_completion_tokens is not None:
+                generation_params["max_tokens"] = params.max_completion_tokens
+            if params.tools:
+                generation_params["tools"] = [t.to_openai() for t in params.tools]
+            if params.thinking:
+                generation_params.update(self._prepare_thinking_params(params.thinking))
+
+        return generation_params
+
+    async def count_tokens(
+        self,
+        messages: MessageHistory,
+        system: Optional[Message] = None,
+        model: Optional[str] = None,
+        tools: List[ToolDefinition] = None,
+        **kwargs,
+    ) -> TokenCount:
+        """Count tokens for a list of messages using tiktoken.
+
+        Provides comprehensive token counting including:
+        - System prompts and messages with formatting overhead
+        - Images with estimation based on base64 data size
+        - Tool definitions with JSON serialization
+        - Tool calls and tool results in message content
+
+        Args:
+            messages: List of messages to count tokens for
+            system: Optional system message
+            model: Optional model name to use for counting (defaults to gpt-4)
+            tools: Optional tool definitions
+
+        Returns:
+            TokenCount object with input token count
+        """
+        encoding = tiktoken.get_encoding("cl100k_base")
+        num_tokens = 0
+
+        # Combine system message with messages if provided
+        all_messages = ([system] + list(messages)) if system else list(messages)
+
+        # Count tokens for each message
+        for message in all_messages:
+            # Base tokens for message formatting
+            num_tokens += 4  # Every message follows <im_start>{role/name}\n{content}<im_end>\n format
+
+            # Add tokens for role
+            if hasattr(message, "role") and message.role:
+                num_tokens += len(encoding.encode(message.role))
+
+            # Add tokens for content
+            if hasattr(message, "content"):
+                if isinstance(message.content, str):
+                    num_tokens += len(encoding.encode(message.content))
+                elif isinstance(message.content, list):
+                    for item in message.content:
+                        # Handle text blocks
+                        if hasattr(item, "text") and item.text:
+                            num_tokens += len(encoding.encode(item.text))
+                        elif isinstance(item, dict) and "text" in item:
+                            num_tokens += len(encoding.encode(item["text"]))
+                        # Handle image blocks
+                        elif isinstance(item, ImageBlock):
+                            num_tokens += self._estimate_image_tokens(len(item.data))
+                        elif hasattr(item, "data") and hasattr(item, "media_type"):
+                            # ImageBlock - estimate tokens based on base64 data size
+                            num_tokens += self._estimate_image_tokens(len(item.data))
+                        # Handle tool calls
+                        elif isinstance(item, ToolCall):
+                            tool_call_json = json.dumps(item.to_openai())
+                            num_tokens += len(encoding.encode(tool_call_json))
+                            num_tokens += 2  # Minimal formatting overhead for tool calls
+                        # Handle tool results
+                        elif isinstance(item, ToolResult):
+                            # Tool results contain content that needs to be counted
+                            if isinstance(item.content, str):
+                                num_tokens += len(encoding.encode(item.content))
+                            elif isinstance(item.content, list):
+                                for result_item in item.content:
+                                    if hasattr(result_item, "text") and result_item.text:
+                                        num_tokens += len(encoding.encode(result_item.text))
+                            num_tokens += 2  # Minimal formatting overhead for tool results
+
+        # Count tool definition tokens
+        if tools:
+            for tool in tools:
+                tool_json = json.dumps(tool.to_openai())
+                # Count JSON tokens
+                json_tokens = len(encoding.encode(tool_json))
+                # OpenAI uses highly optimized internal format (not JSON)
+                # Empirically, their token count is ~79% of raw JSON token count
+                # Apply scaling factor to match API behavior
+                num_tokens += int(json_tokens * 0.79)
+
+        return TokenCount(input_tokens=num_tokens)
+
+    def _estimate_image_tokens(self, base64_data_length: int) -> int:
+        """Estimate image token cost based on base64 data length.
+
+        OpenAI charges for images based on their dimensions after resizing.
+        Since we don't decode images (performance), we estimate based on data size.
+
+        Uses same formula as Anthropic for consistency:
+        tokens ≈ 20 + sqrt(base64_length * 6)
+
+        This gives reasonable estimates:
+        - Tiny images (96 chars base64): ~44 tokens
+        - Small images (~50KB base64): ~659 tokens
+        - Medium images (~200KB base64): ~1285 tokens
+        - Large images (~800KB base64): ~2549 tokens
+
+        Args:
+            base64_data_length: Length of base64 encoded image data
+
+        Returns:
+            Estimated token count for the image
+        """
+        # Use square root scaling for better fit across image sizes
+        # Base cost of 20 tokens + sqrt scaling
+        estimated_tokens = 20 + int(math.sqrt(base64_data_length * 6))
+        return estimated_tokens
+
+    async def stream(
+        self,
+        messages: MessageHistory,
+        system: Optional[Message] = None,
+        params: Optional[GenerationParams] = None,
+        **kwargs,
+    ) -> AsyncContextManager:
+        """Generate a streaming response from OpenAI
+
+        Returns a coroutine that resolves to an async iterator.
+        """
+        generation_params = self._prepare_generation_params(params)
+        generation_params.update(kwargs)
+        generation_params["stream"] = True
+        # Ask provider to include usage in the final stream chunk when supported
+        try:
+            existing_stream_options = generation_params.get("stream_options") or {}
+            existing_stream_options["include_usage"] = True
+            generation_params["stream_options"] = existing_stream_options
+        except Exception:
+            # Best-effort; some providers may not support stream_options
+            pass
+
+        # Swap max_tokens for max_completion_tokens for o3-mini, etc.
+        if generation_params["model"] in self.models_max_completion_tokens:
+            if "max_tokens" in generation_params:
+                generation_params["max_completion_tokens"] = generation_params["max_tokens"]
+                del generation_params["max_tokens"]
+
+        # Combine system message with messages if provided
+        if system:
+            messages = MessageHistory([system] + messages)
+
+        await self.rate_limiter.acquire()
+
+        return OpenAIStreamWrapper(
+            await self.async_client.chat.completions.create(messages=messages.to_openai(), **generation_params),
+            requested_include_usage=True,
+        )
+
+    async def generate(
+        self,
+        messages: MessageHistory,
+        system: Optional[Message] = None,
+        params: Optional[GenerationParams] = None,
+        **kwargs,
+    ) -> Message:
+        generation_params = self._prepare_generation_params(params)
+        generation_params.update(kwargs)
+
+        # Swap max_tokens for max_completion_tokens for o3-mini, etc.
+        if generation_params["model"] in self.models_max_completion_tokens:
+            if "max_tokens" in generation_params:
+                generation_params["max_completion_tokens"] = generation_params["max_tokens"]
+                del generation_params["max_tokens"]
+
+        # Combine system message with messages if provided
+        if system:
+            messages = MessageHistory([system] + messages)
+
+        await self.rate_limiter.acquire()
+        response = await self.async_client.chat.completions.create(messages=messages.to_openai(), **generation_params)
+
+        # Extract message and add usage data
+        message = Message.from_openai(response.choices[0].message)
+
+        # Add usage data from the response
+        if hasattr(response, "usage") and response.usage:
+            message.usage_metadata.update(
+                {
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                }
+            )
+            # Capture cached prompt tokens if available (e.g., DashScope/Qwen)
+            details = getattr(response.usage, "prompt_tokens_details", None)
+            cached = None
+            if details is not None:
+                cached = getattr(details, "cached_tokens", None)
+                if cached is None and isinstance(details, dict):
+                    cached = details.get("cached_tokens")
+            if cached is not None:
+                message.usage_metadata["cache_read_input_tokens"] = cached
+        else:
+            logging.getLogger(__name__).warning(
+                "OpenAIProvider.generate: response contains no usage metadata; billing may be skipped"
+            )
+
+        return message

kolega_code/llm/ratelimit.py

@@ -0,0 +1,40 @@
+import asyncio
+import time
+from typing import List, Optional
+
+
+class RateLimiter:
+    """Rate limiter for API requests and tokens"""
+
+    def __init__(self, requests_per_minute: Optional[int] = None, tokens_per_minute: Optional[int] = None):
+        self.requests_per_minute = requests_per_minute
+        self.tokens_per_minute = tokens_per_minute
+        self.request_timestamps: List[float] = []
+        self.token_timestamps: List[float] = []
+        self.lock = asyncio.Lock()
+
+    async def acquire(self, tokens: Optional[int] = None):
+        """Acquire permission to make a request"""
+        async with self.lock:
+            current_time = time.time()
+
+            # Clean up old timestamps
+            self.request_timestamps = [ts for ts in self.request_timestamps if current_time - ts < 60]
+            self.token_timestamps = [ts for ts in self.token_timestamps if current_time - ts < 60]
+
+            # Check request rate limit
+            if self.requests_per_minute and len(self.request_timestamps) >= self.requests_per_minute:
+                wait_time = 60 - (current_time - self.request_timestamps[0])
+                if wait_time > 0:
+                    await asyncio.sleep(wait_time)
+
+            # Check token rate limit
+            if self.tokens_per_minute and tokens and len(self.token_timestamps) >= self.tokens_per_minute:
+                wait_time = 60 - (current_time - self.token_timestamps[0])
+                if wait_time > 0:
+                    await asyncio.sleep(wait_time)
+
+            # Record new timestamps
+            self.request_timestamps.append(current_time)
+            if tokens:
+                self.token_timestamps.extend([current_time] * tokens)

kolega_code/llm/specs.py

@@ -0,0 +1,67 @@
+from typing import Dict, Tuple
+
+# Dictionary mapping (provider, model_name) to model specifications
+# Each entry contains context_length (maximum input tokens), max_completion_tokens, default_temperature,
+# and optional model capability flags.
+MODEL_SPECS: Dict[Tuple[str, str], Dict[str, int | float | bool]] = {
+    # Anthropic models
+    ("anthropic", "claude-opus-4-7"): {
+        "context_length": 1000000,
+        "max_completion_tokens": 128000,
+        "default_temperature": 1.0,
+        "supports_temperature": False,
+    },
+    ("anthropic", "claude-sonnet-4-6"): {"context_length": 1000000, "max_completion_tokens": 64000, "default_temperature": 1.0},
+    ("anthropic", "claude-3-7-sonnet-20250219"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    ("anthropic", "claude-3-haiku-20240307"): {"context_length": 200000, "max_completion_tokens": 4096, "default_temperature": 1.0},
+    ("anthropic", "claude-3-5-sonnet-20241022"): {"context_length": 200000, "max_completion_tokens": 8192, "default_temperature": 1.0},
+    ("anthropic", "claude-opus-4-20250514"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    ("anthropic", "claude-sonnet-4-20250514"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    ("anthropic", "claude-sonnet-4-5-20250929"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    ("anthropic", "claude-opus-4-5-20251101"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    ("anthropic", "claude-haiku-4-5-20251001"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    # Moonshot models
+    ("moonshot", "kimi-k2.6"): {"context_length": 262144, "max_completion_tokens": 32768, "default_temperature": 1.0},
+    # DeepSeek models
+    ("deepseek", "deepseek-v4-pro"): {"context_length": 1000000, "max_completion_tokens": 384000, "default_temperature": 1.0},
+    # OpenAI models
+    ("openai", "gpt-4o"): {"context_length": 128000, "max_completion_tokens": 4096, "default_temperature": 1.0},
+    ("openai", "o3-mini"): {"context_length": 200000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    ("openai", "gpt-4.1-2025-04-14"): {"context_length": 1000000, "max_completion_tokens": 32768, "default_temperature": 1.0},
+    ("openai", "gpt-4.1-mini"): {"context_length": 1000000, "max_completion_tokens": 32768, "default_temperature": 1.0},
+    ("openai", "o3-2025-04-16"): {"context_length": 200000, "max_completion_tokens": 100000, "default_temperature": 1.0},
+    ("openai", "o3"): {"context_length": 200000, "max_completion_tokens": 100000, "default_temperature": 1.0},
+    ("openai", "o4-mini"): {"context_length": 200000, "max_completion_tokens": 100000, "default_temperature": 1.0},
+    # Together.ai models
+    ("together", "deepseek-ai/DeepSeek-R1"): {"context_length": 64000, "max_completion_tokens": 8000, "default_temperature": 1.0},
+    # Google models
+    ("google", "gemini-2.0-flash"): {"context_length": 1000000, "max_completion_tokens": 8192, "default_temperature": 1.0},
+    ("google", "gemini-2.5-pro-exp-03-25"): {"context_length": 1000000, "max_completion_tokens": 65536, "default_temperature": 1.0},
+    ("google", "gemini-2.5-pro"): {"context_length": 1000000, "max_completion_tokens": 65536, "default_temperature": 1.0},
+    # X.ai models
+    ("xai", "grok-3-beta"): {"context_length": 128000, "max_completion_tokens": 16384, "default_temperature": 1.0},
+    # Fireworks models
+    ("fireworks", "accounts/fireworks/models/glm-4p5"): {"context_length": 128000, "max_completion_tokens": 16384, "default_temperature": 0.6},
+    ("dashscope", "qwen3-coder-plus"): {"context_length": 1000000, "max_completion_tokens": 16384, "default_temperature": 0.7},
+}
+
+
+def get_model_specs(provider: str, model_name: str) -> Dict[str, int | float | bool]:
+    """
+    Get the specifications for a given model.
+
+    Args:
+        provider: The LLM provider (e.g., 'anthropic', 'openai') - can be string or enum
+        model_name: The name of the model
+
+    Returns:
+        Dictionary containing context_length, max_completion_tokens, and default_temperature
+    """
+    # Handle both string and enum provider types
+    provider_str = provider.value if hasattr(provider, "value") else provider
+    key = (provider_str, model_name)
+
+    if key not in MODEL_SPECS:
+        raise ValueError(f"Model {model_name} from provider {provider_str} is not supported.")
+
+    return MODEL_SPECS.get(key)

kolega_code/llm/tool_execution_ids.py

@@ -0,0 +1,18 @@
+import uuid
+
+
+def new_tool_execution_id() -> str:
+    """Create an app-level identifier for one tool execution."""
+    return f"tool_exec_{uuid.uuid4().hex}"
+
+
+class ToolExecutionIdRegistry:
+    """Response-scoped mapping from provider tool call IDs to app execution IDs."""
+
+    def __init__(self) -> None:
+        self._by_provider_tool_call_id: dict[str, str] = {}
+
+    def get_or_create(self, provider_tool_call_id: str) -> str:
+        if provider_tool_call_id not in self._by_provider_tool_call_id:
+            self._by_provider_tool_call_id[provider_tool_call_id] = new_tool_execution_id()
+        return self._by_provider_tool_call_id[provider_tool_call_id]

kolega_code/models/__init__.py

@@ -0,0 +1,9 @@
+"""Shared models used by the agent runtime."""
+
+from .sandbox_terminal_state import SandboxTerminalState, TerminalInfo, TerminalOutput
+
+__all__ = [
+    "SandboxTerminalState",
+    "TerminalInfo",
+    "TerminalOutput",
+]

kolega_code/models/sandbox_terminal_state.py

@@ -0,0 +1,47 @@
+"""Sandbox terminal state model for persisting terminal sessions."""
+
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+import uuid
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class TerminalInfo(BaseModel):
+    """Information about a single terminal."""
+
+    terminal_id: str = Field(..., description="Terminal identifier")
+    created_at: datetime = Field(..., description="When the terminal was created")
+    cwd: str = Field(..., description="Working directory of the terminal")
+    env: Dict[str, str] = Field(default_factory=dict, description="Environment variables")
+    last_command: str = Field(default="", description="Last command executed")
+    last_command_purpose: str = Field(default="", description="Purpose of last command")
+
+
+class TerminalOutput(BaseModel):
+    """Single output entry from a terminal."""
+
+    type: str = Field(..., description="Type of output: command, stdout, stderr, exit")
+    data: str = Field(..., description="Output data")
+    timestamp: datetime = Field(..., description="When the output was generated")
+    purpose: Optional[str] = Field(None, description="Purpose for commands")
+    exit_code: Optional[int] = Field(None, description="Exit code for exit type")
+
+
+class SandboxTerminalState(BaseModel):
+    """Model for persisting sandbox terminal state."""
+
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)
+
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Unique identifier")
+    workspace_id: str = Field(..., description="Associated workspace ID")
+    sandbox_id: str = Field(..., description="Associated sandbox ID")
+    terminals: Dict[str, TerminalInfo] = Field(default_factory=dict)
+    outputs: Dict[str, List[TerminalOutput]] = Field(default_factory=dict)
+    default_terminal_id: Optional[str] = None
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    total_output_size: int = 0
+    MAX_OUTPUT_SIZE: int = 1048576
+    MAX_OUTPUT_PER_TERMINAL: int = 262144
+

kolega_code/runtime.py

@@ -0,0 +1,50 @@
+"""Runtime adapter registry for host-provided services.
+
+The shared agent package intentionally does not own product databases,
+background job runners, or app-specific MCP/environment services. Host
+applications can register those objects directly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+_registry: dict[str, Any] = {}
+
+
+class RuntimeAdapterError(RuntimeError):
+    """Raised when shared agent code needs a host service that was not registered."""
+
+
+def register_runtime_adapter(name: str, value: Any) -> None:
+    """Register a host-provided runtime dependency."""
+    _registry[name] = value
+
+
+def get_runtime_adapter(name: str) -> Any:
+    """Return a registered host adapter."""
+    if name in _registry:
+        return _registry[name]
+
+    raise RuntimeAdapterError(
+        f"Host runtime adapter '{name}' is not registered. "
+        "Call kolega_code.runtime.register_runtime_adapter during app startup."
+    )
+
+
+class RuntimeProxy:
+    """Lazy proxy for a host-provided object."""
+
+    def __init__(self, name: str) -> None:
+        self._name = name
+
+    def _target(self) -> Any:
+        return get_runtime_adapter(self._name)
+
+    def __getattr__(self, item: str) -> Any:
+        if item.startswith("_"):
+            raise AttributeError(item)
+        return getattr(self._target(), item)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._target()(*args, **kwargs)