baml-py 0.89.0__cp38-abi3-manylinux_2_24_aarch64.whl → 0.90.1__cp38-abi3-manylinux_2_24_aarch64.whl

baml_py/baml_py.pyi

@@ -113,6 +113,7 @@ class BamlRuntime:
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
         collectors: List[Collector],
+        env_vars: Dict[str, str],
     ) -> FunctionResult: ...
     def call_function_sync(
         self,
@@ -122,6 +123,7 @@ class BamlRuntime:
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
         collectors: List[Collector],
+        env_vars: Dict[str, str],
     ) -> FunctionResult: ...
     @staticmethod
     def from_files(
@@ -139,6 +141,7 @@ class BamlRuntime:
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
         collectors: List[Collector],
+        env_vars: Dict[str, str],
     ) -> FunctionResultStream: ...
     def stream_function_sync(
         self,
@@ -149,6 +152,7 @@ class BamlRuntime:
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
         collectors: List[Collector],
+        env_vars: Dict[str, str],
     ) -> SyncFunctionResultStream: ...
     def create_context_manager(self) -> RuntimeContextManager: ...
     def flush(self) -> None: ...
@@ -163,6 +167,7 @@ class BamlRuntime:
         ctx: RuntimeContextManager,
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
+        env_vars: Dict[str, str],
         is_stream: bool,
     ) -> HTTPRequest: ...
     def build_request_sync(
@@ -172,6 +177,7 @@ class BamlRuntime:
         ctx: RuntimeContextManager,
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
+        env_vars: Dict[str, str],
         is_stream: bool,
     ) -> HTTPRequest: ...
     def parse_llm_response(
@@ -185,6 +191,7 @@ class BamlRuntime:
         ctx: RuntimeContextManager,
         tb: Optional[TypeBuilder],
         cr: Optional[ClientRegistry],
+        env_vars: Dict[str, str],
     ) -> Any: ...
 
 
@@ -266,7 +273,7 @@ class Collector:
     def id(self, function_log_id: str) -> Optional[FunctionLog]: ...
     # For debugging
     @staticmethod
-    def __function_span_count() -> int: ...
+    def __function_call_count() -> int: ...
     @staticmethod
     def __print_storage() -> None: ...
 

baml_py/ctx_manager.py

@@ -5,10 +5,12 @@ import asyncio
 import contextvars
 import functools
 import inspect
+import os
 import typing
 from .baml_py import BamlLogEvent, RuntimeContextManager, BamlRuntime, BamlSpan
 import atexit
 import threading
+from typing import Dict
 
 F = typing.TypeVar("F", bound=typing.Callable[..., typing.Any])
 
@@ -25,11 +27,28 @@ def current_thread_id() -> int:
     return current_thread.ident or 0
 
 
+prev_ctx_manager: typing.Optional["CtxManager"] = None
+
+
 class CtxManager:
+    def __new__(cls, *args, **kwargs):
+        if prev_ctx_manager is not None:
+            return prev_ctx_manager
+        return super().__new__(cls)
+
     def __init__(self, rt: BamlRuntime):
+        global prev_ctx_manager
+        if prev_ctx_manager is not None:
+            self.rt = prev_ctx_manager.rt
+            self.ctx = prev_ctx_manager.ctx
+            return
+
+        prev_ctx_manager = self
+
         self.rt = rt
+
         self.ctx = contextvars.ContextVar[typing.Dict[int, RuntimeContextManager]](
-            "baml_ctx", default={current_thread_id(): rt.create_context_manager()}
+            "baml_ctx", default={}
         )
         atexit.register(self.rt.flush)
 
@@ -70,21 +89,36 @@ class CtxManager:
         return self.__ctx()
 
     def start_trace_sync(
-        self, name: str, args: typing.Dict[str, typing.Any]
+        self,
+        name: str,
+        args: typing.Dict[str, typing.Any],
+        env_vars: typing.Dict[str, str],
     ) -> BamlSpan:
+        # Clone the current context before creating the span
         mng = self.__ctx()
-        return BamlSpan.new(self.rt, name, args, mng)
+        return BamlSpan.new(self.rt, name, args, mng, env_vars)
 
     def start_trace_async(
-        self, name: str, args: typing.Dict[str, typing.Any]
+        self,
+        name: str,
+        args: typing.Dict[str, typing.Any],
+        env_vars: typing.Dict[str, str],
     ) -> BamlSpan:
         mng = self.__ctx()
         cln = mng.deep_clone()
         self.ctx.set({current_thread_id(): cln})
-        return BamlSpan.new(self.rt, name, args, cln)
+        return BamlSpan.new(self.rt, name, args, cln, env_vars)
 
-    def end_trace(self, span: BamlSpan, response: typing.Any) -> None:
-        span.finish(response, self.__ctx())
+    def clone_context(self) -> RuntimeContextManager:
+        mng = self.__ctx()
+        cln = mng.deep_clone()
+        self.ctx.set({current_thread_id(): cln})
+        return cln
+
+    def end_trace(
+        self, span: BamlSpan, response: typing.Any, env_vars: typing.Dict[str, str]
+    ) -> None:
+        span.finish(response, self.__ctx(), env_vars)
 
     def flush(self) -> None:
         self.rt.flush()
@@ -110,13 +144,13 @@ class CtxManager:
                     for i, arg in enumerate(args)
                 }
                 params.update(kwargs)
-                span = self.start_trace_async(func_name, params)
+                span = self.start_trace_async(func_name, params, os.environ.copy())
                 try:
                     response = await func(*args, **kwargs)
-                    self.end_trace(span, response)
+                    self.end_trace(span, response, os.environ.copy())
                     return response
                 except Exception as e:
-                    self.end_trace(span, e)
+                    self.end_trace(span, e, os.environ.copy())
                     raise e
 
             return typing.cast(F, async_wrapper)
@@ -130,14 +164,14 @@ class CtxManager:
                     for i, arg in enumerate(args)
                 }
                 params.update(kwargs)
-                span = self.start_trace_sync(func_name, params)
+                span = self.start_trace_sync(func_name, params, os.environ.copy())
                 try:
                     response = func(*args, **kwargs)
-                    self.end_trace(span, response)
+                    self.end_trace(span, response, os.environ.copy())
                     return response
                 except Exception as e:
                     print("Except but ending trace!")
-                    self.end_trace(span, e)
+                    self.end_trace(span, e, os.environ.copy())
                     raise e
 
             return typing.cast(F, wrapper)

{baml_py-0.89.0.dist-info → baml_py-0.90.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: baml-py
-Version: 0.89.0
+Version: 0.90.1
 License-File: LICENSE
 Summary: BAML python bindings (pyproject.toml)
 Author: Boundary <contact@boundaryml.com>

{baml_py-0.89.0.dist-info → baml_py-0.90.1.dist-info}/RECORD

@@ -1,11 +1,11 @@
-baml_py-0.89.0.dist-info/METADATA,sha256=YcVnDmFDWI2Yqiew-dgvy3fvqlvIYYtY61768125k9k,334
-baml_py-0.89.0.dist-info/WHEEL,sha256=K117A0tDd2F8c7mIWaUscf9O2xe1i_Jp1R1DM_vsioU,107
-baml_py-0.89.0.dist-info/entry_points.txt,sha256=9Uu_VcUjoI2qQMjVb0PRjEgI6pQ55WqBbzNparAPJyA,54
-baml_py-0.89.0.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+baml_py-0.90.1.dist-info/METADATA,sha256=dfqOnYNMNtHAl8j8oo_CqkZI1KaPu5fS0WkoERvr3YI,334
+baml_py-0.90.1.dist-info/WHEEL,sha256=jxb68-CzmMn8UbkjRfFQEt7fCQGUlaLj4slHZQcBCq8,107
+baml_py-0.90.1.dist-info/entry_points.txt,sha256=9Uu_VcUjoI2qQMjVb0PRjEgI6pQ55WqBbzNparAPJyA,54
+baml_py-0.90.1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
 baml_py/__init__.py,sha256=QWyuR-eLmqkiRW5v9RKF83knRj_HG8fyTlp6Sse9D3o,827
-baml_py/baml_py.abi3.so,sha256=TG3dM1bUQ9uhA80QFnvLTkb5pxR7N4WD2vRf45OPflk,44847392
-baml_py/baml_py.pyi,sha256=MKg8t9dh1cpc9OPV40LTGpT04-iq9-XOYZPGQJxDmOE,13620
-baml_py/ctx_manager.py,sha256=nzOJ-EfpSuahQ7hcrEbkDiNPBZo5q8oprVJzM9eZpX8,4702
+baml_py/baml_py.abi3.so,sha256=jIw7yE-oy-4nGW5CmSurfwYcbX_dT7H9MWH3kkn0yxI,45668400
+baml_py/baml_py.pyi,sha256=YYyXU1tLGLhz2OMRm9lroIAaKmQ0RZfATQ5Dea8q73s,13858
+baml_py/ctx_manager.py,sha256=p6Y4X1qfUExLeBeJKcVuf4WS9NrSabjRcRFrPIasTxY,5660
 baml_py/errors.py,sha256=wqv7xT_-pVXQNxD5JbOrrr_CABCFuNrLrEhmEX8RVJ8,389
 baml_py/internal_monkeypatch.py,sha256=JDwBPw4S8veD3nvJ13lFw8P5p29UOmDvvkgOy8eKA58,2106
 baml_py/logging.py,sha256=zM-yKPJ3LF4qpIptYQVr5zw_Gjimy3deWlTt4dOzVp0,190
@@ -13,4 +13,4 @@ baml_py/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 baml_py/safe_import.py,sha256=turgUpn9B4G273ZuDVjfZ_CkA2qmFQyiP-ZCPhtJO4M,2888
 baml_py/stream.py,sha256=RoHvdlYi1lap7DN0sCUA-H5HtAfxxePnm1nIe6BRTTs,6892
 baml_py/type_builder.py,sha256=HIAlses70C5DWNWgx3ZwsLeGt5-tviWXCXZiyyWedSg,6374
-baml_py-0.89.0.dist-info/RECORD,,
+baml_py-0.90.1.dist-info/RECORD,,

{baml_py-0.89.0.dist-info → baml_py-0.90.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.8.6)
+Generator: maturin (1.8.7)
 Root-Is-Purelib: false
 Tag: cp38-abi3-manylinux_2_24_aarch64