419 lines
18 KiB
Python
419 lines
18 KiB
Python
"""Small, credential-free request snapshots in the existing task trace store."""
|
|
|
|
import base64
|
|
import asyncio
|
|
from contextvars import ContextVar
|
|
from copy import deepcopy
|
|
import io
|
|
import hashlib
|
|
import logging
|
|
import re
|
|
import time
|
|
from uuid import uuid4
|
|
|
|
from PIL import Image
|
|
import httpx
|
|
import httpx2
|
|
|
|
from artemis.data_engine.context_vars import CURRENT_NODE_NAME, CURRENT_TRACE_ID
|
|
from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_TRACKED_TOOLS
|
|
from artemis.llm.final_review import ACTIVE_FINAL_REVIEW, FinalReviewBudgetExceeded
|
|
from artemis.llm.reliability import LLMExhaustedError
|
|
from artemis.llm.usage import usage_from_message
|
|
from artemis.llm.stream_guard import ToolArgumentOverflow
|
|
from artemis.llm.response_guard import ResponseIntegrityError
|
|
from artemis.runtime.cancellation import CancellationReason, cancellation_reason
|
|
from artemis.llm.responses_diagnostics import ObservedResponseStream
|
|
from artemis.llm.tool_constraints import request_constraints
|
|
|
|
logger = logging.getLogger(__name__)
|
|
ACTIVE_ATTEMPT: ContextVar["RequestAttempt | None"] = ContextVar("llm_attempt", default=None)
|
|
|
|
|
|
def _identifier(value):
|
|
"""Keep bounded, printable correlation IDs; never arbitrary header text."""
|
|
if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_.:/-]{1,256}", value):
|
|
return value
|
|
return None
|
|
|
|
|
|
def error_text(error: BaseException, timeout: float | None = None) -> str:
|
|
if isinstance(error, FinalReviewBudgetExceeded):
|
|
return str(error)
|
|
if isinstance(error, TimeoutError):
|
|
return f"模型请求在 {timeout:g} 秒内未完成" if timeout else "模型请求超时"
|
|
text = str(error).strip() or type(error).__name__
|
|
text = re.sub(
|
|
r"(?i)(bearer\s+|(?:api[_-]?key|token|authorization)[\"']?\s*[:=]\s*[\"']?)[^\s\"',;]+",
|
|
r"\1[REDACTED]",
|
|
text,
|
|
)
|
|
text = re.sub(r"\bsk-[\w-]+", "[REDACTED]", text)
|
|
text = re.sub(r"https?://[^\s\"']+", "[endpoint]", text)
|
|
return text[:1000]
|
|
|
|
|
|
def final_review_error_text(error: Exception) -> str:
|
|
if isinstance(error, FinalReviewBudgetExceeded):
|
|
return str(error)
|
|
if isinstance(error, LLMExhaustedError):
|
|
return f"验收模型请求重试已用尽（{error.failure.category.value}）：{error_text(error)}"
|
|
return f"{type(error).__name__}: {error_text(error)}"
|
|
|
|
|
|
def input_images(args: tuple, kwargs: dict) -> list[dict]:
|
|
"""Inspect only inline image headers; never fetch URLs or retain image data."""
|
|
images = []
|
|
messages = args[0] if args else kwargs.get("input", [])
|
|
if not isinstance(messages, (list, tuple)):
|
|
return images
|
|
for message in messages:
|
|
content = getattr(message, "content", None)
|
|
if not isinstance(content, list):
|
|
continue
|
|
for block in content:
|
|
if not isinstance(block, dict) or block.get("type") not in {
|
|
"image_url",
|
|
"input_image",
|
|
"image",
|
|
}:
|
|
continue
|
|
value = block.get("image_url", block.get("url", ""))
|
|
value = value.get("url", "") if isinstance(value, dict) else value
|
|
info = {"width": None, "height": None}
|
|
if isinstance(value, str) and value.startswith("data:image/") and ";base64," in value:
|
|
try:
|
|
# JPEG/PNG dimensions normally fit in this prefix. Unknown is
|
|
# preferable to decoding large images just for diagnostics.
|
|
encoded = value.split(",", 1)[1][:131072]
|
|
with Image.open(io.BytesIO(base64.b64decode(encoded))) as image:
|
|
info.update(width=image.width, height=image.height)
|
|
except (ValueError, OSError):
|
|
pass
|
|
images.append(info)
|
|
return images
|
|
|
|
|
|
class RequestAttempt:
|
|
def __init__(
|
|
self,
|
|
engine,
|
|
*,
|
|
request_id,
|
|
attempt,
|
|
timeout,
|
|
max_retries,
|
|
metadata,
|
|
previous_attempt_id=None,
|
|
):
|
|
self.engine = engine
|
|
self.trace_id = uuid4()
|
|
self.step_id = getattr(engine, "current_step_id", None)
|
|
self.parent_id = CURRENT_TRACE_ID.get()
|
|
self.clock = time.monotonic()
|
|
self.last_publish = self.clock
|
|
self.tools: dict[int | str, dict] = {}
|
|
self.tool_fingerprints: dict[int | str, ArgumentFingerprint] = {}
|
|
self.wire_streams: list[ObservedResponseStream] = []
|
|
self.review = ACTIVE_FINAL_REVIEW.get()
|
|
if self.review:
|
|
self.review.request_started(request_id, attempt)
|
|
self.data = {
|
|
"images": [],
|
|
**metadata,
|
|
"request_id": request_id,
|
|
"attempt_id": str(self.trace_id),
|
|
"previous_attempt_id": previous_attempt_id,
|
|
"attempt": attempt,
|
|
"max_retries": max_retries,
|
|
"timeout_seconds": timeout,
|
|
"role": CURRENT_NODE_NAME.get() or metadata.get("role") or "model",
|
|
"started_at": time.time(),
|
|
"phase": "waiting_response",
|
|
"sequence": 0,
|
|
"usage": None,
|
|
"output_chunks": 0,
|
|
"tool_argument_chars": 0,
|
|
"tool_observation_layer": "langchain_tool_call_chunks",
|
|
"argument_fingerprint_version": 1,
|
|
"tool_stats_truncated": False,
|
|
"tools": [],
|
|
"step_id": str(self.step_id) if self.step_id else None,
|
|
"parent_trace_id": str(self.parent_id) if self.parent_id else None,
|
|
}
|
|
self.publish()
|
|
|
|
def publish(self):
|
|
if self.wire_streams:
|
|
self.data["wire_responses"] = [s.observation.snapshot() for s in self.wire_streams]
|
|
self.data["sequence"] += 1
|
|
self.data["updated_at"] = time.time()
|
|
self.data["elapsed_seconds"] = round(time.monotonic() - self.clock, 3)
|
|
self.last_publish = time.monotonic()
|
|
if self.review:
|
|
self.data["final_review"] = self.review.snapshot()
|
|
if not self.engine or not getattr(self.engine, "current_session_id", None):
|
|
return
|
|
try:
|
|
phase = self.data["phase"]
|
|
status = (
|
|
"success"
|
|
if phase == "completed"
|
|
else "failed"
|
|
if self.data.get("ended_at")
|
|
else "running"
|
|
)
|
|
self.engine.record_trace(
|
|
type="llm_diagnostic",
|
|
name="llm_request",
|
|
trace_id=self.trace_id,
|
|
step_id=self.step_id,
|
|
parent_trace_id=self.parent_id,
|
|
payload=deepcopy(self.data),
|
|
status=status,
|
|
duration=self.data["elapsed_seconds"],
|
|
)
|
|
except (OSError, RuntimeError, ValueError, TypeError, AttributeError):
|
|
logger.warning("Could not persist model request diagnostics", exc_info=True)
|
|
|
|
def publish_if_due(self):
|
|
if time.monotonic() - self.last_publish >= 5:
|
|
self.publish()
|
|
|
|
def stream_cancelled(self):
|
|
self.data.setdefault("cancel_observed_at", time.time())
|
|
self.data.setdefault("cancel_observed_phase", self.data["phase"])
|
|
self.data.setdefault("cancel_phase", self.data["phase"])
|
|
|
|
def stream_work(self, stage: str, elapsed: float):
|
|
stats = self.data.setdefault("stream_work", {}).setdefault(
|
|
stage, {"calls": 0, "total_seconds": 0.0, "max_seconds": 0.0}
|
|
)
|
|
stats["calls"] += 1
|
|
stats["total_seconds"] += elapsed
|
|
stats["max_seconds"] = max(stats["max_seconds"], elapsed)
|
|
|
|
async def close_streams(self):
|
|
try:
|
|
async with asyncio.timeout(1):
|
|
for stream in self.wire_streams:
|
|
await stream.aclose()
|
|
except (TimeoutError, OSError, RuntimeError, httpx.HTTPError, httpx2.HTTPError) as error:
|
|
self.data["stream_close_error_type"] = type(error).__name__
|
|
logger.warning("Model stream cleanup failed: %s", type(error).__name__)
|
|
finally:
|
|
if self.wire_streams:
|
|
self.publish()
|
|
|
|
def headers(self, response):
|
|
self.data.update(
|
|
headers_at=time.time(),
|
|
http_status=response.status_code,
|
|
server_request_id=_identifier(response.headers.get("x-oneapi-request-id"))
|
|
or _identifier(response.headers.get("x-request-id"))
|
|
or _identifier(response.headers.get("request-id")),
|
|
response_x_request_id=_identifier(response.headers.get("x-request-id")),
|
|
server_trace_id=_identifier(response.headers.get("x-trace-id")),
|
|
phase="waiting_output",
|
|
)
|
|
traceparent = response.headers.get("traceparent", "")
|
|
self.data.pop("traceparent", None)
|
|
if re.fullmatch(r"[0-9a-f]{2}-[0-9a-f]{32}-[0-9a-f]{16}-[0-9a-f]{2}", traceparent):
|
|
self.data["traceparent"] = traceparent
|
|
self.publish()
|
|
|
|
def response_metadata(self, response):
|
|
raw = response.get("raw") if isinstance(response, dict) else response
|
|
metadata = getattr(raw, "response_metadata", {})
|
|
metadata = metadata if isinstance(metadata, dict) else {}
|
|
if response_id := _identifier(metadata.get("id")):
|
|
self.data["response_id"] = response_id
|
|
for field in ("status", "finish_reason"):
|
|
value = metadata.get(field)
|
|
if isinstance(value, str) and re.fullmatch(r"[\w.-]{1,64}", value):
|
|
self.data["response_" + field] = value
|
|
|
|
def chunk(self, chunk):
|
|
self.response_metadata(chunk)
|
|
if (usage := usage_from_message(chunk)) is not None:
|
|
self.data["usage"] = usage
|
|
tool_chunks = getattr(chunk, "tool_call_chunks", None) or []
|
|
# Responses exposes the same argument delta in content and tool_call_chunks.
|
|
# Count only the latter; retain bounded fingerprints, never argument bodies.
|
|
for tool in tool_chunks:
|
|
arguments = tool.get("args")
|
|
if isinstance(arguments, str):
|
|
self.data["tool_argument_chars"] += len(arguments)
|
|
if arguments:
|
|
self.data.setdefault("first_tool_arguments_at", time.time())
|
|
self.data["last_tool_arguments_at"] = time.time()
|
|
index = tool.get("index")
|
|
call_id = _identifier(tool.get("id"))
|
|
key = index if index is not None else call_id or "unknown"
|
|
if key not in self.tools and len(self.tools) >= MAX_TRACKED_TOOLS:
|
|
self.data["tool_stats_truncated"] = True
|
|
continue
|
|
stats = self.tools.setdefault(key, {"name": None})
|
|
if call_id:
|
|
stats["call_id"] = call_id
|
|
if index is not None:
|
|
stats["chunk_index"] = index
|
|
name = tool.get("name")
|
|
if isinstance(name, str) and re.fullmatch(r"[\w.:-]{1,128}", name):
|
|
stats["name"] = name
|
|
fingerprint = self.tool_fingerprints.get(key)
|
|
if fingerprint is None:
|
|
fingerprint = self.tool_fingerprints[key] = ArgumentFingerprint()
|
|
if isinstance(arguments, str):
|
|
fingerprint.observe(arguments)
|
|
stats.update(fingerprint.snapshot())
|
|
self.data["tools"] = list(self.tools.values())
|
|
content = getattr(chunk, "content", None)
|
|
substantive = isinstance(content, str) and bool(content)
|
|
if isinstance(content, list):
|
|
substantive = any(
|
|
isinstance(b, dict)
|
|
and (
|
|
b.get("text")
|
|
or b.get("thinking")
|
|
or b.get("summary")
|
|
or b.get("type") in {"tool_call", "function_call"}
|
|
)
|
|
for b in content
|
|
)
|
|
substantive = substantive or bool(tool_chunks) or bool(getattr(chunk, "tool_calls", None))
|
|
if not substantive:
|
|
return
|
|
now = time.time()
|
|
first = not self.data.get("first_output_at")
|
|
self.data.setdefault("first_output_at", now)
|
|
phase = "generating_tool_arguments" if tool_chunks else "receiving_output"
|
|
changed = self.data["phase"] != phase
|
|
self.data.update(last_output_at=now, phase=phase)
|
|
self.data["output_chunks"] += 1
|
|
if first or changed or time.monotonic() - self.last_publish >= 5:
|
|
self.publish()
|
|
|
|
def finish(self, phase, *, response=None, error=None, category=None):
|
|
if category is not None:
|
|
self.data["category"] = category
|
|
if error is not None:
|
|
if isinstance(error, ToolArgumentOverflow):
|
|
self.data["stream_guard"] = error.details
|
|
if isinstance(error, ResponseIntegrityError):
|
|
self.data["response_guard"] = error.details
|
|
self.data.update(
|
|
error=error_text(error, self.data["timeout_seconds"]),
|
|
error_type=type(error).__name__,
|
|
failure_phase=self.data["phase"],
|
|
)
|
|
status = getattr(error, "status_code", None)
|
|
if isinstance(status, int):
|
|
self.data["http_status"] = status
|
|
if request_id := _identifier(getattr(error, "request_id", None)):
|
|
self.data["error_request_id"] = request_id
|
|
if not self.data.get("server_request_id"):
|
|
self.data["server_request_id"] = request_id
|
|
if category == "timeout":
|
|
self.data.setdefault("timeout_scope", "transport")
|
|
if response is not None:
|
|
self.response_metadata(response)
|
|
# The aggregate repeats all streaming deltas; do not count them twice.
|
|
if not self.data.get("first_output_at"):
|
|
self.chunk(response)
|
|
if (usage := usage_from_message(response)) is not None:
|
|
self.data["usage"] = usage
|
|
self.data.update(phase=phase, ended_at=time.time())
|
|
if self.review:
|
|
self.review.request_finished(phase, category)
|
|
self.publish()
|
|
|
|
def retry(self, delay):
|
|
self.data.update(
|
|
phase="retry_wait", retry_at=time.time() + delay, retry_delay_seconds=delay
|
|
)
|
|
if self.review:
|
|
self.review.request_finished("retry_wait")
|
|
self.publish()
|
|
|
|
def cancel(self, error: BaseException | None = None):
|
|
self.stream_cancelled()
|
|
reason = cancellation_reason(error)
|
|
if self.review and self.review.expired and reason is CancellationReason.UNKNOWN:
|
|
self.data["timeout_scope"] = "final_review"
|
|
self.finish(
|
|
"failed",
|
|
error=FinalReviewBudgetExceeded(self.review.total_seconds),
|
|
category="final_review_budget_exhausted",
|
|
)
|
|
else:
|
|
self.data.update(cancel_reason=reason.value, cancel_phase=self.data["phase"])
|
|
if reason is CancellationReason.OUTER_TIMEOUT:
|
|
self.data["timeout_scope"] = "outer"
|
|
else:
|
|
self.data.pop("timeout_scope", None)
|
|
# Preserve a retry's last failure, but never retain arbitrary text
|
|
# from CancelledError.args (which may contain user guidance).
|
|
self.finish("cancelled", category="cancelled")
|
|
|
|
|
|
async def observe_request_headers(request):
|
|
"""Correlate a wire request and record only hashes of routing headers."""
|
|
attempt = ACTIVE_ATTEMPT.get()
|
|
if attempt is None:
|
|
return
|
|
try:
|
|
attempt.data["wire_request_constraints"] = request_constraints(request.content)
|
|
except (httpx.RequestNotRead, httpx2.RequestNotRead):
|
|
attempt.data["wire_request_constraints"] = {"observation": "streaming_body_unavailable"}
|
|
request.headers.setdefault("X-Client-Request-Id", attempt.data["request_id"])
|
|
request.headers.setdefault("X-Client-Attempt-Id", attempt.data["attempt_id"])
|
|
routing = {}
|
|
for name in ("thread-id", "thread_id", "session-id", "session_id"):
|
|
value = request.headers.get(name)
|
|
if value is not None:
|
|
routing[name] = hashlib.sha256(value.encode()).hexdigest()[:16]
|
|
attempt.data.update(
|
|
request_sent_at=time.time(),
|
|
client_request_id=_identifier(request.headers.get("X-Client-Request-Id")),
|
|
client_attempt_id=_identifier(request.headers.get("X-Client-Attempt-Id")),
|
|
routing_header_fingerprints=routing,
|
|
client_request_id_matches=request.headers.get("X-Client-Request-Id")
|
|
== attempt.data["request_id"],
|
|
client_attempt_id_matches=request.headers.get("X-Client-Attempt-Id")
|
|
== attempt.data["attempt_id"],
|
|
)
|
|
attempt.publish()
|
|
|
|
|
|
async def observe_response_headers(response):
|
|
"""Correlate allowlisted headers and passively observe bounded SSE frames."""
|
|
attempt = ACTIVE_ATTEMPT.get()
|
|
if attempt is not None:
|
|
websocket = response.extensions.get("artemis_websocket", False)
|
|
if websocket:
|
|
# 200/SSE is an in-process SDK facade, not a wire HTTP response.
|
|
attempt.data.update(headers_at=time.time(), phase="waiting_output")
|
|
else:
|
|
attempt.headers(response)
|
|
content_type = response.headers.get("content-type", "").split(";", 1)[0].lower()
|
|
if (
|
|
response.status_code == 200
|
|
and content_type == "text/event-stream"
|
|
and isinstance(response.stream, (httpx.AsyncByteStream, httpx2.AsyncByteStream))
|
|
and not isinstance(response.stream, ObservedResponseStream)
|
|
):
|
|
if len(attempt.wire_streams) < 4:
|
|
stream = ObservedResponseStream(
|
|
response.stream,
|
|
attempt,
|
|
encoded=response.headers.get("content-encoding", "identity").lower().strip()
|
|
not in {"", "identity"},
|
|
)
|
|
response.stream = stream
|
|
attempt.wire_streams.append(stream)
|
|
if websocket:
|
|
stream.observation.data["observation_boundary"] = "websocket_event_sse_adapter"
|
|
else:
|
|
attempt.data["wire_streams_truncated"] = True
|