Nodes/ComfyUI_LiteLLM/AddDataModelToLLLm
ComfyUI Node

AddDataModelToLLLm

Force the LLM into a JSON schema, no matter how much it wants to ramble

By Hopping-Mad-Games·Created 2 years ago·Updated 11 months ago· 7
AddDataModelToLLLm
  • model
  • data_model
  • Litellm model
codeclass UserModel(BaseModel): name: str= Field(..., description='The name of the person') age: int

If you've ever told an LLM "return JSON" and gotten back a paragraph that starts with a code fence and ends with a thank-you note, you know the pain this node fixes. AddDataModelToLLLm bolts a structured-output schema onto your model config so the API is contractually obligated to hand you clean, typed JSON - not prose that happens to contain JSON.

The name is a typo - it's "LLLm", not "LLM", and yes, the pack author shipped it that way. It's part of the LiteLLM pack, so it's an ETK/LLM/LiteLLM node and it's the "give me parseable output" layer between your model provider and your completion node.

What it does

It takes a model (from LiteLLMModelProvider or a custom endpoint), and attaches a Pydantic model as the response_format. In LiteLLM/OpenAI terms that's structured output: the API uses the schema to constrain the model's reply. You define the schema in the node's code field as a Pydantic class, and the node injects it into the model's kwargs so the next completion call returns something you can parse.

The default code gives you the shape to copy:

class UserModel(BaseModel):
    name: str = Field(..., description='The name of the person')
    age: int

The inputs and output

  • model - required, a LITELLM_MODEL from a provider node.
  • code (optional) - multiline Python defining a class named exactly UserModel. This is the schema source of truth.
  • data_model (optional) - a DATA_MODEL input, so you can pass a pre-built schema from another node instead of writing code here.

Output is a single Litellm model - the same model, now carrying response_format. Wire it into LiteLLMCompletion (or the provider nodes) and the structured output flows through.

How the code actually runs

This is the part to read twice. The node executes your code string with exec() inside a restricted namespace: a handful of builtins (print, range, type constructors), the typing names, and explicitly BaseModel, Field, and conlist from Pydantic. It then grabs the class named UserModel out of the namespace and uses it. If you define any other name, you get a "No model named 'UserModel'" error.

So: it's a sandbox, but it's a restricted sandbox, not a jail. exec() still runs code on your machine. Keep the schema definition to Pydantic classes and you're fine - just don't paste anything into code that you wouldn't run yourself. And keep in mind this is a custom node; the usual rule applies, don't expose a ComfyUI running this to the public internet.

Installing

It ships with ComfyUI_LiteLLM - ComfyUI Manager, search "ComfyUI_LiteLLM", or:

cd ComfyUI/custom_nodes
git clone https://github.com/Hopping-Mad-Games/ComfyUI_LiteLLM
cd ComfyUI_LiteLLM
pip install -r requirements.txt

Restart ComfyUI, find it under ETK/LLM/LiteLLM. No model downloads; you just need a provider key (e.g. OPENAI_API_KEY) in your environment.

Where people get burned

Structured output support is not universal. OpenAI and most OpenAI-compatible providers honor response_format; some Anthropic routes and older models are picky or ignore it, and then you're back to freeform text. Also, the schema is only as good as your field descriptions - the description= strings are what steer the model, so write them like the API docs you wish you had. And if you're getting JSON-in-a-code-fence anyway, check that the completion node you're using actually forwarded the model's response_format; if it swallowed the kwargs, the schema never reached the API.

CategoryETK/LLM/LiteLLM

Inputs (3)

NameTypeDefaultDescription
modelLITELLM_MODELanthropic/claude-3-haiku-20240307
codeoptSTRINGclass UserModel(BaseModel): name: str= Field(..., description='The name of the person') age: int
data_modeloptDATA_MODEL

Outputs (1)

NameTypeDescription
Litellm modelLITELLM_MODEL