feat: metadata context manager

Author: brian14708

examples/agent.py

@@ -111,7 +111,7 @@ async def agent_fn(
                         for tool_call in result.choices[0].message.tool_calls:
                             await output.send(("call", tool_call.model_dump_json()))
                             tasks.append(
-                                asyncio.create_task(ctx.run(call_tool, None, tool_call))
+                                asyncio.create_task(ctx.run(call_tool, tool_call))
                             )
                         for id_, tool_result in await asyncio.gather(*tasks):
                             await output.send(("tool", tool_result))
@@ -233,7 +233,6 @@ async def _completion(
 
     return await ctx.run(
         _completion,
-        None,
         messages,
     )
 

examples/hello_world.py

@@ -28,7 +28,7 @@ async def generate_lucky_number() -> int:
 @duron.fn
 async def greeting_flow(ctx: duron.Context, name: str) -> str:
     message, lucky_number = await asyncio.gather(
-        ctx.run(work, None, name), ctx.run(generate_lucky_number)
+        ctx.run(work, name), ctx.run(generate_lucky_number)
     )
     return f"{message} Your lucky number is {lucky_number}."
 

src/duron/__init__.py

@@ -1,6 +1,5 @@
 from ._core.config import set_config as set_config
 from ._core.context import Context as Context
-from ._core.options import RunOptions as RunOptions
 from ._core.signal import Signal as Signal
 from ._core.signal import SignalInterrupt as SignalInterrupt
 from ._core.signal import SignalWriter as SignalWriter

src/duron/_core/config.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -12,9 +13,13 @@
 @dataclass(slots=True)
 class _Config:
     codec: Codec
+    debug: bool
 
 
-config = _Config(codec=DefaultCodec())
+config = _Config(
+    codec=DefaultCodec(),
+    debug=os.getenv("DURON_DEBUG", "0").lower() in {"1", "true"},
+)
 
 
 def set_config(

src/duron/_core/context.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import binascii
+from contextlib import contextmanager
 from contextvars import ContextVar
 from random import Random
 from typing import TYPE_CHECKING, cast
@@ -21,11 +22,10 @@
 from duron.typing import inspect_function
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Coroutine
+    from collections.abc import Callable, Coroutine, Generator
     from contextvars import Token
     from types import TracebackType
 
-    from duron._core.options import RunOptions
     from duron._core.signal import Signal, SignalWriter
     from duron._core.stream import Stream, StreamWriter
     from duron._decorator.fn import Fn
@@ -38,6 +38,9 @@
     _P = ParamSpec("_P")
 
 _context: ContextVar[Context | None] = ContextVar("duron_context", default=None)
+_metadata: ContextVar[dict[str, JSONValue] | None] = ContextVar(
+    "duron_metadata", default=None
+)
 
 
 @final
@@ -76,7 +79,6 @@ def current() -> Context:
     async def run(
         self,
         fn: Callable[_P, Coroutine[Any, Any, _T]] | Op[_P, _T],
-        options: RunOptions | None = ...,
         /,
         *args: _P.args,
         **kwargs: _P.kwargs,
@@ -85,7 +87,6 @@ async def run(
     async def run(
         self,
         fn: Callable[_P, _T] | CheckpointOp[_P, _T, Any],
-        options: RunOptions | None = ...,
         /,
         *args: _P.args,
         **kwargs: _P.kwargs,
@@ -95,7 +96,6 @@ async def run(
         fn: Callable[_P, Coroutine[Any, Any, _T] | _T]
         | Op[_P, _T]
         | CheckpointOp[_P, _T, Any],
-        options: RunOptions | None = None,
         /,
         *args: _P.args,
         **kwargs: _P.kwargs,
@@ -105,7 +105,7 @@ async def run(
             raise RuntimeError(msg)
 
         if isinstance(fn, CheckpointOp):
-            async with self.run_stream(fn, options, *args, **kwargs) as stream:
+            async with self.run_stream(fn, *args, **kwargs) as stream:
                 await stream.discard()
                 return await stream
 
@@ -123,20 +123,18 @@ async def run(
                 args=args,
                 kwargs=kwargs,
                 return_type=return_type,
-                metadata=_merge(options.metadata if options else None, metadata),
+                metadata=self._get_metadata(metadata),
             ),
         )
         return cast("_T", await op)
 
     def run_stream(
         self,
         fn: CheckpointOp[_P, _T, _S],
-        options: RunOptions | None = None,
         /,
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> AsyncContextManager[Stream[_S, _T]]:
-        _ = options
         if asyncio.get_running_loop() is not self._loop:
             msg = "Context time can only be used in the context loop"
             raise RuntimeError(msg)
@@ -155,7 +153,6 @@ async def create_stream(
         dtype: TypeHint[_T],
         *,
         external: bool = False,
-        metadata: dict[str, JSONValue] | None = None,
     ) -> tuple[Stream[_T, None], StreamWriter[_T]]:
         if asyncio.get_running_loop() is not self._loop:
             msg = "Context time can only be used in the context loop"
@@ -164,32 +161,28 @@ async def create_stream(
             self._loop,
             dtype,
             external=external,
-            metadata=metadata,
+            metadata=self._get_metadata(None),
         )
 
     async def create_signal(
         self,
         dtype: TypeHint[_T],
-        *,
-        metadata: dict[str, JSONValue] | None = None,
     ) -> tuple[Signal[_T], SignalWriter[_T]]:
         if asyncio.get_running_loop() is not self._loop:
             msg = "Context time can only be used in the context loop"
             raise RuntimeError(msg)
-        return await create_signal(self._loop, dtype, metadata=metadata)
+        return await create_signal(self._loop, dtype, metadata=self._get_metadata(None))
 
     async def create_promise(
         self,
         dtype: type[_T],
-        *,
-        metadata: dict[str, JSONValue] | None = None,
     ) -> tuple[str, asyncio.Future[_T]]:
         if asyncio.get_running_loop() is not self._loop:
             msg = "Context time can only be used in the context loop"
             raise RuntimeError(msg)
         fut = create_op(
             self._loop,
-            ExternalPromiseCreate(metadata=metadata, return_type=dtype),
+            ExternalPromiseCreate(metadata=self._get_metadata(None), return_type=dtype),
         )
         return (
             binascii.b2a_base64(fut.id, newline=False).decode(),
@@ -220,12 +213,30 @@ def random(self) -> Random:
             raise RuntimeError(msg)
         return Random(self._loop.generate_op_id())  # noqa: S311
 
+    @contextmanager
+    def metadata(self, metadata: dict[str, JSONValue]) -> Generator[None, None, None]:
+        if asyncio.get_running_loop() is not self._loop:
+            msg = "Context time can only be used in the context loop"
+            raise RuntimeError(msg)
+        if not metadata:
+            yield
+            return
+
+        current = _metadata.get()
+        merged = {**current, **metadata} if current is not None else metadata
+        token = _metadata.set(merged)
+        try:
+            yield
+        finally:
+            _metadata.reset(token)
 
-def _merge(
-    d1: dict[str, JSONValue] | None, d2: dict[str, JSONValue] | None
-) -> dict[str, JSONValue] | None:
-    if d1 is None:
-        return d2
-    if d2 is None:
-        return d1
-    return {**d1, **d2}
+    @staticmethod
+    def _get_metadata(
+        merge: dict[str, JSONValue] | None,
+    ) -> dict[str, JSONValue] | None:
+        current = _metadata.get()
+        if merge is None:
+            return current
+        if current is None:
+            return merge
+        return {**current, **merge}