judgeval 0.0.11__py3-none-any.whl → 0.22.2__py3-none-any.whl

judgeval/tracer/llm/llm_together/chat_completions.py

@@ -0,0 +1,406 @@
+from __future__ import annotations
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterator,
+    AsyncIterator,
+    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 (
+    immutable_wrap_async,
+    immutable_wrap_sync,
+    mutable_wrap_sync,
+    mutable_wrap_async,
+    immutable_wrap_sync_iterator,
+    immutable_wrap_async_iterator,
+)
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from together import Together, AsyncTogether  # type: ignore[import-untyped]
+    from together.types import ChatCompletionResponse, ChatCompletionChunk  # type: ignore[import-untyped]
+    from together.types.common import UsageData  # type: ignore[import-untyped]
+
+
+def _extract_together_tokens(usage: UsageData) -> tuple[int, int, int, int]:
+    prompt_tokens = usage.prompt_tokens if usage.prompt_tokens is not None else 0
+    completion_tokens = (
+        usage.completion_tokens if usage.completion_tokens is not None else 0
+    )
+    cache_read_input_tokens = 0
+    cache_creation_input_tokens = 0
+    return (
+        prompt_tokens,
+        completion_tokens,
+        cache_read_input_tokens,
+        cache_creation_input_tokens,
+    )
+
+
+def wrap_chat_completions_create_sync(tracer: Tracer, client: Together) -> None:
+    original_func = client.chat.completions.create
+
+    def dispatcher(*args: Any, **kwargs: Any) -> Any:
+        if kwargs.get("stream", False):
+            return _wrap_streaming_sync(tracer, original_func)(*args, **kwargs)  # type: ignore[arg-type]
+        return _wrap_non_streaming_sync(tracer, original_func)(*args, **kwargs)  # type: ignore[arg-type]
+
+    setattr(client.chat.completions, "create", dispatcher)
+
+
+def _wrap_non_streaming_sync(
+    tracer: Tracer, original_func: Callable[..., ChatCompletionResponse]
+) -> Callable[..., ChatCompletionResponse]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "TOGETHER_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", "")
+        prefixed_model_name = (
+            f"together_ai/{ctx['model_name']}" if ctx["model_name"] else ""
+        )
+        ctx["model_name"] = prefixed_model_name
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
+        )
+
+    def post_hook(ctx: Dict[str, Any], result: ChatCompletionResponse) -> None:
+        span = ctx.get("span")
+        if not span:
+            return
+
+        set_span_attribute(
+            span, AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result)
+        )
+
+        if result.usage:
+            prompt_tokens, completion_tokens, _, _ = _extract_together_tokens(
+                result.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_METADATA,
+                safe_serialize(result.usage),
+            )
+
+        set_span_attribute(
+            span,
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
+            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()
+
+    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: Tracer, original_func: Callable[..., Iterator[ChatCompletionChunk]]
+) -> Callable[..., Iterator[ChatCompletionChunk]]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "TOGETHER_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", "")
+        prefixed_model_name = (
+            f"together_ai/{ctx['model_name']}" if ctx["model_name"] else ""
+        )
+        ctx["model_name"] = prefixed_model_name
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
+        )
+        ctx["accumulated_content"] = ""
+
+    def mutate_hook(
+        ctx: Dict[str, Any], result: Iterator[ChatCompletionChunk]
+    ) -> Iterator[ChatCompletionChunk]:
+        def traced_generator() -> Generator[ChatCompletionChunk, None, None]:
+            for chunk in result:
+                yield chunk
+
+        def yield_hook(inner_ctx: Dict[str, Any], chunk: ChatCompletionChunk) -> None:
+            span = ctx.get("span")
+            if not span:
+                return
+
+            if chunk.choices and len(chunk.choices) > 0:
+                delta = chunk.choices[0].delta
+                if delta and hasattr(delta, "content") and delta.content:
+                    ctx["accumulated_content"] = (
+                        ctx.get("accumulated_content", "") + delta.content
+                    )
+
+            if chunk.usage:
+                prompt_tokens, completion_tokens, _, _ = _extract_together_tokens(
+                    chunk.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_METADATA,
+                    safe_serialize(chunk.usage),
+                )
+
+        def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+            span = ctx.get("span")
+            if span:
+                accumulated = ctx.get("accumulated_content", "")
+                set_span_attribute(span, 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)
+
+        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)
+
+    return mutable_wrap_sync(
+        original_func,
+        pre_hook=pre_hook,
+        mutate_hook=mutate_hook,
+        error_hook=error_hook,
+    )
+
+
+def wrap_chat_completions_create_async(tracer: Tracer, client: AsyncTogether) -> None:
+    original_func = client.chat.completions.create
+
+    async def dispatcher(*args: Any, **kwargs: Any) -> Any:
+        if kwargs.get("stream", False):
+            return await _wrap_streaming_async(tracer, original_func)(*args, **kwargs)  # type: ignore[arg-type]
+        return await _wrap_non_streaming_async(tracer, original_func)(*args, **kwargs)  # type: ignore[arg-type]
+
+    setattr(client.chat.completions, "create", dispatcher)
+
+
+def _wrap_non_streaming_async(
+    tracer: Tracer, original_func: Callable[..., Awaitable[ChatCompletionResponse]]
+) -> Callable[..., Awaitable[ChatCompletionResponse]]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "TOGETHER_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", "")
+        prefixed_model_name = (
+            f"together_ai/{ctx['model_name']}" if ctx["model_name"] else ""
+        )
+        ctx["model_name"] = prefixed_model_name
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
+        )
+
+    def post_hook(ctx: Dict[str, Any], result: ChatCompletionResponse) -> None:
+        span = ctx.get("span")
+        if not span:
+            return
+
+        set_span_attribute(
+            span, AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result)
+        )
+
+        if result.usage:
+            prompt_tokens, completion_tokens, _, _ = _extract_together_tokens(
+                result.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_METADATA,
+                safe_serialize(result.usage),
+            )
+
+        set_span_attribute(
+            span,
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
+            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()
+
+    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: Tracer,
+    original_func: Callable[..., Awaitable[AsyncIterator[ChatCompletionChunk]]],
+) -> Callable[..., Awaitable[AsyncIterator[ChatCompletionChunk]]]:
+    def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
+        ctx["span"] = tracer.get_tracer().start_span(
+            "TOGETHER_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", "")
+        prefixed_model_name = (
+            f"together_ai/{ctx['model_name']}" if ctx["model_name"] else ""
+        )
+        ctx["model_name"] = prefixed_model_name
+        set_span_attribute(
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
+        )
+        ctx["accumulated_content"] = ""
+
+    def mutate_hook(
+        ctx: Dict[str, Any], result: AsyncIterator[ChatCompletionChunk]
+    ) -> AsyncIterator[ChatCompletionChunk]:
+        async def traced_generator() -> AsyncGenerator[ChatCompletionChunk, None]:
+            async for chunk in result:
+                yield chunk
+
+        def yield_hook(inner_ctx: Dict[str, Any], chunk: ChatCompletionChunk) -> None:
+            span = ctx.get("span")
+            if not span:
+                return
+
+            if chunk.choices and len(chunk.choices) > 0:
+                delta = chunk.choices[0].delta
+                if delta and hasattr(delta, "content") and delta.content:
+                    ctx["accumulated_content"] = (
+                        ctx.get("accumulated_content", "") + delta.content
+                    )
+
+            if chunk.usage:
+                prompt_tokens, completion_tokens, _, _ = _extract_together_tokens(
+                    chunk.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_METADATA,
+                    safe_serialize(chunk.usage),
+                )
+
+        def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
+            span = ctx.get("span")
+            if span:
+                accumulated = ctx.get("accumulated_content", "")
+                set_span_attribute(span, 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)
+
+        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)
+
+    return mutable_wrap_async(
+        original_func,
+        pre_hook=pre_hook,
+        mutate_hook=mutate_hook,
+        error_hook=error_hook,
+    )

judgeval/tracer/llm/llm_together/config.py

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

judgeval/tracer/llm/llm_together/wrapper.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Union
+import typing
+
+from judgeval.tracer.llm.llm_together.chat_completions import (
+    wrap_chat_completions_create_sync,
+    wrap_chat_completions_create_async,
+)
+
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from together import Together, AsyncTogether  # type: ignore[import-untyped]
+
+    TClient = Union[Together, AsyncTogether]
+
+
+def wrap_together_client_sync(tracer: Tracer, client: Together) -> Together:
+    wrap_chat_completions_create_sync(tracer, client)
+    return client
+
+
+def wrap_together_client_async(tracer: Tracer, client: AsyncTogether) -> AsyncTogether:
+    wrap_chat_completions_create_async(tracer, client)
+    return client
+
+
+@typing.overload
+def wrap_together_client(tracer: Tracer, client: Together) -> Together: ...
+@typing.overload
+def wrap_together_client(tracer: Tracer, client: AsyncTogether) -> AsyncTogether: ...  # type: ignore[overload-cannot-match]
+
+
+def wrap_together_client(tracer: Tracer, client: TClient) -> TClient:
+    from judgeval.tracer.llm.llm_together.config import HAS_TOGETHER
+    from judgeval.logger import judgeval_logger
+
+    if not HAS_TOGETHER:
+        judgeval_logger.error(
+            "Cannot wrap Together client: 'together' library not installed. "
+            "Install it with: pip install together"
+        )
+        return client
+
+    from together import Together, AsyncTogether  # type: ignore[import-untyped]
+
+    if isinstance(client, AsyncTogether):
+        return wrap_together_client_async(tracer, client)
+    elif isinstance(client, Together):
+        return wrap_together_client_sync(tracer, client)
+    else:
+        raise TypeError(f"Invalid client type: {type(client)}")

judgeval/tracer/llm/providers.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+from typing import Any, TypeAlias
+
+from judgeval.tracer.llm.llm_openai.config import HAS_OPENAI
+from judgeval.tracer.llm.llm_together.config import HAS_TOGETHER
+from judgeval.tracer.llm.llm_anthropic.config import HAS_ANTHROPIC
+from judgeval.tracer.llm.llm_google.config import HAS_GOOGLE_GENAI
+
+# TODO: if we support dependency groups we can have this better type, but during runtime, we do
+# not know which clients an end user might have installed.
+ApiClient: TypeAlias = Any
+
+__all__ = [
+    "ApiClient",
+    "HAS_OPENAI",
+    "HAS_TOGETHER",
+    "HAS_ANTHROPIC",
+    "HAS_GOOGLE_GENAI",
+]

judgeval/tracer/managers.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager, contextmanager
+from typing import TYPE_CHECKING, Dict, Optional, List, Any
+from judgeval.tracer.keys import InternalAttributeKeys
+import uuid
+from judgeval.exceptions import JudgmentRuntimeError
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+
+
+@contextmanager
+def sync_span_context(
+    tracer: Tracer,
+    name: str,
+    span_attributes: Optional[Dict[str, str]] = None,
+    disable_partial_emit: bool = False,
+    end_on_exit: bool = False,
+):
+    if span_attributes is None:
+        span_attributes = {}
+
+    with tracer.get_tracer().start_as_current_span(
+        name=name,
+        attributes=span_attributes,
+        end_on_exit=end_on_exit,
+    ) as span:
+        if disable_partial_emit:
+            tracer.judgment_processor.set_internal_attribute(
+                span_context=span.get_span_context(),
+                key=InternalAttributeKeys.DISABLE_PARTIAL_EMIT,
+                value=True,
+            )
+        yield span
+
+
+@asynccontextmanager
+async def async_span_context(
+    tracer: Tracer,
+    name: str,
+    span_attributes: Optional[Dict[str, str]] = None,
+    disable_partial_emit: bool = False,
+    end_on_exit: bool = False,
+):
+    if span_attributes is None:
+        span_attributes = {}
+
+    with tracer.get_tracer().start_as_current_span(
+        name=name,
+        attributes=span_attributes,
+        end_on_exit=end_on_exit,
+    ) as span:
+        if disable_partial_emit:
+            tracer.judgment_processor.set_internal_attribute(
+                span_context=span.get_span_context(),
+                key=InternalAttributeKeys.DISABLE_PARTIAL_EMIT,
+                value=True,
+            )
+        yield span
+
+
+def create_agent_context(
+    tracer: Tracer,
+    args: tuple,
+    class_name: Optional[str] = None,
+    identifier: Optional[str] = None,
+    track_state: bool = False,
+    track_attributes: Optional[List[str]] = None,
+    field_mappings: Optional[Dict[str, str]] = None,
+):
+    """Create agent context and return token for cleanup"""
+    agent_id = str(uuid.uuid4())
+    agent_context: Dict[str, Any] = {"agent_id": agent_id}
+
+    if class_name:
+        agent_context["class_name"] = class_name
+    else:
+        agent_context["class_name"] = None
+
+    agent_context["track_state"] = track_state
+    agent_context["track_attributes"] = track_attributes or []
+    agent_context["field_mappings"] = field_mappings or {}
+
+    instance = args[0] if args else None
+    agent_context["instance"] = instance
+
+    if identifier:
+        if not class_name or not instance or not isinstance(instance, object):
+            raise JudgmentRuntimeError(
+                "'identifier' is set but no class name or instance is available. 'identifier' can only be specified when using the agent() decorator on a class method."
+            )
+        if (
+            instance
+            and hasattr(instance, identifier)
+            and not callable(getattr(instance, identifier))
+        ):
+            instance_name = str(getattr(instance, identifier))
+            agent_context["instance_name"] = instance_name
+        else:
+            raise JudgmentRuntimeError(
+                f"Attribute {identifier} does not exist for {class_name}. Check your agent() decorator."
+            )
+    else:
+        agent_context["instance_name"] = None
+
+    current_agent_context = tracer.get_current_agent_context().get()
+    if current_agent_context and "agent_id" in current_agent_context:
+        agent_context["parent_agent_id"] = current_agent_context["agent_id"]
+    else:
+        agent_context["parent_agent_id"] = None
+
+    agent_context["is_agent_entry_point"] = True
+    token = tracer.get_current_agent_context().set(agent_context)  # type: ignore
+    return token
+
+
+@contextmanager
+def sync_agent_context(
+    tracer: Tracer,
+    args: tuple,
+    class_name: Optional[str] = None,
+    identifier: Optional[str] = None,
+    track_state: bool = False,
+    track_attributes: Optional[List[str]] = None,
+    field_mappings: Optional[Dict[str, str]] = None,
+):
+    """Context manager for synchronous agent context"""
+    token = create_agent_context(
+        tracer=tracer,
+        args=args,
+        class_name=class_name,
+        identifier=identifier,
+        track_state=track_state,
+        track_attributes=track_attributes,
+        field_mappings=field_mappings,
+    )
+    try:
+        yield
+    finally:
+        tracer.get_current_agent_context().reset(token)
+
+
+@asynccontextmanager
+async def async_agent_context(
+    tracer: Tracer,
+    args: tuple,
+    class_name: Optional[str] = None,
+    identifier: Optional[str] = None,
+    track_state: bool = False,
+    track_attributes: Optional[List[str]] = None,
+    field_mappings: Optional[Dict[str, str]] = None,
+):
+    """Context manager for asynchronous agent context"""
+    token = create_agent_context(
+        tracer=tracer,
+        args=args,
+        class_name=class_name,
+        identifier=identifier,
+        track_state=track_state,
+        track_attributes=track_attributes,
+        field_mappings=field_mappings,
+    )
+    try:
+        yield
+    finally:
+        tracer.get_current_agent_context().reset(token)