judgeval 0.16.7__py3-none-any.whl → 0.16.8__py3-none-any.whl

judgeval/tracer/llm/llm_groq/wrapper.py

@@ -1,498 +0,0 @@
-from __future__ import annotations
-import functools
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Optional,
-    Protocol,
-    Tuple,
-    Union,
-    Iterator,
-    AsyncIterator,
-    Sequence,
-    runtime_checkable,
-)
-
-from judgeval.tracer.llm.llm_groq.config import (
-    groq_Groq,
-    groq_AsyncGroq,
-)
-from judgeval.tracer.managers import sync_span_context, async_span_context
-from judgeval.logger import judgeval_logger
-from judgeval.tracer.keys import AttributeKeys
-from judgeval.tracer.utils import set_span_attribute
-from judgeval.utils.serialize import safe_serialize
-
-if TYPE_CHECKING:
-    from judgeval.tracer import Tracer
-    from opentelemetry.trace import Span
-
-# Keep the original client type for runtime compatibility
-GroqClientType = Union[groq_Groq, groq_AsyncGroq]
-
-
-# Usage protocols
-@runtime_checkable
-class GroqPromptTokensDetails(Protocol):
-    cached_tokens: Optional[int]
-
-
-@runtime_checkable
-class GroqUsage(Protocol):
-    prompt_tokens: Optional[int]
-    completion_tokens: Optional[int]
-    total_tokens: Optional[int]
-    prompt_tokens_details: Optional[GroqPromptTokensDetails]
-
-
-# Message protocols
-@runtime_checkable
-class GroqMessage(Protocol):
-    content: Optional[str]
-    role: str
-
-
-@runtime_checkable
-class GroqChoice(Protocol):
-    index: int
-    message: GroqMessage
-    finish_reason: Optional[str]
-
-
-@runtime_checkable
-class GroqChatCompletion(Protocol):
-    id: str
-    object: str
-    created: int
-    model: str
-    choices: Sequence[GroqChoice]
-    usage: Optional[GroqUsage]
-
-
-# Stream protocols
-@runtime_checkable
-class GroqStreamDelta(Protocol):
-    content: Optional[str]
-
-
-@runtime_checkable
-class GroqStreamChoice(Protocol):
-    index: int
-    delta: GroqStreamDelta
-
-
-@runtime_checkable
-class GroqStreamChunk(Protocol):
-    choices: Sequence[GroqStreamChoice]
-    usage: Optional[GroqUsage]
-
-
-# Client protocols
-@runtime_checkable
-class GroqClient(Protocol):
-    pass
-
-
-@runtime_checkable
-class GroqAsyncClient(Protocol):
-    pass
-
-
-# Union types
-GroqResponseType = GroqChatCompletion
-GroqStreamType = Union[Iterator[GroqStreamChunk], AsyncIterator[GroqStreamChunk]]
-
-
-def _extract_groq_content(chunk: GroqStreamChunk) -> str:
-    if chunk.choices and len(chunk.choices) > 0:
-        delta_content = chunk.choices[0].delta.content
-        if delta_content:
-            return delta_content
-    return ""
-
-
-def _extract_groq_tokens(usage_data: GroqUsage) -> Tuple[int, int, int, int]:
-    prompt_tokens = usage_data.prompt_tokens or 0
-    completion_tokens = usage_data.completion_tokens or 0
-    cache_read_input_tokens = 0
-    if (
-        hasattr(usage_data, "prompt_tokens_details")
-        and usage_data.prompt_tokens_details
-        and hasattr(usage_data.prompt_tokens_details, "cached_tokens")
-        and usage_data.prompt_tokens_details.cached_tokens is not None
-    ):
-        cache_read_input_tokens = usage_data.prompt_tokens_details.cached_tokens
-    cache_creation_input_tokens = 0  # Groq doesn't have cache creation tokens
-    return (
-        prompt_tokens,
-        completion_tokens,
-        cache_read_input_tokens,
-        cache_creation_input_tokens,
-    )
-
-
-def _format_groq_output(
-    response: GroqChatCompletion,
-) -> Tuple[Optional[Union[str, list[dict[str, Any]]]], Optional[GroqUsage]]:
-    message_content: Optional[Union[str, list[dict[str, Any]]]] = None
-    usage_data: Optional[GroqUsage] = None
-
-    try:
-        if isinstance(response, GroqChatCompletion):
-            usage_data = response.usage
-            if response.choices and len(response.choices) > 0:
-                content = response.choices[0].message.content
-                if content:
-                    # Return structured data for consistency with other providers
-                    message_content = [{"type": "text", "text": str(content)}]
-    except (AttributeError, IndexError, TypeError):
-        pass
-
-    return message_content, usage_data
-
-
-class TracedGroqGenerator:
-    def __init__(
-        self,
-        tracer: Tracer,
-        generator: Iterator[GroqStreamChunk],
-        client: GroqClientType,
-        span: Span,
-        model_name: str,
-    ):
-        self.tracer = tracer
-        self.generator = generator
-        self.client = client
-        self.span = span
-        self.model_name = model_name
-        self.accumulated_content = ""
-
-    def __iter__(self) -> Iterator[GroqStreamChunk]:
-        return self
-
-    def __next__(self) -> GroqStreamChunk:
-        try:
-            chunk = next(self.generator)
-            content = _extract_groq_content(chunk)
-            if content:
-                self.accumulated_content += content
-            if chunk.usage:
-                prompt_tokens, completion_tokens, cache_read, cache_creation = (
-                    _extract_groq_tokens(chunk.usage)
-                )
-                set_span_attribute(
-                    self.span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
-                )
-                set_span_attribute(
-                    self.span,
-                    AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS,
-                    completion_tokens,
-                )
-                set_span_attribute(
-                    self.span,
-                    AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
-                    cache_read,
-                )
-                set_span_attribute(
-                    self.span,
-                    AttributeKeys.JUDGMENT_USAGE_METADATA,
-                    safe_serialize(chunk.usage),
-                )
-            return chunk
-        except StopIteration:
-            set_span_attribute(
-                self.span, AttributeKeys.GEN_AI_COMPLETION, self.accumulated_content
-            )
-            self.span.end()
-            raise
-        except Exception as e:
-            if self.span:
-                self.span.record_exception(e)
-                self.span.end()
-            raise
-
-
-class TracedGroqAsyncGenerator:
-    def __init__(
-        self,
-        tracer: Tracer,
-        async_generator: AsyncIterator[GroqStreamChunk],
-        client: GroqClientType,
-        span: Span,
-        model_name: str,
-    ):
-        self.tracer = tracer
-        self.async_generator = async_generator
-        self.client = client
-        self.span = span
-        self.model_name = model_name
-        self.accumulated_content = ""
-
-    def __aiter__(self) -> AsyncIterator[GroqStreamChunk]:
-        return self
-
-    async def __anext__(self) -> GroqStreamChunk:
-        try:
-            chunk = await self.async_generator.__anext__()
-            content = _extract_groq_content(chunk)
-            if content:
-                self.accumulated_content += content
-            if chunk.usage:
-                prompt_tokens, completion_tokens, cache_read, cache_creation = (
-                    _extract_groq_tokens(chunk.usage)
-                )
-                set_span_attribute(
-                    self.span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
-                )
-                set_span_attribute(
-                    self.span,
-                    AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS,
-                    completion_tokens,
-                )
-                set_span_attribute(
-                    self.span,
-                    AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
-                    cache_read,
-                )
-                set_span_attribute(
-                    self.span,
-                    AttributeKeys.JUDGMENT_USAGE_METADATA,
-                    safe_serialize(chunk.usage),
-                )
-            return chunk
-        except StopAsyncIteration:
-            set_span_attribute(
-                self.span, AttributeKeys.GEN_AI_COMPLETION, self.accumulated_content
-            )
-            self.span.end()
-            raise
-        except Exception as e:
-            if self.span:
-                self.span.record_exception(e)
-                self.span.end()
-            raise
-
-
-def wrap_groq_client(tracer: Tracer, client: GroqClientType) -> GroqClientType:
-    def wrapped(function: Callable, span_name: str):
-        @functools.wraps(function)
-        def wrapper(*args, **kwargs):
-            if kwargs.get("stream", False):
-                span = tracer.get_tracer().start_span(
-                    span_name, attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
-                )
-                tracer.add_agent_attributes_to_span(span)
-                set_span_attribute(
-                    span, AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
-                )
-                model_name = kwargs.get("model", "")
-                set_span_attribute(span, AttributeKeys.GEN_AI_REQUEST_MODEL, model_name)
-                stream_response = function(*args, **kwargs)
-                return TracedGroqGenerator(
-                    tracer, stream_response, client, span, model_name
-                )
-            else:
-                with sync_span_context(
-                    tracer, span_name, {AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
-                ) as span:
-                    try:
-                        tracer.add_agent_attributes_to_span(span)
-                        set_span_attribute(
-                            span, AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
-                        )
-                        model_name = kwargs.get("model", "")
-                        # Add groq/ prefix for server-side cost calculation
-                        prefixed_model_name = f"groq/{model_name}" if model_name else ""
-                        set_span_attribute(
-                            span,
-                            AttributeKeys.GEN_AI_REQUEST_MODEL,
-                            prefixed_model_name,
-                        )
-                    except Exception as e:
-                        judgeval_logger.error(
-                            f"[groq wrapped] Error adding span metadata: {e}"
-                        )
-
-                    response = function(*args, **kwargs)
-
-                    try:
-                        if isinstance(response, GroqChatCompletion):
-                            output, usage_data = _format_groq_output(response)
-                            # Serialize structured data to JSON for span attribute
-                            if output:
-                                if isinstance(output, list):
-                                    output_str = safe_serialize(output)
-                                else:
-                                    output_str = str(output)
-                                set_span_attribute(
-                                    span, AttributeKeys.GEN_AI_COMPLETION, output_str
-                                )
-                            if usage_data:
-                                (
-                                    prompt_tokens,
-                                    completion_tokens,
-                                    cache_read,
-                                    cache_creation,
-                                ) = _extract_groq_tokens(usage_data)
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS,
-                                    prompt_tokens,
-                                )
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS,
-                                    completion_tokens,
-                                )
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
-                                    cache_read,
-                                )
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.JUDGMENT_USAGE_METADATA,
-                                    safe_serialize(usage_data),
-                                )
-                            # Add groq/ prefix to response model for server-side cost calculation
-                            response_model = getattr(response, "model", model_name)
-                            prefixed_response_model = (
-                                f"groq/{response_model}" if response_model else ""
-                            )
-                            set_span_attribute(
-                                span,
-                                AttributeKeys.GEN_AI_RESPONSE_MODEL,
-                                prefixed_response_model,
-                            )
-                    except Exception as e:
-                        judgeval_logger.error(
-                            f"[groq wrapped] Error adding span metadata: {e}"
-                        )
-                    finally:
-                        return response
-
-        return wrapper
-
-    def wrapped_async(function: Callable, span_name: str):
-        @functools.wraps(function)
-        async def wrapper(*args, **kwargs):
-            if kwargs.get("stream", False):
-                span = tracer.get_tracer().start_span(
-                    span_name, attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
-                )
-                tracer.add_agent_attributes_to_span(span)
-                set_span_attribute(
-                    span, AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
-                )
-                model_name = kwargs.get("model", "")
-                # Add groq/ prefix for server-side cost calculation
-                prefixed_model_name = f"groq/{model_name}" if model_name else ""
-                set_span_attribute(
-                    span, AttributeKeys.GEN_AI_REQUEST_MODEL, prefixed_model_name
-                )
-                stream_response = await function(*args, **kwargs)
-                return TracedGroqAsyncGenerator(
-                    tracer, stream_response, client, span, model_name
-                )
-            else:
-                async with async_span_context(
-                    tracer, span_name, {AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
-                ) as span:
-                    try:
-                        tracer.add_agent_attributes_to_span(span)
-                        set_span_attribute(
-                            span, AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
-                        )
-                        model_name = kwargs.get("model", "")
-                        # Add groq/ prefix for server-side cost calculation
-                        prefixed_model_name = f"groq/{model_name}" if model_name else ""
-                        set_span_attribute(
-                            span,
-                            AttributeKeys.GEN_AI_REQUEST_MODEL,
-                            prefixed_model_name,
-                        )
-                    except Exception as e:
-                        judgeval_logger.error(
-                            f"[groq wrapped_async] Error adding span metadata: {e}"
-                        )
-
-                    response = await function(*args, **kwargs)
-
-                    try:
-                        if isinstance(response, GroqChatCompletion):
-                            output, usage_data = _format_groq_output(response)
-                            # Serialize structured data to JSON for span attribute
-                            if output:
-                                if isinstance(output, list):
-                                    output_str = safe_serialize(output)
-                                else:
-                                    output_str = str(output)
-                                set_span_attribute(
-                                    span, AttributeKeys.GEN_AI_COMPLETION, output_str
-                                )
-                            if usage_data:
-                                (
-                                    prompt_tokens,
-                                    completion_tokens,
-                                    cache_read,
-                                    cache_creation,
-                                ) = _extract_groq_tokens(usage_data)
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS,
-                                    prompt_tokens,
-                                )
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS,
-                                    completion_tokens,
-                                )
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
-                                    cache_read,
-                                )
-                                set_span_attribute(
-                                    span,
-                                    AttributeKeys.JUDGMENT_USAGE_METADATA,
-                                    safe_serialize(usage_data),
-                                )
-                            # Add groq/ prefix to response model for server-side cost calculation
-                            response_model = getattr(response, "model", model_name)
-                            prefixed_response_model = (
-                                f"groq/{response_model}" if response_model else ""
-                            )
-                            set_span_attribute(
-                                span,
-                                AttributeKeys.GEN_AI_RESPONSE_MODEL,
-                                prefixed_response_model,
-                            )
-                    except Exception as e:
-                        judgeval_logger.error(
-                            f"[groq wrapped_async] Error adding span metadata: {e}"
-                        )
-                    finally:
-                        return response
-
-        return wrapper
-
-    span_name = "GROQ_API_CALL"
-    if groq_Groq is not None and isinstance(client, groq_Groq):
-        # Type narrowing for mypy
-        groq_client = client  # type: ignore[assignment]
-        setattr(
-            groq_client.chat.completions,
-            "create",
-            wrapped(groq_client.chat.completions.create, span_name),
-        )
-    elif groq_AsyncGroq is not None and isinstance(client, groq_AsyncGroq):
-        # Type narrowing for mypy
-        async_groq_client = client  # type: ignore[assignment]
-        setattr(
-            async_groq_client.chat.completions,
-            "create",
-            wrapped_async(async_groq_client.chat.completions.create, span_name),
-        )
-
-    return client

judgeval/tracer/local_eval_queue.py

@@ -1,199 +0,0 @@
-"""Local evaluation queue for batching custom scorer evaluations.
-
-This module provides a simple in-memory queue for EvaluationRun objects that contain
-only local (BaseScorer) scorers. Useful for batching evaluations and processing them
-either synchronously or in a background thread.
-"""
-
-import queue
-import threading
-from typing import Callable, List, Optional
-import time
-
-from judgeval.logger import judgeval_logger
-from judgeval.env import JUDGMENT_MAX_CONCURRENT_EVALUATIONS
-from judgeval.data import ScoringResult
-from judgeval.data.evaluation_run import ExampleEvaluationRun
-from judgeval.utils.async_utils import safe_run_async
-from judgeval.scorers.score import a_execute_scoring
-from judgeval.api import JudgmentSyncClient
-from judgeval.env import JUDGMENT_API_KEY, JUDGMENT_ORG_ID
-
-
-class LocalEvaluationQueue:
-    """Lightweight in-memory queue for local evaluation runs.
-
-    Only supports EvaluationRuns with local scorers (BaseScorer instances).
-    API scorers (ExampleAPIScorerConfig) are not supported as they have their own queue.
-    """
-
-    def __init__(
-        self,
-        max_concurrent: int = JUDGMENT_MAX_CONCURRENT_EVALUATIONS,
-        num_workers: int = 4,
-    ):
-        if num_workers <= 0:
-            raise ValueError("num_workers must be a positive integer.")
-        self._queue: queue.Queue[Optional[ExampleEvaluationRun]] = queue.Queue()
-        self._max_concurrent = max_concurrent
-        self._num_workers = num_workers  # Number of worker threads
-        self._worker_threads: List[threading.Thread] = []
-        self._shutdown_event = threading.Event()
-        self._api_client = JudgmentSyncClient(
-            api_key=JUDGMENT_API_KEY,
-            organization_id=JUDGMENT_ORG_ID,
-        )
-
-    def enqueue(self, evaluation_run: ExampleEvaluationRun) -> None:
-        """Add evaluation run to the queue."""
-        self._queue.put(evaluation_run)
-
-    def _process_run(self, evaluation_run: ExampleEvaluationRun) -> List[ScoringResult]:
-        """Execute evaluation run locally and return results."""
-
-        if not evaluation_run.custom_scorers:
-            raise ValueError(
-                "LocalEvaluationQueue only supports runs with local scorers (BaseScorer). "
-                "Found only ExampleAPIScorerConfig instances."
-            )
-
-        return safe_run_async(
-            a_execute_scoring(
-                evaluation_run.examples,
-                evaluation_run.custom_scorers,
-                model=evaluation_run.model,
-                throttle_value=0,
-                max_concurrent=self._max_concurrent // self._num_workers,
-                show_progress=False,
-            )
-        )
-
-    def run_all(
-        self,
-        callback: Optional[
-            Callable[[ExampleEvaluationRun, List[ScoringResult]], None]
-        ] = None,
-    ) -> None:
-        """Process all queued runs synchronously.
-
-        Args:
-            callback: Optional function called after each run with (run, results).
-        """
-        while not self._queue.empty():
-            run = self._queue.get()
-            if run is None:  # Sentinel for worker shutdown
-                self._queue.put(None)
-                break
-            results = self._process_run(run)
-            if callback:
-                callback(run, results)
-            self._queue.task_done()
-
-    def start_workers(
-        self,
-    ) -> List[threading.Thread]:
-        """Start multiple background threads to process runs in parallel.
-        Returns:
-            List of started worker threads.
-        """
-
-        def _worker(worker_id: int) -> None:
-            while not self._shutdown_event.is_set():
-                try:
-                    # Use timeout so workers can check shutdown event periodically
-                    run = self._queue.get(timeout=1.0)
-                    if run is None:  # Sentinel to stop worker
-                        # Put sentinel back for other workers
-                        self._queue.put(None)
-                        self._queue.task_done()
-                        break
-
-                    try:
-                        results = self._process_run(run)
-                        results_dict = [result.model_dump() for result in results]
-                        self._api_client.log_eval_results(
-                            payload={"results": results_dict, "run": run.model_dump()}
-                        )
-                    except Exception as exc:
-                        judgeval_logger.error(
-                            f"Worker {worker_id} error processing {run.eval_name}: {exc}"
-                        )
-                        # Continue processing other runs instead of shutting down all workers
-                    finally:
-                        self._queue.task_done()
-
-                except queue.Empty:
-                    # Timeout - check shutdown event and continue
-                    continue
-
-        # Start worker threads
-        for i in range(self._num_workers):
-            thread = threading.Thread(target=_worker, args=(i,), daemon=True)
-            thread.start()
-            self._worker_threads.append(thread)
-
-        return self._worker_threads
-
-    def start_worker(
-        self,
-        callback: Optional[
-            Callable[[ExampleEvaluationRun, List[ScoringResult]], None]
-        ] = None,
-    ) -> Optional[threading.Thread]:
-        """Start a single background thread to process runs (backward compatibility).
-
-        Args:
-            callback: Optional function called after each run with (run, results).
-
-        Returns:
-            The started thread, or None if no threads were started.
-        """
-        threads = self.start_workers()
-        return threads[0] if threads else None
-
-    def wait_for_completion(self, timeout: Optional[float] = None) -> bool:
-        """Wait for all queued tasks to complete.
-
-        Args:
-            timeout: Maximum time to wait in seconds. None means wait indefinitely.
-
-        Returns:
-            True if all tasks completed, False if timeout occurred.
-        """
-        try:
-            if timeout is None:
-                self._queue.join()
-                return True
-            else:
-                start_time = time.time()
-                while not self._queue.empty() or self._queue.unfinished_tasks > 0:
-                    if time.time() - start_time > timeout:
-                        return False
-                    time.sleep(0.1)
-                return True
-        except Exception:
-            return False
-
-    def stop_workers(self) -> None:
-        """Signal all background workers to stop after current tasks complete."""
-        if not self._worker_threads:
-            return
-
-        # Signal shutdown
-        self._shutdown_event.set()
-
-        # Send sentinel to wake up any blocking workers
-        for _ in range(self._num_workers):
-            self._queue.put(None)
-
-        # Wait for all workers to finish with timeout
-        for thread in self._worker_threads:
-            if thread.is_alive():
-                thread.join(timeout=5.0)
-                if thread.is_alive():
-                    judgeval_logger.warning(
-                        f"Worker thread {thread.name} did not shut down gracefully"
-                    )
-
-        self._worker_threads.clear()
-        self._shutdown_event.clear()