This repository has no description
1"""Persist AI-solve questionnaires in Postgres."""
2
3from __future__ import annotations
4
5import json
6import os
7from typing import Any
8
9import psycopg
10from psycopg.rows import dict_row
11from psycopg.types.json import Jsonb
12
13_UPSERT = """
14 insert into tangled_issue_questionnaires (issue_uri, payload, updated_at)
15 values (%s, %s, now())
16 on conflict (issue_uri) do update set
17 payload = excluded.payload,
18 updated_at = now()
19 returning issue_uri, created_at, updated_at
20"""
21
22_GET = """
23 select issue_uri, payload, created_at, updated_at
24 from tangled_issue_questionnaires
25 where issue_uri = %s
26"""
27
28
29def _connection_string() -> str:
30 dsn = os.getenv("DB_CONNECTION_STRING", "").strip()
31 if not dsn:
32 raise RuntimeError("DB_CONNECTION_STRING is not set")
33 return dsn
34
35
36def parse_questionnaire_json(raw: str) -> dict[str, Any]:
37 """Parse model output into a questionnaire dict (tolerates fences and preamble)."""
38 import re
39 from json import JSONDecoder
40
41 text = raw.strip()
42 if not text:
43 raise ValueError("Empty model response — expected questionnaire JSON")
44
45 fence = re.search(r"```(?:json)?\s*(\{.*\})\s*```", text, re.DOTALL)
46 if fence:
47 text = fence.group(1).strip()
48
49 decoder = JSONDecoder()
50 try:
51 data, _ = decoder.raw_decode(text)
52 except json.JSONDecodeError:
53 start = text.find("{")
54 if start < 0:
55 preview = text[:300].replace("\n", " ")
56 raise ValueError(
57 f"No JSON object in model response (preview: {preview!r})"
58 ) from None
59 data, _ = decoder.raw_decode(text[start:])
60
61 if not isinstance(data, dict) or not isinstance(data.get("items"), list):
62 raise ValueError("Invalid questionnaire: expected object with items[]")
63 return data
64
65
66def save_questionnaire(issue_uri: str, payload: dict[str, Any]) -> dict[str, Any]:
67 """Insert or replace the questionnaire for an issue. Returns row metadata."""
68 if payload.get("issue") and payload["issue"] != issue_uri:
69 raise ValueError(
70 f"payload.issue ({payload['issue']!r}) does not match issue_uri ({issue_uri!r})"
71 )
72 with psycopg.connect(_connection_string(), row_factory=dict_row) as conn:
73 row = conn.execute(
74 _UPSERT,
75 (issue_uri, Jsonb(payload)),
76 ).fetchone()
77 conn.commit()
78 return dict(row)
79
80
81def get_questionnaire(issue_uri: str) -> dict[str, Any] | None:
82 """Load cached questionnaire JSON, or None if missing."""
83 with psycopg.connect(_connection_string(), row_factory=dict_row) as conn:
84 row = conn.execute(_GET, (issue_uri,)).fetchone()
85 if not row:
86 return None
87 payload = row["payload"]
88 if isinstance(payload, str):
89 payload = json.loads(payload)
90 return {
91 "issue_uri": row["issue_uri"],
92 "payload": payload,
93 "created_at": row["created_at"],
94 "updated_at": row["updated_at"],
95 }