"""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)