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

explicit type for trace

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-27 09:24:23 -08:00
parent dfe152cb97
commit 2dfbb9744d
3 changed files with 23 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

13
llama_stack/apis/telemetry/telemetry.py
View file

@ -137,6 +137,13 @@ Event = Annotated[
]
@json_schema_type
class EvalTrace(BaseModel):
step: str
input: str
output: str
@runtime_checkable
class Telemetry(Protocol):
@webmethod(route="/telemetry/log-event")
@ -145,11 +152,11 @@ class Telemetry(Protocol):
@webmethod(route="/telemetry/get-trace", method="POST")
async def get_trace(self, trace_id: str) -> Trace: ...
@webmethod(route="/telemetry/get-traces-for-eval", method="POST")
async def get_traces_for_eval(
@webmethod(route="/telemetry/get-traces-for-agent-eval", method="POST")
async def get_traces_for_agent_eval(
self,
session_ids: List[str],
lookback: str = "1h",
limit: int = 100,
dataset_id: Optional[str] = None,
) -> List[Trace]: ...
) -> List[EvalTrace]: ...

6
llama_stack/providers/inline/meta_reference/telemetry/console.py
View file

@ -54,9 +54,9 @@ class ConsoleTelemetryImpl(Telemetry):
async def get_trace(self, trace_id: str) -> Trace:
raise NotImplementedError()
async def get_traces_for_session(
self, session_id: str, lookback: str = "1h", limit: int = 100
) -> List[Trace]:
async def get_traces_for_agent_eval(
self, session_ids: List[str], lookback: str = "1h", limit: int = 100
) -> List[EvalTrace]:
raise NotImplementedError()

22
llama_stack/providers/remote/telemetry/opentelemetry/opentelemetry.py
View file

@ -199,13 +199,13 @@ class OpenTelemetryAdapter(Telemetry):
span.end()
_GLOBAL_STORAGE["active_spans"].pop(span_id, None)
async def get_traces_for_eval(
async def get_traces_for_agent_eval(
self,
session_ids: List[str],
lookback: str = "1h",
limit: int = 100,
dataset_id: Optional[str] = None,
) -> List[Dict[str, Any]]:
) -> List[EvalTrace]:
traces = []
# Fetch traces for each session ID individually
@ -311,7 +311,7 @@ class OpenTelemetryAdapter(Telemetry):
except Exception as e:
raise Exception(f"Error querying Jaeger trace structure: {str(e)}") from e
async def get_trace_for_eval(self, trace_id: str) -> List[Dict[str, Any]]:
async def get_trace_for_eval(self, trace_id: str) -> List[EvalTrace]:
"""
Get simplified trace information focusing on first-level children of create_and_execute_turn operations.
Returns a list of spans with name, input, and output information, sorted by start time.
@ -320,21 +320,19 @@ class OpenTelemetryAdapter(Telemetry):
if not trace_data:
return []
def find_execute_turn_children(
spans: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
results = []
def find_execute_turn_children(spans: List[Dict[str, Any]]) -> List[EvalTrace]:
results: List[EvalTrace] = []
for span in spans:
if span["name"] == "create_and_execute_turn":
# Extract and format children spans
children = sorted(span["children"], key=lambda x: x["start_time"])
for child in children:
results.append(
{
"name": child["name"],
"input": child["tags"].get("input", ""),
"output": child["tags"].get("output", ""),
}
EvalTrace(
step=child["name"],
input=child["tags"].get("input", ""),
output=child["tags"].get("output", ""),
)
)
# Recursively search in children
results.extend(find_execute_turn_children(span["children"]))