Spaces:
Running
Running
File size: 1,502 Bytes
178738b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from any_agent import AnyAgent
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from collections.abc import Sequence
from typing import TYPE_CHECKING, Callable
from opentelemetry.sdk.trace.export import (
SpanExporter,
SpanExportResult,
)
from any_agent import AgentFramework
from any_agent.tracing import TracingProcessor
from any_agent.tracing.trace import AgentSpan
if TYPE_CHECKING:
from opentelemetry.sdk.trace import ReadableSpan
class StreamlitExporter(SpanExporter):
"""Build an `AgentTrace` and export to the different outputs."""
def __init__( # noqa: D107
self, agent_framework: AgentFramework, callback: Callable
):
self.agent_framework = agent_framework
self.processor: TracingProcessor | None = TracingProcessor.create(
agent_framework
)
self.callback = callback
def export(self, spans: Sequence["ReadableSpan"]) -> SpanExportResult: # noqa: D102
if not self.processor:
return SpanExportResult.SUCCESS
for readable_span in spans:
# Check if this span belongs to our run
span = AgentSpan.from_readable_span(readable_span)
self.callback(span)
return SpanExportResult.SUCCESS
def export_logs(agent: AnyAgent, callback: Callable) -> None:
exporter = StreamlitExporter(agent.framework, callback)
span_processor = SimpleSpanProcessor(exporter)
agent._tracer_provider.add_span_processor(span_processor)
|