codekingpro/portable-devtools
114k
1"""Decorator for creating a run tree from functions."""2 3from __future__ import annotations4 5import asyncio6import contextlib7import contextvars8import datetime9import functools10import inspect11import logging12import threading13import warnings14import weakref15from collections.abc import (16 AsyncGenerator,17 AsyncIterator,18 Awaitable,19 Generator,20 Iterator,21 Mapping,22 Sequence,23)24from contextvars import copy_context25from typing import (26 TYPE_CHECKING,27 Annotated,28 Any,29 Callable,30 Generic,31 Literal,32 Optional,33 Protocol,34 TypedDict,35 TypeVar,36 Union,37 cast,38 get_type_hints,39 overload,40 runtime_checkable,41)42 43from typing_extensions import ParamSpec, TypeGuard, get_args, get_origin44 45import langsmith._internal._context as _context46from langsmith import client as ls_client47from langsmith import run_trees, schemas, utils48from langsmith._internal import _aiter as aitertools49from langsmith._runtime_overrides import (50 _aio_to_thread_override_active as _runtime_override_active,51)52from langsmith.env import _runtime_env53from langsmith.run_trees import WriteReplica54 55if TYPE_CHECKING:56 from types import TracebackType57 58 from langchain_core.runnables import Runnable59 60LOGGER = logging.getLogger(__name__)61_CONTEXT_KEYS: dict[str, contextvars.ContextVar] = {62 "parent_ref": _context._PARENT_RUN_TREE_REF,63 "project_name": _context._PROJECT_NAME,64 "tags": _context._TAGS,65 "metadata": _context._METADATA,66 "enabled": _context._TRACING_ENABLED,67 "client": _context._CLIENT,68 "replicas": run_trees._REPLICAS,69 "distributed_parent_id": run_trees._DISTRIBUTED_PARENT_ID,70}71 72_EXCLUDED_FRAME_FNAME = "langsmith/run_helpers.py"73 74_OTEL_AVAILABLE: Optional[bool] = None75 76 77def get_current_run_tree() -> Optional[run_trees.RunTree]:78 """Get the current run tree.79 80 Uses a weakref-based lookup to avoid memory leaks from captured contexts.81 The RunTree may return None if it has been garbage collected.82 """83 return _context.get_current_run_tree()84 85 86@contextlib.contextmanager87def set_tracing_parent(88 run_tree: run_trees.RunTree,89) -> Generator[None, None, None]:90 """Set a RunTree as the active tracing parent within this block.91 92 Unlike `tracing_context`, this only sets the parent run tree and nothing93 else, making it safe to use in isolated threads where you want precise94 control over which run acts as the parent without inheriting or overwriting95 other context variables.96 97 Args:98 run_tree: The RunTree to use as the active parent.99 """100 token = _context._PARENT_RUN_TREE_REF.set(weakref.ref(run_tree))101 try:102 yield103 finally:104 _context._PARENT_RUN_TREE_REF.reset(token)105 106 107def set_run_metadata(**metadata: Any) -> None:108 """Update metadata on the current run tree."""109 run_tree = get_current_run_tree()110 if run_tree is None:111 LOGGER.warning(112 "No active run tree found. Call `set_run_metadata` inside a traced run."113 )114 else:115 run_tree.metadata.update(metadata)116 return117 118 119def get_tracing_context(120 context: Optional[contextvars.Context] = None,121) -> dict[str, Any]:122 """Get the current tracing context."""123 if context is None:124 parent = _context.get_current_run_tree()125 return {126 "parent": parent,127 "project_name": _context._PROJECT_NAME.get(),128 "tags": _context._TAGS.get(),129 "metadata": _context._METADATA.get(),130 "enabled": _context._TRACING_ENABLED.get(),131 "client": _context._CLIENT.get(),132 "replicas": run_trees._REPLICAS.get(),133 "distributed_parent_id": run_trees._DISTRIBUTED_PARENT_ID.get(),134 }135 # When reading from a copied context, dereference the weakref136 result = {k: context.get(v) for k, v in _CONTEXT_KEYS.items()}137 parent_ref = result.pop("parent_ref", None)138 result["parent"] = parent_ref() if parent_ref is not None else None139 return result140 141 142@contextlib.contextmanager143def tracing_context(144 *,145 project_name: Optional[str] = None,146 tags: Optional[list[str]] = None,147 metadata: Optional[dict[str, Any]] = None,148 parent: Optional[Union[run_trees.RunTree, Mapping, str, Literal[False]]] = None,149 enabled: Optional[Union[bool, Literal["local"]]] = None,150 client: Optional[ls_client.Client] = None,151 replicas: Optional[Sequence[WriteReplica]] = None,152 distributed_parent_id: Optional[str] = None,153 **kwargs: Any,154) -> Generator[None, None, None]:155 """Set the tracing context for a block of code.156 157 Args:158 project_name: The name of the project to log the run to.159 tags: The tags to add to the run.160 metadata: The metadata to add to the run.161 parent: The parent run to use for the context.162 163 Can be a Run/`RunTree` object, request headers (for distributed tracing),164 or the dotted order string.165 client: The client to use for logging the run to LangSmith.166 enabled: Whether tracing is enabled.167 168 Defaults to `None`, meaning it will use the current context value or environment variables.169 replicas: A sequence of `WriteReplica` dictionaries to send runs to.170 171 Example: `[{"api_url": "https://api.example.com", "auth": {"api_key": "key"}, "project_name": "proj"}]`172 or `[{"project_name": "my_experiment", "updates": {"reference_example_id": None}}]`173 distributed_parent_id: The distributed parent ID for distributed tracing. Defaults to None.174 """175 if kwargs:176 # warn177 warnings.warn(178 f"Unrecognized keyword arguments: {kwargs}.",179 DeprecationWarning,180 )181 current_context = get_tracing_context()182 parent_run = (183 _get_parent_run({"parent": parent or kwargs.get("parent_run")})184 if parent is not False185 else None186 )187 distributed_parent_id_to_use = distributed_parent_id188 if distributed_parent_id_to_use is None and parent_run is not None:189 # TODO(angus): decide if we want to merge tags and metadata190 tags = sorted(set(tags or []) | set(parent_run.tags or []))191 metadata = {**parent_run.metadata, **(metadata or {})}192 distributed_parent_id_to_use = parent_run.id # type: ignore[assignment]193 enabled = enabled if enabled is not None else current_context.get("enabled")194 _set_tracing_context(195 {196 "parent": parent_run,197 "project_name": project_name,198 "tags": tags,199 "metadata": metadata,200 "enabled": enabled,201 "client": client,202 "replicas": replicas,203 "distributed_parent_id": distributed_parent_id_to_use,204 }205 )206 try:207 yield208 finally:209 _set_tracing_context(current_context)210 211 212# Alias for backwards compatibility213get_run_tree_context = get_current_run_tree214 215 216def is_traceable_function(func: Any) -> TypeGuard[SupportsLangsmithExtra[P, R]]:217 """Check if a function is `@traceable` decorated."""218 return (219 _is_traceable_function(func)220 or (isinstance(func, functools.partial) and _is_traceable_function(func.func))221 or (hasattr(func, "__call__") and _is_traceable_function(func.__call__))222 )223 224 225def ensure_traceable(226 func: Callable[P, R],227 *,228 name: Optional[str] = None,229 metadata: Optional[Mapping[str, Any]] = None,230 tags: Optional[list[str]] = None,231 client: Optional[ls_client.Client] = None,232 reduce_fn: Optional[Callable[[Sequence], Union[dict, str]]] = None,233 project_name: Optional[str] = None,234 process_inputs: Optional[Callable[[dict], dict]] = None,235 process_outputs: Optional[Callable[..., dict]] = None,236 process_chunk: Optional[Callable] = None,237) -> SupportsLangsmithExtra[P, R]:238 """Ensure that a function is traceable."""239 if is_traceable_function(func):240 return func241 return traceable(242 name=name,243 metadata=metadata,244 tags=tags,245 client=client,246 reduce_fn=reduce_fn,247 project_name=project_name,248 process_inputs=process_inputs,249 process_outputs=process_outputs,250 process_chunk=process_chunk,251 )(func)252 253 254def is_async(func: Callable) -> bool:255 """Inspect function or wrapped function to see if it is async."""256 return inspect.iscoroutinefunction(func) or (257 hasattr(func, "__wrapped__") and inspect.iscoroutinefunction(func.__wrapped__)258 )259 260 261class LangSmithExtra(TypedDict, total=False):262 """Any additional info to be injected into the run dynamically."""263 264 name: Optional[str]265 """Optional name for the run."""266 reference_example_id: Optional[ls_client.ID_TYPE]267 """Optional ID of a reference example."""268 run_extra: Optional[dict]269 """Optional additional run information."""270 parent: Optional[Union[run_trees.RunTree, str, Mapping]]271 """Optional parent run, can be a RunTree, string, or mapping."""272 run_tree: Optional[run_trees.RunTree] # TODO: Deprecate273 """Optional run tree (deprecated)."""274 project_name: Optional[str]275 """Optional name of the project."""276 metadata: Optional[dict[str, Any]]277 """Optional metadata for the run."""278 tags: Optional[list[str]]279 """Optional list of tags for the run."""280 run_id: Optional[ls_client.ID_TYPE]281 """Optional ID for the run."""282 client: Optional[ls_client.Client]283 """Optional LangSmith client."""284 # Optional callback function to be called if the run succeeds and before it is sent.285 _on_success: Optional[Callable[[run_trees.RunTree], None]]286 on_end: Optional[Callable[[run_trees.RunTree], Any]]287 """Optional callback function to be called after the run ends and is sent."""288 289 290R = TypeVar("R", covariant=True)291P = ParamSpec("P")292 293 294@runtime_checkable295class SupportsLangsmithExtra(Protocol, Generic[P, R]):296 """Implementations of this Protocol accept an optional langsmith_extra parameter."""297 298 def __call__( # type: ignore[valid-type]299 self,300 *args: P.args,301 langsmith_extra: Optional[LangSmithExtra] = None,302 **kwargs: P.kwargs,303 ) -> R:304 """Call the instance when it is called as a function.305 306 Args:307 *args: Variable length argument list.308 langsmith_extra: Optional dictionary containing additional309 parameters specific to Langsmith.310 **kwargs: Arbitrary keyword arguments.311 312 Returns:313 R: The return value of the method.314 315 """316 ...317 318 319def _extract_usage(320 *,321 run_tree: run_trees.RunTree,322 outputs: Optional[dict] = None,323 **kwargs: Any,324) -> Optional[schemas.ExtractedUsageMetadata]:325 from_metadata = (run_tree.metadata or {}).get("usage_metadata")326 return (outputs or {}).get("usage_metadata") or from_metadata327 328 329@overload330def traceable(331 func: Callable[P, R],332) -> SupportsLangsmithExtra[P, R]: ...333 334 335@overload336def traceable(337 run_type: ls_client.RUN_TYPE_T = "chain",338 *,339 name: Optional[str] = None,340 metadata: Optional[Mapping[str, Any]] = None,341 tags: Optional[list[str]] = None,342 client: Optional[ls_client.Client] = None,343 reduce_fn: Optional[Callable[[Sequence], Union[dict, str]]] = None,344 project_name: Optional[str] = None,345 process_inputs: Optional[Callable[[dict], dict]] = None,346 process_outputs: Optional[Callable[..., dict]] = None,347 process_chunk: Optional[Callable] = None,348 _invocation_params_fn: Optional[Callable[[dict], dict]] = None,349 dangerously_allow_filesystem: bool = False,350 enabled: Optional[bool] = None,351 exceptions_to_handle: Optional[tuple[type[BaseException], ...]] = None,352) -> Callable[[Callable[P, R]], SupportsLangsmithExtra[P, R]]: ...353 354 355def traceable(356 *args: Any,357 **kwargs: Any,358) -> Union[Callable, Callable[[Callable], Callable]]:359 """Trace a function with langsmith.360 361 Args:362 run_type: The type of run (span) to create.363 364 Examples: `llm`, `chain`, `tool`, `prompt`, `retriever`, etc.365 366 Defaults to "chain".367 name: The name of the run. Defaults to the function name.368 metadata: The metadata to add to the run. Defaults to `None`.369 tags: The tags to add to the run. Defaults to `None`.370 client: The client to use for logging the run to LangSmith. Defaults to371 `None`, which will use the default client.372 reduce_fn: A function to reduce the output of the function if the function373 returns a generator.374 375 Defaults to `None`, which means the values will be logged as a list.376 377 !!! note378 379 If the iterator is never exhausted (e.g. the function returns an380 infinite generator), this will never be called, and the run itself will381 be stuck in a pending state.382 project_name: The name of the project to log the run to.383 384 Defaults to `None`, which will use the default project.385 process_inputs: Custom serialization / processing function for inputs.386 387 Defaults to `None`.388 process_outputs: Custom serialization / processing function for outputs.389 390 Defaults to `None`.391 dangerously_allow_filesystem: Whether to allow filesystem access for attachments.392 393 Defaults to `False`.394 395 Traces that reference local filepaths will be uploaded to LangSmith.396 In general, network-hosted applications should not be using this because397 referenced files are usually on the user's machine, not the host machine.398 enabled: Whether tracing is enabled for this function.399 400 Defaults to `None`, which will use the default value from the current context.401 exceptions_to_handle: Exception types to ignore when logging errors.402 403 If an exception of one of these types is raised, the run will still be recorded404 but the error field will be `None` instead of containing the full traceback.405 406 Defaults to `None`.407 408 Returns:409 The decorated function.410 411 !!! note412 413 Requires that `LANGSMITH_TRACING_V2` be set to 'true' in the environment.414 415 Examples:416 !!! example "Basic usage"417 418 ```python419 @traceable420 def my_function(x: float, y: float) -> float:421 return x + y422 423 424 my_function(5, 6)425 426 427 @traceable428 async def my_async_function(query_params: dict) -> dict:429 async with httpx.AsyncClient() as http_client:430 response = await http_client.get(431 "https://api.example.com/data",432 params=query_params,433 )434 return response.json()435 436 437 asyncio.run(my_async_function({"param": "value"}))438 ```439 440 !!! example "Streaming data with a generator"441 442 ```python443 @traceable444 def my_generator(n: int) -> Iterable:445 for i in range(n):446 yield i447 448 449 for item in my_generator(5):450 print(item)451 ```452 453 !!! example "Async streaming data"454 455 ```python456 @traceable457 async def my_async_generator(query_params: dict) -> Iterable:458 async with httpx.AsyncClient() as http_client:459 response = await http_client.get(460 "https://api.example.com/data",461 params=query_params,462 )463 for item in response.json():464 yield item465 466 467 async def async_code():468 async for item in my_async_generator({"param": "value"}):469 print(item)470 471 472 asyncio.run(async_code())473 ```474 475 !!! example "Specifying a run type and name"476 477 ```python478 @traceable(name="CustomName", run_type="tool")479 def another_function(a: float, b: float) -> float:480 return a * b481 482 483 another_function(5, 6)484 ```485 486 !!! example "Logging with custom metadata and tags"487 488 ```python489 @traceable(490 metadata={"version": "1.0", "author": "John Doe"}, tags=["beta", "test"]491 )492 def tagged_function(x):493 return x**2494 495 496 tagged_function(5)497 ```498 499 !!! example "Specifying a custom client and project name"500 501 ```python502 custom_client = Client(api_key="your_api_key")503 504 505 @traceable(client=custom_client, project_name="My Special Project")506 def project_specific_function(data):507 return data508 509 510 project_specific_function({"data": "to process"})511 ```512 513 !!! example "Manually passing `langsmith_extra`"514 515 ```python516 @traceable517 def manual_extra_function(x):518 return x**2519 520 521 manual_extra_function(5, langsmith_extra={"metadata": {"version": "1.0"}})522 ```523 524 !!! example "Handling specific exceptions"525 526 ```python527 @traceable(exceptions_to_handle=(ValueError, TypeError))528 def function_with_handled_exceptions(x):529 if x < 0:530 raise ValueError("Negative value") # Won't send error in the trace531 return x**2532 533 534 function_with_handled_exceptions(-5)535 ```536 """537 run_type = cast(538 ls_client.RUN_TYPE_T,539 (540 args[0]541 if args and isinstance(args[0], str)542 else (kwargs.pop("run_type", None) or "chain")543 ),544 )545 if run_type not in _VALID_RUN_TYPES:546 warnings.warn(547 f"Unrecognized run_type: {run_type}. Must be one of: {_VALID_RUN_TYPES}."548 f" Did you mean @traceable(name='{run_type}')?"549 )550 if len(args) > 1:551 warnings.warn(552 "The `traceable()` decorator only accepts one positional argument, "553 "which should be the run_type. All other arguments should be passed "554 "as keyword arguments."555 )556 if "extra" in kwargs:557 warnings.warn(558 "The `extra` keyword argument is deprecated. Please use `metadata` "559 "instead.",560 DeprecationWarning,561 )562 reduce_fn = kwargs.pop("reduce_fn", None)563 enabled = kwargs.pop("enabled", None)564 container_input = _ContainerInput(565 # TODO: Deprecate raw extra566 extra_outer=kwargs.pop("extra", None),567 name=kwargs.pop("name", None),568 metadata=kwargs.pop("metadata", None),569 tags=kwargs.pop("tags", None),570 client=kwargs.pop("client", None),571 project_name=kwargs.pop("project_name", None),572 run_type=run_type,573 process_inputs=kwargs.pop("process_inputs", None),574 process_chunk=kwargs.pop("process_chunk", None),575 invocation_params_fn=kwargs.pop("_invocation_params_fn", None),576 dangerously_allow_filesystem=kwargs.pop("dangerously_allow_filesystem", False),577 enabled=enabled,578 exceptions_to_handle=kwargs.pop("exceptions_to_handle", None),579 )580 outputs_processor = kwargs.pop("process_outputs", None)581 _on_run_end = functools.partial(582 _handle_container_end,583 outputs_processor=outputs_processor,584 )585 586 if kwargs:587 warnings.warn(588 f"The following keyword arguments are not recognized and will be ignored: "589 f"{sorted(kwargs.keys())}.",590 DeprecationWarning,591 )592 593 def decorator(func: Callable):594 func_sig = inspect.signature(func)595 func_accepts_parent_run = func_sig.parameters.get("run_tree", None) is not None596 func_accepts_config = func_sig.parameters.get("config", None) is not None597 598 @functools.wraps(func)599 async def async_wrapper(600 *args: Any,601 langsmith_extra: Optional[LangSmithExtra] = None,602 **kwargs: Any,603 ) -> Any:604 """Async version of wrapper function."""605 if not func_accepts_config:606 kwargs.pop("config", None)607 run_container = await aitertools.aio_to_thread(608 copy_context(),609 _setup_run,610 func,611 container_input=container_input,612 langsmith_extra=langsmith_extra,613 args=args,614 kwargs=kwargs,615 )616 617 try:618 accepts_context = aitertools.asyncio_accepts_context()619 if func_accepts_parent_run:620 kwargs["run_tree"] = run_container["new_run"]621 622 otel_context_manager = _maybe_create_otel_context(623 run_container["new_run"]624 )625 use_ctx_task = accepts_context and not _runtime_override_active()626 if otel_context_manager:627 628 async def run_with_otel_context():629 with otel_context_manager:630 return await func(*args, **kwargs)631 632 if use_ctx_task:633 function_result = await asyncio.create_task( # type: ignore[call-arg]634 run_with_otel_context(), context=run_container["context"]635 )636 else:637 with tracing_context(638 **get_tracing_context(run_container["context"])639 ):640 function_result = await run_with_otel_context()641 else:642 fr_coro = func(*args, **kwargs)643 if use_ctx_task:644 function_result = await asyncio.create_task( # type: ignore[call-arg]645 fr_coro, context=run_container["context"]646 )647 else:648 with tracing_context(649 **get_tracing_context(run_container["context"])650 ):651 function_result = await fr_coro652 except BaseException as e:653 # shield from cancellation, given we're catching all exceptions654 _cleanup_traceback(e)655 await asyncio.shield(656 aitertools.aio_to_thread(657 copy_context(), _on_run_end, run_container, error=e658 )659 )660 raise661 await aitertools.aio_to_thread(662 copy_context(), _on_run_end, run_container, outputs=function_result663 )664 return function_result665 666 @functools.wraps(func)667 async def async_generator_wrapper(668 *args: Any, langsmith_extra: Optional[LangSmithExtra] = None, **kwargs: Any669 ) -> AsyncGenerator:670 if not func_accepts_config:671 kwargs.pop("config", None)672 run_container = await aitertools.aio_to_thread(673 copy_context(),674 _setup_run,675 func,676 container_input=container_input,677 langsmith_extra=langsmith_extra,678 args=args,679 kwargs=kwargs,680 )681 results: list[Any] = []682 try:683 if func_accepts_parent_run:684 kwargs["run_tree"] = run_container["new_run"]685 # TODO: Nesting is ambiguous if a nested traceable function is only686 # called mid-generation. Need to explicitly accept run_tree to get687 # around this.688 689 otel_context_manager = _maybe_create_otel_context(690 run_container["new_run"]691 )692 693 async_gen_result = func(*args, **kwargs)694 # Can't iterate through if it's a coroutine695 accepts_context = aitertools.asyncio_accepts_context()696 use_ctx_task = accepts_context and not _runtime_override_active()697 if inspect.iscoroutine(async_gen_result):698 if use_ctx_task:699 async_gen_result = await asyncio.create_task(700 async_gen_result, context=run_container["context"]701 ) # type: ignore702 else:703 with tracing_context(704 **get_tracing_context(run_container["context"])705 ):706 async_gen_result = await async_gen_result707 708 async for item in _process_async_iterator(709 generator=async_gen_result,710 run_container=run_container,711 is_llm_run=(712 run_container["new_run"].run_type == "llm"713 if run_container["new_run"]714 else False715 ),716 accepts_context=use_ctx_task,717 results=results,718 process_chunk=container_input.get("process_chunk"),719 otel_context_manager=otel_context_manager,720 ):721 yield item722 except BaseException as e:723 _cleanup_traceback(e)724 await asyncio.shield(725 aitertools.aio_to_thread(726 copy_context(),727 _on_run_end,728 run_container,729 error=e,730 outputs=_get_function_result(results, reduce_fn),731 )732 )733 raise734 await aitertools.aio_to_thread(735 copy_context(),736 _on_run_end,737 run_container,738 outputs=_get_function_result(results, reduce_fn),739 )740 741 @functools.wraps(func)742 def wrapper(743 *args: Any,744 langsmith_extra: Optional[LangSmithExtra] = None,745 **kwargs: Any,746 ) -> Any:747 """Create a new run or create_child() if run is passed in kwargs."""748 if not func_accepts_config:749 kwargs.pop("config", None)750 run_container = _setup_run(751 func,752 container_input=container_input,753 langsmith_extra=langsmith_extra,754 args=args,755 kwargs=kwargs,756 )757 func_accepts_parent_run = (758 inspect.signature(func).parameters.get("run_tree", None) is not None759 )760 try:761 if func_accepts_parent_run:762 kwargs["run_tree"] = run_container["new_run"]763 764 otel_context_manager = _maybe_create_otel_context(765 run_container["new_run"]766 )767 if otel_context_manager:768 769 def run_with_otel_context():770 with otel_context_manager:771 return func(*args, **kwargs)772 773 function_result = run_container["context"].run(774 run_with_otel_context775 )776 else:777 function_result = run_container["context"].run(778 func, *args, **kwargs779 )780 except BaseException as e:781 _cleanup_traceback(e)782 _on_run_end(run_container, error=e)783 raise784 _on_run_end(run_container, outputs=function_result)785 return function_result786 787 @functools.wraps(func)788 def generator_wrapper(789 *args: Any, langsmith_extra: Optional[LangSmithExtra] = None, **kwargs: Any790 ) -> Any:791 if not func_accepts_config:792 kwargs.pop("config", None)793 run_container = _setup_run(794 func,795 container_input=container_input,796 langsmith_extra=langsmith_extra,797 args=args,798 kwargs=kwargs,799 )800 func_accepts_parent_run = (801 inspect.signature(func).parameters.get("run_tree", None) is not None802 )803 results: list[Any] = []804 function_return: Any = None805 806 try:807 if func_accepts_parent_run:808 kwargs["run_tree"] = run_container["new_run"]809 810 generator_result = run_container["context"].run(func, *args, **kwargs)811 812 otel_context_manager = _maybe_create_otel_context(813 run_container["new_run"]814 )815 816 function_return = yield from _process_iterator(817 generator_result,818 run_container,819 is_llm_run=run_type == "llm",820 results=results,821 process_chunk=container_input.get("process_chunk"),822 otel_context_manager=otel_context_manager,823 )824 825 if function_return is not None:826 results.append(function_return)827 828 except BaseException as e:829 _cleanup_traceback(e)830 _on_run_end(831 run_container,832 error=e,833 outputs=_get_function_result(results, reduce_fn),834 )835 raise836 _on_run_end(run_container, outputs=_get_function_result(results, reduce_fn))837 838 return function_return839 840 # "Stream" functions (used in methods like OpenAI/Anthropic's SDKs)841 # are functions that return iterable responses and should not be842 # considered complete until the streaming is completed843 @functools.wraps(func)844 def stream_wrapper(845 *args: Any, langsmith_extra: Optional[LangSmithExtra] = None, **kwargs: Any846 ) -> Any:847 if not func_accepts_config:848 kwargs.pop("config", None)849 trace_container = _setup_run(850 func,851 container_input=container_input,852 langsmith_extra=langsmith_extra,853 args=args,854 kwargs=kwargs,855 )856 857 try:858 if func_accepts_parent_run:859 kwargs["run_tree"] = trace_container["new_run"]860 stream = trace_container["context"].run(func, *args, **kwargs)861 except Exception as e:862 _cleanup_traceback(e)863 _on_run_end(trace_container, error=e)864 raise865 866 if hasattr(stream, "__iter__"):867 return _TracedStream(stream, trace_container, reduce_fn)868 elif hasattr(stream, "__aiter__"):869 # sync function -> async iterable (unexpected)870 return _TracedAsyncStream(stream, trace_container, reduce_fn)871 872 # If it's not iterable, end the trace immediately873 _on_run_end(trace_container, outputs=stream)874 return stream875 876 @functools.wraps(func)877 async def async_stream_wrapper(878 *args: Any, langsmith_extra: Optional[LangSmithExtra] = None, **kwargs: Any879 ) -> Any:880 if not func_accepts_config:881 kwargs.pop("config", None)882 trace_container = await aitertools.aio_to_thread(883 copy_context(),884 _setup_run,885 func,886 container_input=container_input,887 langsmith_extra=langsmith_extra,888 args=args,889 kwargs=kwargs,890 )891 892 try:893 if func_accepts_parent_run:894 kwargs["run_tree"] = trace_container["new_run"]895 stream = await func(*args, **kwargs)896 except Exception as e:897 await aitertools.aio_to_thread(898 copy_context(), _on_run_end, trace_container, error=e899 )900 raise901 902 if hasattr(stream, "__aiter__"):903 return _TracedAsyncStream(stream, trace_container, reduce_fn)904 elif hasattr(stream, "__iter__"):905 # Async function -> sync iterable906 return _TracedStream(stream, trace_container, reduce_fn)907 908 # If it's not iterable, end the trace immediately909 await aitertools.aio_to_thread(910 copy_context(), _on_run_end, trace_container, outputs=stream911 )912 return stream913 914 if inspect.isasyncgenfunction(func):915 selected_wrapper: Callable = async_generator_wrapper916 elif inspect.isgeneratorfunction(func):917 selected_wrapper = generator_wrapper918 elif is_async(func):919 if reduce_fn:920 selected_wrapper = async_stream_wrapper921 else:922 selected_wrapper = async_wrapper923 else:924 if reduce_fn:925 selected_wrapper = stream_wrapper926 else:927 selected_wrapper = wrapper928 setattr(selected_wrapper, "__langsmith_traceable__", True)929 setattr(930 selected_wrapper,931 "__traceable_config__",932 {933 "process_inputs": container_input.get("process_inputs"),934 "process_outputs": outputs_processor,935 "enabled": enabled,936 "tags": container_input.get("tags"),937 "metadata": container_input.get("metadata"),938 "wrapped": func,939 },940 )941 sig = inspect.signature(selected_wrapper)942 if not sig.parameters.get("config"):943 sig = sig.replace(944 parameters=[945 *(946 param947 for param in sig.parameters.values()948 if param.kind != inspect.Parameter.VAR_KEYWORD949 ),950 inspect.Parameter(951 "config", inspect.Parameter.KEYWORD_ONLY, default=None952 ),953 *(954 param955 for param in sig.parameters.values()956 if param.kind == inspect.Parameter.VAR_KEYWORD957 ),958 ]959 )960 selected_wrapper.__signature__ = sig # type: ignore[attr-defined]961 return selected_wrapper962 963 # If the decorator is called with no arguments, then it's being used as a964 # decorator, so we return the decorator function965 if len(args) == 1 and callable(args[0]) and not kwargs:966 return decorator(args[0])967 # Else it's being used as a decorator factory, so we return the decorator968 return decorator969 970 971class trace:972 """Manage a LangSmith run in context.973 974 This class can be used as both a synchronous and asynchronous context manager.975 976 Args:977 name: Name of the run.978 run_type: Type of run (e.g., `'chain'`, `'llm'`, `'tool'`).979 inputs: Initial input data for the run.980 project_name: Project name to associate the run with.981 parent: Parent run.982 983 Can be a `RunTree`, dotted order string, or tracing headers.984 tags: List of tags for the run.985 metadata: Additional metadata for the run.986 client: LangSmith client for custom settings.987 run_id: Preset identifier for the run.988 reference_example_id: Associates run with a dataset example.989 990 Only for root runs in evaluation.991 exceptions_to_handle: Exception types to ignore.992 extra: Extra data to send to LangSmith.993 994 Use 'metadata' instead.995 996 Examples:997 Synchronous usage:998 999 ```python1000 with trace("My Operation", run_type="tool", tags=["important"]) as run:1001 result = "foo" # Perform operation1002 run.metadata["some-key"] = "some-value"1003 run.end(outputs={"result": result})1004 ```1005 1006 Asynchronous usage:1007 1008 ```python1009 async def main():1010 async with trace("Async Operation", run_type="tool", tags=["async"]) as run:1011 result = "foo" # Await async operation1012 run.metadata["some-key"] = "some-value"1013 # "end" just adds the outputs and sets error to None1014 # The actual patching of the run happens when the context exits1015 run.end(outputs={"result": result})1016 1017 1018 asyncio.run(main())1019 ```1020 1021 Handling specific exceptions:1022 1023 ```python1024 import pytest1025 import sys1026 1027 with trace("Test", exceptions_to_handle=(pytest.skip.Exception,)):1028 if sys.platform == "win32": # Just an example1029 pytest.skip("Skipping test for windows")1030 result = "foo" # Perform test operation1031 ```1032 """1033 1034 def __init__(1035 self,1036 name: str,1037 run_type: ls_client.RUN_TYPE_T = "chain",1038 *,1039 inputs: Optional[dict] = None,1040 extra: Optional[dict] = None,1041 project_name: Optional[str] = None,1042 parent: Optional[1043 Union[run_trees.RunTree, str, Mapping, Literal["ignore"]]1044 ] = None,1045 tags: Optional[list[str]] = None,1046 metadata: Optional[Mapping[str, Any]] = None,1047 client: Optional[ls_client.Client] = None,1048 run_id: Optional[ls_client.ID_TYPE] = None,1049 reference_example_id: Optional[ls_client.ID_TYPE] = None,1050 exceptions_to_handle: Optional[tuple[type[BaseException], ...]] = None,1051 attachments: Optional[schemas.Attachments] = None,1052 **kwargs: Any,1053 ):1054 """Initialize the trace context manager.1055 1056 Warns if unsupported kwargs are passed.1057 """1058 self._end_on_exit = kwargs.pop("_end_on_exit", True)1059 if kwargs:1060 warnings.warn(1061 "The `trace` context manager no longer supports the following kwargs: "1062 f"{sorted(kwargs.keys())}.",1063 DeprecationWarning,1064 )1065 self.name = name1066 self.run_type = run_type1067 self.inputs = inputs1068 self.attachments = attachments1069 self.extra = extra1070 self.project_name = project_name1071 self.parent = parent1072 # The run tree is deprecated. Keeping for backwards compat.1073 # Will fully merge within parent later.1074 self.run_tree = kwargs.get("run_tree")1075 self.tags = tags1076 self.metadata = metadata1077 self.client = client1078 self.run_id = run_id1079 self.reference_example_id = reference_example_id1080 self.exceptions_to_handle = exceptions_to_handle1081 self.new_run: Optional[run_trees.RunTree] = None1082 self.old_ctx: Optional[dict] = None1083 1084 def _setup(self) -> run_trees.RunTree:1085 """Set up the tracing context and create a new run.1086 1087 This method initializes the tracing context, merges tags and metadata,1088 creates a new run (either as a child of an existing run or as a new root run),1089 and sets up the necessary context variables.1090 1091 Returns:1092 run_trees.RunTree: The newly created run.1093 """1094 self.old_ctx = get_tracing_context()1095 enabled = utils.tracing_is_enabled(self.old_ctx)1096 1097 outer_tags = _context._TAGS.get() or _context._GLOBAL_TAGS1098 outer_metadata = _context._METADATA.get() or _context._GLOBAL_METADATA1099 client_ = self.client or self.old_ctx.get("client")1100 parent_run_ = _get_parent_run(1101 {1102 "parent": self.parent,1103 "run_tree": self.run_tree,1104 "client": client_,1105 "project_name": self.project_name,1106 }1107 )1108 1109 tags_ = sorted(set((self.tags or []) + (outer_tags or [])))1110 metadata = {1111 **(self.metadata or {}),1112 **(outer_metadata or {}),1113 "ls_method": "trace",1114 }1115 1116 extra_outer = self.extra or {}1117 extra_outer["metadata"] = metadata1118 1119 project_name_ = _get_project_name(self.project_name)1120 1121 if parent_run_ is not None and enabled:1122 self.new_run = parent_run_.create_child(1123 name=self.name,1124 run_id=self.run_id,1125 run_type=self.run_type,1126 extra=extra_outer,1127 inputs=self.inputs,1128 tags=tags_,1129 attachments=self.attachments,1130 )1131 else:1132 self.new_run = run_trees.RunTree(1133 name=self.name,1134 id=ls_client._ensure_uuid(self.run_id),1135 reference_example_id=ls_client._ensure_uuid(1136 self.reference_example_id, accept_null=True1137 ),1138 run_type=self.run_type,1139 extra=extra_outer,1140 project_name=project_name_ or "default",1141 replicas=run_trees._REPLICAS.get(),1142 inputs=self.inputs or {},1143 tags=tags_,1144 client=client_, # type: ignore1145 attachments=self.attachments or {}, # type: ignore1146 )1147 1148 if enabled is True:1149 self.new_run.post()1150 if enabled:1151 _context._TAGS.set(tags_)1152 _context._METADATA.set(metadata)1153 if self.new_run is not None:1154 _context._PARENT_RUN_TREE_REF.set(weakref.ref(self.new_run))1155 else:1156 _context._PARENT_RUN_TREE_REF.set(None)1157 _context._PROJECT_NAME.set(project_name_)1158 _context._CLIENT.set(client_)1159 1160 return self.new_run1161 1162 def _teardown(1163 self,1164 exc_type: Optional[type[BaseException]],1165 exc_value: Optional[BaseException],1166 traceback: Optional[TracebackType],1167 ) -> None:1168 """Clean up the tracing context and finalize the run.1169 1170 This method handles exceptions, ends the run if necessary,1171 patches the run if it's not disabled, and resets the tracing context.1172 1173 Args:1174 exc_type: The type of the exception that occurred, if any.1175 exc_value: The exception instance that occurred, if any.1176 traceback: The traceback object associated with the exception, if any.1177 """1178 if self.new_run is None:1179 return1180 if exc_type is not None:1181 tb = _format_error_with_exceptions_to_handle(1182 exc_value, self.exceptions_to_handle1183 )1184 self.new_run.end(error=tb)1185 if self.old_ctx is not None:1186 enabled = utils.tracing_is_enabled(self.old_ctx)1187 if enabled is True and self._end_on_exit:1188 self.new_run.patch()1189 1190 _set_tracing_context(self.old_ctx)1191 else:1192 warnings.warn("Tracing context was not set up properly.", RuntimeWarning)1193 1194 def __enter__(self) -> run_trees.RunTree:1195 """Enter the context manager synchronously.1196 1197 Returns:1198 run_trees.RunTree: The newly created run.1199 """1200 return self._setup()