llmcode-cli 1.0.0__py3-none-any.whl

llm_code/runtime/streaming_executor.py

@@ -0,0 +1,216 @@
+"""StreamingToolCollector and StreamingToolExecutor: route and execute tools during streaming."""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any
+
+from llm_code.api.types import ToolResultBlock
+from llm_code.tools.parsing import ParsedToolCall
+from llm_code.tools.registry import ToolRegistry
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+# Thread pool shared for background read-only tool execution during streaming
+_STREAMING_EXECUTOR = ThreadPoolExecutor(max_workers=8, thread_name_prefix="streaming-tool")
+
+
+def _attempt_partial_json_recovery(partial: str) -> dict:
+    """Try to recover a valid dict from partial/malformed JSON.
+
+    Attempts several repair strategies in order:
+    1. Direct parse (already complete)
+    2. Append ``}``
+    3. Append ``"}``
+    4. Append ``"}``  (for unclosed string + object)
+    5. Return empty dict as fallback
+    """
+    candidates = [
+        partial,
+        partial + "}",
+        partial + '"}',
+        partial + '"}}',
+        partial + "}}",
+    ]
+    for candidate in candidates:
+        try:
+            result = json.loads(candidate)
+            if isinstance(result, dict):
+                return result
+        except (json.JSONDecodeError, ValueError):
+            continue
+    return {}
+
+
+class StreamingToolExecutor:
+    """Execute read-only tools concurrently while the model is still streaming.
+
+    Usage pattern (mirrors the streaming loop in conversation.py):
+
+    .. code-block:: python
+
+        executor = StreamingToolExecutor(registry, permission_policy)
+
+        # Inside streaming loop:
+        if isinstance(event, StreamToolUseStart):
+            executor.start_tool(event.id, event.name)
+        elif isinstance(event, StreamToolUseInputDelta):
+            executor.submit(event.id, event.partial_json)
+
+        # After each tool input is complete (StreamToolUseStop or next tool start):
+        executor.finalize(tool_use_id)   # triggers background execution for reads
+
+        # After stream ends, collect all results:
+        results = await executor.collect_results()
+
+    The executor decides at :meth:`finalize` time whether a tool is read-only and
+    concurrency-safe.  If yes, it starts a background ``asyncio.Task``.  Write
+    tools are queued and returned unfired so conversation.py can execute them via
+    the normal ``_execute_tool_with_streaming`` path.
+    """
+
+    def __init__(self, tool_registry: ToolRegistry, permission_policy: Any = None) -> None:
+        self._registry = tool_registry
+        self._permissions = permission_policy
+
+        # json_parts accumulation: tool_use_id -> list[str]
+        self._json_parts: dict[str, list[str]] = {}
+        # tool names: tool_use_id -> name
+        self._tool_names: dict[str, str] = {}
+
+        # background tasks for read-only tools: tool_use_id -> Task
+        self._read_tasks: dict[str, asyncio.Task] = {}
+        # pending write calls (not yet executed)
+        self._write_calls: list[ParsedToolCall] = []
+
+    def start_tool(self, tool_use_id: str, name: str) -> None:
+        """Register a new tool use beginning (StreamToolUseStart event)."""
+        self._tool_names[tool_use_id] = name
+        self._json_parts[tool_use_id] = []
+
+    def submit(self, tool_use_id: str, partial_json: str) -> None:
+        """Accumulate a partial JSON chunk (StreamToolUseInputDelta event)."""
+        if tool_use_id in self._json_parts:
+            self._json_parts[tool_use_id].append(partial_json)
+
+    def finalize(self, tool_use_id: str) -> None:
+        """Mark tool input as complete; launch background execution if read-only.
+
+        For read-only + concurrency-safe tools: starts an asyncio.Task immediately.
+        For write tools: queues a ParsedToolCall for later sequential execution.
+        """
+        name = self._tool_names.get(tool_use_id)
+        if name is None:
+            logger.debug("finalize called for unknown tool_use_id %s", tool_use_id)
+            return
+
+        raw_json = "".join(self._json_parts.get(tool_use_id, []))
+        try:
+            args = json.loads(raw_json) if raw_json.strip() else {}
+        except (json.JSONDecodeError, ValueError):
+            args = _attempt_partial_json_recovery(raw_json)
+
+        tool = self._registry.get(name)
+        if tool is not None and tool.is_read_only(args) and tool.is_concurrency_safe(args):
+            # Start background execution immediately
+            call = ParsedToolCall(id=tool_use_id, name=name, args=args, source="native")
+            task = asyncio.get_event_loop().create_task(
+                self._execute_read_tool(tool_use_id, call, tool, args),
+                name=f"streaming-read-{name}-{tool_use_id[:8]}",
+            )
+            self._read_tasks[tool_use_id] = task
+        else:
+            # Queue for sequential execution after stream completes
+            call = ParsedToolCall(id=tool_use_id, name=name, args=args, source="native")
+            self._write_calls.append(call)
+
+    async def _execute_read_tool(
+        self,
+        tool_use_id: str,
+        call: ParsedToolCall,
+        tool: Any,
+        args: dict,
+    ) -> ToolResultBlock:
+        """Run the tool in a thread pool and return a ToolResultBlock."""
+        loop = asyncio.get_running_loop()
+        try:
+            result = await loop.run_in_executor(
+                _STREAMING_EXECUTOR,
+                lambda: tool.execute(args),
+            )
+            return ToolResultBlock(
+                tool_use_id=tool_use_id,
+                content=result.output,
+                is_error=result.is_error,
+            )
+        except Exception as exc:
+            logger.warning("Background read tool %s failed: %s", call.name, exc)
+            return ToolResultBlock(
+                tool_use_id=tool_use_id,
+                content=f"Tool execution failed: {exc}",
+                is_error=True,
+            )
+
+    async def collect_results(self) -> tuple[list[ToolResultBlock], list[ParsedToolCall]]:
+        """Wait for all background read tasks; return (read_results, write_calls).
+
+        - ``read_results``: ToolResultBlocks for all read-only tools that ran concurrently
+        - ``write_calls``: ParsedToolCalls for write tools that still need execution
+        """
+        read_results: list[ToolResultBlock] = []
+        if self._read_tasks:
+            done = await asyncio.gather(*self._read_tasks.values(), return_exceptions=True)
+            for item in done:
+                if isinstance(item, ToolResultBlock):
+                    read_results.append(item)
+                elif isinstance(item, BaseException):
+                    logger.error("Unexpected error in background read task: %s", item)
+
+        return read_results, list(self._write_calls)
+
+    def pending_write_count(self) -> int:
+        """Return number of write calls waiting for sequential execution."""
+        return len(self._write_calls)
+
+
+class StreamingToolCollector:
+    """Collects completed tool calls and decides whether they can run immediately.
+
+    A tool call is eligible for immediate (concurrent) execution when *both*:
+    - ``tool.is_read_only(args)`` returns True
+    - ``tool.is_concurrency_safe(args)`` returns True
+
+    All other calls (write operations, unknown tools, or tools that are not
+    concurrency-safe) are buffered and returned together via :meth:`flush_pending`.
+    """
+
+    def __init__(self, tool_registry: ToolRegistry) -> None:
+        self._registry = tool_registry
+        self._pending_writes: list[ParsedToolCall] = []
+
+    def on_tool_complete(self, call: ParsedToolCall) -> ParsedToolCall | None:
+        """A tool call finished parsing.
+
+        If the tool is read-only and concurrency-safe, return it immediately
+        for parallel execution.  Otherwise buffer it and return None.
+        """
+        tool = self._registry.get(call.name)
+        if tool is not None and tool.is_read_only(call.args) and tool.is_concurrency_safe(call.args):
+            return call
+        self._pending_writes.append(call)
+        return None
+
+    def flush_pending(self) -> list[ParsedToolCall]:
+        """Return all buffered calls and clear the internal buffer."""
+        pending = self._pending_writes
+        self._pending_writes = []
+        return pending
+
+    def has_pending(self) -> bool:
+        """Return True if there are buffered (write/unsafe) calls waiting."""
+        return len(self._pending_writes) > 0

llm_code/runtime/telemetry.py

@@ -0,0 +1,196 @@
+"""OpenTelemetry observability for llm-code.
+
+All OpenTelemetry imports are lazy — the module works as a no-op when the
+``opentelemetry-*`` packages are not installed.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class TelemetryConfig:
+    enabled: bool = False
+    endpoint: str = "http://localhost:4318"  # OTLP HTTP default
+    service_name: str = "llm-code"
+
+
+# ---------------------------------------------------------------------------
+# Telemetry class
+# ---------------------------------------------------------------------------
+
+class Telemetry:
+    """Thin wrapper around OpenTelemetry tracing and metrics.
+
+    When ``enabled=False`` or the ``opentelemetry-*`` packages are not
+    installed every method is a no-op so callers need no guard clauses.
+    """
+
+    def __init__(self, config: TelemetryConfig) -> None:
+        self._enabled = config.enabled
+        self._tracer: Any = None
+        self._cost_counter: Any = None
+        self._error_counter: Any = None
+        # Status/kind enums stored at setup so methods never re-import
+        self._SpanKind: Any = None
+        self._StatusCode: Any = None
+        self._otel_available = False
+
+        if not self._enabled:
+            return
+
+        try:
+            self._setup(config)
+            self._otel_available = True
+        except Exception:
+            # If setup fails for any reason (e.g., ImportError), degrade gracefully
+            self._enabled = False
+
+    # ------------------------------------------------------------------
+    # Setup (only called when enabled and packages are present)
+    # ------------------------------------------------------------------
+
+    def _setup(self, config: TelemetryConfig) -> None:
+        from opentelemetry import trace, metrics
+        from opentelemetry.trace import SpanKind, StatusCode
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+        from opentelemetry.sdk.metrics import MeterProvider
+        from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+        from opentelemetry.sdk.resources import Resource
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+        from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+
+        # Store enums so methods don't need to re-import
+        self._SpanKind = SpanKind
+        self._StatusCode = StatusCode
+
+        resource = Resource.create({"service.name": config.service_name})
+
+        # Tracer
+        span_exporter = OTLPSpanExporter(endpoint=f"{config.endpoint}/v1/traces")
+        tracer_provider = TracerProvider(resource=resource)
+        tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
+        trace.set_tracer_provider(tracer_provider)
+        self._tracer = trace.get_tracer(config.service_name)
+
+        # Meter
+        metric_exporter = OTLPMetricExporter(endpoint=f"{config.endpoint}/v1/metrics")
+        reader = PeriodicExportingMetricReader(metric_exporter, export_interval_millis=60_000)
+        meter_provider = MeterProvider(resource=resource, metric_readers=[reader])
+        metrics.set_meter_provider(meter_provider)
+        meter = metrics.get_meter(config.service_name)
+
+        self._cost_counter = meter.create_counter(
+            name="llm.cost.usd",
+            unit="USD",
+            description="Cumulative LLM cost in US dollars",
+        )
+        self._error_counter = meter.create_counter(
+            name="llm.errors",
+            description="Count of errors by type",
+        )
+
+    # ------------------------------------------------------------------
+    # Public API — all methods are safe to call unconditionally
+    # ------------------------------------------------------------------
+
+    def trace_turn(
+        self,
+        session_id: str,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        duration_ms: float,
+    ) -> None:
+        """Record a completed LLM conversation turn as an OTel span."""
+        if not self._enabled or self._tracer is None:
+            return
+        try:
+            with self._tracer.start_as_current_span(
+                "llm.turn",
+                kind=self._SpanKind.CLIENT,
+            ) as span:
+                span.set_attribute("session.id", session_id)
+                span.set_attribute("llm.model", model)
+                span.set_attribute("llm.tokens.input", input_tokens)
+                span.set_attribute("llm.tokens.output", output_tokens)
+                span.set_attribute("llm.duration_ms", duration_ms)
+                span.set_status(self._StatusCode.OK)
+        except Exception:
+            pass
+
+    def trace_tool(
+        self,
+        tool_name: str,
+        duration_ms: float,
+        is_error: bool = False,
+    ) -> None:
+        """Record a tool execution as an OTel span."""
+        if not self._enabled or self._tracer is None:
+            return
+        try:
+            with self._tracer.start_as_current_span(
+                f"tool.{tool_name}",
+                kind=self._SpanKind.INTERNAL,
+            ) as span:
+                span.set_attribute("tool.name", tool_name)
+                span.set_attribute("tool.duration_ms", duration_ms)
+                span.set_attribute("tool.is_error", is_error)
+                if is_error:
+                    span.set_status(self._StatusCode.ERROR)
+                else:
+                    span.set_status(self._StatusCode.OK)
+        except Exception:
+            pass
+
+    def record_cost(
+        self,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        cost_usd: float,
+    ) -> None:
+        """Record LLM cost as an OTel metric counter."""
+        if not self._enabled or self._cost_counter is None:
+            return
+        try:
+            self._cost_counter.add(
+                cost_usd,
+                attributes={
+                    "llm.model": model,
+                    "llm.tokens.input": input_tokens,
+                    "llm.tokens.output": output_tokens,
+                },
+            )
+        except Exception:
+            pass
+
+    def record_error(self, error_type: str, message: str) -> None:
+        """Record an error event as an OTel counter increment."""
+        if not self._enabled or self._error_counter is None:
+            return
+        try:
+            self._error_counter.add(
+                1,
+                attributes={"error.type": error_type, "error.message": message[:256]},
+            )
+        except Exception:
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Module-level singleton helpers
+# ---------------------------------------------------------------------------
+
+_NOOP = Telemetry(TelemetryConfig(enabled=False))
+
+
+def get_noop_telemetry() -> Telemetry:
+    """Return the shared no-op Telemetry instance."""
+    return _NOOP

llm_code/runtime/token_budget.py

@@ -0,0 +1,26 @@
+"""Token budget tracking for agentic conversation turns."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class TokenBudget:
+    target: int
+    consumed: int = 0
+
+    def add(self, output_tokens: int) -> None:
+        self.consumed += output_tokens
+
+    def remaining(self) -> int:
+        return max(0, self.target - self.consumed)
+
+    def should_nudge(self) -> bool:
+        return self.consumed < self.target
+
+    def is_exhausted(self) -> bool:
+        return self.consumed >= self.target
+
+    def nudge_message(self) -> str:
+        rem = self.remaining()
+        return f"[Token budget: {rem:,} tokens remaining out of {self.target:,}. Continue working toward the goal.]"

llm_code/runtime/vcr.py

@@ -0,0 +1,142 @@
+"""VCR session recording and playback for llm-code."""
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator
+
+
+EVENT_TYPES = (
+    "user_input",
+    "llm_request",
+    "llm_response",
+    "tool_call",
+    "tool_result",
+    "stream_event",
+    "error",
+)
+
+
+@dataclass(frozen=True)
+class VCREvent:
+    """A single recorded event with timestamp, type, and payload."""
+
+    ts: float
+    type: str
+    data: dict
+
+
+class VCRRecorder:
+    """Records session events as JSONL lines to a file."""
+
+    def __init__(self, path: Path) -> None:
+        self._path = path
+        self._closed = False
+        self._file = None
+
+    def _ensure_open(self) -> None:
+        if self._file is None:
+            self._path.parent.mkdir(parents=True, exist_ok=True)
+            self._file = self._path.open("w", encoding="utf-8")
+
+    def record(self, event_type: str, data: dict) -> None:
+        """Write a single JSONL event line with the current timestamp."""
+        if self._closed:
+            raise RuntimeError("VCRRecorder is closed")
+        self._ensure_open()
+        entry = {"ts": time.time(), "type": event_type, "data": data}
+        self._file.write(json.dumps(entry, ensure_ascii=False) + "\n")
+        self._file.flush()
+
+    def close(self) -> None:
+        """Close the recording file."""
+        if self._file is not None:
+            self._file.close()
+            self._file = None
+        self._closed = True
+
+    # Context manager support
+    def __enter__(self) -> "VCRRecorder":
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()
+
+
+class VCRPlayer:
+    """Replays a JSONL recording file and provides summary statistics."""
+
+    def __init__(self, path: Path) -> None:
+        self._path = path
+
+    def _read_events(self) -> list[VCREvent]:
+        """Parse all valid JSONL events from the file."""
+        events: list[VCREvent] = []
+        try:
+            text = self._path.read_text(encoding="utf-8")
+        except (FileNotFoundError, OSError):
+            return events
+
+        for line in text.splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+                events.append(VCREvent(
+                    ts=float(obj["ts"]),
+                    type=str(obj["type"]),
+                    data=obj.get("data", {}),
+                ))
+            except (json.JSONDecodeError, KeyError, TypeError, ValueError):
+                continue
+        return events
+
+    def replay(self, speed: float = 1.0) -> Iterator[VCREvent]:
+        """Yield events from the recording, optionally with timing.
+
+        Args:
+            speed: Playback speed multiplier. 0.0 means instant (no sleep).
+                   1.0 means real-time. 2.0 means double speed.
+        """
+        events = self._read_events()
+        if not events:
+            return
+
+        prev_ts: float | None = None
+        for event in events:
+            if speed > 0.0 and prev_ts is not None:
+                delay = (event.ts - prev_ts) / speed
+                if delay > 0:
+                    time.sleep(delay)
+            prev_ts = event.ts
+            yield event
+
+    def summary(self) -> dict:
+        """Return summary statistics for the recording.
+
+        Returns:
+            dict with keys:
+              - event_count: total number of events
+              - duration: wall-clock seconds from first to last event
+              - tool_calls: dict mapping tool name to call count
+        """
+        events = self._read_events()
+        if not events:
+            return {"event_count": 0, "duration": 0.0, "tool_calls": {}}
+
+        duration = events[-1].ts - events[0].ts if len(events) > 1 else 0.0
+
+        tool_calls: dict[str, int] = {}
+        for event in events:
+            if event.type == "tool_call":
+                name = event.data.get("name", "unknown")
+                tool_calls[name] = tool_calls.get(name, 0) + 1
+
+        return {
+            "event_count": len(events),
+            "duration": duration,
+            "tool_calls": tool_calls,
+        }

llm_code/runtime/vision.py

@@ -0,0 +1,102 @@
+"""Vision fallback: replace ImageBlocks when the primary model has no vision support."""
+from __future__ import annotations
+
+import dataclasses
+import os
+
+from llm_code.api.types import ImageBlock, Message, TextBlock
+from llm_code.runtime.config import VisionConfig
+
+
+class VisionFallback:
+    """Wraps VisionConfig and can optionally call a vision API to describe images."""
+
+    def __init__(self, config: VisionConfig) -> None:
+        self._config = config
+
+    def is_configured(self) -> bool:
+        """Return True if both vision_model and vision_api are non-empty."""
+        return bool(self._config.vision_model) and bool(self._config.vision_api)
+
+    async def describe_image(self, image: ImageBlock) -> str:
+        """Call the vision API and return a text description of the image."""
+        from llm_code.api.openai_compat import OpenAICompatProvider
+        from llm_code.api.types import Message, MessageRequest, TextBlock
+
+        api_key = ""
+        if self._config.vision_api_key_env:
+            api_key = os.environ.get(self._config.vision_api_key_env, "")
+
+        provider = OpenAICompatProvider(
+            base_url=self._config.vision_api,
+            api_key=api_key,
+            model_name=self._config.vision_model,
+        )
+        try:
+            request = MessageRequest(
+                model=self._config.vision_model,
+                messages=(
+                    Message(
+                        role="user",
+                        content=(
+                            image,
+                            TextBlock(text="Describe this image in detail."),
+                        ),
+                    ),
+                ),
+                stream=False,
+            )
+            response = await provider.send_message(request)
+            for block in response.content:
+                if isinstance(block, TextBlock):
+                    return block.text
+            return ""
+        finally:
+            await provider.close()
+
+
+def preprocess_images(
+    msg: Message,
+    supports_images: bool,
+    vision_fallback: "VisionFallback | None",
+    return_warnings: bool = False,
+) -> "Message | tuple[Message, list[str]]":
+    """Pre-process a message's image blocks based on vision support.
+
+    - If model supports images OR message has no images → passthrough unchanged.
+    - If vision_fallback is configured → replace each ImageBlock with a placeholder
+      TextBlock (sync version; describe_image is not called here).
+    - Otherwise → strip ImageBlocks, emit a warning.
+    """
+    has_images = any(isinstance(b, ImageBlock) for b in msg.content)
+
+    if supports_images or not has_images:
+        if return_warnings:
+            return msg, []
+        return msg
+
+    warnings: list[str] = []
+
+    if vision_fallback is not None and vision_fallback.is_configured():
+        # Replace each image with a placeholder; actual async describe not called here
+        new_blocks = []
+        for block in msg.content:
+            if isinstance(block, ImageBlock):
+                new_blocks.append(
+                    TextBlock(text="[image: vision description not yet available]")
+                )
+            else:
+                new_blocks.append(block)
+        new_msg = dataclasses.replace(msg, content=tuple(new_blocks))
+    else:
+        # Strip images, keep other blocks
+        new_blocks = [b for b in msg.content if not isinstance(b, ImageBlock)]
+        new_msg = dataclasses.replace(msg, content=tuple(new_blocks))
+        warnings.append(
+            "One or more images were stripped because the model does not support "
+            "vision and no vision fallback is configured."
+        )
+
+    if return_warnings:
+        return new_msg, warnings
+    return new_msg

llm_code/swarm/__init__.py

@@ -0,0 +1 @@
+"""Swarm: multi-agent collaboration module."""