| 311 | |
| 312 | |
| 313 | class ApiServerSpanExporter(export_lib.SpanExporter): |
| 314 | |
| 315 | def __init__(self, trace_dict): |
| 316 | self.trace_dict = trace_dict |
| 317 | |
| 318 | def export( |
| 319 | self, spans: typing.Sequence[ReadableSpan] |
| 320 | ) -> export_lib.SpanExportResult: |
| 321 | for span in spans: |
| 322 | if ( |
| 323 | span.name == "call_llm" |
| 324 | or span.name == "send_data" |
| 325 | or span.name.startswith("execute_tool") |
| 326 | ): |
| 327 | attributes = dict(span.attributes) |
| 328 | attributes["trace_id"] = span.get_span_context().trace_id |
| 329 | attributes["span_id"] = span.get_span_context().span_id |
| 330 | if attributes.get("gcp.vertex.agent.event_id", None): |
| 331 | self.trace_dict[attributes["gcp.vertex.agent.event_id"]] = attributes |
| 332 | return export_lib.SpanExportResult.SUCCESS |
| 333 | |
| 334 | def force_flush(self, timeout_millis: int = 30000) -> bool: |
| 335 | return True |
| 336 | |
| 337 | |
| 338 | class InMemoryExporter(export_lib.SpanExporter): |