orca-sdk 0.0.96__py3-none-any.whl → 0.0.97__py3-none-any.whl

orca_sdk/embedding_model_test.py

@@ -179,3 +179,30 @@ def test_drop_finetuned_model_not_found():
 def test_drop_finetuned_model_unauthorized(unauthorized, finetuned_model: FinetunedEmbeddingModel):
     with pytest.raises(LookupError):
         FinetunedEmbeddingModel.drop(finetuned_model.id)
+
+
+def test_default_instruction_with_memoryset_creation():
+    """Test that embedding models work correctly with instruction support."""
+    # Test with an instruction-supporting model
+    model = PretrainedEmbeddingModel.open("E5_LARGE")
+
+    # Verify the model properties
+    assert model.supports_instructions
+
+    # Test that prompt parameter is passed through correctly (orcalib handles the default)
+    embeddings_explicit_instruction = model.embed("Hello world", prompt="Represent this sentence for retrieval:")
+    embeddings_no_instruction = model.embed("Hello world")
+
+    # These should be different since one uses a prompt and the other doesn't
+    assert embeddings_explicit_instruction != embeddings_no_instruction
+
+
+def test_default_instruction_error_cases():
+    """Test basic embedding model functionality."""
+    # Test that model opens correctly and has instruction support information
+    model = PretrainedEmbeddingModel.open("GTE_BASE")
+    assert not model.supports_instructions
+
+    # Test instruction-supporting model
+    instruction_model = PretrainedEmbeddingModel.open("E5_LARGE")
+    assert instruction_model.supports_instructions

orca_sdk/job.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import time
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Callable, Generic, TypedDict, TypeVar, cast
+from typing import Any, Callable, Generic, TypedDict, TypeVar, cast
 
 from tqdm.auto import tqdm
 
@@ -117,7 +117,7 @@ class Job(Generic[TResult]):
         cls,
         status: Status | list[Status] | None = None,
         type: str | list[str] | None = None,
-        limit: int | None = None,
+        limit: int = 100,
         offset: int = 0,
         start: datetime | None = None,
         end: datetime | None = None,
@@ -136,7 +136,7 @@ class Job(Generic[TResult]):
         Returns:
             List of jobs matching the given filters
         """
-        tasks = list_tasks(
+        paginated_tasks = list_tasks(
             status=(
                 [TaskStatus(s.value) for s in status]
                 if isinstance(status, list)
@@ -167,7 +167,7 @@ class Job(Generic[TResult]):
                     obj,
                 )[-1]
             )(t)
-            for t in tasks
+            for t in paginated_tasks.items
         ]
 
     def __init__(self, id: str, get_value: Callable[[], TResult | None] | None = None):
@@ -316,14 +316,15 @@ class Job(Generic[TResult]):
         assert self.value is not None
         return self.value
 
-    def abort(self, show_progress: bool = False, refresh_interval: int = 1, max_wait: int = 20) -> None:
-        """
-        Abort the job
 
-        Params:
-            show_progress: Optionally show a progress bar while waiting for the job to abort
-            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
-        """
-        abort_task(self.id)
-        self.wait(show_progress, refresh_interval, max_wait)
+def abort(self, show_progress: bool = False, refresh_interval: int = 1, max_wait: int = 20) -> None:
+    """
+    Abort the job
+
+    Params:
+        show_progress: Optionally show a progress bar while waiting for the job to abort
+        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
+    """
+    abort_task(self.id)
+    self.wait(show_progress, refresh_interval, max_wait)

orca_sdk/job_test.py

@@ -46,3 +46,37 @@ def test_job_refresh(classification_model: ClassificationModel, datasource: Data
     # calling refresh() should immediately refresh the job
     job.refresh()
     assert job.refreshed_at > last_refreshed_at
+
+
+def test_job_query_pagination(classification_model: ClassificationModel, datasource: Datasource):
+    """Test pagination with Job.query() method"""
+    # Create multiple jobs to test pagination
+    jobs_created = []
+    for i in range(3):
+        job = classification_model.evaluate(datasource, background=True)
+        jobs_created.append(job.id)
+
+    # Test basic pagination with limit
+    jobs_page1 = Job.query(type="EVALUATE_MODEL", limit=2)
+    assert len(jobs_page1) == 2
+
+    # Test pagination with offset
+    jobs_page2 = Job.query(type="EVALUATE_MODEL", limit=2, offset=1)
+    assert len(jobs_page2) == 2
+
+    # Verify different pages contain different jobs (allowing for some overlap due to timing)
+    page1_ids = {job.id for job in jobs_page1}
+    page2_ids = {job.id for job in jobs_page2}
+
+    # At least one job should be different between pages
+    assert len(page1_ids.symmetric_difference(page2_ids)) > 0
+
+    # Test filtering by status
+    jobs_by_status = Job.query(status=Status.PROCESSING, limit=10)
+    for job in jobs_by_status:
+        assert job.status == Status.PROCESSING
+
+    # Test filtering by multiple statuses
+    multi_status_jobs = Job.query(status=[Status.PROCESSING, Status.COMPLETED], limit=10)
+    for job in multi_status_jobs:
+        assert job.status in [Status.PROCESSING, Status.COMPLETED]

orca_sdk/memoryset.py

@@ -73,7 +73,10 @@ from ._generated_api_client.models import (
     PretrainedEmbeddingModelName,
 )
 from ._generated_api_client.models import ScoredMemory as ScoredMemoryResponse
-from ._generated_api_client.models import ScoredMemoryInsert, ScoredMemoryInsertMetadata
+from ._generated_api_client.models import (
+    ScoredMemoryInsert,
+    ScoredMemoryInsertMetadata,
+)
 from ._generated_api_client.models import (
     ScoredMemoryLookup as ScoredMemoryLookupResponse,
 )
@@ -719,6 +722,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
+        prompt: str | None = None,
         remove_duplicates: bool = True,
         index_type: IndexType = "FLAT",
         index_params: dict[str, Any] = {},
@@ -742,6 +746,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
+        prompt: str | None = None,
         remove_duplicates: bool = True,
         index_type: IndexType = "FLAT",
         index_params: dict[str, Any] = {},
@@ -764,6 +769,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
+        prompt: str | None = None,
         remove_duplicates: bool = True,
         index_type: IndexType = "FLAT",
         index_params: dict[str, Any] = {},
@@ -797,6 +803,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             max_seq_length_override: Maximum sequence length of values in the memoryset, if the
                 value is longer than this it will be truncated, will default to the model's max
                 sequence length if not provided
+            prompt: Optional prompt to use when embedding documents/memories for storage
             remove_duplicates: Whether to remove duplicates from the datasource before inserting
                 into the memoryset
             index_type: Type of vector index to use for the memoryset, defaults to `"FLAT"`. Valid
@@ -849,6 +856,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                 remove_duplicates=remove_duplicates,
                 index_type=CreateMemorysetRequestIndexType[index_type],
                 index_params=CreateMemorysetRequestIndexParams.from_dict(index_params),
+                prompt=prompt,
             ),
         )
         job = Job(response.insertion_task_id, lambda: cls.open(response.id))
@@ -1351,6 +1359,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         *,
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
         max_seq_length_override: int | None = None,
+        prompt: str | None = None,
         if_exists: CreateMode = "error",
         background: Literal[True],
     ) -> Job[Self]:
@@ -1363,6 +1372,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         *,
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
         max_seq_length_override: int | None = None,
+        prompt: str | None = None,
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
     ) -> Self:
@@ -1374,6 +1384,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         *,
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
         max_seq_length_override: int | None = None,
+        prompt: str | None = None,
         if_exists: CreateMode = "error",
         background: bool = False,
     ) -> Self | Job[Self]:
@@ -1385,6 +1396,10 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             embedding_model: Optional new embedding model to use for re-embedding the memory values
                 value is longer than this it will be truncated, will default to the model's max
                 sequence length if not provided
+            document_prompt_override: Optional custom document prompt to use for the cloned memoryset.
+                If not provided, will use the source memoryset's document prompt.
+            query_prompt_override: Optional custom query prompt to use for the cloned memoryset.
+                If not provided, will use the source memoryset's query prompt.
             if_exists: What to do if a memoryset with the same name already exists, defaults to
                 `"error"`. Other option is `"open"` to open the existing memoryset.
 
@@ -1399,6 +1414,13 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             >>> new_memoryset = memoryset.clone(
             ...     "my_memoryset_finetuned", embedding_model=finetuned_embedding_model,
             ... )
+
+            >>> # Clone with custom prompts
+            >>> new_memoryset = memoryset.clone(
+            ...     "my_memoryset_with_prompts",
+            ...     document_prompt_override="Represent this document for retrieval:",
+            ...     query_prompt_override="Represent this query for retrieval:",
+            ... )
         """
         if self.exists(name):
             if if_exists == "error":
@@ -1421,6 +1443,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                     embedding_model.id if isinstance(embedding_model, FinetunedEmbeddingModel) else None
                 ),
                 max_seq_length_override=max_seq_length_override,
+                prompt=prompt,
             ),
         )
         job = Job(
@@ -1503,20 +1526,24 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             raise ValueError(f"Invalid index type: {type(index)}")
 
     @overload
-    def search(self, query: str, *, count: int = 1) -> list[MemoryLookupT]:
+    def search(self, query: str, *, count: int = 1, prompt: str | None = None) -> list[MemoryLookupT]:
         pass
 
     @overload
-    def search(self, query: list[str], *, count: int = 1) -> list[list[MemoryLookupT]]:
+    def search(self, query: list[str], *, count: int = 1, prompt: str | None = None) -> list[list[MemoryLookupT]]:
         pass
 
-    def search(self, query: str | list[str], *, count: int = 1) -> list[MemoryLookupT] | list[list[MemoryLookupT]]:
+    def search(
+        self, query: str | list[str], *, count: int = 1, prompt: str | None = None
+    ) -> list[MemoryLookupT] | list[list[MemoryLookupT]]:
         """
         Search for memories that are semantically similar to the query
 
         Params:
             query: Query to lookup memories in the memoryset, can be a single query or a list
             count: Number of memories to return for each query
+            prompt: Optional prompt for query embedding during search.
+                If not provided, the memoryset's default query prompt will be used if available.
 
         Returns:
             List of memories from the memoryset that match the query. If a single query is provided,
@@ -1531,6 +1558,13 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                 LabeledMemoryLookup({ label: <positive: 1>, value: 'I am content' }),
             ]
 
+            Search with custom query prompt for instruction-following models:
+            >>> memoryset.search("I am happy", count=2, query_prompt="Represent this query for sentiment retrieval:")
+            [
+                LabeledMemoryLookup({ label: <positive: 1>, value: 'I am happy' }),
+                LabeledMemoryLookup({ label: <positive: 1>, value: 'I am content' }),
+            ]
+
             Search for similar memories for multiple queries:
             >>> memoryset.search(["I am happy", "I am sad"], count=1)
             [
@@ -1544,7 +1578,11 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         """
         response = memoryset_lookup_gpu(
             name_or_id=self.id,
-            body=LookupRequest(query=query if isinstance(query, list) else [query], count=count),
+            body=LookupRequest(
+                query=query if isinstance(query, list) else [query],
+                count=count,
+                prompt=prompt,
+            ),
         )
         lookups = [
             [
@@ -2062,9 +2100,11 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         """
 
         if embedding_models is not None:
-            embedding_model_enums = [PretrainedEmbeddingModelName(model) for model in embedding_models]
+            embedding_model_enums: list[PretrainedEmbeddingModelName | str] | None = [
+                PretrainedEmbeddingModelName(model) for model in embedding_models
+            ]
         else:
-            embedding_model_enums = None
+            embedding_model_enums: list[PretrainedEmbeddingModelName | str] | None = None
 
         request = EmbeddingEvaluationRequest(
             value_column=value_column,

orca_sdk/regression_model.py

@@ -11,7 +11,7 @@ import numpy as np
 from datasets import Dataset
 
 from ._generated_api_client.api import (
-    create_regression_model_gpu,
+    create_regression_model,
     delete_regression_model,
     evaluate_regression_model,
     get_regression_model,
@@ -178,7 +178,7 @@ class RegressionModel:
 
                 return existing
 
-        metadata = create_regression_model_gpu(
+        metadata = create_regression_model(
             body=CreateRegressionModelRequest(
                 name=name,
                 memoryset_id=memoryset.id,

orca_sdk/telemetry.py

@@ -1,9 +1,18 @@
 from __future__ import annotations
 
 import logging
-from abc import ABC
+from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Generator, Iterable, Self, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generator,
+    Iterable,
+    Literal,
+    Self,
+    cast,
+    overload,
+)
 from uuid import UUID
 
 from orca_sdk._utils.common import UNSET
@@ -16,6 +25,12 @@ from ._generated_api_client.api import (
     record_prediction_feedback,
     update_prediction,
 )
+from ._generated_api_client.api.telemetry.generate_memory_suggestions_telemetry_prediction_prediction_id_memory_suggestions_post import (
+    sync as generate_memory_suggestions,
+)
+from ._generated_api_client.api.telemetry.get_action_recommendation_telemetry_prediction_prediction_id_action_get import (
+    sync as get_action_recommendation,
+)
 from ._generated_api_client.client import get_client
 from ._generated_api_client.errors import get_error_for_response
 from ._generated_api_client.models import (
@@ -77,7 +92,6 @@ class FeedbackCategory:
     created_at: datetime
 
     def __init__(self, category: PredictionFeedbackCategory):
-        # for internal use only, do not document
         self.id = category.id
         self.name = category.name
         self.value_type = bool if category.type == FeedbackType.BINARY else float
@@ -499,6 +513,83 @@ class ClassificationPrediction(_Prediction):
         """
         self._update(tags=tags, expected_label=expected_label)
 
+    def recommend_action(self, *, refresh: bool = False) -> tuple[str, str]:
+        """
+        Get an action recommendation for improving this prediction.
+
+        Analyzes the prediction and suggests the most effective action to improve model
+        performance, such as adding memories, detecting mislabels, removing duplicates,
+        or finetuning.
+
+        Params:
+            refresh: Force the action recommendation agent to re-run even if a recommendation already exists
+
+        Returns:
+            Tuple of (action, rationale) where:
+            - action: The recommended action ("add_memories", "detect_mislabels", "remove_duplicates", or "finetuning") that would resolve the mislabeling
+            - rationale: Explanation for why this action was recommended
+
+        Raises:
+            ValueError: If the prediction has no prediction ID
+            RuntimeError: If the lighthouse API key is not configured
+
+        Examples:
+            Get action recommendation for an incorrect prediction:
+            >>> action, rationale = prediction.recommend_action()
+            >>> print(f"Recommended action: {action}")
+            >>> print(f"Rationale: {rationale}")
+        """
+        if self.prediction_id is None:
+            raise ValueError("Cannot get action recommendation with no prediction ID")
+
+        api_recommendation = get_action_recommendation(
+            prediction_id=UUID(self.prediction_id),
+            refresh=refresh,
+        )
+        return (api_recommendation.action.value, api_recommendation.rationale)
+
+    def generate_memory_suggestions(self, *, num_memories: int = 3) -> list[dict[str, str | int]]:
+        """
+        Generate synthetic memory suggestions to improve this prediction.
+
+        Creates new example memories that are similar to the input but have clearer
+        signals for the expected label. These can be added to the memoryset to improve
+        model performance on similar inputs.
+
+        Params:
+            num_memories: Number of memory suggestions to generate (default: 3)
+
+        Returns:
+            List of dictionaries that can be directly passed to memoryset.insert().
+            Each dictionary contains:
+            - "value": The suggested memory text
+            - "label": The suggested label as an integer
+
+        Raises:
+            ValueError: If the prediction has no prediction ID
+            RuntimeError: If the lighthouse API key is not configured
+
+        Examples:
+            Generate memory suggestions for an incorrect prediction:
+            >>> suggestions = prediction.generate_memory_suggestions(num_memories=3)
+            >>> for suggestion in suggestions:
+            ...     print(f"Value: {suggestion['value']}, Label: {suggestion['label']}")
+            >>>
+            >>> # Add suggestions directly to memoryset
+            >>> model.memoryset.insert(suggestions)
+        """
+        if self.prediction_id is None:
+            raise ValueError("Cannot generate memory suggestions with no prediction ID")
+
+        api_recommendations = generate_memory_suggestions(
+            prediction_id=UUID(self.prediction_id),
+            num_memories=num_memories,
+        )
+
+        return [
+            {"value": suggestion.value, "label": int(suggestion.label)} for suggestion in api_recommendations.memories
+        ]
+
 
 class RegressionPrediction(_Prediction):
     """

{orca_sdk-0.0.96.dist-info → orca_sdk-0.0.97.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: orca_sdk
-Version: 0.0.96
+Version: 0.0.97
 Summary: SDK for interacting with Orca Services
 License: Apache-2.0
 Author: Orca DB Inc.
@@ -71,6 +71,23 @@ model = ClassificationModel("my_model", memoryset)
 
 # Make predictions
 prediction = model.predict("my input")
+
+# Get Action Recommendation
+action, rationale = prediction.recommend_action()
+print(f"Recommended action: {action}")
+print(f"Rationale: {rationale}")
+
+# Generate and add synthetic memory suggestions
+if action == "add_memories":
+    suggestions = prediction.generate_memory_suggestions(num_memories=3)
+
+    # Review suggestions
+    for suggestion in suggestions:
+        print(f"Suggested: '{suggestion['value']}' -> {suggestion['label']}")
+
+    # Add suggestions to memoryset
+    model.memoryset.insert(suggestions)
+    print(f"Added {len(suggestions)} new memories to improve model performance!")
 ```
 
 For a more detailed walkthrough, check out our [Quick Start Guide](https://docs.orcadb.ai/quickstart-sdk/).