phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-09 11:20:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

memory bank registration fixes

Browse source
This commit is contained in:
Ashwin Bharambe 2024-10-06 22:00:54 -07:00 • committed by Ashwin Bharambe
parent 099a95b614
commit 3725e74906
8 changed files with 108 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -80,12 +80,21 @@ class CommonRoutingTableImpl(RoutingTable):
return obj
return None
async def register_object(self, obj: RoutableObject) -> Any:
async def register_object(self, obj: RoutableObject):
if obj.identifier in self.routing_key_to_object:
raise ValueError(f"Object `{obj.identifier}` already registered")
print(f"Object `{obj.identifier}` is already registered")
return
if obj.provider_id not in self.impls_by_provider_id:
raise ValueError(f"Provider `{obj.provider_id}` not found")
if not obj.provider_id:
provider_ids = list(self.impls_by_provider_id.keys())
if not provider_ids:
raise ValueError("No providers found")
print(f"Picking provider `{provider_ids[0]}` for {obj.identifier}")
obj.provider_id = provider_ids[0]
else:
if obj.provider_id not in self.impls_by_provider_id:
raise ValueError(f"Provider `{obj.provider_id}` not found")
p = self.impls_by_provider_id[obj.provider_id]
await register_object_with_provider(obj, p)
@ -93,6 +102,8 @@ class CommonRoutingTableImpl(RoutingTable):
self.routing_key_to_object[obj.identifier] = obj
self.registry.append(obj)
# TODO: persist this to a store
class ModelsRoutingTable(CommonRoutingTableImpl, Models):
async def list_models(self) -> List[ModelDef]: