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

orca_sdk/__init__.py

@@ -4,6 +4,7 @@ OrcaSDK is a Python library for building and using retrieval augmented models in
 
 from ._utils.common import UNSET, CreateMode, DropMode
 from .classification_model import ClassificationMetrics, ClassificationModel
+from .client import orca_api
 from .credentials import OrcaCredentials
 from .datasource import Datasource
 from .embedding_model import (
@@ -23,11 +24,7 @@ from .memoryset import (
     ScoredMemoryset,
 )
 from .regression_model import RegressionModel
-from .telemetry import (
-    ClassificationPrediction,
-    FeedbackCategory,
-    RegressionPrediction,
-)
+from .telemetry import ClassificationPrediction, FeedbackCategory, RegressionPrediction
 
 # only specify things that should show up on the root page of the reference docs because they are in private modules
 __all__ = ["UNSET", "CreateMode", "DropMode"]

orca_sdk/_shared/__init__.py

@@ -5,5 +5,6 @@ from .metrics import (
     ROCCurve,
     calculate_classification_metrics,
     calculate_pr_curve,
+    calculate_regression_metrics,
     calculate_roc_curve,
 )

orca_sdk/_shared/metrics.py

@@ -24,7 +24,7 @@ def softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
 
 
 # We don't want to depend on transformers just for the eval_pred type in orca_sdk
-def transform_eval_pred(eval_pred: Any) -> tuple[NDArray[np.int64], NDArray[np.float32]]:
+def transform_eval_pred(eval_pred: Any) -> tuple[NDArray, NDArray[np.float32]]:
     # convert results from Trainer compute_metrics param for use in calculate_classification_metrics
     logits, references = eval_pred  # transformers.trainer_utils.EvalPrediction
     if isinstance(logits, tuple):

orca_sdk/_utils/analysis_ui.py

@@ -28,7 +28,7 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
         filters=[("metrics.neighbor_predicted_label_matches_current_label", "==", False)]
     )
     # Sort memories by confidence score (higher confidence first)
-    suggested_relabels.sort(key=lambda x: (x.metrics["neighbor_predicted_label_confidence"] or 0.0), reverse=True)
+    suggested_relabels.sort(key=lambda x: (x.metrics.get("neighbor_predicted_label_confidence", 0.0)), reverse=True)
 
     def update_approved(memory_id: str, selected: bool, current_memory_relabel_map: dict[str, RelabelStatus]):
         current_memory_relabel_map[memory_id]["approved"] = selected
@@ -70,9 +70,9 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
             current_memory_relabel_map[mem_id]["new_label"] = new_label
             confidence = "--"
             current_metrics = current_memory_relabel_map[mem_id]["full_memory"].metrics
-            if current_metrics and new_label == current_metrics["neighbor_predicted_label"]:
+            if current_metrics and new_label == current_metrics.get("neighbor_predicted_label"):
                 confidence = (
-                    round(current_metrics["neighbor_predicted_label_confidence"] or 0.0, 2) if current_metrics else 0
+                    round(current_metrics.get("neighbor_predicted_label_confidence", 0.0), 2) if current_metrics else 0
                 )
             return (
                 gr.HTML(
@@ -99,8 +99,8 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
                     memory_id=mem.memory_id,
                     approved=False,
                     new_label=(
-                        mem.metrics["neighbor_predicted_label"]
-                        if (mem.metrics and isinstance(mem.metrics["neighbor_predicted_label"], int))
+                        mem.metrics.get("neighbor_predicted_label")
+                        if (mem.metrics and isinstance(mem.metrics.get("neighbor_predicted_label"), int))
                         else None
                     ),
                     full_memory=mem,

orca_sdk/_utils/auth.py

@@ -2,23 +2,12 @@
 
 import logging
 import os
-from typing import List
+from typing import List, Literal
 
 from dotenv import load_dotenv
 
-from .._generated_api_client.api import (
-    check_authentication,
-    create_api_key,
-    delete_api_key,
-    delete_org,
-    list_api_keys,
-)
-from .._generated_api_client.client import headers_context, set_base_url, set_headers
-from .._generated_api_client.models import (
-    CreateApiKeyRequest,
-    CreateApiKeyRequestScopeItem,
-)
-from .._generated_api_client.models.api_key_metadata import ApiKeyMetadata
+from ..client import ApiKeyMetadata, orca_api
+from ..credentials import OrcaCredentials
 from .common import DropMode
 
 load_dotenv()  # this needs to be here to ensure env is populated before accessing it
@@ -28,43 +17,44 @@ _ORCA_ROOT_ACCESS_API_KEY = os.environ.get("ORCA_ROOT_ACCESS_API_KEY", "00000000
 _DEFAULT_ORG_ID = os.environ.get("DEFAULT_ORG_ID", "10e50000-0000-4000-a000-a78dca14af3a")
 
 
-def _create_api_key(org_id: str, name: str, scopes: list[str] = ["ADMINISTER"]) -> str:
+def _create_api_key(org_id: str, name: str, scopes: list[Literal["ADMINISTER", "PREDICT"]] = ["ADMINISTER"]) -> str:
     """Creates an API key for the given organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        res = create_api_key(
-            body=CreateApiKeyRequest(name=name, scope=[CreateApiKeyRequestScopeItem(scope) for scope in scopes])
-        )
-        return res.api_key
+    response = orca_api.POST(
+        "/auth/api_key",
+        json={"name": name, "scope": scopes},
+        headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id},
+    )
+    return response["api_key"]
 
 
 def _list_api_keys(org_id: str) -> List[ApiKeyMetadata]:
     """Lists all API keys for the given organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        return list_api_keys()
+    return orca_api.GET("/auth/api_key", headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id})
 
 
 def _delete_api_key(org_id: str, name: str, if_not_exists: DropMode = "error") -> None:
     """Deletes the API key with the given name from the organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        try:
-            delete_api_key(name_or_id=name)
-        except LookupError:
-            if if_not_exists == "error":
-                raise
+    try:
+        orca_api.DELETE(
+            "/auth/api_key/{name_or_id}",
+            params={"name_or_id": name},
+            headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id},
+        )
+    except LookupError:
+        if if_not_exists == "error":
+            raise
 
 
 def _delete_org(org_id: str) -> None:
     """Deletes the organization"""
-    with headers_context({"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id}):
-        delete_org()
+    orca_api.DELETE("/auth/org", headers={"Api-Key": _ORCA_ROOT_ACCESS_API_KEY, "Org-Id": org_id})
 
 
 def _authenticate_local_api(org_id: str = _DEFAULT_ORG_ID, api_key_name: str = "local") -> None:
     """Connect to the local API at http://localhost:1584/ and authenticate with a new API key"""
-    set_base_url("http://localhost:1584/")
     _delete_api_key(org_id, api_key_name, if_not_exists="ignore")
-    set_headers({"Api-Key": _create_api_key(org_id, api_key_name)})
-    check_authentication()
+    OrcaCredentials.set_base_url("http://localhost:1584")
+    OrcaCredentials.set_api_key(_create_api_key(org_id, api_key_name))
     logging.info(f"Authenticated against local API at 'http://localhost:1584' with '{api_key_name}' API key")
 
 

orca_sdk/_utils/pagination.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from typing import Callable, Generic, Iterator, TypedDict, TypeVar, cast, overload
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+class Page(TypedDict, Generic[T]):
+    items: list[T]
+    count: int
+
+
+class _PagedIterable(Generic[T, R]):
+    def __init__(
+        self,
+        fetch: Callable[[int, int], Page[T]],
+        *,
+        transform: Callable[[T], R] | None = None,
+        page_size: int = 100,
+    ) -> None:
+        """
+        Iterate over a paginated endpoint.
+
+        Parameters:
+            fetch: function to fetch a page from the endpoint `(limit: int, offset: int) -> TypedDict[{items: list[T], count: int}]`
+            transform: Optional function to transforms item types `(item: T) -> R`, defaults to identity
+            limit: maximum number of items to fetch per page
+        """
+        self.fetch = fetch
+        self.transform = transform or (lambda x: cast(R, x))
+        self.page_size = page_size
+        self.offset = 0  # tracks how much has been yielded, not fetched
+        self.page = fetch(self.page_size, self.offset)  # fetch first page to populate count
+        self.count = self.page["count"]
+
+    def __iter__(self) -> Iterator[R]:
+        if self.offset >= self.count:
+            self.offset = 0
+            if len(self.page["items"]) < self.count:
+                # refetch first page unless we are still on the first page
+                self.page = self.fetch(self.page_size, self.offset)
+
+        # yield prefetched first page
+        if self.offset == 0:
+            yield from map(self.transform, self.page["items"])
+            self.offset += len(self.page["items"])
+
+        # yield remaining pages one by one
+        while self.offset < self.count:
+            self.page = self.fetch(self.page_size, self.offset)
+            yield from map(self.transform, self.page["items"])
+            self.offset += len(self.page["items"])
+
+    @overload
+    def __getitem__(self, key: int) -> R:
+        pass
+
+    @overload
+    def __getitem__(self, key: slice) -> list[R]:
+        pass
+
+    def __getitem__(self, key: int | slice) -> R | list[R]:
+        if isinstance(key, int):
+            effective_key = key
+            if effective_key < 0:
+                effective_key += self.count
+            if not 0 <= effective_key < self.count:
+                raise IndexError(f"Index {key} out of range")
+            # if key is on current page, return item
+            if self.offset <= effective_key < self.offset + len(self.page["items"]):
+                return self.transform(self.page["items"][effective_key - self.offset])
+            # otherwise, fetch and return the single item
+            return self.transform(self.fetch(1, effective_key)["items"][0])
+
+        elif isinstance(key, slice):
+            start, stop, step = key.indices(self.count)
+            if step != 1:
+                raise ValueError("Stepped slicing is not supported")
+            start = start + self.count if start < 0 else start or 0
+            stop = stop + self.count if stop < 0 else stop or self.count
+            if start >= self.count or stop > self.count:
+                raise IndexError(f"Slice {key} out of range")
+            limit = min(self.page_size, stop - start)
+            if limit <= 0:
+                return []
+            items = []
+            for i in range(start, stop, limit):
+                page = self.fetch(limit, i)
+                items.extend(map(self.transform, page["items"]))
+            return items
+
+    def __len__(self) -> int:
+        return self.count
+
+
+# type checking workaround until python 3.13 allows declaring the class as PagedIterable[T, R = T]
+
+
+@overload
+def PagedIterable(
+    fetch: Callable[[int, int], Page[T]],
+    *,
+    transform: None = None,
+    page_size: int = 100,
+) -> _PagedIterable[T, T]:
+    pass
+
+
+@overload
+def PagedIterable(
+    fetch: Callable[[int, int], Page[T]],
+    *,
+    transform: Callable[[T], R],
+    page_size: int = 100,
+) -> _PagedIterable[T, R]:
+    pass
+
+
+def PagedIterable(
+    fetch: Callable[[int, int], Page[T]],
+    *,
+    transform: Callable[[T], R] | None = None,
+    page_size: int = 100,
+) -> _PagedIterable[T, R]:
+    return _PagedIterable(fetch, transform=transform, page_size=page_size)

orca_sdk/_utils/pagination_test.py

@@ -0,0 +1,132 @@
+import pytest
+
+from .pagination import Page, PagedIterable
+
+
+class MockEndpoint:
+    """Mock paginated endpoint for testing"""
+
+    def __init__(self, total_items: int):
+        self.items = list(range(total_items))
+        self.fetch_count = 0
+
+    def fetch(self, limit: int, offset: int) -> Page[int]:
+        self.fetch_count += 1
+        end_index = min(offset + limit, len(self.items))
+        items = self.items[offset:end_index]
+        return {"items": items, "count": len(self.items)}
+
+
+def test_basic_pagination():
+    # Given a mock endpoint with 5 items
+    endpoint = MockEndpoint(5)
+    # When doing a paginated iteration
+    paginated = PagedIterable(endpoint.fetch, page_size=2)
+    # Then we should be able to iterate through all items
+    assert list(paginated) == [0, 1, 2, 3, 4]
+    # And the length should be correct
+    assert len(paginated) == 5
+    # And 3 requests: [0,1], [2,3], [4] should have been made, one for each page
+    assert endpoint.fetch_count == 3
+
+
+def test_empty_results():
+    # Given an empty mock endpoint
+    endpoint = MockEndpoint(0)
+    # When doing a paginated iteration
+    paginated = PagedIterable(endpoint.fetch, page_size=5)
+    # Then we should get an empty list
+    assert list(paginated) == []
+    # And the length should be 0
+    assert len(paginated) == 0
+    # And only one request should have been made, for the first page
+    assert endpoint.fetch_count == 1
+
+
+def test_transform_function():
+    # Given a mock endpoint with 4 items
+    endpoint = MockEndpoint(4)
+    # And a transform function that doubles the items
+    transform = lambda x: f"2x={2*x}"
+    # When doing a paginated iteration with a transform function
+    paginated = PagedIterable(endpoint.fetch, transform=transform, page_size=2)
+    # Then we should get the transformed items
+    assert list(paginated) == ["2x=0", "2x=2", "2x=4", "2x=6"]
+
+
+def test_multiple_iterations():
+    # Given a mock endpoint with 5 items
+    endpoint = MockEndpoint(5)
+    # When we do 2 paginated iterations
+    paginated = PagedIterable(endpoint.fetch, page_size=2)
+    result1 = list(paginated)
+    result2 = list(paginated)
+    # Then we should get the same items twice
+    assert result1 == result2 == [0, 1, 2, 3, 4]
+    # And 6 requests should have been made, 3 for each iteration
+    assert endpoint.fetch_count == 6
+
+
+def test_single_page_optimization():
+    # Given a mock endpoint with 5 items
+    endpoint = MockEndpoint(5)
+    # When doing a paginated iteration with a limit that is greater than the number of items
+    paginated = PagedIterable(endpoint.fetch, page_size=10)
+    # Then we should get all items
+    assert list(paginated) == [0, 1, 2, 3, 4]
+    # And the length should be 5
+    assert len(paginated) == 5
+    # And only one request should have been made
+    assert endpoint.fetch_count == 1
+    # And a second iteration should not make any additional requests
+    assert list(paginated) == [0, 1, 2, 3, 4]
+    assert endpoint.fetch_count == 1
+
+
+def test_indexing():
+    # Given a mock endpoint with 7 items
+    endpoint = MockEndpoint(7)
+    # When creating a paginated iterable with page size 3
+    paginated = PagedIterable(endpoint.fetch, page_size=3)
+    # Then we should be able to access items by index
+    assert paginated[0] == 0
+    assert paginated[2] == 2
+    assert paginated[6] == 6
+    # And negative indices should work
+    assert paginated[-1] == 6
+    # And accessing out of bounds should raise IndexError
+    with pytest.raises(IndexError):
+        paginated[7]
+    with pytest.raises(IndexError):
+        paginated[-8]
+    # And transforms are applied
+    assert PagedIterable(endpoint.fetch, transform=lambda x: x * 10, page_size=3)[1] == 10
+
+
+def test_slicing():
+    # Given a mock endpoint with 10 items
+    endpoint = MockEndpoint(10)
+    # When creating a paginated iterable
+    paginated = PagedIterable(endpoint.fetch, page_size=3)
+    # Then we should be able to slice it
+    assert list(paginated[2:5]) == [2, 3, 4]
+    assert list(paginated[:3]) == [0, 1, 2]
+    assert list(paginated[7:]) == [7, 8, 9]
+    # And negative indices should work
+    assert list(paginated[:-5]) == [0, 1, 2, 3, 4, 5]
+    assert list(paginated[-3:]) == [7, 8, 9]
+    assert list(paginated[-5:-2]) == [5, 6, 7]
+    # And empty slices should work
+    assert list(paginated[5:5]) == []
+    # And slicing with a start and stop that are out of bounds should raise IndexError
+    with pytest.raises(IndexError):
+        list(paginated[20:25])
+    # And slicing with a step other than 1 should raise ValueError
+    with pytest.raises(ValueError):
+        list(paginated[::2])
+    with pytest.raises(ValueError):
+        list(paginated[1:8:3])
+    with pytest.raises(ValueError):
+        list(paginated[::-1])
+    # And transforms are applied
+    assert list(PagedIterable(endpoint.fetch, transform=lambda x: x * 10, page_size=3)[1:3]) == [10, 20]