vantage-litellm-callback 0.0.3__py3-none-any.whl

vantage_callback.py

@@ -0,0 +1,3 @@
+from vantage_litellm_callback import callback_instance
+
+__all__ = ["callback_instance"]

vantage_litellm_callback/__init__.py

@@ -0,0 +1,11 @@
+from .callback import VantageLiteLLMCallback
+from .delivery import DeliveryGapTracker, DeliveryWorker
+
+callback_instance = VantageLiteLLMCallback()
+
+__all__ = [
+    "DeliveryGapTracker",
+    "DeliveryWorker",
+    "VantageLiteLLMCallback",
+    "callback_instance",
+]

vantage_litellm_callback/callback.py

@@ -0,0 +1,361 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections import OrderedDict
+from collections.abc import AsyncGenerator, AsyncIterable, Mapping
+from dataclasses import dataclass
+from datetime import datetime
+from time import monotonic
+from typing import Any
+from uuid import uuid4
+
+from litellm.integrations.custom_logger import CustomLogger
+
+from .client import EventSink, UnixSocketClient
+from .delivery import DeliveryGapTracker, DeliveryWorker
+from .projection import (
+    is_terminal_stream_chunk,
+    merge_usage,
+    project_event,
+    resolve_event_timestamp,
+    usage_from_standard_logging,
+    usage_from_stream_chunk,
+)
+from .schema import Usage
+
+ATTEMPT_TTL_SECONDS = 60 * 60
+MAX_TRACKED_ATTEMPTS = 10_000
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class _Attempt:
+    event_id: str
+    created_at: float
+    request_start_time: datetime
+
+
+class VantageLiteLLMCallback(CustomLogger):
+    def __init__(
+        self,
+        *,
+        event_sink: EventSink | None = None,
+        queue_size: int | None = None,
+        recovery_probe_seconds: float | None = None,
+        gap_tracker: DeliveryGapTracker | None = None,
+    ) -> None:
+        super().__init__(turn_off_message_logging=True)
+        self._delivery = DeliveryWorker(
+            event_sink or UnixSocketClient(),
+            queue_size=queue_size,
+            recovery_probe_seconds=recovery_probe_seconds,
+            gap_tracker=gap_tracker,
+        )
+        self._attempts: OrderedDict[str, _Attempt] = OrderedDict()
+
+    @property
+    def gap_tracker(self) -> DeliveryGapTracker:
+        return self._delivery.gaps
+
+    async def wait_for_delivery(self) -> None:
+        await self._delivery.wait_for_idle()
+
+    async def close(self) -> None:
+        await self._delivery.close()
+
+    async def async_pre_call_deployment_hook(
+        self, kwargs: dict[str, Any], call_type: Any
+    ) -> dict[str, Any]:
+        del call_type
+        event_id = f"lta_{uuid4().hex}"
+        try:
+            request_start_time = resolve_event_timestamp(kwargs)
+            call_id = _litellm_call_id(kwargs)
+            if call_id is not None:
+                self._remember_attempt(call_id, event_id, request_start_time)
+            event = project_event(
+                kwargs,
+                event_id=event_id,
+                status="started",
+                request_start_time=request_start_time,
+            )
+            self._delivery.enqueue(event)
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - callback hooks must fail open
+            self._record_gap("projection", event_id=event_id, status="started")
+        return kwargs
+
+    async def async_post_call_success_deployment_hook(
+        self,
+        request_data: dict[str, Any],
+        response: Any,
+        call_type: Any,
+    ) -> Any:
+        del call_type
+        attempt: _Attempt | None = None
+        try:
+            attempt = self._take_attempt(request_data)
+            if attempt is None:
+                return response
+            event = project_event(
+                request_data,
+                event_id=attempt.event_id,
+                status="success",
+                response=response,
+                request_start_time=attempt.request_start_time,
+            )
+            self._delivery.enqueue(event)
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - callback hooks must fail open
+            event_id = attempt.event_id if attempt is not None else None
+            self._record_gap("projection", event_id=event_id, status="success")
+        return response
+
+    async def async_log_success_event(
+        self,
+        kwargs: dict[str, Any],
+        response_obj: Any,
+        start_time: Any,
+        end_time: Any,
+    ) -> None:
+        """Close attempts the deployment and iterator hooks could not.
+
+        LiteLLM builds ``standard_logging_object`` only after those hooks
+        fire, and hands it exclusively to this hook (``kwargs`` is the logging
+        object's ``model_call_details``). The LiteLLM call ID finds the
+        in-memory attempt. Removing that attempt keeps this a no-op whenever
+        a primary hook already emitted.
+        """
+        del start_time, end_time
+        attempt: _Attempt | None = None
+        try:
+            attempt = self._take_attempt(kwargs)
+            if attempt is None:
+                return
+            standard_logging_object = kwargs.get("standard_logging_object")
+            event = project_event(
+                kwargs,
+                event_id=attempt.event_id,
+                status="success",
+                response=response_obj,
+                usage_override=usage_from_standard_logging(standard_logging_object),
+                standard_logging_object=standard_logging_object,
+                request_start_time=attempt.request_start_time,
+            )
+            self._delivery.enqueue(event)
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - callback hooks must fail open
+            event_id = attempt.event_id if attempt is not None else None
+            self._record_gap("projection", event_id=event_id, status="success")
+
+    async def async_post_call_failure_hook(
+        self,
+        request_data: dict[str, Any],
+        original_exception: Exception,
+        user_api_key_dict: Any,
+        traceback_str: str | None = None,
+    ) -> None:
+        del original_exception, user_api_key_dict, traceback_str
+        attempt: _Attempt | None = None
+        try:
+            attempt = self._take_attempt(request_data)
+            if attempt is None:
+                return
+            event = project_event(
+                request_data,
+                event_id=attempt.event_id,
+                status="failure",
+                request_start_time=attempt.request_start_time,
+            )
+            self._delivery.enqueue(event)
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - callback hooks must fail open
+            event_id = attempt.event_id if attempt is not None else None
+            self._record_gap("projection", event_id=event_id, status="failure")
+
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: Any,
+        response: AsyncIterable[object],
+        request_data: dict[str, Any],
+    ) -> AsyncGenerator[object, None]:
+        del user_api_key_dict
+        try:
+            attempt = self._find_attempt(request_data)
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - callback hooks must fail open
+            attempt = None
+        if attempt is None:
+            async for chunk in response:
+                yield chunk
+            return
+
+        terminal_seen = False
+        usage = Usage()
+        try:
+            async for chunk in response:
+                try:
+                    usage = merge_usage(usage, usage_from_stream_chunk(chunk))
+                    if not terminal_seen and is_terminal_stream_chunk(chunk):
+                        terminal_seen = True
+                        status = _stream_terminal_status(chunk)
+                        terminal_attempt = self._take_attempt(
+                            request_data, expected=attempt
+                        )
+                        if terminal_attempt is None:
+                            yield chunk
+                            continue
+                        event = project_event(
+                            request_data,
+                            event_id=terminal_attempt.event_id,
+                            status=status,
+                            response=chunk,
+                            usage_override=usage,
+                            request_start_time=terminal_attempt.request_start_time,
+                        )
+                        self._delivery.enqueue(event)
+                except asyncio.CancelledError:
+                    raise
+                except Exception:  # noqa: BLE001 - callback hooks must fail open
+                    self._record_gap(
+                        "projection",
+                        event_id=attempt.event_id,
+                        status="success",
+                    )
+                yield chunk
+        finally:
+            if not terminal_seen:
+                self._close_unterminated_stream(request_data, attempt, usage)
+
+    def _close_unterminated_stream(
+        self,
+        request_data: Mapping[str, Any],
+        attempt: _Attempt,
+        usage: Usage,
+    ) -> None:
+        """Close a stream that ended without a terminal chunk.
+
+        With no usage accumulated the attempt stays in the map so
+        ``async_log_success_event`` can close it with LiteLLM's assembled
+        usage, or the failure hook can close an error. With usage in hand
+        this emits success itself.
+        """
+        if usage == Usage():
+            logger.warning(
+                "Vantage callback: stream ended without a terminal marker or "
+                "usage; deferring to the success-event safety net"
+            )
+            return
+
+        logger.warning(
+            "Vantage callback: stream ended without a terminal marker; "
+            "emitting success with usage merged from chunks"
+        )
+        try:
+            terminal_attempt = self._take_attempt(request_data, expected=attempt)
+            if terminal_attempt is None:
+                return
+            event = project_event(
+                request_data,
+                event_id=terminal_attempt.event_id,
+                status="success",
+                usage_override=usage,
+                request_start_time=terminal_attempt.request_start_time,
+            )
+            self._delivery.enqueue(event)
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - callback hooks must fail open
+            self._record_gap(
+                "projection",
+                event_id=attempt.event_id,
+                status="success",
+            )
+
+    def _record_gap(
+        self,
+        category: str,
+        *,
+        event_id: str | None,
+        status: str,
+    ) -> None:
+        try:
+            self._delivery.record_gap(
+                category,  # type: ignore[arg-type]
+                event_id=event_id,
+                status=status,  # type: ignore[arg-type]
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception:  # noqa: BLE001 - gap tracking must fail open
+            # Gap accounting itself must not affect provider traffic.
+            return
+
+    def _remember_attempt(
+        self, call_id: str, event_id: str, request_start_time: datetime
+    ) -> None:
+        now = monotonic()
+        self._evict_attempts(now)
+        self._attempts.pop(call_id, None)
+        while len(self._attempts) >= MAX_TRACKED_ATTEMPTS:
+            self._attempts.popitem(last=False)
+        self._attempts[call_id] = _Attempt(
+            event_id=event_id,
+            created_at=now,
+            request_start_time=request_start_time,
+        )
+
+    def _find_attempt(self, request_data: Mapping[str, Any]) -> _Attempt | None:
+        self._evict_attempts(monotonic())
+        call_id = _litellm_call_id(request_data)
+        if call_id is None:
+            return None
+        return self._attempts.get(call_id)
+
+    def _take_attempt(
+        self,
+        request_data: Mapping[str, Any],
+        *,
+        expected: _Attempt | None = None,
+    ) -> _Attempt | None:
+        attempt = self._find_attempt(request_data)
+        if attempt is None or (expected is not None and attempt is not expected):
+            return None
+        call_id = _litellm_call_id(request_data)
+        if call_id is None:
+            return None
+        return self._attempts.pop(call_id, None)
+
+    def _evict_attempts(self, now: float) -> None:
+        while self._attempts:
+            _, attempt = next(iter(self._attempts.items()))
+            if now - attempt.created_at < ATTEMPT_TTL_SECONDS:
+                break
+            self._attempts.popitem(last=False)
+
+
+def _litellm_call_id(request_data: Mapping[str, Any]) -> str | None:
+    call_id = request_data.get("litellm_call_id")
+    if isinstance(call_id, str) and call_id:
+        return call_id
+    return None
+
+
+def _stream_terminal_status(response: object) -> str:
+    if isinstance(response, Mapping):
+        status = response.get("status")
+        event_type = response.get("type")
+        if (
+            status in {"failed", "cancelled", "incomplete"}
+            or event_type in {"response.failed", "response.incomplete"}
+            or response.get("error") is not None
+        ):
+            return "failure"
+    return "success"

vantage_litellm_callback/client.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from dataclasses import dataclass
+from typing import Protocol
+
+from pydantic import ValidationError
+
+from .schema import DeliveryGap, DeliveryGapAck, Frame, GapCategory, UsageAck
+
+DEFAULT_SOCKET_PATH = "/var/run/vantage-collector/collector.sock"
+DEFAULT_ACK_TIMEOUT_SECONDS = 5.0
+DEFAULT_CONNECTION_POOL_SIZE = 16
+MAX_ACK_BYTES = 4096
+
+
+class DeliveryError(RuntimeError):
+    def __init__(
+        self,
+        message: str,
+        *,
+        category: GapCategory = "unknown",
+    ) -> None:
+        super().__init__(message)
+        self.category = category
+
+
+class EventSink(Protocol):
+    async def send(self, frame: Frame) -> None: ...
+
+
+@dataclass(slots=True)
+class _Connection:
+    reader: asyncio.StreamReader | None = None
+    writer: asyncio.StreamWriter | None = None
+
+    async def close(self) -> None:
+        writer = self.writer
+        self.reader = None
+        self.writer = None
+        if writer is None:
+            return
+        writer.close()
+        await writer.wait_closed()
+
+
+class UnixSocketClient:
+    def __init__(
+        self,
+        socket_path: str | None = None,
+        ack_timeout_seconds: float | None = None,
+        connection_pool_size: int | None = None,
+    ) -> None:
+        self.socket_path = socket_path or os.getenv(
+            "VANTAGE_COLLECTOR_SOCKET_PATH", DEFAULT_SOCKET_PATH
+        )
+        self.ack_timeout_seconds = ack_timeout_seconds or _positive_float_from_env(
+            "VANTAGE_COLLECTOR_ACK_TIMEOUT_SECONDS",
+            DEFAULT_ACK_TIMEOUT_SECONDS,
+        )
+        pool_size = connection_pool_size or _positive_int_from_env(
+            "VANTAGE_COLLECTOR_CONNECTION_POOL_SIZE",
+            DEFAULT_CONNECTION_POOL_SIZE,
+        )
+        self._connections = asyncio.LifoQueue[_Connection](maxsize=pool_size)
+        for connection in tuple(_Connection() for _ in range(pool_size)):
+            self._connections.put_nowait(connection)
+
+    async def send(self, frame: Frame) -> None:
+        connection = await self._connections.get()
+        try:
+            await self._send_on_connection(connection, frame)
+        except asyncio.CancelledError:
+            await connection.close()
+            raise
+        except asyncio.TimeoutError as error:
+            await connection.close()
+            raise DeliveryError(
+                f"Vantage collector acknowledgement timed out for {_frame_id(frame)}",
+                category="timeout",
+            ) from error
+        except OSError as error:
+            await connection.close()
+            raise DeliveryError(
+                f"Vantage collector socket unavailable for {_frame_id(frame)}",
+                category="socket_unavailable",
+            ) from error
+        except CollectorRejection as error:
+            await connection.close()
+            raise DeliveryError(
+                f"Vantage collector rejected {_frame_id(frame)}",
+                category="collector_rejection",
+            ) from error
+        except (ValidationError, UnicodeDecodeError, ValueError) as error:
+            await connection.close()
+            raise DeliveryError(
+                f"Invalid Vantage collector acknowledgement for {_frame_id(frame)}",
+                category="protocol_ack",
+            ) from error
+        finally:
+            self._connections.put_nowait(connection)
+
+    async def close(self) -> None:
+        connections = tuple(
+            await asyncio.gather(
+                *(self._connections.get() for _ in range(self._connections.maxsize))
+            )
+        )
+        await asyncio.gather(*(connection.close() for connection in connections))
+        for connection in connections:
+            self._connections.put_nowait(connection)
+
+    async def _send_on_connection(self, connection: _Connection, frame: Frame) -> None:
+        if (
+            connection.reader is None
+            or connection.writer is None
+            or connection.writer.is_closing()
+        ):
+            connection.reader, connection.writer = await asyncio.wait_for(
+                asyncio.open_unix_connection(self.socket_path),
+                timeout=self.ack_timeout_seconds,
+            )
+
+        payload = (
+            frame.model_dump_json(exclude_none=True, by_alias=True).encode("utf-8")
+            + b"\n"
+        )
+        connection.writer.write(payload)
+        await asyncio.wait_for(
+            connection.writer.drain(), timeout=self.ack_timeout_seconds
+        )
+        ack_line = await asyncio.wait_for(
+            connection.reader.readline(),
+            timeout=self.ack_timeout_seconds,
+        )
+        if (
+            not ack_line
+            or len(ack_line) > MAX_ACK_BYTES
+            or not ack_line.endswith(b"\n")
+        ):
+            raise ValueError("invalid collector acknowledgement frame")
+
+        if isinstance(frame, DeliveryGap):
+            ack = DeliveryGapAck.model_validate_json(ack_line)
+            expected_id = frame.gap_id
+            acknowledged_id = ack.gap_id
+        else:
+            ack = UsageAck.model_validate_json(ack_line)
+            expected_id = frame.event_id
+            acknowledged_id = ack.event_id
+        if ack.error:
+            raise CollectorRejection(ack.error)
+        if acknowledged_id != expected_id:
+            raise ValueError("collector acknowledgement ID mismatch")
+        if ack.collector_sequence is None:
+            raise ValueError("collector acknowledgement omitted durable sequence")
+
+
+class CollectorRejection(ValueError):
+    pass
+
+
+def _frame_id(frame: Frame) -> str:
+    return frame.gap_id if isinstance(frame, DeliveryGap) else frame.event_id
+
+
+def _positive_float_from_env(name: str, default: float) -> float:
+    raw_value = os.getenv(name)
+    if raw_value is None:
+        return default
+    value = float(raw_value)
+    if value <= 0:
+        raise ValueError(f"{name} must be greater than zero")
+    return value
+
+
+def _positive_int_from_env(name: str, default: int) -> int:
+    raw_value = os.getenv(name)
+    if raw_value is None:
+        return default
+    value = int(raw_value)
+    if value <= 0:
+        raise ValueError(f"{name} must be greater than zero")
+    return value