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

judgeval/tracer/llm/llm_anthropic/messages_stream.py

@@ -0,0 +1,322 @@
+from __future__ import annotations
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generator,
+    AsyncGenerator,
+)
+
+from judgeval.tracer.keys import AttributeKeys
+from judgeval.tracer.utils import set_span_attribute
+from judgeval.utils.serialize import safe_serialize
+from judgeval.utils.wrappers import (
+    mutable_wrap_sync,
+    immutable_wrap_sync_iterator,
+    immutable_wrap_async_iterator,
+)
+from judgeval.tracer.llm.llm_anthropic.messages import (
+    _extract_anthropic_tokens,
+)
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from anthropic import Anthropic, AsyncAnthropic
+    from anthropic.lib.streaming import (
+        MessageStreamManager,
+        AsyncMessageStreamManager,
+        MessageStream,
+        AsyncMessageStream,
+    )
+
+
+def wrap_messages_stream_sync(tracer: Tracer, client: Anthropic) -> None:
+    original_func = client.messages.stream
+
+    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"}
+        )
+        tracer._inject_judgment_context(ctx["span"])
+        set_span_attribute(
+            ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
+        )
+
+        ctx["model_name"] = kwargs.get("model", "")
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+        ctx["accumulated_content"] = ""
+
+    def mutate_hook(
+        ctx: Dict[str, Any], result: MessageStreamManager
+    ) -> MessageStreamManager:
+        original_manager = result
+
+        class WrappedMessageStreamManager:
+            def __init__(self, manager: MessageStreamManager):
+                self._manager = manager
+
+            def __enter__(self) -> MessageStream:
+                stream = self._manager.__enter__()
+                post_hook_enter_impl(stream)
+                return stream
+
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                result = self._manager.__exit__(exc_type, exc_val, exc_tb)
+                post_hook_exit_impl()
+                return result
+
+            def __getattr__(self, name):
+                return getattr(self._manager, name)
+
+        def post_hook_enter_impl(stream: MessageStream) -> None:
+            ctx["stream"] = stream
+            original_text_stream = stream.text_stream
+
+            def traced_text_stream() -> Generator[str, None, None]:
+                for text_chunk in original_text_stream:
+                    yield text_chunk
+
+            def yield_hook(inner_ctx: Dict[str, Any], text_chunk: str) -> None:
+                span = ctx.get("span")
+                if span and text_chunk:
+                    ctx["accumulated_content"] = (
+                        ctx.get("accumulated_content", "") + text_chunk
+                    )
+
+            def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+                pass
+
+            def error_hook_inner(inner_ctx: Dict[str, Any], error: Exception) -> None:
+                span = ctx.get("span")
+                if span:
+                    span.record_exception(error)
+
+            def finally_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+                pass
+
+            wrapped_text_stream = immutable_wrap_sync_iterator(
+                traced_text_stream,
+                yield_hook=yield_hook,
+                post_hook=post_hook_inner,
+                error_hook=error_hook_inner,
+                finally_hook=finally_hook_inner,
+            )
+
+            stream.text_stream = wrapped_text_stream()
+
+        def post_hook_exit_impl() -> None:
+            span = ctx.get("span")
+            if span:
+                accumulated = ctx.get("accumulated_content", "")
+                set_span_attribute(span, AttributeKeys.GEN_AI_COMPLETION, accumulated)
+
+                stream: MessageStream | None = ctx.get("stream")
+                if stream:
+                    try:
+                        final_message = stream.get_final_message()
+                        if final_message.usage:
+                            (
+                                prompt_tokens,
+                                completion_tokens,
+                                cache_read,
+                                cache_creation,
+                            ) = _extract_anthropic_tokens(final_message.usage)
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                                prompt_tokens,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS,
+                                completion_tokens,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS,
+                                cache_read,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                                cache_creation,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_METADATA,
+                                safe_serialize(final_message.usage),
+                            )
+
+                        set_span_attribute(
+                            span,
+                            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
+                            final_message.model,
+                        )
+                    except Exception:
+                        pass
+
+                span.end()
+
+        return WrappedMessageStreamManager(original_manager)  # type: ignore[return-value]
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+
+    wrapped = mutable_wrap_sync(
+        original_func,
+        pre_hook=pre_hook,
+        mutate_hook=mutate_hook,
+        error_hook=error_hook,
+    )
+
+    setattr(client.messages, "stream", wrapped)
+
+
+def wrap_messages_stream_async(tracer: Tracer, client: AsyncAnthropic) -> None:
+    original_func = client.messages.stream
+
+    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"}
+        )
+        tracer._inject_judgment_context(ctx["span"])
+        set_span_attribute(
+            ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
+        )
+
+        ctx["model_name"] = kwargs.get("model", "")
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+        ctx["accumulated_content"] = ""
+
+    def mutate_hook(
+        ctx: Dict[str, Any], result: AsyncMessageStreamManager
+    ) -> AsyncMessageStreamManager:
+        original_manager = result
+
+        class WrappedAsyncMessageStreamManager:
+            def __init__(self, manager: AsyncMessageStreamManager):
+                self._manager = manager
+
+            async def __aenter__(self) -> AsyncMessageStream:
+                stream = await self._manager.__aenter__()
+                post_hook_aenter_impl(stream)
+                return stream
+
+            async def __aexit__(self, exc_type, exc_val, exc_tb):
+                result = await self._manager.__aexit__(exc_type, exc_val, exc_tb)
+                await post_hook_aexit_impl()
+                return result
+
+            def __getattr__(self, name):
+                return getattr(self._manager, name)
+
+        def post_hook_aenter_impl(stream: AsyncMessageStream) -> None:
+            ctx["stream"] = stream
+            original_text_stream = stream.text_stream
+
+            async def traced_text_stream() -> AsyncGenerator[str, None]:
+                async for text_chunk in original_text_stream:
+                    yield text_chunk
+
+            def yield_hook(inner_ctx: Dict[str, Any], text_chunk: str) -> None:
+                span = ctx.get("span")
+                if span and text_chunk:
+                    ctx["accumulated_content"] = (
+                        ctx.get("accumulated_content", "") + text_chunk
+                    )
+
+            def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+                pass
+
+            def error_hook_inner(inner_ctx: Dict[str, Any], error: Exception) -> None:
+                span = ctx.get("span")
+                if span:
+                    span.record_exception(error)
+
+            def finally_hook_inner_sync(inner_ctx: Dict[str, Any]) -> None:
+                pass
+
+            wrapped_text_stream = immutable_wrap_async_iterator(
+                traced_text_stream,
+                yield_hook=yield_hook,
+                post_hook=post_hook_inner,
+                error_hook=error_hook_inner,
+                finally_hook=finally_hook_inner_sync,
+            )
+
+            stream.text_stream = wrapped_text_stream()
+
+        async def post_hook_aexit_impl() -> None:
+            span = ctx.get("span")
+            if span:
+                accumulated = ctx.get("accumulated_content", "")
+                set_span_attribute(span, AttributeKeys.GEN_AI_COMPLETION, accumulated)
+
+                stream: AsyncMessageStream | None = ctx.get("stream")
+                if stream:
+                    try:
+                        final_message = await stream.get_final_message()
+                        if final_message.usage:
+                            (
+                                prompt_tokens,
+                                completion_tokens,
+                                cache_read,
+                                cache_creation,
+                            ) = _extract_anthropic_tokens(final_message.usage)
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                                prompt_tokens,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS,
+                                completion_tokens,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS,
+                                cache_read,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                                cache_creation,
+                            )
+                            set_span_attribute(
+                                span,
+                                AttributeKeys.JUDGMENT_USAGE_METADATA,
+                                safe_serialize(final_message.usage),
+                            )
+
+                        set_span_attribute(
+                            span,
+                            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
+                            final_message.model,
+                        )
+                    except Exception:
+                        pass
+
+                span.end()
+
+        return WrappedAsyncMessageStreamManager(original_manager)  # type: ignore[return-value]
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+
+    wrapped = mutable_wrap_sync(
+        original_func,
+        pre_hook=pre_hook,
+        mutate_hook=mutate_hook,
+        error_hook=error_hook,
+    )
+
+    setattr(client.messages, "stream", wrapped)

judgeval/tracer/llm/llm_anthropic/wrapper.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Union
+import typing
+
+from judgeval.tracer.llm.llm_anthropic.messages import (
+    wrap_messages_create_sync,
+    wrap_messages_create_async,
+)
+from judgeval.tracer.llm.llm_anthropic.messages_stream import (
+    wrap_messages_stream_sync,
+    wrap_messages_stream_async,
+)
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from anthropic import Anthropic, AsyncAnthropic
+
+    TClient = Union[Anthropic, AsyncAnthropic]
+
+
+def wrap_anthropic_client_sync(tracer: Tracer, client: Anthropic) -> Anthropic:
+    wrap_messages_create_sync(tracer, client)
+    wrap_messages_stream_sync(tracer, client)
+    return client
+
+
+def wrap_anthropic_client_async(
+    tracer: Tracer, client: AsyncAnthropic
+) -> AsyncAnthropic:
+    wrap_messages_create_async(tracer, client)
+    wrap_messages_stream_async(tracer, client)
+    return client
+
+
+@typing.overload
+def wrap_anthropic_client(tracer: Tracer, client: Anthropic) -> Anthropic: ...
+@typing.overload
+def wrap_anthropic_client(tracer: Tracer, client: AsyncAnthropic) -> AsyncAnthropic: ...
+
+
+def wrap_anthropic_client(tracer: Tracer, client: TClient) -> TClient:
+    from judgeval.tracer.llm.llm_anthropic.config import HAS_ANTHROPIC
+    from judgeval.logger import judgeval_logger
+
+    if not HAS_ANTHROPIC:
+        judgeval_logger.error(
+            "Cannot wrap Anthropic client: 'anthropic' library not installed. "
+            "Install it with: pip install anthropic"
+        )
+        return client
+
+    from anthropic import Anthropic, AsyncAnthropic
+
+    if isinstance(client, AsyncAnthropic):
+        return wrap_anthropic_client_async(tracer, client)
+    elif isinstance(client, Anthropic):
+        return wrap_anthropic_client_sync(tracer, client)
+    else:
+        raise TypeError(f"Invalid client type: {type(client)}")

judgeval/tracer/llm/llm_google/__init__.py

@@ -0,0 +1,3 @@
+from .wrapper import wrap_google_client
+
+__all__ = ["wrap_google_client"]

judgeval/tracer/llm/llm_google/config.py

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

judgeval/tracer/llm/llm_google/generate_content.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+    Tuple,
+)
+
+from judgeval.tracer.keys import AttributeKeys
+from judgeval.tracer.utils import set_span_attribute
+from judgeval.utils.serialize import safe_serialize
+from judgeval.utils.wrappers import immutable_wrap_sync
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from google.genai import Client
+    from google.genai.types import (
+        GenerateContentResponse,
+        GenerateContentResponseUsageMetadata,
+    )
+
+
+def _extract_google_tokens(
+    usage: GenerateContentResponseUsageMetadata,
+) -> Tuple[int, int, int, int]:
+    prompt_tokens = (
+        usage.prompt_token_count if usage.prompt_token_count is not None else 0
+    )
+    completion_tokens = (
+        usage.candidates_token_count if usage.candidates_token_count is not None else 0
+    )
+    cache_read_input_tokens = (
+        usage.cached_content_token_count
+        if usage.cached_content_token_count is not None
+        else 0
+    )
+    cache_creation_input_tokens = 0
+    return (
+        prompt_tokens,
+        completion_tokens,
+        cache_read_input_tokens,
+        cache_creation_input_tokens,
+    )
+
+
+def _format_google_output(
+    response: GenerateContentResponse,
+) -> Tuple[Optional[str], Optional[GenerateContentResponseUsageMetadata]]:
+    return response.text, response.usage_metadata
+
+
+def wrap_generate_content_sync(tracer: Tracer, client: Client) -> None:
+    original_func = client.models.generate_content
+
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "GOOGLE_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
+        )
+        tracer._inject_judgment_context(ctx["span"])
+        set_span_attribute(
+            ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
+        )
+        ctx["model_name"] = kwargs.get("model", "")
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, ctx["model_name"]
+        )
+
+    def post_hook(ctx: Dict[str, Any], result: GenerateContentResponse) -> None:
+        span = ctx.get("span")
+        if not span:
+            return
+
+        output, usage_data = _format_google_output(result)
+        set_span_attribute(span, AttributeKeys.GEN_AI_COMPLETION, output)
+
+        if usage_data:
+            prompt_tokens, completion_tokens, cache_read, cache_creation = (
+                _extract_google_tokens(usage_data)
+            )
+            set_span_attribute(
+                span,
+                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                prompt_tokens,
+            )
+            set_span_attribute(
+                span, AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
+            )
+            set_span_attribute(
+                span, AttributeKeys.JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
+            )
+            set_span_attribute(
+                span,
+                AttributeKeys.JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                cache_creation,
+            )
+            set_span_attribute(
+                span,
+                AttributeKeys.JUDGMENT_USAGE_METADATA,
+                safe_serialize(usage_data),
+            )
+
+        set_span_attribute(
+            span,
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
+            result.model_version if result.model_version else ctx["model_name"],
+        )
+
+    def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
+        span = ctx.get("span")
+        if span:
+            span.record_exception(error)
+
+    def finally_hook(ctx: Dict[str, Any]) -> None:
+        span = ctx.get("span")
+        if span:
+            span.end()
+
+    wrapped = immutable_wrap_sync(
+        original_func,
+        pre_hook=pre_hook,
+        post_hook=post_hook,
+        error_hook=error_hook,
+        finally_hook=finally_hook,
+    )
+
+    setattr(client.models, "generate_content", wrapped)

judgeval/tracer/llm/llm_google/wrapper.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from judgeval.tracer.llm.llm_google.generate_content import (
+    wrap_generate_content_sync,
+)
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from google.genai import Client
+
+
+def wrap_google_client(tracer: Tracer, client: Client) -> Client:
+    from judgeval.tracer.llm.llm_google.config import HAS_GOOGLE_GENAI
+    from judgeval.logger import judgeval_logger
+
+    if not HAS_GOOGLE_GENAI:
+        judgeval_logger.error(
+            "Cannot wrap Google GenAI client: 'google-genai' library not installed. "
+            "Install it with: pip install google-genai"
+        )
+        return client
+
+    from google.genai import Client
+
+    if isinstance(client, Client):
+        wrap_generate_content_sync(tracer, client)
+        return client
+    else:
+        raise TypeError(f"Invalid client type: {type(client)}")

judgeval/tracer/llm/llm_openai/__init__.py

@@ -0,0 +1,3 @@
+from .wrapper import wrap_openai_client
+
+__all__ = ["wrap_openai_client"]