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

judgeval/dataset/__init__.py

@@ -0,0 +1,335 @@
+import datetime
+import orjson
+import os
+import yaml
+from dataclasses import dataclass
+from typing import List, Literal, Optional, Iterable, Iterator
+from itertools import islice
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    BarColumn,
+    TaskProgressColumn,
+)
+
+from judgeval.data import Example
+from judgeval.data.trace import Trace
+from judgeval.utils.file_utils import get_examples_from_yaml, get_examples_from_json
+from judgeval.api import JudgmentSyncClient
+from judgeval.logger import judgeval_logger
+from judgeval.env import JUDGMENT_API_KEY, JUDGMENT_ORG_ID
+
+from judgeval.data.judgment_types import DatasetKind
+
+
+def _batch_examples(
+    examples: Iterable[Example], batch_size: int = 100
+) -> Iterator[List[Example]]:
+    """Generator that yields batches of examples for efficient memory usage.
+
+    Works with any iterable including generators, consuming only batch_size items at a time.
+    """
+    iterator = iter(examples)
+    while True:
+        batch = list(islice(iterator, batch_size))
+        if not batch:
+            break
+        yield batch
+
+
+@dataclass
+class DatasetInfo:
+    dataset_id: str
+    name: str
+    created_at: str
+    kind: DatasetKind
+    entries: int
+    creator: str
+
+
+@dataclass
+class Dataset:
+    name: str
+    project_name: str
+    dataset_kind: DatasetKind = DatasetKind.example
+    examples: Optional[List[Example]] = None
+    traces: Optional[List[Trace]] = None
+    judgment_api_key: str | None = JUDGMENT_API_KEY
+    organization_id: str | None = JUDGMENT_ORG_ID
+
+    @classmethod
+    def get(
+        cls,
+        name: str,
+        project_name: str,
+    ):
+        if not cls.judgment_api_key or not cls.organization_id:
+            raise ValueError("Judgment API key and organization ID are required")
+        client = JudgmentSyncClient(cls.judgment_api_key, cls.organization_id)
+        dataset = client.datasets_pull_for_judgeval(
+            {
+                "dataset_name": name,
+                "project_name": project_name,
+            },
+        )
+        if not dataset:
+            raise ValueError(f"Dataset {name} not found in project {project_name}")
+
+        dataset_kind = DatasetKind(dataset.get("dataset_kind", "example"))
+
+        if dataset_kind == DatasetKind.example:
+            examples = dataset.get("examples", [])
+            if examples is None:
+                examples = []
+
+            for e in examples:
+                if isinstance(e, dict) and isinstance(e.get("data", {}), dict):
+                    e.update(e.pop("data"))  # type: ignore
+                    e.pop(
+                        "example_id"
+                    )  # TODO: remove once scorer data migration is complete
+            judgeval_logger.info(f"Successfully retrieved example dataset {name}!")
+            return cls(
+                name=name,
+                project_name=project_name,
+                dataset_kind=dataset_kind,
+                examples=[Example(**e) for e in examples],
+            )
+
+        elif dataset_kind == DatasetKind.trace:
+            trace_data = dataset.get("traces", [])
+            if trace_data is None:
+                trace_data = []
+
+            traces = []
+            for trace_item in trace_data:
+                if isinstance(trace_item, dict):
+                    trace = Trace.from_dataset_trace_with_spans(trace_item)
+                    traces.append(trace)
+
+            judgeval_logger.info(f"Successfully retrieved trace dataset {name}!")
+            return cls(
+                name=name,
+                project_name=project_name,
+                dataset_kind=dataset_kind,
+                traces=traces,
+            )
+
+        else:
+            raise ValueError(f"Unsupported dataset kind: {dataset_kind}")
+
+    @classmethod
+    def create(
+        cls,
+        name: str,
+        project_name: str,
+        examples: Iterable[Example] = [],
+        overwrite: bool = False,
+        batch_size: int = 100,
+    ):
+        """Create a dataset with batched example uploads for large datasets.
+
+        Args:
+            name: Dataset name
+            project_name: Project name
+            examples: Iterable of examples to add (can be a list, generator, etc.)
+            overwrite: Whether to overwrite existing dataset
+            batch_size: Number of examples to upload per batch (default: 100)
+        """
+        if not cls.judgment_api_key or not cls.organization_id:
+            raise ValueError("Judgment API key and organization ID are required")
+
+        client = JudgmentSyncClient(cls.judgment_api_key, cls.organization_id)
+
+        client.datasets_create_for_judgeval(
+            {
+                "name": name,
+                "project_name": project_name,
+                "examples": [],  # type: ignore
+                "dataset_kind": "example",
+                "overwrite": overwrite,
+            }
+        )
+        judgeval_logger.info(f"Created dataset {name}")
+
+        if not isinstance(examples, list):
+            examples = list(examples)
+
+        dataset = cls(
+            name=name,
+            project_name=project_name,
+            examples=examples,
+        )
+        dataset.add_examples(examples, batch_size=batch_size)
+
+        return dataset
+
+    @classmethod
+    def list(cls, project_name: str):
+        if not cls.judgment_api_key or not cls.organization_id:
+            raise ValueError("Judgment API key and organization ID are required")
+        client = JudgmentSyncClient(cls.judgment_api_key, cls.organization_id)
+        datasets = client.datasets_pull_all_for_judgeval({"project_name": project_name})
+
+        judgeval_logger.info(f"Fetched all datasets for project {project_name}!")
+
+        return [DatasetInfo(**dataset_info) for dataset_info in datasets]
+
+    def add_from_json(self, file_path: str) -> None:
+        """
+        Adds examples from a JSON file.
+
+        The JSON file is expected to have the following format:
+        [
+            {
+                "key_01": "value_01",
+                "key_02": "value_02"
+            },
+            {
+                "key_11": "value_11",
+                "key_12": "value_12",
+                "key_13": "value_13"
+            },
+            ...
+        ]
+        """
+        examples = get_examples_from_json(file_path)
+        self.add_examples(examples)
+
+    def add_from_yaml(self, file_path: str) -> None:
+        """
+        Adds examples from a YAML file.
+
+        The YAML file is expected to have the following format:
+        - key_01: value_01
+          key_02: value_02
+        - key_11: value_11
+          key_12: value_12
+          key_13: value_13
+        ...
+        """
+
+        examples = get_examples_from_yaml(file_path)
+        self.add_examples(examples)
+
+    def add_examples(self, examples: Iterable[Example], batch_size: int = 100) -> None:
+        if not self.judgment_api_key or not self.organization_id:
+            raise ValueError("Judgment API key and organization ID are required")
+
+        client = JudgmentSyncClient(self.judgment_api_key, self.organization_id)
+
+        batches = _batch_examples(examples, batch_size)
+        total_uploaded = 0
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]{task.description}"),
+            BarColumn(pulse_style="green"),
+            TaskProgressColumn(),
+            TextColumn("[dim]{task.fields[info]}"),
+        ) as progress:
+            task = progress.add_task(
+                f"Uploading to {self.name}",
+                total=None,
+                info="",
+            )
+
+            batch_num = 0
+            for batch in batches:
+                if len(batch) > 0 and not isinstance(batch[0], Example):
+                    raise TypeError("Examples must be a list of Example objects")
+
+                batch_num += 1
+                batch_size_actual = len(batch)
+                total_uploaded += batch_size_actual
+
+                progress.update(
+                    task,
+                    advance=1,
+                    info=f"Batch {batch_num} ({batch_size_actual} examples, {total_uploaded} total)",
+                )
+
+                client.datasets_insert_examples_for_judgeval(
+                    {
+                        "dataset_name": self.name,
+                        "project_name": self.project_name,
+                        "examples": batch,  # type: ignore
+                    }
+                )
+
+        judgeval_logger.info(
+            f"Successfully added {total_uploaded} examples to dataset {self.name}"
+        )
+
+    def save_as(
+        self,
+        file_type: Literal["json", "yaml"],
+        dir_path: str,
+        save_name: str | None = None,
+    ) -> None:
+        """
+        Saves the dataset as a file. Save only the examples.
+
+        Args:
+            file_type (Literal["json", "csv"]): The file type to save the dataset as.
+            dir_path (str): The directory path to save the file to.
+            save_name (str, optional): The name of the file to save. Defaults to None.
+        """
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+        file_name = (
+            datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            if save_name is None
+            else save_name
+        )
+        complete_path = os.path.join(dir_path, f"{file_name}.{file_type}")
+        if file_type == "json":
+            with open(complete_path, "wb") as file:
+                file.write(
+                    orjson.dumps(
+                        {
+                            "examples": [e.to_dict() for e in self.examples]
+                            if self.examples
+                            else [],
+                        },
+                        option=orjson.OPT_INDENT_2,
+                    )
+                )
+        elif file_type == "yaml":
+            with open(complete_path, "w") as file:
+                yaml_data = {
+                    "examples": [e.to_dict() for e in self.examples]
+                    if self.examples
+                    else [],
+                }
+                yaml.dump(yaml_data, file, default_flow_style=False)
+        else:
+            ACCEPTABLE_FILE_TYPES = ["json", "yaml"]
+            raise TypeError(
+                f"Invalid file type: {file_type}. Please choose from {ACCEPTABLE_FILE_TYPES}"
+            )
+
+    def __iter__(self):
+        if self.dataset_kind == DatasetKind.example and self.examples:
+            return iter(self.examples)
+        elif self.dataset_kind == DatasetKind.trace and self.traces:
+            return iter(self.traces)
+        else:
+            return iter([])
+
+    def __len__(self):
+        if self.dataset_kind == DatasetKind.example and self.examples:
+            return len(self.examples)
+        elif self.dataset_kind == DatasetKind.trace and self.traces:
+            return len(self.traces)
+        else:
+            return 0
+
+    def __str__(self):
+        if self.dataset_kind == DatasetKind.example:
+            return (
+                f"{self.__class__.__name__}(examples={self.examples}, name={self.name})"
+            )
+        else:
+            return f"{self.__class__.__name__}(traces={self.traces}, name={self.name})"

judgeval/env.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+from dotenv import load_dotenv
+
+load_dotenv()
+
+import os
+from typing import overload
+
+
+@overload
+def optional_env_var(var_name: str) -> str | None: ...
+
+
+@overload
+def optional_env_var(var_name: str, default: str) -> str: ...
+
+
+def optional_env_var(var_name: str, default: str | None = None) -> str | None:
+    return os.getenv(var_name, default)
+
+
+JUDGMENT_API_KEY = optional_env_var("JUDGMENT_API_KEY")
+JUDGMENT_ORG_ID = optional_env_var("JUDGMENT_ORG_ID")
+JUDGMENT_API_URL = optional_env_var("JUDGMENT_API_URL", "https://api.judgmentlabs.ai")
+
+JUDGMENT_DEFAULT_GPT_MODEL = optional_env_var(
+    "JUDGMENT_DEFAULT_GPT_MODEL", "gpt-5-mini"
+)
+JUDGMENT_DEFAULT_TOGETHER_MODEL = optional_env_var(
+    "JUDGMENT_DEFAULT_TOGETHER_MODEL", "meta-llama/Meta-Llama-3-8B-Instruct-Lite"
+)
+JUDGMENT_MAX_CONCURRENT_EVALUATIONS = int(
+    optional_env_var("JUDGMENT_MAX_CONCURRENT_EVALUATIONS", "10")
+)
+
+
+JUDGMENT_ENABLE_MONITORING = optional_env_var("JUDGMENT_ENABLE_MONITORING", "true")
+JUDGMENT_ENABLE_EVALUATIONS = optional_env_var("JUDGMENT_ENABLE_EVALUATIONS", "true")
+
+JUDGMENT_S3_ACCESS_KEY_ID = optional_env_var("JUDGMENT_S3_ACCESS_KEY_ID")
+JUDGMENT_S3_SECRET_ACCESS_KEY = optional_env_var("JUDGMENT_S3_SECRET_ACCESS_KEY")
+JUDGMENT_S3_REGION_NAME = optional_env_var("JUDGMENT_S3_REGION_NAME")
+JUDGMENT_S3_BUCKET_NAME = optional_env_var("JUDGMENT_S3_BUCKET_NAME")
+JUDGMENT_S3_PREFIX = optional_env_var("JUDGMENT_S3_PREFIX", "spans/")
+JUDGMENT_S3_ENDPOINT_URL = optional_env_var("JUDGMENT_S3_ENDPOINT_URL")
+JUDGMENT_S3_SIGNATURE_VERSION = optional_env_var("JUDGMENT_S3_SIGNATURE_VERSION", "s3")
+JUDGMENT_S3_ADDRESSING_STYLE = optional_env_var("JUDGMENT_S3_ADDRESSING_STYLE", "auto")
+
+
+JUDGMENT_NO_COLOR = optional_env_var("JUDGMENT_NO_COLOR")
+JUDGMENT_LOG_LEVEL = optional_env_var("JUDGMENT_LOG_LEVEL", "WARNING")
+
+
+TOGETHERAI_API_KEY = optional_env_var("TOGETHERAI_API_KEY")
+TOGETHER_API_KEY = optional_env_var("TOGETHER_API_KEY")