mcp-efficiency-engine 0.1.0 → 0.1.2

package/telemetry/collector/engine.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import uuid
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any
+
+from telemetry.context.propagation import (
+    TelemetryContext,
+    child_context,
+    get_current_context,
+    new_execution_context,
+    reset_current_context,
+    set_current_context,
+)
+from telemetry.events.catalog import CANONICAL_EVENTS
+from telemetry.events.factory import make_event
+from telemetry.metrics.aggregator import MetricsAggregator
+from telemetry.scoring.efficiency import compute_efficiency_score
+from telemetry.tracing.span import SpanRecord
+
+
+def utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+@dataclass
+class _ExecutionScope(AbstractContextManager[TelemetryContext]):
+    collector: "TelemetryCollector"
+    context: TelemetryContext
+    token: Any
+    operation: str
+    start_ts: str
+
+    def __enter__(self) -> TelemetryContext:
+        self.collector.record_event("ExecutionStarted", {"operation": self.operation, "started_at": self.start_ts})
+        return self.context
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        if exc is not None:
+            self.collector.record_event("ExceptionThrown", {"operation": self.operation, "error": str(exc)}, level="ERROR")
+            self.collector.metrics.bump("error_count")
+        self.collector.metrics.snapshot.execution_time_ms = max(
+            0.0,
+            (datetime.fromisoformat(utc_now()) - datetime.fromisoformat(self.start_ts)).total_seconds() * 1000.0,
+        )
+        self.collector.record_metric("efficiency_score", compute_efficiency_score(self.collector.metrics.to_dict()), unit="score")
+        self.collector.record_event("ExecutionFinished", {"operation": self.operation, "finished_at": utc_now()})
+        self.collector.flush()
+        reset_current_context(self.token)
+        return False
+
+
+@dataclass
+class _SpanScope(AbstractContextManager[SpanRecord]):
+    collector: "TelemetryCollector"
+    span: SpanRecord
+    token: Any
+
+    def __enter__(self) -> SpanRecord:
+        self.collector.record_event(
+            "SpanStarted",
+            {
+                "name": self.span.name,
+                "span_id": self.span.span_id,
+                "parent_span_id": self.span.parent_span_id,
+                "kind": self.span.kind,
+            },
+        )
+        return self.span
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        if exc is not None:
+            self.span.end(status="ERROR", error=str(exc))
+            self.collector.metrics.bump("error_count")
+            self.collector.record_event("ExceptionThrown", {"span": self.span.name, "error": str(exc)}, level="ERROR")
+        else:
+            self.span.end(status="OK")
+
+        if self.span.duration_ms is not None:
+            self.collector.metrics.add_duration(self.span.duration_ms)
+
+        self.collector._emit(
+            make_event(
+                event_name="SpanFinished",
+                event_type="span",
+                context=get_current_context(),
+                payload={
+                    "name": self.span.name,
+                    "kind": self.span.kind,
+                    "span_id": self.span.span_id,
+                    "parent_span_id": self.span.parent_span_id,
+                    "trace_id": self.span.trace_id,
+                    "execution_id": self.span.execution_id,
+                    "started_at": self.span.started_at,
+                    "ended_at": self.span.ended_at,
+                    "duration_ms": self.span.duration_ms,
+                    "status": self.span.status,
+                    "attributes": self.span.attributes,
+                    "events": self.span.events,
+                    "error": self.span.error,
+                },
+            )
+        )
+        reset_current_context(self.token)
+        return False
+
+
+class TelemetryCollector:
+    def __init__(
+        self,
+        *,
+        exporters: list[Any],
+        enabled: bool = True,
+        batch_size: int = 25,
+    ) -> None:
+        self.enabled = bool(enabled)
+        self.exporters = exporters
+        self.batch_size = max(1, int(batch_size))
+        self.queue: list[dict[str, Any]] = []
+        self.metrics = MetricsAggregator()
+        self.warning_buffer: list[str] = []
+
+    def start_execution(
+        self,
+        *,
+        operation: str,
+        request_id: str | None = None,
+        session_id: str = "default-session",
+        correlation_id: str | None = None,
+        execution_id: str | None = None,
+        user: str | None = None,
+        provider: str | None = None,
+        model: str | None = None,
+    ) -> AbstractContextManager[TelemetryContext]:
+        request = request_id or str(uuid.uuid4())
+        correlation = correlation_id or request
+        execution = execution_id or str(uuid.uuid4())
+        context = new_execution_context(
+            request_id=request,
+            session_id=session_id,
+            correlation_id=correlation,
+            execution_id=execution,
+            user=user,
+            provider=provider,
+            model=model,
+        )
+        token = set_current_context(context)
+        return _ExecutionScope(self, context, token, operation, utc_now())
+
+    def start_span(
+        self,
+        *,
+        name: str,
+        kind: str = "INTERNAL",
+        attributes: dict[str, Any] | None = None,
+    ) -> AbstractContextManager[SpanRecord]:
+        current = get_current_context()
+        if current is None:
+            current = new_execution_context(
+                request_id=str(uuid.uuid4()),
+                session_id="implicit-session",
+                correlation_id=str(uuid.uuid4()),
+                execution_id=str(uuid.uuid4()),
+            )
+            token = set_current_context(current)
+        else:
+            child = child_context(current)
+            token = set_current_context(child)
+            current = child
+
+        span = SpanRecord(
+            name=name,
+            kind=kind,
+            span_id=current.span_id,
+            parent_span_id=current.parent_span_id,
+            trace_id=current.trace_id,
+            execution_id=current.execution_id,
+            started_at=utc_now(),
+            attributes=attributes or {},
+        )
+        return _SpanScope(self, span, token)
+
+    def record_event(self, event_name: str, payload: dict[str, Any] | None = None, *, level: str = "INFO") -> None:
+        canonical = event_name if event_name in CANONICAL_EVENTS else "WarningGenerated"
+        if canonical == "WarningGenerated" and event_name not in CANONICAL_EVENTS:
+            payload = {"original_event": event_name, **(payload or {})}
+        if level.upper() == "WARNING":
+            self.metrics.bump("warning_count")
+        self._emit(make_event(event_name=canonical, event_type="event", context=get_current_context(), payload=payload, level=level.upper()))
+
+    def record_metric(self, metric_name: str, value: float, *, unit: str = "count", tags: dict[str, Any] | None = None) -> None:
+        payload = {"name": metric_name, "value": value, "unit": unit, "tags": tags or {}}
+        if metric_name == "estimated_cost":
+            self.metrics.add_cost(float(value))
+        self._emit(make_event(event_name="MetricCalculated", event_type="metric", context=get_current_context(), payload=payload))
+
+    def record_usage(self, *, input_tokens: int, output_tokens: int, estimated_cost_usd: float) -> None:
+        self.metrics.add_tokens(input_tokens, output_tokens)
+        self.metrics.add_cost(estimated_cost_usd)
+        self.record_metric("tokens_input", float(input_tokens), unit="tokens")
+        self.record_metric("tokens_output", float(output_tokens), unit="tokens")
+        self.record_metric("total_tokens", float(input_tokens + output_tokens), unit="tokens")
+        self.record_metric("estimated_cost", float(estimated_cost_usd), unit="usd")
+
+    def _emit(self, record: dict[str, Any]) -> None:
+        if not self.enabled:
+            return
+        self.queue.append(record)
+        if len(self.queue) >= self.batch_size:
+            self.flush()
+
+    def flush(self) -> None:
+        if not self.enabled or not self.queue:
+            return
+
+        batch = self.queue[:]
+        self.queue = []
+        for exporter in self.exporters:
+            try:
+                exporter_started = make_event(
+                    event_name="ExporterStarted",
+                    event_type="event",
+                    context=get_current_context(),
+                    payload={"exporter": exporter.name},
+                )
+                batch_with_started = [exporter_started, *batch]
+                exporter.export(batch_with_started)
+                exporter.flush()
+                exporter_finished = make_event(
+                    event_name="ExporterFinished",
+                    event_type="event",
+                    context=get_current_context(),
+                    payload={"exporter": exporter.name, "success": True},
+                )
+                exporter.export([exporter_finished])
+            except Exception as exc:
+                self.warning_buffer.append(f"exporter_failed:{exporter.name}:{exc}")
+                warning = make_event(
+                    event_name="WarningGenerated",
+                    event_type="event",
+                    context=get_current_context(),
+                    payload={"warning": "exporter_failed", "exporter": exporter.name, "error": str(exc)},
+                    level="WARNING",
+                )
+                for backup in self.exporters:
+                    if backup.name == exporter.name:
+                        continue
+                    try:
+                        backup.export([warning])
+                    except Exception:
+                        continue
+
+    def shutdown(self) -> None:
+        self.flush()
+        for exporter in self.exporters:
+            try:
+                exporter.shutdown()
+            except Exception:
+                continue

package/telemetry/config.json

@@ -0,0 +1,20 @@
+{
+  "telemetry": {
+    "enabled": true,
+    "batch_size": 100,
+    "telemetry_dir": ".telemetry",
+    "exporters": [
+      "console",
+      "json"
+    ]
+  },
+  "langsmith": {
+    "enabled": false,
+    "api_key": "",
+    "project": "",
+    "endpoint": "",
+    "high_signal_only": true,
+    "min_span_duration_ms": 100,
+    "emit_execution_summary": true
+  }
+}

package/telemetry/config.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+def _parse_bool(value: str | bool | None, default: bool) -> bool:
+    if isinstance(value, bool):
+        return value
+    if value is None:
+        return default
+    return str(value).strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _parse_float(value: str | float | int | None, default: float) -> float:
+    if isinstance(value, (int, float)):
+        return float(value)
+    if value is None:
+        return default
+    try:
+        return float(str(value).strip())
+    except Exception:
+        return default
+
+
+@dataclass
+class LangSmithConfig:
+    enabled: bool = False
+    api_key: str = ""
+    project: str = ""
+    endpoint: str = ""
+    workspace_id: str = ""
+    high_signal_only: bool = True
+    min_span_duration_ms: float = 100.0
+    emit_execution_summary: bool = True
+
+
+@dataclass
+class TelemetryConfig:
+    enabled: bool = True
+    exporters: list[str] = field(default_factory=lambda: ["console", "json"])
+    batch_size: int = 25
+    telemetry_dir: str = ".telemetry"
+    langsmith: LangSmithConfig = field(default_factory=LangSmithConfig)
+
+
+def _load_json_file(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        return {}
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        return {}
+    return payload if isinstance(payload, dict) else {}
+
+
+def _load_dotenv_file(path: Path) -> dict[str, str]:
+    if not path.exists():
+        return {}
+
+    values: dict[str, str] = {}
+    for raw_line in path.read_text(encoding="utf-8").splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+
+        key, value = line.split("=", 1)
+        env_key = key.strip()
+        if not env_key:
+            continue
+
+        env_value = value.strip()
+        if len(env_value) >= 2 and env_value[0] == env_value[-1] and env_value[0] in {'"', "'"}:
+            env_value = env_value[1:-1]
+        values[env_key] = env_value
+
+    return values
+
+
+def load_config(repo_root: Path) -> TelemetryConfig:
+    raw = _load_json_file(repo_root / "telemetry" / "config.json")
+    dotenv = _load_dotenv_file(repo_root / ".env")
+    telemetry_raw = raw.get("telemetry", {}) if isinstance(raw.get("telemetry"), dict) else {}
+    langsmith_raw = raw.get("langsmith", {}) if isinstance(raw.get("langsmith"), dict) else {}
+
+    def _get_env(name: str, default: str | None = None) -> str | None:
+        if name in os.environ:
+            value = os.environ.get(name)
+            return None if value is None else str(value)
+        return dotenv.get(name, default)
+
+    exporters = telemetry_raw.get("exporters", ["console", "json"])
+    if not isinstance(exporters, list):
+        exporters = ["console", "json"]
+
+    env_exporters = (_get_env("TELEMETRY_EXPORTERS", "") or "").strip()
+    if env_exporters:
+        exporters = [item.strip().lower() for item in env_exporters.split(",") if item.strip()]
+
+    cfg = TelemetryConfig(
+        enabled=_parse_bool(_get_env("TELEMETRY_ENABLED"), _parse_bool(telemetry_raw.get("enabled"), True)),
+        exporters=[str(item).lower() for item in exporters],
+        batch_size=max(1, int((_get_env("TELEMETRY_BATCH_SIZE", str(telemetry_raw.get("batch_size", 25))) or str(telemetry_raw.get("batch_size", 25))))),
+        telemetry_dir=str((_get_env("TELEMETRY_DIR", str(telemetry_raw.get("telemetry_dir", ".telemetry"))) or str(telemetry_raw.get("telemetry_dir", ".telemetry")))),
+        langsmith=LangSmithConfig(
+            enabled=_parse_bool(
+                _get_env("LANGSMITH_ENABLED", _get_env("LANGSMITH_TRACING")),
+                _parse_bool(langsmith_raw.get("enabled"), False),
+            ),
+            api_key=str((_get_env("LANGSMITH_API_KEY", str(langsmith_raw.get("api_key", ""))) or str(langsmith_raw.get("api_key", "")))),
+            project=str((_get_env("LANGSMITH_PROJECT", str(langsmith_raw.get("project", ""))) or str(langsmith_raw.get("project", "")))),
+            endpoint=str((_get_env("LANGSMITH_ENDPOINT", str(langsmith_raw.get("endpoint", ""))) or str(langsmith_raw.get("endpoint", "")))),
+            workspace_id=str((_get_env("LANGSMITH_WORKSPACE_ID", str(langsmith_raw.get("workspace_id", ""))) or str(langsmith_raw.get("workspace_id", "")))),
+            high_signal_only=_parse_bool(
+                _get_env("LANGSMITH_HIGH_SIGNAL_ONLY"),
+                _parse_bool(langsmith_raw.get("high_signal_only"), True),
+            ),
+            min_span_duration_ms=_parse_float(
+                _get_env("LANGSMITH_MIN_SPAN_DURATION_MS"),
+                _parse_float(langsmith_raw.get("min_span_duration_ms"), 100.0),
+            ),
+            emit_execution_summary=_parse_bool(
+                _get_env("LANGSMITH_EMIT_EXECUTION_SUMMARY"),
+                _parse_bool(langsmith_raw.get("emit_execution_summary"), True),
+            ),
+        ),
+    )
+    return cfg

package/telemetry/context/propagation.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import contextvars
+import secrets
+from dataclasses import dataclass, replace
+from datetime import datetime, timezone
+
+
+def utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def new_trace_id() -> str:
+    # 16 bytes -> 32 lowercase hex chars, aligned with OTel expectations.
+    return secrets.token_hex(16)
+
+
+def new_span_id() -> str:
+    # 8 bytes -> 16 lowercase hex chars, aligned with OTel expectations.
+    return secrets.token_hex(8)
+
+
+@dataclass(frozen=True)
+class TelemetryContext:
+    execution_id: str
+    trace_id: str
+    span_id: str
+    parent_span_id: str | None
+    request_id: str
+    session_id: str
+    correlation_id: str
+    timestamp: str
+    user: str | None = None
+    provider: str | None = None
+    model: str | None = None
+
+
+_context_var: contextvars.ContextVar[TelemetryContext | None] = contextvars.ContextVar("telemetry_context", default=None)
+
+
+def get_current_context() -> TelemetryContext | None:
+    return _context_var.get()
+
+
+def set_current_context(context: TelemetryContext | None) -> contextvars.Token[TelemetryContext | None]:
+    return _context_var.set(context)
+
+
+def reset_current_context(token: contextvars.Token[TelemetryContext | None]) -> None:
+    _context_var.reset(token)
+
+
+def new_execution_context(
+    *,
+    request_id: str,
+    session_id: str,
+    correlation_id: str,
+    execution_id: str,
+    user: str | None = None,
+    provider: str | None = None,
+    model: str | None = None,
+) -> TelemetryContext:
+    root_span_id = new_span_id()
+    return TelemetryContext(
+        execution_id=execution_id,
+        trace_id=new_trace_id(),
+        span_id=root_span_id,
+        parent_span_id=None,
+        request_id=request_id,
+        session_id=session_id,
+        correlation_id=correlation_id,
+        timestamp=utc_now(),
+        user=user,
+        provider=provider,
+        model=model,
+    )
+
+
+def child_context(parent: TelemetryContext) -> TelemetryContext:
+    return replace(
+        parent,
+        span_id=new_span_id(),
+        parent_span_id=parent.span_id,
+        timestamp=utc_now(),
+    )

package/telemetry/events/catalog.py

@@ -0,0 +1,24 @@
+CANONICAL_EVENTS = {
+    "ExecutionStarted",
+    "ExecutionFinished",
+    "ToolStarted",
+    "ToolFinished",
+    "PromptGenerated",
+    "PromptOptimized",
+    "ContextOptimized",
+    "CacheHit",
+    "CacheMiss",
+    "LLMRequest",
+    "LLMResponse",
+    "ProviderRequest",
+    "ProviderResponse",
+    "ExporterStarted",
+    "ExporterFinished",
+    "ExceptionThrown",
+    "WarningGenerated",
+    "MetricCalculated",
+    "KnowledgeGenerated",
+    "SpanStarted",
+    "SpanFinished",
+    "RoutingResolved",
+}

package/telemetry/events/factory.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any
+
+from telemetry.context.propagation import TelemetryContext
+
+
+def utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def make_event(
+    *,
+    event_name: str,
+    event_type: str,
+    context: TelemetryContext | None,
+    payload: dict[str, Any] | None = None,
+    level: str = "INFO",
+) -> dict[str, Any]:
+    data: dict[str, Any] = {
+        "event_schema_version": "1.0.0",
+        "timestamp": utc_now(),
+        "type": event_type,
+        "event_name": event_name,
+        "level": level,
+        "payload": payload or {},
+    }
+    if context is not None:
+        data["context"] = {
+            "execution_id": context.execution_id,
+            "trace_id": context.trace_id,
+            "span_id": context.span_id,
+            "parent_span_id": context.parent_span_id,
+            "request_id": context.request_id,
+            "session_id": context.session_id,
+            "correlation_id": context.correlation_id,
+            "user": context.user,
+            "provider": context.provider,
+            "model": context.model,
+        }
+    return data

package/telemetry/exporters/base.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class ITelemetryExporter(Protocol):
+    name: str
+
+    def export(self, records: list[dict]) -> None:
+        ...
+
+    def flush(self) -> None:
+        ...
+
+    def shutdown(self) -> None:
+        ...

package/telemetry/exporters/console/exporter.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import json
+
+
+class ConsoleExporter:
+    name = "console"
+
+    def export(self, records: list[dict]) -> None:
+        for record in records:
+            event_name = str(record.get("event_name", "unknown"))
+            event_type = str(record.get("type", "event"))
+            level = str(record.get("level", "INFO"))
+            context = record.get("context", {})
+            execution_id = context.get("execution_id", "n/a")
+            print(f"[telemetry][{level}][{event_type}] {event_name} execution={execution_id}")
+            if event_type == "span":
+                payload = record.get("payload", {})
+                summary = {
+                    "name": payload.get("name"),
+                    "duration_ms": payload.get("duration_ms"),
+                    "status": payload.get("status"),
+                }
+                print(f"[telemetry][span] {json.dumps(summary, ensure_ascii=False)}")
+
+    def flush(self) -> None:
+        return
+
+    def shutdown(self) -> None:
+        return

package/telemetry/exporters/json/exporter.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+class JsonExporter:
+    name = "json"
+
+    def __init__(self, telemetry_dir: Path) -> None:
+        self.telemetry_dir = telemetry_dir
+        self.telemetry_dir.mkdir(parents=True, exist_ok=True)
+        self.traces_file = self.telemetry_dir / "traces.jsonl"
+        self.metrics_file = self.telemetry_dir / "metrics.jsonl"
+
+    def _append_lines(self, file_path: Path, entries: list[dict]) -> None:
+        if not entries:
+            return
+        with file_path.open("a", encoding="utf-8") as handle:
+            for entry in entries:
+                handle.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")))
+                handle.write("\n")
+
+    def export(self, records: list[dict]) -> None:
+        traces = [r for r in records if str(r.get("type", "")) in {"event", "span"}]
+        metrics = [r for r in records if str(r.get("type", "")) == "metric"]
+        self._append_lines(self.traces_file, traces)
+        self._append_lines(self.metrics_file, metrics)
+
+    def flush(self) -> None:
+        return
+
+    def shutdown(self) -> None:
+        return