diff --git a/python_modules/libraries/dagster-openai/dagster_openai/resources.py b/python_modules/libraries/dagster-openai/dagster_openai/resources.py index 2c8afbd07c251..c1a2dbcd5d07b 100644 --- a/python_modules/libraries/dagster-openai/dagster_openai/resources.py +++ b/python_modules/libraries/dagster-openai/dagster_openai/resources.py @@ -12,6 +12,7 @@ InitResourceContext, OpExecutionContext, ) +from dagster._annotations import experimental from dagster._core.errors import ( DagsterInvariantViolationError, ) @@ -36,7 +37,7 @@ class ApiEndpointClassesEnum(Enum): context_to_counters = WeakKeyDictionary() -def add_to_asset_metadata( +def _add_to_asset_metadata( context: AssetExecutionContext, usage_metadata: dict, output_name: Optional[str] ): if context not in context_to_counters: @@ -48,6 +49,7 @@ def add_to_asset_metadata( context.add_output_metadata(dict(counters), output_name) +@experimental def with_usage_metadata(context: AssetExecutionContext, output_name: Optional[str], func): """This wrapper can be used on any endpoint of the `openai library ` @@ -132,13 +134,14 @@ def wrapper(*args, **kwargs): } if hasattr(usage, "completion_tokens"): usage_metadata["openai.completion_tokens"] = usage.completion_tokens - add_to_asset_metadata(context, usage_metadata, output_name) + _add_to_asset_metadata(context, usage_metadata, output_name) return response return wrapper +@experimental class OpenAIResource(ConfigurableResource): """This resource is wrapper over the `openai library `_.