File size: 7,162 Bytes
200cb0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
"""JSON parsing helpers for model outputs."""

from __future__ import annotations

import json
import ast
import re
from typing import Any

from iris.errors import IrisResponseError


FENCE_RE = re.compile(r"^```(?:json)?\s*|\s*```$", re.IGNORECASE)


def parse_json_object(text: str) -> dict[str, Any]:
    """Parse the first JSON object from a model response."""

    if not isinstance(text, str):
        raise IrisResponseError(f"Expected text response, got: {type(text).__name__}")

    cleaned = FENCE_RE.sub("", text.strip())
    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError:
        try:
            parsed = _parse_embedded_object(cleaned)
        except IrisResponseError:
            parsed = _parse_repaired_json(cleaned)
            if parsed is None:
                parsed = _parse_unkeyed_bite(cleaned)
                if parsed is None:
                    parsed = _parse_python_literal(cleaned)
                    if parsed is None:
                        raise

    if not isinstance(parsed, dict):
        raise IrisResponseError(f"Expected a JSON object, got: {type(parsed).__name__}")
    return parsed


def parse_json_object_with_key(
    text: str,
    key: str,
    aliases: tuple[str, ...] = (),
) -> dict[str, Any]:
    """Parse the last embedded JSON object containing a requested key."""

    if not isinstance(text, str):
        raise IrisResponseError(f"Expected text response, got: {type(text).__name__}")

    normalized_keys = {_normalize_key(item) for item in (key, *aliases)}
    cleaned = FENCE_RE.sub("", text.strip())
    for candidate in reversed(_json_object_candidates(cleaned)):
        if _has_any_key(candidate, normalized_keys):
            return candidate
    return parse_json_object(text)


def require_string(data: dict[str, Any], key: str, aliases: tuple[str, ...] = ()) -> str:
    value = _get_alias_value(data, key, aliases)
    if not isinstance(value, str) or not value.strip():
        raise IrisResponseError(f"Expected non-empty string field: {key}")
    return value.strip()


def _parse_embedded_object(text: str) -> dict[str, Any]:
    candidate = _first_json_object(text)
    if candidate is None:
        raise IrisResponseError(f"Model did not return JSON: {text[:500]}")

    try:
        parsed = json.loads(candidate)
    except json.JSONDecodeError as exc:
        parsed = _parse_repaired_json(candidate)
        if parsed is None:
            parsed = _parse_unkeyed_bite(candidate)
            if parsed is None:
                parsed = _parse_python_literal(candidate)
                if parsed is None:
                    raise IrisResponseError(f"Could not parse model JSON: {candidate[:500]}") from exc
    if not isinstance(parsed, dict):
        raise IrisResponseError(f"Expected a JSON object, got: {type(parsed).__name__}")
    return parsed


def _first_json_object(text: str) -> str | None:
    start = text.find("{")
    if start == -1:
        return None

    depth = 0
    in_string = False
    escape = False

    for index in range(start, len(text)):
        char = text[index]
        if in_string:
            if escape:
                escape = False
            elif char == "\\":
                escape = True
            elif char == '"':
                in_string = False
            continue

        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]

    return None


def _json_object_candidates(text: str) -> list[dict[str, Any]]:
    candidates: list[dict[str, Any]] = []
    for candidate_text in _json_object_strings(text):
        candidate = _parse_object_candidate(candidate_text)
        if candidate is not None:
            candidates.append(candidate)
    return candidates


def _json_object_strings(text: str) -> list[str]:
    objects: list[str] = []
    start: int | None = None
    depth = 0
    in_string = False
    escape = False

    for index, char in enumerate(text):
        if in_string:
            if escape:
                escape = False
            elif char == "\\":
                escape = True
            elif char == '"':
                in_string = False
            continue

        if char == '"':
            in_string = True
        elif char == "{":
            if depth == 0:
                start = index
            depth += 1
        elif char == "}":
            if depth == 0:
                continue
            depth -= 1
            if depth == 0 and start is not None:
                objects.append(text[start : index + 1])
                start = None

    return objects


def _parse_object_candidate(text: str) -> dict[str, Any] | None:
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        parsed = _parse_repaired_json(text)
        if parsed is None:
            parsed = _parse_unkeyed_bite(text)
            if parsed is None:
                parsed = _parse_python_literal(text)
    return parsed if isinstance(parsed, dict) else None


def _parse_python_literal(text: str) -> dict[str, Any] | None:
    try:
        parsed = ast.literal_eval(text)
    except (ValueError, SyntaxError):
        return None
    if isinstance(parsed, dict):
        return parsed
    return None


def _parse_repaired_json(text: str) -> dict[str, Any] | None:
    repaired = re.sub(
        r'"\s+"([A-Za-z_][A-Za-z0-9_ ]*)"\s*:',
        r'", "\1":',
        text,
    )
    repaired = re.sub(
        r':(?!\s*")\s*([^"{\[\]-].*?)"\s*([,}])',
        lambda match: f': "{match.group(1).strip()}"{match.group(2)}',
        repaired,
    )
    if repaired == text:
        return None
    try:
        parsed = json.loads(repaired)
    except json.JSONDecodeError:
        return None
    if isinstance(parsed, dict):
        return parsed
    return None


def _parse_unkeyed_bite(text: str) -> dict[str, Any] | None:
    match = re.match(
        r'^\{\s*"pressure"\s*:\s*"(?P<pressure>.*?)"\s*,?\s*"(?P<why>.*?)"\s*\}\s*$',
        text,
        re.IGNORECASE | re.DOTALL,
    )
    if not match:
        return None
    why = match.group("why").strip()
    why = re.sub(r"^why it bites:\s*", "", why, flags=re.IGNORECASE)
    why = re.sub(r"^this bites because\s*", "", why, flags=re.IGNORECASE)
    why = re.sub(r"^because\s*", "", why, flags=re.IGNORECASE)
    return {"pressure": match.group("pressure").strip(), "why_it_bites": why.strip()}


def _get_alias_value(data: dict[str, Any], key: str, aliases: tuple[str, ...]) -> Any:
    if key in data:
        return data[key]

    normalized = {_normalize_key(item_key): value for item_key, value in data.items()}
    for candidate in (key, *aliases):
        value = normalized.get(_normalize_key(candidate))
        if value is not None:
            return value
    return None


def _has_any_key(data: dict[str, Any], normalized_keys: set[str]) -> bool:
    return any(_normalize_key(item_key) in normalized_keys for item_key in data)


def _normalize_key(key: str) -> str:
    return re.sub(r"[^a-z0-9]", "", key.lower())