mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-17 12:59:47 +00:00
move the save to dataset to telemetry
This commit is contained in:
parent
4c78432bc8
commit
f5d427c178
9 changed files with 79 additions and 64 deletions
|
|
@ -16,7 +16,6 @@ from llama_stack.apis.datasets import Datasets
|
|||
from llama_stack.apis.eval_tasks import EvalTask
|
||||
from llama_stack.apis.inference import Inference
|
||||
from llama_stack.apis.scoring import Scoring
|
||||
from llama_stack.apis.telemetry import QueryCondition, SpanWithChildren, Telemetry
|
||||
from llama_stack.providers.datatypes import EvalTasksProtocolPrivate
|
||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
||||
from tqdm import tqdm
|
||||
|
|
@ -43,7 +42,6 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
|
|||
scoring_api: Scoring,
|
||||
inference_api: Inference,
|
||||
agents_api: Agents,
|
||||
telemetry_api: Telemetry,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.datasetio_api = datasetio_api
|
||||
|
|
@ -51,7 +49,6 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
|
|||
self.scoring_api = scoring_api
|
||||
self.inference_api = inference_api
|
||||
self.agents_api = agents_api
|
||||
self.telemetry_api = telemetry_api
|
||||
|
||||
# TODO: assume sync job, will need jobs API for async scheduling
|
||||
self.jobs = {}
|
||||
|
|
@ -272,50 +269,3 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
|
|||
raise ValueError(f"Job is not completed, Status: {status.value}")
|
||||
|
||||
return self.jobs[job_id]
|
||||
|
||||
async def create_annotation_dataset(self, session_id: str, dataset_id: str) -> None:
|
||||
traces = await self.telemetry_api.query_traces(
|
||||
attribute_filters=[
|
||||
QueryCondition(key="session_id", op="eq", value=session_id),
|
||||
]
|
||||
)
|
||||
|
||||
annotation_rows = []
|
||||
|
||||
for trace in traces:
|
||||
span_tree = await self.telemetry_api.get_span_tree(
|
||||
span_id=trace.root_span_id,
|
||||
attributes_to_return=[
|
||||
"input",
|
||||
"output",
|
||||
"name",
|
||||
],
|
||||
)
|
||||
|
||||
def extract_spans(span: SpanWithChildren) -> List[Dict[str, Any]]:
|
||||
rows = []
|
||||
if (
|
||||
span.attributes
|
||||
and "input" in span.attributes
|
||||
and "output" in span.attributes
|
||||
):
|
||||
row = {
|
||||
"input_query": span.attributes.get("input", ""),
|
||||
"generated_answer": span.attributes.get("output", ""),
|
||||
"trace_id": trace.root_span_id,
|
||||
"span_id": span.span_id,
|
||||
"step_name": span.name,
|
||||
}
|
||||
rows.append(row)
|
||||
|
||||
for child in span.children:
|
||||
rows.extend(extract_spans(child))
|
||||
|
||||
return rows
|
||||
|
||||
annotation_rows.extend(extract_spans(span_tree))
|
||||
|
||||
if annotation_rows:
|
||||
await self.datasetio_api.append_rows(
|
||||
dataset_id=dataset_id, rows=annotation_rows
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue