diff --git a/docs/resources/llama-stack-spec.html b/docs/resources/llama-stack-spec.html index a0b4bccca..e487daac3 100644 --- a/docs/resources/llama-stack-spec.html +++ b/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": "" }, + { + "name": "DeleteMemoryBankRequest", + "description": "" + }, { "name": "DeleteModelRequest", "description": "" @@ -8566,6 +8679,10 @@ "name": "UnstructuredLogEvent", "description": "" }, + { + "name": "UpdateMemoryBankRequest", + "description": "" + }, { "name": "UpdateModelRequest", "description": "" @@ -8657,6 +8774,7 @@ "Dataset", "DeleteAgentsRequest", "DeleteAgentsSessionRequest", + "DeleteMemoryBankRequest", "DeleteModelRequest", "DoraFinetuningConfig", "EmbeddingsRequest", @@ -8755,6 +8873,7 @@ "Turn", "URL", "UnstructuredLogEvent", + "UpdateMemoryBankRequest", "UpdateModelRequest", "UserMessage", "VectorMemoryBank", diff --git a/docs/resources/llama-stack-spec.yaml b/docs/resources/llama-stack-spec.yaml index 2ca26f759..f29945396 100644 --- a/docs/resources/llama-stack-spec.yaml +++ b/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: name: DeleteAgentsSessionRequest +- description: + name: DeleteMemoryBankRequest - description: name: DeleteModelRequest @@ -5260,6 +5330,9 @@ tags: - description: name: UnstructuredLogEvent +- description: + name: UpdateMemoryBankRequest - description: 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 diff --git a/llama_stack/apis/memory_banks/memory_banks.py b/llama_stack/apis/memory_banks/memory_banks.py index c0a0c643a..a754a0818 100644 --- a/llama_stack/apis/memory_banks/memory_banks.py +++ b/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: ... diff --git a/llama_stack/distribution/routers/routing_tables.py b/llama_stack/distribution/routers/routing_tables.py index a940dbae6..b0082db18 100644 --- a/llama_stack/distribution/routers/routing_tables.py +++ b/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]: