""" Configurable log sink for execution-trace / audit persistence. The agent already accumulates a full execution trace in-memory (`WorkflowEngine.trace_logs` + `message_history`, surfaced by `CodeAgent.get_trace()`): prompts, tool invocations, generated code, and dataset-load steps. This module abstracts *where* that trace JSON is written so the destination is a CONFIG value, not hardcoded — prep for a future migration to OHSU-managed AWS. Selection is by the `LOG_SINK` env var (`local` | `hf` | `s3`), default `local`. Every sink implements one method:: persist_trace(run_id: str, trace: dict) -> str | None returning a location string (path / URL) on success, or None on failure / no-op. Callers wrap the call so a logging failure NEVER crashes a user request. Env vars (all optional; sensible defaults): LOG_SINK local | hf | s3 (default: local) local sink: LOG_SINK_LOCAL_DIR output directory (default: ./run_logs) hf sink (reproduces today's HuggingFace-dataset behavior): LOG_SINK_HF_DATASET dataset repo id (default: anne-voigt/decoupleRpy_results) decouplerpy_results_token HF write token (existing var; no token => no-op) s3 sink (OHSU-managed AWS migration target; see ADR-0009): LOG_SINK_S3_BUCKET bucket name (required; no bucket => error) LOG_SINK_S3_REGION AWS region (default: us-west-2) LOG_SINK_S3_PREFIX key prefix (default: runs) The s3 sink writes trace blobs only. At-rest encryption (SSE-KMS), versioning, lifecycle/retention, and access are enforced by the BUCKET (bucket policy / IAM / default encryption) per OHSU standard — the sink manages neither keys nor lifecycle (ADR-0008 posture: the store governs retention). Credentials come from the task's IAM role via boto3's default credential chain — never long-lived keys in env. boto3 is imported lazily so this module imports even where boto3 is absent (the HF-hosted `hf`/`local` posture does not need it; boto3 becomes a runtime requirement at AWS cutover). """ import json import os from datetime import datetime # --------------------------------------------------------------------------- # # Base interface # --------------------------------------------------------------------------- # class LogSink: """Common interface for trace persistence.""" name = "base" def persist_trace(self, run_id: str, trace: dict) -> "str | None": """Persist the trace dict for `run_id`. Return a location string or None.""" raise NotImplementedError # --------------------------------------------------------------------------- # # local (DEFAULT) — always-on, no external deps # --------------------------------------------------------------------------- # class LocalLogSink(LogSink): """Write the trace JSON to a configurable directory on local disk.""" name = "local" def __init__(self, directory: str = None): self.directory = directory or os.environ.get("LOG_SINK_LOCAL_DIR", "./run_logs") def persist_trace(self, run_id: str, trace: dict) -> "str | None": os.makedirs(self.directory, exist_ok=True) path = os.path.join(self.directory, f"{run_id}_trace.json") with open(path, "w", encoding="utf-8") as f: json.dump(trace, f, indent=2, ensure_ascii=False, default=str) print(f"[log_sink:local] Trace saved: {path}") return path # --------------------------------------------------------------------------- # # hf — refactor of the existing HuggingFace-dataset behavior # --------------------------------------------------------------------------- # class HFLogSink(LogSink): """Persist the trace JSON to an HF Dataset repo under runs//trace.json. Preserves today's behavior: token comes from `decouplerpy_results_token` (same var `HFResultsStorage` uses); no token => no-op (returns None). huggingface_hub is imported lazily so a missing dep never breaks import. """ name = "hf" def __init__(self, repo_id: str = None): self.repo_id = repo_id or os.environ.get( "LOG_SINK_HF_DATASET", "anne-voigt/decoupleRpy_results" ) self._api = None def _get_api(self): if self._api is None: token = os.environ.get("decouplerpy_results_token") if not token: print("[log_sink:hf] decouplerpy_results_token not set — uploads disabled") return None from huggingface_hub import HfApi # lazy self._api = HfApi(token=token) return self._api def persist_trace(self, run_id: str, trace: dict) -> "str | None": import tempfile api = self._get_api() if api is None: return None remote_path = f"runs/{run_id}/trace.json" tmp_path = None try: with tempfile.NamedTemporaryFile( mode="w", suffix=".json", delete=False, encoding="utf-8" ) as f: json.dump(trace, f, indent=2, ensure_ascii=False, default=str) tmp_path = f.name api.upload_file( path_or_fileobj=tmp_path, path_in_repo=remote_path, repo_id=self.repo_id, repo_type="dataset", ) finally: if tmp_path and os.path.exists(tmp_path): os.unlink(tmp_path) url = f"https://huggingface.co/datasets/{self.repo_id}/blob/main/{remote_path}" print(f"[log_sink:hf] Trace saved: {url}") return url # --------------------------------------------------------------------------- # # s3 — OHSU-managed AWS migration target (ADR-0009) # --------------------------------------------------------------------------- # class S3LogSink(LogSink): """Write the trace JSON to S3 via boto3 `put_object`. Target key: ``{LOG_SINK_S3_PREFIX}/{run_id}/trace.json`` in ``LOG_SINK_S3_BUCKET``; returns the ``s3://…`` URI on success. boto3 is imported LAZILY (and the client cached), so this module imports cleanly where boto3 is absent — the HF-hosted `hf`/`local` posture never touches it. Encryption / retention / access are the bucket's job, not the sink's (ADR-0009): no SSE or ACL params are set here, so the bucket's default encryption and lifecycle policy govern the object. Credentials come from boto3's default chain (the task IAM role), never env keys. """ name = "s3" def __init__(self): self.bucket = os.environ.get("LOG_SINK_S3_BUCKET") self.region = os.environ.get("LOG_SINK_S3_REGION", "us-west-2") # Normalize the prefix: drop surrounding slashes so key joining never # yields a leading/double slash (e.g. "runs/" or "/runs" -> "runs"). self.prefix = os.environ.get("LOG_SINK_S3_PREFIX", "runs").strip("/") self._client = None def _get_client(self): if self._client is None: import boto3 # lazy — only needed on the s3 path self._client = boto3.client("s3", region_name=self.region) return self._client def persist_trace(self, run_id: str, trace: dict) -> "str | None": if not self.bucket: # LOG_SINK=s3 was selected but no bucket configured — a config # error. Raise loudly; persist_trace_safe still swallows it so a # run never crashes on a logging misconfiguration. raise ValueError( "LOG_SINK=s3 requires LOG_SINK_S3_BUCKET to be set " "(the OHSU-managed audit bucket). See ADR-0009." ) key = f"{self.prefix}/{run_id}/trace.json" if self.prefix else f"{run_id}/trace.json" body = json.dumps(trace, indent=2, ensure_ascii=False, default=str).encode("utf-8") self._get_client().put_object( Bucket=self.bucket, Key=key, Body=body, ContentType="application/json", ) uri = f"s3://{self.bucket}/{key}" print(f"[log_sink:s3] Trace saved: {uri}") return uri # --------------------------------------------------------------------------- # # Factory # --------------------------------------------------------------------------- # _SINKS = { "local": LocalLogSink, "hf": HFLogSink, "s3": S3LogSink, } def get_log_sink(kind: str = None) -> LogSink: """Return the configured LogSink. `kind` overrides the `LOG_SINK` env var; default is `local`. An unknown value raises ValueError (fail loud on misconfiguration, before any run). """ kind = (kind or os.environ.get("LOG_SINK", "local")).strip().lower() if kind not in _SINKS: raise ValueError(f"Unknown LOG_SINK={kind!r}. Valid values: {sorted(_SINKS)}.") return _SINKS[kind]() def persist_trace_safe(sink: LogSink, run_id: str, trace: dict) -> "str | None": """Persist a trace through `sink`, swallowing any error (warn, don't crash). A logging failure must NEVER take down a user request, so this wraps the write in try/except and returns None on any failure. """ if run_id is None: run_id = datetime.now().strftime("%Y%m%d_%H%M%S") try: return sink.persist_trace(run_id, trace) except Exception as e: # noqa: BLE001 — logging must never crash a request print(f"[log_sink:{getattr(sink, 'name', '?')}] persist_trace failed: {e}") return None