From ba702d0bf8fd59e754ee4468bf79d06beddcf26c Mon Sep 17 00:00:00 2001 From: aaronvg Date: Fri, 7 Jun 2024 11:14:09 -0700 Subject: [PATCH] Fix python partial dynamic streaming 0.36.0 (#651) --- .../src/python/generate_types.rs | 2 + .../src/python/templates/partial_types.py.j2 | 8 +++- engine/language_client_python/pyproject.toml | 2 +- .../language_client_typescript/package.json | 2 +- integ-tests/python/app/test_functions.py | 24 +++++++++++- .../python/baml_client/partial_types.py | 37 +++++++++++++++++-- typescript/vscode-ext/packages/package.json | 2 +- 7 files changed, 68 insertions(+), 9 deletions(-) diff --git a/engine/language-client-codegen/src/python/generate_types.rs b/engine/language-client-codegen/src/python/generate_types.rs index 5276d9375..99bb0dba8 100644 --- a/engine/language-client-codegen/src/python/generate_types.rs +++ b/engine/language-client-codegen/src/python/generate_types.rs @@ -41,6 +41,7 @@ pub(crate) struct PythonStreamTypes<'ir> { /// The Python class corresponding to Partial 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> 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 diff --git a/engine/language-client-codegen/src/python/templates/partial_types.py.j2 b/engine/language-client-codegen/src/python/templates/partial_types.py.j2 index 123690324..03aa6c519 100644 --- a/engine/language-client-codegen/src/python/templates/partial_types.py.j2 +++ b/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 %} diff --git a/engine/language_client_python/pyproject.toml b/engine/language_client_python/pyproject.toml index 4907784c6..fbcb24f6a 100644 --- a/engine/language_client_python/pyproject.toml +++ b/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", "contact@boundaryml.com"]] diff --git a/engine/language_client_typescript/package.json b/engine/language_client_typescript/package.json index 20da9fd96..955ed578c 100644 --- a/engine/language_client_typescript/package.json +++ b/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", diff --git a/integ-tests/python/app/test_functions.py b/integ-tests/python/app/test_functions.py index 6bfdef076..04fa1e981 100644 --- a/integ-tests/python/app/test_functions.py +++ b/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" \ No newline at end of file + 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" \ No newline at end of file diff --git a/integ-tests/python/baml_client/partial_types.py b/integ-tests/python/baml_client/partial_types.py index dc4204cee..64d161d27 100644 --- a/integ-tests/python/baml_client/partial_types.py +++ b/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,12 +82,14 @@ 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 @@ -84,43 +97,52 @@ class Event(BaseModel): 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 diff --git a/typescript/vscode-ext/packages/package.json b/typescript/vscode-ext/packages/package.json index f54530404..9efe17520 100644 --- a/typescript/vscode-ext/packages/package.json +++ b/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",