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

orca_sdk/client.py

@@ -60,10 +60,12 @@ class ActionRecommendation(TypedDict):
 class AddMemorySuggestion(TypedDict):
     value: str
     label_name: str
+    similarity: NotRequired[float | None]
 
 
 class AliveResponse(TypedDict):
     ok: bool
+    checks: dict[str, bool]
 
 
 class ApiKeyMetadata(TypedDict):
@@ -290,7 +292,7 @@ class JobStatusInfo(TypedDict):
 class LabelClassMetrics(TypedDict):
     label: int | None
     label_name: NotRequired[str | None]
-    average_lookup_score: float
+    average_lookup_score: float | None
     memory_count: int
 
 
@@ -344,7 +346,7 @@ class MemoryMetrics(TypedDict):
     cluster: NotRequired[int]
     embedding_2d: NotRequired[list]
     anomaly_score: NotRequired[float]
-    neighbor_label_logits: NotRequired[list[float]]
+    neighbor_label_logits: NotRequired[list[float] | None]
     neighbor_predicted_label: NotRequired[int | None]
     neighbor_predicted_label_ambiguity: NotRequired[float]
     neighbor_predicted_label_confidence: NotRequired[float]
@@ -1155,6 +1157,9 @@ class FieldValidationError(TypedDict):
 
 class AddMemoryRecommendations(TypedDict):
     memories: list[AddMemorySuggestion]
+    attempts_used: NotRequired[int]
+    partial: NotRequired[bool]
+    rejection_counts: NotRequired[dict[str, int]]
 
 
 class BootstrapClassificationModelRequest(TypedDict):
@@ -1270,10 +1275,19 @@ class CreateClassificationModelRequest(TypedDict):
     num_classes: NotRequired[int | None]
 
 
-class CreateMemorysetRequest(TypedDict):
+class CreateMemorysetFromDatasourceRequest(TypedDict):
     name: str
     description: NotRequired[str | None]
     notes: NotRequired[str | None]
+    pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
+    finetuned_embedding_model_name_or_id: NotRequired[str | None]
+    max_seq_length_override: NotRequired[int | None]
+    label_names: NotRequired[list[str] | None]
+    index_type: NotRequired[Literal["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "DISKANN"]]
+    index_params: NotRequired[dict[str, int | float | str]]
+    prompt: NotRequired[str]
+    hidden: NotRequired[bool]
+    memory_type: NotRequired[MemoryType | None]
     datasource_name_or_id: str
     datasource_label_column: NotRequired[str | None]
     datasource_score_column: NotRequired[str | None]
@@ -1281,6 +1295,14 @@ class CreateMemorysetRequest(TypedDict):
     datasource_source_id_column: NotRequired[str | None]
     datasource_partition_id_column: NotRequired[str | None]
     remove_duplicates: NotRequired[bool]
+    batch_size: NotRequired[int]
+    subsample: NotRequired[int | float | None]
+
+
+class CreateMemorysetRequest(TypedDict):
+    name: str
+    description: NotRequired[str | None]
+    notes: NotRequired[str | None]
     pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
     finetuned_embedding_model_name_or_id: NotRequired[str | None]
     max_seq_length_override: NotRequired[int | None]
@@ -1289,9 +1311,7 @@ class CreateMemorysetRequest(TypedDict):
     index_params: NotRequired[dict[str, int | float | str]]
     prompt: NotRequired[str]
     hidden: NotRequired[bool]
-    batch_size: NotRequired[int]
-    subsample: NotRequired[int | float | None]
-    memory_type: NotRequired[MemoryType]
+    memory_type: NotRequired[MemoryType | None]
 
 
 class CreateRegressionModelRequest(TypedDict):
@@ -1651,8 +1671,8 @@ class MemorysetMetadata(TypedDict):
     created_at: str
     updated_at: str
     memories_updated_at: str
-    insertion_job_id: str
-    insertion_status: JobStatus
+    insertion_job_id: str | None
+    insertion_status: JobStatus | None
     metrics: MemorysetMetrics
     memory_type: MemoryType
     label_names: list[str] | None
@@ -1662,7 +1682,7 @@ class MemorysetMetadata(TypedDict):
     document_prompt_override: str | None
     query_prompt_override: str | None
     hidden: bool
-    insertion_task_id: str
+    insertion_task_id: str | None
 
 
 class PaginatedWorkerInfo(TypedDict):
@@ -1733,6 +1753,22 @@ class OrcaClient(Client):
     ) -> AliveResponse:
         pass
 
+    @overload
+    def GET(
+        self,
+        path: Literal["/gpu/check/alive"],
+        *,
+        params: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> AliveResponse:
+        pass
+
     @overload
     def GET(
         self,
@@ -1752,7 +1788,7 @@ class OrcaClient(Client):
     @overload
     def GET(
         self,
-        path: Literal["/gpu/check/healthy"],
+        path: Literal["/check/healthy"],
         *,
         params: None = None,
         parse_as: Literal["json"] = "json",
@@ -1768,7 +1804,7 @@ class OrcaClient(Client):
     @overload
     def GET(
         self,
-        path: Literal["/check/healthy"],
+        path: Literal["/gpu/check/healthy"],
         *,
         params: None = None,
         parse_as: Literal["json"] = "json",
@@ -2893,6 +2929,26 @@ class OrcaClient(Client):
         path: Literal["/memoryset"],
         *,
         params: None = None,
+        json: CreateMemorysetFromDatasourceRequest,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> MemorysetMetadata:
+        pass
+
+    @overload
+    def POST(
+        self,
+        path: Literal["/memoryset/empty"],
+        *,
+        params: None = None,
         json: CreateMemorysetRequest,
         data: None = None,
         files: None = None,

orca_sdk/embedding_model.py

@@ -26,14 +26,23 @@ if TYPE_CHECKING:
 class EmbeddingModelBase(ABC):
     embedding_dim: int
     max_seq_length: int
+    num_params: int
     uses_context: bool
     supports_instructions: bool
 
     def __init__(
-        self, *, name: str, embedding_dim: int, max_seq_length: int, uses_context: bool, supports_instructions: bool
+        self,
+        *,
+        name: str,
+        embedding_dim: int,
+        max_seq_length: int,
+        num_params: int,
+        uses_context: bool,
+        supports_instructions: bool,
     ):
         self.embedding_dim = embedding_dim
         self.max_seq_length = max_seq_length
+        self.num_params = num_params
         self.uses_context = uses_context
         self.supports_instructions = supports_instructions
 
@@ -359,6 +368,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         name: Name of the pretrained embedding model
         embedding_dim: Dimension of the embeddings that are generated by the model
         max_seq_length: Maximum input length (in tokens not characters) that this model can process. Inputs that are longer will be truncated during the embedding process
+        num_params: Number of parameters in the model
         uses_context: Whether the pretrained embedding model uses context
         supports_instructions: Whether this model supports instruction-following
     """
@@ -383,6 +393,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
             name=metadata["name"],
             embedding_dim=metadata["embedding_dim"],
             max_seq_length=metadata["max_seq_length"],
+            num_params=metadata["num_params"],
             uses_context=metadata["uses_context"],
             supports_instructions=(
                 bool(metadata["supports_instructions"]) if "supports_instructions" in metadata else False
@@ -393,7 +404,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         return isinstance(other, PretrainedEmbeddingModel) and self.name == other.name
 
     def __repr__(self) -> str:
-        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}}})"
+        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {self.num_params/1000000:.0f}M}})"
 
     @classmethod
     def all(cls) -> list[PretrainedEmbeddingModel]:
@@ -541,7 +552,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         elif exists and if_exists == "open":
             existing = FinetunedEmbeddingModel.open(name)
 
-            if existing.base_model_name != self.name:
+            if existing.base_model.name != self.name:
                 raise ValueError(f"Finetuned embedding model '{name}' already exists, but with different base model")
 
             return existing
@@ -596,7 +607,7 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
     name: str
     created_at: datetime
     updated_at: datetime
-    base_model_name: PretrainedEmbeddingModelName
+    base_model: PretrainedEmbeddingModel
     _status: Status
 
     def __init__(self, metadata: FinetunedEmbeddingModelMetadata):
@@ -605,13 +616,14 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
         self.name = metadata["name"]
         self.created_at = datetime.fromisoformat(metadata["created_at"])
         self.updated_at = datetime.fromisoformat(metadata["updated_at"])
-        self.base_model_name = metadata["base_model"]
+        self.base_model = PretrainedEmbeddingModel._get(metadata["base_model"])
         self._status = Status(metadata["finetuning_status"])
 
         super().__init__(
             name=metadata["name"],
             embedding_dim=metadata["embedding_dim"],
             max_seq_length=metadata["max_seq_length"],
+            num_params=self.base_model.num_params,
             uses_context=metadata["uses_context"],
             supports_instructions=self.base_model.supports_instructions,
         )
@@ -625,15 +637,10 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
             f"    name: {self.name},\n"
             f"    embedding_dim: {self.embedding_dim},\n"
             f"    max_seq_length: {self.max_seq_length},\n"
-            f"    base_model: PretrainedEmbeddingModel.{self.base_model_name}\n"
+            f"    base_model: PretrainedEmbeddingModel.{self.base_model.name}\n"
             "})"
         )
 
-    @property
-    def base_model(self) -> PretrainedEmbeddingModel:
-        """Pretrained model the finetuned embedding model was based on"""
-        return PretrainedEmbeddingModel._get(self.base_model_name)
-
     @classmethod
     def all(cls) -> list[FinetunedEmbeddingModel]:
         """
@@ -700,6 +707,6 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
                 "/finetuned_embedding_model/{name_or_id}",
                 params={"name_or_id": name_or_id},
             )
-        except (LookupError, RuntimeError):
+        except LookupError:
             if if_not_exists == "error":
                 raise

orca_sdk/embedding_model_test.py

@@ -119,7 +119,7 @@ def test_finetune_model_unauthenticated(unauthenticated_client, datasource: Data
 def test_use_finetuned_model_in_memoryset(datasource: Datasource, finetuned_model: FinetunedEmbeddingModel):
     memoryset = LabeledMemoryset.create(
         "test_memoryset_finetuned_model",
-        datasource,
+        datasource=datasource,
         embedding_model=finetuned_model,
     )
     assert memoryset is not None