add update and delete for memory banks

docs/resources/llama-stack-spec.html

@@ -21,7 +21,7 @@
     "info": {
         "title": "[DRAFT] Llama Stack Specification",
         "version": "0.0.1",
-        "description": "This is the specification of the llama stack that provides\n                a set of endpoints and their corresponding interfaces that are tailored to\n                best leverage Llama Models. The specification is still in draft and subject to change.\n                Generated at 2024-11-14 12:51:12.176325"
+        "description": "This is the specification of the llama stack that provides\n                a set of endpoints and their corresponding interfaces that are tailored to\n                best leverage Llama Models. The specification is still in draft and subject to change.\n                Generated at 2024-11-14 14:47:24.435917"
     },
     "servers": [
         {
@@ -429,6 +429,39 @@
                 }
             }
         },
+        "/memory_banks/delete": {
+            "post": {
+                "responses": {
+                    "200": {
+                        "description": "OK"
+                    }
+                },
+                "tags": [
+                    "MemoryBanks"
+                ],
+                "parameters": [
+                    {
+                        "name": "X-LlamaStack-ProviderData",
+                        "in": "header",
+                        "description": "JSON-encoded provider data which will be made available to the adapter servicing the API",
+                        "required": false,
+                        "schema": {
+                            "type": "string"
+                        }
+                    }
+                ],
+                "requestBody": {
+                    "content": {
+                        "application/json": {
+                            "schema": {
+                                "$ref": "#/components/schemas/DeleteMemoryBankRequest"
+                            }
+                        }
+                    },
+                    "required": true
+                }
+            }
+        },
         "/models/delete": {
             "post": {
                 "responses": {
@@ -2259,6 +2292,35 @@
                 }
             }
         },
+        "/memory_banks/update": {
+            "post": {
+                "responses": {},
+                "tags": [
+                    "MemoryBanks"
+                ],
+                "parameters": [
+                    {
+                        "name": "X-LlamaStack-ProviderData",
+                        "in": "header",
+                        "description": "JSON-encoded provider data which will be made available to the adapter servicing the API",
+                        "required": false,
+                        "schema": {
+                            "type": "string"
+                        }
+                    }
+                ],
+                "requestBody": {
+                    "content": {
+                        "application/json": {
+                            "schema": {
+                                "$ref": "#/components/schemas/UpdateMemoryBankRequest"
+                            }
+                        }
+                    },
+                    "required": true
+                }
+            }
+        },
         "/models/update": {
             "post": {
                 "responses": {
@@ -4622,6 +4684,18 @@
                     "session_id"
                 ]
             },
+            "DeleteMemoryBankRequest": {
+                "type": "object",
+                "properties": {
+                    "memory_bank_id": {
+                        "type": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "memory_bank_id"
+                ]
+            },
             "DeleteModelRequest": {
                 "type": "object",
                 "properties": {
@@ -7912,6 +7986,41 @@
                 ],
                 "title": "Response from the synthetic data generation. Batch of (prompt, response, score) tuples that pass the threshold."
             },
+            "UpdateMemoryBankRequest": {
+                "type": "object",
+                "properties": {
+                    "memory_bank_id": {
+                        "type": "string"
+                    },
+                    "params": {
+                        "oneOf": [
+                            {
+                                "$ref": "#/components/schemas/VectorMemoryBankParams"
+                            },
+                            {
+                                "$ref": "#/components/schemas/KeyValueMemoryBankParams"
+                            },
+                            {
+                                "$ref": "#/components/schemas/KeywordMemoryBankParams"
+                            },
+                            {
+                                "$ref": "#/components/schemas/GraphMemoryBankParams"
+                            }
+                        ]
+                    },
+                    "provider_id": {
+                        "type": "string"
+                    },
+                    "provider_memory_bank_id": {
+                        "type": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "memory_bank_id",
+                    "params"
+                ]
+            },
             "UpdateModelRequest": {
                 "type": "object",
                 "properties": {
@@ -8132,6 +8241,10 @@
             "name": "DeleteAgentsSessionRequest",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/DeleteAgentsSessionRequest\" />"
         },
+        {
+            "name": "DeleteMemoryBankRequest",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/DeleteMemoryBankRequest\" />"
+        },
         {
             "name": "DeleteModelRequest",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/DeleteModelRequest\" />"
@@ -8566,6 +8679,10 @@
             "name": "UnstructuredLogEvent",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/UnstructuredLogEvent\" />"
         },
+        {
+            "name": "UpdateMemoryBankRequest",
+            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/UpdateMemoryBankRequest\" />"
+        },
         {
             "name": "UpdateModelRequest",
             "description": "<SchemaDefinition schemaRef=\"#/components/schemas/UpdateModelRequest\" />"
@@ -8657,6 +8774,7 @@
                 "Dataset",
                 "DeleteAgentsRequest",
                 "DeleteAgentsSessionRequest",
+                "DeleteMemoryBankRequest",
                 "DeleteModelRequest",
                 "DoraFinetuningConfig",
                 "EmbeddingsRequest",
@@ -8755,6 +8873,7 @@
                 "Turn",
                 "URL",
                 "UnstructuredLogEvent",
+                "UpdateMemoryBankRequest",
                 "UpdateModelRequest",
                 "UserMessage",
                 "VectorMemoryBank",

docs/resources/llama-stack-spec.yaml

@@ -867,6 +867,14 @@ components:
       - agent_id
       - session_id
       type: object
+    DeleteMemoryBankRequest:
+      additionalProperties: false
+      properties:
+        memory_bank_id:
+          type: string
+      required:
+      - memory_bank_id
+      type: object
     DeleteModelRequest:
       additionalProperties: false
       properties:
@@ -3280,6 +3288,25 @@ components:
       - message
       - severity
       type: object
+    UpdateMemoryBankRequest:
+      additionalProperties: false
+      properties:
+        memory_bank_id:
+          type: string
+        params:
+          oneOf:
+          - $ref: '#/components/schemas/VectorMemoryBankParams'
+          - $ref: '#/components/schemas/KeyValueMemoryBankParams'
+          - $ref: '#/components/schemas/KeywordMemoryBankParams'
+          - $ref: '#/components/schemas/GraphMemoryBankParams'
+        provider_id:
+          type: string
+        provider_memory_bank_id:
+          type: string
+      required:
+      - memory_bank_id
+      - params
+      type: object
     UpdateModelRequest:
       additionalProperties: false
       properties:
@@ -3414,7 +3441,7 @@ info:
   description: "This is the specification of the llama stack that provides\n     \
     \           a set of endpoints and their corresponding interfaces that are tailored\
     \ to\n                best leverage Llama Models. The specification is still in\
-    \ draft and subject to change.\n                Generated at 2024-11-14 12:51:12.176325"
+    \ draft and subject to change.\n                Generated at 2024-11-14 14:47:24.435917"
   title: '[DRAFT] Llama Stack Specification'
   version: 0.0.1
 jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema
@@ -4144,6 +4171,27 @@ paths:
           description: OK
       tags:
       - Memory
+  /memory_banks/delete:
+    post:
+      parameters:
+      - description: JSON-encoded provider data which will be made available to the
+          adapter servicing the API
+        in: header
+        name: X-LlamaStack-ProviderData
+        required: false
+        schema:
+          type: string
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/DeleteMemoryBankRequest'
+        required: true
+      responses:
+        '200':
+          description: OK
+      tags:
+      - MemoryBanks
   /memory_banks/get:
     get:
       parameters:
@@ -4216,6 +4264,25 @@ paths:
       responses: {}
       tags:
       - MemoryBanks
+  /memory_banks/update:
+    post:
+      parameters:
+      - description: JSON-encoded provider data which will be made available to the
+          adapter servicing the API
+        in: header
+        name: X-LlamaStack-ProviderData
+        required: false
+        schema:
+          type: string
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/UpdateMemoryBankRequest'
+        required: true
+      responses: {}
+      tags:
+      - MemoryBanks
   /models/delete:
     post:
       parameters:
@@ -4960,6 +5027,9 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/DeleteAgentsSessionRequest"
     />
   name: DeleteAgentsSessionRequest
+- description: <SchemaDefinition schemaRef="#/components/schemas/DeleteMemoryBankRequest"
+    />
+  name: DeleteMemoryBankRequest
 - description: <SchemaDefinition schemaRef="#/components/schemas/DeleteModelRequest"
     />
   name: DeleteModelRequest
@@ -5260,6 +5330,9 @@ tags:
 - description: <SchemaDefinition schemaRef="#/components/schemas/UnstructuredLogEvent"
     />
   name: UnstructuredLogEvent
+- description: <SchemaDefinition schemaRef="#/components/schemas/UpdateMemoryBankRequest"
+    />
+  name: UpdateMemoryBankRequest
 - description: <SchemaDefinition schemaRef="#/components/schemas/UpdateModelRequest"
     />
   name: UpdateModelRequest
@@ -5338,6 +5411,7 @@ x-tagGroups:
   - Dataset
   - DeleteAgentsRequest
   - DeleteAgentsSessionRequest
+  - DeleteMemoryBankRequest
   - DeleteModelRequest
   - DoraFinetuningConfig
   - EmbeddingsRequest
@@ -5436,6 +5510,7 @@ x-tagGroups:
   - Turn
   - URL
   - UnstructuredLogEvent
+  - UpdateMemoryBankRequest
   - UpdateModelRequest
   - UserMessage
   - VectorMemoryBank

llama_stack/apis/memory_banks/memory_banks.py

@@ -144,3 +144,15 @@ class MemoryBanks(Protocol):
         provider_id: Optional[str] = None,
         provider_memory_bank_id: Optional[str] = None,
     ) -> MemoryBank: ...
+
+    @webmethod(route="/memory_banks/update", method="POST")
+    async def update_memory_bank(
+        self,
+        memory_bank_id: str,
+        params: BankParams,
+        provider_id: Optional[str] = None,
+        provider_memory_bank_id: Optional[str] = None,
+    ) -> MemoryBank: ...
+
+    @webmethod(route="/memory_banks/delete", method="POST")
+    async def delete_memory_bank(self, memory_bank_id: str) -> None: ...

llama_stack/distribution/routers/routing_tables.py

@@ -158,7 +158,7 @@ class CommonRoutingTableImpl(RoutingTable):
         registered_obj = await register_object_with_provider(
             obj, self.impls_by_provider_id[obj.provider_id]
         )
-        return await self.dist_registry.update(registered_obj)
+        return await self.dist_registry.update(registered_obj or obj)
 
     async def register_object(
         self, obj: RoutableObjectWithProvider
@@ -333,6 +333,37 @@ class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
         await self.register_object(memory_bank)
         return memory_bank
 
+    async def update_memory_bank(
+        self,
+        memory_bank_id: str,
+        params: BankParams,
+        provider_id: Optional[str] = None,
+        provider_memory_bank_id: Optional[str] = None,
+    ) -> MemoryBank:
+        existing_bank = await self.get_memory_bank(memory_bank_id)
+        if existing_bank is None:
+            raise ValueError(f"Memory bank {memory_bank_id} not found")
+
+        updated_bank = parse_obj_as(
+            MemoryBank,
+            {
+                "identifier": memory_bank_id,
+                "type": ResourceType.memory_bank.value,
+                "provider_id": provider_id or existing_bank.provider_id,
+                "provider_resource_id": provider_memory_bank_id
+                or existing_bank.provider_resource_id,
+                **params.model_dump(),
+            },
+        )
+        registered_bank = await self.update_object(updated_bank)
+        return registered_bank
+
+    async def delete_memory_bank(self, memory_bank_id: str) -> None:
+        existing_bank = await self.get_memory_bank(memory_bank_id)
+        if existing_bank is None:
+            raise ValueError(f"Memory bank {memory_bank_id} not found")
+        await self.delete_object(existing_bank)
+
 
 class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
     async def list_datasets(self) -> List[Dataset]: