Skip to content

Commit d2b13e2

Browse files
jeromechooclaude
jeromechoo
and
claude
committed
Add structured output support to ask
The Diffbot LLM endpoint supports response_format, but the client had no way to pass it: the payload in ask.py was hardcoded to model/messages/stream. Adds response_format to ask() and a new ask_json() that returns a parsed object, both sync and async, plus json_schema_format() to build the payload. The CLI gains --schema, and --json now uses real structured output instead of nudging the prompt and slicing between the first { and last }. Server-side (see ../diffbot-llm) the schema is compiled to an EBNF grammar and the final answer is constrained to it, so this is enforced decoding rather than a request the model may ignore. Two rough edges are handled client-side: - The schema must be nested at json_schema.schema. Anywhere else and the server returns 200 with the constraint silently dropped, so we reject that locally. - The grammar permits a <think> block before the final answer, so ask_json strips think blocks and markdown fences before parsing. ask_json without a schema deliberately sends a permissive {"type": "object"} json_schema rather than {"type": "json_object"}. json_object gets no grammar applied, and the RAG loop's internal tool call is itself a valid JSON object, so it can be returned as the final answer -- reproducible with a system message on some queries (5/5 on one, 0/12 without a system message). A guard raises ValidationError if a tool call ever comes back. Tested with 44 unit tests and 7 live tests against the real endpoint, covering flat and nested schemas, streaming with a schema, async, and a regression test for the system-message case. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 parent b64f6a1 commit d2b13e2

7 files changed

Lines changed: 655 additions & 29 deletions

File tree

‎README.md‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,42 @@ for chunk in db.ask([{"role": "user", "content": "What's the capital of France?"
7575
print(chunk, end="")
7676
```
7777

78+
### Structured output
79+
Pass a JSON Schema to constrain the answer. The model is held to the schema by a grammar during decoding, so the result parses reliably even though the answer is retrieved live from the web.
80+
81+
```python
82+
from diffbot import Diffbot
83+
84+
db = Diffbot(token="YOUR_TOKEN")
85+
schema = {
86+
"type": "object",
87+
"properties": {
88+
"country": {"type": "string"},
89+
"capital": {"type": "string"},
90+
},
91+
"required": ["country", "capital"],
92+
}
93+
answer = db.ask_json([{"role": "user", "content": "What's the capital of France?"}], schema)
94+
print(answer["capital"])
95+
```
96+
97+
Omit the schema to let the model choose the shape, or use `ask` with `response_format` to stream a constrained answer:
98+
99+
```python
100+
from diffbot import Diffbot, json_schema_format
101+
102+
db = Diffbot(token="YOUR_TOKEN")
103+
answer = db.ask_json([{"role": "user", "content": "What's the capital of France?"}])
104+
105+
for chunk in db.ask(
106+
[{"role": "user", "content": "What's the capital of France?"}],
107+
response_format=json_schema_format(schema),
108+
):
109+
print(chunk, end="")
110+
```
111+
112+
> **Avoid `response_format={"type": "json_object"}`.** The endpoint accepts it, but applies no grammar to it — the RAG loop's internal tool call is itself a JSON object, so it can be returned as the final answer. This is reproducible whenever the request includes a system message. `ask_json` therefore defaults to a permissive JSON Schema instead, and raises `ValidationError` if it ever sees a tool call come back.
113+
78114
### Crawl a site for structured content
79115
```python
80116
from diffbot import Diffbot
@@ -142,6 +178,30 @@ async def main():
142178
asyncio.run(main())
143179
```
144180

181+
### Structured output
182+
```python
183+
import asyncio
184+
from diffbot import DiffbotAsync
185+
186+
schema = {
187+
"type": "object",
188+
"properties": {
189+
"country": {"type": "string"},
190+
"capital": {"type": "string"},
191+
},
192+
"required": ["country", "capital"],
193+
}
194+
195+
async def main():
196+
async with DiffbotAsync(token="YOUR_TOKEN") as db:
197+
answer = await db.ask_json(
198+
[{"role": "user", "content": "What's the capital of France?"}], schema
199+
)
200+
print(answer["capital"])
201+
202+
asyncio.run(main())
203+
```
204+
145205
### Crawl a site for structured content
146206
```python
147207
import asyncio
@@ -231,6 +291,8 @@ export DIFFBOT_API_TOKEN=your-token-here
231291

232292
db extract https://www.example.com
233293
db ask "What's the capital of France?"
294+
db ask "What's the capital of France?" --json
295+
db ask "What's the capital of France?" --schema capital.json
234296
db crawl https://www.example.com --hops 1
235297
db crawl-list-jobs
236298
db crawl-delete-job crawl-1234567890

‎src/diffbot/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
__version__ = "0.0.0"
1111

1212
from ._auth import resolve_token
13+
from .ask import json_schema_format
1314
from .client import Diffbot, DiffbotAsync
1415
from .crawl import CrawlEvent, CrawlEventType
1516
from .errors import (
@@ -26,6 +27,7 @@
2627
"Diffbot",
2728
"DiffbotAsync",
2829
"resolve_token",
30+
"json_schema_format",
2931
"CrawlEvent",
3032
"CrawlEventType",
3133
"Ontology",

‎src/diffbot/ask.py‎

Lines changed: 156 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,75 @@
11
"""Diffbot LLM RAG API: stream a chat completion."""
22

33
import json
4-
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List
4+
import re
5+
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List, Optional
6+
7+
from .errors import ValidationError
58

69
if TYPE_CHECKING:
710
from .client import Diffbot, DiffbotAsync
811

12+
MODEL = "diffbot-small-xl"
13+
14+
#: `response_format` types accepted by the Diffbot LLM endpoint.
15+
RESPONSE_FORMAT_TYPES = ("text", "json_object", "json_schema")
16+
17+
# The RAG loop may prefix its final answer with a think block, and the model
18+
# occasionally wraps JSON in a markdown fence despite being told not to.
19+
_THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL)
20+
_JSON_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$", re.MULTILINE)
21+
22+
23+
def json_schema_format(schema: Dict[str, Any], *, name: str = "response") -> Dict[str, Any]:
24+
"""Build a ``response_format`` value that constrains output to ``schema``.
25+
26+
The endpoint requires the schema nested under ``json_schema.schema``; passing
27+
it anywhere else is ignored server-side without an error, so prefer this
28+
helper over hand-building the dict.
29+
30+
Example:
31+
>>> json_schema_format({"type": "object", "properties": {"city": {"type": "string"}}})
32+
{'type': 'json_schema', 'json_schema': {'name': 'response', 'schema': {...}}}
33+
"""
34+
if not isinstance(schema, dict):
35+
raise ValidationError("schema must be a JSON Schema dict")
36+
return {"type": "json_schema", "json_schema": {"name": name, "schema": schema}}
37+
938

10-
def _build_payload(client: Any, messages: List[Dict[str, str]]) -> tuple:
39+
def _validate_response_format(response_format: Optional[Dict[str, Any]]) -> None:
40+
"""Reject shapes the endpoint would accept but silently not enforce."""
41+
if response_format is None:
42+
return
43+
if not isinstance(response_format, dict):
44+
raise ValidationError("response_format must be a dict")
45+
46+
fmt_type = response_format.get("type", "text")
47+
if fmt_type not in RESPONSE_FORMAT_TYPES:
48+
raise ValidationError(
49+
f"response_format type must be one of {', '.join(RESPONSE_FORMAT_TYPES)}; got {fmt_type!r}"
50+
)
51+
52+
if fmt_type == "json_schema":
53+
json_schema = response_format.get("json_schema")
54+
if not isinstance(json_schema, dict) or "schema" not in json_schema:
55+
raise ValidationError(
56+
'response_format {"type": "json_schema"} requires the schema nested as '
57+
'{"json_schema": {"schema": {...}}}. Without it the server returns 200 and '
58+
"ignores the constraint. Use diffbot.json_schema_format(schema) to build it."
59+
)
60+
61+
62+
def _build_payload(
63+
client: Any,
64+
messages: List[Dict[str, str]],
65+
*,
66+
response_format: Optional[Dict[str, Any]] = None,
67+
) -> tuple:
68+
_validate_response_format(response_format)
1169
headers = {"Authorization": f"Bearer {client.token}"}
12-
payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
70+
payload: Dict[str, Any] = {"model": MODEL, "messages": messages, "stream": True}
71+
if response_format is not None:
72+
payload["response_format"] = response_format
1373
return headers, payload
1474

1575

@@ -24,9 +84,68 @@ def _parse_chunk(line: str):
2484
return None
2585

2686

27-
def ask(client: "Diffbot", messages: List[Dict[str, str]]) -> Iterator[str]:
28-
headers = {"Authorization": f"Bearer {client.token}"}
29-
payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
87+
def _extract_json(text: str) -> Any:
88+
"""Parse the model's final answer as JSON, tolerating think blocks and fences."""
89+
cleaned = _THINK_BLOCK.sub("", text).strip()
90+
cleaned = _JSON_FENCE.sub("", cleaned).strip()
91+
92+
try:
93+
return json.loads(cleaned)
94+
except json.JSONDecodeError:
95+
pass
96+
97+
# Fall back to the outermost object or array span in the response.
98+
spans = []
99+
for opener, closer in (("{", "}"), ("[", "]")):
100+
start, end = cleaned.find(opener), cleaned.rfind(closer)
101+
if start != -1 and end > start:
102+
spans.append((start, cleaned[start : end + 1]))
103+
for _, span in sorted(spans):
104+
try:
105+
return json.loads(span)
106+
except json.JSONDecodeError:
107+
continue
108+
109+
raise ValidationError(f"could not parse JSON from the model response: {text[:200]!r}")
110+
111+
112+
#: Schema used when the caller wants JSON but has no shape in mind. This goes
113+
#: through the json_schema path rather than {"type": "json_object"} on purpose:
114+
#: json_object applies no server-side grammar, so the RAG loop's internal
115+
#: <functioncall> JSON satisfies it and gets returned as the final answer. That
116+
#: is reproducible whenever the request carries a system message.
117+
ANY_OBJECT_SCHEMA = {"type": "object"}
118+
119+
120+
def _resolve_format(
121+
schema: Optional[Dict[str, Any]],
122+
response_format: Optional[Dict[str, Any]],
123+
) -> Dict[str, Any]:
124+
if schema is not None and response_format is not None:
125+
raise ValidationError("pass either schema or response_format, not both")
126+
if response_format is not None:
127+
return response_format
128+
return json_schema_format(schema if schema is not None else ANY_OBJECT_SCHEMA)
129+
130+
131+
def _check_tool_call_leak(parsed: Any) -> Any:
132+
"""Catch the internal tool call surfacing as the answer (see ANY_OBJECT_SCHEMA)."""
133+
if isinstance(parsed, dict) and parsed.get("name") == "functioncall" and "arguments" in parsed:
134+
raise ValidationError(
135+
"the model returned its internal tool call instead of an answer; this happens with "
136+
'response_format {"type": "json_object"} because the server applies no grammar to it. '
137+
"Pass a schema instead."
138+
)
139+
return parsed
140+
141+
142+
def ask(
143+
client: "Diffbot",
144+
messages: List[Dict[str, str]],
145+
*,
146+
response_format: Optional[Dict[str, Any]] = None,
147+
) -> Iterator[str]:
148+
headers, payload = _build_payload(client, messages, response_format=response_format)
30149
with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
31150
client._raise_for_status(response)
32151
for line in response.iter_lines():
@@ -36,13 +155,41 @@ def ask(client: "Diffbot", messages: List[Dict[str, str]]) -> Iterator[str]:
36155
yield content
37156

38157

39-
async def ask_async(client: "DiffbotAsync", messages: List[Dict[str, str]]) -> AsyncIterator[str]:
40-
headers = {"Authorization": f"Bearer {client.token}"}
41-
payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
158+
async def ask_async(
159+
client: "DiffbotAsync",
160+
messages: List[Dict[str, str]],
161+
*,
162+
response_format: Optional[Dict[str, Any]] = None,
163+
) -> AsyncIterator[str]:
164+
headers, payload = _build_payload(client, messages, response_format=response_format)
42165
async with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
43166
client._raise_for_status(response)
44167
async for line in response.aiter_lines():
45168
if line:
46169
content = _parse_chunk(line)
47170
if content:
48171
yield content
172+
173+
174+
def ask_json(
175+
client: "Diffbot",
176+
messages: List[Dict[str, str]],
177+
schema: Optional[Dict[str, Any]] = None,
178+
*,
179+
response_format: Optional[Dict[str, Any]] = None,
180+
) -> Any:
181+
fmt = _resolve_format(schema, response_format)
182+
text = "".join(ask(client, messages, response_format=fmt))
183+
return _check_tool_call_leak(_extract_json(text))
184+
185+
186+
async def ask_json_async(
187+
client: "DiffbotAsync",
188+
messages: List[Dict[str, str]],
189+
schema: Optional[Dict[str, Any]] = None,
190+
*,
191+
response_format: Optional[Dict[str, Any]] = None,
192+
) -> Any:
193+
fmt = _resolve_format(schema, response_format)
194+
chunks = [chunk async for chunk in ask_async(client, messages, response_format=fmt)]
195+
return _check_tool_call_leak(_extract_json("".join(chunks)))

‎src/diffbot/cli/__init__.py‎

Lines changed: 42 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,15 @@
1111
from rich.progress import Progress, SpinnerColumn, TextColumn
1212
from datetime import datetime
1313

14-
from diffbot import __version__, CrawlEventType, AuthError, ExtractionError, APIError
14+
from diffbot import (
15+
__version__,
16+
CrawlEventType,
17+
AuthError,
18+
ExtractionError,
19+
APIError,
20+
ValidationError,
21+
json_schema_format,
22+
)
1523

1624
from ._common import get_client
1725

@@ -86,10 +94,29 @@ def extract(url, output, fmt, api):
8694
@click.argument("prompt")
8795
@click.option("-o", "--output", type=click.Path(), help="Write output to file instead of stdout")
8896
@click.option("--json", "as_json", is_flag=True, help="Output from LLM as JSON")
89-
def ask(prompt: str, output: str = None, as_json: bool = False):
97+
@click.option(
98+
"--schema",
99+
"schema_path",
100+
type=click.Path(exists=True, dir_okay=False),
101+
help="JSON Schema file constraining the output (implies --json)",
102+
)
103+
def ask(prompt: str, output: str = None, as_json: bool = False, schema_path: str = None):
90104
"""Ask a question to the Diffbot LLM"""
91105
db = get_client()
92106
stdin_content = sys.stdin.read() if not sys.stdin.isatty() else None
107+
108+
response_format = None
109+
if schema_path:
110+
as_json = True
111+
try:
112+
with open(schema_path) as f:
113+
response_format = json_schema_format(json.load(f))
114+
except json.JSONDecodeError as e:
115+
click.echo(f"Error: {schema_path} is not valid JSON: {e}", err=True)
116+
raise click.Abort()
117+
118+
# Without --schema, response_format stays None and ask_json picks its own
119+
# permissive default; --json alone must not become {"type": "json_object"}.
93120
interactive_mode = is_interactive and not output and not as_json
94121

95122
messages = [
@@ -100,9 +127,6 @@ def ask(prompt: str, output: str = None, as_json: bool = False):
100127
{"role": "user", "content": prompt},
101128
]
102129

103-
if as_json:
104-
messages[1]["content"] += "\nReturn the output as a JSON object. Do not include any other text outside of the JSON object."
105-
106130
if stdin_content:
107131
messages[1]["content"] = f"<input>{stdin_content}</input>\n" + messages[1]["content"]
108132

@@ -137,17 +161,20 @@ def ask(prompt: str, output: str = None, as_json: bool = False):
137161
else:
138162
response = ""
139163
if output:
140-
for chunk in db.ask(messages):
141-
response += chunk
164+
if as_json:
165+
response = json.dumps(
166+
db.ask_json(messages, response_format=response_format), indent=2
167+
)
168+
else:
169+
for chunk in db.ask(messages):
170+
response += chunk
142171
with open(output, "w") as f:
143172
f.write(response)
144173
click.echo(f"Output written to {output}")
145174
elif as_json:
146-
buffer = ""
147-
for chunk in db.ask(messages):
148-
buffer += chunk
149-
buffer = buffer[buffer.find("{") : buffer.rfind("}") + 1]
150-
sys.stdout.write(buffer)
175+
sys.stdout.write(
176+
json.dumps(db.ask_json(messages, response_format=response_format), indent=2)
177+
)
151178
else:
152179
for chunk in db.ask(messages):
153180
sys.stdout.write(chunk)
@@ -156,6 +183,9 @@ def ask(prompt: str, output: str = None, as_json: bool = False):
156183
except AuthError:
157184
click.echo("Error: Invalid or unauthorized API token.", err=True)
158185
raise click.Abort()
186+
except ValidationError as e:
187+
click.echo(f"Error: {e}", err=True)
188+
raise click.Abort()
159189
except APIError as e:
160190
click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
161191
raise click.Abort()

0 commit comments

Comments
 (0)