aitracer 0.1.0__py3-none-any.whl

python/wrappers/gemini_wrapper.py

@@ -0,0 +1,350 @@
+"""Google Gemini 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 google.generativeai import GenerativeModel
+    from google.generativeai.types import GenerateContentResponse
+    from aitracer.client import AITracer
+
+
+def wrap_gemini_model(model: "GenerativeModel", tracer: "AITracer") -> "GenerativeModel":
+    """
+    Wrap a Google Gemini GenerativeModel to automatically log all API calls.
+
+    Args:
+        model: GenerativeModel instance.
+        tracer: AITracer instance.
+
+    Returns:
+        Wrapped GenerativeModel (same instance, modified in place).
+
+    Example:
+        >>> import google.generativeai as genai
+        >>> from aitracer import AITracer
+        >>> from aitracer.wrappers import wrap_gemini_model
+        >>>
+        >>> genai.configure(api_key="your-api-key")
+        >>> model = genai.GenerativeModel("gemini-1.5-flash")
+        >>> tracer = AITracer(api_key="your-aitracer-key")
+        >>> model = wrap_gemini_model(model, tracer)
+        >>>
+        >>> response = model.generate_content("Hello!")
+    """
+    model_name = getattr(model, "model_name", "gemini-unknown")
+
+    # Wrap generate_content
+    original_generate = model.generate_content
+
+    @wraps(original_generate)
+    def wrapped_generate_content(
+        contents: Any,
+        *args: Any,
+        stream: bool = False,
+        **kwargs: Any,
+    ) -> Any:
+        """Wrapped generate_content method."""
+        start_time = time.time()
+        span_id = str(uuid.uuid4())
+
+        try:
+            response = original_generate(contents, *args, stream=stream, **kwargs)
+
+            if stream:
+                return _wrap_stream_response(
+                    response=response,
+                    tracer=tracer,
+                    model=model_name,
+                    contents=contents,
+                    start_time=start_time,
+                    span_id=span_id,
+                )
+            else:
+                latency_ms = int((time.time() - start_time) * 1000)
+                _log_generation(
+                    tracer=tracer,
+                    model=model_name,
+                    contents=contents,
+                    response=response,
+                    latency_ms=latency_ms,
+                    span_id=span_id,
+                )
+                return response
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            tracer.log(
+                model=model_name,
+                provider="google",
+                input_data={"contents": _serialize_contents(contents)},
+                output_data=None,
+                latency_ms=latency_ms,
+                status="error",
+                error_message=str(e),
+                span_id=span_id,
+            )
+            raise
+
+    model.generate_content = wrapped_generate_content  # type: ignore
+
+    # Wrap generate_content_async if available
+    if hasattr(model, "generate_content_async"):
+        original_generate_async = model.generate_content_async
+
+        @wraps(original_generate_async)
+        async def wrapped_generate_content_async(
+            contents: Any,
+            *args: Any,
+            stream: bool = False,
+            **kwargs: Any,
+        ) -> Any:
+            """Wrapped async generate_content method."""
+            start_time = time.time()
+            span_id = str(uuid.uuid4())
+
+            try:
+                response = await original_generate_async(
+                    contents, *args, stream=stream, **kwargs
+                )
+
+                if stream:
+                    return _wrap_async_stream_response(
+                        response=response,
+                        tracer=tracer,
+                        model=model_name,
+                        contents=contents,
+                        start_time=start_time,
+                        span_id=span_id,
+                    )
+                else:
+                    latency_ms = int((time.time() - start_time) * 1000)
+                    _log_generation(
+                        tracer=tracer,
+                        model=model_name,
+                        contents=contents,
+                        response=response,
+                        latency_ms=latency_ms,
+                        span_id=span_id,
+                    )
+                    return response
+
+            except Exception as e:
+                latency_ms = int((time.time() - start_time) * 1000)
+                tracer.log(
+                    model=model_name,
+                    provider="google",
+                    input_data={"contents": _serialize_contents(contents)},
+                    output_data=None,
+                    latency_ms=latency_ms,
+                    status="error",
+                    error_message=str(e),
+                    span_id=span_id,
+                )
+                raise
+
+        model.generate_content_async = wrapped_generate_content_async  # type: ignore
+
+    return model
+
+
+def _wrap_stream_response(
+    response: Iterator["GenerateContentResponse"],
+    tracer: "AITracer",
+    model: str,
+    contents: Any,
+    start_time: float,
+    span_id: str,
+) -> Iterator["GenerateContentResponse"]:
+    """Wrap streaming response to log after completion."""
+    content_parts: list[str] = []
+    input_tokens = 0
+    output_tokens = 0
+    last_response = None
+
+    try:
+        for chunk in response:
+            last_response = chunk
+
+            # Accumulate content
+            if chunk.text:
+                content_parts.append(chunk.text)
+
+            # Get usage if available
+            if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
+                input_tokens = getattr(chunk.usage_metadata, "prompt_token_count", 0) or 0
+                output_tokens = getattr(chunk.usage_metadata, "candidates_token_count", 0) or 0
+
+            yield chunk
+
+        # Log after stream completes
+        latency_ms = int((time.time() - start_time) * 1000)
+        full_content = "".join(content_parts)
+
+        tracer.log(
+            model=model,
+            provider="google",
+            input_data={"contents": _serialize_contents(contents)},
+            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="google",
+            input_data={"contents": _serialize_contents(contents)},
+            output_data=None,
+            latency_ms=latency_ms,
+            status="error",
+            error_message=str(e),
+            span_id=span_id,
+        )
+        raise
+
+
+async def _wrap_async_stream_response(
+    response: Any,
+    tracer: "AITracer",
+    model: str,
+    contents: Any,
+    start_time: float,
+    span_id: str,
+) -> Any:
+    """Wrap async streaming response to log after completion."""
+    content_parts: list[str] = []
+    input_tokens = 0
+    output_tokens = 0
+
+    try:
+        async for chunk in response:
+            # Accumulate content
+            if chunk.text:
+                content_parts.append(chunk.text)
+
+            # Get usage if available
+            if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
+                input_tokens = getattr(chunk.usage_metadata, "prompt_token_count", 0) or 0
+                output_tokens = getattr(chunk.usage_metadata, "candidates_token_count", 0) or 0
+
+            yield chunk
+
+        # Log after stream completes
+        latency_ms = int((time.time() - start_time) * 1000)
+        full_content = "".join(content_parts)
+
+        tracer.log(
+            model=model,
+            provider="google",
+            input_data={"contents": _serialize_contents(contents)},
+            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="google",
+            input_data={"contents": _serialize_contents(contents)},
+            output_data=None,
+            latency_ms=latency_ms,
+            status="error",
+            error_message=str(e),
+            span_id=span_id,
+        )
+        raise
+
+
+def _log_generation(
+    tracer: "AITracer",
+    model: str,
+    contents: Any,
+    response: "GenerateContentResponse",
+    latency_ms: int,
+    span_id: str,
+) -> None:
+    """Log a non-streaming generation."""
+    # Extract response data
+    output_content = None
+    try:
+        output_content = response.text
+    except Exception:
+        # Response may not have text (e.g., safety blocked)
+        if response.candidates:
+            parts = []
+            for candidate in response.candidates:
+                if candidate.content and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if hasattr(part, "text"):
+                            parts.append(part.text)
+            output_content = "".join(parts) if parts else None
+
+    # Extract token usage
+    input_tokens = 0
+    output_tokens = 0
+    if hasattr(response, "usage_metadata") and response.usage_metadata:
+        input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0) or 0
+        output_tokens = getattr(response.usage_metadata, "candidates_token_count", 0) or 0
+
+    tracer.log(
+        model=model,
+        provider="google",
+        input_data={"contents": _serialize_contents(contents)},
+        output_data={"content": output_content},
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        latency_ms=latency_ms,
+        status="success",
+        span_id=span_id,
+    )
+
+
+def _serialize_contents(contents: Any) -> Any:
+    """Serialize contents to JSON-compatible format."""
+    if contents is None:
+        return None
+
+    if isinstance(contents, str):
+        return contents
+
+    if isinstance(contents, list):
+        result = []
+        for item in contents:
+            result.append(_serialize_content_item(item))
+        return result
+
+    return _serialize_content_item(contents)
+
+
+def _serialize_content_item(item: Any) -> Any:
+    """Serialize a single content item."""
+    if isinstance(item, str):
+        return item
+
+    if isinstance(item, dict):
+        return item
+
+    if hasattr(item, "to_dict"):
+        return item.to_dict()
+
+    if hasattr(item, "model_dump"):
+        return item.model_dump()
+
+    if hasattr(item, "__dict__"):
+        return {k: v for k, v in item.__dict__.items() if not k.startswith("_")}
+
+    return str(item)

python/wrappers/openai_wrapper.py

@@ -0,0 +1,193 @@
+"""OpenAI client wrapper for automatic logging."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, Iterator
+
+if TYPE_CHECKING:
+    from openai import OpenAI
+    from openai.types.chat import ChatCompletion, ChatCompletionChunk
+    from aitracer.client import AITracer
+
+
+def wrap_openai_client(client: "OpenAI", tracer: "AITracer") -> "OpenAI":
+    """
+    Wrap an OpenAI client to automatically log all API calls.
+
+    Args:
+        client: OpenAI client instance.
+        tracer: AITracer instance.
+
+    Returns:
+        Wrapped OpenAI client (same instance, modified in place).
+    """
+    # Store original method
+    original_create = client.chat.completions.create
+
+    @wraps(original_create)
+    def wrapped_create(*args: Any, **kwargs: Any) -> Any:
+        """Wrapped chat.completions.create method."""
+        start_time = time.time()
+        span_id = str(uuid.uuid4())
+
+        # Extract request data
+        model = kwargs.get("model", args[0] if args else "unknown")
+        messages = kwargs.get("messages", args[1] if len(args) > 1 else [])
+        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,
+                    start_time=start_time,
+                    span_id=span_id,
+                )
+            else:
+                # Handle non-streaming response
+                latency_ms = int((time.time() - start_time) * 1000)
+                _log_completion(
+                    tracer=tracer,
+                    model=model,
+                    messages=messages,
+                    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="openai",
+                input_data={"messages": _serialize_messages(messages)},
+                output_data=None,
+                latency_ms=latency_ms,
+                status="error",
+                error_message=str(e),
+                span_id=span_id,
+            )
+            raise
+
+    # Replace method
+    client.chat.completions.create = wrapped_create  # type: ignore
+
+    return client
+
+
+def _wrap_stream_response(
+    response: Iterator["ChatCompletionChunk"],
+    tracer: "AITracer",
+    model: str,
+    messages: list,
+    start_time: float,
+    span_id: str,
+) -> Iterator["ChatCompletionChunk"]:
+    """Wrap streaming response to log after completion."""
+    chunks: list[Any] = []
+    content_parts: list[str] = []
+    input_tokens = 0
+    output_tokens = 0
+
+    try:
+        for chunk in response:
+            chunks.append(chunk)
+
+            # Accumulate content
+            if chunk.choices and chunk.choices[0].delta.content:
+                content_parts.append(chunk.choices[0].delta.content)
+
+            # Get usage if available (some models include it in final chunk)
+            if hasattr(chunk, "usage") and chunk.usage:
+                input_tokens = chunk.usage.prompt_tokens or 0
+                output_tokens = chunk.usage.completion_tokens or 0
+
+            yield chunk
+
+        # Log after stream completes
+        latency_ms = int((time.time() - start_time) * 1000)
+        full_content = "".join(content_parts)
+
+        tracer.log(
+            model=model,
+            provider="openai",
+            input_data={"messages": _serialize_messages(messages)},
+            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="openai",
+            input_data={"messages": _serialize_messages(messages)},
+            output_data=None,
+            latency_ms=latency_ms,
+            status="error",
+            error_message=str(e),
+            span_id=span_id,
+        )
+        raise
+
+
+def _log_completion(
+    tracer: "AITracer",
+    model: str,
+    messages: list,
+    response: "ChatCompletion",
+    latency_ms: int,
+    span_id: str,
+) -> None:
+    """Log a non-streaming completion."""
+    # Extract response data
+    output_content = None
+    if response.choices and response.choices[0].message:
+        output_content = response.choices[0].message.content
+
+    input_tokens = 0
+    output_tokens = 0
+    if response.usage:
+        input_tokens = response.usage.prompt_tokens
+        output_tokens = response.usage.completion_tokens
+
+    tracer.log(
+        model=model,
+        provider="openai",
+        input_data={"messages": _serialize_messages(messages)},
+        output_data={"content": output_content},
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        latency_ms=latency_ms,
+        status="success",
+        span_id=span_id,
+    )
+
+
+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