orca-sdk 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (186) hide show
  1. orca_sdk/__init__.py +10 -4
  2. orca_sdk/_shared/__init__.py +10 -0
  3. orca_sdk/_shared/metrics.py +393 -0
  4. orca_sdk/_shared/metrics_test.py +273 -0
  5. orca_sdk/_utils/analysis_ui.py +12 -10
  6. orca_sdk/_utils/analysis_ui_style.css +0 -3
  7. orca_sdk/_utils/auth.py +31 -29
  8. orca_sdk/_utils/data_parsing.py +28 -2
  9. orca_sdk/_utils/data_parsing_test.py +15 -15
  10. orca_sdk/_utils/pagination.py +126 -0
  11. orca_sdk/_utils/pagination_test.py +132 -0
  12. orca_sdk/_utils/prediction_result_ui.py +67 -21
  13. orca_sdk/_utils/tqdm_file_reader.py +12 -0
  14. orca_sdk/_utils/value_parser.py +45 -0
  15. orca_sdk/_utils/value_parser_test.py +39 -0
  16. orca_sdk/async_client.py +3795 -0
  17. orca_sdk/classification_model.py +601 -129
  18. orca_sdk/classification_model_test.py +415 -117
  19. orca_sdk/client.py +3787 -0
  20. orca_sdk/conftest.py +184 -38
  21. orca_sdk/credentials.py +162 -20
  22. orca_sdk/credentials_test.py +100 -16
  23. orca_sdk/datasource.py +268 -68
  24. orca_sdk/datasource_test.py +266 -18
  25. orca_sdk/embedding_model.py +434 -82
  26. orca_sdk/embedding_model_test.py +66 -33
  27. orca_sdk/job.py +343 -0
  28. orca_sdk/job_test.py +108 -0
  29. orca_sdk/memoryset.py +1690 -324
  30. orca_sdk/memoryset_test.py +456 -119
  31. orca_sdk/regression_model.py +694 -0
  32. orca_sdk/regression_model_test.py +378 -0
  33. orca_sdk/telemetry.py +460 -143
  34. orca_sdk/telemetry_test.py +43 -24
  35. {orca_sdk-0.1.1.dist-info → orca_sdk-0.1.3.dist-info}/METADATA +34 -16
  36. orca_sdk-0.1.3.dist-info/RECORD +41 -0
  37. {orca_sdk-0.1.1.dist-info → orca_sdk-0.1.3.dist-info}/WHEEL +1 -1
  38. orca_sdk/_generated_api_client/__init__.py +0 -3
  39. orca_sdk/_generated_api_client/api/__init__.py +0 -193
  40. orca_sdk/_generated_api_client/api/auth/__init__.py +0 -0
  41. orca_sdk/_generated_api_client/api/auth/check_authentication_auth_get.py +0 -128
  42. orca_sdk/_generated_api_client/api/auth/create_api_key_auth_api_key_post.py +0 -170
  43. orca_sdk/_generated_api_client/api/auth/delete_api_key_auth_api_key_name_or_id_delete.py +0 -156
  44. orca_sdk/_generated_api_client/api/auth/delete_org_auth_org_delete.py +0 -130
  45. orca_sdk/_generated_api_client/api/auth/list_api_keys_auth_api_key_get.py +0 -127
  46. orca_sdk/_generated_api_client/api/classification_model/__init__.py +0 -0
  47. orca_sdk/_generated_api_client/api/classification_model/create_evaluation_classification_model_model_name_or_id_evaluation_post.py +0 -183
  48. orca_sdk/_generated_api_client/api/classification_model/create_model_classification_model_post.py +0 -170
  49. orca_sdk/_generated_api_client/api/classification_model/delete_evaluation_classification_model_model_name_or_id_evaluation_task_id_delete.py +0 -168
  50. orca_sdk/_generated_api_client/api/classification_model/delete_model_classification_model_name_or_id_delete.py +0 -154
  51. orca_sdk/_generated_api_client/api/classification_model/get_evaluation_classification_model_model_name_or_id_evaluation_task_id_get.py +0 -170
  52. orca_sdk/_generated_api_client/api/classification_model/get_model_classification_model_name_or_id_get.py +0 -156
  53. orca_sdk/_generated_api_client/api/classification_model/list_evaluations_classification_model_model_name_or_id_evaluation_get.py +0 -161
  54. orca_sdk/_generated_api_client/api/classification_model/list_models_classification_model_get.py +0 -127
  55. orca_sdk/_generated_api_client/api/classification_model/predict_gpu_classification_model_name_or_id_prediction_post.py +0 -190
  56. orca_sdk/_generated_api_client/api/datasource/__init__.py +0 -0
  57. orca_sdk/_generated_api_client/api/datasource/create_datasource_datasource_post.py +0 -167
  58. orca_sdk/_generated_api_client/api/datasource/delete_datasource_datasource_name_or_id_delete.py +0 -156
  59. orca_sdk/_generated_api_client/api/datasource/get_datasource_datasource_name_or_id_get.py +0 -156
  60. orca_sdk/_generated_api_client/api/datasource/list_datasources_datasource_get.py +0 -127
  61. orca_sdk/_generated_api_client/api/default/__init__.py +0 -0
  62. orca_sdk/_generated_api_client/api/default/healthcheck_get.py +0 -118
  63. orca_sdk/_generated_api_client/api/default/healthcheck_gpu_get.py +0 -118
  64. orca_sdk/_generated_api_client/api/finetuned_embedding_model/__init__.py +0 -0
  65. orca_sdk/_generated_api_client/api/finetuned_embedding_model/create_finetuned_embedding_model_finetuned_embedding_model_post.py +0 -168
  66. orca_sdk/_generated_api_client/api/finetuned_embedding_model/delete_finetuned_embedding_model_finetuned_embedding_model_name_or_id_delete.py +0 -156
  67. orca_sdk/_generated_api_client/api/finetuned_embedding_model/embed_with_finetuned_model_gpu_finetuned_embedding_model_name_or_id_embedding_post.py +0 -189
  68. orca_sdk/_generated_api_client/api/finetuned_embedding_model/get_finetuned_embedding_model_finetuned_embedding_model_name_or_id_get.py +0 -156
  69. orca_sdk/_generated_api_client/api/finetuned_embedding_model/list_finetuned_embedding_models_finetuned_embedding_model_get.py +0 -127
  70. orca_sdk/_generated_api_client/api/memoryset/__init__.py +0 -0
  71. orca_sdk/_generated_api_client/api/memoryset/clone_memoryset_memoryset_name_or_id_clone_post.py +0 -181
  72. orca_sdk/_generated_api_client/api/memoryset/create_analysis_memoryset_name_or_id_analysis_post.py +0 -183
  73. orca_sdk/_generated_api_client/api/memoryset/create_memoryset_memoryset_post.py +0 -168
  74. orca_sdk/_generated_api_client/api/memoryset/delete_memories_memoryset_name_or_id_memories_delete_post.py +0 -181
  75. orca_sdk/_generated_api_client/api/memoryset/delete_memory_memoryset_name_or_id_memory_memory_id_delete.py +0 -167
  76. orca_sdk/_generated_api_client/api/memoryset/delete_memoryset_memoryset_name_or_id_delete.py +0 -156
  77. orca_sdk/_generated_api_client/api/memoryset/get_analysis_memoryset_name_or_id_analysis_analysis_task_id_get.py +0 -169
  78. orca_sdk/_generated_api_client/api/memoryset/get_memories_memoryset_name_or_id_memories_get_post.py +0 -188
  79. orca_sdk/_generated_api_client/api/memoryset/get_memory_memoryset_name_or_id_memory_memory_id_get.py +0 -169
  80. orca_sdk/_generated_api_client/api/memoryset/get_memoryset_memoryset_name_or_id_get.py +0 -156
  81. orca_sdk/_generated_api_client/api/memoryset/insert_memories_gpu_memoryset_name_or_id_memory_post.py +0 -184
  82. orca_sdk/_generated_api_client/api/memoryset/list_analyses_memoryset_name_or_id_analysis_get.py +0 -260
  83. orca_sdk/_generated_api_client/api/memoryset/list_memorysets_memoryset_get.py +0 -127
  84. orca_sdk/_generated_api_client/api/memoryset/memoryset_lookup_gpu_memoryset_name_or_id_lookup_post.py +0 -193
  85. orca_sdk/_generated_api_client/api/memoryset/query_memoryset_memoryset_name_or_id_memories_post.py +0 -188
  86. orca_sdk/_generated_api_client/api/memoryset/update_memories_gpu_memoryset_name_or_id_memories_patch.py +0 -191
  87. orca_sdk/_generated_api_client/api/memoryset/update_memory_gpu_memoryset_name_or_id_memory_patch.py +0 -187
  88. orca_sdk/_generated_api_client/api/pretrained_embedding_model/__init__.py +0 -0
  89. orca_sdk/_generated_api_client/api/pretrained_embedding_model/embed_with_pretrained_model_gpu_pretrained_embedding_model_model_name_embedding_post.py +0 -188
  90. orca_sdk/_generated_api_client/api/pretrained_embedding_model/get_pretrained_embedding_model_pretrained_embedding_model_model_name_get.py +0 -157
  91. orca_sdk/_generated_api_client/api/pretrained_embedding_model/list_pretrained_embedding_models_pretrained_embedding_model_get.py +0 -127
  92. orca_sdk/_generated_api_client/api/task/__init__.py +0 -0
  93. orca_sdk/_generated_api_client/api/task/abort_task_task_task_id_abort_delete.py +0 -154
  94. orca_sdk/_generated_api_client/api/task/get_task_status_task_task_id_status_get.py +0 -156
  95. orca_sdk/_generated_api_client/api/task/list_tasks_task_get.py +0 -243
  96. orca_sdk/_generated_api_client/api/telemetry/__init__.py +0 -0
  97. orca_sdk/_generated_api_client/api/telemetry/drop_feedback_category_with_data_telemetry_feedback_category_name_or_id_delete.py +0 -162
  98. orca_sdk/_generated_api_client/api/telemetry/get_feedback_category_telemetry_feedback_category_name_or_id_get.py +0 -156
  99. orca_sdk/_generated_api_client/api/telemetry/get_prediction_telemetry_prediction_prediction_id_get.py +0 -157
  100. orca_sdk/_generated_api_client/api/telemetry/list_feedback_categories_telemetry_feedback_category_get.py +0 -127
  101. orca_sdk/_generated_api_client/api/telemetry/list_predictions_telemetry_prediction_post.py +0 -175
  102. orca_sdk/_generated_api_client/api/telemetry/record_prediction_feedback_telemetry_prediction_feedback_put.py +0 -171
  103. orca_sdk/_generated_api_client/api/telemetry/update_prediction_telemetry_prediction_prediction_id_patch.py +0 -181
  104. orca_sdk/_generated_api_client/client.py +0 -216
  105. orca_sdk/_generated_api_client/errors.py +0 -38
  106. orca_sdk/_generated_api_client/models/__init__.py +0 -159
  107. orca_sdk/_generated_api_client/models/analyze_neighbor_labels_result.py +0 -84
  108. orca_sdk/_generated_api_client/models/api_key_metadata.py +0 -118
  109. orca_sdk/_generated_api_client/models/base_model.py +0 -55
  110. orca_sdk/_generated_api_client/models/body_create_datasource_datasource_post.py +0 -176
  111. orca_sdk/_generated_api_client/models/classification_evaluation_result.py +0 -114
  112. orca_sdk/_generated_api_client/models/clone_labeled_memoryset_request.py +0 -150
  113. orca_sdk/_generated_api_client/models/column_info.py +0 -114
  114. orca_sdk/_generated_api_client/models/column_type.py +0 -14
  115. orca_sdk/_generated_api_client/models/conflict_error_response.py +0 -80
  116. orca_sdk/_generated_api_client/models/create_api_key_request.py +0 -99
  117. orca_sdk/_generated_api_client/models/create_api_key_response.py +0 -126
  118. orca_sdk/_generated_api_client/models/create_labeled_memoryset_request.py +0 -259
  119. orca_sdk/_generated_api_client/models/create_rac_model_request.py +0 -209
  120. orca_sdk/_generated_api_client/models/datasource_metadata.py +0 -142
  121. orca_sdk/_generated_api_client/models/delete_memories_request.py +0 -70
  122. orca_sdk/_generated_api_client/models/embed_request.py +0 -127
  123. orca_sdk/_generated_api_client/models/embedding_finetuning_method.py +0 -9
  124. orca_sdk/_generated_api_client/models/evaluation_request.py +0 -180
  125. orca_sdk/_generated_api_client/models/evaluation_response.py +0 -140
  126. orca_sdk/_generated_api_client/models/feedback_type.py +0 -9
  127. orca_sdk/_generated_api_client/models/field_validation_error.py +0 -103
  128. orca_sdk/_generated_api_client/models/filter_item.py +0 -231
  129. orca_sdk/_generated_api_client/models/filter_item_field_type_0_item.py +0 -15
  130. orca_sdk/_generated_api_client/models/filter_item_field_type_2_item_type_1.py +0 -16
  131. orca_sdk/_generated_api_client/models/filter_item_op.py +0 -16
  132. orca_sdk/_generated_api_client/models/find_duplicates_analysis_result.py +0 -70
  133. orca_sdk/_generated_api_client/models/finetune_embedding_model_request.py +0 -259
  134. orca_sdk/_generated_api_client/models/finetune_embedding_model_request_training_args.py +0 -66
  135. orca_sdk/_generated_api_client/models/finetuned_embedding_model_metadata.py +0 -166
  136. orca_sdk/_generated_api_client/models/get_memories_request.py +0 -70
  137. orca_sdk/_generated_api_client/models/internal_server_error_response.py +0 -80
  138. orca_sdk/_generated_api_client/models/label_class_metrics.py +0 -108
  139. orca_sdk/_generated_api_client/models/label_prediction_memory_lookup.py +0 -274
  140. orca_sdk/_generated_api_client/models/label_prediction_memory_lookup_metadata.py +0 -68
  141. orca_sdk/_generated_api_client/models/label_prediction_result.py +0 -101
  142. orca_sdk/_generated_api_client/models/label_prediction_with_memories_and_feedback.py +0 -232
  143. orca_sdk/_generated_api_client/models/labeled_memory.py +0 -197
  144. orca_sdk/_generated_api_client/models/labeled_memory_insert.py +0 -108
  145. orca_sdk/_generated_api_client/models/labeled_memory_insert_metadata.py +0 -68
  146. orca_sdk/_generated_api_client/models/labeled_memory_lookup.py +0 -258
  147. orca_sdk/_generated_api_client/models/labeled_memory_lookup_metadata.py +0 -68
  148. orca_sdk/_generated_api_client/models/labeled_memory_metadata.py +0 -68
  149. orca_sdk/_generated_api_client/models/labeled_memory_metrics.py +0 -277
  150. orca_sdk/_generated_api_client/models/labeled_memory_update.py +0 -171
  151. orca_sdk/_generated_api_client/models/labeled_memory_update_metadata_type_0.py +0 -68
  152. orca_sdk/_generated_api_client/models/labeled_memoryset_metadata.py +0 -195
  153. orca_sdk/_generated_api_client/models/list_analyses_memoryset_name_or_id_analysis_get_type_type_0.py +0 -9
  154. orca_sdk/_generated_api_client/models/list_memories_request.py +0 -104
  155. orca_sdk/_generated_api_client/models/list_predictions_request.py +0 -234
  156. orca_sdk/_generated_api_client/models/list_predictions_request_sort_item_item_type_0.py +0 -9
  157. orca_sdk/_generated_api_client/models/list_predictions_request_sort_item_item_type_1.py +0 -9
  158. orca_sdk/_generated_api_client/models/lookup_request.py +0 -81
  159. orca_sdk/_generated_api_client/models/memoryset_analysis_request.py +0 -83
  160. orca_sdk/_generated_api_client/models/memoryset_analysis_request_type.py +0 -9
  161. orca_sdk/_generated_api_client/models/memoryset_analysis_response.py +0 -180
  162. orca_sdk/_generated_api_client/models/memoryset_analysis_response_config.py +0 -66
  163. orca_sdk/_generated_api_client/models/memoryset_analysis_response_type.py +0 -9
  164. orca_sdk/_generated_api_client/models/not_found_error_response.py +0 -100
  165. orca_sdk/_generated_api_client/models/not_found_error_response_resource_type_0.py +0 -20
  166. orca_sdk/_generated_api_client/models/prediction_feedback.py +0 -157
  167. orca_sdk/_generated_api_client/models/prediction_feedback_category.py +0 -115
  168. orca_sdk/_generated_api_client/models/prediction_feedback_request.py +0 -122
  169. orca_sdk/_generated_api_client/models/prediction_feedback_result.py +0 -102
  170. orca_sdk/_generated_api_client/models/prediction_request.py +0 -169
  171. orca_sdk/_generated_api_client/models/pretrained_embedding_model_metadata.py +0 -97
  172. orca_sdk/_generated_api_client/models/pretrained_embedding_model_name.py +0 -11
  173. orca_sdk/_generated_api_client/models/rac_head_type.py +0 -11
  174. orca_sdk/_generated_api_client/models/rac_model_metadata.py +0 -191
  175. orca_sdk/_generated_api_client/models/service_unavailable_error_response.py +0 -80
  176. orca_sdk/_generated_api_client/models/task.py +0 -198
  177. orca_sdk/_generated_api_client/models/task_status.py +0 -14
  178. orca_sdk/_generated_api_client/models/task_status_info.py +0 -133
  179. orca_sdk/_generated_api_client/models/unauthenticated_error_response.py +0 -72
  180. orca_sdk/_generated_api_client/models/unauthorized_error_response.py +0 -80
  181. orca_sdk/_generated_api_client/models/unprocessable_input_error_response.py +0 -94
  182. orca_sdk/_generated_api_client/models/update_prediction_request.py +0 -93
  183. orca_sdk/_generated_api_client/py.typed +0 -1
  184. orca_sdk/_generated_api_client/types.py +0 -56
  185. orca_sdk/_utils/task.py +0 -73
  186. orca_sdk-0.1.1.dist-info/RECORD +0 -175
orca_sdk/conftest.py CHANGED
@@ -4,20 +4,42 @@ from typing import Generator
4
4
  from uuid import uuid4
5
5
 
6
6
  import pytest
7
+ import pytest_asyncio
7
8
  from datasets import ClassLabel, Dataset, Features, Value
8
9
 
9
- from ._generated_api_client.client import set_headers
10
10
  from ._utils.auth import _create_api_key, _delete_org
11
+ from .async_client import OrcaAsyncClient
11
12
  from .classification_model import ClassificationModel
13
+ from .client import OrcaClient
12
14
  from .credentials import OrcaCredentials
13
15
  from .datasource import Datasource
14
16
  from .embedding_model import PretrainedEmbeddingModel
15
- from .memoryset import LabeledMemoryset
17
+ from .memoryset import LabeledMemoryset, ScoredMemoryset
18
+ from .regression_model import RegressionModel
16
19
 
17
- logging.basicConfig(level=logging.INFO)
20
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
18
21
 
19
22
  os.environ["ORCA_API_URL"] = os.environ.get("ORCA_API_URL", "http://localhost:1584/")
20
23
 
24
+ os.environ["ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY"] = "true"
25
+
26
+
27
+ def skip_in_prod(reason: str):
28
+ """Custom decorator to skip tests when running against production API"""
29
+ PROD_API_URLs = ["https://api.orcadb.ai", "https://api.staging.orcadb.ai"]
30
+ return pytest.mark.skipif(
31
+ os.environ["ORCA_API_URL"] in PROD_API_URLs,
32
+ reason=reason,
33
+ )
34
+
35
+
36
+ def skip_in_ci(reason: str):
37
+ """Custom decorator to skip tests when running in CI"""
38
+ return pytest.mark.skipif(
39
+ os.environ.get("GITHUB_ACTIONS", "false") == "true",
40
+ reason=reason,
41
+ )
42
+
21
43
 
22
44
  def _create_org_id():
23
45
  # UUID start to identify test data (0xtest...)
@@ -29,39 +51,46 @@ def org_id():
29
51
  return _create_org_id()
30
52
 
31
53
 
54
+ @pytest.fixture(scope="session")
55
+ def other_org_id():
56
+ return _create_org_id()
57
+
58
+
32
59
  @pytest.fixture(autouse=True, scope="session")
33
60
  def api_key(org_id) -> Generator[str, None, None]:
34
61
  api_key = _create_api_key(org_id=org_id, name="orca_sdk_test")
35
- OrcaCredentials.set_api_key(api_key, check_validity=True)
36
- yield api_key
62
+ with OrcaClient(api_key=api_key).use():
63
+ yield api_key
37
64
  _delete_org(org_id)
38
65
 
39
66
 
67
+ # We cannot use a session scoped fixture because async pytest tears down the client after each test
40
68
  @pytest.fixture(autouse=True)
41
- def authenticated(api_key):
42
- OrcaCredentials.set_api_key(api_key, check_validity=False)
69
+ def authenticate_async_client(api_key) -> Generator[None, None, None]:
70
+ with OrcaAsyncClient(api_key=api_key).use():
71
+ yield
43
72
 
44
73
 
45
- @pytest.fixture()
46
- def unauthenticated(api_key):
47
- OrcaCredentials.set_api_key(str(uuid4()), check_validity=False)
48
- yield
49
- # Need to reset the api key to the original api key so following tests don't fail
50
- OrcaCredentials.set_api_key(api_key, check_validity=False)
74
+ @pytest.fixture(scope="session")
75
+ def unauthenticated_client() -> OrcaClient:
76
+ return OrcaClient(api_key=str(uuid4()))
51
77
 
52
78
 
53
- @pytest.fixture()
54
- def other_org_id():
55
- return _create_org_id()
79
+ @pytest_asyncio.fixture()
80
+ def unauthenticated_async_client() -> OrcaAsyncClient:
81
+ return OrcaAsyncClient(api_key=str(uuid4()))
56
82
 
57
83
 
58
- @pytest.fixture()
59
- def unauthorized(api_key, other_org_id):
84
+ @pytest.fixture(scope="session")
85
+ def unauthorized_client(other_org_id):
60
86
  different_api_key = _create_api_key(org_id=other_org_id, name="orca_sdk_test_other_org")
61
- OrcaCredentials.set_api_key(different_api_key, check_validity=False)
62
- yield
63
- OrcaCredentials.set_api_key(api_key, check_validity=False)
64
- _delete_org(other_org_id)
87
+ return OrcaClient(api_key=different_api_key)
88
+
89
+
90
+ @pytest.fixture(scope="session")
91
+ def predict_only_client() -> OrcaClient:
92
+ predict_api_key = OrcaCredentials.create_api_key("orca_sdk_test_predict", scopes={"PREDICT"})
93
+ return OrcaClient(api_key=predict_api_key)
65
94
 
66
95
 
67
96
  @pytest.fixture(scope="session")
@@ -70,22 +99,44 @@ def label_names():
70
99
 
71
100
 
72
101
  SAMPLE_DATA = [
73
- {"text": "i love soup", "label": 0, "key": "val1", "score": 0.1, "source_id": "s1"},
74
- {"text": "cats are cute", "label": 1, "key": "val2", "score": 0.2, "source_id": "s2"},
75
- {"text": "soup is good", "label": 0, "key": "val3", "score": 0.3, "source_id": "s3"},
76
- {"text": "i love cats", "label": 1, "key": "val4", "score": 0.4, "source_id": "s4"},
77
- {"text": "everyone loves cats", "label": 1, "key": "val5", "score": 0.5, "source_id": "s5"},
78
- {"text": "soup is great for the winter", "label": 0, "key": "val6", "score": 0.6, "source_id": "s6"},
102
+ {"value": "i love soup", "label": 0, "key": "g1", "score": 0.1, "source_id": "s1"},
103
+ {"value": "cats are cute", "label": 1, "key": "g1", "score": 0.9, "source_id": "s2"},
104
+ {"value": "soup is good", "label": 0, "key": "g1", "score": 0.1, "source_id": "s3"},
105
+ {"value": "i love cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s4"},
106
+ {"value": "everyone loves cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s5"},
107
+ {"value": "soup is great for the winter", "label": 0, "key": "g1", "score": 0.1, "source_id": "s6"},
108
+ {"value": "hot soup on a rainy day!", "label": 0, "key": "g1", "score": 0.1, "source_id": "s7"},
109
+ {"value": "cats sleep all day", "label": 1, "key": "g1", "score": 0.9, "source_id": "s8"},
110
+ {"value": "homemade soup recipes", "label": 0, "key": "g1", "score": 0.1, "source_id": "s9"},
111
+ {"value": "cats purr when happy", "label": 1, "key": "g2", "score": 0.9, "source_id": "s10"},
112
+ {"value": "chicken noodle soup is classic", "label": 0, "key": "g1", "score": 0.1, "source_id": "s11"},
113
+ {"value": "kittens are baby cats", "label": 1, "key": "g2", "score": 0.9, "source_id": "s12"},
114
+ {"value": "soup can be served cold too", "label": 0, "key": "g1", "score": 0.1, "source_id": "s13"},
115
+ {"value": "cats have nine lives", "label": 1, "key": "g2", "score": 0.9, "source_id": "s14"},
116
+ {"value": "tomato soup with grilled cheese", "label": 0, "key": "g1", "score": 0.1, "source_id": "s15"},
117
+ {"value": "cats are independent animals", "label": 1, "key": "g2", "score": 0.9, "source_id": "s16"},
118
+ {"value": "the beach is always fun", "label": None, "key": "g3", "score": None, "source_id": "s17"},
119
+ {"value": "i love the beach", "label": None, "key": "g3", "score": None, "source_id": "s18"},
120
+ {"value": "the ocean is healing", "label": None, "key": "g3", "score": None, "source_id": "s19"},
121
+ {
122
+ "value": "sandy feet, sand between my toes at the beach",
123
+ "label": None,
124
+ "key": "g3",
125
+ "score": None,
126
+ "source_id": "s20",
127
+ },
128
+ {"value": "i am such a beach bum", "label": None, "key": "g3", "score": None, "source_id": "s21"},
129
+ {"value": "i will always want to be at the beach", "label": None, "key": "g3", "score": None, "source_id": "s22"},
79
130
  ]
80
131
 
81
132
 
82
133
  @pytest.fixture(scope="session")
83
- def hf_dataset(label_names):
134
+ def hf_dataset(label_names: list[str]) -> Dataset:
84
135
  return Dataset.from_list(
85
136
  SAMPLE_DATA,
86
137
  features=Features(
87
138
  {
88
- "text": Value("string"),
139
+ "value": Value("string"),
89
140
  "label": ClassLabel(names=label_names),
90
141
  "key": Value("string"),
91
142
  "score": Value("float"),
@@ -96,22 +147,117 @@ def hf_dataset(label_names):
96
147
 
97
148
 
98
149
  @pytest.fixture(scope="session")
99
- def datasource(hf_dataset) -> Datasource:
100
- return Datasource.from_hf_dataset("test_datasource", hf_dataset)
150
+ def datasource(hf_dataset: Dataset) -> Datasource:
151
+ datasource = Datasource.from_hf_dataset("test_datasource", hf_dataset)
152
+ return datasource
153
+
154
+
155
+ EVAL_DATASET = [
156
+ {"value": "chicken noodle soup is the best", "label": 1, "score": 0.9}, # mislabeled
157
+ {"value": "cats are cute", "label": 0, "score": 0.1}, # mislabeled
158
+ {"value": "soup is great for the winter", "label": 0, "score": 0.1},
159
+ {"value": "i love cats", "label": 1, "score": 0.9},
160
+ ]
161
+
162
+
163
+ @pytest.fixture(scope="session")
164
+ def eval_datasource() -> Datasource:
165
+ eval_datasource = Datasource.from_list("eval_datasource", EVAL_DATASET)
166
+ return eval_datasource
167
+
168
+
169
+ @pytest.fixture(scope="session")
170
+ def eval_dataset() -> Dataset:
171
+ eval_dataset = Dataset.from_list(EVAL_DATASET)
172
+ return eval_dataset
173
+
174
+
175
+ @pytest.fixture(scope="session")
176
+ def readonly_memoryset(datasource: Datasource) -> LabeledMemoryset:
177
+ memoryset = LabeledMemoryset.create(
178
+ "test_readonly_memoryset",
179
+ datasource=datasource,
180
+ embedding_model=PretrainedEmbeddingModel.GTE_BASE,
181
+ source_id_column="source_id",
182
+ max_seq_length_override=32,
183
+ index_type="IVF_FLAT",
184
+ index_params={"n_lists": 100},
185
+ )
186
+ return memoryset
187
+
188
+
189
+ @pytest.fixture(scope="function")
190
+ def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[LabeledMemoryset, None, None]:
191
+ """
192
+ Function-scoped fixture that provides a writable memoryset for tests that mutate state.
193
+
194
+ This fixture creates a fresh `LabeledMemoryset` named 'test_writable_memoryset' before each test.
195
+ After the test, it attempts to restore the memoryset to its initial state by deleting any added entries
196
+ and reinserting sample data — unless the memoryset has been dropped by the test itself, in which case
197
+ it will be recreated on the next invocation.
198
+
199
+ Note: Re-creating the memoryset from scratch is surprisingly more expensive than cleaning it up.
200
+ """
201
+ # It shouldn't be possible for this memoryset to already exist
202
+ memoryset = LabeledMemoryset.create(
203
+ "test_writable_memoryset",
204
+ datasource=datasource,
205
+ embedding_model=PretrainedEmbeddingModel.GTE_BASE,
206
+ source_id_column="source_id",
207
+ max_seq_length_override=32,
208
+ if_exists="open",
209
+ )
210
+ try:
211
+ yield memoryset
212
+ finally:
213
+ # Restore the memoryset to a clean state for the next test.
214
+ with OrcaClient(api_key=api_key).use():
215
+ if LabeledMemoryset.exists("test_writable_memoryset"):
216
+ memoryset.refresh()
217
+
218
+ memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
219
+
220
+ if memory_ids:
221
+ memoryset.delete(memory_ids)
222
+ memoryset.refresh()
223
+ assert len(memoryset) == 0
224
+ memoryset.insert(SAMPLE_DATA)
225
+ # If the test dropped the memoryset, do nothing — it will be recreated on the next use.
226
+
227
+
228
+ @pytest.fixture(scope="session")
229
+ def classification_model(readonly_memoryset: LabeledMemoryset) -> ClassificationModel:
230
+ model = ClassificationModel.create(
231
+ "test_classification_model",
232
+ readonly_memoryset,
233
+ num_classes=2,
234
+ memory_lookup_count=3,
235
+ description="test_description",
236
+ )
237
+ return model
101
238
 
102
239
 
240
+ # Add scored memoryset and regression model fixtures
103
241
  @pytest.fixture(scope="session")
104
- def memoryset(datasource) -> LabeledMemoryset:
105
- return LabeledMemoryset.create(
106
- "test_memoryset",
242
+ def scored_memoryset(datasource: Datasource) -> ScoredMemoryset:
243
+ memoryset = ScoredMemoryset.create(
244
+ "test_scored_memoryset",
107
245
  datasource=datasource,
108
246
  embedding_model=PretrainedEmbeddingModel.GTE_BASE,
109
- value_column="text",
110
247
  source_id_column="source_id",
111
248
  max_seq_length_override=32,
249
+ index_type="IVF_FLAT",
250
+ index_params={"n_lists": 100},
112
251
  )
252
+ return memoryset
113
253
 
114
254
 
115
255
  @pytest.fixture(scope="session")
116
- def model(memoryset) -> ClassificationModel:
117
- return ClassificationModel.create("test_model", memoryset, num_classes=2, memory_lookup_count=3)
256
+ def regression_model(scored_memoryset: ScoredMemoryset) -> RegressionModel:
257
+ model = RegressionModel.create(
258
+ "test_regression_model",
259
+ scored_memoryset,
260
+ memory_lookup_count=3,
261
+ description="test_regression_description",
262
+ )
263
+ return model
orca_sdk/credentials.py CHANGED
@@ -1,21 +1,47 @@
1
+ import os
1
2
  from datetime import datetime
2
- from typing import NamedTuple
3
+ from typing import Literal, NamedTuple
3
4
 
4
- from ._generated_api_client.api import check_authentication, list_api_keys
5
- from ._generated_api_client.client import get_base_url, get_headers, set_headers
5
+ import httpx
6
+ from httpx import ConnectError, Headers, HTTPTransport
7
+ from typing_extensions import deprecated
6
8
 
9
+ from .async_client import OrcaAsyncClient
10
+ from .client import OrcaClient
7
11
 
8
- class ApiKeyInfo(NamedTuple):
12
+ Scope = Literal["ADMINISTER", "PREDICT"]
13
+ """
14
+ The scopes of an API key.
15
+
16
+ - `ADMINISTER`: Can do anything, including creating and deleting organizations, models, and API keys.
17
+ - `PREDICT`: Can only call model.predict and perform CRUD operations on predictions.
18
+ """
19
+
20
+
21
+ class ApiKeyInfo:
9
22
  """
10
- Named tuple containing information about an API key
23
+ Information about an API key
24
+
25
+ Note:
26
+ The value of the API key is only available at creation time.
11
27
 
12
28
  Attributes:
13
29
  name: Unique name of the API key
14
30
  created_at: When the API key was created
31
+ scopes: The scopes of the API key
15
32
  """
16
33
 
17
34
  name: str
18
35
  created_at: datetime
36
+ scopes: set[Scope]
37
+
38
+ def __init__(self, name: str, created_at: datetime, scopes: set[Scope]):
39
+ self.name = name
40
+ self.created_at = created_at
41
+ self.scopes = scopes
42
+
43
+ def __repr__(self) -> str:
44
+ return "ApiKey({ " + f"name: '{self.name}', scopes: <{'|'.join(self.scopes)}>" + "})"
19
45
 
20
46
 
21
47
  class OrcaCredentials:
@@ -24,11 +50,36 @@ class OrcaCredentials:
24
50
  """
25
51
 
26
52
  @staticmethod
27
- def get_api_url() -> str:
53
+ def is_authenticated() -> bool:
28
54
  """
29
- Get the Orca API base URL that is currently being used
55
+ Check if you are authenticated to interact with the Orca API
56
+
57
+ Returns:
58
+ True if you are authenticated, False otherwise
30
59
  """
31
- return get_base_url()
60
+ client = OrcaClient._resolve_client()
61
+ try:
62
+ return client.GET("/auth")
63
+ except ValueError as e:
64
+ if "Invalid API key" in str(e):
65
+ return False
66
+ raise e
67
+
68
+ @staticmethod
69
+ def is_healthy() -> bool:
70
+ """
71
+ Check whether the API is healthy
72
+
73
+ Returns:
74
+ True if the API is healthy, False otherwise
75
+ """
76
+ client = OrcaClient._resolve_client()
77
+ try:
78
+ # we don't want a retry transport here, so we use httpx directly
79
+ httpx.get(f"{client.base_url}/check/healthy")
80
+ except Exception:
81
+ return False
82
+ return True
32
83
 
33
84
  @staticmethod
34
85
  def list_api_keys() -> list[ApiKeyInfo]:
@@ -38,23 +89,52 @@ class OrcaCredentials:
38
89
  Returns:
39
90
  A list of named tuples, with the name and creation date time of the API key
40
91
  """
41
- return [ApiKeyInfo(name=api_key.name, created_at=api_key.created_at) for api_key in list_api_keys()]
92
+ client = OrcaClient._resolve_client()
93
+ return [
94
+ ApiKeyInfo(
95
+ name=api_key["name"],
96
+ created_at=datetime.fromisoformat(api_key["created_at"]),
97
+ scopes=set(api_key["scope"]),
98
+ )
99
+ for api_key in client.GET("/auth/api_key")
100
+ ]
42
101
 
43
102
  @staticmethod
44
- def is_authenticated() -> bool:
103
+ def create_api_key(name: str, scopes: set[Scope] = {"ADMINISTER"}) -> str:
45
104
  """
46
- Check if you are authenticated to interact with the Orca API
105
+ Create a new API key with the given name and scopes
106
+
107
+ Params:
108
+ name: The name of the API key
109
+ scopes: The scopes of the API key
47
110
 
48
111
  Returns:
49
- True if you are authenticated, False otherwise
112
+ The secret value of the API key. Make sure to save this value as it will not be shown again.
50
113
  """
51
- try:
52
- return check_authentication()
53
- except ValueError as e:
54
- if "Invalid API key" in str(e):
55
- return False
56
- raise e
114
+ client = OrcaClient._resolve_client()
115
+ res = client.POST(
116
+ "/auth/api_key",
117
+ json={"name": name, "scope": list(scopes)},
118
+ )
119
+ return res["api_key"]
57
120
 
121
+ @staticmethod
122
+ def revoke_api_key(name: str) -> None:
123
+ """
124
+ Delete an API key
125
+
126
+ Params:
127
+ name: The name of the API key to delete
128
+
129
+ Raises:
130
+ ValueError: if the API key is not found
131
+ """
132
+ client = OrcaClient._resolve_client()
133
+ client.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
134
+
135
+ # TODO: remove deprecated methods after 2026-01-01
136
+
137
+ @deprecated("Use `OrcaClient.api_key` instead")
58
138
  @staticmethod
59
139
  def set_api_key(api_key: str, check_validity: bool = True):
60
140
  """
@@ -70,6 +150,68 @@ class OrcaCredentials:
70
150
  Raises:
71
151
  ValueError: if the API key is invalid and `check_validity` is True
72
152
  """
73
- set_headers(get_headers() | {"Api-Key": api_key})
153
+ sync_client = OrcaClient._resolve_client()
154
+ sync_client.api_key = api_key
155
+ if check_validity:
156
+ sync_client.GET("/auth")
157
+
158
+ async_client = OrcaAsyncClient._resolve_client()
159
+ async_client.api_key = api_key
160
+
161
+ @deprecated("Use `OrcaClient.base_url` instead")
162
+ @staticmethod
163
+ def get_api_url() -> str:
164
+ """
165
+ Get the base URL of the Orca API that is currently being used
166
+ """
167
+ client = OrcaClient._resolve_client()
168
+ return str(client.base_url)
169
+
170
+ @deprecated("Use `OrcaClient.base_url` instead")
171
+ @staticmethod
172
+ def set_api_url(url: str, check_validity: bool = True):
173
+ """
174
+ Set the base URL for the Orca API
175
+
176
+ Args:
177
+ url: The base URL to set
178
+ check_validity: Whether to check if there is an API running at the given base URL
179
+
180
+ Raises:
181
+ ValueError: if there is no healthy API running at the given base URL and `check_validity` is True
182
+ """
183
+ # check if the base url is reachable before setting it
184
+ if check_validity:
185
+ try:
186
+ httpx.get(url, timeout=1)
187
+ except ConnectError as e:
188
+ raise ValueError(f"No API found at {url}") from e
189
+
190
+ sync_client = OrcaClient._resolve_client()
191
+ sync_client.base_url = url
192
+
193
+ async_client = OrcaAsyncClient._resolve_client()
194
+ async_client.base_url = url
195
+
196
+ # check if the api passes the health check
74
197
  if check_validity:
75
- check_authentication()
198
+ OrcaCredentials.is_healthy()
199
+
200
+ @deprecated("Use `OrcaClient.headers` instead")
201
+ @staticmethod
202
+ def set_api_headers(headers: dict[str, str]):
203
+ """
204
+ Add or override default HTTP headers for all Orca API requests.
205
+
206
+ Params:
207
+ headers: Mapping of header names to their string values
208
+
209
+ Notes:
210
+ New keys are merged into the existing headers, this will overwrite headers with the
211
+ same name, but leave other headers untouched.
212
+ """
213
+ sync_client = OrcaClient._resolve_client()
214
+ sync_client.headers.update(Headers(headers))
215
+
216
+ async_client = OrcaAsyncClient._resolve_client()
217
+ async_client.headers.update(Headers(headers))
@@ -2,36 +2,120 @@ from uuid import uuid4
2
2
 
3
3
  import pytest
4
4
 
5
+ from .client import OrcaClient
5
6
  from .credentials import OrcaCredentials
6
7
 
7
8
 
9
+ def test_is_authenticated():
10
+ assert OrcaCredentials.is_authenticated()
11
+
12
+
13
+ def test_is_authenticated_false(unauthenticated_client):
14
+ with unauthenticated_client.use():
15
+ assert not OrcaCredentials.is_authenticated()
16
+
17
+
18
+ def test_is_healthy():
19
+ assert OrcaCredentials.is_healthy()
20
+
21
+
22
+ def test_is_healthy_false(api_key):
23
+ with OrcaClient(api_key=api_key, base_url="http://localhost:1582").use():
24
+ assert not OrcaCredentials.is_healthy()
25
+
26
+
8
27
  def test_list_api_keys():
9
28
  api_keys = OrcaCredentials.list_api_keys()
10
29
  assert len(api_keys) >= 1
11
30
  assert "orca_sdk_test" in [api_key.name for api_key in api_keys]
12
31
 
13
32
 
14
- def test_list_api_keys_unauthenticated(unauthenticated):
15
- with pytest.raises(ValueError, match="Invalid API key"):
16
- OrcaCredentials.list_api_keys()
33
+ def test_list_api_keys_unauthenticated(unauthenticated_client):
34
+ with unauthenticated_client.use():
35
+ with pytest.raises(ValueError, match="Invalid API key"):
36
+ OrcaCredentials.list_api_keys()
17
37
 
18
38
 
19
- def test_is_authenticated():
20
- assert OrcaCredentials.is_authenticated()
39
+ def test_manage_api_key():
40
+ api_key_name = f"orca_sdk_test_{uuid4().hex[:8]}"
41
+ api_key = OrcaCredentials.create_api_key(api_key_name)
42
+ assert api_key is not None
43
+ assert len(api_key) > 0
44
+ assert api_key_name in [aki.name for aki in OrcaCredentials.list_api_keys()]
45
+ OrcaCredentials.revoke_api_key(api_key_name)
46
+ assert api_key_name not in [aki.name for aki in OrcaCredentials.list_api_keys()]
21
47
 
22
48
 
23
- def test_is_authenticated_false(unauthenticated):
24
- assert not OrcaCredentials.is_authenticated()
49
+ def test_create_api_key_unauthenticated(unauthenticated_client):
50
+ with unauthenticated_client.use():
51
+ with pytest.raises(ValueError, match="Invalid API key"):
52
+ OrcaCredentials.create_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
25
53
 
26
54
 
27
- def test_set_api_key(api_key, unauthenticated):
28
- assert not OrcaCredentials.is_authenticated()
29
- OrcaCredentials.set_api_key(api_key)
30
- assert OrcaCredentials.is_authenticated()
55
+ def test_create_api_key_unauthorized(predict_only_client):
56
+ with predict_only_client.use():
57
+ with pytest.raises(PermissionError):
58
+ OrcaCredentials.create_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
31
59
 
32
60
 
33
- def test_set_invalid_api_key(api_key):
34
- assert OrcaCredentials.is_authenticated()
35
- with pytest.raises(ValueError, match="Invalid API key"):
36
- OrcaCredentials.set_api_key(str(uuid4()))
37
- assert not OrcaCredentials.is_authenticated()
61
+ def test_revoke_api_key_unauthenticated(unauthenticated_client):
62
+ with unauthenticated_client.use():
63
+ with pytest.raises(ValueError, match="Invalid API key"):
64
+ OrcaCredentials.revoke_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
65
+
66
+
67
+ def test_revoke_api_key_unauthorized(predict_only_client):
68
+ with predict_only_client.use():
69
+ with pytest.raises(PermissionError):
70
+ OrcaCredentials.revoke_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
71
+
72
+
73
+ def test_create_api_key_already_exists():
74
+ with pytest.raises(ValueError, match="API key with this name already exists"):
75
+ OrcaCredentials.create_api_key("orca_sdk_test")
76
+
77
+
78
+ def test_set_api_key(api_key):
79
+ client = OrcaClient(api_key=str(uuid4()))
80
+ with client.use():
81
+ assert not OrcaCredentials.is_authenticated()
82
+ client.api_key = api_key
83
+ assert client.api_key == api_key
84
+ assert OrcaCredentials.is_authenticated()
85
+
86
+
87
+ def test_set_base_url(api_key):
88
+ client = OrcaClient(base_url="http://localhost:1582")
89
+ assert client.base_url == "http://localhost:1582"
90
+ client.base_url = "http://localhost:1583"
91
+ assert client.base_url == "http://localhost:1583"
92
+
93
+
94
+ # deprecated methods:
95
+
96
+
97
+ def test_deprecated_set_api_key(api_key):
98
+ with OrcaClient(api_key=str(uuid4())).use():
99
+ assert not OrcaCredentials.is_authenticated()
100
+ OrcaCredentials.set_api_key(api_key)
101
+ assert OrcaCredentials.is_authenticated()
102
+
103
+
104
+ def test_deprecated_set_invalid_api_key(api_key):
105
+ with OrcaClient(api_key=api_key).use():
106
+ assert OrcaCredentials.is_authenticated()
107
+ with pytest.raises(ValueError, match="Invalid API key"):
108
+ OrcaCredentials.set_api_key(str(uuid4()))
109
+ assert not OrcaCredentials.is_authenticated()
110
+
111
+
112
+ def test_deprecated_set_api_url(api_key):
113
+ with OrcaClient(api_key=api_key).use():
114
+ OrcaCredentials.set_api_url("http://api.orcadb.ai")
115
+ assert str(OrcaClient._resolve_client().base_url) == "http://api.orcadb.ai"
116
+
117
+
118
+ def test_deprecated_set_invalid_api_url(api_key):
119
+ with OrcaClient(api_key=api_key).use():
120
+ with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
121
+ OrcaCredentials.set_api_url("http://localhost:1582")