-
Notifications
You must be signed in to change notification settings - Fork 39
Expand file tree
/
Copy pathask.py
More file actions
195 lines (158 loc) · 7.24 KB
/
Copy pathask.py
File metadata and controls
195 lines (158 loc) · 7.24 KB
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
"""Diffbot LLM RAG API: stream a chat completion."""
import json
import re
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List, Optional
from .errors import ValidationError
if TYPE_CHECKING:
from .client import Diffbot, DiffbotAsync
MODEL = "diffbot-small-xl"
#: `response_format` types accepted by the Diffbot LLM endpoint.
RESPONSE_FORMAT_TYPES = ("text", "json_object", "json_schema")
# The RAG loop may prefix its final answer with a think block, and the model
# occasionally wraps JSON in a markdown fence despite being told not to.
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)
_JSON_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$", re.MULTILINE)
def json_schema_format(schema: Dict[str, Any], *, name: str = "response") -> Dict[str, Any]:
"""Build a ``response_format`` value that constrains output to ``schema``.
The endpoint requires the schema nested under ``json_schema.schema``; passing
it anywhere else is ignored server-side without an error, so prefer this
helper over hand-building the dict.
Example:
>>> json_schema_format({"type": "object", "properties": {"city": {"type": "string"}}})
{'type': 'json_schema', 'json_schema': {'name': 'response', 'schema': {...}}}
"""
if not isinstance(schema, dict):
raise ValidationError("schema must be a JSON Schema dict")
return {"type": "json_schema", "json_schema": {"name": name, "schema": schema}}
def _validate_response_format(response_format: Optional[Dict[str, Any]]) -> None:
"""Reject shapes the endpoint would accept but silently not enforce."""
if response_format is None:
return
if not isinstance(response_format, dict):
raise ValidationError("response_format must be a dict")
fmt_type = response_format.get("type", "text")
if fmt_type not in RESPONSE_FORMAT_TYPES:
raise ValidationError(
f"response_format type must be one of {', '.join(RESPONSE_FORMAT_TYPES)}; got {fmt_type!r}"
)
if fmt_type == "json_schema":
json_schema = response_format.get("json_schema")
if not isinstance(json_schema, dict) or "schema" not in json_schema:
raise ValidationError(
'response_format {"type": "json_schema"} requires the schema nested as '
'{"json_schema": {"schema": {...}}}. Without it the server returns 200 and '
"ignores the constraint. Use diffbot.json_schema_format(schema) to build it."
)
def _build_payload(
client: Any,
messages: List[Dict[str, str]],
*,
response_format: Optional[Dict[str, Any]] = None,
) -> tuple:
_validate_response_format(response_format)
headers = {"Authorization": f"Bearer {client.token}"}
payload: Dict[str, Any] = {"model": MODEL, "messages": messages, "stream": True}
if response_format is not None:
payload["response_format"] = response_format
return headers, payload
def _parse_chunk(line: str):
try:
chunk = json.loads(line.replace("data: ", ""))
except json.JSONDecodeError:
return None
choices = chunk.get("choices")
if choices and choices[0].get("delta", {}).get("content"):
return choices[0]["delta"]["content"]
return None
def _extract_json(text: str) -> Any:
"""Parse the model's final answer as JSON, tolerating think blocks and fences."""
cleaned = _THINK_BLOCK.sub("", text).strip()
cleaned = _JSON_FENCE.sub("", cleaned).strip()
try:
return json.loads(cleaned)
except json.JSONDecodeError:
pass
# Fall back to the outermost object or array span in the response.
spans = []
for opener, closer in (("{", "}"), ("[", "]")):
start, end = cleaned.find(opener), cleaned.rfind(closer)
if start != -1 and end > start:
spans.append((start, cleaned[start : end + 1]))
for _, span in sorted(spans):
try:
return json.loads(span)
except json.JSONDecodeError:
continue
raise ValidationError(f"could not parse JSON from the model response: {text[:200]!r}")
#: Schema used when the caller wants JSON but has no shape in mind. This goes
#: through the json_schema path rather than {"type": "json_object"} on purpose:
#: json_object applies no server-side grammar, so the RAG loop's internal
#: <functioncall> JSON satisfies it and gets returned as the final answer. That
#: is reproducible whenever the request carries a system message.
ANY_OBJECT_SCHEMA = {"type": "object"}
def _resolve_format(
schema: Optional[Dict[str, Any]],
response_format: Optional[Dict[str, Any]],
) -> Dict[str, Any]:
if schema is not None and response_format is not None:
raise ValidationError("pass either schema or response_format, not both")
if response_format is not None:
return response_format
return json_schema_format(schema if schema is not None else ANY_OBJECT_SCHEMA)
def _check_tool_call_leak(parsed: Any) -> Any:
"""Catch the internal tool call surfacing as the answer (see ANY_OBJECT_SCHEMA)."""
if isinstance(parsed, dict) and parsed.get("name") == "functioncall" and "arguments" in parsed:
raise ValidationError(
"the model returned its internal tool call instead of an answer; this happens with "
'response_format {"type": "json_object"} because the server applies no grammar to it. '
"Pass a schema instead."
)
return parsed
def ask(
client: "Diffbot",
messages: List[Dict[str, str]],
*,
response_format: Optional[Dict[str, Any]] = None,
) -> Iterator[str]:
headers, payload = _build_payload(client, messages, response_format=response_format)
with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
client._raise_for_status(response)
for line in response.iter_lines():
if line:
content = _parse_chunk(line)
if content:
yield content
async def ask_async(
client: "DiffbotAsync",
messages: List[Dict[str, str]],
*,
response_format: Optional[Dict[str, Any]] = None,
) -> AsyncIterator[str]:
headers, payload = _build_payload(client, messages, response_format=response_format)
async with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
client._raise_for_status(response)
async for line in response.aiter_lines():
if line:
content = _parse_chunk(line)
if content:
yield content
def ask_json(
client: "Diffbot",
messages: List[Dict[str, str]],
schema: Optional[Dict[str, Any]] = None,
*,
response_format: Optional[Dict[str, Any]] = None,
) -> Any:
fmt = _resolve_format(schema, response_format)
text = "".join(ask(client, messages, response_format=fmt))
return _check_tool_call_leak(_extract_json(text))
async def ask_json_async(
client: "DiffbotAsync",
messages: List[Dict[str, str]],
schema: Optional[Dict[str, Any]] = None,
*,
response_format: Optional[Dict[str, Any]] = None,
) -> Any:
fmt = _resolve_format(schema, response_format)
chunks = [chunk async for chunk in ask_async(client, messages, response_format=fmt)]
return _check_tool_call_leak(_extract_json("".join(chunks)))