Skip to content

feat(llm): add AI_EXTRACT_RECORD for schema-driven structured extraction - #133

Open
Mrhs121 wants to merge 1 commit into
LadybugDB:mainfrom
Mrhs121:ai_extract_record
Open

Mrhs121 wants to merge 1 commit into
LadybugDB:mainfrom
Mrhs121:ai_extract_record

Conversation

@Mrhs121

@Mrhs121 Mrhs121 commented Oct 18, 2026 •

Copy link
Copy Markdown

Motivation

The existing AI_EXTRACT returns an untyped string and takes the API key, model and endpoint as SQL arguments. That works for ad-hoc queries, but not for two common deployments:

  1. Typed results. Callers want triage.urgency_score >= 7, not json_extract calls in every downstream query.
  2. Configuration outside SQL. When credentials are injected by the deployment layer and never stored in query text, they must not appear as SQL arguments.

Like DuckDB's ai_extract_record, this function derives the result type from a constant JSON Schema at bind time. Unlike DuckDB's version, provider configuration is read from environment variables so credentials stay out of SQL:

Variable Meaning
OPENAI_COMPATIBLE_API_KEY Bearer token. Required. Read at execution time.
OPENAI_COMPATIBLE_BASE_URL Default endpoint. Optional; overridden by argument 4.
OPENAI_COMPATIBLE_MODEL Default model. Optional; overridden by argument 3.

If a model or endpoint argument is omitted or empty, the corresponding environment variable is used, falling back to gpt-4o-mini and https://api.openai.com/v1. The variable names are the same ones DuckDB's AI functions use, so one deployment configuration works for
both engines.

Schema to type mapping

json_schema must be a constant string whose top-level "type" is "object". Each property becomes a STRUCT field:

JSON Schema type Lbug type
integer INT64
number DOUBLE
boolean BOOL
array LIST(STRING)
string (and anything else) STRING

description, enum, minimum, required and similar annotations are forwarded to the model but do not affect the bound type. This is the same projection rule as DuckDB's ai_extract_record.

Example

Support ticket triage, the same scenario used in the DuckDB AI cookbook:

LOAD EXTENSION 'libllm.lbug_extension';

CREATE NODE TABLE SupportTicket(ticket_id INT64, subject STRING, body STRING, PRIMARY KEY(ticket_id));
CREATE (:SupportTicket {ticket_id: 1001, subject: 'Login fails after upgrade',
        body: 'Since we upgraded to 6.2 last Friday, about 40 users in the Shanghai office cannot log in. The page reports an invalid session token. We have a customer demo on Wednesday.'}),
       (:SupportTicket {ticket_id: 1002, subject: 'Invoice shows the wrong currency',
        body: 'Our October invoice was issued in USD instead of EUR. Accounting cannot process it before the quarter closes on the 30th.'}),
       (:SupportTicket {ticket_id: 1003, subject: 'Feature request: dark mode',
        body: 'It would be nice to have a dark theme in the admin console. No rush, just a suggestion from our team.'});

MATCH (t:SupportTicket)
WITH t.ticket_id AS ticket_id,
     AI_EXTRACT_RECORD(
       t.subject + ': ' + t.body,
       '{
         "type": "object",
         "properties": {
           "product_area":       {"type": "string"},
           "needs_engineering":  {"type": "boolean"},
           "urgency_score":      {"type": "integer", "minimum": 0, "maximum": 10},
           "recommended_owner":  {"type": "string"}
         },
         "required": ["product_area", "needs_engineering", "urgency_score"]
       }'
     ) AS extracted
RETURN ticket_id,
       extracted.product_area       AS product_area,
       extracted.needs_engineering  AS needs_engineering,
       extracted.urgency_score      AS urgency_score,
       extracted.recommended_owner  AS recommended_owner
ORDER BY ticket_id;

Actual output (model: glm-5.3-flash against an OpenAI-compatible gateway):

┌───────────┬─────────────────────────────────────┬───────────────────┬───────────────┬───────────────────┐
│ ticket_id │ product_area                        │ needs_engineering │ urgency_score │ recommended_owner │
│ INT64     │ STRING                              │ BOOL              │ INT64         │ STRING            │
├───────────┼─────────────────────────────────────┼───────────────────┼───────────────┼───────────────────┤
│ 1001      │ Authentication / Session Management │ True              │ 8             │ Engineering       │
│ 1002      │ Billing                             │ True              │ 8             │ Billing team      │
│ 1003      │ admin console                       │ True              │ 2             │ Frontend/UI team  │
└───────────┴─────────────────────────────────────┴───────────────────┴───────────────┴───────────────────┘

Fields are ordinary typed columns, so results can be filtered and aggregated in Cypher without JSON parsing:

MATCH (t:SupportTicket)
WITH t, AI_EXTRACT_RECORD(t.subject + ': ' + t.body, $schema) AS x
WHERE x.needs_engineering AND x.urgency_score >= 7
RETURN t.ticket_id, x.product_area, x.urgency_score
ORDER BY x.urgency_score DESC;

As with DuckDB, call the function once per row and project the fields in an outer WITH/RETURN, so each row costs exactly one request.

Behavior notes

  • Rows are dispatched in parallel with up to 4 workers per chunk, the same ceiling AI_EXTRACT uses.
  • NULL text produces a NULL STRUCT without contacting the provider; the API key is only required when at least one non-null row is sent.
  • If the model output is not parseable as a JSON object, the whole STRUCT is set to NULL for that row instead of failing the statement (equivalent to DuckDB's on_error := 'null'). Markdown code fences around the JSON are stripped.
  • The source text is wrapped in <source_text> tags, and the system prompt instructs the model to treat it as data rather than instructions.
  • Nested objects currently map to STRING; nested STRUCT projection can be added later without changing the function signature.

Adds AI_EXTRACT_RECORD(text, json_schema[, model[, endpoint]]), which
returns a STRUCT derived from the JSON Schema's properties at bind time,
mirroring DuckDB's ai_extract_record. Configuration is read from
OPENAI_COMPATIBLE_API_KEY / _BASE_URL / _MODEL environment variables so
credentials never appear in SQL.

- schema -> STRUCT type mapping: integer/INT64, number/DOUBLE,
  boolean/BOOL, array/LIST(STRING), string/STRING
- parallel per-row dispatch, same worker ceiling as AI_EXTRACT
- NULL input yields NULL; unparseable model output yields a NULL struct
- add ai_extract_record_error.test
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant