add update and delete for memory banks

This commit is contained in:
Dinesh Yeduguru 2024-11-14 14:44:37 -08:00
parent bba6edd06b
commit 9b75e92852
4 changed files with 240 additions and 3 deletions

View file

@ -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",

View file

@ -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

View file

@ -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: ...

View file

@ -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]: