phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-16 06:53:47 +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:45:41 -08:00
parent c82f13bf9e
commit bd43d37f4e
4 changed files with 99 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

16
llama_stack/distribution/routers/routers.py
View file

@ -7,8 +7,8 @@
from typing import Any, AsyncGenerator, Dict, List, Optional
from llama_stack.apis.datasetio.datasetio import DatasetIO
from llama_stack.apis.memory_banks.memory_banks import BankParams, MemoryBankType
from llama_stack.distribution.datatypes import RoutingTable
from llama_stack.apis.memory import * # noqa: F403
from llama_stack.apis.inference import * # noqa: F403
from llama_stack.apis.safety import * # noqa: F403
@ -34,9 +34,19 @@ class MemoryRouter(Memory):
async def register_memory_bank(
self,
request: RegistrationRequest,
memory_bank_id: str,
memory_bank_type: MemoryBankType,
provider_id: Optional[str] = None,
provider_memorybank_id: Optional[str] = None,
params: Optional[BankParams] = None,
) -> None:
await self.routing_table.register_memory_bank(request)
await self.routing_table.register_memory_bank(
memory_bank_id,
memory_bank_type,
provider_id,
provider_memorybank_id,
params,
)
async def insert_documents(
self,

18
llama_stack/distribution/routers/routing_tables.py
View file

@ -17,6 +17,7 @@ from llama_stack.apis.eval_tasks import * # noqa: F403
from llama_stack.distribution.store import DistributionRegistry
from llama_stack.distribution.datatypes import * # noqa: F403
from llama_stack.distribution.utils.memory_bank_utils import build_memory_bank
def get_impl_api(p: Any) -> Api:
@ -272,19 +273,24 @@ class MemoryBanksRoutingTable(CommonRoutingTableImpl, MemoryBanks):
async def register_memory_bank(
self,
request: RegistrationRequest,
memory_bank_id: str,
provider_id: str,
provider_memorybank_id: str,
params: BankParams,
) -> MemoryBank:
if request.provider_resource_id is None:
request.provider_resource_id = request.memory_bank_id
if request.provider_id is None:
if provider_memorybank_id is None:
provider_memorybank_id = memory_bank_id
if provider_id is None:
# If provider_id not specified, use the only provider if it supports this shield type
if len(self.impls_by_provider_id) == 1:
request.provider_id = list(self.impls_by_provider_id.keys())[0]
provider_id = list(self.impls_by_provider_id.keys())[0]
else:
raise ValueError(
"No provider specified and multiple providers available. Please specify a provider_id."
)
memory_bank = registration_request_to_memory_bank(request)
memory_bank = build_memory_bank(
memory_bank_id, params.type, provider_id, provider_memorybank_id, params
)
await self.register_object(memory_bank)
return memory_bank