|
|
@@ -0,0 +1,487 @@
|
|
|
+"""PostgreSQL repository and best-effort writer for pipeline tracing."""
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import logging
|
|
|
+from typing import Any
|
|
|
+from uuid import UUID
|
|
|
+
|
|
|
+import psycopg2.extras
|
|
|
+
|
|
|
+from core.config import CreationDbConfig
|
|
|
+from core.db_session import transaction
|
|
|
+from pipeline.models import PipelineJob, PipelineRun, PipelineRunEvent, PipelineStage
|
|
|
+from pipeline.tracing import TraceConfig, TraceContext, sanitize_payload
|
|
|
+
|
|
|
+Json = psycopg2.extras.Json
|
|
|
+psycopg2.extras.register_uuid()
|
|
|
+logger = logging.getLogger(__name__)
|
|
|
+
|
|
|
+
|
|
|
+class PostgresPipelineRepository:
|
|
|
+ """Repository for orchestration state and trace ledger rows.
|
|
|
+
|
|
|
+ This class owns no transaction boundary; API dependencies pass a pooled
|
|
|
+ connection, while the best-effort trace writer opens short transactions.
|
|
|
+ """
|
|
|
+
|
|
|
+ def __init__(self, conn: Any):
|
|
|
+ self.conn = conn
|
|
|
+
|
|
|
+ def _one(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any]:
|
|
|
+ with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
|
|
+ cur.execute(sql, params)
|
|
|
+ row = cur.fetchone()
|
|
|
+ if row is None:
|
|
|
+ raise RuntimeError("expected one row, got none")
|
|
|
+ return dict(row)
|
|
|
+
|
|
|
+ def _one_or_none(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any] | None:
|
|
|
+ with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
|
|
+ cur.execute(sql, params)
|
|
|
+ row = cur.fetchone()
|
|
|
+ return dict(row) if row else None
|
|
|
+
|
|
|
+ def _all(self, sql: str, params: tuple[Any, ...]) -> list[dict[str, Any]]:
|
|
|
+ with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
|
|
+ cur.execute(sql, params)
|
|
|
+ return [dict(row) for row in cur.fetchall()]
|
|
|
+
|
|
|
+ def create_pipeline_run(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ run_key: str | None = None,
|
|
|
+ batch_id: UUID | None = None,
|
|
|
+ status: str = "pending",
|
|
|
+ current_stage: str | None = None,
|
|
|
+ config: dict[str, Any] | None = None,
|
|
|
+ metadata: dict[str, Any] | None = None,
|
|
|
+ ) -> PipelineRun:
|
|
|
+ row = self._one(
|
|
|
+ """
|
|
|
+ INSERT INTO pipeline_runs(
|
|
|
+ run_key, batch_id, status, current_stage, config, metadata, started_at
|
|
|
+ )
|
|
|
+ VALUES (%s, %s, %s, %s, %s, %s, CASE WHEN %s = 'running' THEN now() ELSE NULL END)
|
|
|
+ ON CONFLICT (run_key) DO UPDATE SET
|
|
|
+ batch_id = COALESCE(EXCLUDED.batch_id, pipeline_runs.batch_id),
|
|
|
+ status = EXCLUDED.status,
|
|
|
+ current_stage = COALESCE(EXCLUDED.current_stage, pipeline_runs.current_stage),
|
|
|
+ config = pipeline_runs.config || EXCLUDED.config,
|
|
|
+ metadata = pipeline_runs.metadata || EXCLUDED.metadata,
|
|
|
+ started_at = COALESCE(pipeline_runs.started_at, EXCLUDED.started_at)
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ run_key,
|
|
|
+ batch_id,
|
|
|
+ status,
|
|
|
+ current_stage,
|
|
|
+ Json(config or {}),
|
|
|
+ Json(metadata or {}),
|
|
|
+ status,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ return _pipeline_run(row)
|
|
|
+
|
|
|
+ def mark_pipeline_run_status(
|
|
|
+ self,
|
|
|
+ run_id: UUID,
|
|
|
+ *,
|
|
|
+ status: str,
|
|
|
+ current_stage: str | None = None,
|
|
|
+ error_message: str | None = None,
|
|
|
+ metadata: dict[str, Any] | None = None,
|
|
|
+ ) -> PipelineRun:
|
|
|
+ row = self._one(
|
|
|
+ """
|
|
|
+ UPDATE pipeline_runs SET
|
|
|
+ status = %s,
|
|
|
+ current_stage = COALESCE(%s, current_stage),
|
|
|
+ error_message = %s,
|
|
|
+ metadata = CASE WHEN %s THEN metadata || %s ELSE metadata END,
|
|
|
+ finished_at = CASE WHEN %s IN ('done', 'partial', 'failed') THEN now() ELSE finished_at END
|
|
|
+ WHERE id = %s
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ status,
|
|
|
+ current_stage,
|
|
|
+ error_message,
|
|
|
+ metadata is not None,
|
|
|
+ Json(metadata or {}),
|
|
|
+ status,
|
|
|
+ run_id,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ return _pipeline_run(row)
|
|
|
+
|
|
|
+ def get_pipeline_run(self, run_id: UUID) -> PipelineRun:
|
|
|
+ return _pipeline_run(self._one("SELECT * FROM pipeline_runs WHERE id = %s", (run_id,)))
|
|
|
+
|
|
|
+ def get_pipeline_run_by_acquisition_run(self, acquisition_run_id: UUID) -> PipelineRun:
|
|
|
+ row = self._one(
|
|
|
+ """
|
|
|
+ SELECT pr.* FROM pipeline_runs pr
|
|
|
+ JOIN pipeline_run_events pre ON pre.pipeline_run_id = pr.id
|
|
|
+ WHERE pre.acquisition_run_id = %s
|
|
|
+ ORDER BY pre.created_at DESC
|
|
|
+ LIMIT 1
|
|
|
+ """,
|
|
|
+ (acquisition_run_id,),
|
|
|
+ )
|
|
|
+ return _pipeline_run(row)
|
|
|
+
|
|
|
+ def save_pipeline_job(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ run_id: UUID,
|
|
|
+ stage: PipelineStage | str,
|
|
|
+ target_id: UUID | None = None,
|
|
|
+ target_table: str | None = None,
|
|
|
+ status: str = "pending",
|
|
|
+ metadata: dict[str, Any] | None = None,
|
|
|
+ ) -> PipelineJob:
|
|
|
+ row = self._one(
|
|
|
+ """
|
|
|
+ INSERT INTO pipeline_jobs(
|
|
|
+ pipeline_run_id, stage, target_table, target_id, status, metadata, started_at
|
|
|
+ )
|
|
|
+ VALUES (%s, %s, %s, %s, %s, %s, CASE WHEN %s = 'running' THEN now() ELSE NULL END)
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (run_id, stage, target_table, target_id, status, Json(metadata or {}), status),
|
|
|
+ )
|
|
|
+ return _pipeline_job(row)
|
|
|
+
|
|
|
+ def mark_job_status(
|
|
|
+ self,
|
|
|
+ job_id: UUID,
|
|
|
+ *,
|
|
|
+ status: str,
|
|
|
+ error_message: str | None = None,
|
|
|
+ metadata: dict[str, Any] | None = None,
|
|
|
+ ) -> PipelineJob:
|
|
|
+ row = self._one(
|
|
|
+ """
|
|
|
+ UPDATE pipeline_jobs SET
|
|
|
+ status = %s,
|
|
|
+ error_message = %s,
|
|
|
+ metadata = CASE WHEN %s THEN metadata || %s ELSE metadata END,
|
|
|
+ finished_at = CASE WHEN %s IN ('done', 'partial', 'failed', 'skipped') THEN now() ELSE finished_at END
|
|
|
+ WHERE id = %s
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ status,
|
|
|
+ error_message,
|
|
|
+ metadata is not None,
|
|
|
+ Json(metadata or {}),
|
|
|
+ status,
|
|
|
+ job_id,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ return _pipeline_job(row)
|
|
|
+
|
|
|
+ def append_event(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ context: TraceContext,
|
|
|
+ stage: str,
|
|
|
+ event_type: str,
|
|
|
+ status: str | None = None,
|
|
|
+ severity: str = "info",
|
|
|
+ target_table: str | None = None,
|
|
|
+ target_id: UUID | None = None,
|
|
|
+ message: str | None = None,
|
|
|
+ payload: dict[str, Any] | None = None,
|
|
|
+ error_message: str | None = None,
|
|
|
+ duration_ms: int | None = None,
|
|
|
+ attempt_index: int | None = None,
|
|
|
+ ) -> PipelineRunEvent:
|
|
|
+ row = self._one(
|
|
|
+ """
|
|
|
+ INSERT INTO pipeline_run_events(
|
|
|
+ pipeline_run_id, pipeline_job_id, stage, event_type, status, severity,
|
|
|
+ target_table, target_id, acquisition_run_id, acquisition_job_id,
|
|
|
+ query_id, item_id, decode_job_id, decode_result_id, payload_draft_id,
|
|
|
+ ingest_record_id, platform, attempt_index, message, payload,
|
|
|
+ error_message, duration_ms
|
|
|
+ )
|
|
|
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ context.pipeline_run_id,
|
|
|
+ context.pipeline_job_id,
|
|
|
+ stage,
|
|
|
+ event_type,
|
|
|
+ status,
|
|
|
+ severity,
|
|
|
+ target_table,
|
|
|
+ target_id,
|
|
|
+ context.acquisition_run_id,
|
|
|
+ context.acquisition_job_id,
|
|
|
+ context.query_id,
|
|
|
+ context.item_id,
|
|
|
+ context.decode_job_id,
|
|
|
+ context.decode_result_id,
|
|
|
+ context.payload_draft_id,
|
|
|
+ context.ingest_record_id,
|
|
|
+ context.platform,
|
|
|
+ attempt_index,
|
|
|
+ message,
|
|
|
+ Json(payload or {}),
|
|
|
+ error_message,
|
|
|
+ duration_ms,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ return PipelineRunEvent.model_validate(_event_row(row))
|
|
|
+
|
|
|
+ def record_candidate_hit(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ context: TraceContext,
|
|
|
+ item_id: UUID | None,
|
|
|
+ platform: str,
|
|
|
+ unique_key: str | None = None,
|
|
|
+ platform_item_id: str | None = None,
|
|
|
+ search_provider: str | None = None,
|
|
|
+ detail_provider: str | None = None,
|
|
|
+ attempt_index: int | None = None,
|
|
|
+ page_index: int | None = None,
|
|
|
+ page_rank: int | None = None,
|
|
|
+ candidate_rank: int | None = None,
|
|
|
+ source_cursor: str | None = None,
|
|
|
+ is_duplicate_hit: bool = False,
|
|
|
+ raw_candidate: dict[str, Any] | None = None,
|
|
|
+ metadata: dict[str, Any] | None = None,
|
|
|
+ ) -> dict[str, Any]:
|
|
|
+ return self._one(
|
|
|
+ """
|
|
|
+ INSERT INTO candidate_item_hits(
|
|
|
+ pipeline_run_id, acquisition_run_id, acquisition_job_id,
|
|
|
+ query_id, item_id, platform, unique_key, platform_item_id,
|
|
|
+ search_provider, detail_provider, attempt_index, page_index,
|
|
|
+ page_rank, candidate_rank, source_cursor, is_duplicate_hit,
|
|
|
+ raw_candidate, metadata
|
|
|
+ )
|
|
|
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ context.pipeline_run_id,
|
|
|
+ context.acquisition_run_id,
|
|
|
+ context.acquisition_job_id,
|
|
|
+ context.query_id,
|
|
|
+ item_id,
|
|
|
+ platform,
|
|
|
+ unique_key,
|
|
|
+ platform_item_id,
|
|
|
+ search_provider,
|
|
|
+ detail_provider,
|
|
|
+ attempt_index,
|
|
|
+ page_index,
|
|
|
+ page_rank,
|
|
|
+ candidate_rank,
|
|
|
+ source_cursor,
|
|
|
+ is_duplicate_hit,
|
|
|
+ Json(raw_candidate or {}),
|
|
|
+ Json(metadata or {}),
|
|
|
+ ),
|
|
|
+ )
|
|
|
+
|
|
|
+ def save_llm_call_trace(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ context: TraceContext,
|
|
|
+ stage: str,
|
|
|
+ substage: str | None = None,
|
|
|
+ provider: str | None = None,
|
|
|
+ model_name: str | None = None,
|
|
|
+ endpoint: str | None = None,
|
|
|
+ prompt_name: str | None = None,
|
|
|
+ prompt_hash: str | None = None,
|
|
|
+ request_payload: dict[str, Any] | None = None,
|
|
|
+ response_payload: dict[str, Any] | None = None,
|
|
|
+ parsed_payload: dict[str, Any] | None = None,
|
|
|
+ status: str = "pending",
|
|
|
+ error_message: str | None = None,
|
|
|
+ latency_ms: int | None = None,
|
|
|
+ attempt_index: int | None = None,
|
|
|
+ ) -> dict[str, Any]:
|
|
|
+ return self._one(
|
|
|
+ """
|
|
|
+ INSERT INTO llm_call_traces(
|
|
|
+ pipeline_run_id, pipeline_job_id, stage, substage, provider,
|
|
|
+ model_name, endpoint, prompt_name, prompt_hash, trace_context,
|
|
|
+ request_payload, response_payload, parsed_payload, status,
|
|
|
+ error_message, latency_ms, attempt_index
|
|
|
+ )
|
|
|
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
|
+ RETURNING *
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ context.pipeline_run_id,
|
|
|
+ context.pipeline_job_id,
|
|
|
+ stage,
|
|
|
+ substage,
|
|
|
+ provider,
|
|
|
+ model_name,
|
|
|
+ endpoint,
|
|
|
+ prompt_name,
|
|
|
+ prompt_hash,
|
|
|
+ Json(context.as_payload()),
|
|
|
+ Json(request_payload or {}),
|
|
|
+ Json(response_payload or {}),
|
|
|
+ Json(parsed_payload or {}),
|
|
|
+ status,
|
|
|
+ error_message,
|
|
|
+ latency_ms,
|
|
|
+ attempt_index,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+
|
|
|
+ def list_timeline(self, run_id: UUID) -> list[dict[str, Any]]:
|
|
|
+ return self._all(
|
|
|
+ """
|
|
|
+ SELECT * FROM pipeline_run_events
|
|
|
+ WHERE pipeline_run_id = %s
|
|
|
+ ORDER BY created_at, id
|
|
|
+ """,
|
|
|
+ (run_id,),
|
|
|
+ )
|
|
|
+
|
|
|
+ def list_jobs(self, run_id: UUID) -> list[dict[str, Any]]:
|
|
|
+ return self._all(
|
|
|
+ """
|
|
|
+ SELECT * FROM pipeline_jobs
|
|
|
+ WHERE pipeline_run_id = %s
|
|
|
+ ORDER BY created_at, id
|
|
|
+ """,
|
|
|
+ (run_id,),
|
|
|
+ )
|
|
|
+
|
|
|
+ def list_candidate_hits(self, run_id: UUID) -> list[dict[str, Any]]:
|
|
|
+ return self._all(
|
|
|
+ """
|
|
|
+ SELECT * FROM candidate_item_hits
|
|
|
+ WHERE pipeline_run_id = %s
|
|
|
+ ORDER BY created_at, page_index NULLS LAST, page_rank NULLS LAST, id
|
|
|
+ """,
|
|
|
+ (run_id,),
|
|
|
+ )
|
|
|
+
|
|
|
+ def list_llm_call_traces(self, run_id: UUID, *, limit: int = 500) -> list[dict[str, Any]]:
|
|
|
+ return self._all(
|
|
|
+ """
|
|
|
+ SELECT * FROM llm_call_traces
|
|
|
+ WHERE pipeline_run_id = %s
|
|
|
+ ORDER BY created_at, id
|
|
|
+ LIMIT %s
|
|
|
+ """,
|
|
|
+ (run_id, limit),
|
|
|
+ )
|
|
|
+
|
|
|
+ def get_timeline_bundle(self, run_id: UUID) -> dict[str, Any]:
|
|
|
+ return {
|
|
|
+ "events": self.list_timeline(run_id),
|
|
|
+ "jobs": self.list_jobs(run_id),
|
|
|
+ "candidate_hits": self.list_candidate_hits(run_id),
|
|
|
+ "llm_call_traces": self.list_llm_call_traces(run_id),
|
|
|
+ }
|
|
|
+
|
|
|
+ def get_resume_cursor(self, run_id: UUID):
|
|
|
+ return None
|
|
|
+
|
|
|
+
|
|
|
+class PostgresTraceWriter:
|
|
|
+ """Best-effort writer that never lets tracing break the business flow."""
|
|
|
+
|
|
|
+ def __init__(self, *, db_config: CreationDbConfig, config: TraceConfig):
|
|
|
+ self.db_config = db_config
|
|
|
+ self.config = config
|
|
|
+
|
|
|
+ def event(self, *, context: TraceContext, stage: str, event_type: str, **kwargs: Any) -> None:
|
|
|
+ self._best_effort(
|
|
|
+ lambda repo: repo.append_event(
|
|
|
+ context=context,
|
|
|
+ stage=stage,
|
|
|
+ event_type=event_type,
|
|
|
+ payload=sanitize_payload(kwargs.pop("payload", {}) or {}, max_chars=self.config.max_payload_chars),
|
|
|
+ **kwargs,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+ def candidate_hit(self, *, context: TraceContext, **kwargs: Any) -> None:
|
|
|
+ self._best_effort(
|
|
|
+ lambda repo: repo.record_candidate_hit(
|
|
|
+ context=context,
|
|
|
+ raw_candidate=sanitize_payload(kwargs.pop("raw_candidate", {}) or {}, max_chars=self.config.max_payload_chars),
|
|
|
+ metadata=sanitize_payload(kwargs.pop("metadata", {}) or {}, max_chars=self.config.max_payload_chars),
|
|
|
+ **kwargs,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+ def llm_call(self, *, context: TraceContext, stage: str, substage: str | None = None, **kwargs: Any) -> None:
|
|
|
+ request_payload = kwargs.pop("request_payload", {}) if self.config.capture_request else {}
|
|
|
+ response_payload = kwargs.pop("response_payload", {}) if self.config.capture_response else {}
|
|
|
+ parsed_payload = kwargs.pop("parsed_payload", {})
|
|
|
+ self._best_effort(
|
|
|
+ lambda repo: repo.save_llm_call_trace(
|
|
|
+ context=context,
|
|
|
+ stage=stage,
|
|
|
+ substage=substage,
|
|
|
+ request_payload=sanitize_payload(request_payload or {}, max_chars=self.config.max_payload_chars),
|
|
|
+ response_payload=sanitize_payload(response_payload or {}, max_chars=self.config.max_payload_chars),
|
|
|
+ parsed_payload=sanitize_payload(parsed_payload or {}, max_chars=self.config.max_payload_chars),
|
|
|
+ **kwargs,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+ def _best_effort(self, fn) -> None:
|
|
|
+ if not self.config.enabled:
|
|
|
+ return
|
|
|
+ try:
|
|
|
+ with transaction(self.db_config) as conn:
|
|
|
+ fn(PostgresPipelineRepository(conn))
|
|
|
+ except Exception as exc: # pragma: no cover - defensive isolation
|
|
|
+ logger.warning("pipeline trace write skipped: %s", exc)
|
|
|
+
|
|
|
+
|
|
|
+def _pipeline_run(row: dict[str, Any]) -> PipelineRun:
|
|
|
+ return PipelineRun(
|
|
|
+ id=row.get("id"),
|
|
|
+ run_key=row.get("run_key"),
|
|
|
+ batch_id=row.get("batch_id"),
|
|
|
+ status=row.get("status") or "pending",
|
|
|
+ current_stage=row.get("current_stage"),
|
|
|
+ metadata={
|
|
|
+ **(row.get("config") or {}),
|
|
|
+ **(row.get("metadata") or {}),
|
|
|
+ },
|
|
|
+ error_message=row.get("error_message"),
|
|
|
+ started_at=row.get("started_at"),
|
|
|
+ finished_at=row.get("finished_at"),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _pipeline_job(row: dict[str, Any]) -> PipelineJob:
|
|
|
+ return PipelineJob(
|
|
|
+ id=row.get("id"),
|
|
|
+ run_id=row.get("pipeline_run_id"),
|
|
|
+ stage=row.get("stage"),
|
|
|
+ status=row.get("status") or "pending",
|
|
|
+ target_table=row.get("target_table"),
|
|
|
+ target_id=row.get("target_id"),
|
|
|
+ attempt_count=row.get("attempt_count") or 0,
|
|
|
+ metadata=row.get("metadata") or {},
|
|
|
+ error_message=row.get("error_message"),
|
|
|
+ started_at=row.get("started_at"),
|
|
|
+ finished_at=row.get("finished_at"),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _event_row(row: dict[str, Any]) -> dict[str, Any]:
|
|
|
+ out = dict(row)
|
|
|
+ out["pipeline_run_id"] = out.pop("pipeline_run_id", None)
|
|
|
+ return out
|