Redo the { models, shields, memory_banks } typeset

llama_stack/apis/memory/memory.py

@@ -13,7 +13,6 @@ from typing import List, Optional, Protocol
 from llama_models.schema_utils import json_schema_type, webmethod
 
 from pydantic import BaseModel, Field
-from typing_extensions import Annotated
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 
@@ -26,44 +25,6 @@ class MemoryBankDocument(BaseModel):
     metadata: Dict[str, Any] = Field(default_factory=dict)
 
 
-@json_schema_type
-class MemoryBankType(Enum):
-    vector = "vector"
-    keyvalue = "keyvalue"
-    keyword = "keyword"
-    graph = "graph"
-
-
-class VectorMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.vector.value] = MemoryBankType.vector.value
-    embedding_model: str
-    chunk_size_in_tokens: int
-    overlap_size_in_tokens: Optional[int] = None
-
-
-class KeyValueMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.keyvalue.value] = MemoryBankType.keyvalue.value
-
-
-class KeywordMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.keyword.value] = MemoryBankType.keyword.value
-
-
-class GraphMemoryBankConfig(BaseModel):
-    type: Literal[MemoryBankType.graph.value] = MemoryBankType.graph.value
-
-
-MemoryBankConfig = Annotated[
-    Union[
-        VectorMemoryBankConfig,
-        KeyValueMemoryBankConfig,
-        KeywordMemoryBankConfig,
-        GraphMemoryBankConfig,
-    ],
-    Field(discriminator="type"),
-]
-
-
 class Chunk(BaseModel):
     content: InterleavedTextMedia
     token_count: int
@@ -76,46 +37,7 @@ class QueryDocumentsResponse(BaseModel):
     scores: List[float]
 
 
-@json_schema_type
-class QueryAPI(Protocol):
-    @webmethod(route="/query_documents")
-    def query_documents(
-        self,
-        query: InterleavedTextMedia,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse: ...
-
-
-@json_schema_type
-class MemoryBank(BaseModel):
-    bank_id: str
-    name: str
-    config: MemoryBankConfig
-    # if there's a pre-existing (reachable-from-distribution) store which supports QueryAPI
-    url: Optional[URL] = None
-
-
 class Memory(Protocol):
-    @webmethod(route="/memory/create")
-    async def create_memory_bank(
-        self,
-        name: str,
-        config: MemoryBankConfig,
-        url: Optional[URL] = None,
-    ) -> MemoryBank: ...
-
-    @webmethod(route="/memory/list", method="GET")
-    async def list_memory_banks(self) -> List[MemoryBank]: ...
-
-    @webmethod(route="/memory/get", method="GET")
-    async def get_memory_bank(self, bank_id: str) -> Optional[MemoryBank]: ...
-
-    @webmethod(route="/memory/drop", method="DELETE")
-    async def drop_memory_bank(
-        self,
-        bank_id: str,
-    ) -> str: ...
-
     # this will just block now until documents are inserted, but it should
     # probably return a Job instance which can be polled for completion
     @webmethod(route="/memory/insert")