orca-sdk 0.1.9__py3-none-any.whl

orca_sdk/_utils/analysis_ui.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import logging
+import re
+from pathlib import Path
+from typing import TypedDict, cast
+
+import gradio as gr
+
+from ..memoryset import LabeledMemory, LabeledMemoryset
+
+# Suppress all httpx logs
+logging.getLogger("httpx").setLevel(logging.CRITICAL)
+
+# Optionally suppress other libraries Gradio might use
+logging.getLogger("gradio").setLevel(logging.CRITICAL)
+
+
+class RelabelStatus(TypedDict):
+    memory_id: str
+    approved: bool
+    new_label: int | None
+    full_memory: LabeledMemory
+
+
+def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
+    suggested_relabels = memoryset.query(
+        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.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
+        return current_memory_relabel_map
+
+    def approve_all(current_all_approved, selected: bool):
+        for mem_id in current_all_approved:
+            current_all_approved[mem_id]["approved"] = selected
+        return current_all_approved, selected
+
+    def apply_selected(current_memory_relabel_map: dict[str, RelabelStatus], progress=gr.Progress(track_tqdm=True)):
+        progress(0, desc="Processing label updates...")
+        to_be_deleted = []
+        approved_relabels = [mem for mem in current_memory_relabel_map.values() if mem["approved"]]
+        for memory in progress.tqdm(approved_relabels, desc="Applying label updates..."):
+            memory = cast(RelabelStatus, memory)
+            new_label = memory["new_label"]
+            assert isinstance(new_label, int)
+            memoryset.update(
+                {
+                    "memory_id": memory["memory_id"],
+                    "label": new_label,
+                }
+            )
+            to_be_deleted.append(memory["memory_id"])
+        for mem_id in to_be_deleted:
+            del current_memory_relabel_map[mem_id]
+        return (
+            current_memory_relabel_map,
+            gr.HTML(
+                f"<h1 style='display: inline-block; position: fixed; z-index: 1000; left: 36px; top: 14px;'>Suggested Label Updates: {len(current_memory_relabel_map)}</h1>",
+            ),
+        )
+
+    def update_label(mem_id: str, label: str, current_memory_relabel_map: dict[str, RelabelStatus]):
+        match = re.search(r".*\((\d+)\)$", label)
+        if match:
+            new_label = int(match.group(1))
+            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.get("neighbor_predicted_label"):
+                confidence = (
+                    round(current_metrics.get("neighbor_predicted_label_confidence", 0.0), 2) if current_metrics else 0
+                )
+            return (
+                gr.HTML(
+                    f"<p style='font-size: 10px; color: #888;'>Confidence: {confidence}</p>",
+                    elem_classes="no-padding",
+                ),
+                current_memory_relabel_map,
+            )
+        else:
+            logging.error(f"Invalid label format: {label}")
+
+    with gr.Blocks(
+        fill_width=True,
+        title="Suggested Label Updates",
+        css_paths=str(Path(__file__).parent / "analysis_ui_style.css"),
+    ) as demo:
+        label_names = memoryset.label_names
+
+        refresh = gr.State(False)
+        all_approved = gr.State(False)
+        memory_relabel_map = gr.State(
+            {
+                mem.memory_id: RelabelStatus(
+                    memory_id=mem.memory_id,
+                    approved=False,
+                    new_label=(
+                        mem.metrics.get("neighbor_predicted_label")
+                        if (mem.metrics and isinstance(mem.metrics.get("neighbor_predicted_label"), int))
+                        else None
+                    ),
+                    full_memory=mem,
+                )
+                for mem in suggested_relabels
+            }
+        )
+
+        @gr.render(
+            inputs=[memory_relabel_map, all_approved],
+            triggers=[demo.load, refresh.change, all_approved.change, memory_relabel_map.change],  # type: ignore[arg-type]
+        )
+        def render_table(current_memory_relabel_map, current_all_approved):
+            if len(current_memory_relabel_map):
+                with gr.Group(elem_classes="header"):
+                    title = gr.HTML(
+                        f"<h1 style='display: inline-block; position: fixed; z-index: 1000; left: 36px; top: 14px;'>Suggested Label Updates: {len(current_memory_relabel_map)}</h1>"
+                    )
+                    apply_selected_button = gr.Button("Apply Selected", elem_classes="button")
+                    apply_selected_button.click(
+                        apply_selected,
+                        inputs=[memory_relabel_map],
+                        outputs=[memory_relabel_map, title],
+                        show_progress="full",
+                    )
+                with gr.Row(equal_height=True, variant="panel", elem_classes="margin-top"):
+                    with gr.Column(scale=9):
+                        gr.Markdown("**Value**")
+                    with gr.Column(scale=2, min_width=90):
+                        gr.Markdown("**Current Label**")
+                    with gr.Column(scale=3, min_width=150):
+                        gr.Markdown("**Suggested Label**", elem_classes="centered")
+                    with gr.Column(scale=2, min_width=50):
+                        approve_all_checkbox = gr.Checkbox(
+                            show_label=False,
+                            value=current_all_approved,
+                            label="",
+                            container=False,
+                            elem_classes="centered",
+                        )
+                        approve_all_checkbox.change(
+                            approve_all,
+                            inputs=[memory_relabel_map, approve_all_checkbox],
+                            outputs=[memory_relabel_map, all_approved],
+                        )
+                for i, memory_relabel in enumerate(current_memory_relabel_map.values()):
+                    mem = memory_relabel["full_memory"]
+                    predicted_label = mem.metrics["neighbor_predicted_label"]
+                    predicted_label_name = label_names[predicted_label]
+                    predicted_label_confidence = mem.metrics.get("neighbor_predicted_label_confidence", 0)
+
+                    with gr.Row(equal_height=True, variant="panel"):
+                        with gr.Column(scale=9):
+                            assert isinstance(mem.value, str)
+                            gr.Markdown(mem.value, label="Value", height=50)
+                        with gr.Column(scale=2, min_width=90):
+                            gr.Markdown(f"{mem.label_name} ({mem.label})", label="Current Label", height=50)
+                        with gr.Column(scale=3, min_width=150):
+                            dropdown = gr.Dropdown(
+                                choices=[f"{label_name} ({i})" for i, label_name in enumerate(label_names)],
+                                label="SuggestedLabel",
+                                value=f"{predicted_label_name} ({predicted_label})",
+                                interactive=True,
+                                container=False,
+                            )
+                            confidence = gr.HTML(
+                                f"<p style='font-size: 10px; color: #888;'>Confidence: {predicted_label_confidence:.2f}</p>",
+                                elem_classes="no-padding",
+                            )
+                            dropdown.change(
+                                lambda val, map, mem_id=mem.memory_id: update_label(mem_id, val, map),
+                                inputs=[dropdown, memory_relabel_map],
+                                outputs=[confidence, memory_relabel_map],
+                            )
+                        with gr.Column(scale=2, min_width=50):
+                            checkbox = gr.Checkbox(
+                                show_label=False,
+                                label="",
+                                value=current_memory_relabel_map[mem.memory_id]["approved"],
+                                container=False,
+                                elem_classes="centered",
+                                interactive=True,
+                            )
+                            checkbox.input(
+                                lambda selected, map, mem_id=mem.memory_id: update_approved(mem_id, selected, map),
+                                inputs=[checkbox, memory_relabel_map],
+                                outputs=[memory_relabel_map],
+                            )
+
+            else:
+                gr.HTML("<h1>No suggested label updates</h1>")
+
+    demo.launch()

orca_sdk/_utils/analysis_ui_style.css

@@ -0,0 +1,51 @@
+.centered input {
+  margin: auto;
+}
+.centered p {
+  text-align: center;
+}
+.button {
+  display: inline-block;
+  max-width: 250px;
+  background-color: #2b9a66;
+  color: white;
+  position: fixed;
+  z-index: 1000;
+  right: 36px;
+  border-radius: 8px;
+  top: 12px;
+}
+.margin-top {
+  margin-top: 60px;
+}
+.header {
+  position: fixed;
+  z-index: 1000;
+  height: 64px;
+  left: 0;
+  top: 0;
+  border-radius: 0;
+}
+
+input[type='checkbox']:checked,
+input[type='checkbox']:checked:hover,
+input[type='checkbox']:checked:focus {
+  background-color: #2b9a66;
+  border-color: #2b9a66;
+}
+input[type='checkbox']:focus {
+  border-color: #2b9a66;
+}
+.html-container:has(.no-padding) {
+  padding: 0;
+}
+
+.progress-bar {
+  background-color: #2b9a66;
+}
+.header .full {
+  position: fixed !important;
+  z-index: 1100;
+  background-color: #e4e4e7;
+  height: 68px;
+}

orca_sdk/_utils/auth.py

@@ -0,0 +1,65 @@
+"""This module contains internal utils for managing api keys in tests"""
+
+import logging
+import os
+from typing import List, Literal
+
+from dotenv import load_dotenv
+
+from ..client import ApiKeyMetadata, OrcaClient
+from .common import DropMode
+
+load_dotenv()  # this needs to be here to ensure env is populated before accessing it
+
+# the defaults here must match nautilus and lighthouse config defaults
+_ORCA_ROOT_ACCESS_API_KEY = os.environ.get("ORCA_ROOT_ACCESS_API_KEY", "00000000-0000-0000-0000-000000000000")
+_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[Literal["ADMINISTER", "PREDICT"]] = ["ADMINISTER"]) -> str:
+    """Creates an API key for the given organization"""
+    client = OrcaClient._resolve_client()
+    response = client.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"""
+    client = OrcaClient._resolve_client()
+    return client.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"""
+    try:
+        client = OrcaClient._resolve_client()
+        client.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"""
+    client = OrcaClient._resolve_client()
+    client.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"""
+    _delete_api_key(org_id, api_key_name, if_not_exists="ignore")
+    client = OrcaClient._resolve_client()
+    client.base_url = "http://localhost:1584"
+    client.headers.update({"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")
+
+
+__all__ = ["_create_api_key", "_delete_api_key", "_delete_org", "_list_api_keys", "_authenticate_local_api"]

orca_sdk/_utils/auth_test.py

@@ -0,0 +1,31 @@
+from uuid import uuid4
+
+from ..credentials import OrcaCredentials
+from .auth import _create_api_key, _delete_api_key, _delete_org, _list_api_keys
+
+
+def test_list_api_keys(org_id):
+    assert len(_list_api_keys(org_id)) >= 1
+
+
+def test_create_api_key(org_id):
+    name = f"test-{uuid4().hex[:8]}"
+    api_key = _create_api_key(org_id=org_id, name=name)
+    assert api_key is not None
+    assert name in [api_key.name for api_key in OrcaCredentials.list_api_keys()]
+
+
+def test_delete_api_key(org_id):
+    name = f"test-{uuid4().hex[:8]}"
+    api_key = _create_api_key(org_id=org_id, name=name)
+    assert api_key is not None
+    assert name in [api_key.name for api_key in OrcaCredentials.list_api_keys()]
+    _delete_api_key(org_id=org_id, name=name)
+    assert name not in [api_key.name for api_key in OrcaCredentials.list_api_keys()]
+
+
+def test_delete_org(other_org_id):
+    _create_api_key(org_id=other_org_id, name="test")
+    assert len(_list_api_keys(other_org_id)) >= 1
+    _delete_org(other_org_id)
+    assert len(_list_api_keys(other_org_id)) == 0

orca_sdk/_utils/common.py

@@ -0,0 +1,37 @@
+from typing import Any, Literal
+
+CreateMode = Literal["error", "open"]
+"""
+Mode for creating a resource.
+
+**Options:**
+
+- `"error"`: raise an error if a resource with the same name already exists
+- `"open"`: open the resource with the same name if it exists
+"""
+
+DropMode = Literal["error", "ignore"]
+"""
+Mode for deleting a resource.
+
+**Options:**
+
+- `"error"`: raise an error if the resource does not exist
+- `"ignore"`: do nothing if the resource does not exist
+"""
+
+
+class _UnsetSentinel:
+    """See corresponding class in orcalib.pydantic_utils"""
+
+    def __bool__(self) -> bool:
+        return False
+
+    def __repr__(self) -> str:
+        return "UNSET"
+
+
+UNSET: Any = _UnsetSentinel()
+"""
+Default value to indicate that no update should be applied to a field and it should not be set to None
+"""

orca_sdk/_utils/data_parsing.py

@@ -0,0 +1,129 @@
+import pickle
+from dataclasses import asdict, is_dataclass
+from os import PathLike
+from typing import Any, cast
+
+from datasets import Dataset
+from datasets.exceptions import DatasetGenerationError
+from torch.utils.data import DataLoader as TorchDataLoader
+from torch.utils.data import Dataset as TorchDataset
+
+
+def parse_dict_like(item: Any, column_names: list[str] | None = None) -> dict:
+    if isinstance(item, dict):
+        return item
+
+    if isinstance(item, tuple):
+        if column_names is not None:
+            if len(item) != len(column_names):
+                raise ValueError(
+                    f"Tuple length ({len(item)}) does not match number of column names ({len(column_names)})"
+                )
+            return {column_names[i]: item[i] for i in range(len(item))}
+        elif hasattr(item, "_fields") and all(isinstance(field, str) for field in item._fields):  # type: ignore
+            return {field: getattr(item, field) for field in item._fields}  # type: ignore
+        else:
+            raise ValueError("For datasets that return unnamed tuples, please provide column_names argument")
+
+    if is_dataclass(item) and not isinstance(item, type):
+        return asdict(item)
+
+    raise ValueError(f"Cannot parse {type(item)}")
+
+
+def parse_batch(batch: Any, column_names: list[str] | None = None) -> list[dict]:
+    if isinstance(batch, list):
+        return [parse_dict_like(item, column_names) for item in batch]
+
+    batch = parse_dict_like(batch, column_names)
+    keys = list(batch.keys())
+    batch_size = len(batch[keys[0]])
+    for key in keys:
+        if not len(batch[key]) == batch_size:
+            raise ValueError(f"Batch must consist of values of the same length, but {key} has length {len(batch[key])}")
+    return [{key: batch[key][idx] for key in keys} for idx in range(batch_size)]
+
+
+def hf_dataset_from_torch(
+    torch_data: TorchDataLoader | TorchDataset,
+    column_names: list[str] | None = None,
+) -> Dataset:
+    """
+    Create a HuggingFace Dataset from a PyTorch DataLoader or Dataset.
+
+    NOTE:  It's important to ignore the cached files when testing (i.e., ignore_cache=Ture), because
+    cached results can ignore changes you've made to tests. This can make a test appear to succeed
+    when it's actually broken or vice versa.
+
+    Params:
+        torch_data: A PyTorch DataLoader or Dataset object to create the HuggingFace Dataset from.
+        column_names: Optional list of column names to use for the dataset. If not provided,
+            the column names will be inferred from the data.
+    Returns:
+        A HuggingFace Dataset object containing the data from the PyTorch DataLoader or Dataset.
+    """
+    if isinstance(torch_data, TorchDataLoader):
+        dataloader = torch_data
+    else:
+        dataloader = TorchDataLoader(torch_data, batch_size=1, collate_fn=lambda x: x)
+
+    # Collect data from the dataloader into a list to avoid serialization issues
+    # with Dataset.from_generator in Python 3.14 (see datasets issue #7839)
+    data_list = []
+    try:
+        for batch in dataloader:
+            data_list.extend(parse_batch(batch, column_names=column_names))
+    except ValueError as e:
+        raise DatasetGenerationError(str(e)) from e
+
+    ds = Dataset.from_list(data_list)
+
+    if not isinstance(ds, Dataset):
+        raise ValueError(f"Failed to create dataset from list: {type(ds)}")
+    return ds
+
+
+def hf_dataset_from_disk(file_path: str | PathLike) -> Dataset:
+    """
+    Load a dataset from disk into a HuggingFace Dataset object.
+
+    Params:
+        file_path: Path to the file on disk to create the memoryset from. The file type will
+                be inferred from the file extension. The following file types are supported:
+
+                - .pkl: [`Pickle`][pickle] files containing lists of dictionaries or dictionaries of columns
+                - .json/.jsonl: [`JSON`][json] and [`JSON`] Lines files
+                - .csv: [`CSV`][csv] files
+                - .parquet: [`Parquet`](https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetFile.html#pyarrow.parquet.ParquetFile) files
+                - dataset directory: Directory containing a saved HuggingFace [`Dataset`][datasets.Dataset]
+
+    Returns:
+        A HuggingFace Dataset object containing the loaded data.
+
+    Raises:
+        [`ValueError`][ValueError]: If the pickle file contains unsupported data types or if
+            loading the dataset fails for any reason.
+    """
+    if str(file_path).endswith(".pkl"):
+        data = pickle.load(open(file_path, "rb"))
+        if isinstance(data, list):
+            return Dataset.from_list(data)
+        elif isinstance(data, dict):
+            return Dataset.from_dict(data)
+        else:
+            raise ValueError(f"Unsupported pickle file: {file_path}")
+    elif str(file_path).endswith(".json"):
+        hf_dataset = Dataset.from_json(file_path)
+    elif str(file_path).endswith(".jsonl"):
+        hf_dataset = Dataset.from_json(file_path)
+    elif str(file_path).endswith(".csv"):
+        hf_dataset = Dataset.from_csv(file_path)
+    elif str(file_path).endswith(".parquet"):
+        hf_dataset = Dataset.from_parquet(file_path)
+    else:
+        try:
+            hf_dataset = Dataset.load_from_disk(file_path)
+        except Exception as e:
+            raise ValueError(f"Failed to load dataset from disk: {e}")
+
+    return cast(Dataset, hf_dataset)