judgeval 0.1.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

judgeval/v1/instrumentation/llm/config.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from judgeval.logger import judgeval_logger
+
+from judgeval.v1.instrumentation.llm.constants import ProviderType
+from judgeval.v1.instrumentation.llm.providers import (
+    HAS_OPENAI,
+    HAS_TOGETHER,
+    HAS_ANTHROPIC,
+    HAS_GOOGLE_GENAI,
+    ApiClient,
+)
+
+if TYPE_CHECKING:
+    from judgeval.v1.tracer.base_tracer import BaseTracer
+
+
+def _detect_provider(client: ApiClient) -> ProviderType:
+    if HAS_OPENAI:
+        from openai import OpenAI, AsyncOpenAI
+
+        if isinstance(client, (OpenAI, AsyncOpenAI)):
+            return ProviderType.OPENAI
+
+    if HAS_ANTHROPIC:
+        from anthropic import Anthropic, AsyncAnthropic
+
+        if isinstance(client, (Anthropic, AsyncAnthropic)):
+            return ProviderType.ANTHROPIC
+
+    if HAS_TOGETHER:
+        from together import Together, AsyncTogether  # type: ignore[import-untyped]
+
+        if isinstance(client, (Together, AsyncTogether)):
+            return ProviderType.TOGETHER
+
+    if HAS_GOOGLE_GENAI:
+        from google.genai import Client as GoogleClient
+
+        if isinstance(client, GoogleClient):
+            return ProviderType.GOOGLE
+
+    judgeval_logger.warning(
+        f"Unknown client type {type(client)}, Trying to wrap as OpenAI-compatible. "
+        "If this is a mistake or you think we should support this client, please file an issue at https://github.com/JudgmentLabs/judgeval/issues!"
+    )
+
+    return ProviderType.DEFAULT
+
+
+def wrap_provider(tracer: BaseTracer, client: ApiClient) -> ApiClient:
+    """
+    Wraps an API client to add tracing capabilities.
+    Supports OpenAI, Together, Anthropic, and Google GenAI clients.
+    """
+    provider_type = _detect_provider(client)
+
+    if provider_type == ProviderType.OPENAI:
+        from .llm_openai.wrapper import wrap_openai_client
+
+        return wrap_openai_client(tracer, client)
+    elif provider_type == ProviderType.ANTHROPIC:
+        from .llm_anthropic.wrapper import wrap_anthropic_client
+
+        return wrap_anthropic_client(tracer, client)
+    elif provider_type == ProviderType.TOGETHER:
+        from .llm_together.wrapper import wrap_together_client
+
+        return wrap_together_client(tracer, client)
+    elif provider_type == ProviderType.GOOGLE:
+        from .llm_google.wrapper import wrap_google_client
+
+        return wrap_google_client(tracer, client)
+    else:
+        # Default to OpenAI-compatible wrapping for unknown clients
+        from .llm_openai.wrapper import wrap_openai_client
+
+        return wrap_openai_client(tracer, client)

judgeval/v1/instrumentation/llm/constants.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class ProviderType(Enum):
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    TOGETHER = "together"
+    GOOGLE = "google"
+    DEFAULT = "default"

judgeval/v1/instrumentation/llm/llm_anthropic/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .wrapper import wrap_anthropic_client
+
+__all__ = ["wrap_anthropic_client"]

judgeval/v1/instrumentation/llm/llm_anthropic/config.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+import importlib.util
+
+HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None
+
+__all__ = ["HAS_ANTHROPIC"]

judgeval/v1/instrumentation/llm/llm_anthropic/messages.py

@@ -0,0 +1,414 @@
+from __future__ import annotations
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterator,
+    AsyncIterator,
+    Generator,
+    AsyncGenerator,
+    Tuple,
+)
+
+from opentelemetry.trace import Status, StatusCode
+from judgeval.judgment_attribute_keys import AttributeKeys
+from judgeval.utils.serialize import safe_serialize
+from judgeval.utils.wrappers import (
+    immutable_wrap_sync,
+    immutable_wrap_async,
+    mutable_wrap_sync,
+    mutable_wrap_async,
+    immutable_wrap_sync_iterator,
+    immutable_wrap_async_iterator,
+)
+
+if TYPE_CHECKING:
+    from judgeval.v1.tracer import BaseTracer
+    from anthropic import Anthropic, AsyncAnthropic
+    from anthropic.types import (
+        Message,
+        Usage,
+        MessageDeltaUsage,
+        RawMessageStreamEvent,
+    )
+
+
+def _extract_anthropic_content(chunk: RawMessageStreamEvent) -> str:
+    if chunk.type == "content_block_delta":
+        delta = chunk.delta
+        if delta.type == "text_delta" and delta.text:
+            return delta.text
+    return ""
+
+
+def _extract_anthropic_tokens(
+    usage: Usage | MessageDeltaUsage,
+) -> Tuple[int, int, int, int]:
+    input_tokens = usage.input_tokens if usage.input_tokens is not None else 0
+    output_tokens = usage.output_tokens if usage.output_tokens is not None else 0
+    cache_read = (
+        usage.cache_read_input_tokens
+        if usage.cache_read_input_tokens is not None
+        else 0
+    )
+    cache_creation = (
+        usage.cache_creation_input_tokens
+        if usage.cache_creation_input_tokens is not None
+        else 0
+    )
+    return (input_tokens, output_tokens, cache_read, cache_creation)
+
+
+def _extract_anthropic_chunk_usage(
+    chunk: RawMessageStreamEvent,
+) -> Usage | MessageDeltaUsage | None:
+    if chunk.type == "message_start":
+        return chunk.message.usage if chunk.message else None
+    elif chunk.type == "message_delta":
+        return chunk.usage if hasattr(chunk, "usage") else None
+    return None
+
+
+def wrap_messages_create_sync(tracer: BaseTracer, client: Anthropic) -> None:
+    original_func = client.messages.create
+
+    def dispatcher(*args: Any, **kwargs: Any) -> Any:
+        if kwargs.get("stream", False):
+            return _wrap_streaming_sync(tracer, original_func)(*args, **kwargs)
+        return _wrap_non_streaming_sync(tracer, original_func)(*args, **kwargs)
+
+    setattr(client.messages, "create", dispatcher)
+
+
+def _wrap_non_streaming_sync(
+    tracer: BaseTracer, original_func: Callable[..., Message]
+) -> Callable[..., Message]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "ANTHROPIC_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
+        )
+        ctx["span"].set_attribute(AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs))
+        ctx["model_name"] = kwargs.get("model", "")
+        ctx["span"].set_attribute(
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+
+    def post_hook(ctx: Dict[str, Any], result: Message) -> None:
+        span = ctx.get("span")
+        if not span:
+            return
+
+        span.set_attribute(AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result))
+
+        if result.usage:
+            prompt_tokens, completion_tokens, cache_read, cache_creation = (
+                _extract_anthropic_tokens(result.usage)
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                prompt_tokens,
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                cache_creation,
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_METADATA,
+                safe_serialize(result.usage),
+            )
+
+        span.set_attribute(AttributeKeys.JUDGMENT_LLM_MODEL_NAME, result.model)
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR))
+
+    def finally_hook(ctx: Dict[str, Any]) -> None:
+        span = ctx.get("span")
+        if span:
+            span.end()
+
+    return immutable_wrap_sync(
+        original_func,
+        pre_hook=pre_hook,
+        post_hook=post_hook,
+        error_hook=error_hook,
+        finally_hook=finally_hook,
+    )
+
+
+def _wrap_streaming_sync(
+    tracer: BaseTracer, original_func: Callable[..., Iterator[RawMessageStreamEvent]]
+) -> Callable[..., Iterator[RawMessageStreamEvent]]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "ANTHROPIC_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
+        )
+        ctx["span"].set_attribute(AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs))
+        ctx["model_name"] = kwargs.get("model", "")
+        ctx["span"].set_attribute(
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+        ctx["accumulated_content"] = ""
+
+    def mutate_hook(
+        ctx: Dict[str, Any], result: Iterator[RawMessageStreamEvent]
+    ) -> Iterator[RawMessageStreamEvent]:
+        def traced_generator() -> Generator[RawMessageStreamEvent, None, None]:
+            for chunk in result:
+                yield chunk
+
+        def yield_hook(inner_ctx: Dict[str, Any], chunk: RawMessageStreamEvent) -> None:
+            span = ctx.get("span")
+            if not span:
+                return
+
+            content = _extract_anthropic_content(chunk)
+            if content:
+                ctx["accumulated_content"] = (
+                    ctx.get("accumulated_content", "") + content
+                )
+
+            usage_data = _extract_anthropic_chunk_usage(chunk)
+            if usage_data:
+                prompt_tokens, completion_tokens, cache_read, cache_creation = (
+                    _extract_anthropic_tokens(usage_data)
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS, prompt_tokens
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                    cache_creation,
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_METADATA, safe_serialize(usage_data)
+                )
+
+        def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+            span = ctx.get("span")
+            if span:
+                accumulated = ctx.get("accumulated_content", "")
+                span.set_attribute(AttributeKeys.GEN_AI_COMPLETION, accumulated)
+
+        def error_hook_inner(inner_ctx: Dict[str, Any], error: Exception) -> None:
+            span = ctx.get("span")
+            if span:
+                span.record_exception(error)
+                span.set_status(Status(StatusCode.ERROR))
+
+        def finally_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+            span = ctx.get("span")
+            if span:
+                span.end()
+
+        wrapped_generator = immutable_wrap_sync_iterator(
+            traced_generator,
+            yield_hook=yield_hook,
+            post_hook=post_hook_inner,
+            error_hook=error_hook_inner,
+            finally_hook=finally_hook_inner,
+        )
+
+        return wrapped_generator()
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR))
+
+    return mutable_wrap_sync(
+        original_func,
+        pre_hook=pre_hook,
+        mutate_hook=mutate_hook,
+        error_hook=error_hook,
+    )
+
+
+def wrap_messages_create_async(tracer: BaseTracer, client: AsyncAnthropic) -> None:
+    original_func = client.messages.create
+
+    async def dispatcher(*args: Any, **kwargs: Any) -> Any:
+        if kwargs.get("stream", False):
+            return await _wrap_streaming_async(tracer, original_func)(*args, **kwargs)
+        return await _wrap_non_streaming_async(tracer, original_func)(*args, **kwargs)
+
+    setattr(client.messages, "create", dispatcher)
+
+
+def _wrap_non_streaming_async(
+    tracer: BaseTracer, original_func: Callable[..., Awaitable[Message]]
+) -> Callable[..., Awaitable[Message]]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "ANTHROPIC_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
+        )
+        ctx["span"].set_attribute(AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs))
+        ctx["model_name"] = kwargs.get("model", "")
+        ctx["span"].set_attribute(
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+
+    def post_hook(ctx: Dict[str, Any], result: Message) -> None:
+        span = ctx.get("span")
+        if not span:
+            return
+
+        span.set_attribute(AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result))
+
+        if result.usage:
+            prompt_tokens, completion_tokens, cache_read, cache_creation = (
+                _extract_anthropic_tokens(result.usage)
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                prompt_tokens,
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                cache_creation,
+            )
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_USAGE_METADATA,
+                safe_serialize(result.usage),
+            )
+
+        span.set_attribute(AttributeKeys.JUDGMENT_LLM_MODEL_NAME, result.model)
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR))
+
+    def finally_hook(ctx: Dict[str, Any]) -> None:
+        span = ctx.get("span")
+        if span:
+            span.end()
+
+    return immutable_wrap_async(
+        original_func,
+        pre_hook=pre_hook,
+        post_hook=post_hook,
+        error_hook=error_hook,
+        finally_hook=finally_hook,
+    )
+
+
+def _wrap_streaming_async(
+    tracer: BaseTracer,
+    original_func: Callable[..., Awaitable[AsyncIterator[RawMessageStreamEvent]]],
+) -> Callable[..., Awaitable[AsyncIterator[RawMessageStreamEvent]]]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "ANTHROPIC_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
+        )
+        ctx["span"].set_attribute(AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs))
+        ctx["model_name"] = kwargs.get("model", "")
+        ctx["span"].set_attribute(
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+        ctx["accumulated_content"] = ""
+
+    def mutate_hook(
+        ctx: Dict[str, Any], result: AsyncIterator[RawMessageStreamEvent]
+    ) -> AsyncIterator[RawMessageStreamEvent]:
+        async def traced_generator() -> AsyncGenerator[RawMessageStreamEvent, None]:
+            async for chunk in result:
+                yield chunk
+
+        def yield_hook(inner_ctx: Dict[str, Any], chunk: RawMessageStreamEvent) -> None:
+            span = ctx.get("span")
+            if not span:
+                return
+
+            content = _extract_anthropic_content(chunk)
+            if content:
+                ctx["accumulated_content"] = (
+                    ctx.get("accumulated_content", "") + content
+                )
+
+            usage_data = _extract_anthropic_chunk_usage(chunk)
+            if usage_data:
+                prompt_tokens, completion_tokens, cache_read, cache_creation = (
+                    _extract_anthropic_tokens(usage_data)
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS, prompt_tokens
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                    cache_creation,
+                )
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_USAGE_METADATA, safe_serialize(usage_data)
+                )
+
+        def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+            span = ctx.get("span")
+            if span:
+                accumulated = ctx.get("accumulated_content", "")
+                span.set_attribute(AttributeKeys.GEN_AI_COMPLETION, accumulated)
+
+        def error_hook_inner(inner_ctx: Dict[str, Any], error: Exception) -> None:
+            span = ctx.get("span")
+            if span:
+                span.record_exception(error)
+                span.set_status(Status(StatusCode.ERROR))
+
+        def finally_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+            span = ctx.get("span")
+            if span:
+                span.end()
+
+        wrapped_generator = immutable_wrap_async_iterator(
+            traced_generator,
+            yield_hook=yield_hook,
+            post_hook=post_hook_inner,
+            error_hook=error_hook_inner,
+            finally_hook=finally_hook_inner,
+        )
+
+        return wrapped_generator()
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR))
+
+    return mutable_wrap_async(
+        original_func,
+        pre_hook=pre_hook,
+        mutate_hook=mutate_hook,
+        error_hook=error_hook,
+    )