From 06abd7e6c880d3504765d13b0082a7a022b07a31 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Fri, 20 Sep 2024 17:51:53 -0700 Subject: [PATCH] update MemoryToolDefinition --- docs/llama-stack-spec.html | 394 +++++++++++++++--------------- docs/llama-stack-spec.yaml | 256 +++++++++---------- llama_stack/apis/agents/agents.py | 1 + 3 files changed, 333 insertions(+), 318 deletions(-) diff --git a/docs/llama-stack-spec.html b/docs/llama-stack-spec.html index a7ab57343..bc6a7d37f 100644 --- a/docs/llama-stack-spec.html +++ b/docs/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-09-18 19:27:39.955190" + "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-09-20 17:50:36.257743" }, "servers": [ { @@ -2147,183 +2147,7 @@ "$ref": "#/components/schemas/FunctionCallToolDefinition" }, { - "type": "object", - "properties": { - "input_shields": { - "type": "array", - "items": { - "$ref": "#/components/schemas/ShieldDefinition" - } - }, - "output_shields": { - "type": "array", - "items": { - "$ref": "#/components/schemas/ShieldDefinition" - } - }, - "type": { - "type": "string", - "const": "memory" - }, - "memory_bank_configs": { - "type": "array", - "items": { - "oneOf": [ - { - "type": "object", - "properties": { - "bank_id": { - "type": "string" - }, - "type": { - "type": "string", - "const": "vector" - } - }, - "additionalProperties": false, - "required": [ - "bank_id", - "type" - ] - }, - { - "type": "object", - "properties": { - "bank_id": { - "type": "string" - }, - "type": { - "type": "string", - "const": "keyvalue" - }, - "keys": { - "type": "array", - "items": { - "type": "string" - } - } - }, - "additionalProperties": false, - "required": [ - "bank_id", - "type", - "keys" - ] - }, - { - "type": "object", - "properties": { - "bank_id": { - "type": "string" - }, - "type": { - "type": "string", - "const": "keyword" - } - }, - "additionalProperties": false, - "required": [ - "bank_id", - "type" - ] - }, - { - "type": "object", - "properties": { - "bank_id": { - "type": "string" - }, - "type": { - "type": "string", - "const": "graph" - }, - "entities": { - "type": "array", - "items": { - "type": "string" - } - } - }, - "additionalProperties": false, - "required": [ - "bank_id", - "type", - "entities" - ] - } - ] - } - }, - "query_generator_config": { - "oneOf": [ - { - "type": "object", - "properties": { - "type": { - "type": "string", - "const": "default" - }, - "sep": { - "type": "string" - } - }, - "additionalProperties": false, - "required": [ - "type", - "sep" - ] - }, - { - "type": "object", - "properties": { - "type": { - "type": "string", - "const": "llm" - }, - "model": { - "type": "string" - }, - "template": { - "type": "string" - } - }, - "additionalProperties": false, - "required": [ - "type", - "model", - "template" - ] - }, - { - "type": "object", - "properties": { - "type": { - "type": "string", - "const": "custom" - } - }, - "additionalProperties": false, - "required": [ - "type" - ] - } - ] - }, - "max_tokens_in_context": { - "type": "integer" - }, - "max_chunks": { - "type": "integer" - } - }, - "additionalProperties": false, - "required": [ - "type", - "memory_bank_configs", - "query_generator_config", - "max_tokens_in_context", - "max_chunks" - ] + "$ref": "#/components/schemas/MemoryToolDefinition" } ] } @@ -2432,6 +2256,185 @@ "parameters" ] }, + "MemoryToolDefinition": { + "type": "object", + "properties": { + "input_shields": { + "type": "array", + "items": { + "$ref": "#/components/schemas/ShieldDefinition" + } + }, + "output_shields": { + "type": "array", + "items": { + "$ref": "#/components/schemas/ShieldDefinition" + } + }, + "type": { + "type": "string", + "const": "memory" + }, + "memory_bank_configs": { + "type": "array", + "items": { + "oneOf": [ + { + "type": "object", + "properties": { + "bank_id": { + "type": "string" + }, + "type": { + "type": "string", + "const": "vector" + } + }, + "additionalProperties": false, + "required": [ + "bank_id", + "type" + ] + }, + { + "type": "object", + "properties": { + "bank_id": { + "type": "string" + }, + "type": { + "type": "string", + "const": "keyvalue" + }, + "keys": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "additionalProperties": false, + "required": [ + "bank_id", + "type", + "keys" + ] + }, + { + "type": "object", + "properties": { + "bank_id": { + "type": "string" + }, + "type": { + "type": "string", + "const": "keyword" + } + }, + "additionalProperties": false, + "required": [ + "bank_id", + "type" + ] + }, + { + "type": "object", + "properties": { + "bank_id": { + "type": "string" + }, + "type": { + "type": "string", + "const": "graph" + }, + "entities": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "additionalProperties": false, + "required": [ + "bank_id", + "type", + "entities" + ] + } + ] + } + }, + "query_generator_config": { + "oneOf": [ + { + "type": "object", + "properties": { + "type": { + "type": "string", + "const": "default" + }, + "sep": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "type", + "sep" + ] + }, + { + "type": "object", + "properties": { + "type": { + "type": "string", + "const": "llm" + }, + "model": { + "type": "string" + }, + "template": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "type", + "model", + "template" + ] + }, + { + "type": "object", + "properties": { + "type": { + "type": "string", + "const": "custom" + } + }, + "additionalProperties": false, + "required": [ + "type" + ] + } + ] + }, + "max_tokens_in_context": { + "type": "integer" + }, + "max_chunks": { + "type": "integer" + } + }, + "additionalProperties": false, + "required": [ + "type", + "memory_bank_configs", + "query_generator_config", + "max_tokens_in_context", + "max_chunks" + ] + }, "OnViolationAction": { "type": "integer", "enum": [ @@ -5149,37 +5152,37 @@ ], "tags": [ { - "name": "BatchInference" - }, - { - "name": "PostTraining" + "name": "Safety" }, { "name": "Inference" }, - { - "name": "Safety" - }, - { - "name": "RewardScoring" - }, - { - "name": "Telemetry" - }, { "name": "Evaluations" }, { - "name": "SyntheticDataGeneration" + "name": "PostTraining" + }, + { + "name": "BatchInference" }, { "name": "Memory" }, + { + "name": "Datasets" + }, + { + "name": "RewardScoring" + }, { "name": "Agents" }, { - "name": "Datasets" + "name": "Telemetry" + }, + { + "name": "SyntheticDataGeneration" }, { "name": "BuiltinTool", @@ -5317,6 +5320,10 @@ "name": "FunctionCallToolDefinition", "description": "" }, + { + "name": "MemoryToolDefinition", + "description": "" + }, { "name": "OnViolationAction", "description": "" @@ -5763,6 +5770,7 @@ "MemoryBank", "MemoryBankDocument", "MemoryRetrievalStep", + "MemoryToolDefinition", "MetricEvent", "OnViolationAction", "OptimizerConfig", diff --git a/docs/llama-stack-spec.yaml b/docs/llama-stack-spec.yaml index 33d7d9a3a..d4872cf46 100644 --- a/docs/llama-stack-spec.yaml +++ b/docs/llama-stack-spec.yaml @@ -30,123 +30,7 @@ components: - $ref: '#/components/schemas/PhotogenToolDefinition' - $ref: '#/components/schemas/CodeInterpreterToolDefinition' - $ref: '#/components/schemas/FunctionCallToolDefinition' - - additionalProperties: false - properties: - input_shields: - items: - $ref: '#/components/schemas/ShieldDefinition' - type: array - max_chunks: - type: integer - max_tokens_in_context: - type: integer - memory_bank_configs: - items: - oneOf: - - additionalProperties: false - properties: - bank_id: - type: string - type: - const: vector - type: string - required: - - bank_id - - type - type: object - - additionalProperties: false - properties: - bank_id: - type: string - keys: - items: - type: string - type: array - type: - const: keyvalue - type: string - required: - - bank_id - - type - - keys - type: object - - additionalProperties: false - properties: - bank_id: - type: string - type: - const: keyword - type: string - required: - - bank_id - - type - type: object - - additionalProperties: false - properties: - bank_id: - type: string - entities: - items: - type: string - type: array - type: - const: graph - type: string - required: - - bank_id - - type - - entities - type: object - type: array - output_shields: - items: - $ref: '#/components/schemas/ShieldDefinition' - type: array - query_generator_config: - oneOf: - - additionalProperties: false - properties: - sep: - type: string - type: - const: default - type: string - required: - - type - - sep - type: object - - additionalProperties: false - properties: - model: - type: string - template: - type: string - type: - const: llm - type: string - required: - - type - - model - - template - type: object - - additionalProperties: false - properties: - type: - const: custom - type: string - required: - - type - type: object - type: - const: memory - type: string - required: - - type - - memory_bank_configs - - query_generator_config - - max_tokens_in_context - - max_chunks - type: object + - $ref: '#/components/schemas/MemoryToolDefinition' type: array required: - model @@ -1190,6 +1074,124 @@ components: - memory_bank_ids - inserted_context type: object + MemoryToolDefinition: + additionalProperties: false + properties: + input_shields: + items: + $ref: '#/components/schemas/ShieldDefinition' + type: array + max_chunks: + type: integer + max_tokens_in_context: + type: integer + memory_bank_configs: + items: + oneOf: + - additionalProperties: false + properties: + bank_id: + type: string + type: + const: vector + type: string + required: + - bank_id + - type + type: object + - additionalProperties: false + properties: + bank_id: + type: string + keys: + items: + type: string + type: array + type: + const: keyvalue + type: string + required: + - bank_id + - type + - keys + type: object + - additionalProperties: false + properties: + bank_id: + type: string + type: + const: keyword + type: string + required: + - bank_id + - type + type: object + - additionalProperties: false + properties: + bank_id: + type: string + entities: + items: + type: string + type: array + type: + const: graph + type: string + required: + - bank_id + - type + - entities + type: object + type: array + output_shields: + items: + $ref: '#/components/schemas/ShieldDefinition' + type: array + query_generator_config: + oneOf: + - additionalProperties: false + properties: + sep: + type: string + type: + const: default + type: string + required: + - type + - sep + type: object + - additionalProperties: false + properties: + model: + type: string + template: + type: string + type: + const: llm + type: string + required: + - type + - model + - template + type: object + - additionalProperties: false + properties: + type: + const: custom + type: string + required: + - type + type: object + type: + const: memory + type: string + required: + - type + - memory_bank_configs + - query_generator_config + - max_tokens_in_context + - max_chunks + type: object MetricEvent: additionalProperties: false properties: @@ -2362,7 +2364,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-09-18 19:27:39.955190" + \ draft and subject to change.\n Generated at 2024-09-20 17:50:36.257743" title: '[DRAFT] Llama Stack Specification' version: 0.0.1 jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema @@ -3150,17 +3152,17 @@ security: servers: - url: http://any-hosted-llama-stack.com tags: -- name: BatchInference -- name: PostTraining -- name: Inference - name: Safety -- name: RewardScoring -- name: Telemetry +- name: Inference - name: Evaluations -- name: SyntheticDataGeneration +- name: PostTraining +- name: BatchInference - name: Memory -- name: Agents - name: Datasets +- name: RewardScoring +- name: Agents +- name: Telemetry +- name: SyntheticDataGeneration - description: name: BuiltinTool - description: name: FunctionCallToolDefinition +- description: + name: MemoryToolDefinition - description: name: OnViolationAction @@ -3634,6 +3639,7 @@ x-tagGroups: - MemoryBank - MemoryBankDocument - MemoryRetrievalStep + - MemoryToolDefinition - MetricEvent - OnViolationAction - OptimizerConfig diff --git a/llama_stack/apis/agents/agents.py b/llama_stack/apis/agents/agents.py index 5cc9ce242..ca4790456 100644 --- a/llama_stack/apis/agents/agents.py +++ b/llama_stack/apis/agents/agents.py @@ -151,6 +151,7 @@ MemoryQueryGeneratorConfig = Annotated[ ] +@json_schema_type class MemoryToolDefinition(ToolDefinitionCommon): type: Literal[AgentTool.memory.value] = AgentTool.memory.value memory_bank_configs: List[MemoryBankConfig] = Field(default_factory=list)