gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Add baml client changes

Browse source
This commit is contained in:
Igor Ilic 2025-09-08 12:55:33 +02:00
parent 7c27546951
commit 6231882c24
6 changed files with 344 additions and 129 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/__init__.py
View file

@ -10,7 +10,7 @@
# BAML files and re-generate this code using: baml-cli generate # BAML files and re-generate this code using: baml-cli generate
# baml-cli is available with the baml package. # baml-cli is available with the baml package.
__version__ = "0.201.0" __version__ = "0.206.1"
try: try:
from baml_py.safe_import import EnsureBamlPyImport from baml_py.safe_import import EnsureBamlPyImport

185
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/async_client.py
View file

@ -44,6 +44,7 @@ class BamlAsyncClient:
typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]] typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
] = None, ] = None,
env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None, env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
on_tick: typing.Optional[typing.Callable[[str, baml_py.baml_py.FunctionLog], None]] = None,
) -> "BamlAsyncClient": ) -> "BamlAsyncClient":
options: BamlCallOptions = {} options: BamlCallOptions = {}
if tb is not None: if tb is not None:
@ -54,6 +55,8 @@ class BamlAsyncClient:
options["collector"] = collector options["collector"] = collector
if env is not None: if env is not None:
options["env"] = env options["env"] = env
if on_tick is not None:
options["on_tick"] = on_tick
return BamlAsyncClient(self.__options.merge_options(options)) return BamlAsyncClient(self.__options.merge_options(options))
@property @property
@ -83,33 +86,52 @@ class BamlAsyncClient:
user_prompt: str, user_prompt: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DynamicModel: ) -> types.DynamicModel:
result = await self.__options.merge_options(baml_options).call_function_async( # Check if on_tick is provided
function_name="AcreateStructuredOutput", if "on_tick" in baml_options:
args={ # Use streaming internally when on_tick is provided
"content": content, stream = self.stream.AcreateStructuredOutput(
"system_prompt": system_prompt, content=content,
"user_prompt": user_prompt, system_prompt=system_prompt,
}, user_prompt=user_prompt,
) baml_options=baml_options,
return typing.cast( )
types.DynamicModel, result.cast_to(types, types, stream_types, False, __runtime__) return await stream.get_final_response()
) else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="AcreateStructuredOutput",
args={
"content": content,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
)
return typing.cast(
types.DynamicModel, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractCategories( async def ExtractCategories(
self, self,
content: str, content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction: ) -> types.DefaultContentPrediction:
result = await self.__options.merge_options(baml_options).call_function_async( # Check if on_tick is provided
function_name="ExtractCategories", if "on_tick" in baml_options:
args={ # Use streaming internally when on_tick is provided
"content": content, stream = self.stream.ExtractCategories(content=content, baml_options=baml_options)
}, return await stream.get_final_response()
) else:
return typing.cast( # Original non-streaming code
types.DefaultContentPrediction, result = await self.__options.merge_options(baml_options).call_function_async(
result.cast_to(types, types, stream_types, False, __runtime__), function_name="ExtractCategories",
) args={
"content": content,
},
)
return typing.cast(
types.DefaultContentPrediction,
result.cast_to(types, types, stream_types, False, __runtime__),
)
async def ExtractContentGraphGeneric( async def ExtractContentGraphGeneric(
self, self,
@ -126,17 +148,29 @@ class BamlAsyncClient:
custom_prompt_content: typing.Optional[str] = None, custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph: ) -> types.KnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async( # Check if on_tick is provided
function_name="ExtractContentGraphGeneric", if "on_tick" in baml_options:
args={ # Use streaming internally when on_tick is provided
"content": content, stream = self.stream.ExtractContentGraphGeneric(
"mode": mode, content=content,
"custom_prompt_content": custom_prompt_content, mode=mode,
}, custom_prompt_content=custom_prompt_content,
) baml_options=baml_options,
return typing.cast( )
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__) return await stream.get_final_response()
) else:
# Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractContentGraphGeneric",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def ExtractDynamicContentGraph( async def ExtractDynamicContentGraph(
self, self,
@ -153,48 +187,75 @@ class BamlAsyncClient:
custom_prompt_content: typing.Optional[str] = None, custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph: ) -> types.DynamicKnowledgeGraph:
result = await self.__options.merge_options(baml_options).call_function_async( # Check if on_tick is provided
function_name="ExtractDynamicContentGraph", if "on_tick" in baml_options:
args={ # Use streaming internally when on_tick is provided
"content": content, stream = self.stream.ExtractDynamicContentGraph(
"mode": mode, content=content,
"custom_prompt_content": custom_prompt_content, mode=mode,
}, custom_prompt_content=custom_prompt_content,
) baml_options=baml_options,
return typing.cast( )
types.DynamicKnowledgeGraph, return await stream.get_final_response()
result.cast_to(types, types, stream_types, False, __runtime__), else:
) # Original non-streaming code
result = await self.__options.merge_options(baml_options).call_function_async(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.DynamicKnowledgeGraph,
result.cast_to(types, types, stream_types, False, __runtime__),
)
async def SummarizeCode( async def SummarizeCode(
self, self,
content: str, content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedCode: ) -> types.SummarizedCode:
result = await self.__options.merge_options(baml_options).call_function_async( # Check if on_tick is provided
function_name="SummarizeCode", if "on_tick" in baml_options:
args={ # Use streaming internally when on_tick is provided
"content": content, stream = self.stream.SummarizeCode(content=content, baml_options=baml_options)
}, return await stream.get_final_response()
) else:
return typing.cast( # Original non-streaming code
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__) result = await self.__options.merge_options(baml_options).call_function_async(
) function_name="SummarizeCode",
args={
"content": content,
},
)
return typing.cast(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
)
async def SummarizeContent( async def SummarizeContent(
self, self,
content: str, content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedContent: ) -> types.SummarizedContent:
result = await self.__options.merge_options(baml_options).call_function_async( # Check if on_tick is provided
function_name="SummarizeContent", if "on_tick" in baml_options:
args={ # Use streaming internally when on_tick is provided
"content": content, stream = self.stream.SummarizeContent(content=content, baml_options=baml_options)
}, return await stream.get_final_response()
) else:
return typing.cast( # Original non-streaming code
types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__) result = await self.__options.merge_options(baml_options).call_function_async(
) function_name="SummarizeContent",
args={
"content": content,
},
)
return typing.cast(
types.SummarizedContent,
result.cast_to(types, types, stream_types, False, __runtime__),
)
class BamlStreamClient: class BamlStreamClient:

2
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/inlinedbaml.py
View file

File diff suppressed because one or more lines are too long

78
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/runtime.py
View file

@ -30,6 +30,10 @@ class BamlCallOptions(typing.TypedDict, total=False):
collector: typing_extensions.NotRequired[ collector: typing_extensions.NotRequired[
typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]] typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
] ]
abort_controller: typing_extensions.NotRequired[baml_py.baml_py.AbortController]
on_tick: typing_extensions.NotRequired[
typing.Callable[[str, baml_py.baml_py.FunctionLog], None]
]
class _ResolvedBamlOptions: class _ResolvedBamlOptions:
@ -37,6 +41,8 @@ class _ResolvedBamlOptions:
client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] client_registry: typing.Optional[baml_py.baml_py.ClientRegistry]
collectors: typing.List[baml_py.baml_py.Collector] collectors: typing.List[baml_py.baml_py.Collector]
env_vars: typing.Dict[str, str] env_vars: typing.Dict[str, str]
abort_controller: typing.Optional[baml_py.baml_py.AbortController]
on_tick: typing.Optional[typing.Callable[[], None]]
def __init__( def __init__(
self, self,
@ -44,11 +50,15 @@ class _ResolvedBamlOptions:
client_registry: typing.Optional[baml_py.baml_py.ClientRegistry], client_registry: typing.Optional[baml_py.baml_py.ClientRegistry],
collectors: typing.List[baml_py.baml_py.Collector], collectors: typing.List[baml_py.baml_py.Collector],
env_vars: typing.Dict[str, str], env_vars: typing.Dict[str, str],
abort_controller: typing.Optional[baml_py.baml_py.AbortController],
on_tick: typing.Optional[typing.Callable[[], None]],
): ):
self.tb = tb self.tb = tb
self.client_registry = client_registry self.client_registry = client_registry
self.collectors = collectors self.collectors = collectors
self.env_vars = env_vars self.env_vars = env_vars
self.abort_controller = abort_controller
self.on_tick = on_tick
class DoNotUseDirectlyCallManager: class DoNotUseDirectlyCallManager:
@ -85,11 +95,27 @@ class DoNotUseDirectlyCallManager:
else: else:
env_vars.pop(k, None) env_vars.pop(k, None)
abort_controller = self.__baml_options.get("abort_controller")
on_tick = self.__baml_options.get("on_tick")
if on_tick is not None:
collector = baml_py.baml_py.Collector("on-tick-collector")
collectors_as_list.append(collector)
def on_tick_wrapper():
log = collector.last
if log is not None:
on_tick("Unknown", log)
else:
on_tick_wrapper = None
return _ResolvedBamlOptions( return _ResolvedBamlOptions(
baml_tb, baml_tb,
client_registry, client_registry,
collectors_as_list, collectors_as_list,
env_vars, env_vars,
abort_controller,
on_tick_wrapper,
) )
def merge_options(self, options: BamlCallOptions) -> "DoNotUseDirectlyCallManager": def merge_options(self, options: BamlCallOptions) -> "DoNotUseDirectlyCallManager":
@ -99,6 +125,14 @@ class DoNotUseDirectlyCallManager:
self, *, function_name: str, args: typing.Dict[str, typing.Any] self, *, function_name: str, args: typing.Dict[str, typing.Any]
) -> baml_py.baml_py.FunctionResult: ) -> baml_py.baml_py.FunctionResult:
resolved_options = self.__resolve() resolved_options = self.__resolve()
# Check if already aborted
if (
resolved_options.abort_controller is not None
and resolved_options.abort_controller.aborted
):
raise Exception("BamlAbortError: Operation was aborted")
return await __runtime__.call_function( return await __runtime__.call_function(
function_name, function_name,
args, args,
@ -112,12 +146,22 @@ class DoNotUseDirectlyCallManager:
resolved_options.collectors, resolved_options.collectors,
# env_vars # env_vars
resolved_options.env_vars, resolved_options.env_vars,
# abort_controller
resolved_options.abort_controller,
) )
def call_function_sync( def call_function_sync(
self, *, function_name: str, args: typing.Dict[str, typing.Any] self, *, function_name: str, args: typing.Dict[str, typing.Any]
) -> baml_py.baml_py.FunctionResult: ) -> baml_py.baml_py.FunctionResult:
resolved_options = self.__resolve() resolved_options = self.__resolve()
# Check if already aborted
if (
resolved_options.abort_controller is not None
and resolved_options.abort_controller.aborted
):
raise Exception("BamlAbortError: Operation was aborted")
ctx = __ctx__manager__.get() ctx = __ctx__manager__.get()
return __runtime__.call_function_sync( return __runtime__.call_function_sync(
function_name, function_name,
@ -132,6 +176,8 @@ class DoNotUseDirectlyCallManager:
resolved_options.collectors, resolved_options.collectors,
# env_vars # env_vars
resolved_options.env_vars, resolved_options.env_vars,
# abort_controller
resolved_options.abort_controller,
) )
def create_async_stream( def create_async_stream(
@ -158,6 +204,8 @@ class DoNotUseDirectlyCallManager:
resolved_options.collectors, resolved_options.collectors,
# env_vars # env_vars
resolved_options.env_vars, resolved_options.env_vars,
# on_tick
resolved_options.on_tick,
) )
return ctx, result return ctx, result
@ -170,6 +218,10 @@ class DoNotUseDirectlyCallManager:
baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream baml_py.baml_py.RuntimeContextManager, baml_py.baml_py.SyncFunctionResultStream
]: ]:
resolved_options = self.__resolve() resolved_options = self.__resolve()
if resolved_options.on_tick is not None:
raise ValueError(
"on_tick is not supported for sync streams. Please use async streams instead."
)
ctx = __ctx__manager__.get() ctx = __ctx__manager__.get()
result = __runtime__.stream_function_sync( result = __runtime__.stream_function_sync(
function_name, function_name,
@ -187,6 +239,9 @@ class DoNotUseDirectlyCallManager:
resolved_options.collectors, resolved_options.collectors,
# env_vars # env_vars
resolved_options.env_vars, resolved_options.env_vars,
# on_tick
# always None! sync streams don't support on_tick
None,
) )
return ctx, result return ctx, result
@ -264,3 +319,26 @@ class DoNotUseDirectlyCallManager:
# env_vars # env_vars
resolved_options.env_vars, resolved_options.env_vars,
) )
def disassemble(function: typing.Callable) -> None:
import inspect
from . import b
if not callable(function):
print(f"disassemble: object {function} is not a Baml function")
return
is_client_method = False
for method_name, _ in inspect.getmembers(b, predicate=inspect.ismethod):
if method_name == function.__name__:
is_client_method = True
break
if not is_client_method:
print(f"disassemble: function {function.__name__} is not a Baml function")
return
print(f"----- function {function.__name__} -----")
__runtime__.disassemble(function.__name__)

179
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/sync_client.py
View file

@ -57,6 +57,7 @@ class BamlSyncClient:
typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]] typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
] = None, ] = None,
env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None, env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
on_tick: typing.Optional[typing.Callable[[str, baml_py.baml_py.FunctionLog], None]] = None,
) -> "BamlSyncClient": ) -> "BamlSyncClient":
options: BamlCallOptions = {} options: BamlCallOptions = {}
if tb is not None: if tb is not None:
@ -67,6 +68,8 @@ class BamlSyncClient:
options["collector"] = collector options["collector"] = collector
if env is not None: if env is not None:
options["env"] = env options["env"] = env
if on_tick is not None:
options["on_tick"] = on_tick
return BamlSyncClient(self.__options.merge_options(options)) return BamlSyncClient(self.__options.merge_options(options))
@property @property
@ -96,33 +99,50 @@ class BamlSyncClient:
user_prompt: str, user_prompt: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DynamicModel: ) -> types.DynamicModel:
result = self.__options.merge_options(baml_options).call_function_sync( # Check if on_tick is provided
function_name="AcreateStructuredOutput", if "on_tick" in baml_options:
args={ stream = self.stream.AcreateStructuredOutput(
"content": content, content=content,
"system_prompt": system_prompt, system_prompt=system_prompt,
"user_prompt": user_prompt, user_prompt=user_prompt,
}, baml_options=baml_options,
) )
return typing.cast( return stream.get_final_response()
types.DynamicModel, result.cast_to(types, types, stream_types, False, __runtime__) else:
) # Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="AcreateStructuredOutput",
args={
"content": content,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
},
)
return typing.cast(
types.DynamicModel, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractCategories( def ExtractCategories(
self, self,
content: str, content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DefaultContentPrediction: ) -> types.DefaultContentPrediction:
result = self.__options.merge_options(baml_options).call_function_sync( # Check if on_tick is provided
function_name="ExtractCategories", if "on_tick" in baml_options:
args={ stream = self.stream.ExtractCategories(content=content, baml_options=baml_options)
"content": content, return stream.get_final_response()
}, else:
) # Original non-streaming code
return typing.cast( result = self.__options.merge_options(baml_options).call_function_sync(
types.DefaultContentPrediction, function_name="ExtractCategories",
result.cast_to(types, types, stream_types, False, __runtime__), args={
) "content": content,
},
)
return typing.cast(
types.DefaultContentPrediction,
result.cast_to(types, types, stream_types, False, __runtime__),
)
def ExtractContentGraphGeneric( def ExtractContentGraphGeneric(
self, self,
@ -139,17 +159,28 @@ class BamlSyncClient:
custom_prompt_content: typing.Optional[str] = None, custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.KnowledgeGraph: ) -> types.KnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync( # Check if on_tick is provided
function_name="ExtractContentGraphGeneric", if "on_tick" in baml_options:
args={ stream = self.stream.ExtractContentGraphGeneric(
"content": content, content=content,
"mode": mode, mode=mode,
"custom_prompt_content": custom_prompt_content, custom_prompt_content=custom_prompt_content,
}, baml_options=baml_options,
) )
return typing.cast( return stream.get_final_response()
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__) else:
) # Original non-streaming code
result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractContentGraphGeneric",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
)
def ExtractDynamicContentGraph( def ExtractDynamicContentGraph(
self, self,
@ -166,48 +197,72 @@ class BamlSyncClient:
custom_prompt_content: typing.Optional[str] = None, custom_prompt_content: typing.Optional[str] = None,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.DynamicKnowledgeGraph: ) -> types.DynamicKnowledgeGraph:
result = self.__options.merge_options(baml_options).call_function_sync( # Check if on_tick is provided
function_name="ExtractDynamicContentGraph", if "on_tick" in baml_options:
args={ stream = self.stream.ExtractDynamicContentGraph(
"content": content, content=content,
"mode": mode, mode=mode,
"custom_prompt_content": custom_prompt_content, custom_prompt_content=custom_prompt_content,
}, baml_options=baml_options,
) )
return typing.cast( return stream.get_final_response()
types.DynamicKnowledgeGraph, else:
result.cast_to(types, types, stream_types, False, __runtime__), # Original non-streaming code
) result = self.__options.merge_options(baml_options).call_function_sync(
function_name="ExtractDynamicContentGraph",
args={
"content": content,
"mode": mode,
"custom_prompt_content": custom_prompt_content,
},
)
return typing.cast(
types.DynamicKnowledgeGraph,
result.cast_to(types, types, stream_types, False, __runtime__),
)
def SummarizeCode( def SummarizeCode(
self, self,
content: str, content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedCode: ) -> types.SummarizedCode:
result = self.__options.merge_options(baml_options).call_function_sync( # Check if on_tick is provided
function_name="SummarizeCode", if "on_tick" in baml_options:
args={ stream = self.stream.SummarizeCode(content=content, baml_options=baml_options)
"content": content, return stream.get_final_response()
}, else:
) # Original non-streaming code
return typing.cast( result = self.__options.merge_options(baml_options).call_function_sync(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__) function_name="SummarizeCode",
) args={
"content": content,
},
)
return typing.cast(
types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
)
def SummarizeContent( def SummarizeContent(
self, self,
content: str, content: str,
baml_options: BamlCallOptions = {}, baml_options: BamlCallOptions = {},
) -> types.SummarizedContent: ) -> types.SummarizedContent:
result = self.__options.merge_options(baml_options).call_function_sync( # Check if on_tick is provided
function_name="SummarizeContent", if "on_tick" in baml_options:
args={ stream = self.stream.SummarizeContent(content=content, baml_options=baml_options)
"content": content, return stream.get_final_response()
}, else:
) # Original non-streaming code
return typing.cast( result = self.__options.merge_options(baml_options).call_function_sync(
types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__) function_name="SummarizeContent",
) args={
"content": content,
},
)
return typing.cast(
types.SummarizedContent,
result.cast_to(types, types, stream_types, False, __runtime__),
)
class BamlStreamClient: class BamlStreamClient:

27
cognee/infrastructure/llm/structured_output_framework/baml/baml_client/type_builder.py
View file

@ -13,6 +13,9 @@
import typing import typing
from baml_py import type_builder from baml_py import type_builder
from baml_py import baml_py from baml_py import baml_py
# These are exports, not used here, hence the linter is disabled
from baml_py.baml_py import FieldType, EnumValueBuilder, EnumBuilder, ClassBuilder # noqa: F401 # pylint: disable=unused-import
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME
@ -296,7 +299,13 @@ class DynamicKnowledgeGraphBuilder(DynamicKnowledgeGraphAst):
return self._bldr.property(name).type(type) return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]: def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return [(name, self._bldr.property(name)) for name in self._properties] return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class DynamicKnowledgeGraphProperties: class DynamicKnowledgeGraphProperties:
@ -339,7 +348,13 @@ class DynamicModelBuilder(DynamicModelAst):
return self._bldr.property(name).type(type) return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]: def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return [(name, self._bldr.property(name)) for name in self._properties] return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class DynamicModelProperties: class DynamicModelProperties:
@ -619,7 +634,13 @@ class NodeBuilder(NodeAst):
return self._bldr.property(name).type(type) return self._bldr.property(name).type(type)
def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]: def list_properties(self) -> typing.List[typing.Tuple[str, baml_py.ClassPropertyBuilder]]:
return [(name, self._bldr.property(name)) for name in self._properties] return self._bldr.list_properties()
def remove_property(self, name: str) -> None:
self._bldr.remove_property(name)
def reset(self) -> None:
self._bldr.reset()
class NodeProperties: class NodeProperties: