orca-sdk 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

orca_sdk/_utils/analysis_ui.py

@@ -111,7 +111,7 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
 
         @gr.render(
             inputs=[memory_relabel_map, all_approved],
-            triggers=[demo.load, refresh.change, all_approved.change, memory_relabel_map.change],
+            triggers=[demo.load, refresh.change, all_approved.change, memory_relabel_map.change],  # type: ignore[arg-type]
         )
         def render_table(current_memory_relabel_map, current_all_approved):
             if len(current_memory_relabel_map):

orca_sdk/_utils/data_parsing.py

@@ -1,10 +1,10 @@
 import pickle
 from dataclasses import asdict, is_dataclass
 from os import PathLike
-from tempfile import TemporaryDirectory
 from typing import Any, cast
 
 from datasets import Dataset
+from datasets.exceptions import DatasetGenerationError
 from torch.utils.data import DataLoader as TorchDataLoader
 from torch.utils.data import Dataset as TorchDataset
 
@@ -15,7 +15,10 @@ def parse_dict_like(item: Any, column_names: list[str] | None = None) -> dict:
 
     if isinstance(item, tuple):
         if column_names is not None:
-            assert len(item) == len(column_names)
+            if len(item) != len(column_names):
+                raise ValueError(
+                    f"Tuple length ({len(item)}) does not match number of column names ({len(column_names)})"
+                )
             return {column_names[i]: item[i] for i in range(len(item))}
         elif hasattr(item, "_fields") and all(isinstance(field, str) for field in item._fields):  # type: ignore
             return {field: getattr(item, field) for field in item._fields}  # type: ignore
@@ -42,7 +45,8 @@ def parse_batch(batch: Any, column_names: list[str] | None = None) -> list[dict]
 
 
 def hf_dataset_from_torch(
-    torch_data: TorchDataLoader | TorchDataset, column_names: list[str] | None = None, ignore_cache=False
+    torch_data: TorchDataLoader | TorchDataset,
+    column_names: list[str] | None = None,
 ) -> Dataset:
     """
     Create a HuggingFace Dataset from a PyTorch DataLoader or Dataset.
@@ -55,7 +59,6 @@ def hf_dataset_from_torch(
         torch_data: A PyTorch DataLoader or Dataset object to create the HuggingFace Dataset from.
         column_names: Optional list of column names to use for the dataset. If not provided,
             the column names will be inferred from the data.
-        ignore_cache: If True, the dataset will not be cached on disk.
     Returns:
         A HuggingFace Dataset object containing the data from the PyTorch DataLoader or Dataset.
     """
@@ -64,18 +67,19 @@ def hf_dataset_from_torch(
     else:
         dataloader = TorchDataLoader(torch_data, batch_size=1, collate_fn=lambda x: x)
 
-    def generator():
+    # Collect data from the dataloader into a list to avoid serialization issues
+    # with Dataset.from_generator in Python 3.14 (see datasets issue #7839)
+    data_list = []
+    try:
         for batch in dataloader:
-            yield from parse_batch(batch, column_names=column_names)
+            data_list.extend(parse_batch(batch, column_names=column_names))
+    except ValueError as e:
+        raise DatasetGenerationError(str(e)) from e
 
-    if ignore_cache:
-        with TemporaryDirectory() as temp_dir:
-            ds = Dataset.from_generator(generator, cache_dir=temp_dir)
-    else:
-        ds = Dataset.from_generator(generator)
+    ds = Dataset.from_list(data_list)
 
     if not isinstance(ds, Dataset):
-        raise ValueError(f"Failed to create dataset from generator: {type(ds)}")
+        raise ValueError(f"Failed to create dataset from list: {type(ds)}")
     return ds
 
 

orca_sdk/_utils/data_parsing_test.py

@@ -29,7 +29,7 @@ class PytorchDictDataset(TorchDataset):
 def test_hf_dataset_from_torch_dict():
     # Given a Pytorch dataset that returns a dictionary for each item
     dataset = PytorchDictDataset()
-    hf_dataset = hf_dataset_from_torch(dataset, ignore_cache=True)
+    hf_dataset = hf_dataset_from_torch(dataset)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
@@ -51,7 +51,7 @@ def test_hf_dataset_from_torch_tuple():
     # Given a Pytorch dataset that returns a tuple for each item
     dataset = PytorchTupleDataset()
     # And the correct number of column names passed in
-    hf_dataset = hf_dataset_from_torch(dataset, column_names=["value", "label"], ignore_cache=True)
+    hf_dataset = hf_dataset_from_torch(dataset, column_names=["value", "label"])
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
@@ -63,7 +63,7 @@ def test_hf_dataset_from_torch_tuple_error():
     dataset = PytorchTupleDataset()
     # Then the HF dataset should raise an error if no column names are passed in
     with pytest.raises(DatasetGenerationError):
-        hf_dataset_from_torch(dataset, ignore_cache=True)
+        hf_dataset_from_torch(dataset)
 
 
 def test_hf_dataset_from_torch_tuple_error_not_enough_columns():
@@ -71,7 +71,7 @@ def test_hf_dataset_from_torch_tuple_error_not_enough_columns():
     dataset = PytorchTupleDataset()
     # Then the HF dataset should raise an error if not enough column names are passed in
     with pytest.raises(DatasetGenerationError):
-        hf_dataset_from_torch(dataset, column_names=["value"], ignore_cache=True)
+        hf_dataset_from_torch(dataset, column_names=["value"])
 
 
 DatasetTuple = namedtuple("DatasetTuple", ["value", "label"])
@@ -92,7 +92,7 @@ def test_hf_dataset_from_torch_named_tuple():
     # Given a Pytorch dataset that returns a namedtuple for each item
     dataset = PytorchNamedTupleDataset()
     # And no column names are passed in
-    hf_dataset = hf_dataset_from_torch(dataset, ignore_cache=True)
+    hf_dataset = hf_dataset_from_torch(dataset)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
@@ -119,7 +119,7 @@ class PytorchDataclassDataset(TorchDataset):
 def test_hf_dataset_from_torch_dataclass():
     # Given a Pytorch dataset that returns a dataclass for each item
     dataset = PytorchDataclassDataset()
-    hf_dataset = hf_dataset_from_torch(dataset, ignore_cache=True)
+    hf_dataset = hf_dataset_from_torch(dataset)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)
@@ -142,7 +142,7 @@ def test_hf_dataset_from_torch_invalid_dataset():
     dataset = PytorchInvalidDataset()
     # Then the HF dataset should raise an error
     with pytest.raises(DatasetGenerationError):
-        hf_dataset_from_torch(dataset, ignore_cache=True)
+        hf_dataset_from_torch(dataset)
 
 
 def test_hf_dataset_from_torchdataloader():
@@ -153,7 +153,7 @@ def test_hf_dataset_from_torchdataloader():
         return {"value": [item["value"] for item in x], "label": [item["label"] for item in x]}
 
     dataloader = TorchDataLoader(dataset, batch_size=3, collate_fn=collate_fn)
-    hf_dataset = hf_dataset_from_torch(dataloader, ignore_cache=True)
+    hf_dataset = hf_dataset_from_torch(dataloader)
     # Then the HF dataset should be created successfully
     assert isinstance(hf_dataset, Dataset)
     assert len(hf_dataset) == len(dataset)

orca_sdk/async_client.py

@@ -17,6 +17,7 @@ from typing import (
     Mapping,
     NotRequired,
     Self,
+    TypeAlias,
     TypedDict,
     cast,
     overload,
@@ -155,11 +156,14 @@ class ClusterMetrics(TypedDict):
     memory_count: int
 
 
-ColumnType = Literal["STRING", "FLOAT", "INT", "BOOL", "ENUM", "IMAGE", "OTHER"]
+ColumnType: TypeAlias = Literal["STRING", "FLOAT", "INT", "BOOL", "ENUM", "IMAGE", "OTHER"]
+"""
+The type of a column in a datasource
+"""
 
 
 class ConstraintViolationErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[409]
     constraint: str
 
 
@@ -170,6 +174,7 @@ class CountPredictionsRequest(TypedDict):
     start_timestamp: NotRequired[str | None]
     end_timestamp: NotRequired[str | None]
     memory_id: NotRequired[str | None]
+    expected_label_match: NotRequired[bool | None]
 
 
 class CreateApiKeyRequest(TypedDict):
@@ -232,7 +237,7 @@ class EmbeddingEvaluationRequest(TypedDict):
     weigh_memories: NotRequired[bool]
 
 
-EmbeddingFinetuningMethod = Literal["classification", "regression", "batch_triplet_loss"]
+EmbeddingFinetuningMethod: TypeAlias = Literal["classification", "regression", "batch_triplet_loss"]
 
 
 class FeedbackMetrics(TypedDict):
@@ -240,13 +245,55 @@ class FeedbackMetrics(TypedDict):
     count: int
 
 
-FeedbackType = Literal["CONTINUOUS", "BINARY"]
+FeedbackType: TypeAlias = Literal["CONTINUOUS", "BINARY"]
 
 
 class FilterItem(TypedDict):
-    field: list
-    op: Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "like"]
+    field: (
+        tuple[
+            Literal[
+                "memory_id",
+                "value",
+                "label",
+                "metadata",
+                "source_id",
+                "partition_id",
+                "created_at",
+                "updated_at",
+                "edited_at",
+                "metrics",
+                "score",
+                "labels",
+            ]
+        ]
+        | tuple[Literal["metadata"], str]
+        | tuple[
+            Literal["metrics"],
+            Literal[
+                "cluster",
+                "embedding_2d",
+                "is_duplicate",
+                "duplicate_memory_ids",
+                "has_potential_duplicates",
+                "potential_duplicate_memory_ids",
+                "anomaly_score",
+                "neighbor_label_logits",
+                "neighbor_predicted_label",
+                "neighbor_predicted_label_ambiguity",
+                "neighbor_predicted_label_confidence",
+                "current_label_neighbor_confidence",
+                "normalized_neighbor_label_entropy",
+                "neighbor_predicted_label_matches_current_label",
+                "spread",
+                "uniformity",
+                "concept_id",
+                "subconcept_id",
+            ],
+        ]
+    )
+    op: Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "like", "contains all", "contains any"]
     value: str | int | float | bool | list[str | None] | list[int] | list[float] | list[bool] | None
+    transform: NotRequired[Literal["length"]]
 
 
 class GetDatasourceRowCountRequest(TypedDict):
@@ -274,12 +321,17 @@ class HealthyResponse(TypedDict):
 
 
 class InternalServerErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[500]
     message: str
     request_id: str
 
 
-JobStatus = Literal["INITIALIZED", "DISPATCHED", "WAITING", "PROCESSING", "COMPLETED", "FAILED", "ABORTING", "ABORTED"]
+JobStatus: TypeAlias = Literal[
+    "INITIALIZED", "DISPATCHED", "WAITING", "PROCESSING", "COMPLETED", "FAILED", "ABORTING", "ABORTED"
+]
+"""
+Status of job in the job queue
+"""
 
 
 class JobStatusInfo(TypedDict):
@@ -346,7 +398,7 @@ class MemoryMetrics(TypedDict):
     has_potential_duplicates: NotRequired[bool]
     potential_duplicate_memory_ids: NotRequired[list[str] | None]
     cluster: NotRequired[int]
-    embedding_2d: NotRequired[list]
+    embedding_2d: NotRequired[tuple[float, float]]
     anomaly_score: NotRequired[float]
     neighbor_label_logits: NotRequired[list[float] | None]
     neighbor_predicted_label: NotRequired[int | None]
@@ -361,7 +413,7 @@ class MemoryMetrics(TypedDict):
     subconcept_id: NotRequired[int | None]
 
 
-MemoryType = Literal["LABELED", "SCORED"]
+MemoryType: TypeAlias = Literal["LABELED", "SCORED"]
 
 
 class MemorysetClassPatternsAnalysisConfig(TypedDict):
@@ -467,7 +519,7 @@ class MemorysetUpdate(TypedDict):
 
 
 class NotFoundErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[404]
     resource: (
         Literal[
             "org",
@@ -547,7 +599,7 @@ class PredictionFeedbackResult(TypedDict):
     new_category_ids: list[str]
 
 
-PredictionSort = list[list]
+PredictionSort: TypeAlias = list[tuple[Literal["timestamp", "confidence", "anomaly_score"], Literal["asc", "desc"]]]
 
 
 class PredictiveModelUpdate(TypedDict):
@@ -556,15 +608,18 @@ class PredictiveModelUpdate(TypedDict):
     locked: NotRequired[bool]
 
 
-PretrainedEmbeddingModelName = Literal[
+PretrainedEmbeddingModelName: TypeAlias = Literal[
     "CLIP_BASE", "GTE_BASE", "CDE_SMALL", "DISTILBERT", "GTE_SMALL", "MXBAI_LARGE", "E5_LARGE", "BGE_BASE", "GIST_LARGE"
 ]
+"""
+Names of pretrained embedding models that are supported by OrcaCloud
+"""
 
 
-RACHeadType = Literal["KNN", "MMOE", "FF", "BMMOE"]
+RACHeadType: TypeAlias = Literal["KNN", "MMOE", "FF", "BMMOE"]
 
 
-RARHeadType = Literal["MMOE", "KNN"]
+RARHeadType: TypeAlias = Literal["MMOE", "KNN"]
 
 
 class ROCCurve(TypedDict):
@@ -671,6 +726,7 @@ class ScorePredictionWithMemoriesAndFeedback(TypedDict):
     tags: list[str]
     explanation: str | None
     memory_id: str | None
+    is_in_dense_neighborhood: NotRequired[bool | None]
     feedbacks: list[PredictionFeedback]
 
 
@@ -742,7 +798,7 @@ class ScoredMemoryWithFeedbackMetrics(TypedDict):
 
 
 class ServiceUnavailableErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[503]
     service: str
 
 
@@ -754,7 +810,9 @@ class SubConceptMetrics(TypedDict):
     memory_count: int
 
 
-TelemetryField = list
+TelemetryField: TypeAlias = (
+    tuple[Literal["feedback_metrics"], str, Literal["avg", "count"]] | tuple[Literal["lookup"], Literal["count"]]
+)
 
 
 class TelemetryFilterItem(TypedDict):
@@ -769,11 +827,11 @@ class TelemetrySortOptions(TypedDict):
 
 
 class UnauthenticatedErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[401]
 
 
 class UnauthorizedErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[403]
     reason: str
 
 
@@ -794,7 +852,10 @@ class ValidationError(TypedDict):
     type: str
 
 
-WorkerStatus = Literal["IDLE", "BUSY", "DRAINING", "SHUTDOWN", "CRASHED"]
+WorkerStatus: TypeAlias = Literal["IDLE", "BUSY", "DRAINING", "SHUTDOWN", "CRASHED"]
+"""
+Status of worker in the worker pool
+"""
 
 
 class GetTestErrorByStatusCodeParams(TypedDict):
@@ -870,7 +931,7 @@ class PostGpuMemorysetByNameOrIdMemoryParams(TypedDict):
     name_or_id: str
 
 
-PostGpuMemorysetByNameOrIdMemoryRequest = list[LabeledMemoryInsert] | list[ScoredMemoryInsert]
+PostGpuMemorysetByNameOrIdMemoryRequest: TypeAlias = list[LabeledMemoryInsert] | list[ScoredMemoryInsert]
 
 
 class PatchGpuMemorysetByNameOrIdMemoriesParams(TypedDict):
@@ -1106,6 +1167,10 @@ class GetWorkerByWorkerIdParams(TypedDict):
 
 class GetTelemetryPredictionByPredictionIdParams(TypedDict):
     prediction_id: str
+    calc_neighborhood_density: NotRequired[bool]
+    """
+    Calculate neighborhood density
+    """
 
 
 class PatchTelemetryPredictionByPredictionIdParams(TypedDict):
@@ -1144,7 +1209,7 @@ class DeleteTelemetryFeedbackCategoryByNameOrIdParams(TypedDict):
     name_or_id: str
 
 
-PutTelemetryPredictionFeedbackRequest = list[PredictionFeedbackRequest]
+PutTelemetryPredictionFeedbackRequest: TypeAlias = list[PredictionFeedbackRequest]
 
 
 class GetAgentsBootstrapClassificationModelByJobIdParams(TypedDict):
@@ -1197,6 +1262,8 @@ class ClassificationMetrics(TypedDict):
     pr_auc: NotRequired[float | None]
     pr_curve: NotRequired[PRCurve | None]
     roc_curve: NotRequired[ROCCurve | None]
+    confusion_matrix: NotRequired[list[list[int]] | None]
+    warnings: NotRequired[list[str]]
 
 
 class ClassificationModelMetadata(TypedDict):
@@ -1420,7 +1487,7 @@ class HTTPValidationError(TypedDict):
 
 
 class InvalidInputErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[422]
     validation_issues: list[FieldValidationError]
 
 
@@ -1480,6 +1547,7 @@ class LabelPredictionWithMemoriesAndFeedback(TypedDict):
     tags: list[str]
     explanation: str | None
     memory_id: str | None
+    is_in_dense_neighborhood: NotRequired[bool | None]
     feedbacks: list[PredictionFeedback]
 
 
@@ -1551,10 +1619,10 @@ class ListPredictionsRequest(TypedDict):
     start_timestamp: NotRequired[str | None]
     end_timestamp: NotRequired[str | None]
     memory_id: NotRequired[str | None]
+    expected_label_match: NotRequired[bool | None]
     limit: NotRequired[int]
     offset: NotRequired[int | None]
     sort: NotRequired[PredictionSort]
-    expected_label_match: NotRequired[bool | None]
 
 
 class MemorysetAnalysisConfigs(TypedDict):
@@ -1633,10 +1701,10 @@ class WorkerInfo(TypedDict):
     config: dict[str, str | float | int | bool | dict[str, str] | None]
 
 
-PatchGpuMemorysetByNameOrIdMemoryRequest = LabeledMemoryUpdate | ScoredMemoryUpdate
+PatchGpuMemorysetByNameOrIdMemoryRequest: TypeAlias = LabeledMemoryUpdate | ScoredMemoryUpdate
 
 
-PatchGpuMemorysetByNameOrIdMemoriesRequest = list[LabeledMemoryUpdate] | list[ScoredMemoryUpdate]
+PatchGpuMemorysetByNameOrIdMemoriesRequest: TypeAlias = list[LabeledMemoryUpdate] | list[ScoredMemoryUpdate]
 
 
 class CascadingEditSuggestion(TypedDict):
@@ -1864,7 +1932,7 @@ class OrcaAsyncClient(AsyncClient):
         follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
-    ) -> bool:
+    ) -> Literal[True]:
         """Return true only when called with a valid root API key; otherwise 401 Unauthenticated."""
         pass
 
@@ -1898,7 +1966,7 @@ class OrcaAsyncClient(AsyncClient):
         follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
-    ) -> bool:
+    ) -> Literal[True]:
         """Returns true if the api key header is valid for the org (will be false for admin api key)"""
         pass