11import type { DiffbotClient } from "./client.js" ;
2+ import { ValidationError } from "./errors.js" ;
23import { readBodyText , raiseForStatus } from "./http.js" ;
4+ import type { JsonObject } from "./json.js" ;
35
46export interface ChatMessage {
57 role : string ;
68 content : string ;
79}
810
11+ const MODEL = "diffbot-small-xl" ;
12+
13+ /** `response_format` types accepted by the Diffbot LLM endpoint. */
14+ export const RESPONSE_FORMAT_TYPES = [ "text" , "json_object" , "json_schema" ] as const ;
15+ export type ResponseFormatType = ( typeof RESPONSE_FORMAT_TYPES ) [ number ] ;
16+
17+ export interface ResponseFormat {
18+ type : ResponseFormatType ;
19+ json_schema ?: { name ?: string ; schema : JsonObject } ;
20+ }
21+
22+ export interface AskOptions {
23+ /**
24+ * Constrain the output, e.g. `{ type: "json_object" }` or `jsonSchemaFormat(schema)`.
25+ * Prefer a schema: `json_object` applies no server-side grammar.
26+ */
27+ responseFormat ?: ResponseFormat ;
28+ }
29+
30+ /**
31+ * Build a `response_format` that constrains output to `schema`. The endpoint requires
32+ * the schema nested under `json_schema.schema`; anywhere else is ignored server-side
33+ * without an error, so prefer this helper over hand-building the object.
34+ */
35+ export function jsonSchemaFormat ( schema : JsonObject , name = "response" ) : ResponseFormat {
36+ if ( schema === null || typeof schema !== "object" || Array . isArray ( schema ) ) {
37+ throw new ValidationError ( "schema must be a JSON Schema object" ) ;
38+ }
39+ return { type : "json_schema" , json_schema : { name, schema } } ;
40+ }
41+
42+ /** Reject shapes the endpoint would accept but silently not enforce. */
43+ function validateResponseFormat ( responseFormat : ResponseFormat | undefined ) : void {
44+ if ( responseFormat === undefined ) return ;
45+ if ( responseFormat === null || typeof responseFormat !== "object" || Array . isArray ( responseFormat ) ) {
46+ throw new ValidationError ( "responseFormat must be an object" ) ;
47+ }
48+ const fmtType = responseFormat . type ?? "text" ;
49+ if ( ! ( RESPONSE_FORMAT_TYPES as readonly string [ ] ) . includes ( fmtType ) ) {
50+ throw new ValidationError (
51+ `responseFormat type must be one of ${ RESPONSE_FORMAT_TYPES . join ( ", " ) } ${ JSON . stringify ( fmtType ) } ,
52+ ) ;
53+ }
54+ if ( fmtType === "json_schema" ) {
55+ const jsonSchema = responseFormat . json_schema ;
56+ if ( jsonSchema === null || typeof jsonSchema !== "object" || ! ( "schema" in jsonSchema ) ) {
57+ throw new ValidationError (
58+ 'responseFormat { type: "json_schema" } requires the schema nested as ' +
59+ "{ json_schema: { schema: {...} } }. Without it the server returns 200 and " +
60+ "ignores the constraint. Use jsonSchemaFormat(schema) to build it." ,
61+ ) ;
62+ }
63+ }
64+ }
65+
66+ // The RAG loop may prefix its final answer with a think block, and the model
67+ // occasionally wraps JSON in a markdown fence despite being told not to.
68+ const THINK_BLOCK = / < t h i n k > [ \s \S ] * ?< \/ t h i n k > / g;
69+ const JSON_FENCE = / ^ ` ` ` (?: j s o n ) ? \s * | \s * ` ` ` $ / gm;
70+
71+ /** Parse the model's final answer as JSON, tolerating think blocks and fences. */
72+ export function extractJson ( text : string ) : unknown {
73+ let cleaned = text . replace ( THINK_BLOCK , "" ) . trim ( ) ;
74+ cleaned = cleaned . replace ( JSON_FENCE , "" ) . trim ( ) ;
75+
76+ try {
77+ return JSON . parse ( cleaned ) ;
78+ } catch {
79+ // fall through to span search
80+ }
81+
82+ // Fall back to the outermost object or array span in the response.
83+ const spans : Array < [ number , string ] > = [ ] ;
84+ for ( const [ opener , closer ] of [
85+ [ "{" , "}" ] ,
86+ [ "[" , "]" ] ,
87+ ] as const ) {
88+ const start = cleaned . indexOf ( opener ) ;
89+ const end = cleaned . lastIndexOf ( closer ) ;
90+ if ( start !== - 1 && end > start ) spans . push ( [ start , cleaned . slice ( start , end + 1 ) ] ) ;
91+ }
92+ spans . sort ( ( a , b ) => a [ 0 ] - b [ 0 ] ) ;
93+ for ( const [ , span ] of spans ) {
94+ try {
95+ return JSON . parse ( span ) ;
96+ } catch {
97+ continue ;
98+ }
99+ }
100+
101+ throw new ValidationError (
102+ `could not parse JSON from the model response: ${ JSON . stringify ( text . slice ( 0 , 200 ) ) } ,
103+ ) ;
104+ }
105+
106+ /**
107+ * Schema used when the caller wants JSON but has no shape in mind. This goes through
108+ * the json_schema path rather than `{ type: "json_object" }` on purpose: json_object
109+ * applies no server-side grammar, so the RAG loop's internal <functioncall> JSON
110+ * satisfies it and gets returned as the final answer. That is reproducible whenever
111+ * the request carries a system message.
112+ */
113+ export const ANY_OBJECT_SCHEMA : JsonObject = { type : "object" } ;
114+
115+ function resolveFormat (
116+ schema : JsonObject | undefined ,
117+ responseFormat : ResponseFormat | undefined ,
118+ ) : ResponseFormat {
119+ if ( schema !== undefined && responseFormat !== undefined ) {
120+ throw new ValidationError ( "pass either schema or responseFormat, not both" ) ;
121+ }
122+ if ( responseFormat !== undefined ) return responseFormat ;
123+ return jsonSchemaFormat ( schema ?? ANY_OBJECT_SCHEMA ) ;
124+ }
125+
126+ /** Catch the internal tool call surfacing as the answer (see ANY_OBJECT_SCHEMA). */
127+ function checkToolCallLeak < T > ( parsed : T ) : T {
128+ if (
129+ parsed !== null &&
130+ typeof parsed === "object" &&
131+ ! Array . isArray ( parsed ) &&
132+ ( parsed as JsonObject ) . name === "functioncall" &&
133+ "arguments" in ( parsed as JsonObject )
134+ ) {
135+ throw new ValidationError (
136+ "the model returned its internal tool call instead of an answer; this happens with " +
137+ 'responseFormat { type: "json_object" } because the server applies no grammar to it. ' +
138+ "Pass a schema instead." ,
139+ ) ;
140+ }
141+ return parsed ;
142+ }
143+
9144function parseChunk ( line : string ) : string | undefined {
10145 try {
11146 const chunk = JSON . parse ( line . replace ( / ^ d a t a : / , "" ) ) as {
@@ -45,9 +180,12 @@ async function* linesFromStream(body: ReadableStream<Uint8Array>): AsyncGenerato
45180export async function * ask (
46181 client : DiffbotClient ,
47182 messages : ChatMessage [ ] ,
183+ options : AskOptions = { } ,
48184) : AsyncGenerator < string > {
185+ validateResponseFormat ( options . responseFormat ) ;
49186 const headers = { Authorization : `Bearer ${ client . token } } ;
50- const payload = { model : "diffbot-small-xl" , messages, stream : true } ;
187+ const payload : JsonObject = { model : MODEL , messages, stream : true } ;
188+ if ( options . responseFormat !== undefined ) payload . response_format = options . responseFormat ;
51189
52190 const response = await client . http . post ( client . llmUrl , { headers, json : payload } ) ;
53191 if ( ! response . ok ) {
@@ -70,3 +208,20 @@ export async function* ask(
70208 if ( content ) yield content ;
71209 }
72210}
211+
212+ /**
213+ * Ask the Diffbot LLM and return the answer parsed as JSON. With `schema` (a JSON Schema
214+ * object) the model is constrained to it by grammar; without one it falls back to
215+ * unconstrained JSON output. Pass `responseFormat` instead of `schema` for full control.
216+ */
217+ export async function askJson < T = unknown > (
218+ client : DiffbotClient ,
219+ messages : ChatMessage [ ] ,
220+ schema ?: JsonObject ,
221+ options : AskOptions = { } ,
222+ ) : Promise < T > {
223+ const responseFormat = resolveFormat ( schema , options . responseFormat ) ;
224+ let text = "" ;
225+ for await ( const chunk of ask ( client , messages , { responseFormat } ) ) text += chunk ;
226+ return checkToolCallLeak ( extractJson ( text ) ) as T ;
227+ }
0 commit comments