aitracer 0.1.0__py3-none-any.whl

python/queue.py

@@ -0,0 +1,219 @@
+"""Log queue for batching and async sending."""
+
+from __future__ import annotations
+
+import json
+import logging
+import queue
+import threading
+import time
+from typing import Any, Optional
+
+import httpx
+
+logger = logging.getLogger("aitracer")
+
+
+class LogQueue:
+    """
+    Queue for batching and asynchronously sending logs to AITracer API.
+
+    Features:
+    - Batches logs for efficient sending
+    - Background thread for async sending
+    - Automatic flush on interval
+    - Graceful shutdown with flush
+    - Sync mode for serverless environments
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str,
+        batch_size: int = 10,
+        flush_interval: float = 5.0,
+        sync: bool = False,
+        max_queue_size: int = 10000,
+        timeout: float = 30.0,
+    ):
+        """
+        Initialize the log queue.
+
+        Args:
+            api_key: AITracer API key.
+            base_url: AITracer API base URL.
+            batch_size: Number of logs to batch before sending.
+            flush_interval: Seconds between automatic flushes.
+            sync: If True, send logs synchronously.
+            max_queue_size: Maximum queue size before dropping logs.
+            timeout: HTTP request timeout in seconds.
+        """
+        self.api_key = api_key
+        self.base_url = base_url
+        self.batch_size = batch_size
+        self.flush_interval = flush_interval
+        self.sync = sync
+        self.max_queue_size = max_queue_size
+        self.timeout = timeout
+
+        self._queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=max_queue_size)
+        self._shutdown = threading.Event()
+        self._worker_thread: Optional[threading.Thread] = None
+
+        # HTTP client
+        self._client = httpx.Client(
+            base_url=base_url,
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=timeout,
+        )
+
+        # Start background worker if async mode
+        if not sync:
+            self._start_worker()
+
+    def _start_worker(self) -> None:
+        """Start the background worker thread."""
+        self._worker_thread = threading.Thread(
+            target=self._worker_loop,
+            daemon=True,
+            name="aitracer-worker",
+        )
+        self._worker_thread.start()
+
+    def _worker_loop(self) -> None:
+        """Background worker loop for processing queued logs."""
+        batch: list[dict[str, Any]] = []
+        last_flush = time.time()
+
+        while not self._shutdown.is_set():
+            try:
+                # Try to get an item with timeout
+                try:
+                    item = self._queue.get(timeout=0.1)
+                    batch.append(item)
+                    self._queue.task_done()
+                except queue.Empty:
+                    pass
+
+                # Check if we should flush
+                should_flush = (
+                    len(batch) >= self.batch_size
+                    or (batch and time.time() - last_flush >= self.flush_interval)
+                )
+
+                if should_flush and batch:
+                    self._send_batch(batch)
+                    batch = []
+                    last_flush = time.time()
+
+            except Exception as e:
+                logger.exception("Error in worker loop: %s", e)
+                time.sleep(1)
+
+        # Final flush on shutdown
+        if batch:
+            self._send_batch(batch)
+
+        # Drain remaining items
+        while not self._queue.empty():
+            try:
+                item = self._queue.get_nowait()
+                batch.append(item)
+                self._queue.task_done()
+                if len(batch) >= self.batch_size:
+                    self._send_batch(batch)
+                    batch = []
+            except queue.Empty:
+                break
+
+        if batch:
+            self._send_batch(batch)
+
+    def _send_batch(self, batch: list[dict[str, Any]]) -> None:
+        """Send a batch of logs to the API."""
+        if not batch:
+            return
+
+        try:
+            response = self._client.post(
+                "/v1/logs/batch",
+                json={"logs": batch},
+            )
+            response.raise_for_status()
+            logger.debug("Sent %d logs successfully", len(batch))
+        except httpx.HTTPStatusError as e:
+            logger.error(
+                "Failed to send logs: HTTP %d - %s",
+                e.response.status_code,
+                e.response.text,
+            )
+        except httpx.RequestError as e:
+            logger.error("Failed to send logs: %s", e)
+        except Exception as e:
+            logger.exception("Unexpected error sending logs: %s", e)
+
+    def _send_single(self, log: dict[str, Any]) -> None:
+        """Send a single log synchronously."""
+        try:
+            response = self._client.post(
+                "/v1/logs",
+                json=log,
+            )
+            response.raise_for_status()
+            logger.debug("Sent log successfully")
+        except httpx.HTTPStatusError as e:
+            logger.error(
+                "Failed to send log: HTTP %d - %s",
+                e.response.status_code,
+                e.response.text,
+            )
+        except httpx.RequestError as e:
+            logger.error("Failed to send log: %s", e)
+        except Exception as e:
+            logger.exception("Unexpected error sending log: %s", e)
+
+    def add(self, log: dict[str, Any]) -> None:
+        """
+        Add a log to the queue.
+
+        Args:
+            log: Log entry dictionary.
+        """
+        if self.sync:
+            # Sync mode: send immediately
+            self._send_single(log)
+        else:
+            # Async mode: add to queue
+            try:
+                self._queue.put_nowait(log)
+            except queue.Full:
+                logger.warning("Log queue is full, dropping log")
+
+    def flush(self) -> None:
+        """Flush all pending logs immediately."""
+        if self.sync:
+            return  # Nothing to flush in sync mode
+
+        # Collect all items
+        batch: list[dict[str, Any]] = []
+        while not self._queue.empty():
+            try:
+                item = self._queue.get_nowait()
+                batch.append(item)
+                self._queue.task_done()
+            except queue.Empty:
+                break
+
+        # Send in batches
+        for i in range(0, len(batch), self.batch_size):
+            self._send_batch(batch[i : i + self.batch_size])
+
+    def shutdown(self) -> None:
+        """Shutdown the queue and wait for pending logs."""
+        self._shutdown.set()
+        if self._worker_thread:
+            self._worker_thread.join(timeout=10.0)
+        self._client.close()

python/session.py

@@ -0,0 +1,144 @@
+"""Session tracking for AITracer."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from aitracer.client import AITracer
+
+
+@dataclass
+class Session:
+    """
+    Session for tracking user interactions over time.
+
+    A session groups multiple logs/traces that belong to the same user interaction flow.
+
+    Usage:
+        with tracer.session(user_id="user-123") as session:
+            response = client.chat.completions.create(...)
+            session.event("user_action", data={"action": "clicked_button"})
+            session.feedback(log_id="...", type="thumbs_up")
+    """
+
+    session_id: str
+    user_id: Optional[str] = None
+    name: Optional[str] = None
+    metadata: dict = field(default_factory=dict)
+    tags: list[str] = field(default_factory=list)
+
+    # Internal
+    _tracer: Optional["AITracer"] = field(default=None, repr=False)
+    _started_at: datetime = field(default_factory=datetime.utcnow)
+    _events: list[dict] = field(default_factory=list)
+    _feedbacks: list[dict] = field(default_factory=list)
+    _last_log_id: Optional[str] = field(default=None, repr=False)
+
+    def set_metadata(self, key: str, value: Any) -> None:
+        """Set a metadata value."""
+        self.metadata[key] = value
+
+    def add_tag(self, tag: str) -> None:
+        """Add a tag to the session."""
+        if tag not in self.tags:
+            self.tags.append(tag)
+
+    def event(
+        self,
+        event_type: str,
+        *,
+        name: Optional[str] = None,
+        data: Optional[dict] = None,
+        log_id: Optional[str] = None,
+    ) -> None:
+        """
+        Record a custom event in the session.
+
+        Args:
+            event_type: Type of event (e.g., "message", "response", "tool_call", "error", "custom")
+            name: Optional event name
+            data: Optional event data
+            log_id: Optional associated log ID
+        """
+        event = {
+            "event_type": event_type,
+            "name": name,
+            "data": data or {},
+            "log_id": log_id or self._last_log_id,
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+        }
+        self._events.append(event)
+
+        # Send event immediately if tracer is available
+        if self._tracer and self._tracer.enabled:
+            self._tracer._send_session_event(self.session_id, event)
+
+    def feedback(
+        self,
+        feedback_type: str,
+        *,
+        log_id: Optional[str] = None,
+        score: Optional[int] = None,
+        comment: Optional[str] = None,
+        tags: Optional[list[str]] = None,
+    ) -> None:
+        """
+        Record user feedback for an AI response.
+
+        Args:
+            feedback_type: Type of feedback ("thumbs_up", "thumbs_down", "rating", "text")
+            log_id: Log ID the feedback relates to. Defaults to last log in session.
+            score: Numeric score (1-5 for rating, 1/-1 for thumbs)
+            comment: Text comment
+            tags: Feedback tags (e.g., ["helpful", "incorrect"])
+        """
+        # Auto-set score for thumbs feedback
+        if score is None:
+            if feedback_type == "thumbs_up":
+                score = 1
+            elif feedback_type == "thumbs_down":
+                score = -1
+
+        feedback = {
+            "feedback_type": feedback_type,
+            "log_id": log_id or self._last_log_id,
+            "score": score,
+            "comment": comment,
+            "tags": tags or [],
+            "user_id": self.user_id,
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+        }
+        self._feedbacks.append(feedback)
+
+        # Send feedback immediately if tracer is available
+        if self._tracer and self._tracer.enabled:
+            self._tracer._send_feedback(self.session_id, feedback)
+
+    def thumbs_up(self, log_id: Optional[str] = None, comment: Optional[str] = None) -> None:
+        """Record a thumbs up feedback."""
+        self.feedback("thumbs_up", log_id=log_id, comment=comment)
+
+    def thumbs_down(self, log_id: Optional[str] = None, comment: Optional[str] = None) -> None:
+        """Record a thumbs down feedback."""
+        self.feedback("thumbs_down", log_id=log_id, comment=comment)
+
+    def rate(self, score: int, log_id: Optional[str] = None, comment: Optional[str] = None) -> None:
+        """
+        Record a rating feedback (1-5).
+
+        Args:
+            score: Rating from 1 to 5
+            log_id: Optional log ID
+            comment: Optional comment
+        """
+        if not 1 <= score <= 5:
+            raise ValueError("Score must be between 1 and 5")
+        self.feedback("rating", log_id=log_id, score=score, comment=comment)
+
+    def _set_last_log_id(self, log_id: str) -> None:
+        """Internal: Set the last log ID for feedback association."""
+        self._last_log_id = log_id

python/trace.py

@@ -0,0 +1,65 @@
+"""Trace context for grouping related API calls."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+@dataclass
+class Trace:
+    """
+    Trace context for grouping related API calls.
+
+    Usage:
+        with tracer.trace("user-query-123") as trace:
+            response1 = client.chat.completions.create(...)
+            trace.set_metadata({"user_id": "user-456"})
+            trace.add_tag("production")
+    """
+
+    trace_id: str
+    name: Optional[str] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    tags: list[str] = field(default_factory=list)
+
+    def set_metadata(self, metadata: dict[str, Any]) -> "Trace":
+        """
+        Set metadata for this trace.
+
+        Args:
+            metadata: Dictionary of key-value pairs.
+
+        Returns:
+            Self for chaining.
+        """
+        self.metadata.update(metadata)
+        return self
+
+    def add_tag(self, tag: str) -> "Trace":
+        """
+        Add a tag to this trace.
+
+        Args:
+            tag: Tag string.
+
+        Returns:
+            Self for chaining.
+        """
+        if tag not in self.tags:
+            self.tags.append(tag)
+        return self
+
+    def add_tags(self, tags: list[str]) -> "Trace":
+        """
+        Add multiple tags to this trace.
+
+        Args:
+            tags: List of tag strings.
+
+        Returns:
+            Self for chaining.
+        """
+        for tag in tags:
+            self.add_tag(tag)
+        return self

python/wrappers/__init__.py

@@ -0,0 +1,7 @@
+"""LLM client wrappers."""
+
+from aitracer.wrappers.openai_wrapper import wrap_openai_client
+from aitracer.wrappers.anthropic_wrapper import wrap_anthropic_client
+from aitracer.wrappers.gemini_wrapper import wrap_gemini_model
+
+__all__ = ["wrap_openai_client", "wrap_anthropic_client", "wrap_gemini_model"]

python/wrappers/anthropic_wrapper.py

@@ -0,0 +1,208 @@
+"""Anthropic client wrapper for automatic logging."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Iterator
+
+if TYPE_CHECKING:
+    from anthropic import Anthropic
+    from anthropic.types import Message, RawMessageStreamEvent
+    from aitracer.client import AITracer
+
+
+def wrap_anthropic_client(client: "Anthropic", tracer: "AITracer") -> "Anthropic":
+    """
+    Wrap an Anthropic client to automatically log all API calls.
+
+    Args:
+        client: Anthropic client instance.
+        tracer: AITracer instance.
+
+    Returns:
+        Wrapped Anthropic client (same instance, modified in place).
+    """
+    # Store original method
+    original_create = client.messages.create
+
+    @wraps(original_create)
+    def wrapped_create(*args: Any, **kwargs: Any) -> Any:
+        """Wrapped messages.create method."""
+        start_time = time.time()
+        span_id = str(uuid.uuid4())
+
+        # Extract request data
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+        system = kwargs.get("system")
+        max_tokens = kwargs.get("max_tokens", 1024)
+        stream = kwargs.get("stream", False)
+
+        try:
+            response = original_create(*args, **kwargs)
+
+            if stream:
+                # Handle streaming response
+                return _wrap_stream_response(
+                    response=response,
+                    tracer=tracer,
+                    model=model,
+                    messages=messages,
+                    system=system,
+                    start_time=start_time,
+                    span_id=span_id,
+                )
+            else:
+                # Handle non-streaming response
+                latency_ms = int((time.time() - start_time) * 1000)
+                _log_message(
+                    tracer=tracer,
+                    model=model,
+                    messages=messages,
+                    system=system,
+                    response=response,
+                    latency_ms=latency_ms,
+                    span_id=span_id,
+                )
+                return response
+
+        except Exception as e:
+            # Log error
+            latency_ms = int((time.time() - start_time) * 1000)
+            tracer.log(
+                model=model,
+                provider="anthropic",
+                input_data=_build_input_data(messages, system),
+                output_data=None,
+                latency_ms=latency_ms,
+                status="error",
+                error_message=str(e),
+                span_id=span_id,
+            )
+            raise
+
+    # Replace method
+    client.messages.create = wrapped_create  # type: ignore
+
+    return client
+
+
+def _wrap_stream_response(
+    response: Iterator["RawMessageStreamEvent"],
+    tracer: "AITracer",
+    model: str,
+    messages: list,
+    system: Any,
+    start_time: float,
+    span_id: str,
+) -> Iterator["RawMessageStreamEvent"]:
+    """Wrap streaming response to log after completion."""
+    content_parts: list[str] = []
+    input_tokens = 0
+    output_tokens = 0
+
+    try:
+        for event in response:
+            # Accumulate content from text deltas
+            if hasattr(event, "type"):
+                if event.type == "content_block_delta":
+                    if hasattr(event, "delta") and hasattr(event.delta, "text"):
+                        content_parts.append(event.delta.text)
+                elif event.type == "message_delta":
+                    if hasattr(event, "usage"):
+                        output_tokens = getattr(event.usage, "output_tokens", 0)
+                elif event.type == "message_start":
+                    if hasattr(event, "message") and hasattr(event.message, "usage"):
+                        input_tokens = getattr(event.message.usage, "input_tokens", 0)
+
+            yield event
+
+        # Log after stream completes
+        latency_ms = int((time.time() - start_time) * 1000)
+        full_content = "".join(content_parts)
+
+        tracer.log(
+            model=model,
+            provider="anthropic",
+            input_data=_build_input_data(messages, system),
+            output_data={"content": full_content},
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            latency_ms=latency_ms,
+            status="success",
+            span_id=span_id,
+        )
+
+    except Exception as e:
+        latency_ms = int((time.time() - start_time) * 1000)
+        tracer.log(
+            model=model,
+            provider="anthropic",
+            input_data=_build_input_data(messages, system),
+            output_data=None,
+            latency_ms=latency_ms,
+            status="error",
+            error_message=str(e),
+            span_id=span_id,
+        )
+        raise
+
+
+def _log_message(
+    tracer: "AITracer",
+    model: str,
+    messages: list,
+    system: Any,
+    response: "Message",
+    latency_ms: int,
+    span_id: str,
+) -> None:
+    """Log a non-streaming message."""
+    # Extract response content
+    output_content = None
+    if response.content:
+        text_parts = []
+        for block in response.content:
+            if hasattr(block, "text"):
+                text_parts.append(block.text)
+        output_content = "".join(text_parts)
+
+    input_tokens = getattr(response.usage, "input_tokens", 0) if response.usage else 0
+    output_tokens = getattr(response.usage, "output_tokens", 0) if response.usage else 0
+
+    tracer.log(
+        model=model,
+        provider="anthropic",
+        input_data=_build_input_data(messages, system),
+        output_data={"content": output_content},
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        latency_ms=latency_ms,
+        status="success",
+        span_id=span_id,
+    )
+
+
+def _build_input_data(messages: list, system: Any) -> dict:
+    """Build input data dictionary."""
+    input_data: dict[str, Any] = {"messages": _serialize_messages(messages)}
+    if system:
+        input_data["system"] = system if isinstance(system, str) else str(system)
+    return input_data
+
+
+def _serialize_messages(messages: list) -> list[dict]:
+    """Serialize messages to JSON-compatible format."""
+    result = []
+    for msg in messages:
+        if isinstance(msg, dict):
+            result.append(msg)
+        elif hasattr(msg, "model_dump"):
+            result.append(msg.model_dump())
+        elif hasattr(msg, "__dict__"):
+            result.append(msg.__dict__)
+        else:
+            result.append({"content": str(msg)})
+    return result