271 lines
11 KiB
Python
271 lines
11 KiB
Python
"""Small, credential-free request snapshots in the existing task trace store."""
|
|
|
|
import base64
|
|
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
|
|
|
|
from artemis.data_engine.context_vars import CURRENT_NODE_NAME, CURRENT_TRACE_ID
|
|
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
|
|
|
|
logger = logging.getLogger(__name__)
|
|
ACTIVE_ATTEMPT: ContextVar["RequestAttempt | None"] = ContextVar("llm_attempt", default=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):
|
|
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.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),
|
|
"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,
|
|
"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):
|
|
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 headers(self, response):
|
|
self.data.update(
|
|
headers_at=time.time(),
|
|
http_status=response.status_code,
|
|
server_request_id=response.headers.get("x-request-id")
|
|
or response.headers.get("request-id"),
|
|
phase="waiting_output",
|
|
)
|
|
self.publish()
|
|
|
|
def chunk(self, chunk):
|
|
response_metadata = getattr(chunk, "response_metadata", {}) or {}
|
|
for field in ("status", "finish_reason"):
|
|
value = response_metadata.get(field)
|
|
if isinstance(value, str) and re.fullmatch(r"[\w.-]{1,64}", value):
|
|
self.data["response_" + field] = value
|
|
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 lengths/names, never argument bodies.
|
|
for tool in tool_chunks:
|
|
index = tool.get("index")
|
|
key = index if index is not None else tool.get("id") or "unknown"
|
|
stats = self.tools.setdefault(key, {"name": None, "argument_chars": 0})
|
|
name = tool.get("name")
|
|
if isinstance(name, str) and re.fullmatch(r"[\w.:-]{1,128}", name):
|
|
stats["name"] = name
|
|
arguments = tool.get("args")
|
|
if isinstance(arguments, str):
|
|
stats["argument_chars"] += len(arguments)
|
|
self.data["tool_argument_chars"] += len(arguments)
|
|
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 error is not None:
|
|
if isinstance(error, ToolArgumentOverflow):
|
|
self.data["stream_guard"] = error.details
|
|
self.data.update(
|
|
error=error_text(error, self.data["timeout_seconds"]),
|
|
error_type=type(error).__name__,
|
|
category=category,
|
|
failure_phase=self.data["phase"],
|
|
)
|
|
status = getattr(error, "status_code", None)
|
|
if isinstance(status, int):
|
|
self.data["http_status"] = status
|
|
if response is not None:
|
|
# 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):
|
|
if self.review and self.review.expired and not isinstance(error, KeyboardInterrupt):
|
|
self.finish(
|
|
"failed",
|
|
error=FinalReviewBudgetExceeded(self.review.total_seconds),
|
|
category="final_review_budget_exhausted",
|
|
)
|
|
else:
|
|
self.finish("cancelled", error=error, 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
|
|
request.headers.setdefault("X-Client-Request-Id", attempt.data["request_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(),
|
|
routing_header_fingerprints=routing,
|
|
client_request_id_matches=request.headers.get("X-Client-Request-Id")
|
|
== attempt.data["request_id"],
|
|
)
|
|
attempt.publish()
|
|
|
|
|
|
async def observe_response_headers(response):
|
|
"""httpx hook: only allowlisted headers, never body/auth/request headers."""
|
|
attempt = ACTIVE_ATTEMPT.get()
|
|
if attempt is not None:
|
|
attempt.headers(response)
|