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

orca_sdk/embedding_model.py

@@ -12,9 +12,9 @@ from .client import (
     EmbedRequest,
     FinetunedEmbeddingModelMetadata,
     FinetuneEmbeddingModelRequest,
+    OrcaClient,
     PretrainedEmbeddingModelMetadata,
     PretrainedEmbeddingModelName,
-    orca_api,
 )
 from .datasource import Datasource
 from .job import Job, Status
@@ -82,15 +82,16 @@ class EmbeddingModelBase(ABC):
             "max_seq_length": max_seq_length,
             "prompt": prompt,
         }
+        client = OrcaClient._resolve_client()
         if isinstance(self, PretrainedEmbeddingModel):
-            embeddings = orca_api.POST(
+            embeddings = client.POST(
                 "/gpu/pretrained_embedding_model/{model_name}/embedding",
                 params={"model_name": cast(PretrainedEmbeddingModelName, self.name)},
                 json=payload,
                 timeout=30,  # may be slow in case of cold start
             )
         elif isinstance(self, FinetunedEmbeddingModel):
-            embeddings = orca_api.POST(
+            embeddings = client.POST(
                 "/gpu/finetuned_embedding_model/{name_or_id}/embedding",
                 params={"name_or_id": self.id},
                 json=payload,
@@ -202,14 +203,15 @@ class EmbeddingModelBase(ABC):
             "batch_size": batch_size,
             "weigh_memories": weigh_memories,
         }
+        client = OrcaClient._resolve_client()
         if isinstance(self, PretrainedEmbeddingModel):
-            response = orca_api.POST(
+            response = client.POST(
                 "/pretrained_embedding_model/{model_name}/evaluation",
                 params={"model_name": self.name},
                 json=payload,
             )
         elif isinstance(self, FinetunedEmbeddingModel):
-            response = orca_api.POST(
+            response = client.POST(
                 "/finetuned_embedding_model/{name_or_id}/evaluation",
                 params={"name_or_id": self.id},
                 json=payload,
@@ -218,13 +220,14 @@ class EmbeddingModelBase(ABC):
             raise ValueError("Invalid embedding model")
 
         def get_result(task_id: str) -> ClassificationMetrics | RegressionMetrics:
+            client = OrcaClient._resolve_client()
             if isinstance(self, PretrainedEmbeddingModel):
-                res = orca_api.GET(
+                res = client.GET(
                     "/pretrained_embedding_model/{model_name}/evaluation/{task_id}",
                     params={"model_name": self.name, "task_id": task_id},
                 )["result"]
             elif isinstance(self, FinetunedEmbeddingModel):
-                res = orca_api.GET(
+                res = client.GET(
                     "/finetuned_embedding_model/{name_or_id}/evaluation/{task_id}",
                     params={"name_or_id": self.id, "task_id": task_id},
                 )["result"]
@@ -401,7 +404,8 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         Returns:
             A list of all pretrained embedding models available in the OrcaCloud
         """
-        return [cls(metadata) for metadata in orca_api.GET("/pretrained_embedding_model")]
+        client = OrcaClient._resolve_client()
+        return [cls(metadata) for metadata in client.GET("/pretrained_embedding_model")]
 
     _instances: dict[str, PretrainedEmbeddingModel] = {}
 
@@ -410,7 +414,8 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         # for internal use only, do not document - we want people to use dot notation to get the model
         cache_key = str(name)
         if cache_key not in cls._instances:
-            metadata = orca_api.GET(
+            client = OrcaClient._resolve_client()
+            metadata = client.GET(
                 "/pretrained_embedding_model/{model_name}",
                 params={"model_name": name},
             )
@@ -555,7 +560,8 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         if eval_datasource is not None:
             payload["eval_datasource_name_or_id"] = eval_datasource.id
 
-        res = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        res = client.POST(
             "/finetuned_embedding_model",
             json=payload,
         )
@@ -630,7 +636,8 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
         Returns:
             A list of all finetuned embedding model handles in the OrcaCloud
         """
-        return [cls(metadata) for metadata in orca_api.GET("/finetuned_embedding_model")]
+        client = OrcaClient._resolve_client()
+        return [cls(metadata) for metadata in client.GET("/finetuned_embedding_model")]
 
     @classmethod
     def open(cls, name: str) -> FinetunedEmbeddingModel:
@@ -646,7 +653,8 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
         Raises:
             LookupError: If the finetuned embedding model does not exist
         """
-        metadata = orca_api.GET(
+        client = OrcaClient._resolve_client()
+        metadata = client.GET(
             "/finetuned_embedding_model/{name_or_id}",
             params={"name_or_id": name},
         )
@@ -681,7 +689,8 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
             LookupError: If the finetuned embedding model does not exist and `if_not_exists` is `"error"`
         """
         try:
-            orca_api.DELETE(
+            client = OrcaClient._resolve_client()
+            client.DELETE(
                 "/finetuned_embedding_model/{name_or_id}",
                 params={"name_or_id": name_or_id},
             )

orca_sdk/embedding_model_test.py

@@ -25,9 +25,10 @@ def test_open_pretrained_model():
     assert model is PretrainedEmbeddingModel.GTE_BASE
 
 
-def test_open_pretrained_model_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline")
+def test_open_pretrained_model_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline")
 
 
 def test_open_pretrained_model_not_found():
@@ -52,9 +53,10 @@ def test_embed_text():
     assert isinstance(embedding[0], float)
 
 
-def test_embed_text_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline", max_seq_length=32)
+def test_embed_text_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline", max_seq_length=32)
 
 
 def test_evaluate_pretrained_model(datasource: Datasource):
@@ -108,9 +110,10 @@ def test_finetune_model_already_exists_return(datasource: Datasource, finetuned_
     assert new_model._status == Status.COMPLETED
 
 
-def test_finetune_model_unauthenticated(unauthenticated, datasource: Datasource):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model_unauthenticated", datasource)
+def test_finetune_model_unauthenticated(unauthenticated_client, datasource: Datasource):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model_unauthenticated", datasource)
 
 
 def test_use_finetuned_model_in_memoryset(datasource: Datasource, finetuned_model: FinetunedEmbeddingModel):
@@ -150,13 +153,15 @@ def test_all_finetuned_models(finetuned_model: FinetunedEmbeddingModel):
     assert any(model.name == finetuned_model.name for model in models)
 
 
-def test_all_finetuned_models_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        FinetunedEmbeddingModel.all()
+def test_all_finetuned_models_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            FinetunedEmbeddingModel.all()
 
 
-def test_all_finetuned_models_unauthorized(unauthorized, finetuned_model: FinetunedEmbeddingModel):
-    assert finetuned_model not in FinetunedEmbeddingModel.all()
+def test_all_finetuned_models_unauthorized(unauthorized_client, finetuned_model: FinetunedEmbeddingModel):
+    with unauthorized_client.use():
+        assert finetuned_model not in FinetunedEmbeddingModel.all()
 
 
 def test_drop_finetuned_model(datasource: Datasource):
@@ -167,9 +172,10 @@ def test_drop_finetuned_model(datasource: Datasource):
         FinetunedEmbeddingModel.open("finetuned_model_to_delete")
 
 
-def test_drop_finetuned_model_unauthenticated(unauthenticated, datasource: Datasource):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource)
+def test_drop_finetuned_model_unauthenticated(unauthenticated_client, datasource: Datasource):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource)
 
 
 def test_drop_finetuned_model_not_found():
@@ -179,9 +185,10 @@ def test_drop_finetuned_model_not_found():
     FinetunedEmbeddingModel.drop(str(uuid4()), if_not_exists="ignore")
 
 
-def test_drop_finetuned_model_unauthorized(unauthorized, finetuned_model: FinetunedEmbeddingModel):
-    with pytest.raises(LookupError):
-        FinetunedEmbeddingModel.drop(finetuned_model.id)
+def test_drop_finetuned_model_unauthorized(unauthorized_client, finetuned_model: FinetunedEmbeddingModel):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            FinetunedEmbeddingModel.drop(finetuned_model.id)
 
 
 def test_supports_instructions():

orca_sdk/job.py

@@ -7,7 +7,7 @@ from typing import Callable, Generic, TypedDict, TypeVar, cast
 
 from tqdm.auto import tqdm
 
-from .client import orca_api
+from .client import OrcaClient
 
 
 class JobConfig(TypedDict):
@@ -140,7 +140,8 @@ class Job(Generic[TResult]):
         Returns:
             List of jobs matching the given filters
         """
-        paginated_tasks = orca_api.GET(
+        client = OrcaClient._resolve_client()
+        paginated_tasks = client.GET(
             "/task",
             params={
                 "status": (
@@ -186,11 +187,14 @@ class Job(Generic[TResult]):
             get_value: Optional function to customize how the value is resolved, if not provided the result will be a dict
         """
         self.id = id
-        task = orca_api.GET("/task/{task_id}", params={"task_id": id})
+        client = OrcaClient._resolve_client()
+        task = client.GET("/task/{task_id}", params={"task_id": id})
 
-        self._get_value = get_value or (
-            lambda: cast(TResult | None, orca_api.GET("/task/{task_id}", params={"task_id": id})["result"])
-        )
+        def default_get_value():
+            client = OrcaClient._resolve_client()
+            return cast(TResult | None, client.GET("/task/{task_id}", params={"task_id": id})["result"])
+
+        self._get_value = get_value or default_get_value
         self.type = task["type"]
         self.status = Status(task["status"])
         self.steps_total = task["steps_total"]
@@ -222,7 +226,8 @@ class Job(Generic[TResult]):
             return
         self.refreshed_at = current_time
 
-        status_info = orca_api.GET("/task/{task_id}/status", params={"task_id": self.id})
+        client = OrcaClient._resolve_client()
+        status_info = client.GET("/task/{task_id}/status", params={"task_id": self.id})
         self.status = Status(status_info["status"])
         if status_info["steps_total"] is not None:
             self.steps_total = status_info["steps_total"]
@@ -333,5 +338,6 @@ def abort(self, show_progress: bool = False, refresh_interval: int = 1, max_wait
         refresh_interval: Polling interval in seconds while waiting for the job to abort
         max_wait: Maximum time to wait for the job to abort in seconds
     """
-    orca_api.DELETE("/task/{task_id}/abort", params={"task_id": self.id})
+    client = OrcaClient._resolve_client()
+    client.DELETE("/task/{task_id}/abort", params={"task_id": self.id})
     self.wait(show_progress, refresh_interval, max_wait)