ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
@@ -0,0 +1,135 @@
|
||||
# coding: utf-8
|
||||
# Copyright (c) 2025 inclusionAI.
|
||||
import typing
|
||||
from os import linesep
|
||||
from aworld.trace.base import Span
|
||||
from aworld.trace.span_cosumer import SpanConsumer, get_span_consumers
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.sdk.trace.export import SpanExportResult, SpanExporter
|
||||
from aworld.logs.util import logger
|
||||
|
||||
|
||||
class FileSpanExporter(SpanExporter):
|
||||
"""Implementation of :class:`SpanExporter` that prints spans to the
|
||||
console.
|
||||
|
||||
This class can be used for diagnostic purposes. It prints the exported
|
||||
spans to the console STDOUT.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_path: str = None,
|
||||
formatter: typing.Callable[
|
||||
[ReadableSpan], str
|
||||
] = lambda span: span.to_json() + linesep,
|
||||
):
|
||||
self.formatter = formatter
|
||||
self.file_path = file_path
|
||||
|
||||
def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
|
||||
try:
|
||||
with open(self.file_path, 'a') as f:
|
||||
for span in spans:
|
||||
f.write(self.formatter(span))
|
||||
|
||||
return SpanExportResult.SUCCESS
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
return SpanExportResult.FAILURE
|
||||
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class ReadOnlySpan(Span, ReadableSpan):
|
||||
"""Implementation of :class:`Span` that wraps a :class:`ReadableSpan`.
|
||||
This class can be used to wrap a :class:`ReadableSpan` to make it
|
||||
read-only.
|
||||
Args:
|
||||
span: The span to wrap.
|
||||
"""
|
||||
|
||||
def __init__(self, span: ReadableSpan):
|
||||
self._span = span
|
||||
|
||||
if not typing.TYPE_CHECKING:
|
||||
def __getattr__(self, name: str) -> typing.Any:
|
||||
return getattr(self._span, name)
|
||||
|
||||
def end(self, end_time: typing.Optional[int] = None) -> None:
|
||||
pass
|
||||
|
||||
def set_attribute(self, key: str, value: typing.Any) -> None:
|
||||
pass
|
||||
|
||||
def set_attributes(self, attributes: dict[str, typing.Any]) -> None:
|
||||
pass
|
||||
|
||||
def is_recording(self) -> bool:
|
||||
return False
|
||||
|
||||
def record_exception(
|
||||
self,
|
||||
exception: BaseException,
|
||||
attributes: dict[str, typing.Any] = None,
|
||||
timestamp: typing.Optional[int] = None,
|
||||
escaped: bool = False,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def get_trace_id(self) -> str:
|
||||
return f"{self._span.get_span_context().trace_id:032x}"
|
||||
|
||||
def get_span_id(self) -> str:
|
||||
return f"{self._span.get_span_context().span_id:016x}"
|
||||
|
||||
|
||||
class SpanConsumerExporter(SpanExporter):
|
||||
"""Implementation of :class:`SpanExporter` that exports spans to
|
||||
multiple span consumers.
|
||||
This class can be used for exporting spans to multiple span consumers.
|
||||
It exports the spans to the span consumers in the order they are passed
|
||||
in the constructor.
|
||||
Args:
|
||||
span_consumers: A sequence of span consumers to export spans to.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
span_consumers: typing.Sequence[SpanConsumer] = None,
|
||||
):
|
||||
self._span_consumers = span_consumers or []
|
||||
self._loaded = False
|
||||
|
||||
def _load_span_consumers(self):
|
||||
if not self._loaded:
|
||||
self._span_consumers.extend(get_span_consumers())
|
||||
self._loaded = True
|
||||
|
||||
def export(
|
||||
self, spans: typing.Sequence[ReadableSpan]
|
||||
) -> SpanExportResult:
|
||||
self._load_span_consumers()
|
||||
span_batches = []
|
||||
for span in spans:
|
||||
span_batches.append(ReadOnlySpan(span))
|
||||
for span_consumer in self._span_consumers:
|
||||
try:
|
||||
span_consumer.consume(span_batches)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error consume spans: {e}, span_consumer: {span_consumer.__class__.__name__}")
|
||||
return SpanExportResult.SUCCESS
|
||||
|
||||
|
||||
class NoOpSpanExporter(SpanExporter):
|
||||
"""Implementation of :class:`SpanExporter` that does not export spans."""
|
||||
|
||||
def export(
|
||||
self, spans: typing.Sequence[ReadableSpan]
|
||||
) -> SpanExportResult:
|
||||
return SpanExportResult.SUCCESS
|
||||
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
return True
|
||||
@@ -0,0 +1,238 @@
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Dict, Any, Union
|
||||
from opentelemetry.sdk.trace import Span, SpanContext
|
||||
from opentelemetry.sdk.trace.export import SpanExporter
|
||||
from aworld.logs.util import logger
|
||||
from aworld.trace.constants import ATTRIBUTES_MESSAGE_RUN_TYPE_KEY, RunType
|
||||
|
||||
|
||||
class SpanStatus(BaseModel):
|
||||
code: str = "UNSET"
|
||||
description: Optional[str] = None
|
||||
|
||||
|
||||
class SpanModel(BaseModel):
|
||||
trace_id: str
|
||||
span_id: str
|
||||
name: str
|
||||
start_time: str
|
||||
end_time: str
|
||||
duration_ms: float
|
||||
attributes: Dict[str, Any]
|
||||
status: SpanStatus
|
||||
parent_id: Optional[str]
|
||||
children: list['SpanModel'] = []
|
||||
run_type: Optional[str] = RunType.OTHER.value
|
||||
is_event: bool = False
|
||||
|
||||
@staticmethod
|
||||
def from_span(span):
|
||||
start_timestamp = span.start_time / 1e9
|
||||
end_timestamp = span.end_time / 1e9
|
||||
start_ms = int((span.start_time % 1e9) / 1e6)
|
||||
end_ms = int((span.end_time % 1e9) / 1e6)
|
||||
|
||||
return SpanModel(
|
||||
trace_id=f"{span.get_span_context().trace_id:032x}",
|
||||
span_id=SpanModel.get_span_id(span),
|
||||
name=span.name,
|
||||
start_time=time.strftime(
|
||||
'%Y-%m-%d %H:%M:%S', time.localtime(start_timestamp)) + f'.{start_ms:03d}',
|
||||
end_time=time.strftime(
|
||||
'%Y-%m-%d %H:%M:%S', time.localtime(end_timestamp)) + f'.{end_ms:03d}',
|
||||
duration_ms=(span.end_time - span.start_time)/1e6,
|
||||
attributes={k: v for k, v in span.attributes.items()},
|
||||
status=SpanStatus(
|
||||
code=str(
|
||||
span.status.status_code) if span.status.status_code else "UNSET",
|
||||
description=span.status.description or None
|
||||
),
|
||||
parent_id=SpanModel.get_span_id(
|
||||
span.parent) if span.parent else None,
|
||||
run_type=span.attributes.get(
|
||||
ATTRIBUTES_MESSAGE_RUN_TYPE_KEY, RunType.OTHER.value),
|
||||
is_event=(span.attributes.get("event.id") is not None)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_span_id(span: Union[Span, SpanContext]):
|
||||
if isinstance(span, SpanContext):
|
||||
return f"{span.span_id:016x}"
|
||||
return f"{span.get_span_context().span_id:016x}"
|
||||
|
||||
|
||||
class TraceStorage(ABC):
|
||||
"""
|
||||
Storage for traces.
|
||||
"""
|
||||
@abstractmethod
|
||||
def add_span(self, span: Span) -> None:
|
||||
"""
|
||||
Add a span to the storage.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_all_traces(self) -> list[str]:
|
||||
"""
|
||||
Get all trace ids.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_all_spans(self, trace_id) -> list[SpanModel]:
|
||||
"""
|
||||
Get all spans of a trace.
|
||||
"""
|
||||
|
||||
|
||||
class InMemoryStorage(TraceStorage):
|
||||
"""
|
||||
In-memory storage for spans.
|
||||
"""
|
||||
|
||||
def __init__(self, max_traces=1000):
|
||||
self._traces = defaultdict(list)
|
||||
self._trace_order = []
|
||||
self.max_traces = max_traces
|
||||
|
||||
def add_span(self, span: Span):
|
||||
trace_id = f"{span.get_span_context().trace_id:032x}"
|
||||
if trace_id not in self._traces:
|
||||
self._trace_order.append(trace_id)
|
||||
if len(self._trace_order) > self.max_traces:
|
||||
oldest_trace = self._trace_order.pop(0)
|
||||
del self._traces[oldest_trace]
|
||||
self._traces[trace_id].append(SpanModel.from_span(span))
|
||||
|
||||
def get_all_traces(self):
|
||||
return list(self._traces.keys())
|
||||
|
||||
def get_all_spans(self, trace_id):
|
||||
return self._traces.get(trace_id, [])
|
||||
|
||||
|
||||
class InMemoryWithPersistStorage(TraceStorage):
|
||||
"""
|
||||
In-memory storage for spans with optimized disk persistence.
|
||||
"""
|
||||
|
||||
def __init__(self, storage_dir: str = "./trace_data"):
|
||||
self._traces = defaultdict(list)
|
||||
self._pending_spans = []
|
||||
self.storage_dir = os.path.abspath(storage_dir)
|
||||
os.makedirs(self.storage_dir, exist_ok=True)
|
||||
self._lock = threading.Lock()
|
||||
self._persist_thread = None
|
||||
self._load_today_traces()
|
||||
self.current_filename = None
|
||||
|
||||
def _get_today_filename(self):
|
||||
if not self.current_filename:
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
self.current_filename = f"trace_{timestamp}.json"
|
||||
return self.current_filename
|
||||
|
||||
def _load_today_traces(self):
|
||||
today = datetime.now().strftime("%Y%m%d")
|
||||
for filename in os.listdir(self.storage_dir):
|
||||
if filename.startswith(f"trace_{today}") and filename.endswith(".json"):
|
||||
filepath = os.path.join(self.storage_dir, filename)
|
||||
try:
|
||||
with self._lock, open(filepath, 'r') as f:
|
||||
data = json.load(f)
|
||||
for span_data in data:
|
||||
trace_id = span_data.get("trace_id")
|
||||
span_json = span_data.get("span")
|
||||
self._traces[trace_id].append(
|
||||
SpanModel.parse_raw(span_json))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error loading trace file {filename}: {str(e)}")
|
||||
|
||||
def _start_persist_thread(self):
|
||||
if self._persist_thread is None:
|
||||
self._persist_thread = threading.Thread(
|
||||
target=self._persist_worker, daemon=True)
|
||||
self._persist_thread.start()
|
||||
|
||||
def _persist_worker(self):
|
||||
while True:
|
||||
time.sleep(5)
|
||||
self._persist()
|
||||
|
||||
def _persist(self):
|
||||
if not self._pending_spans:
|
||||
return
|
||||
|
||||
temp_filepath = os.path.join(
|
||||
self.storage_dir, f"temp_{time.time_ns()}.json")
|
||||
final_filepath = os.path.join(
|
||||
self.storage_dir, self._get_today_filename())
|
||||
|
||||
try:
|
||||
spans_to_persist = []
|
||||
with self._lock:
|
||||
spans_to_persist = self._pending_spans.copy()
|
||||
self._pending_spans.clear()
|
||||
|
||||
if spans_to_persist:
|
||||
existing_data = []
|
||||
if os.path.exists(final_filepath):
|
||||
try:
|
||||
with open(final_filepath, 'r') as f:
|
||||
existing_data = json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error reading existing trace file: {str(e)}")
|
||||
|
||||
merged_spans = existing_data + spans_to_persist
|
||||
|
||||
with open(temp_filepath, 'w') as f:
|
||||
json.dump(merged_spans, f, default=str)
|
||||
os.replace(temp_filepath, final_filepath)
|
||||
except Exception as e:
|
||||
logger.error(f"Error persisting traces: {str(e)}")
|
||||
try:
|
||||
os.unlink(temp_filepath)
|
||||
except:
|
||||
pass
|
||||
|
||||
def add_span(self, span: Span):
|
||||
span_model = SpanModel.from_span(span)
|
||||
with self._lock:
|
||||
self._traces[span_model.trace_id].append(span_model)
|
||||
self._pending_spans.append({
|
||||
"trace_id": span_model.trace_id,
|
||||
"span": span_model.json()
|
||||
})
|
||||
self._start_persist_thread()
|
||||
|
||||
def get_all_traces(self):
|
||||
with self._lock:
|
||||
return list(self._traces.keys())
|
||||
|
||||
def get_all_spans(self, trace_id):
|
||||
with self._lock:
|
||||
return self._traces.get(trace_id, [])
|
||||
|
||||
|
||||
class InMemorySpanExporter(SpanExporter):
|
||||
"""
|
||||
Span exporter that stores spans in memory.
|
||||
"""
|
||||
|
||||
def __init__(self, storage: TraceStorage):
|
||||
self._storage = storage
|
||||
|
||||
def export(self, spans):
|
||||
for span in spans:
|
||||
self._storage.add_span(span)
|
||||
|
||||
def shutdown(self):
|
||||
pass
|
||||
@@ -0,0 +1,436 @@
|
||||
import sys
|
||||
import os
|
||||
import traceback
|
||||
import time
|
||||
import datetime
|
||||
import requests
|
||||
from threading import Lock
|
||||
from typing import Any, Iterator, Sequence, Optional, TYPE_CHECKING
|
||||
from contextvars import Token
|
||||
from urllib.parse import urljoin
|
||||
import opentelemetry.context as otlp_context_api
|
||||
from opentelemetry.trace import (
|
||||
SpanKind,
|
||||
set_span_in_context,
|
||||
get_current_span as get_current_otlp_span,
|
||||
NonRecordingSpan,
|
||||
SpanContext,
|
||||
TraceFlags
|
||||
)
|
||||
from opentelemetry.trace.status import StatusCode
|
||||
from opentelemetry.sdk.trace import (
|
||||
ReadableSpan,
|
||||
SynchronousMultiSpanProcessor,
|
||||
Tracer as SDKTracer,
|
||||
Span as SDKSpan,
|
||||
TracerProvider as SDKTracerProvider
|
||||
)
|
||||
from opentelemetry.context import Context as OTLPContext
|
||||
from opentelemetry.semconv.trace import SpanAttributes
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
|
||||
|
||||
from aworld.trace.base import (
|
||||
AttributeValueType,
|
||||
NoOpTracer,
|
||||
SpanType,
|
||||
TraceProvider,
|
||||
Tracer,
|
||||
Span,
|
||||
TraceContext,
|
||||
set_tracer_provider
|
||||
)
|
||||
from aworld.trace.span_cosumer import SpanConsumer
|
||||
from aworld.trace.propagator import get_global_trace_context
|
||||
from aworld.trace.baggage.sofa_tracer import SofaSpanHelper
|
||||
from aworld.logs.util import logger
|
||||
from aworld.utils.common import get_local_ip
|
||||
from .memory_storage import InMemorySpanExporter, InMemoryStorage
|
||||
from ..constants import ATTRIBUTES_MESSAGE_KEY
|
||||
from .export import FileSpanExporter, NoOpSpanExporter, SpanConsumerExporter
|
||||
from ..server import set_trace_server
|
||||
|
||||
|
||||
class OTLPTraceProvider(TraceProvider):
|
||||
"""A TraceProvider that wraps an existing `SDKTracerProvider`.
|
||||
This class provides a way to use a `SDKTracerProvider` as a `TraceProvider`.
|
||||
When the context manager is entered, it returns the `SDKTracerProvider` itself.
|
||||
When the context manager is exited, it calls `shutdown` on the `SDKTracerProvider`.
|
||||
Args:
|
||||
provider: The internal provider to wrap.
|
||||
"""
|
||||
|
||||
def __init__(self, provider: SDKTracerProvider, suppressed_scopes: Optional[set[str]] = None):
|
||||
self._provider: SDKTracerProvider = provider
|
||||
self._suppressed_scopes = set()
|
||||
if suppressed_scopes:
|
||||
self._suppressed_scopes.update(suppressed_scopes)
|
||||
self._lock: Lock = Lock()
|
||||
|
||||
def get_tracer(
|
||||
self,
|
||||
name: str,
|
||||
version: Optional[str] = None
|
||||
):
|
||||
with self._lock:
|
||||
if name in self._suppressed_scopes:
|
||||
return NoOpTracer()
|
||||
else:
|
||||
tracer = self._provider.get_tracer(instrumenting_module_name=name,
|
||||
instrumenting_library_version=version)
|
||||
return OTLPTracer(tracer)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
with self._lock:
|
||||
if isinstance(self._provider, SDKTracerProvider):
|
||||
self._provider.shutdown()
|
||||
|
||||
def force_flush(self, timeout: Optional[float] = None) -> bool:
|
||||
with self._lock:
|
||||
if isinstance(self._provider, SDKTracerProvider):
|
||||
return self._provider.force_flush(timeout)
|
||||
else:
|
||||
return False
|
||||
|
||||
def get_current_span(self) -> Optional["Span"]:
|
||||
otlp_span = get_current_otlp_span()
|
||||
return OTLPSpan(otlp_span, is_new_span=False)
|
||||
|
||||
|
||||
class OTLPTracer(Tracer):
|
||||
"""A Tracer represents a collection of Spans.
|
||||
Args:
|
||||
tracer: The internal tracer to wrap.
|
||||
"""
|
||||
|
||||
def __init__(self, tracer: SDKTracer):
|
||||
self._tracer = tracer
|
||||
|
||||
def start_span(
|
||||
self,
|
||||
name: str,
|
||||
span_type: SpanType = SpanType.INTERNAL,
|
||||
attributes: dict[str, AttributeValueType] = None,
|
||||
start_time: Optional[int] = None,
|
||||
record_exception: bool = True,
|
||||
set_status_on_exception: bool = True,
|
||||
trace_context: Optional[TraceContext] = None
|
||||
) -> "Span":
|
||||
otel_context = None
|
||||
trace_context = trace_context or get_global_trace_context().get_and_clear()
|
||||
if trace_context:
|
||||
otel_context = self._get_otel_context_from_trace_context(
|
||||
trace_context)
|
||||
start_time = start_time or time.time_ns()
|
||||
attributes = {**(attributes or {})}
|
||||
attributes.setdefault(ATTRIBUTES_MESSAGE_KEY, name)
|
||||
SofaSpanHelper.set_sofa_context_to_attr(attributes)
|
||||
attributes = {k: v for k, v in attributes.items(
|
||||
) if is_valid_attribute_value(k, v)}
|
||||
|
||||
span_kind = self._convert_to_span_kind(
|
||||
span_type) if span_type else SpanKind.INTERNAL
|
||||
span = self._tracer.start_span(name=name,
|
||||
kind=span_kind,
|
||||
context=otel_context,
|
||||
attributes=attributes,
|
||||
start_time=start_time,
|
||||
record_exception=record_exception,
|
||||
set_status_on_exception=set_status_on_exception)
|
||||
return OTLPSpan(span)
|
||||
|
||||
def start_as_current_span(
|
||||
self,
|
||||
name: str,
|
||||
span_type: SpanType = SpanType.INTERNAL,
|
||||
attributes: dict[str, AttributeValueType] = None,
|
||||
start_time: Optional[int] = None,
|
||||
record_exception: bool = True,
|
||||
set_status_on_exception: bool = True,
|
||||
end_on_exit: bool = True,
|
||||
trace_context: Optional[TraceContext] = None
|
||||
) -> Iterator["Span"]:
|
||||
|
||||
start_time = start_time or time.time_ns()
|
||||
attributes = {**(attributes or {})}
|
||||
attributes.setdefault(ATTRIBUTES_MESSAGE_KEY, name)
|
||||
SofaSpanHelper.set_sofa_context_to_attr(attributes)
|
||||
attributes = {k: v for k, v in attributes.items(
|
||||
) if is_valid_attribute_value(k, v)}
|
||||
|
||||
span_kind = self._convert_to_span_kind(
|
||||
span_type) if span_type else SpanKind.INTERNAL
|
||||
otel_context = None
|
||||
trace_context = trace_context or get_global_trace_context().get_and_clear()
|
||||
if trace_context:
|
||||
otel_context = self._get_otel_context_from_trace_context(
|
||||
trace_context)
|
||||
|
||||
class _OTLPSpanContextManager:
|
||||
def __init__(self, tracer: SDKTracer):
|
||||
self._span_cm = None
|
||||
self._tracer = tracer
|
||||
|
||||
def __enter__(self):
|
||||
self._span_cm = self._tracer.start_as_current_span(
|
||||
name=name,
|
||||
kind=span_kind,
|
||||
context=otel_context,
|
||||
attributes=attributes,
|
||||
start_time=start_time,
|
||||
record_exception=record_exception,
|
||||
set_status_on_exception=set_status_on_exception,
|
||||
end_on_exit=end_on_exit
|
||||
)
|
||||
inner_span = self._span_cm.__enter__()
|
||||
return OTLPSpan(inner_span)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
return self._span_cm.__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
return _OTLPSpanContextManager(self._tracer)
|
||||
|
||||
def _convert_to_span_kind(self, span_type: SpanType) -> str:
|
||||
if span_type == SpanType.INTERNAL:
|
||||
return SpanKind.INTERNAL
|
||||
elif span_type == SpanType.CLIENT:
|
||||
return SpanKind.CLIENT
|
||||
elif span_type == SpanType.SERVER:
|
||||
return SpanKind.SERVER
|
||||
elif span_type == SpanType.PRODUCER:
|
||||
return SpanKind.PRODUCER
|
||||
elif span_type == SpanType.CONSUMER:
|
||||
return SpanKind.CONSUMER
|
||||
else:
|
||||
return SpanKind.INTERNAL
|
||||
|
||||
def _get_otel_context_from_trace_context(self, trace_context: TraceContext) -> OTLPContext:
|
||||
trace_flags = None
|
||||
if trace_context.trace_flags:
|
||||
trace_flags = TraceFlags(int(trace_context.trace_flags, 16))
|
||||
otel_context = otlp_context_api.Context()
|
||||
return set_span_in_context(
|
||||
NonRecordingSpan(
|
||||
SpanContext(
|
||||
trace_id=int(trace_context.trace_id, 16),
|
||||
span_id=int(trace_context.span_id, 16),
|
||||
is_remote=True,
|
||||
trace_flags=trace_flags
|
||||
)
|
||||
),
|
||||
otel_context,
|
||||
)
|
||||
|
||||
|
||||
class OTLPSpan(Span, ReadableSpan):
|
||||
"""A Span represents a single operation within a trace.
|
||||
"""
|
||||
|
||||
def __init__(self, span: SDKSpan, is_new_span=True):
|
||||
self._span = span
|
||||
self._token: Optional[Token[OTLPContext]] = None
|
||||
if is_new_span:
|
||||
self._attach()
|
||||
self._add_to_open_spans()
|
||||
|
||||
if not TYPE_CHECKING: # pragma: no branch
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._span, name)
|
||||
|
||||
def end(self, end_time: Optional[int] = None) -> None:
|
||||
self._remove_from_open_spans()
|
||||
end_time = end_time or time.time_ns()
|
||||
if not self._span._status or self._span._status.status_code == StatusCode.UNSET:
|
||||
self._span.set_status(
|
||||
status=StatusCode.OK,
|
||||
description="",
|
||||
)
|
||||
self._span.end(end_time=end_time)
|
||||
self._detach()
|
||||
|
||||
def set_attribute(self, key: str, value: Any) -> None:
|
||||
if not is_valid_attribute_value(key, value):
|
||||
return
|
||||
self._span.set_attribute(key=key, value=value)
|
||||
|
||||
def set_attributes(self, attributes: dict[str, Any]) -> None:
|
||||
attributes = {k: v for k, v in attributes.items(
|
||||
) if is_valid_attribute_value(k, v)}
|
||||
self._span.set_attributes(attributes=attributes)
|
||||
|
||||
def is_recording(self) -> bool:
|
||||
return self._span.is_recording()
|
||||
|
||||
def record_exception(
|
||||
self,
|
||||
exception: BaseException,
|
||||
attributes: dict[str, Any] = None,
|
||||
timestamp: Optional[int] = None,
|
||||
escaped: bool = False,
|
||||
) -> None:
|
||||
timestamp = timestamp or time.time_ns()
|
||||
attributes = {**(attributes or {})}
|
||||
|
||||
stacktrace = ''.join(traceback.format_exception(
|
||||
type(exception), exception, exception.__traceback__))
|
||||
self._span.set_attributes({
|
||||
SpanAttributes.EXCEPTION_STACKTRACE: stacktrace,
|
||||
SpanAttributes.EXCEPTION_TYPE: type(exception).__name__,
|
||||
SpanAttributes.EXCEPTION_MESSAGE: str(exception),
|
||||
SpanAttributes.EXCEPTION_ESCAPED: escaped
|
||||
})
|
||||
if exception is not sys.exc_info()[1]:
|
||||
attributes[SpanAttributes.EXCEPTION_STACKTRACE] = stacktrace
|
||||
|
||||
self._span.record_exception(exception=exception,
|
||||
attributes=attributes,
|
||||
timestamp=timestamp,
|
||||
escaped=escaped)
|
||||
self._span.set_status(
|
||||
status=StatusCode.ERROR,
|
||||
description=str(exception),
|
||||
)
|
||||
|
||||
def get_trace_id(self) -> str:
|
||||
"""Get the trace ID of the span.
|
||||
Returns:
|
||||
The trace ID of the span.
|
||||
"""
|
||||
if not self._span or not self._span.get_span_context() or not self.is_recording():
|
||||
return None
|
||||
return f"{self._span.get_span_context().trace_id:032x}"
|
||||
|
||||
def get_span_id(self) -> str:
|
||||
"""Get the span ID of the span.
|
||||
Returns:
|
||||
The span ID of the span.
|
||||
"""
|
||||
if not self._span or not self._span.get_span_context() or not self.is_recording():
|
||||
return None
|
||||
return f"{self._span.get_span_context().span_id:016x}"
|
||||
|
||||
def _attach(self):
|
||||
if self._token is not None:
|
||||
return
|
||||
self._token = otlp_context_api.attach(set_span_in_context(self._span))
|
||||
|
||||
def _detach(self):
|
||||
if self._token is None:
|
||||
return
|
||||
try:
|
||||
otlp_context_api.detach(self._token)
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to detach context: {e}")
|
||||
finally:
|
||||
self._token = None
|
||||
|
||||
|
||||
def configure_otlp_provider(
|
||||
backends: Sequence[str] = None,
|
||||
base_url: str = None,
|
||||
write_token: str = None,
|
||||
span_consumers: Optional[Sequence[SpanConsumer]] = None,
|
||||
**kwargs
|
||||
) -> None:
|
||||
"""Configure the OTLP provider.
|
||||
Args:
|
||||
backend: The backend to use.
|
||||
write_token: The write token to use.
|
||||
**kwargs: Additional keyword arguments to pass to the provider.
|
||||
"""
|
||||
from aworld.metrics.opentelemetry.opentelemetry_adapter import build_otel_resource
|
||||
backends = backends or ["logfire"]
|
||||
processor = SynchronousMultiSpanProcessor()
|
||||
processor.add_span_processor(BatchSpanProcessor(
|
||||
SpanConsumerExporter(span_consumers)))
|
||||
for backend in backends:
|
||||
if backend == "logfire":
|
||||
span_exporter = _configure_logfire_exporter(
|
||||
write_token=write_token, base_url=base_url, **kwargs)
|
||||
processor.add_span_processor(BatchSpanProcessor(span_exporter))
|
||||
elif backend == "console":
|
||||
from opentelemetry.sdk.trace.export import ConsoleSpanExporter
|
||||
processor.add_span_processor(
|
||||
BatchSpanProcessor(ConsoleSpanExporter()))
|
||||
elif backend == "file":
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
file_path = kwargs.get("file_path", f"traces_{timestamp}.json")
|
||||
processor.add_span_processor(
|
||||
BatchSpanProcessor(FileSpanExporter(file_path)))
|
||||
elif backend == "memory":
|
||||
logger.info("Using in-memory storage for traces.")
|
||||
storage = kwargs.get(
|
||||
"storage", InMemoryStorage()) or InMemoryStorage()
|
||||
processor.add_span_processor(
|
||||
SimpleSpanProcessor(InMemorySpanExporter(storage=storage)))
|
||||
server_enabled = str(kwargs.get("server_enabled")) or os.getenv(
|
||||
"START_TRACE_SERVER") or "true"
|
||||
server_port = kwargs.get("server_port") or 7079
|
||||
if (server_enabled.lower() == "true"):
|
||||
logger.info(f"Starting trace server on port {server_port}.")
|
||||
set_trace_server(storage=storage, port=int(
|
||||
server_port), start_server=True)
|
||||
else:
|
||||
logger.info("Trace server is not started.")
|
||||
set_trace_server(storage=storage, port=int(
|
||||
server_port), start_server=False)
|
||||
else:
|
||||
span_exporter = _configure_otlp_exporter(
|
||||
base_url=base_url, **kwargs)
|
||||
processor.add_span_processor(BatchSpanProcessor(span_exporter))
|
||||
|
||||
set_tracer_provider(OTLPTraceProvider(SDKTracerProvider(active_span_processor=processor,
|
||||
resource=build_otel_resource())))
|
||||
|
||||
|
||||
def _configure_logfire_exporter(write_token: str, base_url: str = None) -> None:
|
||||
"""Configure the Logfire exporter.
|
||||
Args:
|
||||
write_token: The write token to use.
|
||||
base_url: The base URL to use.
|
||||
**kwargs: Additional keyword arguments to pass to the exporter.
|
||||
"""
|
||||
from opentelemetry.exporter.otlp.proto.http import Compression
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
|
||||
base_url = base_url or "https://logfire-us.pydantic.dev"
|
||||
headers = {'User-Agent': f'logfire/3.14.0', 'Authorization': write_token}
|
||||
session = requests.Session()
|
||||
session.headers.update(headers)
|
||||
return OTLPSpanExporter(
|
||||
endpoint=urljoin(base_url, '/v1/traces'),
|
||||
session=session,
|
||||
compression=Compression.Gzip,
|
||||
)
|
||||
|
||||
|
||||
def _configure_otlp_exporter(base_url: str = None, **kwargs) -> None:
|
||||
"""Configure the OTLP exporter.
|
||||
Args:
|
||||
write_token: The write token to use.
|
||||
base_url: The base URL to use.
|
||||
**kwargs: Additional keyword arguments to pass to the exporter.
|
||||
"""
|
||||
import requests
|
||||
from opentelemetry.exporter.otlp.proto.http import Compression
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
|
||||
otlp_traces_endpoint = os.getenv("OTLP_TRACES_ENDPOINT")
|
||||
base_url = base_url or otlp_traces_endpoint
|
||||
session = requests.Session()
|
||||
return OTLPSpanExporter(
|
||||
endpoint=base_url,
|
||||
session=session,
|
||||
compression=Compression.Gzip,
|
||||
)
|
||||
|
||||
|
||||
def is_valid_attribute_value(k, v):
|
||||
valid = True
|
||||
if not v:
|
||||
valid = False
|
||||
valid = isinstance(v, (str, bool, int, float)) or \
|
||||
(isinstance(v, Sequence) and
|
||||
all(isinstance(i, (str, bool, int, float)) for i in v))
|
||||
if not valid:
|
||||
logger.debug(f"value of attribute[{k}] is invalid: {v}")
|
||||
return valid
|
||||
Reference in New Issue
Block a user