phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-01 00:05:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

address feedback

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-08 16:51:05 -08:00
parent bd43d37f4e
commit 24b914f0fe
2 changed files with 15 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

19
llama_stack/apis/memory_banks/memory_banks.py
View file

@ -5,11 +5,19 @@
# the root directory of this source tree. # the root directory of this source tree.
from enum import Enum from enum import Enum
from typing import List, Literal, Optional, Protocol, runtime_checkable from typing import (
Annotated,
List,
Literal,
Optional,
Protocol,
runtime_checkable,
Union,
)
from llama_models.schema_utils import json_schema_type, webmethod from llama_models.schema_utils import json_schema_type, webmethod
from pydantic import BaseModel from pydantic import BaseModel, Field
from llama_stack.apis.resource import Resource, ResourceType from llama_stack.apis.resource import Resource, ResourceType
@ -63,9 +71,10 @@ class VectorMemoryBankParams(BaseModel):
overlap_size_in_tokens: Optional[int] = None overlap_size_in_tokens: Optional[int] = None
BankParams = VectorMemoryBankParams # For now, since we only have one type of params BankParams = Annotated[
# If you need to add more types later, you can do: Union[VectorMemoryBankParams],
# BankParams = Union[VectorMemoryBankParams, KeyValueMemoryBankParams, KeywordMemoryBankParams, GraphMemoryBankParams] Field(discriminator="type"),
]
@runtime_checkable @runtime_checkable

1
llama_stack/providers/datatypes.py
View file

@ -13,6 +13,7 @@ from pydantic import BaseModel, Field
from llama_stack.apis.datasets import DatasetDef from llama_stack.apis.datasets import DatasetDef
from llama_stack.apis.eval_tasks import EvalTaskDef from llama_stack.apis.eval_tasks import EvalTaskDef
from llama_stack.apis.memory_banks.memory_banks import MemoryBank
from llama_stack.apis.models import Model from llama_stack.apis.models import Model
from llama_stack.apis.scoring_functions import ScoringFnDef from llama_stack.apis.scoring_functions import ScoringFnDef
from llama_stack.apis.shields import Shield from llama_stack.apis.shields import Shield