fix: add additionalProperties: false to structured output schemas for Databricks compatibility

[debu-sinha]

Summary

This PR fixes a compatibility issue with Databricks AI Gateway that prevents TruLens structured outputs from working with Databricks model serving endpoints.

Fixes: #2307

cc: @joshreini1 @piotrm0 @sfc-gh-dkurokawa

Problem

Databricks model serving endpoints require additionalProperties: false in JSON schemas for structured outputs (when using response_format). Without this property, TruLens feedback evaluations fail with:

BadRequestError: Invalid schema for response_format 'ChainOfThoughtResponse':
In context=(), 'additionalProperties' is required to be supplied and to be false.

This affects users who want to use TruLens with:

• Databricks AI Gateway
• LiteLLM with Databricks endpoints
• Any provider that requires strict JSON schema validation

Solution

Add model_config = ConfigDict(extra="forbid") to Pydantic models in output_schemas.py. This causes Pydantic to include additionalProperties: false in the generated JSON schema.

Code Change

# src/feedback/trulens/feedback/output_schemas.py

from pydantic import BaseModel
from pydantic import ConfigDict  # Added
from pydantic import Field


class BaseFeedbackResponse(BaseModel):
    """
    A base model for feedback responses.
    """

    model_config = ConfigDict(extra="forbid")  # Added - enables additionalProperties: false

    score: int = Field(description="The score based on the given criteria.")


class ChainOfThoughtResponse(BaseModel):
    """
    A model to represent the response from a Chain of Thought (COT) evaluation.
    """

    model_config = ConfigDict(extra="forbid")  # Added - enables additionalProperties: false

    criteria: str = Field(description="The criteria for the evaluation.")
    supporting_evidence: str = Field(
        description="Supporting evidence for the score, detailing the reasoning step by step."
    )
    score: int = Field(description="The score based on the given criteria.")

Schema Before (without the fix):

{
  "properties": {
    "criteria": {"description": "The criteria for the evaluation.", "title": "Criteria", "type": "string"},
    "supporting_evidence": {"description": "Supporting evidence...", "title": "Supporting Evidence", "type": "string"},
    "score": {"description": "The score based on the given criteria.", "title": "Score", "type": "integer"}
  },
  "required": ["criteria", "supporting_evidence", "score"],
  "title": "ChainOfThoughtResponse",
  "type": "object"
}

Schema After (with the fix):

{
  "additionalProperties": false,
  "properties": {
    "criteria": {"description": "The criteria for the evaluation.", "title": "Criteria", "type": "string"},
    "supporting_evidence": {"description": "Supporting evidence...", "title": "Supporting Evidence", "type": "string"},
    "score": {"description": "The score based on the given criteria.", "title": "Score", "type": "integer"}
  },
  "required": ["criteria", "supporting_evidence", "score"],
  "title": "ChainOfThoughtResponse",
  "type": "object"
}

Changes

File	Change
src/feedback/trulens/feedback/output_schemas.py	Added model_config = ConfigDict(extra="forbid") to both models
tests/unit/test_output_schemas.py	Added 5 new unit tests to verify schema and runtime behavior

Testing

# Run new unit tests
poetry run pytest tests/unit/test_output_schemas.py -v
# Result: 5 passed

# Run existing OpenAI capability tests
TEST_OPTIONAL=1 poetry run pytest tests/unit/providers/test_openai_capabilities.py -v
# Result: 11 passed

Test Verification

>>> from trulens.feedback.output_schemas import ChainOfThoughtResponse
>>> import json
>>> print(json.dumps(ChainOfThoughtResponse.model_json_schema(), indent=2))
{
  "additionalProperties": false,  # <-- Now included!
  "properties": {...},
  "required": [...],
  "title": "ChainOfThoughtResponse",
  "type": "object"
}

Backward Compatibility

This change is backward compatible:

• JSON schema output changes from implicit to explicit additionalProperties: false
• Runtime behavior now rejects extra fields at validation time (stricter validation is generally safe)
• No API or interface changes
• Existing code that uses these models correctly will continue to work

Related Issues

• TruLens Issue: LiteLLM provider: Databricks structured outputs fail due to missing additionalProperties: false in JSON schema #2307
• MLflow integration: mlflow/mlflow#19327

[sfc-gh-jreini]

Thanks again for contributing!

Can you update the API tests by running make write-api from root in this branch? That should get the last tests passing

[debu-sinha]

@sfc-gh-jreini I've set up Python 3.11 locally and ran make write-api. The golden files have been updated and pushed. The changes are just version bumps from 2.4.2 to 2.5.0 in the API golden file.

[sfc-gh-nvytla]

lgtm 👍

[sfc-gh-jreini]

works well on my side for non-dbrx providers 👍