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

orca_sdk/memoryset.py

@@ -8,88 +8,43 @@ from typing import Any, Generic, Iterable, Literal, Self, TypeVar, cast, overloa
 
 import pandas as pd
 import pyarrow as pa
-from attrs import fields
 from datasets import Dataset
 from torch.utils.data import DataLoader as TorchDataLoader
 from torch.utils.data import Dataset as TorchDataset
 
-from ._generated_api_client.api import (
-    analyze_memoryset,
-    clone_memoryset,
-    create_embedding_evaluation,
-    create_memoryset,
-    delete_memories,
-    delete_memoryset,
-    get_analysis,
-    get_embedding_evaluation,
-    get_memories,
-    get_memory,
-    get_memoryset,
-    insert_memories_gpu,
-    list_memories_with_feedback,
-    list_memorysets,
-    memoryset_lookup_gpu,
-    potential_duplicate_groups,
-    query_memoryset,
-    suggest_cascading_edits,
-    update_memories_gpu,
-    update_memory_gpu,
-    update_memoryset,
-)
-from ._generated_api_client.models import (
-    CascadeEditSuggestionsRequest,
+from ._utils.common import UNSET, CreateMode, DropMode
+from .client import (
     CascadingEditSuggestion,
     CloneMemorysetRequest,
-    ColumnType,
     CreateMemorysetRequest,
-    CreateMemorysetRequestIndexParams,
-    CreateMemorysetRequestIndexType,
-    DeleteMemoriesRequest,
-    EmbeddingEvaluationRequest,
+    EmbeddingModelResult,
     FilterItem,
-    FilterItemOp,
-    GetMemoriesRequest,
-)
-from ._generated_api_client.models import LabeledMemory as LabeledMemoryResponse
-from ._generated_api_client.models import (
-    LabeledMemoryInsert,
-    LabeledMemoryInsertMetadata,
 )
-from ._generated_api_client.models import (
-    LabeledMemoryLookup as LabeledMemoryLookupResponse,
-)
-from ._generated_api_client.models import (
+from .client import LabeledMemory as LabeledMemoryResponse
+from .client import LabeledMemoryInsert
+from .client import LabeledMemoryLookup as LabeledMemoryLookupResponse
+from .client import (
     LabeledMemoryUpdate,
-    LabeledMemoryUpdateMetadataType0,
     LabeledMemoryWithFeedbackMetrics,
     LabelPredictionMemoryLookup,
-    ListMemoriesRequest,
-    LookupRequest,
+    MemoryMetrics,
     MemorysetAnalysisConfigs,
-    MemorysetAnalysisRequest,
     MemorysetMetadata,
+    MemorysetMetrics,
     MemorysetUpdate,
     MemoryType,
-    PretrainedEmbeddingModelName,
-)
-from ._generated_api_client.models import ScoredMemory as ScoredMemoryResponse
-from ._generated_api_client.models import ScoredMemoryInsert, ScoredMemoryInsertMetadata
-from ._generated_api_client.models import (
-    ScoredMemoryLookup as ScoredMemoryLookupResponse,
 )
-from ._generated_api_client.models import (
+from .client import ScoredMemory as ScoredMemoryResponse
+from .client import ScoredMemoryInsert
+from .client import ScoredMemoryLookup as ScoredMemoryLookupResponse
+from .client import (
     ScoredMemoryUpdate,
-    ScoredMemoryUpdateMetadataType0,
     ScoredMemoryWithFeedbackMetrics,
     ScorePredictionMemoryLookup,
     TelemetryFilterItem,
-    TelemetryFilterItemOp,
-    TelemetryMemoriesRequest,
     TelemetrySortOptions,
-    TelemetrySortOptionsDirection,
+    orca_api,
 )
-from ._generated_api_client.types import UNSET as CLIENT_UNSET
-from ._utils.common import UNSET, CreateMode, DropMode
 from .datasource import Datasource
 from .embedding_model import (
     FinetunedEmbeddingModel,
@@ -151,6 +106,10 @@ FORBIDDEN_METADATA_COLUMN_NAMES = {
 }
 
 
+def _is_metric_column(column: str):
+    return column in ["feedback_metrics", "lookup"]
+
+
 def _parse_filter_item_from_tuple(input: FilterItemTuple) -> FilterItem | TelemetryFilterItem:
     field = input[0].split(".")
     if (
@@ -158,9 +117,11 @@ def _parse_filter_item_from_tuple(input: FilterItemTuple) -> FilterItem | Teleme
         and field[0] not in DEFAULT_COLUMN_NAMES | TYPE_SPECIFIC_COLUMN_NAMES | FORBIDDEN_METADATA_COLUMN_NAMES
     ):
         field = ["metadata", field[0]]
-    op = FilterItemOp(input[1])
+    op = input[1]
     value = input[2]
-    if field[0] in ["feedback_metrics", "lookup"]:
+    if isinstance(value, datetime):
+        value = value.isoformat()
+    if _is_metric_column(field[0]):
         if not (
             (isinstance(value, list) and all(isinstance(v, float) or isinstance(v, int) for v in value))
             or isinstance(value, float)
@@ -173,9 +134,11 @@ def _parse_filter_item_from_tuple(input: FilterItemTuple) -> FilterItem | Teleme
             )
         elif field[0] == "lookup" and (len(field) != 2 or field[1] != "count"):
             raise ValueError("Lookup filters must follow the format `lookup.count`")
-        op = TelemetryFilterItemOp(op)
-
-        return TelemetryFilterItem(field=field, op=op, value=cast(list[float] | list[int] | float | int, value))
+        if op == "like":
+            raise ValueError("Like filters are not supported on metric columns")
+        op = cast(Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in"], op)
+        value = cast(float | int | list[float] | list[int], value)
+        return TelemetryFilterItem(field=field, op=op, value=value)
 
     return FilterItem(field=field, op=op, value=value)
 
@@ -201,7 +164,7 @@ def _parse_sort_item_from_tuple(
             raise ValueError("Lookup must follow the format `lookup.count`")
         if field[1] != "count":
             raise ValueError("Lookup can only be sorted on count")
-    return TelemetrySortOptions(field=field, direction=TelemetrySortOptionsDirection(input[1]))
+    return TelemetrySortOptions(field=field, direction=input[1])
 
 
 def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMemoryInsert | ScoredMemoryInsert:
@@ -212,45 +175,26 @@ def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMem
     if source_id and not isinstance(source_id, str):
         raise ValueError("Memory source_id must be a string")
     match type:
-        case MemoryType.LABELED:
+        case "LABELED":
             label = memory.get("label")
             if not isinstance(label, int):
                 raise ValueError("Memory label must be an integer")
-            metadata = LabeledMemoryInsertMetadata.from_dict(
-                {k: v for k, v in memory.items() if k not in DEFAULT_COLUMN_NAMES | {"label"}}
-            )
+            metadata = {k: v for k, v in memory.items() if k not in DEFAULT_COLUMN_NAMES | {"label"}}
             if any(k in metadata for k in FORBIDDEN_METADATA_COLUMN_NAMES):
                 raise ValueError(
                     f"The following column names are reserved: {', '.join(FORBIDDEN_METADATA_COLUMN_NAMES)}"
                 )
-            return LabeledMemoryInsert(value=value, label=label, source_id=source_id, metadata=metadata)
-        case MemoryType.SCORED:
+            return {"value": value, "label": label, "source_id": source_id, "metadata": metadata}
+        case "SCORED":
             score = memory.get("score")
             if not isinstance(score, (int, float)):
                 raise ValueError("Memory score must be a number")
-            metadata = ScoredMemoryInsertMetadata.from_dict(
-                {k: v for k, v in memory.items() if k not in DEFAULT_COLUMN_NAMES | {"score"}}
-            )
+            metadata = {k: v for k, v in memory.items() if k not in DEFAULT_COLUMN_NAMES | {"score"}}
             if any(k in metadata for k in FORBIDDEN_METADATA_COLUMN_NAMES):
                 raise ValueError(
                     f"The following column names are reserved: {', '.join(FORBIDDEN_METADATA_COLUMN_NAMES)}"
                 )
-            return ScoredMemoryInsert(value=value, score=score, source_id=source_id, metadata=metadata)
-
-
-@overload
-def _parse_memory_update(update: dict[str, Any], type: Literal[MemoryType.LABELED]) -> LabeledMemoryUpdate:
-    pass
-
-
-@overload
-def _parse_memory_update(update: dict[str, Any], type: Literal[MemoryType.SCORED]) -> ScoredMemoryUpdate:
-    pass
-
-
-@overload
-def _parse_memory_update(update: dict[str, Any], type: MemoryType) -> ScoredMemoryUpdate | LabeledMemoryUpdate:
-    pass
+            return {"value": value, "score": score, "source_id": source_id, "metadata": metadata}
 
 
 def _parse_memory_update(update: dict[str, Any], type: MemoryType) -> LabeledMemoryUpdate | ScoredMemoryUpdate:
@@ -259,41 +203,42 @@ def _parse_memory_update(update: dict[str, Any], type: MemoryType) -> LabeledMem
     memory_id = update["memory_id"]
     if not isinstance(memory_id, str):
         raise ValueError("memory_id must be a string")
-    value = update.get("value", CLIENT_UNSET)
-    if value is not CLIENT_UNSET and not isinstance(value, str):
-        raise ValueError("value must be a string or unset")
-    source_id = update.get("source_id", CLIENT_UNSET)
-    if source_id is not CLIENT_UNSET and not isinstance(source_id, str):
-        raise ValueError("source_id must be a string or unset")
+    payload: LabeledMemoryUpdate | ScoredMemoryUpdate = {"memory_id": memory_id}
+    if "value" in update:
+        if not isinstance(update["value"], str):
+            raise ValueError("value must be a string or unset")
+        payload["value"] = update["value"]
+    if "source_id" in update:
+        if not isinstance(update["source_id"], str):
+            raise ValueError("source_id must be a string or unset")
+        payload["source_id"] = update["source_id"]
     match type:
-        case MemoryType.LABELED:
-            label = update.get("label", CLIENT_UNSET)
-            if label is not CLIENT_UNSET and not isinstance(label, int):
-                raise ValueError("label must be an integer or unset")
-            metadata = LabeledMemoryUpdateMetadataType0.from_dict(
-                {k: v for k, v in update.items() if k not in DEFAULT_COLUMN_NAMES | {"memory_id", "label"}}
-            )
+        case "LABELED":
+            payload = cast(LabeledMemoryUpdate, payload)
+            if "label" in update:
+                if not isinstance(update["label"], int):
+                    raise ValueError("label must be an integer or unset")
+                payload["label"] = update["label"]
+            metadata = {k: v for k, v in update.items() if k not in DEFAULT_COLUMN_NAMES | {"memory_id", "label"}}
             if any(k in metadata for k in FORBIDDEN_METADATA_COLUMN_NAMES):
                 raise ValueError(
                     f"Cannot update the following metadata keys: {', '.join(FORBIDDEN_METADATA_COLUMN_NAMES)}"
                 )
-            return LabeledMemoryUpdate(
-                memory_id=memory_id, value=value, label=label, source_id=source_id, metadata=metadata
-            )
-        case MemoryType.SCORED:
-            score = update.get("score", CLIENT_UNSET)
-            if score is not CLIENT_UNSET and not isinstance(score, (int, float)):
-                raise ValueError("score must be a number or unset")
-            metadata = ScoredMemoryUpdateMetadataType0.from_dict(
-                {k: v for k, v in update.items() if k not in DEFAULT_COLUMN_NAMES | {"memory_id", "score"}}
-            )
+            payload["metadata"] = metadata
+            return payload
+        case "SCORED":
+            payload = cast(ScoredMemoryUpdate, payload)
+            if "score" in update:
+                if not isinstance(update["score"], (int, float)):
+                    raise ValueError("score must be a number or unset")
+                payload["score"] = update["score"]
+            metadata = {k: v for k, v in update.items() if k not in DEFAULT_COLUMN_NAMES | {"memory_id", "score"}}
             if any(k in metadata for k in FORBIDDEN_METADATA_COLUMN_NAMES):
                 raise ValueError(
                     f"Cannot update the following metadata keys: {', '.join(FORBIDDEN_METADATA_COLUMN_NAMES)}"
                 )
-            return ScoredMemoryUpdate(
-                memory_id=memory_id, value=value, score=score, source_id=source_id, metadata=metadata
-            )
+            payload["metadata"] = metadata
+            return cast(ScoredMemoryUpdate, payload)
 
 
 class _Memory(ABC):
@@ -303,7 +248,7 @@ class _Memory(ABC):
     created_at: datetime
     updated_at: datetime
     metadata: dict[str, str | float | int | bool | None]
-    metrics: dict[str, Any]
+    metrics: MemoryMetrics
     memory_id: str
     memory_version: int
     feedback_metrics: dict[str, Any]
@@ -326,19 +271,17 @@ class _Memory(ABC):
     ):
         # for internal use only, do not document
         self.memoryset_id = memoryset_id
-        self.memory_id = memory.memory_id
-        self.memory_version = memory.memory_version
-        self.value = memory.value
-        self.embedding = memory.embedding
-        self.source_id = memory.source_id
-        self.created_at = memory.created_at
-        self.updated_at = memory.updated_at
-        self.metadata = memory.metadata.to_dict()
-        self.metrics = memory.metrics.to_dict() if memory.metrics else {}
-        self.feedback_metrics = (
-            getattr(memory, "feedback_metrics").to_dict() if hasattr(memory, "feedback_metrics") else {}
-        )
-        self.lookup_count = getattr(memory, "lookup_count", 0)
+        self.memory_id = memory["memory_id"]
+        self.memory_version = memory["memory_version"]
+        self.value = cast(str, memory["value"])
+        self.embedding = memory["embedding"]
+        self.source_id = memory["source_id"]
+        self.created_at = datetime.fromisoformat(memory["created_at"])
+        self.updated_at = datetime.fromisoformat(memory["updated_at"])
+        self.metadata = memory["metadata"]
+        self.metrics = memory["metrics"] if "metrics" in memory else {}
+        self.feedback_metrics = memory.get("feedback_metrics", {}) or {}
+        self.lookup_count = memory.get("lookup_count", 0)
 
     def __getattr__(self, key: str) -> Any:
         if key.startswith("__") or key not in self.metadata:
@@ -368,9 +311,10 @@ class _Memory(ABC):
         Returns:
             The updated memory
         """
-        response = update_memory_gpu(
-            self.memoryset_id,
-            body=_parse_memory_update(
+        response = orca_api.PATCH(
+            "/gpu/memoryset/{name_or_id}/memory",
+            params={"name_or_id": self.memoryset_id},
+            json=_parse_memory_update(
                 {"memory_id": self.memory_id}
                 | ({"value": value} if value is not UNSET else {})
                 | ({"source_id": source_id} if source_id is not UNSET else {})
@@ -381,6 +325,25 @@ class _Memory(ABC):
         self.__dict__.update(self.__class__(self.memoryset_id, response).__dict__)
         return self
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert the memory to a dictionary
+        """
+        return {
+            "value": self.value,
+            "embedding": self.embedding,
+            "source_id": self.source_id,
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+            "metadata": self.metadata,
+            "metrics": self.metrics,
+            "memory_id": self.memory_id,
+            "memory_version": self.memory_version,
+            "feedback_metrics": self.feedback_metrics,
+            "lookup_count": self.lookup_count,
+            "memory_type": self.memory_type,
+        }
+
 
 class LabeledMemory(_Memory):
     """
@@ -409,7 +372,7 @@ class LabeledMemory(_Memory):
 
     label: int
     label_name: str | None
-    memory_type = MemoryType.LABELED
+    memory_type = "LABELED"
 
     def __init__(
         self,
@@ -423,8 +386,8 @@ class LabeledMemory(_Memory):
     ):
         # for internal use only, do not document
         super().__init__(memoryset_id, memory)
-        self.label = memory.label
-        self.label_name = memory.label_name
+        self.label = memory["label"]
+        self.label_name = memory["label_name"]
 
     def __repr__(self) -> str:
         return (
@@ -464,6 +427,15 @@ class LabeledMemory(_Memory):
         super().update(value=value, label=label, source_id=source_id, **metadata)
         return self
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert the memory to a dictionary
+        """
+        super_dict = super().to_dict()
+        super_dict["label"] = self.label
+        super_dict["label_name"] = self.label_name
+        return super_dict
+
 
 class LabeledMemoryLookup(LabeledMemory):
     """
@@ -499,10 +471,8 @@ class LabeledMemoryLookup(LabeledMemory):
     def __init__(self, memoryset_id: str, memory_lookup: LabeledMemoryLookupResponse | LabelPredictionMemoryLookup):
         # for internal use only, do not document
         super().__init__(memoryset_id, memory_lookup)
-        self.lookup_score = memory_lookup.lookup_score
-        self.attention_weight = (
-            memory_lookup.attention_weight if isinstance(memory_lookup, LabelPredictionMemoryLookup) else None
-        )
+        self.lookup_score = memory_lookup["lookup_score"]
+        self.attention_weight = memory_lookup["attention_weight"] if "attention_weight" in memory_lookup else None
 
     def __repr__(self) -> str:
         return (
@@ -540,7 +510,7 @@ class ScoredMemory(_Memory):
     """
 
     score: float
-    memory_type = MemoryType.SCORED
+    memory_type = "SCORED"
 
     def __init__(
         self,
@@ -554,7 +524,7 @@ class ScoredMemory(_Memory):
     ):
         # for internal use only, do not document
         super().__init__(memoryset_id, memory)
-        self.score = memory.score
+        self.score = memory["score"]
 
     def __repr__(self) -> str:
         return (
@@ -594,6 +564,14 @@ class ScoredMemory(_Memory):
         super().update(value=value, score=score, source_id=source_id, **metadata)
         return self
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert the memory to a dictionary
+        """
+        super_dict = super().to_dict()
+        super_dict["score"] = self.score
+        return super_dict
+
 
 class ScoredMemoryLookup(ScoredMemory):
     """
@@ -625,10 +603,8 @@ class ScoredMemoryLookup(ScoredMemory):
     def __init__(self, memoryset_id: str, memory_lookup: ScoredMemoryLookupResponse | ScorePredictionMemoryLookup):
         # for internal use only, do not document
         super().__init__(memoryset_id, memory_lookup)
-        self.lookup_score = memory_lookup.lookup_score
-        self.attention_weight = (
-            memory_lookup.attention_weight if isinstance(memory_lookup, ScorePredictionMemoryLookup) else None
-        )
+        self.lookup_score = memory_lookup["lookup_score"]
+        self.attention_weight = memory_lookup["attention_weight"] if "attention_weight" in memory_lookup else None
 
     def __repr__(self) -> str:
         return (
@@ -674,23 +650,23 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
 
     def __init__(self, metadata: MemorysetMetadata):
         # for internal use only, do not document
-        if metadata.pretrained_embedding_model_name:
-            self.embedding_model = PretrainedEmbeddingModel._get(metadata.pretrained_embedding_model_name)
-        elif metadata.finetuned_embedding_model_id:
-            self.embedding_model = FinetunedEmbeddingModel.open(metadata.finetuned_embedding_model_id)
+        if metadata["pretrained_embedding_model_name"]:
+            self.embedding_model = PretrainedEmbeddingModel._get(metadata["pretrained_embedding_model_name"])
+        elif metadata["finetuned_embedding_model_id"]:
+            self.embedding_model = FinetunedEmbeddingModel.open(metadata["finetuned_embedding_model_id"])
         else:
             raise ValueError("Either pretrained_embedding_model_name or finetuned_embedding_model_id must be provided")
-        self.id = metadata.id
-        self.name = metadata.name
-        self.description = metadata.description
-        self.length = metadata.length
-        self.created_at = metadata.created_at
-        self.updated_at = metadata.updated_at
-        self.insertion_status = Status(metadata.insertion_status.value)
+        self.id = metadata["id"]
+        self.name = metadata["name"]
+        self.description = metadata["description"]
+        self.length = metadata["length"]
+        self.created_at = datetime.fromisoformat(metadata["created_at"])
+        self.updated_at = datetime.fromisoformat(metadata["updated_at"])
+        self.insertion_status = Status(metadata["insertion_status"])
         self._last_refresh = datetime.now()
-        self.index_type = metadata.index_type.value
-        self.index_params = metadata.index_params.to_dict()
-        self.memory_type = MemoryType(metadata.memory_type.value)
+        self.index_type = metadata["index_type"]
+        self.index_params = metadata["index_params"]
+        self.memory_type = metadata["memory_type"]
 
     def __eq__(self, other) -> bool:
         return isinstance(other, _Memoryset) and self.id == other.id
@@ -719,6 +695,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 +719,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 +742,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 +776,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
@@ -829,29 +809,30 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                         raise ValueError(f"Memoryset with name {name} already exists with a different {attribute}.")
                 return existing
 
-        response = create_memoryset(
-            body=CreateMemorysetRequest(
-                name=name,
-                description=description,
-                datasource_id=datasource.id,
-                datasource_label_column=label_column,
-                datasource_score_column=score_column,
-                datasource_value_column=value_column,
-                datasource_source_id_column=source_id_column,
-                pretrained_embedding_model_name=(
-                    embedding_model._model_name if isinstance(embedding_model, PretrainedEmbeddingModel) else None
-                ),
-                finetuned_embedding_model_id=(
-                    embedding_model.id if isinstance(embedding_model, FinetunedEmbeddingModel) else None
-                ),
-                label_names=label_names,
-                max_seq_length_override=max_seq_length_override,
-                remove_duplicates=remove_duplicates,
-                index_type=CreateMemorysetRequestIndexType[index_type],
-                index_params=CreateMemorysetRequestIndexParams.from_dict(index_params),
-            ),
-        )
-        job = Job(response.insertion_task_id, lambda: cls.open(response.id))
+        payload: CreateMemorysetRequest = {
+            "name": name,
+            "description": description,
+            "datasource_name_or_id": datasource.id,
+            "datasource_label_column": label_column,
+            "datasource_score_column": score_column,
+            "datasource_value_column": value_column,
+            "datasource_source_id_column": source_id_column,
+            "label_names": label_names,
+            "max_seq_length_override": max_seq_length_override,
+            "remove_duplicates": remove_duplicates,
+            "index_type": index_type,
+            "index_params": index_params,
+        }
+        if prompt is not None:
+            payload["prompt"] = prompt
+        if isinstance(embedding_model, PretrainedEmbeddingModel):
+            payload["pretrained_embedding_model_name"] = embedding_model.name
+        elif isinstance(embedding_model, FinetunedEmbeddingModel):
+            payload["finetuned_embedding_model_name_or_id"] = embedding_model.id
+        else:
+            raise ValueError("Invalid embedding model")
+        response = orca_api.POST("/memoryset", json=payload)
+        job = Job(response["insertion_task_id"], lambda: cls.open(response["id"]))
         return job if background else job.result()
 
     @overload
@@ -1272,7 +1253,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         Raises:
             LookupError: If the memoryset does not exist
         """
-        metadata = get_memoryset(name)
+        metadata = orca_api.GET("/memoryset/{name_or_id}", params={"name_or_id": name})
         return cls(metadata)
 
     @classmethod
@@ -1300,7 +1281,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         Returns:
             List of handles to all memorysets in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_memorysets(type=cls.memory_type)]
+        return [cls(metadata) for metadata in orca_api.GET("/memoryset", params={"type": cls.memory_type})]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error"):
@@ -1316,7 +1297,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             LookupError: If the memoryset does not exist and if_not_exists is `"error"`
         """
         try:
-            delete_memoryset(name_or_id)
+            orca_api.DELETE("/memoryset/{name_or_id}", params={"name_or_id": name_or_id})
             logging.info(f"Deleted memoryset {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
@@ -1334,14 +1315,15 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             name: Value to set for the name
             label_names: Value to replace existing label names with
         """
-        update_memoryset(
-            self.id,
-            body=MemorysetUpdate(
-                name=name if name is not UNSET else CLIENT_UNSET,
-                description=description if description is not UNSET else CLIENT_UNSET,
-                label_names=label_names if label_names is not UNSET else CLIENT_UNSET,
-            ),
-        )
+        payload: MemorysetUpdate = {}
+        if name is not UNSET:
+            payload["name"] = name
+        if description is not UNSET:
+            payload["description"] = description
+        if label_names is not UNSET:
+            payload["label_names"] = label_names
+
+        orca_api.PATCH("/memoryset/{name_or_id}", params={"name_or_id": self.id}, json=payload)
         self.refresh()
 
     @overload
@@ -1351,6 +1333,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 +1346,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:
@@ -1373,7 +1357,8 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         name: str,
         *,
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
-        max_seq_length_override: int | None = None,
+        max_seq_length_override: int | None = UNSET,
+        prompt: str | None = None,
         if_exists: CreateMode = "error",
         background: bool = False,
     ) -> Self | Job[Self]:
@@ -1385,6 +1370,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 +1388,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":
@@ -1409,23 +1405,20 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                     if locals()[attribute] is not None and locals()[attribute] != getattr(existing, attribute):
                         raise ValueError(f"Memoryset with name {name} already exists with a different {attribute}.")
                 return existing
-
-        metadata = clone_memoryset(
-            self.id,
-            body=CloneMemorysetRequest(
-                name=name,
-                pretrained_embedding_model_name=(
-                    embedding_model._model_name if isinstance(embedding_model, PretrainedEmbeddingModel) else None
-                ),
-                finetuned_embedding_model_id=(
-                    embedding_model.id if isinstance(embedding_model, FinetunedEmbeddingModel) else None
-                ),
-                max_seq_length_override=max_seq_length_override,
-            ),
-        )
+        payload: CloneMemorysetRequest = {"name": name}
+        if max_seq_length_override is not UNSET:
+            payload["max_seq_length_override"] = max_seq_length_override
+        if prompt is not None:
+            payload["prompt"] = prompt
+        if isinstance(embedding_model, PretrainedEmbeddingModel):
+            payload["pretrained_embedding_model_name"] = embedding_model.name
+        elif isinstance(embedding_model, FinetunedEmbeddingModel):
+            payload["finetuned_embedding_model_name_or_id"] = embedding_model.id
+
+        metadata = orca_api.POST("/memoryset/{name_or_id}/clone", params={"name_or_id": self.id}, json=payload)
         job = Job(
-            metadata.insertion_task_id,
-            lambda: self.open(metadata.id),
+            metadata["insertion_task_id"],
+            lambda: self.open(metadata["id"]),
         )
         return job if background else job.result()
 
@@ -1503,20 +1496,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 +1528,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)
             [
@@ -1542,9 +1546,14 @@ 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),
+        response = orca_api.POST(
+            "/gpu/memoryset/{name_or_id}/lookup",
+            params={"name_or_id": self.id},
+            json={
+                "query": query if isinstance(query, list) else [query],
+                "count": count,
+                "prompt": prompt,
+            },
         )
         lookups = [
             [
@@ -1552,7 +1561,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                     MemoryLookupT,
                     (
                         LabeledMemoryLookup(self.id, lookup_response)
-                        if isinstance(lookup_response, LabeledMemoryLookupResponse)
+                        if "label" in lookup_response
                         else ScoredMemoryLookup(self.id, lookup_response)
                     ),
                 )
@@ -1594,50 +1603,70 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         ]
 
         if with_feedback_metrics:
+            response = orca_api.POST(
+                "/telemetry/memories",
+                json={
+                    "memoryset_id": self.id,
+                    "offset": offset,
+                    "limit": limit,
+                    "filters": parsed_filters,
+                    "sort": [_parse_sort_item_from_tuple(item) for item in sort] if sort else None,
+                },
+            )
             return [
                 cast(
                     MemoryT,
-                    (
-                        LabeledMemory(self.id, memory)
-                        if isinstance(memory, LabeledMemoryWithFeedbackMetrics)
-                        else ScoredMemory(self.id, memory)
-                    ),
+                    (LabeledMemory(self.id, memory) if "label" in memory else ScoredMemory(self.id, memory)),
                 )
-                for memory in list_memories_with_feedback(
-                    body=TelemetryMemoriesRequest(
-                        memoryset_id=self.id,
-                        offset=offset,
-                        limit=limit,
-                        filters=parsed_filters,
-                        sort=[_parse_sort_item_from_tuple(item) for item in sort] if sort else None,
-                    ),
-                ).items
+                for memory in response["items"]
             ]
 
-        if any(isinstance(filter, TelemetryFilterItem) for filter in parsed_filters):
+        if any(_is_metric_column(filter[0]) for filter in filters):
             raise ValueError("Feedback metrics are only supported when the with_feedback_metrics flag is set to True")
 
         if sort:
             logging.warning("Sorting is not supported when with_feedback_metrics is False. Sort value will be ignored.")
 
+        response = orca_api.POST(
+            "/memoryset/{name_or_id}/memories",
+            params={"name_or_id": self.id},
+            json={
+                "offset": offset,
+                "limit": limit,
+                "filters": cast(list[FilterItem], parsed_filters),
+            },
+        )
         return [
             cast(
                 MemoryT,
-                (
-                    LabeledMemory(self.id, memory)
-                    if isinstance(memory, LabeledMemoryResponse)
-                    else ScoredMemory(self.id, memory)
-                ),
+                (LabeledMemory(self.id, memory) if "label" in memory else ScoredMemory(self.id, memory)),
             )
-            for memory in query_memoryset(
-                self.id,
-                body=ListMemoriesRequest(
+            for memory in response
+        ]
+
+    def to_pandas(
+        self,
+        offset: int = 0,
+        limit: int = 100,
+        filters: list[FilterItemTuple] = [],
+        with_feedback_metrics: bool = False,
+        sort: list[TelemetrySortItem] | None = None,
+    ) -> pd.DataFrame:
+        """
+        Convert the memoryset to a pandas DataFrame
+        """
+        return pd.DataFrame(
+            [
+                memory.to_dict()
+                for memory in self.query(
                     offset=offset,
                     limit=limit,
-                    filters=cast(list[FilterItem], parsed_filters),
-                ),
-            )
-        ]
+                    filters=filters,
+                    with_feedback_metrics=with_feedback_metrics,
+                    sort=sort,
+                )
+            ]
+        )
 
     def insert(self, items: Iterable[dict[str, Any]] | dict[str, Any]) -> None:
         """
@@ -1659,16 +1688,15 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             ...     {"value": "I am sad", "label": 0, "source_id": "user_124", "tag": "sad"},
             ... ])
         """
-        insert_memories_gpu(
-            self.id,
-            body=(
-                cast(
-                    list[LabeledMemoryInsert] | list[ScoredMemoryInsert],
-                    [
-                        _parse_memory_insert(memory, type=self.memory_type)
-                        for memory in (cast(list[dict[str, Any]], [items]) if isinstance(items, dict) else items)
-                    ],
-                )
+        orca_api.POST(
+            "/gpu/memoryset/{name_or_id}/memory",
+            params={"name_or_id": self.id},
+            json=cast(
+                list[LabeledMemoryInsert] | list[ScoredMemoryInsert],
+                [
+                    _parse_memory_insert(memory, type=self.memory_type)
+                    for memory in (cast(list[dict[str, Any]], [items]) if isinstance(items, dict) else items)
+                ],
             ),
         )
         self.refresh()
@@ -1710,25 +1738,23 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             ]
         """
         if isinstance(memory_id, str):
-            response = get_memory(self.id, memory_id)
+            response = orca_api.GET(
+                "/memoryset/{name_or_id}/memory/{memory_id}", params={"name_or_id": self.id, "memory_id": memory_id}
+            )
             return cast(
                 MemoryT,
-                (
-                    LabeledMemory(self.id, response)
-                    if isinstance(response, LabeledMemoryResponse)
-                    else ScoredMemory(self.id, response)
-                ),
+                (LabeledMemory(self.id, response) if "label" in response else ScoredMemory(self.id, response)),
             )
         else:
-            response = get_memories(self.id, body=GetMemoriesRequest(memory_ids=list(memory_id)))
+            response = orca_api.POST(
+                "/memoryset/{name_or_id}/memories/get",
+                params={"name_or_id": self.id},
+                json={"memory_ids": list(memory_id)},
+            )
             return [
                 cast(
                     MemoryT,
-                    (
-                        LabeledMemory(self.id, memory)
-                        if isinstance(memory, LabeledMemoryResponse)
-                        else ScoredMemory(self.id, memory)
-                    ),
+                    (LabeledMemory(self.id, memory) if "label" in memory else ScoredMemory(self.id, memory)),
                 )
                 for memory in response
             ]
@@ -1773,9 +1799,10 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             ...     for m in memoryset.query(filters=[("tag", "==", "happy")])
             ... )
         """
-        response = update_memories_gpu(
-            self.id,
-            body=cast(
+        response = orca_api.PATCH(
+            "/gpu/memoryset/{name_or_id}/memories",
+            params={"name_or_id": self.id},
+            json=cast(
                 list[LabeledMemoryUpdate] | list[ScoredMemoryUpdate],
                 [
                     _parse_memory_update(update, type=self.memory_type)
@@ -1786,11 +1813,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         updated_memories = [
             cast(
                 MemoryT,
-                (
-                    LabeledMemory(self.id, memory)
-                    if isinstance(memory, LabeledMemoryResponse)
-                    else ScoredMemory(self.id, memory)
-                ),
+                (LabeledMemory(self.id, memory) if "label" in memory else ScoredMemory(self.id, memory)),
             )
             for memory in response
         ]
@@ -1836,21 +1859,21 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         Returns:
             A list of CascadingEditSuggestion objects, each containing a neighbor and the suggested new label.
         """
-
-        return suggest_cascading_edits(
-            name_or_id=self.id,
-            memory_id=memory.memory_id,
-            body=CascadeEditSuggestionsRequest(
-                old_label=old_label,
-                new_label=new_label,
-                max_neighbors=max_neighbors,
-                max_validation_neighbors=max_validation_neighbors,
-                similarity_threshold=similarity_threshold,
-                only_if_has_old_label=only_if_has_old_label,
-                exclude_if_new_label=exclude_if_new_label,
-                suggestion_cooldown_time=suggestion_cooldown_time,
-                label_confirmation_cooldown_time=label_confirmation_cooldown_time,
-            ),
+        # TODO: properly integrate this with memory edits and return something that can be applied
+        return orca_api.POST(
+            "/memoryset/{name_or_id}/memory/{memory_id}/cascading_edits",
+            params={"name_or_id": self.id, "memory_id": memory.memory_id},
+            json={
+                "old_label": old_label,
+                "new_label": new_label,
+                "max_neighbors": max_neighbors,
+                "max_validation_neighbors": max_validation_neighbors,
+                "similarity_threshold": similarity_threshold,
+                "only_if_has_old_label": only_if_has_old_label,
+                "exclude_if_new_label": exclude_if_new_label,
+                "suggestion_cooldown_time": suggestion_cooldown_time,
+                "label_confirmation_cooldown_time": label_confirmation_cooldown_time,
+            },
         )
 
     def delete(self, memory_id: str | Iterable[str]) -> None:
@@ -1872,37 +1895,39 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
 
         """
         memory_ids = [memory_id] if isinstance(memory_id, str) else list(memory_id)
-        delete_memories(self.id, body=DeleteMemoriesRequest(memory_ids=memory_ids))
+        orca_api.POST(
+            "/memoryset/{name_or_id}/memories/delete", params={"name_or_id": self.id}, json={"memory_ids": memory_ids}
+        )
         logging.info(f"Deleted {len(memory_ids)} memories from memoryset.")
         self.refresh()
 
     @overload
     def analyze(
         self,
-        *analyses: Iterable[dict[str, Any] | str],
+        *analyses: dict[str, Any] | str,
         lookup_count: int = 15,
         clear_metrics: bool = False,
         background: Literal[True],
-    ) -> Job[dict]:
+    ) -> Job[MemorysetMetrics]:
         pass
 
     @overload
     def analyze(
         self,
-        *analyses: Iterable[dict[str, Any] | str],
+        *analyses: dict[str, Any] | str,
         lookup_count: int = 15,
         clear_metrics: bool = False,
         background: Literal[False] = False,
-    ) -> dict:
+    ) -> MemorysetMetrics:
         pass
 
     def analyze(
         self,
-        *analyses: Iterable[dict[str, Any] | str],
+        *analyses: dict[str, Any] | str,
         lookup_count: int = 15,
         clear_metrics: bool = False,
         background: bool = False,
-    ) -> Job[dict] | dict:
+    ) -> Job[MemorysetMetrics] | MemorysetMetrics:
         """
         Run analyses on the memoryset to find duplicates, clusters, mislabelings, and more
 
@@ -1965,11 +1990,9 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         """
 
         # Get valid analysis names from MemorysetAnalysisConfigs
-        valid_analysis_names = {
-            field.name for field in fields(MemorysetAnalysisConfigs) if field.name != "additional_properties"
-        }
+        valid_analysis_names = set(MemorysetAnalysisConfigs.__annotations__)
 
-        configs: dict[str, dict] = {}
+        configs: MemorysetAnalysisConfigs = {}
         for analysis in analyses:
             if isinstance(analysis, str):
                 error_msg = (
@@ -1979,30 +2002,37 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
                     raise ValueError(error_msg)
                 configs[analysis] = {}
             else:
-                name = analysis.pop("name")  # type: ignore
+                name = analysis.pop("name")
                 error_msg = f"Invalid analysis name: {name}. Valid names are: {', '.join(sorted(valid_analysis_names))}"
                 if name not in valid_analysis_names:
                     raise ValueError(error_msg)
-                configs[name] = analysis  # type: ignore
-
-        analysis = analyze_memoryset(
-            self.id,
-            body=MemorysetAnalysisRequest(
-                configs=MemorysetAnalysisConfigs.from_dict(configs),
-                lookup_count=lookup_count,
-                clear_metrics=clear_metrics,
-            ),
+                configs[name] = analysis
+
+        analysis = orca_api.POST(
+            "/memoryset/{name_or_id}/analysis",
+            params={"name_or_id": self.id},
+            json={
+                "configs": configs,
+                "lookup_count": lookup_count,
+                "clear_metrics": clear_metrics,
+            },
         )
         job = Job(
-            analysis.task_id,
-            lambda: (r := get_analysis(self.id, analysis.task_id).results) and r.to_dict(),
+            analysis["task_id"],
+            lambda: orca_api.GET(
+                "/memoryset/{name_or_id}/analysis/{analysis_task_id}",
+                params={"name_or_id": self.id, "analysis_task_id": analysis["task_id"]},
+            )["results"],
         )
         return job if background else job.result()
 
-    def get_potential_duplicate_groups(self):
+    def get_potential_duplicate_groups(self) -> list[list[MemoryT]]:
         """Group potential duplicates in the memoryset"""
-        response = potential_duplicate_groups(self.name)
-        return response
+        response = orca_api.GET("/memoryset/{name_or_id}/potential_duplicate_groups", params={"name_or_id": self.id})
+        return [
+            [cast(MemoryT, LabeledMemory(self.id, m) if "label" in m else ScoredMemory(self.id, m)) for m in ms]
+            for ms in response
+        ]
 
     @overload
     @staticmethod
@@ -2015,7 +2045,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         neighbor_count: int = 5,
         embedding_models: list[str] | None = None,
         background: Literal[True],
-    ) -> Job[dict]:
+    ) -> Job[list[EmbeddingModelResult]]:
         pass
 
     @overload
@@ -2029,7 +2059,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         neighbor_count: int = 5,
         embedding_models: list[str] | None = None,
         background: Literal[False] = False,
-    ) -> dict:
+    ) -> list[EmbeddingModelResult]:
         pass
 
     @staticmethod
@@ -2042,7 +2072,7 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
         neighbor_count: int = 5,
         embedding_models: list[str] | None = None,
         background: bool = False,
-    ) -> Job[dict] | dict:
+    ) -> Job[list[EmbeddingModelResult]] | list[EmbeddingModelResult]:
         """
         Test the quality of embeddings for the datasource by computing metrics such as prediction accuracy.
 
@@ -2061,24 +2091,27 @@ class _Memoryset(Generic[MemoryT, MemoryLookupT], ABC):
             A dictionary containing the results of the embedding evaluation
         """
 
-        if embedding_models is not None:
-            embedding_model_enums = [PretrainedEmbeddingModelName(model) for model in embedding_models]
-        else:
-            embedding_model_enums = None
-
-        request = EmbeddingEvaluationRequest(
-            value_column=value_column,
-            label_column=label_column,
-            source_id_column=source_id_column,
-            neighbor_count=neighbor_count,
-            embedding_models=embedding_model_enums,
+        response = orca_api.POST(
+            "/datasource/{name_or_id}/embedding_evaluation",
+            params={"name_or_id": datasource.id},
+            json={
+                "value_column": value_column,
+                "label_column": label_column,
+                "source_id_column": source_id_column,
+                "neighbor_count": neighbor_count,
+                "embedding_models": embedding_models,
+            },
         )
 
-        response = create_embedding_evaluation(name_or_id=datasource.id, body=request)
-        job = Job(
-            response.task_id,
-            lambda: (r := get_embedding_evaluation(datasource.id, response.task_id).result) and r.to_dict(),
-        )
+        def get_value() -> list[EmbeddingModelResult]:
+            res = orca_api.GET(
+                "/datasource/{name_or_id}/embedding_evaluation/{task_id}",
+                params={"name_or_id": datasource.id, "task_id": response["task_id"]},
+            )
+            assert res["result"] is not None
+            return res["result"]["evaluation_results"]
+
+        job = Job(response["task_id"], get_value)
         return job if background else job.result()
 
 
@@ -2098,12 +2131,12 @@ class LabeledMemoryset(_Memoryset[LabeledMemory, LabeledMemoryLookup]):
     """
 
     label_names: list[str]
-    memory_type: MemoryType = MemoryType.LABELED
+    memory_type: MemoryType = "LABELED"
 
     def __init__(self, metadata: MemorysetMetadata):
         super().__init__(metadata)
-        assert metadata.label_names is not None
-        self.label_names = metadata.label_names
+        assert metadata["label_names"] is not None
+        self.label_names = metadata["label_names"]
 
     def __eq__(self, other) -> bool:
         return isinstance(other, LabeledMemoryset) and self.id == other.id
@@ -2112,12 +2145,6 @@ class LabeledMemoryset(_Memoryset[LabeledMemory, LabeledMemoryLookup]):
     def create(cls, name: str, datasource: Datasource, *, label_column: str | None = "label", **kwargs):
         return super().create(name, datasource, label_column=label_column, score_column=None, **kwargs)
 
-    def display_label_analysis(self):
-        """Display a UI to review and act upon the label analysis results"""
-        from ._utils.analysis_ui import display_suggested_memory_relabels
-
-        display_suggested_memory_relabels(self)
-
 
 class ScoredMemoryset(_Memoryset[ScoredMemory, ScoredMemoryLookup]):
     """
@@ -2133,7 +2160,7 @@ class ScoredMemoryset(_Memoryset[ScoredMemory, ScoredMemoryLookup]):
         updated_at: When the memoryset was last updated, automatically updated on updates
     """
 
-    memory_type: MemoryType = MemoryType.SCORED
+    memory_type: MemoryType = "SCORED"
 
     def __eq__(self, other) -> bool:
         return isinstance(other, ScoredMemoryset) and self.id == other.id