orca-sdk 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

orca_sdk/client.py

@@ -135,6 +135,8 @@ class ClassificationEvaluationRequest(TypedDict):
     telemetry_tags: NotRequired[list[str] | None]
     subsample: NotRequired[int | float | None]
     ignore_unlabeled: NotRequired[bool]
+    datasource_partition_column: NotRequired[str | None]
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class CleanupResponse(TypedDict):
@@ -315,12 +317,16 @@ class ListMemoriesRequest(TypedDict):
     offset: NotRequired[int]
     limit: NotRequired[int]
     filters: NotRequired[list[FilterItem]]
+    partition_id: NotRequired[str | None]
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class LookupRequest(TypedDict):
     query: list[str]
     count: NotRequired[int]
     prompt: NotRequired[str | None]
+    partition_id: NotRequired[str | list[str | None] | None]
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class LookupScoreMetrics(TypedDict):
@@ -586,6 +592,8 @@ class RegressionEvaluationRequest(TypedDict):
     telemetry_tags: NotRequired[list[str] | None]
     subsample: NotRequired[int | float | None]
     ignore_unlabeled: NotRequired[bool]
+    datasource_partition_column: NotRequired[str | None]
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class RegressionMetrics(TypedDict):
@@ -629,6 +637,8 @@ class RegressionPredictionRequest(TypedDict):
     use_lookup_cache: NotRequired[bool]
     consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"] | None]
     ignore_unlabeled: NotRequired[bool]
+    partition_ids: NotRequired[str | list[str | None] | None]
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class ScorePredictionMemoryLookup(TypedDict):
@@ -1216,6 +1226,8 @@ class ClassificationPredictionRequest(TypedDict):
     use_lookup_cache: NotRequired[bool]
     consistency_level: NotRequired[Literal["Bounded", "Session", "Strong", "Eventual"] | None]
     ignore_unlabeled: NotRequired[bool]
+    partition_ids: NotRequired[str | list[str | None] | None]
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class CloneMemorysetRequest(TypedDict):
@@ -1269,6 +1281,7 @@ class CreateMemorysetRequest(TypedDict):
     datasource_score_column: NotRequired[str | None]
     datasource_value_column: str
     datasource_source_id_column: NotRequired[str | None]
+    datasource_partition_id_column: NotRequired[str | None]
     remove_duplicates: NotRequired[bool]
     pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
     finetuned_embedding_model_name_or_id: NotRequired[str | None]
@@ -1539,6 +1552,7 @@ class MemorysetAnalysisRequest(TypedDict):
     batch_size: NotRequired[int]
     clear_metrics: NotRequired[bool]
     configs: MemorysetAnalysisConfigs
+    partition_filter_mode: NotRequired[Literal["ignore_partitions", "include_global", "exclude_global", "only_global"]]
 
 
 class MemorysetConceptMetrics(TypedDict):

orca_sdk/conftest.py

@@ -99,34 +99,105 @@ def label_names():
 
 
 SAMPLE_DATA = [
-    {"value": "i love soup", "label": 0, "key": "g1", "score": 0.1, "source_id": "s1"},
-    {"value": "cats are cute", "label": 1, "key": "g1", "score": 0.9, "source_id": "s2"},
-    {"value": "soup is good", "label": 0, "key": "g1", "score": 0.1, "source_id": "s3"},
-    {"value": "i love cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s4"},
-    {"value": "everyone loves cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s5"},
-    {"value": "soup is great for the winter", "label": 0, "key": "g1", "score": 0.1, "source_id": "s6"},
-    {"value": "hot soup on a rainy day!", "label": 0, "key": "g1", "score": 0.1, "source_id": "s7"},
-    {"value": "cats sleep all day", "label": 1, "key": "g1", "score": 0.9, "source_id": "s8"},
-    {"value": "homemade soup recipes", "label": 0, "key": "g1", "score": 0.1, "source_id": "s9"},
-    {"value": "cats purr when happy", "label": 1, "key": "g2", "score": 0.9, "source_id": "s10"},
-    {"value": "chicken noodle soup is classic", "label": 0, "key": "g1", "score": 0.1, "source_id": "s11"},
-    {"value": "kittens are baby cats", "label": 1, "key": "g2", "score": 0.9, "source_id": "s12"},
-    {"value": "soup can be served cold too", "label": 0, "key": "g1", "score": 0.1, "source_id": "s13"},
-    {"value": "cats have nine lives", "label": 1, "key": "g2", "score": 0.9, "source_id": "s14"},
-    {"value": "tomato soup with grilled cheese", "label": 0, "key": "g1", "score": 0.1, "source_id": "s15"},
-    {"value": "cats are independent animals", "label": 1, "key": "g2", "score": 0.9, "source_id": "s16"},
-    {"value": "the beach is always fun", "label": None, "key": "g3", "score": None, "source_id": "s17"},
-    {"value": "i love the beach", "label": None, "key": "g3", "score": None, "source_id": "s18"},
-    {"value": "the ocean is healing", "label": None, "key": "g3", "score": None, "source_id": "s19"},
+    {"value": "i love soup", "label": 0, "key": "g1", "score": 0.1, "source_id": "s1", "partition_id": "p1"},
+    {"value": "cats are cute", "label": 1, "key": "g1", "score": 0.9, "source_id": "s2", "partition_id": "p1"},
+    {"value": "soup is good", "label": 0, "key": "g1", "score": 0.1, "source_id": "s3", "partition_id": "p1"},
+    {"value": "i love cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s4", "partition_id": "p1"},
+    {"value": "everyone loves cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s5", "partition_id": "p1"},
+    {
+        "value": "soup is great for the winter",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s6",
+        "partition_id": "p1",
+    },
+    {
+        "value": "hot soup on a rainy day!",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s7",
+        "partition_id": "p1",
+    },
+    {"value": "cats sleep all day", "label": 1, "key": "g1", "score": 0.9, "source_id": "s8", "partition_id": "p1"},
+    {"value": "homemade soup recipes", "label": 0, "key": "g1", "score": 0.1, "source_id": "s9", "partition_id": "p2"},
+    {"value": "cats purr when happy", "label": 1, "key": "g2", "score": 0.9, "source_id": "s10", "partition_id": "p2"},
+    {
+        "value": "chicken noodle soup is classic",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s11",
+        "partition_id": "p2",
+    },
+    {"value": "kittens are baby cats", "label": 1, "key": "g2", "score": 0.9, "source_id": "s12", "partition_id": "p2"},
+    {
+        "value": "soup can be served cold too",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s13",
+        "partition_id": "p2",
+    },
+    {"value": "cats have nine lives", "label": 1, "key": "g2", "score": 0.9, "source_id": "s14", "partition_id": "p2"},
+    {
+        "value": "tomato soup with grilled cheese",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s15",
+        "partition_id": "p2",
+    },
+    {
+        "value": "cats are independent animals",
+        "label": 1,
+        "key": "g2",
+        "score": 0.9,
+        "source_id": "s16",
+        "partition_id": None,
+    },
+    {
+        "value": "the beach is always fun",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s17",
+        "partition_id": None,
+    },
+    {"value": "i love the beach", "label": None, "key": "g3", "score": None, "source_id": "s18", "partition_id": None},
+    {
+        "value": "the ocean is healing",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s19",
+        "partition_id": None,
+    },
     {
         "value": "sandy feet, sand between my toes at the beach",
         "label": None,
         "key": "g3",
         "score": None,
         "source_id": "s20",
+        "partition_id": None,
+    },
+    {
+        "value": "i am such a beach bum",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s21",
+        "partition_id": None,
+    },
+    {
+        "value": "i will always want to be at the beach",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s22",
+        "partition_id": None,
     },
-    {"value": "i am such a beach bum", "label": None, "key": "g3", "score": None, "source_id": "s21"},
-    {"value": "i will always want to be at the beach", "label": None, "key": "g3", "score": None, "source_id": "s22"},
 ]
 
 
@@ -141,6 +212,7 @@ def hf_dataset(label_names: list[str]) -> Dataset:
                 "key": Value("string"),
                 "score": Value("float"),
                 "source_id": Value("string"),
+                "partition_id": Value("string"),
             }
         ),
     )
@@ -186,6 +258,18 @@ def readonly_memoryset(datasource: Datasource) -> LabeledMemoryset:
     return memoryset
 
 
+@pytest.fixture(scope="session")
+def readonly_partitioned_memoryset(datasource: Datasource) -> LabeledMemoryset:
+    memoryset = LabeledMemoryset.create(
+        "test_readonly_partitioned_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        partition_id_column="partition_id",
+    )
+    return memoryset
+
+
 @pytest.fixture(scope="function")
 def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[LabeledMemoryset, None, None]:
     """
@@ -237,6 +321,18 @@ def classification_model(readonly_memoryset: LabeledMemoryset) -> Classification
     return model
 
 
+@pytest.fixture(scope="session")
+def partitioned_classification_model(readonly_partitioned_memoryset: LabeledMemoryset) -> ClassificationModel:
+    model = ClassificationModel.create(
+        "test_partitioned_classification_model",
+        readonly_partitioned_memoryset,
+        num_classes=2,
+        memory_lookup_count=3,
+        description="test_partitioned_description",
+    )
+    return model
+
+
 # Add scored memoryset and regression model fixtures
 @pytest.fixture(scope="session")
 def scored_memoryset(datasource: Datasource) -> ScoredMemoryset:
@@ -261,3 +357,26 @@ def regression_model(scored_memoryset: ScoredMemoryset) -> RegressionModel:
         description="test_regression_description",
     )
     return model
+
+
+@pytest.fixture(scope="session")
+def readonly_partitioned_scored_memoryset(datasource: Datasource) -> ScoredMemoryset:
+    memoryset = ScoredMemoryset.create(
+        "test_readonly_partitioned_scored_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        partition_id_column="partition_id",
+    )
+    return memoryset
+
+
+@pytest.fixture(scope="session")
+def partitioned_regression_model(readonly_partitioned_scored_memoryset: ScoredMemoryset) -> RegressionModel:
+    model = RegressionModel.create(
+        "test_partitioned_regression_model",
+        readonly_partitioned_scored_memoryset,
+        memory_lookup_count=3,
+        description="test_partitioned_regression_description",
+    )
+    return model