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

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.

Author: tempusfrangit

integration-tests/tests/pydantic2_output.txtar

@@ -0,0 +1,37 @@
+# Test that Pydantic v2 BaseModel works as prediction output type.
+# Coglet's make_encodeable() must call model_dump() to serialize.
+
+# Build the image
+cog build -t $TEST_IMAGE
+
+# Predict returns structured Pydantic output
+cog predict $TEST_IMAGE -i name=alice -i score=0.95
+stdout '"name": "alice"'
+stdout '"score": 0.95'
+stdout '"tags"'
+stdout 'default'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+  python_packages:
+    - "pydantic>2"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from typing import List
+
+from pydantic import BaseModel as PydanticBaseModel
+
+from cog import BasePredictor
+
+
+class Result(PydanticBaseModel):
+    name: str
+    score: float
+    tags: List[str]
+
+
+class Predictor(BasePredictor):
+    def predict(self, name: str, score: float = 0.5) -> Result:
+        return Result(name=name, score=score, tags=["default"])

pkg/schema/python/parser.go

@@ -161,6 +161,19 @@ func parseImportFrom(node *sitter.Node, source []byte, ctx *schema.ImportContext
 				name := content(child, source)
 				ctx.Names.Set(name, schema.ImportEntry{Module: module, Original: name})
 			}
+		case "aliased_import":
+			// Single aliased import: `from X import name as alias`
+			origNode := child.ChildByFieldName("name")
+			aliasNode := child.ChildByFieldName("alias")
+			orig := ""
+			if origNode != nil {
+				orig = content(origNode, source)
+			}
+			alias := orig
+			if aliasNode != nil {
+				alias = content(aliasNode, source)
+			}
+			ctx.Names.Set(alias, schema.ImportEntry{Module: module, Original: orig})
 		case "import_list":
 			for _, importChild := range allChildren(child) {
 				switch importChild.Type() {
@@ -404,11 +417,18 @@ func inheritsFromBaseModel(classNode *sitter.Node, source []byte, imports *schem
 		return false
 	}
 	for _, child := range allChildren(supers) {
-		if child.Type() == "identifier" {
+		switch child.Type() {
+		case "identifier":
 			name := content(child, source)
 			if imports.IsBaseModel(name) || name == "BaseModel" {
 				return true
 			}
+		case "attribute":
+			// Handle dotted access: pydantic.BaseModel, cog.BaseModel
+			text := content(child, source)
+			if strings.HasSuffix(text, ".BaseModel") {
+				return true
+			}
 		}
 	}
 	return false

pkg/schema/python/parser_test.go

@@ -1002,6 +1002,78 @@ class Predictor(BasePredictor):
 	require.Equal(t, 0.0, conf.Default.Float)
 }
 
+// ---------------------------------------------------------------------------
+// Pydantic BaseModel output
+// ---------------------------------------------------------------------------
+
+func TestPydanticBaseModelOutput(t *testing.T) {
+	source := `
+from pydantic import BaseModel as PydanticBaseModel
+from cog import BasePredictor
+
+class Result(PydanticBaseModel):
+    name: str
+    score: float
+    tags: list[str]
+
+class Predictor(BasePredictor):
+    def predict(self, name: str) -> Result:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	require.Equal(t, schema.OutputObject, info.Output.Kind)
+	require.NotNil(t, info.Output.Fields)
+	require.Equal(t, 3, info.Output.Fields.Len())
+
+	name, ok := info.Output.Fields.Get("name")
+	require.True(t, ok)
+	require.Equal(t, schema.TypeString, name.FieldType.Primitive)
+
+	score, ok := info.Output.Fields.Get("score")
+	require.True(t, ok)
+	require.Equal(t, schema.TypeFloat, score.FieldType.Primitive)
+}
+
+func TestPydanticBaseModelDottedOutput(t *testing.T) {
+	source := `
+import pydantic
+from cog import BasePredictor
+
+class Result(pydantic.BaseModel):
+    text: str
+
+class Predictor(BasePredictor):
+    def predict(self, s: str) -> Result:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	require.Equal(t, schema.OutputObject, info.Output.Kind)
+
+	text, ok := info.Output.Fields.Get("text")
+	require.True(t, ok)
+	require.Equal(t, schema.TypeString, text.FieldType.Primitive)
+}
+
+func TestPydanticBaseModelDirectImport(t *testing.T) {
+	source := `
+from pydantic import BaseModel
+from cog import BasePredictor
+
+class Output(BaseModel):
+    value: int
+
+class Predictor(BasePredictor):
+    def predict(self, x: int) -> Output:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	require.Equal(t, schema.OutputObject, info.Output.Kind)
+
+	val, ok := info.Output.Fields.Get("value")
+	require.True(t, ok)
+	require.Equal(t, schema.TypeInteger, val.FieldType.Primitive)
+}
+
 // ---------------------------------------------------------------------------
 // No-input predictor (only self)
 // ---------------------------------------------------------------------------

pkg/schema/types.go

@@ -340,10 +340,10 @@ func (ctx *ImportContext) IsTypingType(name string) bool {
 	return false
 }
 
-// IsBaseModel returns true if name resolves to cog.BaseModel.
+// IsBaseModel returns true if name resolves to cog.BaseModel or pydantic.BaseModel.
 func (ctx *ImportContext) IsBaseModel(name string) bool {
 	if e, ok := ctx.Names.Get(name); ok {
-		return e.Module == "cog" && e.Original == "BaseModel"
+		return (e.Module == "cog" || e.Module == "pydantic" || e.Module == "pydantic.v1") && e.Original == "BaseModel"
 	}
 	return false
 }