Fix python partial dynamic streaming 0.36.0 (#651)

engine/language-client-codegen/src/python/generate_types.rs

@@ -41,6 +41,7 @@ pub(crate) struct PythonStreamTypes<'ir> {
 /// The Python class corresponding to Partial<TypeDefinedInBaml>
 struct PartialPythonClass<'ir> {
     name: &'ir str,
+    dynamic: bool,
     // the name, and the type of the field
     fields: Vec<(&'ir str, String)>,
 }
@@ -128,6 +129,7 @@ impl<'ir> From<ClassWalker<'ir>> for PartialPythonClass<'ir> {
     fn from(c: ClassWalker<'ir>) -> PartialPythonClass<'ir> {
         PartialPythonClass {
             name: c.name(),
+            dynamic: c.item.attributes.get("dynamic_type").is_some(),
             fields: c
                 .item
                 .elem

engine/language-client-codegen/src/python/templates/partial_types.py.j2

@@ -1,7 +1,7 @@
 {#- baml_py must be imported to enable access to baml_py.Image -#}
 import baml_py
 from enum import Enum
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 from typing import List, Optional, Union
 
 from . import types
@@ -16,7 +16,11 @@ from . import types
 {# Partial classes (used for streaming) -#}
 {% for cls in partial_classes %}
 class {{cls.name}}(BaseModel):
-    {% if cls.fields.is_empty() %}pass{% endif %}
+    {% if cls.dynamic %}
+    model_config = ConfigDict(extra='allow')
+    {%- endif %}
+    {% if cls.fields.is_empty() && !cls.dynamic %}pass{% endif %}
+
     {%- for (name, partial_type) in cls.fields %}
     {{name}}: {{partial_type}}
     {%- endfor %}

engine/language_client_python/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "baml-py"
-version = "0.35.1"
+version = "0.36.0"
 description = "BAML python bindings (pyproject.toml)"
 readme = "README.md"
 authors = [["Boundary", "[email protected]"]]

engine/language_client_typescript/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@boundaryml/baml",
-  "version": "0.35.0",
+  "version": "0.36.0",
   "description": "BAML typescript bindings (package.json)",
   "repository": {
     "type": "git",

integ-tests/python/app/test_functions.py

@@ -312,4 +312,26 @@ async def test_dynamic_class_output():
 
     output = await b.MyFunc(input="My name is Harrison. My hair is black and I'm 6 feet tall.", baml_options={"tb": tb})
     print(output.model_dump_json())
-    assert output.hair_color == "black"
+    assert output.hair_color == "black"
+
+@pytest.mark.asyncio
+async def test_stream_dynamic_class_output():
+    tb = TypeBuilder()
+    tb.DynamicOutput.add_property("hair_color", tb.string())
+    print(tb.DynamicOutput.list_properties())
+    for prop in tb.DynamicOutput.list_properties():
+        print(f"Property: {prop}")
+
+    stream = b.stream.MyFunc(input="My name is Harrison. My hair is black and I'm 6 feet tall.", baml_options={"tb": tb})
+    msgs = []
+    async for msg in stream:
+        print("streamed ", msg)
+        print("streamed ", msg.model_dump())
+        msgs.append(msg)
+    final = await stream.get_final_response()
+
+    assert len(msgs) > 0, "Expected at least one streamed response but got none."
+    print("final ", final)
+    print("final ", final.model_dump())
+    print("final ", final.model_dump_json())
+    assert final.hair_color == "black"

integ-tests/python/baml_client/partial_types.py

@@ -15,7 +15,7 @@
 # fmt: off
 import baml_py
 from enum import Enum
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 from typing import List, Optional, Union
 
 from . import types
@@ -30,39 +30,50 @@
 
 class Blah(BaseModel):
 
+
     prop4: Optional[str] = None
 
 class ClassOptionalOutput(BaseModel):
 
+
     prop1: Optional[str] = None
     prop2: Optional[str] = None
 
 class ClassOptionalOutput2(BaseModel):
 
+
     prop1: Optional[str] = None
     prop2: Optional[str] = None
     prop3: "Blah"
 
 class ClassWithImage(BaseModel):
 
+
     myImage: Optional[baml_py.Image] = None
     param2: Optional[str] = None
     fake_image: "FakeImage"
 
 class DynamicClassOne(BaseModel):
-    pass
+
+    model_config = ConfigDict(extra='allow')
+
 
 class DynamicClassTwo(BaseModel):
 
+    model_config = ConfigDict(extra='allow')
+
     hi: Optional[str] = None
     some_class: "SomeClassNestedDynamic"
     status: Optional[Union[types.DynEnumOne, str]] = None
 
 class DynamicOutput(BaseModel):
-    pass
+
+    model_config = ConfigDict(extra='allow')
+
 
 class Education(BaseModel):
 
+
     institution: Optional[str] = None
     location: Optional[str] = None
     degree: Optional[str] = None
@@ -71,56 +82,67 @@ class Education(BaseModel):
 
 class Email(BaseModel):
 
+
     subject: Optional[str] = None
     body: Optional[str] = None
     from_address: Optional[str] = None
 
 class Event(BaseModel):
 
+
     title: Optional[str] = None
     date: Optional[str] = None
     location: Optional[str] = None
     description: Optional[str] = None
 
 class FakeImage(BaseModel):
 
+
     url: Optional[str] = None
 
 class NamedArgsSingleClass(BaseModel):
 
+
     key: Optional[str] = None
     key_two: Optional[bool] = None
     key_three: Optional[int] = None
 
 class OptionalTest_Prop1(BaseModel):
 
+
     omega_a: Optional[str] = None
     omega_b: Optional[int] = None
 
 class OptionalTest_ReturnType(BaseModel):
 
+
     omega_1: "OptionalTest_Prop1"
     omega_2: Optional[str] = None
     omega_3: List[Optional[types.OptionalTest_CategoryType]]
 
 class OrderInfo(BaseModel):
 
+
     order_status: Optional[types.OrderStatus] = None
     tracking_number: Optional[str] = None
     estimated_arrival_date: Optional[str] = None
 
 class Person(BaseModel):
 
+    model_config = ConfigDict(extra='allow')
+
     name: Optional[str] = None
     hair_color: Optional[Union[types.Color, str]] = None
 
 class RaysData(BaseModel):
 
+
     dataType: Optional[types.DataType] = None
     value: Optional[Union["Resume", "Event"]] = None
 
 class Resume(BaseModel):
 
+
     name: Optional[str] = None
     email: Optional[str] = None
     phone: Optional[str] = None
@@ -130,6 +152,7 @@ class Resume(BaseModel):
 
 class SearchParams(BaseModel):
 
+
     dateRange: Optional[int] = None
     location: List[Optional[str]]
     jobTitle: "WithReasoning"
@@ -139,10 +162,13 @@ class SearchParams(BaseModel):
 
 class SomeClassNestedDynamic(BaseModel):
 
+    model_config = ConfigDict(extra='allow')
+
     hi: Optional[str] = None
 
 class TestClassAlias(BaseModel):
 
+
     key: Optional[str] = None
     key2: Optional[str] = None
     key3: Optional[str] = None
@@ -151,27 +177,32 @@ class TestClassAlias(BaseModel):
 
 class TestClassWithEnum(BaseModel):
 
+
     prop1: Optional[str] = None
     prop2: Optional[types.EnumInClass] = None
 
 class TestOutputClass(BaseModel):
 
+
     prop1: Optional[str] = None
     prop2: Optional[int] = None
 
 class TestOutputClassNested(BaseModel):
 
+
     prop1: Optional[str] = None
     prop2: Optional[int] = None
     prop3: "TestOutputClass"
 
 class UnionTest_ReturnType(BaseModel):
 
+
     prop1: Optional[Union[Optional[str], Optional[bool]]] = None
     prop2: List[Optional[Union[Optional[float], Optional[bool]]]]
     prop3: Optional[Union[List[Optional[float]], List[Optional[bool]]]] = None
 
 class WithReasoning(BaseModel):
 
+
     value: Optional[str] = None
     reasoning: Optional[str] = None

typescript/vscode-ext/packages/package.json

@@ -2,7 +2,7 @@
   "name": "baml-extension",
   "displayName": "Baml",
   "description": "BAML is a DSL for AI applications.",
-  "version": "0.35.1",
+  "version": "0.36.0",
   "publisher": "Boundary",
   "repository": "https://github.com/BoundaryML/baml",
   "homepage": "https://www.boundaryml.com",