aitracer 0.1.0__py3-none-any.whl

python/client.py

@@ -0,0 +1,437 @@
+"""AITracer main client."""
+
+from __future__ import annotations
+
+import atexit
+import os
+import threading
+import time
+import uuid
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Optional
+
+from aitracer.pii import PIIDetector
+from aitracer.queue import LogQueue
+from aitracer.session import Session
+from aitracer.trace import Trace
+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
+
+if TYPE_CHECKING:
+    from openai import OpenAI
+    from anthropic import Anthropic
+    from google.generativeai import GenerativeModel
+
+
+class AITracer:
+    """
+    AITracer client for monitoring AI/LLM applications.
+
+    Usage:
+        tracer = AITracer(api_key="at-xxxxxxxx", project="my-chatbot")
+        client = tracer.wrap_openai(OpenAI())
+
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    _instance: Optional["AITracer"] = None
+    _lock = threading.Lock()
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        project: Optional[str] = None,
+        *,
+        base_url: str = "https://api.aitracer.co",
+        sync: bool = False,
+        flush_on_exit: bool = True,
+        batch_size: int = 10,
+        flush_interval: float = 5.0,
+        enabled: bool = True,
+        # PII settings
+        pii_detection: bool = False,
+        pii_action: str = "mask",
+        pii_types: Optional[list[str]] = None,
+    ):
+        """
+        Initialize AITracer client.
+
+        Args:
+            api_key: AITracer API key. Can also be set via AITRACER_API_KEY env var.
+            project: Project name or ID. Can also be set via AITRACER_PROJECT env var.
+            base_url: AITracer API base URL.
+            sync: If True, send logs synchronously (useful for Lambda/serverless).
+            flush_on_exit: If True, flush pending logs on program exit.
+            batch_size: Number of logs to batch before sending.
+            flush_interval: Seconds between automatic flushes.
+            enabled: If False, disable all logging (useful for testing).
+            pii_detection: If True, enable PII detection and masking.
+            pii_action: Action for detected PII - "mask", "redact", "hash", "none".
+            pii_types: List of PII types to detect. Default: ["email", "phone", "credit_card", "ssn"].
+        """
+        self.api_key = api_key or os.environ.get("AITRACER_API_KEY")
+        self.project = project or os.environ.get("AITRACER_PROJECT")
+        self.base_url = base_url.rstrip("/")
+        self.sync = sync
+        self.enabled = enabled
+
+        # PII detection settings
+        self.pii_detection = pii_detection
+        self.pii_action = pii_action
+        self.pii_types = pii_types or ["email", "phone", "credit_card", "ssn"]
+        self._pii_detector: Optional[PIIDetector] = None
+        if self.pii_detection:
+            self._pii_detector = PIIDetector(
+                enabled_types=self.pii_types,
+                action=self.pii_action,
+            )
+
+        if not self.api_key:
+            raise ValueError(
+                "API key is required. Set it via api_key parameter or AITRACER_API_KEY env var."
+            )
+
+        if not self.project:
+            raise ValueError(
+                "Project is required. Set it via project parameter or AITRACER_PROJECT env var."
+            )
+
+        # Log queue for async sending
+        self._queue = LogQueue(
+            api_key=self.api_key,
+            base_url=self.base_url,
+            batch_size=batch_size,
+            flush_interval=flush_interval,
+            sync=sync,
+        )
+
+        # Current trace context (thread-local)
+        self._trace_context = threading.local()
+
+        # Current session context (thread-local)
+        self._session_context = threading.local()
+
+        # Register flush on exit
+        if flush_on_exit:
+            atexit.register(self.flush)
+
+        # Set as singleton instance
+        with AITracer._lock:
+            AITracer._instance = self
+
+    @classmethod
+    def get_instance(cls) -> Optional["AITracer"]:
+        """Get the current AITracer instance."""
+        return cls._instance
+
+    def wrap_openai(self, client: "OpenAI") -> "OpenAI":
+        """
+        Wrap an OpenAI client to automatically log all API calls.
+
+        Args:
+            client: OpenAI client instance.
+
+        Returns:
+            Wrapped OpenAI client.
+        """
+        return wrap_openai_client(client, self)
+
+    def wrap_anthropic(self, client: "Anthropic") -> "Anthropic":
+        """
+        Wrap an Anthropic client to automatically log all API calls.
+
+        Args:
+            client: Anthropic client instance.
+
+        Returns:
+            Wrapped Anthropic client.
+        """
+        return wrap_anthropic_client(client, self)
+
+    def wrap_gemini(self, model: "GenerativeModel") -> "GenerativeModel":
+        """
+        Wrap a Google Gemini GenerativeModel to automatically log all API calls.
+
+        Args:
+            model: GenerativeModel instance.
+
+        Returns:
+            Wrapped GenerativeModel.
+
+        Example:
+            import google.generativeai as genai
+            genai.configure(api_key="your-api-key")
+            model = genai.GenerativeModel("gemini-1.5-flash")
+            model = tracer.wrap_gemini(model)
+        """
+        return wrap_gemini_model(model, self)
+
+    @contextmanager
+    def trace(
+        self,
+        trace_id: Optional[str] = None,
+        name: Optional[str] = None,
+    ):
+        """
+        Create a trace context for grouping related API calls.
+
+        Usage:
+            with tracer.trace("user-query-123") as trace:
+                response1 = client.chat.completions.create(...)
+                response2 = client.chat.completions.create(...)
+                trace.set_metadata({"user_id": "user-456"})
+
+        Args:
+            trace_id: Optional trace ID. Auto-generated if not provided.
+            name: Optional name for the trace.
+
+        Yields:
+            Trace object for setting metadata.
+        """
+        trace = Trace(
+            trace_id=trace_id or str(uuid.uuid4()),
+            name=name,
+        )
+
+        # Store in thread-local context
+        previous_trace = getattr(self._trace_context, "current", None)
+        self._trace_context.current = trace
+
+        try:
+            yield trace
+        finally:
+            self._trace_context.current = previous_trace
+
+    def get_current_trace(self) -> Optional[Trace]:
+        """Get the current trace context."""
+        return getattr(self._trace_context, "current", None)
+
+    @contextmanager
+    def session(
+        self,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        name: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ):
+        """
+        Create a session context 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.thumbs_up()  # Record positive feedback for last response
+
+        Args:
+            session_id: Optional session ID. Auto-generated if not provided.
+            user_id: Optional user identifier.
+            name: Optional name for the session.
+            metadata: Optional metadata dictionary.
+
+        Yields:
+            Session object for recording events and feedback.
+        """
+        sess = Session(
+            session_id=session_id or str(uuid.uuid4()),
+            user_id=user_id,
+            name=name,
+            metadata=metadata or {},
+            _tracer=self,
+        )
+
+        # Store in thread-local context
+        previous_session = getattr(self._session_context, "current", None)
+        self._session_context.current = sess
+
+        # Start session on server
+        if self.enabled:
+            self._start_session(sess)
+
+        try:
+            yield sess
+        finally:
+            # End session on server
+            if self.enabled:
+                self._end_session(sess)
+            self._session_context.current = previous_session
+
+    def get_current_session(self) -> Optional[Session]:
+        """Get the current session context."""
+        return getattr(self._session_context, "current", None)
+
+    def _start_session(self, session: Session) -> None:
+        """Internal: Start a session on the server."""
+        import httpx
+
+        try:
+            with httpx.Client() as client:
+                client.post(
+                    f"{self.base_url}/api/v1/sessions/",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "session_id": session.session_id,
+                        "user_id": session.user_id,
+                        "name": session.name,
+                        "metadata": session.metadata,
+                        "tags": session.tags,
+                        "project": self.project,
+                    },
+                    timeout=10.0,
+                )
+        except Exception:
+            pass  # Silently fail - don't break user code
+
+    def _end_session(self, session: Session) -> None:
+        """Internal: End a session on the server."""
+        import httpx
+
+        try:
+            with httpx.Client() as client:
+                client.patch(
+                    f"{self.base_url}/api/v1/sessions/{session.session_id}/end/",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "metadata": session.metadata,
+                        "tags": session.tags,
+                    },
+                    timeout=10.0,
+                )
+        except Exception:
+            pass
+
+    def _send_session_event(self, session_id: str, event: dict) -> None:
+        """Internal: Send a session event to the server."""
+        import httpx
+
+        try:
+            with httpx.Client() as client:
+                client.post(
+                    f"{self.base_url}/api/v1/sessions/{session_id}/events/",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json=event,
+                    timeout=10.0,
+                )
+        except Exception:
+            pass
+
+    def _send_feedback(self, session_id: str, feedback: dict) -> None:
+        """Internal: Send user feedback to the server."""
+        import httpx
+
+        try:
+            with httpx.Client() as client:
+                client.post(
+                    f"{self.base_url}/api/v1/sessions/{session_id}/feedback/",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json=feedback,
+                    timeout=10.0,
+                )
+        except Exception:
+            pass
+
+    def log(
+        self,
+        *,
+        model: str,
+        provider: str,
+        input_data: Any,
+        output_data: Any,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        latency_ms: int = 0,
+        status: str = "success",
+        error_message: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        tags: Optional[list[str]] = None,
+        trace_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+    ) -> None:
+        """
+        Log an LLM API call.
+
+        This is called automatically when using wrapped clients,
+        but can also be called manually for custom logging.
+        """
+        if not self.enabled:
+            return
+
+        # Apply PII detection and masking if enabled
+        processed_input = input_data
+        processed_output = output_data
+        if self._pii_detector:
+            processed_input, _ = self._pii_detector.process_json(input_data)
+            if output_data:
+                processed_output, _ = self._pii_detector.process_json(output_data)
+
+        # Get trace context
+        current_trace = self.get_current_trace()
+        if current_trace and not trace_id:
+            trace_id = current_trace.trace_id
+            if current_trace.metadata:
+                metadata = {**(metadata or {}), **current_trace.metadata}
+            if current_trace.tags:
+                tags = list(set((tags or []) + current_trace.tags))
+
+        # Get session context
+        session_id = None
+        current_session = self.get_current_session()
+        if current_session:
+            session_id = current_session.session_id
+            if current_session.user_id:
+                metadata = {**(metadata or {}), "user_id": current_session.user_id}
+
+        # Generate span_id
+        generated_span_id = span_id or str(uuid.uuid4())
+
+        log_entry = {
+            "project": self.project,
+            "model": model,
+            "provider": provider,
+            "input": processed_input,
+            "output": processed_output,
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "latency_ms": latency_ms,
+            "status": status,
+            "error_message": error_message,
+            "metadata": metadata,
+            "tags": tags,
+            "trace_id": trace_id,
+            "span_id": generated_span_id,
+            "parent_span_id": parent_span_id,
+            "session_id": session_id,
+        }
+
+        self._queue.add(log_entry)
+
+        # Update session's last log ID for feedback association
+        if current_session:
+            current_session._set_last_log_id(generated_span_id)
+
+    def flush(self) -> None:
+        """Flush all pending logs immediately."""
+        self._queue.flush()
+
+    def shutdown(self) -> None:
+        """Shutdown the client and flush pending logs."""
+        self.flush()
+        self._queue.shutdown()

python/integrations/__init__.py

@@ -0,0 +1,5 @@
+"""AITracer integrations with popular LLM frameworks."""
+
+from .langchain import AITracerCallbackHandler
+
+__all__ = ["AITracerCallbackHandler"]