Spaces:
Sleeping
Sleeping
| """One event stream, two consumers. The live log and the trace attributes come from here.""" | |
| import queue | |
| from contextlib import contextmanager | |
| _DONE = object() | |
| _tracer = None | |
| _checked = False | |
| def init_tracing(): | |
| """ADOT installs the tracer provider in the AgentCore container, so we only borrow it.""" | |
| global _tracer, _checked | |
| if _checked: | |
| return _tracer | |
| _checked = True | |
| from opentelemetry import trace | |
| provider = trace.get_tracer_provider() | |
| # without a real SDK provider nothing is collecting, so our spans stay off rather than buffer | |
| if not hasattr(provider, "add_span_processor"): | |
| return None | |
| _tracer = trace.get_tracer("easeapply") | |
| return _tracer | |
| class RunTrace: | |
| """One span per run, kept current so every stage span becomes a child of it. | |
| Without this each stage is its own root trace and a collector shows twenty unrelated | |
| entries instead of one run you can open. | |
| """ | |
| def __init__(self, run_id: str, kind: str = "full"): | |
| self._span = None | |
| self._token = None | |
| tracer = init_tracing() | |
| if tracer is None: | |
| return | |
| from opentelemetry import context, trace | |
| self._context = context | |
| self._span = tracer.start_span(f"easeapply run {run_id}") | |
| self._span.set_attribute("easeapply.run_id", run_id) | |
| self._span.set_attribute("easeapply.kind", kind) | |
| self._span.set_attribute("session.id", run_id) | |
| self._token = context.attach(trace.set_span_in_context(self._span)) | |
| def finish(self, **attributes) -> None: | |
| if self._span is None: | |
| return | |
| for key, value in attributes.items(): | |
| self._span.set_attribute(f"easeapply.{key}", value) | |
| self._context.detach(self._token) | |
| self._span.end() | |
| def stage_span(stage: str, run_id: str | None = None, **attributes): | |
| """Adds the funnel to the default Strands spans, which only show that agents ran.""" | |
| tracer = init_tracing() | |
| if tracer is None: | |
| yield None | |
| return | |
| with tracer.start_as_current_span(stage) as span: | |
| span.set_attribute("easeapply.stage", stage) | |
| if run_id: | |
| # session.id is what Phoenix groups a run by, so one run reads as one session | |
| span.set_attribute("session.id", run_id) | |
| span.set_attribute("easeapply.run_id", run_id) | |
| for key, value in attributes.items(): | |
| span.set_attribute(f"easeapply.{key}", value) | |
| yield span | |
| class EventBus: | |
| """In process queue. The pipeline writes, the Gradio generator drains.""" | |
| def __init__(self): | |
| self._queue: queue.Queue = queue.Queue() | |
| self.funnel: dict = {} | |
| self.run_id: str | None = None | |
| def emit(self, message: str, stage: str | None = None, **attributes) -> None: | |
| if attributes: | |
| self.funnel.update(attributes) | |
| with stage_span(stage or "event", run_id=self.run_id, **attributes): | |
| pass | |
| self._queue.put({"message": message, "stage": stage, **attributes}) | |
| def close(self) -> None: | |
| self._queue.put(_DONE) | |
| def stream(self): | |
| """Yields events until the producer closes the stream.""" | |
| while True: | |
| event = self._queue.get() | |
| if event is _DONE: | |
| return | |
| yield event | |