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

orca_sdk/regression_model_test.py

@@ -193,6 +193,140 @@ def test_evaluate_with_telemetry(regression_model, eval_dataset: Dataset):
     assert all(np.allclose(p.expected_score, s) for p, s in zip(predictions, eval_dataset["score"]))
 
 
+def test_evaluate_with_partition_column_dataset(partitioned_regression_model: RegressionModel):
+    """Test evaluate with partition_column on a Dataset"""
+    # Create a test dataset with partition_id column
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+            {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
+            {"value": "homemade soup recipes", "score": 0.1, "partition_id": "p2"},
+            {"value": "cats purr when happy", "score": 0.9, "partition_id": "p2"},
+        ]
+    )
+
+    # Evaluate with partition_column
+    result = partitioned_regression_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="exclude_global",
+    )
+    assert result is not None
+    assert isinstance(result, RegressionMetrics)
+    assert isinstance(result.mae, float)
+    assert isinstance(result.mse, float)
+    assert isinstance(result.rmse, float)
+    assert result.r2 is not None
+
+
+def test_evaluate_with_partition_column_include_global(partitioned_regression_model: RegressionModel):
+    """Test evaluate with partition_column and include_global mode"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+            {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
+        ]
+    )
+
+    # Evaluate with partition_column and include_global (default)
+    result = partitioned_regression_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="include_global",
+    )
+    assert result is not None
+    assert isinstance(result, RegressionMetrics)
+
+
+def test_evaluate_with_partition_column_exclude_global(partitioned_regression_model: RegressionModel):
+    """Test evaluate with partition_column and exclude_global mode"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+            {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
+        ]
+    )
+
+    # Evaluate with partition_column and exclude_global
+    result = partitioned_regression_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="exclude_global",
+    )
+    assert result is not None
+    assert isinstance(result, RegressionMetrics)
+
+
+def test_evaluate_with_partition_column_only_global(partitioned_regression_model: RegressionModel):
+    """Test evaluate with partition_filter_mode only_global"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "cats are independent animals", "score": 0.9, "partition_id": None},
+            {"value": "i love the beach", "score": 0.5, "partition_id": None},
+        ]
+    )
+
+    # Evaluate with only_global mode
+    result = partitioned_regression_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="only_global",
+    )
+    assert result is not None
+    assert isinstance(result, RegressionMetrics)
+
+
+def test_evaluate_with_partition_column_ignore_partitions(partitioned_regression_model: RegressionModel):
+    """Test evaluate with partition_filter_mode ignore_partitions"""
+    eval_dataset_with_partition = Dataset.from_list(
+        [
+            {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+            {"value": "cats are cute", "score": 0.9, "partition_id": "p2"},
+        ]
+    )
+
+    # Evaluate with ignore_partitions mode
+    result = partitioned_regression_model.evaluate(
+        eval_dataset_with_partition,
+        partition_column="partition_id",
+        partition_filter_mode="ignore_partitions",
+    )
+    assert result is not None
+    assert isinstance(result, RegressionMetrics)
+
+
+@pytest.mark.parametrize("data_type", ["dataset", "datasource"])
+def test_evaluate_with_partition_column_datasource(partitioned_regression_model: RegressionModel, data_type):
+    """Test evaluate with partition_column on a Datasource"""
+    # Create a test datasource with partition_id column
+    eval_data_with_partition = [
+        {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+        {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
+        {"value": "homemade soup recipes", "score": 0.1, "partition_id": "p2"},
+        {"value": "cats purr when happy", "score": 0.9, "partition_id": "p2"},
+    ]
+
+    if data_type == "dataset":
+        eval_data = Dataset.from_list(eval_data_with_partition)
+        result = partitioned_regression_model.evaluate(
+            eval_data,
+            partition_column="partition_id",
+            partition_filter_mode="exclude_global",
+        )
+    else:
+        eval_datasource = Datasource.from_list("eval_datasource_with_partition_regression", eval_data_with_partition)
+        result = partitioned_regression_model.evaluate(
+            eval_datasource,
+            partition_column="partition_id",
+            partition_filter_mode="exclude_global",
+        )
+
+    assert result is not None
+    assert isinstance(result, RegressionMetrics)
+    assert isinstance(result.mae, float)
+    assert isinstance(result.mse, float)
+
+
 def test_predict(regression_model: RegressionModel):
     predictions = regression_model.predict(["Do you love soup?", "Are cats cute?"])
     assert len(predictions) == 2
@@ -248,6 +382,85 @@ def test_predict_with_prompt(regression_model: RegressionModel):
     assert 0 <= prediction_without_prompt.confidence <= 1
 
 
+def test_predict_with_partition_id(partitioned_regression_model: RegressionModel):
+    """Test predict with a specific partition_id"""
+    # Predict with partition_id p1 - should use memories from p1
+    prediction = partitioned_regression_model.predict("soup", partition_id="p1", partition_filter_mode="exclude_global")
+    assert prediction.score is not None
+    assert 0 <= prediction.confidence <= 1
+
+    # Predict with partition_id p2 - should use memories from p2
+    prediction_p2 = partitioned_regression_model.predict(
+        "cats", partition_id="p2", partition_filter_mode="exclude_global"
+    )
+    assert prediction_p2.score is not None
+    assert 0 <= prediction_p2.confidence <= 1
+
+
+def test_predict_with_partition_id_include_global(partitioned_regression_model: RegressionModel):
+    """Test predict with partition_id and include_global mode (default)"""
+    # Predict with partition_id p1 and include_global (default) - should include both p1 and global memories
+    prediction = partitioned_regression_model.predict("soup", partition_id="p1", partition_filter_mode="include_global")
+    assert prediction.score is not None
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_with_partition_id_exclude_global(partitioned_regression_model: RegressionModel):
+    """Test predict with partition_id and exclude_global mode"""
+    # Predict with partition_id p1 and exclude_global - should only use p1 memories
+    prediction = partitioned_regression_model.predict("soup", partition_id="p1", partition_filter_mode="exclude_global")
+    assert prediction.score is not None
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_with_partition_id_only_global(partitioned_regression_model: RegressionModel):
+    """Test predict with partition_filter_mode only_global"""
+    # Predict with only_global mode - should only use global memories
+    prediction = partitioned_regression_model.predict("cats", partition_filter_mode="only_global")
+    assert prediction.score is not None
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_with_partition_id_ignore_partitions(partitioned_regression_model: RegressionModel):
+    """Test predict with partition_filter_mode ignore_partitions"""
+    # Predict with ignore_partitions mode - should ignore partition filtering
+    prediction = partitioned_regression_model.predict("soup", partition_filter_mode="ignore_partitions")
+    assert prediction.score is not None
+    assert 0 <= prediction.confidence <= 1
+
+
+def test_predict_batch_with_partition_id(partitioned_regression_model: RegressionModel):
+    """Test batch predict with partition_id"""
+    # Batch predict with partition_id p1
+    predictions = partitioned_regression_model.predict(
+        ["soup is good", "cats are cute"],
+        partition_id="p1",
+        partition_filter_mode="exclude_global",
+    )
+    assert len(predictions) == 2
+    assert all(p.score is not None for p in predictions)
+    assert all(0 <= p.confidence <= 1 for p in predictions)
+
+
+def test_predict_batch_with_list_of_partition_ids(partitioned_regression_model: RegressionModel):
+    """Test batch predict with a list of partition_ids (one for each query input)"""
+    # Batch predict with a list of partition_ids - one for each input
+    # First input uses p1, second input uses p2
+    predictions = partitioned_regression_model.predict(
+        ["soup is good", "cats are cute"],
+        partition_id=["p1", "p2"],
+        partition_filter_mode="exclude_global",
+    )
+    assert len(predictions) == 2
+    assert all(p.score is not None for p in predictions)
+    assert all(0 <= p.confidence <= 1 for p in predictions)
+
+    # Verify that predictions were made using the correct partitions
+    # Each prediction should use memories from its respective partition
+    assert predictions[0].input_value == "soup is good"
+    assert predictions[1].input_value == "cats are cute"
+
+
 def test_record_prediction_feedback(regression_model: RegressionModel):
     predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
     expected_scores = [0.9, 0.1]

{orca_sdk-0.1.4.dist-info → orca_sdk-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orca_sdk
-Version: 0.1.4
+Version: 0.1.6
 Summary: SDK for interacting with Orca Services
 License-Expression: Apache-2.0
 Author: Orca DB Inc.

{orca_sdk-0.1.4.dist-info → orca_sdk-0.1.6.dist-info}/RECORD

@@ -1,7 +1,7 @@
 orca_sdk/__init__.py,sha256=xyjNwkLQXaX8A-UYgGwYDjv2btOXArT_yiMTfmW7KA8,1003
 orca_sdk/_shared/__init__.py,sha256=3Kt0Hu3QLI5FEp9nqGTxqAm3hAoBJKcagfaGQZ-lbJQ,223
-orca_sdk/_shared/metrics.py,sha256=LEZfAUWUtUWv_WWy9F_yjGLlUQHQpmR9WxG2fbKxa7U,14419
-orca_sdk/_shared/metrics_test.py,sha256=Rw1MaH37FppNsMnW8Ir9vMd8xxnZt3eo2Iypx1igtBI,9440
+orca_sdk/_shared/metrics.py,sha256=a_FdsPGDjR3CMOEBaEhEBqMfWUg7sqz9Jeh26XzAeg0,19756
+orca_sdk/_shared/metrics_test.py,sha256=n7eEAT8e6RqbI94ftEDljTBzOuh-YkFpXfF3DOoZA10,12905
 orca_sdk/_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 orca_sdk/_utils/analysis_ui.py,sha256=nT-M_YcNRCVPQzvuqYNFKnNHhYkADYBvq1GlIUePrWw,9232
 orca_sdk/_utils/analysis_ui_style.css,sha256=q_ba_-_KtgztepHg829zLzypaxKayl7ySC1-oYDzV3k,836
@@ -9,7 +9,7 @@ orca_sdk/_utils/auth.py,sha256=nC252O171_3_wn4KBAN7kg8GNvoZFiQ5Xtzkrm5dWDo,2645
 orca_sdk/_utils/auth_test.py,sha256=ygVWv1Ex53LaxIP7p2hzPHl8l9qYyBD5IGmEFJMps6s,1056
 orca_sdk/_utils/common.py,sha256=wUm2pNDWytEecC5WiDWd02-yCZw3Akx0bIutG4lHsFA,805
 orca_sdk/_utils/data_parsing.py,sha256=gkAwWEC8qRt3vRUObe7n7Pr0azOayNwc2yFY04WFp7E,5220
-orca_sdk/_utils/data_parsing_test.py,sha256=fNEYzPzE1jt3KWE2Kj91KqIeuv-L5REHFAa98zkNGSQ,8962
+orca_sdk/_utils/data_parsing_test.py,sha256=mJIPDOfs9fTsrQO_oMGonVAtNCA3_gcw_3qSmOX5Q4g,8978
 orca_sdk/_utils/pagination.py,sha256=986z0QPZixrZeurJWorF6eMgnTRdDF84AagEA6qNbMw,4245
 orca_sdk/_utils/pagination_test.py,sha256=BUylCrcHnwoKEBmMUzVr0lwLpA35ivcCwdBK4rMw9y8,4887
 orca_sdk/_utils/prediction_result_ui.css,sha256=sqBlkRLnovb5X5EcUDdB6iGpH63nVRlTW4uAmXuD0WM,258
@@ -17,25 +17,25 @@ orca_sdk/_utils/prediction_result_ui.py,sha256=Ur_FY7dz3oWNmtPiP3Wl3yRlEMgK8q9Uf
 orca_sdk/_utils/tqdm_file_reader.py,sha256=Lw7Cg1UgNuRUoN6jjqZb-IlV00H-kbRcrZLdudr1GxE,324
 orca_sdk/_utils/value_parser.py,sha256=c3qMABCCDQcIjn9N1orYYnlRwDW9JWdGwW_2TDZPLdI,1286
 orca_sdk/_utils/value_parser_test.py,sha256=OybsiC-Obi32RRi9NIuwrVBRAnlyPMV1xVAaevSrb7M,1079
-orca_sdk/async_client.py,sha256=mBd8z5xuHpE8-7Zd0D7YjH2e1OHSO-sRm2tSddqsc9Q,130387
-orca_sdk/classification_model.py,sha256=iYUrGjeYHvPvXwYXjXU_LGL7Dn2XxUcGCt6w93DlJO8,41702
-orca_sdk/classification_model_test.py,sha256=_gaDg8QB0h0ByN4UwTk2fIIDXE4UzahuJBjz7NSPK28,23605
-orca_sdk/client.py,sha256=voNo4NPsc-rsZQ3lZO2fsFuFLw4DC4Dl9REVJQEyKhY,129454
-orca_sdk/conftest.py,sha256=RtINF1xea2iMycMkpMXIOOqRbfWeIZsceSAemhBmgNE,9761
+orca_sdk/async_client.py,sha256=y2D3fPQZmbwmtYWAk5acJ45atSZen9MNfjP2tKjpP6Q,132737
+orca_sdk/classification_model.py,sha256=4AcQvAm0EN7w0qx0WpgEs7VUoIIPTqIVE86wtkaIAYs,46249
+orca_sdk/classification_model_test.py,sha256=vBn7KBb9-ACuJEdzW50n54Fn6Mh9iEYbn1197lE8-yI,36997
+orca_sdk/client.py,sha256=oQd8Lm0agetLyAdVRP8IZqe6S5mjxhFSnbVHqhT7dmI,131798
+orca_sdk/conftest.py,sha256=0O1VY-SPKNAvi9fBLdY1RMnYVgZvMjP92y99bNAqqiw,12461
 orca_sdk/credentials.py,sha256=80_1r8n5jruEvN_E629SaRrRhKvF_NhWUEZyZzPXkqQ,6620
 orca_sdk/credentials_test.py,sha256=TLbXJMz3IlThvtSrHeLM7jRsKnrncA_ahOTpHg15Ei4,4089
 orca_sdk/datasource.py,sha256=6QaccghiyFEUSFcqnwjIJzpgIh9Id0snJk2EqViqPsU,22356
 orca_sdk/datasource_test.py,sha256=sCk3IcQJbDut5oN4Wf7PXhTxyMwalxMuCXJekSxy9wk,16665
-orca_sdk/embedding_model.py,sha256=bZhbNJBimWc9Ryklza3q9HS0MRWsiH5Lhn6p7pff0RI,28165
+orca_sdk/embedding_model.py,sha256=4xxfo26b5X_YJtU8KyqoMmJQ6VgfHEcYftVSz-RfDng,27920
 orca_sdk/embedding_model_test.py,sha256=-NItbNb3tTVj5jAvSi3WjV3FP448q08lmT5iObg9vwA,8133
 orca_sdk/job.py,sha256=wHwVt-s7i-v8udhLGybB-90Kp4dwOLrY806bE4Tam5Q,13092
 orca_sdk/job_test.py,sha256=nRSWxd_1UIfrj9oMVvrXjt6OBkBpddYAjb2y6P-DTUg,4327
-orca_sdk/memoryset.py,sha256=QSnHA2SpAJkGdpVd8wQX2weAhLu9Iw-lfpeQvJxLedg,111690
-orca_sdk/memoryset_test.py,sha256=7wGOtbVa3MEu91fN8DTjiYgB6QIObuA3cTchHmddTIk,33551
-orca_sdk/regression_model.py,sha256=GIL-KgKtGzdb5dFraOKu6OD8yrcavc-CeXASPsKGLGM,28086
-orca_sdk/regression_model_test.py,sha256=slwxbty_vL9d24OCn5xN61eKyri5GS7Jv2YmpEOMTrM,15856
+orca_sdk/memoryset.py,sha256=OtIZH14bN7tejevdsHJpBeje5Ovc58vtkw1inRMWdw4,117709
+orca_sdk/memoryset_test.py,sha256=YrMP_Z7BPhKAIZq1FKar83dNzR_2b56A89pnEiUbRrs,46509
+orca_sdk/regression_model.py,sha256=vXdY2Fbfc0MyECUR3fa_IR-nETPrDN7VFAdjvsgHPrs,31382
+orca_sdk/regression_model_test.py,sha256=DfWLkqxB835jjwM-sj1uxQ6Yz_ZBMnt8EHjdfnHsRnU,25103
 orca_sdk/telemetry.py,sha256=ZyCMiyyo_SchjadWZH55TlLrC4Ucq5S316NbW26LL4Y,27834
 orca_sdk/telemetry_test.py,sha256=eT66C5lFdNg-pQdo2I__BP7Tn5fTc9aTkVo9ZhWwhU0,5519
-orca_sdk-0.1.4.dist-info/METADATA,sha256=AQBTSp780409HcaGN9ozHCYUCElgNqP30XP8u4fyBiw,3659
-orca_sdk-0.1.4.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-orca_sdk-0.1.4.dist-info/RECORD,,
+orca_sdk-0.1.6.dist-info/METADATA,sha256=85QDZDP9Uxda4oZ3BMPP_kI5T4GPy1mFMYtWh1-nI54,3659
+orca_sdk-0.1.6.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+orca_sdk-0.1.6.dist-info/RECORD,,