voidx 1.0.0__py3-none-any.whl

voidx/llm/usage.py

@@ -0,0 +1,341 @@
+"""Token usage accounting for context and LLM calls."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from voidx.llm.context import count_messages_tokens
+
+
+@dataclass
+class TokenUsage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    reasoning_tokens: int = 0
+    cache_read_tokens: int = 0
+    cache_write_tokens: int = 0
+    cache_tokens_reported: bool = False
+
+
+@dataclass
+class UsageStats:
+    context_tokens: int = 0
+    context_limit: int = 128_000
+    last_input_tokens: int = 0
+    last_output_tokens: int = 0
+    last_cache_read_tokens: int = 0
+    last_cache_write_tokens: int = 0
+    last_estimated_cache_read_tokens: int = 0
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    total_cache_read_tokens: int = 0
+    total_cache_write_tokens: int = 0
+    total_estimated_cache_read_tokens: int = 0
+    total_cache_metric_calls: int = 0
+    estimated_cache_calls: int = 0
+    total_calls: int = 0
+    _cache_context_history: dict[str, list[list[dict[str, str]]]] = field(
+        default_factory=dict,
+        repr=False,
+        compare=False,
+    )
+
+    @property
+    def total_tokens(self) -> int:
+        return self.total_input_tokens + self.total_output_tokens
+
+    @property
+    def cache_observed_tokens(self) -> int:
+        return self.total_cache_read_tokens + self.total_cache_write_tokens
+
+    @property
+    def cache_hit_rate(self) -> float | None:
+        actual = self.actual_cache_hit_rate
+        if actual is not None:
+            return actual
+        return self.estimated_cache_hit_rate
+
+    @property
+    def actual_cache_hit_rate(self) -> float | None:
+        if self.total_cache_metric_calls <= 0 and self.cache_observed_tokens <= 0:
+            return None
+        denominator = max(self.total_input_tokens, self.cache_observed_tokens)
+        if denominator <= 0:
+            return None
+        return self.total_cache_read_tokens / denominator
+
+    @property
+    def estimated_cache_hit_rate(self) -> float | None:
+        if self.estimated_cache_calls <= 0 or self.total_input_tokens <= 0:
+            return None
+        return self.total_estimated_cache_read_tokens / self.total_input_tokens
+
+    @property
+    def cache_hit_rate_is_estimated(self) -> bool:
+        return self.actual_cache_hit_rate is None and self.estimated_cache_hit_rate is not None
+
+    def reset(self) -> None:
+        self.context_tokens = 0
+        self.last_input_tokens = 0
+        self.last_output_tokens = 0
+        self.last_cache_read_tokens = 0
+        self.last_cache_write_tokens = 0
+        self.last_estimated_cache_read_tokens = 0
+        self.total_input_tokens = 0
+        self.total_output_tokens = 0
+        self.total_cache_read_tokens = 0
+        self.total_cache_write_tokens = 0
+        self.total_estimated_cache_read_tokens = 0
+        self.total_cache_metric_calls = 0
+        self.estimated_cache_calls = 0
+        self.total_calls = 0
+        self._cache_context_history.clear()
+
+    def update_context(self, tokens: int, limit: int | None = None) -> None:
+        self.context_tokens = max(tokens, 0)
+        if limit is not None:
+            self.context_limit = max(limit, 0)
+
+    def record_call(
+        self,
+        usage: TokenUsage,
+        *,
+        fallback_input_tokens: int = 0,
+        fallback_output_tokens: int = 0,
+        messages: list | None = None,
+        model: str = "",
+        cache_key: str = "",
+    ) -> None:
+        cache_key = cache_key or model or "default"
+        current_context = _messages_for_count(messages) if messages is not None else None
+        estimated_cache_read_tokens = (
+            self._estimate_cache_reuse_tokens(current_context, model, cache_key)
+            if current_context is not None else None
+        )
+        input_tokens = usage.input_tokens or fallback_input_tokens
+        output_tokens = usage.output_tokens or fallback_output_tokens
+        self.last_input_tokens = max(input_tokens, 0)
+        self.last_output_tokens = max(output_tokens, 0)
+        self.last_cache_read_tokens = max(usage.cache_read_tokens, 0)
+        self.last_cache_write_tokens = max(usage.cache_write_tokens, 0)
+        self.last_estimated_cache_read_tokens = 0
+        self.total_input_tokens += self.last_input_tokens
+        self.total_output_tokens += self.last_output_tokens
+        self.total_cache_read_tokens += self.last_cache_read_tokens
+        self.total_cache_write_tokens += self.last_cache_write_tokens
+        if usage.cache_tokens_reported:
+            self.total_cache_metric_calls += 1
+        elif estimated_cache_read_tokens is not None and self.last_input_tokens > 0:
+            self.last_estimated_cache_read_tokens = min(
+                max(estimated_cache_read_tokens, 0),
+                self.last_input_tokens,
+            )
+            self.total_estimated_cache_read_tokens += self.last_estimated_cache_read_tokens
+            self.estimated_cache_calls += 1
+        self.total_calls += 1
+        if usage.input_tokens:
+            self.context_tokens = usage.input_tokens
+        if current_context is not None:
+            history = self._cache_context_history.setdefault(cache_key, [])
+            history.append(current_context)
+            del history[:-8]
+
+    def _estimate_cache_reuse_tokens(
+        self,
+        current_context: list[dict[str, str]],
+        model: str,
+        cache_key: str,
+    ) -> int:
+        best = 0
+        for previous_context in self._cache_context_history.get(cache_key, []):
+            common_prefix = _common_prefix_context(previous_context, current_context)
+            if common_prefix:
+                best = max(best, count_messages_tokens(common_prefix, model))
+        return best
+
+
+def extract_token_usage(message: object) -> TokenUsage:
+    """Extract provider token usage from LangChain message metadata."""
+    usage = TokenUsage()
+    sources = _usage_sources(message)
+    for source in sources:
+        usage.input_tokens = usage.input_tokens or _first_int(
+            source,
+            "input_tokens",
+            "prompt_tokens",
+            "input",
+        )
+        usage.output_tokens = usage.output_tokens or _first_int(
+            source,
+            "output_tokens",
+            "completion_tokens",
+            "output",
+        )
+        usage.total_tokens = usage.total_tokens or _first_int(source, "total_tokens", "total")
+        usage.reasoning_tokens = usage.reasoning_tokens or _nested_first_int(
+            source,
+            ("output_token_details", "reasoning"),
+            ("completion_tokens_details", "reasoning_tokens"),
+            ("completion_token_details", "reasoning_tokens"),
+            ("reasoning",),
+        )
+        usage.cache_read_tokens = usage.cache_read_tokens or _nested_first_int(
+            source,
+            ("input_token_details", "cache_read"),
+            ("input_token_details", "cached_tokens"),
+            ("prompt_tokens_details", "cached_tokens"),
+            ("cache_read_input_tokens",),
+        )
+        usage.cache_write_tokens = usage.cache_write_tokens or _nested_first_int(
+            source,
+            ("input_token_details", "cache_creation"),
+            ("cache_creation_input_tokens",),
+        )
+        usage.cache_tokens_reported = usage.cache_tokens_reported or _has_nested_key(
+            source,
+            ("input_token_details", "cache_read"),
+            ("input_token_details", "cached_tokens"),
+            ("prompt_tokens_details", "cached_tokens"),
+            ("cache_read_input_tokens",),
+            ("input_token_details", "cache_creation"),
+            ("cache_creation_input_tokens",),
+        )
+    if not usage.total_tokens and (usage.input_tokens or usage.output_tokens):
+        usage.total_tokens = usage.input_tokens + usage.output_tokens + usage.reasoning_tokens
+    return usage
+
+
+def estimate_context_tokens(messages: list, model: str = "") -> int:
+    return count_messages_tokens(_messages_for_count(messages), model)
+
+
+def estimate_message_tokens(message: object, model: str = "") -> int:
+    return count_messages_tokens([_message_for_count(message)], model)
+
+
+def format_token_count(value: int | None) -> str:
+    value = max(int(value or 0), 0)
+    if value >= 1_000_000:
+        return _format_scaled(value, 1_000_000, "m")
+    if value >= 1_000:
+        return _format_scaled(value, 1_000, "k")
+    return str(value)
+
+
+def format_cache_hit_rate(stats: UsageStats) -> str:
+    rate = stats.cache_hit_rate
+    if rate is None:
+        return "--"
+    percent = max(0, min(round(rate * 100), 100))
+    prefix = "~" if stats.cache_hit_rate_is_estimated else ""
+    return f"{prefix}{percent}%"
+
+
+def _format_scaled(value: int, divisor: int, suffix: str) -> str:
+    if value % divisor == 0:
+        return f"{value // divisor}{suffix}"
+    return f"{value / divisor:.1f}{suffix}"
+
+
+def _usage_sources(message: object) -> list[dict[str, Any]]:
+    sources: list[dict[str, Any]] = []
+    for value in (
+        getattr(message, "usage_metadata", None),
+        getattr(message, "response_metadata", None),
+        getattr(message, "additional_kwargs", None),
+    ):
+        if isinstance(value, dict):
+            sources.append(value)
+            for key in ("usage", "token_usage"):
+                nested = value.get(key)
+                if isinstance(nested, dict):
+                    sources.append(nested)
+    return sources
+
+
+def _first_int(source: dict[str, Any], *keys: str) -> int:
+    for key in keys:
+        value = source.get(key)
+        if isinstance(value, int):
+            return value
+    return 0
+
+
+def _nested_first_int(source: dict[str, Any], *paths: tuple[str, ...]) -> int:
+    for path in paths:
+        current: Any = source
+        for part in path:
+            if not isinstance(current, dict):
+                current = None
+                break
+            current = current.get(part)
+        if isinstance(current, int):
+            return current
+    return 0
+
+
+def _has_nested_key(source: dict[str, Any], *paths: tuple[str, ...]) -> bool:
+    for path in paths:
+        current: Any = source
+        for part in path:
+            if not isinstance(current, dict) or part not in current:
+                current = None
+                break
+            current = current[part]
+        if current is not None:
+            return True
+    return False
+
+
+def _messages_for_count(messages: list) -> list[dict[str, str]]:
+    return [_message_for_count(message) for message in messages]
+
+
+def _common_prefix_context(
+    previous_context: list[dict[str, str]],
+    current_context: list[dict[str, str]],
+) -> list[dict[str, str]]:
+    common: list[dict[str, str]] = []
+    for previous, current in zip(previous_context, current_context):
+        if previous != current:
+            break
+        common.append(current)
+    return common
+
+
+def _message_for_count(message: object) -> dict[str, str]:
+    role = getattr(message, "type", "") or message.__class__.__name__.removesuffix("Message").lower()
+    parts = [str(_message_content_for_count(message))]
+    tool_calls = getattr(message, "tool_calls", None)
+    if tool_calls:
+        parts.append(str(tool_calls))
+    return {"role": role, "content": "\n".join(parts)}
+
+
+def _message_content_for_count(message: object) -> object:
+    if isinstance(message, dict):
+        return message.get("content", "")
+    return _content_for_count(getattr(message, "content", str(message)))
+
+
+def _content_for_count(content: object) -> str:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for item in content:
+            if isinstance(item, str):
+                parts.append(item)
+            elif isinstance(item, dict):
+                if item.get("type") == "text":
+                    parts.append(str(item.get("text", "")))
+                elif item.get("type") in {"image", "image_url"}:
+                    parts.append("[image]")
+                else:
+                    parts.append(str(item))
+            else:
+                parts.append(str(item))
+        return "\n".join(parts)
+    return str(content)

voidx/lsp/__init__.py

@@ -0,0 +1,30 @@
+"""Language Server Protocol integration."""
+
+from __future__ import annotations
+
+from voidx.lsp.config import LSP_CONFIG_FILE, load_lsp_servers
+from voidx.lsp.manager import LspManager
+from voidx.lsp.schema import (
+    LspDiagnostic,
+    LspDoctorCheck,
+    LspLocation,
+    LspPosition,
+    LspRange,
+    LspRuntimeStatus,
+    LspServerConfig,
+    LspSymbol,
+)
+
+__all__ = [
+    "LSP_CONFIG_FILE",
+    "LspDiagnostic",
+    "LspDoctorCheck",
+    "LspLocation",
+    "LspManager",
+    "LspPosition",
+    "LspRange",
+    "LspRuntimeStatus",
+    "LspServerConfig",
+    "LspSymbol",
+    "load_lsp_servers",
+]

voidx/lsp/client.py

@@ -0,0 +1,259 @@
+"""Async stdio JSON-RPC client for Language Server Protocol."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from collections.abc import Callable
+from typing import Any
+
+from voidx.lsp.errors import LspConnectionError, LspRequestError
+from voidx.lsp.schema import LspServerConfig, parse_diagnostics
+
+log = logging.getLogger(__name__)
+
+NotificationHandler = Callable[[str, dict[str, Any]], None]
+
+
+def encode_lsp_message(payload: dict[str, Any]) -> bytes:
+    body = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+    header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
+    return header + body
+
+
+async def read_lsp_message(reader: asyncio.StreamReader) -> dict[str, Any] | None:
+    headers: dict[str, str] = {}
+    while True:
+        line = await reader.readline()
+        if line == b"":
+            return None
+        if line in (b"\r\n", b"\n"):
+            break
+        key, sep, value = line.decode("ascii", errors="replace").partition(":")
+        if sep:
+            headers[key.strip().lower()] = value.strip()
+
+    raw_length = headers.get("content-length")
+    if raw_length is None:
+        raise LspConnectionError("LSP message missing Content-Length header")
+    try:
+        length = int(raw_length)
+    except ValueError as exc:
+        raise LspConnectionError(f"Invalid Content-Length: {raw_length}") from exc
+    body = await reader.readexactly(length)
+    return json.loads(body.decode("utf-8"))
+
+
+class LspClient:
+    def __init__(
+        self,
+        config: LspServerConfig,
+        *,
+        cwd: str,
+        notification_handler: NotificationHandler | None = None,
+    ) -> None:
+        self.config = config
+        self.cwd = cwd
+        self._notification_handler = notification_handler
+        self._process: asyncio.subprocess.Process | None = None
+        self._reader_task: asyncio.Task | None = None
+        self._stderr_task: asyncio.Task | None = None
+        self._pending: dict[int, asyncio.Future] = {}
+        self._next_id = 1
+        self._diagnostics: dict[str, list] = {}
+        self._capabilities: dict[str, Any] = {}
+        self._error_message = ""
+
+    @property
+    def pid(self) -> int | None:
+        return self._process.pid if self._process is not None else None
+
+    @property
+    def connected(self) -> bool:
+        return self._process is not None and self._process.returncode is None
+
+    @property
+    def status(self) -> str:
+        if self.connected:
+            return "connected"
+        if self._error_message:
+            return "error"
+        return "disconnected"
+
+    @property
+    def error_message(self) -> str:
+        return self._error_message
+
+    @property
+    def capabilities(self) -> dict[str, Any]:
+        return self._capabilities
+
+    async def start(self, *, root_uri: str, timeout: float = 10.0) -> None:
+        if self.connected:
+            return
+        try:
+            self._process = await asyncio.create_subprocess_exec(
+                self.config.command,
+                *self.config.args,
+                cwd=self.cwd,
+                env=os.environ.copy(),
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+        except OSError as exc:
+            self._error_message = str(exc)
+            raise LspConnectionError(f"Could not start {self.config.language} LSP: {exc}") from exc
+
+        self._reader_task = asyncio.create_task(self._read_loop())
+        self._stderr_task = asyncio.create_task(self._drain_stderr())
+        result = await self.request("initialize", {
+            "processId": None,
+            "rootUri": root_uri,
+            "capabilities": _client_capabilities(),
+            "workspaceFolders": [{"uri": root_uri, "name": "workspace"}],
+        }, timeout=timeout)
+        self._capabilities = result.get("capabilities", {}) if isinstance(result, dict) else {}
+        await self.notify("initialized", {})
+
+    async def stop(self) -> None:
+        if self._process is None:
+            return
+        if self.connected:
+            try:
+                await self.request("shutdown", None, timeout=2.0)
+            except Exception:
+                pass
+            try:
+                await self.notify("exit", {})
+            except Exception:
+                pass
+        try:
+            await asyncio.wait_for(self._process.wait(), timeout=2.0)
+        except asyncio.TimeoutError:
+            self._process.terminate()
+            try:
+                await asyncio.wait_for(self._process.wait(), timeout=2.0)
+            except asyncio.TimeoutError:
+                self._process.kill()
+                await self._process.wait()
+        self._process = None
+        self._cancel_tasks()
+
+    async def request(self, method: str, params: Any = None, *, timeout: float = 10.0) -> Any:
+        if self._process is None or self._process.stdin is None:
+            raise LspConnectionError("LSP client is not started")
+        if self._process.returncode is not None:
+            raise LspConnectionError(f"LSP server exited with code {self._process.returncode}")
+        req_id = self._next_id
+        self._next_id += 1
+        future = asyncio.get_running_loop().create_future()
+        self._pending[req_id] = future
+        await self._send({"jsonrpc": "2.0", "id": req_id, "method": method, "params": params})
+        try:
+            response = await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError as exc:
+            self._pending.pop(req_id, None)
+            raise LspConnectionError(f"LSP request timed out: {method}") from exc
+        if isinstance(response, dict) and response.get("error"):
+            error = response["error"]
+            message = error.get("message", str(error)) if isinstance(error, dict) else str(error)
+            raise LspRequestError(f"{method}: {message}")
+        return response.get("result") if isinstance(response, dict) else response
+
+    async def notify(self, method: str, params: Any = None) -> None:
+        await self._send({"jsonrpc": "2.0", "method": method, "params": params})
+
+    def diagnostics_for(self, uri: str) -> list:
+        return list(self._diagnostics.get(uri, []))
+
+    def all_diagnostics(self) -> list:
+        result: list = []
+        for diagnostics in self._diagnostics.values():
+            result.extend(diagnostics)
+        return result
+
+    async def _send(self, payload: dict[str, Any]) -> None:
+        if self._process is None or self._process.stdin is None:
+            raise LspConnectionError("LSP client is not started")
+        self._process.stdin.write(encode_lsp_message(payload))
+        await self._process.stdin.drain()
+
+    async def _read_loop(self) -> None:
+        assert self._process is not None and self._process.stdout is not None
+        try:
+            while True:
+                message = await read_lsp_message(self._process.stdout)
+                if message is None:
+                    break
+                self._handle_message(message)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            self._error_message = str(exc)
+            log.debug("LSP read loop failed for %s: %s", self.config.language, exc)
+        finally:
+            for future in self._pending.values():
+                if not future.done():
+                    future.set_exception(LspConnectionError("LSP connection closed"))
+            self._pending.clear()
+
+    def _handle_message(self, message: dict[str, Any]) -> None:
+        msg_id = message.get("id")
+        if msg_id in self._pending:
+            future = self._pending.pop(msg_id)
+            if not future.done():
+                future.set_result(message)
+            return
+
+        method = message.get("method")
+        params = message.get("params")
+        if method == "textDocument/publishDiagnostics" and isinstance(params, dict):
+            uri = params.get("uri")
+            if isinstance(uri, str):
+                self._diagnostics[uri] = parse_diagnostics(uri, params.get("diagnostics", []))
+        if isinstance(method, str) and isinstance(params, dict) and self._notification_handler:
+            self._notification_handler(method, params)
+
+    async def _drain_stderr(self) -> None:
+        assert self._process is not None and self._process.stderr is not None
+        try:
+            while True:
+                line = await self._process.stderr.readline()
+                if not line:
+                    break
+                if not self._error_message:
+                    self._error_message = line.decode("utf-8", errors="replace").strip()
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            pass
+
+    def _cancel_tasks(self) -> None:
+        for task in (self._reader_task, self._stderr_task):
+            if task is not None and not task.done():
+                task.cancel()
+        self._reader_task = None
+        self._stderr_task = None
+
+
+def _client_capabilities() -> dict[str, Any]:
+    return {
+        "textDocument": {
+            "documentSymbol": {"hierarchicalDocumentSymbolSupport": True},
+            "definition": {"linkSupport": True},
+            "references": {},
+            "formatting": {},
+            "publishDiagnostics": {"relatedInformation": True},
+            "synchronization": {
+                "didSave": True,
+                "dynamicRegistration": False,
+            },
+        },
+        "workspace": {
+            "symbol": {"resolveSupport": {"properties": []}},
+            "workspaceFolders": True,
+        },
+    }