passing tests

This commit is contained in:
Dinesh Yeduguru 2025-01-16 11:56:51 -08:00
parent bb2adef7e4
commit 97d017de1c
7 changed files with 887 additions and 655 deletions

File diff suppressed because it is too large Load diff

View file

@ -836,10 +836,7 @@ components:
oneOf: oneOf:
- $ref: '#/components/schemas/BenchmarkEvalTaskConfig' - $ref: '#/components/schemas/BenchmarkEvalTaskConfig'
- $ref: '#/components/schemas/AppEvalTaskConfig' - $ref: '#/components/schemas/AppEvalTaskConfig'
task_id:
type: string
required: required:
- task_id
- input_rows - input_rows
- scoring_functions - scoring_functions
- task_config - task_config
@ -1003,17 +1000,6 @@ components:
required: required:
- job_id - job_id
type: object type: object
JobCancelRequest:
additionalProperties: false
properties:
job_id:
type: string
task_id:
type: string
required:
- job_id
- task_id
type: object
JobStatus: JobStatus:
enum: enum:
- completed - completed
@ -1178,6 +1164,16 @@ components:
required: required:
- data - data
type: object type: object
ListRoutesResponse:
additionalProperties: false
properties:
data:
items:
$ref: '#/components/schemas/RouteInfo'
type: array
required:
- data
type: object
ListScoringFunctionsResponse: ListScoringFunctionsResponse:
additionalProperties: false additionalProperties: false
properties: properties:
@ -1781,52 +1777,35 @@ components:
- chunks - chunks
- scores - scores
type: object type: object
QuerySpanTreeRequest: QuerySpanTreeResponse:
additionalProperties: false additionalProperties: false
properties: properties:
attributes_to_return: data:
items: additionalProperties:
type: string $ref: '#/components/schemas/SpanWithStatus'
type: array type: object
max_depth:
type: integer
span_id:
type: string
required: required:
- span_id - data
type: object type: object
QuerySpansRequest: QuerySpansResponse:
additionalProperties: false additionalProperties: false
properties: properties:
attribute_filters: data:
items: items:
$ref: '#/components/schemas/QueryCondition' $ref: '#/components/schemas/Span'
type: array type: array
attributes_to_return:
items:
type: string
type: array
max_depth:
type: integer
required: required:
- attribute_filters - data
- attributes_to_return
type: object type: object
QueryTracesRequest: QueryTracesResponse:
additionalProperties: false additionalProperties: false
properties: properties:
attribute_filters: data:
items: items:
$ref: '#/components/schemas/QueryCondition' $ref: '#/components/schemas/Trace'
type: array
limit:
type: integer
offset:
type: integer
order_by:
items:
type: string
type: array type: array
required:
- data
type: object type: object
RegexParserScoringFnParams: RegexParserScoringFnParams:
additionalProperties: false additionalProperties: false
@ -2082,10 +2061,7 @@ components:
oneOf: oneOf:
- $ref: '#/components/schemas/BenchmarkEvalTaskConfig' - $ref: '#/components/schemas/BenchmarkEvalTaskConfig'
- $ref: '#/components/schemas/AppEvalTaskConfig' - $ref: '#/components/schemas/AppEvalTaskConfig'
task_id:
type: string
required: required:
- task_id
- task_config - task_config
type: object type: object
RunShieldRequest: RunShieldRequest:
@ -3916,9 +3892,14 @@ paths:
description: OK description: OK
tags: tags:
- EvalTasks - EvalTasks
/v1/eval/evaluate-rows: /v1/eval/tasks/{task_id}/evaluations:
post: post:
parameters: parameters:
- in: path
name: task_id
required: true
schema:
type: string
- description: JSON-encoded provider data which will be made available to the - description: JSON-encoded provider data which will be made available to the
adapter servicing the API adapter servicing the API
in: header in: header
@ -3948,9 +3929,14 @@ paths:
description: OK description: OK
tags: tags:
- Eval - Eval
/v1/eval/jobs/cancel: /v1/eval/tasks/{task_id}/jobs:
post: post:
parameters: parameters:
- in: path
name: task_id
required: true
schema:
type: string
- description: JSON-encoded provider data which will be made available to the - description: JSON-encoded provider data which will be made available to the
adapter servicing the API adapter servicing the API
in: header in: header
@ -3969,26 +3955,61 @@ paths:
content: content:
application/json: application/json:
schema: schema:
$ref: '#/components/schemas/JobCancelRequest' $ref: '#/components/schemas/RunEvalRequest'
required: true required: true
responses: responses:
'200': '200':
content:
application/json:
schema:
$ref: '#/components/schemas/Job'
description: OK description: OK
tags: tags:
- Eval - Eval
/v1/eval/jobs/{job_id}: /v1/eval/tasks/{task_id}/jobs/{job_id}:
get: delete:
parameters: parameters:
- in: path
name: task_id
required: true
schema:
type: string
- in: path - in: path
name: job_id name: job_id
required: true required: true
schema: schema:
type: string type: string
- in: query - description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
description: OK
tags:
- Eval
get:
parameters:
- in: path
name: task_id name: task_id
required: true required: true
schema: schema:
type: string type: string
- in: path
name: job_id
required: true
schema:
type: string
- description: JSON-encoded provider data which will be made available to the - description: JSON-encoded provider data which will be made available to the
adapter servicing the API adapter servicing the API
in: header in: header
@ -4014,7 +4035,7 @@ paths:
description: OK description: OK
tags: tags:
- Eval - Eval
/v1/eval/jobs/{job_id}/result: /v1/eval/tasks/{task_id}/jobs/{job_id}/result:
get: get:
parameters: parameters:
- in: path - in: path
@ -4022,7 +4043,7 @@ paths:
required: true required: true
schema: schema:
type: string type: string
- in: query - in: path
name: task_id name: task_id
required: true required: true
schema: schema:
@ -4050,38 +4071,6 @@ paths:
description: OK description: OK
tags: tags:
- Eval - Eval
/v1/eval/run:
post:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/RunEvalRequest'
required: true
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/Job'
description: OK
tags:
- Eval
/v1/health: /v1/health:
get: get:
parameters: parameters:
@ -4208,6 +4197,58 @@ paths:
description: OK description: OK
tags: tags:
- Inference - Inference
/v1/inspect/providers:
get:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/ListProvidersResponse'
description: OK
tags:
- Inspect
/v1/inspect/routes:
get:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/ListRoutesResponse'
description: OK
tags:
- Inspect
/v1/memory-banks: /v1/memory-banks:
get: get:
parameters: parameters:
@ -4688,62 +4729,6 @@ paths:
description: OK description: OK
tags: tags:
- PostTraining (Coming Soon) - PostTraining (Coming Soon)
/v1/providers/list:
get:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/ListProvidersResponse'
description: OK
tags:
- Inspect
/v1/routes/list:
get:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
additionalProperties:
items:
$ref: '#/components/schemas/RouteInfo'
type: array
type: object
description: OK
tags:
- Inspect
/v1/safety/run-shield: /v1/safety/run-shield:
post: post:
parameters: parameters:
@ -5048,7 +5033,7 @@ paths:
description: OK description: OK
tags: tags:
- SyntheticDataGeneration (Coming Soon) - SyntheticDataGeneration (Coming Soon)
/v1/telemetry/log-event: /v1/telemetry/events:
post: post:
parameters: parameters:
- description: JSON-encoded provider data which will be made available to the - description: JSON-encoded provider data which will be made available to the
@ -5076,9 +5061,28 @@ paths:
description: OK description: OK
tags: tags:
- Telemetry - Telemetry
/v1/telemetry/query-span-tree: /v1/telemetry/spans:
post: get:
parameters: parameters:
- in: query
name: attribute_filters
required: true
schema:
items:
$ref: '#/components/schemas/QueryCondition'
type: array
- in: query
name: attributes_to_return
required: true
schema:
items:
type: string
type: array
- in: query
name: max_depth
required: false
schema:
type: integer
- description: JSON-encoded provider data which will be made available to the - description: JSON-encoded provider data which will be made available to the
adapter servicing the API adapter servicing the API
in: header in: header
@ -5093,88 +5097,16 @@ paths:
required: false required: false
schema: schema:
type: string type: string
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/QuerySpanTreeRequest'
required: true
responses: responses:
'200': '200':
content: content:
application/json: application/json:
schema: schema:
additionalProperties: $ref: '#/components/schemas/QuerySpansResponse'
$ref: '#/components/schemas/SpanWithStatus'
type: object
description: OK description: OK
tags: tags:
- Telemetry - Telemetry
/v1/telemetry/query-spans: /v1/telemetry/spans/export:
post:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/QuerySpansRequest'
required: true
responses:
'200':
content:
application/jsonl:
schema:
$ref: '#/components/schemas/Span'
description: OK
tags:
- Telemetry
/v1/telemetry/query-traces:
post:
parameters:
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
requestBody:
content:
application/json:
schema:
$ref: '#/components/schemas/QueryTracesRequest'
required: true
responses:
'200':
content:
application/jsonl:
schema:
$ref: '#/components/schemas/Trace'
description: OK
tags:
- Telemetry
/v1/telemetry/save-spans-to-dataset:
post: post:
parameters: parameters:
- description: JSON-encoded provider data which will be made available to the - description: JSON-encoded provider data which will be made available to the
@ -5202,6 +5134,166 @@ paths:
description: OK description: OK
tags: tags:
- Telemetry - Telemetry
/v1/telemetry/spans/{span_id}/tree:
get:
parameters:
- in: path
name: span_id
required: true
schema:
type: string
- in: query
name: attributes_to_return
required: false
schema:
items:
type: string
type: array
- in: query
name: max_depth
required: false
schema:
type: integer
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/QuerySpanTreeResponse'
description: OK
tags:
- Telemetry
/v1/telemetry/traces:
get:
parameters:
- in: query
name: attribute_filters
required: false
schema:
items:
$ref: '#/components/schemas/QueryCondition'
type: array
- in: query
name: limit
required: false
schema:
type: integer
- in: query
name: offset
required: false
schema:
type: integer
- in: query
name: order_by
required: false
schema:
items:
type: string
type: array
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/QueryTracesResponse'
description: OK
tags:
- Telemetry
/v1/telemetry/traces/{trace_id}:
get:
parameters:
- in: path
name: trace_id
required: true
schema:
type: string
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/Trace'
description: OK
tags:
- Telemetry
/v1/telemetry/traces/{trace_id}/spans/{span_id}:
get:
parameters:
- in: path
name: trace_id
required: true
schema:
type: string
- in: path
name: span_id
required: true
schema:
type: string
- description: JSON-encoded provider data which will be made available to the
adapter servicing the API
in: header
name: X-LlamaStack-Provider-Data
required: false
schema:
type: string
- description: Version of the client making the request. This is used to ensure
that the client and server are compatible.
in: header
name: X-LlamaStack-Client-Version
required: false
schema:
type: string
responses:
'200':
content:
application/json:
schema:
$ref: '#/components/schemas/Span'
description: OK
tags:
- Telemetry
/v1/tool-runtime/invoke: /v1/tool-runtime/invoke:
post: post:
parameters: parameters:
@ -5673,9 +5765,6 @@ tags:
name: InvokeToolRequest name: InvokeToolRequest
- description: <SchemaDefinition schemaRef="#/components/schemas/Job" /> - description: <SchemaDefinition schemaRef="#/components/schemas/Job" />
name: Job name: Job
- description: <SchemaDefinition schemaRef="#/components/schemas/JobCancelRequest"
/>
name: JobCancelRequest
- description: <SchemaDefinition schemaRef="#/components/schemas/JobStatus" /> - description: <SchemaDefinition schemaRef="#/components/schemas/JobStatus" />
name: JobStatus name: JobStatus
- description: <SchemaDefinition schemaRef="#/components/schemas/KeyValueMemoryBank" - description: <SchemaDefinition schemaRef="#/components/schemas/KeyValueMemoryBank"
@ -5711,6 +5800,9 @@ tags:
- description: <SchemaDefinition schemaRef="#/components/schemas/ListProvidersResponse" - description: <SchemaDefinition schemaRef="#/components/schemas/ListProvidersResponse"
/> />
name: ListProvidersResponse name: ListProvidersResponse
- description: <SchemaDefinition schemaRef="#/components/schemas/ListRoutesResponse"
/>
name: ListRoutesResponse
- description: <SchemaDefinition schemaRef="#/components/schemas/ListScoringFunctionsResponse" - description: <SchemaDefinition schemaRef="#/components/schemas/ListScoringFunctionsResponse"
/> />
name: ListScoringFunctionsResponse name: ListScoringFunctionsResponse
@ -5797,15 +5889,15 @@ tags:
- description: <SchemaDefinition schemaRef="#/components/schemas/QueryDocumentsResponse" - description: <SchemaDefinition schemaRef="#/components/schemas/QueryDocumentsResponse"
/> />
name: QueryDocumentsResponse name: QueryDocumentsResponse
- description: <SchemaDefinition schemaRef="#/components/schemas/QuerySpanTreeRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/QuerySpanTreeResponse"
/> />
name: QuerySpanTreeRequest name: QuerySpanTreeResponse
- description: <SchemaDefinition schemaRef="#/components/schemas/QuerySpansRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/QuerySpansResponse"
/> />
name: QuerySpansRequest name: QuerySpansResponse
- description: <SchemaDefinition schemaRef="#/components/schemas/QueryTracesRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/QueryTracesResponse"
/> />
name: QueryTracesRequest name: QueryTracesResponse
- description: <SchemaDefinition schemaRef="#/components/schemas/RegexParserScoringFnParams" - description: <SchemaDefinition schemaRef="#/components/schemas/RegexParserScoringFnParams"
/> />
name: RegexParserScoringFnParams name: RegexParserScoringFnParams
@ -6085,7 +6177,6 @@ x-tagGroups:
- InterleavedContentItem - InterleavedContentItem
- InvokeToolRequest - InvokeToolRequest
- Job - Job
- JobCancelRequest
- JobStatus - JobStatus
- KeyValueMemoryBank - KeyValueMemoryBank
- KeyValueMemoryBankParams - KeyValueMemoryBankParams
@ -6098,6 +6189,7 @@ x-tagGroups:
- ListModelsResponse - ListModelsResponse
- ListPostTrainingJobsResponse - ListPostTrainingJobsResponse
- ListProvidersResponse - ListProvidersResponse
- ListRoutesResponse
- ListScoringFunctionsResponse - ListScoringFunctionsResponse
- ListShieldsResponse - ListShieldsResponse
- ListToolGroupsResponse - ListToolGroupsResponse
@ -6127,9 +6219,9 @@ x-tagGroups:
- QueryConditionOp - QueryConditionOp
- QueryDocumentsRequest - QueryDocumentsRequest
- QueryDocumentsResponse - QueryDocumentsResponse
- QuerySpanTreeRequest - QuerySpanTreeResponse
- QuerySpansRequest - QuerySpansResponse
- QueryTracesRequest - QueryTracesResponse
- RegexParserScoringFnParams - RegexParserScoringFnParams
- RegisterDatasetRequest - RegisterDatasetRequest
- RegisterEvalTaskRequest - RegisterEvalTaskRequest

View file

@ -204,7 +204,7 @@ class Telemetry(Protocol):
async def get_span(self, trace_id: str, span_id: str) -> Span: ... async def get_span(self, trace_id: str, span_id: str) -> Span: ...
@webmethod(route="/telemetry/spans/{span_id}/tree", method="GET") @webmethod(route="/telemetry/spans/{span_id}/tree", method="GET")
async def query_span_tree( async def get_span_tree(
self, self,
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,

View file

@ -259,7 +259,7 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
async def get_span(self, trace_id: str, span_id: str) -> Span: async def get_span(self, trace_id: str, span_id: str) -> Span:
return await self.trace_store.get_span(trace_id, span_id) return await self.trace_store.get_span(trace_id, span_id)
async def query_span_tree( async def get_span_tree(
self, self,
span_id: str, span_id: str,
attributes_to_return: Optional[List[str]] = None, attributes_to_return: Optional[List[str]] = None,

View file

@ -49,17 +49,17 @@ class TelemetryDatasetMixin:
attributes_to_return: List[str], attributes_to_return: List[str],
max_depth: Optional[int] = None, max_depth: Optional[int] = None,
) -> QuerySpansResponse: ) -> QuerySpansResponse:
traces = await self.query_traces(attribute_filters=attribute_filters).data traces = await self.query_traces(attribute_filters=attribute_filters)
spans = [] spans = []
for trace in traces: for trace in traces.data:
spans_by_id = await self.query_span_tree( spans_by_id_resp = await self.get_span_tree(
span_id=trace.root_span_id, span_id=trace.root_span_id,
attributes_to_return=attributes_to_return, attributes_to_return=attributes_to_return,
max_depth=max_depth, max_depth=max_depth,
) )
for span in spans_by_id.values(): for span in spans_by_id_resp.data.values():
if span.attributes and all( if span.attributes and all(
attr in span.attributes and span.attributes[attr] is not None attr in span.attributes and span.attributes[attr] is not None
for attr in attributes_to_return for attr in attributes_to_return
@ -76,4 +76,4 @@ class TelemetryDatasetMixin:
) )
) )
return spans return QuerySpansResponse(data=spans)

View file

@ -174,6 +174,8 @@ class SQLiteTraceStore(TraceStore):
conn.row_factory = aiosqlite.Row conn.row_factory = aiosqlite.Row
async with conn.execute(query, (trace_id,)) as cursor: async with conn.execute(query, (trace_id,)) as cursor:
row = await cursor.fetchone() row = await cursor.fetchone()
if row is None:
raise ValueError(f"Trace {trace_id} not found")
return Trace(**row) return Trace(**row)
async def get_span(self, trace_id: str, span_id: str) -> Span: async def get_span(self, trace_id: str, span_id: str) -> Span:
@ -182,4 +184,6 @@ class SQLiteTraceStore(TraceStore):
conn.row_factory = aiosqlite.Row conn.row_factory = aiosqlite.Row
async with conn.execute(query, (trace_id, span_id)) as cursor: async with conn.execute(query, (trace_id, span_id)) as cursor:
row = await cursor.fetchone() row = await cursor.fetchone()
if row is None:
raise ValueError(f"Span {span_id} not found")
return Span(**row) return Span(**row)

View file

@ -178,14 +178,6 @@ def test_builtin_tool_web_search(llama_stack_client, agent_config):
assert "Tool:brave_search Response:" in logs_str assert "Tool:brave_search Response:" in logs_str
assert "mark zuckerberg" in logs_str.lower() assert "mark zuckerberg" in logs_str.lower()
assert "No Violation" in logs_str assert "No Violation" in logs_str
spans = llama_stack_client.telemetry.query_spans(
attribute_filters=[
{"key": "session_id", "op": "eq", "value": session_id},
],
attributes_to_return=["input", "output"],
)
breakpoint()
assert len(spans) > 1
def test_builtin_tool_code_execution(llama_stack_client, agent_config): def test_builtin_tool_code_execution(llama_stack_client, agent_config):