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

⏪ ensuring formatting changes are reverted in routing_tables.py while maintaining functional changes in the ShieldsRoutingTable

Browse source
This commit is contained in:
m-misiura 2025-03-25 10:32:08 +00:00
parent 87d209d6ef
commit e0ce15b8b5
1 changed files with 25 additions and 67 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -118,9 +118,7 @@ class CommonRoutingTableImpl(RoutingTable):
self.dist_registry = dist_registry self.dist_registry = dist_registry
async def initialize(self) -> None: async def initialize(self) -> None:
async def add_objects( async def add_objects(objs: List[RoutableObjectWithProvider], provider_id: str, cls) -> None:
objs: List[RoutableObjectWithProvider], provider_id: str, cls
) -> None:
for obj in objs: for obj in objs:
if cls is None: if cls is None:
obj.provider_id = provider_id obj.provider_id = provider_id
@ -155,9 +153,7 @@ class CommonRoutingTableImpl(RoutingTable):
for p in self.impls_by_provider_id.values(): for p in self.impls_by_provider_id.values():
await p.shutdown() await p.shutdown()
def get_provider_impl( def get_provider_impl(self, routing_key: str, provider_id: Optional[str] = None) -> Any:
self, routing_key: str, provider_id: Optional[str] = None
) -> Any:
def apiname_object(): def apiname_object():
if isinstance(self, ModelsRoutingTable): if isinstance(self, ModelsRoutingTable):
return ("Inference", "model") return ("Inference", "model")
@ -195,32 +191,24 @@ class CommonRoutingTableImpl(RoutingTable):
raise ValueError(f"Provider not found for `{routing_key}`") raise ValueError(f"Provider not found for `{routing_key}`")
async def get_object_by_identifier( async def get_object_by_identifier(self, type: str, identifier: str) -> Optional[RoutableObjectWithProvider]:
self, type: str, identifier: str
) -> Optional[RoutableObjectWithProvider]:
# Get from disk registry # Get from disk registry
obj = await self.dist_registry.get(type, identifier) obj = await self.dist_registry.get(type, identifier)
if not obj: if not obj:
return None return None
# Check if user has permission to access this object # Check if user has permission to access this object
if not check_access(obj, get_auth_attributes()): if not check_access(obj.identifier, getattr(obj, "access_attributes", None), get_auth_attributes()):
logger.debug( logger.debug(f"Access denied to {type} '{identifier}' based on attribute mismatch")
f"Access denied to {type} '{identifier}' based on attribute mismatch"
)
return None return None
return obj return obj
async def unregister_object(self, obj: RoutableObjectWithProvider) -> None: async def unregister_object(self, obj: RoutableObjectWithProvider) -> None:
await self.dist_registry.delete(obj.type, obj.identifier) await self.dist_registry.delete(obj.type, obj.identifier)
await unregister_object_from_provider( await unregister_object_from_provider(obj, self.impls_by_provider_id[obj.provider_id])
obj, self.impls_by_provider_id[obj.provider_id]
)
async def register_object( async def register_object(self, obj: RoutableObjectWithProvider) -> RoutableObjectWithProvider:
self, obj: RoutableObjectWithProvider
) -> RoutableObjectWithProvider:
# if provider_id is not specified, pick an arbitrary one from existing entries # if provider_id is not specified, pick an arbitrary one from existing entries
if not obj.provider_id and len(self.impls_by_provider_id) > 0: if not obj.provider_id and len(self.impls_by_provider_id) > 0:
obj.provider_id = list(self.impls_by_provider_id.keys())[0] obj.provider_id = list(self.impls_by_provider_id.keys())[0]
@ -235,9 +223,7 @@ class CommonRoutingTableImpl(RoutingTable):
creator_attributes = get_auth_attributes() creator_attributes = get_auth_attributes()
if creator_attributes: if creator_attributes:
obj.access_attributes = AccessAttributes(**creator_attributes) obj.access_attributes = AccessAttributes(**creator_attributes)
logger.info( logger.info(f"Setting access attributes for {obj.type} '{obj.identifier}' based on creator's identity")
f"Setting access attributes for {obj.type} '{obj.identifier}' based on creator's identity"
)
registered_obj = await register_object_with_provider(obj, p) registered_obj = await register_object_with_provider(obj, p)
# TODO: This needs to be fixed for all APIs once they return the registered object # TODO: This needs to be fixed for all APIs once they return the registered object
@ -256,7 +242,9 @@ class CommonRoutingTableImpl(RoutingTable):
# Apply attribute-based access control filtering # Apply attribute-based access control filtering
if filtered_objs: if filtered_objs:
filtered_objs = [ filtered_objs = [
obj for obj in filtered_objs if check_access(obj, get_auth_attributes()) obj
for obj in filtered_objs
if check_access(obj.identifier, getattr(obj, "access_attributes", None), get_auth_attributes())
] ]
return filtered_objs return filtered_objs
@ -295,9 +283,7 @@ class ModelsRoutingTable(CommonRoutingTableImpl, Models):
if model_type is None: if model_type is None:
model_type = ModelType.llm model_type = ModelType.llm
if "embedding_dimension" not in metadata and model_type == ModelType.embedding: if "embedding_dimension" not in metadata and model_type == ModelType.embedding:
raise ValueError( raise ValueError("Embedding model must have an embedding dimension in its metadata")
"Embedding model must have an embedding dimension in its metadata"
)
model = ModelWithACL( model = ModelWithACL(
identifier=model_id, identifier=model_id,
provider_resource_id=provider_model_id, provider_resource_id=provider_model_id,
@ -331,39 +317,25 @@ class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
# Fetch shields after initialization - with robust error handling # Fetch shields after initialization - with robust error handling
try: try:
# Check if the provider implements list_shields # Check if the provider implements list_shields
if hasattr(provider, "list_shields") and callable( if hasattr(provider, "list_shields") and callable(getattr(provider, "list_shields")):
getattr(provider, "list_shields")
):
shields_response = await provider.list_shields() shields_response = await provider.list_shields()
if ( if shields_response and hasattr(shields_response, "data") and shields_response.data:
shields_response
and hasattr(shields_response, "data")
and shields_response.data
):
for shield in shields_response.data: for shield in shields_response.data:
# Ensure type is set # Ensure type is set
if not hasattr(shield, "type") or not shield.type: if not hasattr(shield, "type") or not shield.type:
shield.type = ResourceType.shield.value shield.type = ResourceType.shield.value
await self.dist_registry.register(shield) await self.dist_registry.register(shield)
logger.info( logger.info(f"Registered {len(shields_response.data)} shields from provider {provider_id}")
f"Registered {len(shields_response.data)} shields from provider {provider_id}"
)
else: else:
logger.info(f"No shields found for provider {provider_id}") logger.info(f"No shields found for provider {provider_id}")
else: else:
logger.info( logger.info(f"Provider {provider_id} does not support listing shields")
f"Provider {provider_id} does not support listing shields"
)
except Exception as e: except Exception as e:
# Log the error but continue initialization # Log the error but continue initialization
logger.warning( logger.warning(f"Error listing shields from provider {provider_id}: {str(e)}")
f"Error listing shields from provider {provider_id}: {str(e)}"
)
async def list_shields(self) -> ListShieldsResponse: async def list_shields(self) -> ListShieldsResponse:
return ListShieldsResponse( return ListShieldsResponse(data=await self.get_all_with_type(ResourceType.shield.value))
data=await self.get_all_with_type(ResourceType.shield.value)
)
async def get_shield(self, identifier: str) -> Shield: async def get_shield(self, identifier: str) -> Shield:
shield = await self.get_object_by_identifier("shield", identifier) shield = await self.get_object_by_identifier("shield", identifier)
@ -428,18 +400,14 @@ class VectorDBsRoutingTable(CommonRoutingTableImpl, VectorDBs):
f"No provider specified and multiple providers available. Arbitrarily selected the first provider {provider_id}." f"No provider specified and multiple providers available. Arbitrarily selected the first provider {provider_id}."
) )
else: else:
raise ValueError( raise ValueError("No provider available. Please configure a vector_io provider.")
"No provider available. Please configure a vector_io provider."
)
model = await self.get_object_by_identifier("model", embedding_model) model = await self.get_object_by_identifier("model", embedding_model)
if model is None: if model is None:
raise ValueError(f"Model {embedding_model} not found") raise ValueError(f"Model {embedding_model} not found")
if model.model_type != ModelType.embedding: if model.model_type != ModelType.embedding:
raise ValueError(f"Model {embedding_model} is not an embedding model") raise ValueError(f"Model {embedding_model} is not an embedding model")
if "embedding_dimension" not in model.metadata: if "embedding_dimension" not in model.metadata:
raise ValueError( raise ValueError(f"Model {embedding_model} does not have an embedding dimension")
f"Model {embedding_model} does not have an embedding dimension"
)
vector_db_data = { vector_db_data = {
"identifier": vector_db_id, "identifier": vector_db_id,
"type": ResourceType.vector_db.value, "type": ResourceType.vector_db.value,
@ -461,9 +429,7 @@ class VectorDBsRoutingTable(CommonRoutingTableImpl, VectorDBs):
class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets): class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
async def list_datasets(self) -> ListDatasetsResponse: async def list_datasets(self) -> ListDatasetsResponse:
return ListDatasetsResponse( return ListDatasetsResponse(data=await self.get_all_with_type(ResourceType.dataset.value))
data=await self.get_all_with_type(ResourceType.dataset.value)
)
async def get_dataset(self, dataset_id: str) -> Dataset: async def get_dataset(self, dataset_id: str) -> Dataset:
dataset = await self.get_object_by_identifier("dataset", dataset_id) dataset = await self.get_object_by_identifier("dataset", dataset_id)
@ -525,14 +491,10 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
class ScoringFunctionsRoutingTable(CommonRoutingTableImpl, ScoringFunctions): class ScoringFunctionsRoutingTable(CommonRoutingTableImpl, ScoringFunctions):
async def list_scoring_functions(self) -> ListScoringFunctionsResponse: async def list_scoring_functions(self) -> ListScoringFunctionsResponse:
return ListScoringFunctionsResponse( return ListScoringFunctionsResponse(data=await self.get_all_with_type(ResourceType.scoring_function.value))
data=await self.get_all_with_type(ResourceType.scoring_function.value)
)
async def get_scoring_function(self, scoring_fn_id: str) -> ScoringFn: async def get_scoring_function(self, scoring_fn_id: str) -> ScoringFn:
scoring_fn = await self.get_object_by_identifier( scoring_fn = await self.get_object_by_identifier("scoring_function", scoring_fn_id)
"scoring_function", scoring_fn_id
)
if scoring_fn is None: if scoring_fn is None:
raise ValueError(f"Scoring function '{scoring_fn_id}' not found") raise ValueError(f"Scoring function '{scoring_fn_id}' not found")
return scoring_fn return scoring_fn
@ -635,12 +597,8 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
args: Optional[Dict[str, Any]] = None, args: Optional[Dict[str, Any]] = None,
) -> None: ) -> None:
tools = [] tools = []
tool_defs = await self.impls_by_provider_id[provider_id].list_runtime_tools( tool_defs = await self.impls_by_provider_id[provider_id].list_runtime_tools(toolgroup_id, mcp_endpoint)
toolgroup_id, mcp_endpoint tool_host = ToolHost.model_context_protocol if mcp_endpoint else ToolHost.distribution
)
tool_host = (
ToolHost.model_context_protocol if mcp_endpoint else ToolHost.distribution
)
for tool_def in tool_defs: for tool_def in tool_defs:
tools.append( tools.append(