fenra 0.1.0__py3-none-any.whl

fenra/integrations/anthropic/__init__.py

@@ -0,0 +1,677 @@
+"""Anthropic integration with auto-instrumentation via monkey patching."""
+
+import functools
+import logging
+from typing import Any
+
+from fenra._context import get_context
+from fenra._core import enqueue_transaction
+
+logger = logging.getLogger(__name__)
+
+_patched_sync = False
+_patched_async = False
+_patched_stream_sync = False
+_patched_stream_async = False
+
+# Store original methods for unpatching
+_original_create: Any = None
+_original_create_async: Any = None
+_original_stream: Any = None
+_original_stream_async: Any = None
+
+
+def _detect_web_search(response: Any, tools: list[Any] | None = None) -> dict[str, Any] | None:
+    """
+    Detect web search tool usage from response content blocks.
+
+    Returns a requests usage entry if web search was used, None otherwise.
+    """
+    if not hasattr(response, "content") or response.content is None:
+        return None
+
+    # Check for web_search_tool_result blocks in response content
+    search_count = sum(
+        1
+        for block in response.content
+        if getattr(block, "type", None) == "web_search_tool_result"
+    )
+
+    # Fallback: check for tool_use blocks with name == "web_search"
+    if search_count == 0:
+        search_count = sum(
+            1
+            for block in response.content
+            if getattr(block, "type", None) == "tool_use"
+            and getattr(block, "name", None) == "web_search"
+        )
+
+    # If still no count but web search tool was in the request, default to 1
+    if search_count == 0 and tools:
+        for tool in tools:
+            tool_type = tool.get("type", "") if isinstance(tool, dict) else getattr(tool, "type", "")
+            if "web_search" in str(tool_type):
+                search_count = 1
+                break
+
+    if search_count > 0:
+        return {
+            "type": "requests",
+            "metrics": {"count": search_count, "request_type": "web_search"},
+        }
+
+    return None
+
+
+def _track_message_response(
+    response: Any,
+    model: str,
+    context: dict[str, Any],
+    tools: list[Any] | None = None,
+) -> None:
+    """Extract usage from Anthropic Message response and queue transaction."""
+    if not hasattr(response, "usage") or response.usage is None:
+        return
+
+    usage = response.usage
+
+    input_tokens = getattr(usage, "input_tokens", 0) or 0
+    output_tokens = getattr(usage, "output_tokens", 0) or 0
+
+    metrics: dict[str, Any] = {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": input_tokens + output_tokens,
+    }
+
+    # Include cached tokens if present (prompt caching)
+    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
+    cache_creation = getattr(usage, "cache_creation_input_tokens", 0) or 0
+    if cache_read or cache_creation:
+        metrics["cached_tokens"] = cache_read + cache_creation
+
+    usage_entries: list[dict[str, Any]] = [
+        {
+            "type": "tokens",
+            "metrics": metrics,
+        }
+    ]
+
+    # Detect web search tool usage
+    web_search_usage = _detect_web_search(response, tools)
+    if web_search_usage:
+        usage_entries.append(web_search_usage)
+
+    transaction: dict[str, Any] = {
+        "provider": "anthropic",
+        "model": model,
+        "usage": usage_entries,
+        "context": context,
+    }
+
+    # Include raw usage for debugging if available
+    try:
+        transaction["provider_usage_raw"] = {
+            "input_tokens": getattr(usage, "input_tokens", None),
+            "output_tokens": getattr(usage, "output_tokens", None),
+            "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None),
+            "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None),
+        }
+    except Exception:
+        pass
+
+    enqueue_transaction(transaction)
+
+
+def _track_stream(
+    stream: Any,
+    model: str,
+    context: dict[str, Any],
+    tools: list[Any] | None = None,
+) -> Any:
+    """
+    Wrap a streaming response to track usage from the final chunk.
+
+    For Anthropic's Stream[RawMessageStreamEvent], we need to accumulate
+    usage from both message_start and message_delta events. The message_start
+    contains initial input_tokens, while message_delta at the end contains
+    the final output_tokens.
+    """
+    accumulated_message = None
+    final_usage = None
+
+    def wrapped_stream() -> Any:
+        nonlocal accumulated_message, final_usage
+        try:
+            for event in stream:
+                if hasattr(event, "type"):
+                    if event.type == "message_start" and hasattr(event, "message"):
+                        accumulated_message = event.message
+                    elif event.type == "message_delta" and hasattr(event, "usage"):
+                        # Capture usage from delta - this has final output_tokens
+                        if event.usage:
+                            final_usage = event.usage
+                yield event
+        finally:
+            # After stream is exhausted, track usage
+            if accumulated_message:
+                # Merge final_usage into accumulated_message if available
+                if final_usage and hasattr(accumulated_message, "usage") and accumulated_message.usage:
+                    # Update output_tokens from delta
+                    output_tokens = getattr(final_usage, "output_tokens", None)
+                    if output_tokens is not None:
+                        accumulated_message.usage.output_tokens = output_tokens
+                try:
+                    _track_message_response(accumulated_message, model, context, tools)
+                except Exception as e:
+                    logger.error(
+                        f"Error tracking Anthropic streaming usage: {e}", exc_info=True
+                    )
+
+    return wrapped_stream()
+
+
+async def _track_stream_async(
+    stream: Any,
+    model: str,
+    context: dict[str, Any],
+    tools: list[Any] | None = None,
+) -> Any:
+    """
+    Wrap an async streaming response to track usage from the final chunk.
+
+    For Anthropic's async stream, we need to accumulate usage from both
+    message_start and message_delta events. The message_start contains
+    initial input_tokens, while message_delta at the end contains
+    the final output_tokens.
+    """
+    accumulated_message = None
+    final_usage = None
+
+    try:
+        async for event in stream:
+            if hasattr(event, "type"):
+                if event.type == "message_start" and hasattr(event, "message"):
+                    accumulated_message = event.message
+                elif event.type == "message_delta" and hasattr(event, "usage"):
+                    # Capture usage from delta - this has final output_tokens
+                    if event.usage:
+                        final_usage = event.usage
+            yield event
+    finally:
+        # After stream is exhausted, track usage
+        if accumulated_message:
+            # Merge final_usage into accumulated_message if available
+            if final_usage and hasattr(accumulated_message, "usage") and accumulated_message.usage:
+                # Update output_tokens from delta
+                output_tokens = getattr(final_usage, "output_tokens", None)
+                if output_tokens is not None:
+                    accumulated_message.usage.output_tokens = output_tokens
+            try:
+                _track_message_response(accumulated_message, model, context, tools)
+            except Exception as e:
+                logger.error(
+                    f"Error tracking Anthropic async streaming usage: {e}", exc_info=True
+                )
+
+
+class _TrackedMessageStreamManager:
+    """
+    Wrapper for MessageStreamManager that tracks usage after stream completion.
+    """
+
+    def __init__(
+        self,
+        manager: Any,
+        model: str,
+        context: dict[str, Any],
+        tools: list[Any] | None = None,
+    ):
+        self._manager = manager
+        self._model = model
+        self._context = context
+        self._tools = tools
+
+    def __enter__(self) -> Any:
+        self._stream = self._manager.__enter__()
+        return _TrackedMessageStream(
+            self._stream, self._model, self._context, self._tools
+        )
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
+        return self._manager.__exit__(exc_type, exc_val, exc_tb)
+
+
+class _TrackedMessageStream:
+    """
+    Wrapper for MessageStream that tracks usage when get_final_message is called.
+    """
+
+    def __init__(
+        self,
+        stream: Any,
+        model: str,
+        context: dict[str, Any],
+        tools: list[Any] | None = None,
+    ):
+        self._stream = stream
+        self._model = model
+        self._context = context
+        self._tools = tools
+        self._tracked = False
+
+    def __iter__(self) -> Any:
+        return iter(self._stream)
+
+    def __getattr__(self, name: str) -> Any:
+        attr = getattr(self._stream, name)
+
+        if name == "get_final_message" and not self._tracked:
+
+            @functools.wraps(attr)
+            def tracked_get_final_message(*args: Any, **kwargs: Any) -> Any:
+                result = attr(*args, **kwargs)
+                if not self._tracked:
+                    self._tracked = True
+                    try:
+                        _track_message_response(
+                            result, self._model, self._context, self._tools
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Error tracking Anthropic stream usage: {e}", exc_info=True
+                        )
+                return result
+
+            return tracked_get_final_message
+
+        return attr
+
+
+class _TrackedAsyncMessageStreamManager:
+    """
+    Wrapper for AsyncMessageStreamManager that tracks usage after stream completion.
+    """
+
+    def __init__(
+        self,
+        manager: Any,
+        model: str,
+        context: dict[str, Any],
+        tools: list[Any] | None = None,
+    ):
+        self._manager = manager
+        self._model = model
+        self._context = context
+        self._tools = tools
+
+    async def __aenter__(self) -> Any:
+        self._stream = await self._manager.__aenter__()
+        return _TrackedAsyncMessageStream(
+            self._stream, self._model, self._context, self._tools
+        )
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
+        return await self._manager.__aexit__(exc_type, exc_val, exc_tb)
+
+
+class _TrackedAsyncMessageStream:
+    """
+    Wrapper for AsyncMessageStream that tracks usage when get_final_message is called.
+    """
+
+    def __init__(
+        self,
+        stream: Any,
+        model: str,
+        context: dict[str, Any],
+        tools: list[Any] | None = None,
+    ):
+        self._stream = stream
+        self._model = model
+        self._context = context
+        self._tools = tools
+        self._tracked = False
+
+    def __aiter__(self) -> Any:
+        return self._stream.__aiter__()
+
+    def __getattr__(self, name: str) -> Any:
+        attr = getattr(self._stream, name)
+
+        if name == "get_final_message" and not self._tracked:
+
+            @functools.wraps(attr)
+            async def tracked_get_final_message(*args: Any, **kwargs: Any) -> Any:
+                result = await attr(*args, **kwargs)
+                if not self._tracked:
+                    self._tracked = True
+                    try:
+                        _track_message_response(
+                            result, self._model, self._context, self._tools
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Error tracking Anthropic async stream usage: {e}",
+                            exc_info=True,
+                        )
+                return result
+
+            return tracked_get_final_message
+
+        return attr
+
+
+# ============================================================================
+# Messages.create Patching (sync)
+# ============================================================================
+
+
+def patch_anthropic() -> None:
+    """
+    Patch Anthropic Messages client to auto-track usage.
+
+    This patches the synchronous `Messages.create` method.
+    """
+    global _patched_sync, _original_create
+
+    if _patched_sync:
+        logger.debug("Anthropic Messages sync already patched, skipping")
+        return
+
+    try:
+        from anthropic.resources.messages.messages import Messages
+    except ImportError:
+        logger.warning(
+            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
+        )
+        return
+
+    original_create = Messages.create
+    _original_create = original_create
+
+    @functools.wraps(original_create)
+    def patched_create(self: Messages, *args: Any, **kwargs: Any) -> Any:
+        is_streaming = kwargs.get("stream", False)
+
+        response = original_create(self, *args, **kwargs)
+
+        # Don't track streaming responses here - they need special handling
+        if is_streaming:
+            context = get_context()
+            model = kwargs.get("model") or getattr(response, "model", "unknown")
+            tools = kwargs.get("tools")
+            try:
+                return _track_stream(response, model, context, tools)
+            except Exception as e:
+                logger.error(f"Error wrapping Anthropic stream: {e}", exc_info=True)
+                return response
+
+        context = get_context()
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+        tools = kwargs.get("tools")
+
+        try:
+            _track_message_response(response, model, context, tools)
+        except Exception as e:
+            logger.error(f"Error tracking Anthropic usage: {e}", exc_info=True)
+
+        return response
+
+    Messages.create = patched_create  # type: ignore[assignment]
+    _patched_sync = True
+    logger.info("Anthropic Messages SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# AsyncMessages.create Patching (async)
+# ============================================================================
+
+
+def patch_anthropic_async() -> None:
+    """
+    Patch Anthropic Messages async client to auto-track usage.
+
+    This patches the asynchronous `AsyncMessages.create` method.
+    """
+    global _patched_async, _original_create_async
+
+    if _patched_async:
+        logger.debug("Anthropic Messages async already patched, skipping")
+        return
+
+    try:
+        from anthropic.resources.messages.messages import AsyncMessages
+    except ImportError:
+        logger.warning(
+            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
+        )
+        return
+
+    original_create = AsyncMessages.create
+    _original_create_async = original_create
+
+    @functools.wraps(original_create)
+    async def patched_create(self: AsyncMessages, *args: Any, **kwargs: Any) -> Any:
+        is_streaming = kwargs.get("stream", False)
+
+        response = await original_create(self, *args, **kwargs)
+
+        # Don't track streaming responses here - they need special handling
+        if is_streaming:
+            context = get_context()
+            model = kwargs.get("model") or getattr(response, "model", "unknown")
+            tools = kwargs.get("tools")
+            try:
+                return _track_stream_async(response, model, context, tools)
+            except Exception as e:
+                logger.error(
+                    f"Error wrapping Anthropic async stream: {e}", exc_info=True
+                )
+                return response
+
+        context = get_context()
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+        tools = kwargs.get("tools")
+
+        try:
+            _track_message_response(response, model, context, tools)
+        except Exception as e:
+            logger.error(f"Error tracking Anthropic async usage: {e}", exc_info=True)
+
+        return response
+
+    AsyncMessages.create = patched_create  # type: ignore[assignment]
+    _patched_async = True
+    logger.info("Anthropic Messages async SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Messages.stream Patching (sync)
+# ============================================================================
+
+
+def patch_anthropic_stream() -> None:
+    """
+    Patch Anthropic Messages.stream to auto-track usage.
+
+    This patches the synchronous `Messages.stream` method which returns
+    a MessageStreamManager context manager.
+    """
+    global _patched_stream_sync, _original_stream
+
+    if _patched_stream_sync:
+        logger.debug("Anthropic Messages.stream sync already patched, skipping")
+        return
+
+    try:
+        from anthropic.resources.messages.messages import Messages
+    except ImportError:
+        logger.warning(
+            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
+        )
+        return
+
+    if not hasattr(Messages, "stream"):
+        logger.debug("Anthropic Messages has no stream method, skipping")
+        return
+
+    original_stream = Messages.stream
+    _original_stream = original_stream
+
+    @functools.wraps(original_stream)
+    def patched_stream(self: Messages, *args: Any, **kwargs: Any) -> Any:
+        manager = original_stream(self, *args, **kwargs)
+
+        context = get_context()
+        model = kwargs.get("model", "unknown")
+        tools = kwargs.get("tools")
+
+        try:
+            return _TrackedMessageStreamManager(manager, model, context, tools)
+        except Exception as e:
+            logger.error(f"Error wrapping Anthropic stream manager: {e}", exc_info=True)
+            return manager
+
+    Messages.stream = patched_stream  # type: ignore[assignment]
+    _patched_stream_sync = True
+    logger.info("Anthropic Messages.stream SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# AsyncMessages.stream Patching (async)
+# ============================================================================
+
+
+def patch_anthropic_stream_async() -> None:
+    """
+    Patch Anthropic AsyncMessages.stream to auto-track usage.
+
+    This patches the asynchronous `AsyncMessages.stream` method which returns
+    an AsyncMessageStreamManager context manager.
+    """
+    global _patched_stream_async, _original_stream_async
+
+    if _patched_stream_async:
+        logger.debug("Anthropic AsyncMessages.stream async already patched, skipping")
+        return
+
+    try:
+        from anthropic.resources.messages.messages import AsyncMessages
+    except ImportError:
+        logger.warning(
+            "Anthropic SDK not installed. Install with: pip install 'fenra[anthropic]'"
+        )
+        return
+
+    if not hasattr(AsyncMessages, "stream"):
+        logger.debug("Anthropic AsyncMessages has no stream method, skipping")
+        return
+
+    original_stream = AsyncMessages.stream
+    _original_stream_async = original_stream
+
+    @functools.wraps(original_stream)
+    def patched_stream(self: AsyncMessages, *args: Any, **kwargs: Any) -> Any:
+        manager = original_stream(self, *args, **kwargs)
+
+        context = get_context()
+        model = kwargs.get("model", "unknown")
+        tools = kwargs.get("tools")
+
+        try:
+            return _TrackedAsyncMessageStreamManager(manager, model, context, tools)
+        except Exception as e:
+            logger.error(
+                f"Error wrapping Anthropic async stream manager: {e}", exc_info=True
+            )
+            return manager
+
+    AsyncMessages.stream = patched_stream  # type: ignore[assignment]
+    _patched_stream_async = True
+    logger.info("Anthropic AsyncMessages.stream SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Unpatch Functions
+# ============================================================================
+
+
+def unpatch_anthropic() -> None:
+    """Restore original Anthropic Messages create method."""
+    global _patched_sync, _original_create
+
+    if not _patched_sync or _original_create is None:
+        return
+
+    try:
+        from anthropic.resources.messages.messages import Messages
+
+        Messages.create = _original_create  # type: ignore[assignment]
+        _patched_sync = False
+        _original_create = None
+        logger.info("Anthropic Messages SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_anthropic_async() -> None:
+    """Restore original Anthropic Messages async create method."""
+    global _patched_async, _original_create_async
+
+    if not _patched_async or _original_create_async is None:
+        return
+
+    try:
+        from anthropic.resources.messages.messages import AsyncMessages
+
+        AsyncMessages.create = _original_create_async  # type: ignore[assignment]
+        _patched_async = False
+        _original_create_async = None
+        logger.info("Anthropic Messages async SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_anthropic_stream() -> None:
+    """Restore original Anthropic Messages.stream method."""
+    global _patched_stream_sync, _original_stream
+
+    if not _patched_stream_sync or _original_stream is None:
+        return
+
+    try:
+        from anthropic.resources.messages.messages import Messages
+
+        Messages.stream = _original_stream  # type: ignore[assignment]
+        _patched_stream_sync = False
+        _original_stream = None
+        logger.info("Anthropic Messages.stream SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_anthropic_stream_async() -> None:
+    """Restore original Anthropic AsyncMessages.stream method."""
+    global _patched_stream_async, _original_stream_async
+
+    if not _patched_stream_async or _original_stream_async is None:
+        return
+
+    try:
+        from anthropic.resources.messages.messages import AsyncMessages
+
+        AsyncMessages.stream = _original_stream_async  # type: ignore[assignment]
+        _patched_stream_async = False
+        _original_stream_async = None
+        logger.info("Anthropic AsyncMessages.stream SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_anthropic_all() -> None:
+    """Restore all original Anthropic SDK methods."""
+    unpatch_anthropic()
+    unpatch_anthropic_async()
+    unpatch_anthropic_stream()
+    unpatch_anthropic_stream_async()