phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2026-01-03 11:22:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Revert "fix: Revert "feat: record token usage for inference API (#1300)" (#1476)"

Browse source 
This reverts commit 60e7f3d705.
This commit is contained in:
Dinesh Yeduguru 2025-03-11 08:35:27 -07:00
parent 2370e826bc
commit 0cac53ab30
5 changed files with 162 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

12
llama_stack/distribution/routers/__init__.py
View file

@ -45,7 +45,7 @@ async def get_routing_table_impl(
return impl
async def get_auto_router_impl(api: Api, routing_table: RoutingTable, _deps) -> Any:
async def get_auto_router_impl(api: Api, routing_table: RoutingTable, deps: Dict[str, Any]) -> Any:
from .routers import (
DatasetIORouter,
EvalRouter,
@ -65,9 +65,17 @@ async def get_auto_router_impl(api: Api, routing_table: RoutingTable, _deps) ->
"eval": EvalRouter,
"tool_runtime": ToolRuntimeRouter,
}
api_to_deps = {
"inference": {"telemetry": Api.telemetry},
}
if api.value not in api_to_routers:
raise ValueError(f"API {api.value} not found in router map")
impl = api_to_routers[api.value](routing_table)
api_to_dep_impl = {}
for dep_name, dep_api in api_to_deps.get(api.value, {}).items():
if dep_api in deps:
api_to_dep_impl[dep_name] = deps[dep_api]
impl = api_to_routers[api.value](routing_table, **api_to_dep_impl)
await impl.initialize()
return impl