Reduce namespace pollution in codegen (#2016)

The more relative imports we add to our codegen, the more we pollute the
namespace. This PR attempts to use all the internal baml codegen under a
single namespace `_baml`.

Related issues:
- #2003

Author: antoniosarosi

engine/language_client_codegen/src/python/mod.rs

@@ -17,6 +17,10 @@ use internal_baml_core::{
 use self::python_language_features::{PythonLanguageFeatures, ToPython};
 use crate::{dir_writer::FileCollector, field_type_attributes};
 
+#[derive(askama::Template)]
+#[template(path = "_baml.py.j2", escape = "none")]
+struct BamlNamespace {}
+
 #[derive(askama::Template)]
 #[template(path = "config.py.j2", escape = "none")]
 struct PythonConfig {}
@@ -127,6 +131,7 @@ pub(crate) fn generate(
         .add_template::<generate_types::PythonStreamTypes>("partial_types.py", (ir, generator))?;
     collector.add_template::<generate_types::PythonTypes>("types.py", (ir, generator))?;
     collector.add_template::<generate_types::TypeBuilder>("type_builder.py", (ir, generator))?;
+    collector.add_template::<BamlNamespace>("_baml.py", (ir, generator))?;
     collector.add_template::<AsyncPythonClient>("async_client.py", (ir, generator))?;
     collector.add_template::<SyncPythonClient>("sync_client.py", (ir, generator))?;
     collector.add_template::<PythonGlobals>("globals.py", (ir, generator))?;
@@ -149,6 +154,14 @@ impl TryFrom<(&'_ IntermediateRepr, &'_ crate::GeneratorArgs)> for PythonConfig
     }
 }
 
+impl TryFrom<(&'_ IntermediateRepr, &'_ crate::GeneratorArgs)> for BamlNamespace {
+    type Error = anyhow::Error;
+
+    fn try_from(_: (&'_ IntermediateRepr, &'_ crate::GeneratorArgs)) -> Result<Self> {
+        Ok(BamlNamespace {})
+    }
+}
+
 impl TryFrom<(&'_ IntermediateRepr, &'_ crate::GeneratorArgs)> for PythonTracing {
     type Error = anyhow::Error;
 
@@ -277,14 +290,14 @@ impl ToTypeReferenceInClientDefinition for FieldType {
                     .map(|e| e.item.attributes.get("dynamic_type").is_some())
                     .unwrap_or(false)
                 {
-                    format!("Union[types.{name}, str]")
+                    format!("Union[_baml.types.{name}, str]")
                 } else {
-                    format!("types.{name}")
+                    format!("_baml.types.{name}")
                 }
             }
             FieldType::Literal(value) => to_python_literal(value),
-            FieldType::RecursiveTypeAlias(name) => format!("types.{name}"),
-            FieldType::Class(name) => format!("types.{name}"),
+            FieldType::RecursiveTypeAlias(name) => format!("_baml.types.{name}"),
+            FieldType::Class(name) => format!("_baml.types.{name}"),
             FieldType::List(inner) => format!("List[{}]", inner.to_type_ref(ir)),
             FieldType::Map(key, value) => {
                 format!("Dict[{}, {}]", key.to_type_ref(ir), value.to_type_ref(ir))
@@ -330,9 +343,9 @@ impl ToTypeReferenceInClientDefinition for FieldType {
         let with_state = metadata.1.state;
         let constraints = metadata.0;
         let module_prefix = if is_partial_type {
-            "partial_types."
+            "_baml.partial_types."
         } else {
-            "types."
+            "_baml.types."
         };
 
         let base_rep = match &base_type {
@@ -347,19 +360,19 @@ impl ToTypeReferenceInClientDefinition for FieldType {
                     // wrap primitives in `Optional` when generating partial types,
                     // although we should probably only do this when `!needed`.
                     if false {
-                        format!("Union[types.{name}, str]")
+                        format!("Union[_baml.types.{name}, str]")
                     } else {
-                        format!("Optional[Union[types.{name}, str]]")
+                        format!("Optional[Union[_baml.types.{name}, str]]")
                     }
                 } else {
                     // Note: The `false` here preserves potentially bugged codegen
                     // from before this commit. As the `false` implies, we always
                     // wrap primitives in `Optional` when generating partial types,
                     // although we should probably only do this when `!needed`.
                     if false {
-                        format!("types.{name}")
+                        format!("_baml.types.{name}")
                     } else {
-                        format!("Optional[types.{name}]")
+                        format!("Optional[_baml.types.{name}]")
                     }
                 }
             }
@@ -372,9 +385,9 @@ impl ToTypeReferenceInClientDefinition for FieldType {
             }
             FieldType::RecursiveTypeAlias(name) => {
                 if needed {
-                    format!("types.{name}")
+                    format!("_baml.types.{name}")
                 } else {
-                    format!("Optional[types.{name}]")
+                    format!("Optional[_baml.types.{name}]")
                 }
             }
             FieldType::Literal(value) => {

engine/language_client_codegen/src/python/templates/_baml.py.j2

@@ -0,0 +1,44 @@
+
+"""Private internal Baml namespace and utilities."""
+
+import os
+from typing import TypedDict, Dict, Optional, Union, List
+from typing_extensions import NotRequired
+
+import baml_py
+
+from . import types
+from . import partial_types
+from . import type_builder
+
+
+class BamlCallOptions(TypedDict, total=False):
+  """Additional parameters for Baml function calls."""
+  tb: NotRequired[type_builder.TypeBuilder]
+  client_registry: NotRequired[baml_py.baml_py.ClientRegistry]
+  collector: NotRequired[Union[baml_py.baml_py.Collector, List[baml_py.baml_py.Collector]]]
+  env: NotRequired[Dict[str, Optional[str]]]
+
+
+class BamlCallOptionsModApi(TypedDict, total=False):
+  """Additional parameters for modular API calls (doesn't take a collector)."""
+  tb: NotRequired[type_builder.TypeBuilder]
+  client_registry: NotRequired[baml_py.baml_py.ClientRegistry]
+  env: NotRequired[Dict[str, Optional[str]]]
+
+
+def env_vars_to_dict(overrides: Dict[str, Optional[str]]) -> Dict[str, str]:
+  base = os.environ.copy()
+  for k, v in overrides.items():
+    if v is not None:
+      base[k] = v
+    else:
+      base.pop(k, None)
+  return base
+
+
+__all__ = [
+  "types",
+  "partial_types",
+  "type_builder",
+]

engine/language_client_codegen/src/python/templates/async_client.py.j2

@@ -1,48 +1,29 @@
-from typing import Any, Dict, List, Optional, TypeVar, Union, TypedDict, Type, cast
-from typing_extensions import NotRequired, Literal
-import pprint
-import os
+from typing import Dict, List, Optional, TypeVar, Union, cast
+from typing_extensions import Literal
+
 import baml_py
-from pydantic import BaseModel, ValidationError, create_model
 
-from . import partial_types, types
+from . import _baml
 from .types import Checked, Check
-from .type_builder import TypeBuilder
 from .parser import LlmResponseParser, LlmStreamParser
 from .async_request import AsyncHttpRequest, AsyncHttpStreamRequest
 from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX, DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME
 
-OutputType = TypeVar('OutputType')
-
 
-# Define the TypedDict with optional parameters having default values
-class BamlCallOptions(TypedDict, total=False):
-    tb: NotRequired[TypeBuilder]
-    client_registry: NotRequired[baml_py.baml_py.ClientRegistry]
-    collector: NotRequired[Union[baml_py.baml_py.Collector, List[baml_py.baml_py.Collector]]]
-    env: NotRequired[Dict[str, Optional[str]]]
-
-def env_vars_to_dict(overrides: Dict[str, Optional[str]]) -> Dict[str, str]:
-    base = os.environ.copy()
-    for k, v in overrides.items():
-        if v is not None:
-            base[k] = v
-        else:
-            base.pop(k, None)
-    return base
+OutputType = TypeVar('OutputType')
 
 
 class BamlAsyncClient:
     __runtime: baml_py.BamlRuntime
     __ctx_manager: baml_py.BamlCtxManager
     __stream_client: "BamlStreamClient"
-    __http_request: "AsyncHttpRequest"
-    __http_stream_request: "AsyncHttpStreamRequest"
+    __http_request: AsyncHttpRequest
+    __http_stream_request: AsyncHttpStreamRequest
     __llm_response_parser: LlmResponseParser
     __llm_stream_parser: LlmStreamParser
-    __baml_options: BamlCallOptions
+    __baml_options: _baml.BamlCallOptions
 
-    def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[BamlCallOptions] = None):
+    def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[_baml.BamlCallOptions] = None):
       self.__runtime = runtime
       self.__ctx_manager = ctx_manager
       self.__stream_client = BamlStreamClient(self.__runtime, self.__ctx_manager, baml_options)
@@ -54,7 +35,7 @@ class BamlAsyncClient:
 
     def with_options(
       self,
-      tb: Optional[TypeBuilder] = None,
+      tb: Optional[_baml.type_builder.TypeBuilder] = None,
       client_registry: Optional[baml_py.baml_py.ClientRegistry] = None,
       collector: Optional[Union[baml_py.baml_py.Collector, List[baml_py.baml_py.Collector]]] = None,
       env: Optional[Dict[str, Optional[str]]] = None,
@@ -103,9 +84,9 @@ class BamlAsyncClient:
         {% for (name, type, default_value) in fn.args -%}
         {{name}}: {{type}}{% if let Some(d) = default_value %} = {{d}}{% endif %},
         {%- endfor %}
-        baml_options: BamlCallOptions = {},
+        baml_options: _baml.BamlCallOptions = {},
     ) -> {{fn.return_type}}:
-      options: BamlCallOptions = {**self.__baml_options, **(baml_options or {})}
+      options: _baml.BamlCallOptions = {**self.__baml_options, **(baml_options or {})}
 
       __tb__ = options.get("tb", None)
       if __tb__ is not None:
@@ -115,7 +96,7 @@ class BamlAsyncClient:
       __cr__ = options.get("client_registry", None)
       collector = options.get("collector", None)
       collectors = collector if isinstance(collector, list) else [collector] if collector is not None else []
-      env = env_vars_to_dict(options.get("env", {}))
+      env = _baml.env_vars_to_dict(options.get("env", {}))
       raw = await self.__runtime.call_function(
         "{{fn.name}}",
         {
@@ -129,15 +110,15 @@ class BamlAsyncClient:
         collectors,
         env,
       )
-      return cast({{fn.return_type}}, raw.cast_to(types, types, partial_types, False))
+      return cast({{fn.return_type}}, raw.cast_to(_baml.types, _baml.types, _baml.partial_types, False))
     {% endfor %}
 
 
 class BamlStreamClient:
     __runtime: baml_py.BamlRuntime
     __ctx_manager: baml_py.BamlCtxManager
-    __baml_options: BamlCallOptions
-    def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[BamlCallOptions] = None):
+    __baml_options: _baml.BamlCallOptions
+    def __init__(self, runtime: baml_py.BamlRuntime, ctx_manager: baml_py.BamlCtxManager, baml_options: Optional[_baml.BamlCallOptions] = None):
       self.__runtime = runtime
       self.__ctx_manager = ctx_manager
       self.__baml_options = baml_options or {}
@@ -148,9 +129,9 @@ class BamlStreamClient:
         {% for (name, type, default_value) in fn.args -%}
         {{name}}: {{type}}{% if let Some(d) = default_value %} = {{d}}{% endif %},
         {%- endfor %}
-        baml_options: BamlCallOptions = {},
+        baml_options: _baml.BamlCallOptions = {},
     ) -> baml_py.BamlStream[{{ fn.partial_return_type }}, {{ fn.return_type }}]:
-      options: BamlCallOptions = {**self.__baml_options, **(baml_options or {})}
+      options: _baml.BamlCallOptions = {**self.__baml_options, **(baml_options or {})}
       __tb__ = options.get("tb", None)
       if __tb__ is not None:
         tb = __tb__._tb # type: ignore (we know how to use this private attribute)
@@ -159,7 +140,7 @@ class BamlStreamClient:
       __cr__ = options.get("client_registry", None)
       collector = options.get("collector", None)
       collectors = collector if isinstance(collector, list) else [collector] if collector is not None else []
-      env = env_vars_to_dict(options.get("env", {}))
+      env = _baml.env_vars_to_dict(options.get("env", {}))
       raw = self.__runtime.stream_function(
         "{{fn.name}}",
         {
@@ -177,8 +158,8 @@ class BamlStreamClient:
 
       return baml_py.BamlStream[{{ fn.partial_return_type }}, {{ fn.return_type }}](
         raw,
-        lambda x: cast({{fn.partial_return_type}}, x.cast_to(types, types, partial_types, True)),
-        lambda x: cast({{fn.return_type}}, x.cast_to(types, types, partial_types, False)),
+        lambda x: cast({{fn.partial_return_type}}, x.cast_to(_baml.types, _baml.types, _baml.partial_types, True)),
+        lambda x: cast({{fn.return_type}}, x.cast_to(_baml.types, _baml.types, _baml.partial_types, False)),
         self.__ctx_manager.get(),
       )
     {% endfor %}

engine/language_client_codegen/src/python/templates/async_request.py.j2

@@ -1,28 +1,11 @@
-from typing import Any, Dict, List, Optional, Union, TypedDict, Type
-from typing_extensions import NotRequired, Literal
-import os
+from typing import Dict, List, Optional, Union
+from typing_extensions import Literal
 
 import baml_py
 
-from . import types
-from .types import Checked, Check
-from .type_builder import TypeBuilder
+from . import _baml
 
 
-class BamlCallOptions(TypedDict, total=False):
-    tb: NotRequired[TypeBuilder]
-    client_registry: NotRequired[baml_py.baml_py.ClientRegistry]
-    env: NotRequired[Dict[str, Optional[str]]]
-
-def env_vars_to_dict(overrides: Dict[str, Optional[str]]) -> Dict[str, str]:
-    base = os.environ.copy()
-    for k, v in overrides.items():
-        if v is not None:
-            base[k] = v
-        else:
-            base.pop(k, None)
-    return base
-
 class AsyncHttpRequest:
     __runtime: baml_py.BamlRuntime
     __ctx_manager: baml_py.BamlCtxManager
@@ -37,15 +20,15 @@ class AsyncHttpRequest:
         {% for (name, type, default_value) in fn.args -%}
         {{name}}: {{type}}{% if let Some(d) = default_value %} = {{d}}{% endif %},
         {%- endfor %}
-        baml_options: BamlCallOptions = {},
+        baml_options: _baml.BamlCallOptionsModApi = {},
     ) -> baml_py.HTTPRequest:
       __tb__ = baml_options.get("tb", None)
       if __tb__ is not None:
         tb = __tb__._tb # type: ignore (we know how to use this private attribute)
       else:
         tb = None
       __cr__ = baml_options.get("client_registry", None)
-      env = env_vars_to_dict(baml_options.get("env", {}))
+      env = _baml.env_vars_to_dict(baml_options.get("env", {}))
 
       return await self.__runtime.build_request(
         "{{fn.name}}",
@@ -77,15 +60,15 @@ class AsyncHttpStreamRequest:
         {% for (name, type, default_value) in fn.args -%}
         {{name}}: {{type}}{% if let Some(d) = default_value %} = {{d}}{% endif %},
         {%- endfor %}
-        baml_options: BamlCallOptions = {},
+        baml_options: _baml.BamlCallOptionsModApi = {},
     ) -> baml_py.HTTPRequest:
       __tb__ = baml_options.get("tb", None)
       if __tb__ is not None:
         tb = __tb__._tb # type: ignore (we know how to use this private attribute)
       else:
         tb = None
       __cr__ = baml_options.get("client_registry", None)
-      env = env_vars_to_dict(baml_options.get("env", {}))
+      env = _baml.env_vars_to_dict(baml_options.get("env", {}))
 
       return await self.__runtime.build_request(
         "{{fn.name}}",