Skip to content

fix: support Pydantic BaseModel as prediction output type in schema gen#2785

Merged
tempusfrangit merged 1 commit intomainfrom
test/pydantic-output-integration
Feb 27, 2026
Merged

fix: support Pydantic BaseModel as prediction output type in schema gen#2785
tempusfrangit merged 1 commit intomainfrom
test/pydantic-output-integration

Conversation

@tempusfrangit
Copy link
Member

The tree-sitter schema parser only recognized cog.BaseModel subclasses. Pydantic BaseModel subclasses failed with 'unsupported type' because:

  1. IsBaseModel() only checked module=='cog', not 'pydantic'
  2. inheritsFromBaseModel() only handled 'identifier' AST nodes, missing 'attribute' nodes for dotted access (pydantic.BaseModel)
  3. parseImportFrom() only handled aliased_import inside import_list, missing top-level aliased imports (from X import Y as Z)

Now supports all three import styles:

  • from pydantic import BaseModel
  • from pydantic import BaseModel as PydanticBaseModel
  • import pydantic; class Foo(pydantic.BaseModel)

Includes 3 unit tests and 1 integration test.

The tree-sitter schema parser only recognized cog.BaseModel subclasses.
Pydantic BaseModel subclasses failed with 'unsupported type' because:

1. IsBaseModel() only checked module=='cog', not 'pydantic'
2. inheritsFromBaseModel() only handled 'identifier' AST nodes,
   missing 'attribute' nodes for dotted access (pydantic.BaseModel)
3. parseImportFrom() only handled aliased_import inside import_list,
   missing top-level aliased imports (from X import Y as Z)

Now supports all three import styles:
  - from pydantic import BaseModel
  - from pydantic import BaseModel as PydanticBaseModel
  - import pydantic; class Foo(pydantic.BaseModel)

Includes 3 unit tests and 1 integration test.
@tempusfrangit tempusfrangit requested a review from a team as a code owner February 27, 2026 23:48
@tempusfrangit tempusfrangit enabled auto-merge (squash) February 27, 2026 23:50
@tempusfrangit tempusfrangit added this to the 0.17.0 Release milestone Feb 27, 2026
@tempusfrangit tempusfrangit merged commit 5418c4d into main Feb 27, 2026
31 checks passed
@tempusfrangit tempusfrangit deleted the test/pydantic-output-integration branch February 27, 2026 23:56
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.

2 participants