aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

aiq/builder/builder.py

@@ -20,22 +20,32 @@ from abc import abstractmethod
 from collections.abc import Sequence
 from pathlib import Path
 
+from aiq.authentication.interfaces import AuthProviderBase
 from aiq.builder.context import AIQContext
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.builder.function import Function
+from aiq.data_models.authentication import AuthProviderBaseConfig
+from aiq.data_models.component_ref import AuthenticationRef
 from aiq.data_models.component_ref import EmbedderRef
 from aiq.data_models.component_ref import FunctionRef
+from aiq.data_models.component_ref import ITSStrategyRef
 from aiq.data_models.component_ref import LLMRef
 from aiq.data_models.component_ref import MemoryRef
+from aiq.data_models.component_ref import ObjectStoreRef
 from aiq.data_models.component_ref import RetrieverRef
 from aiq.data_models.embedder import EmbedderBaseConfig
 from aiq.data_models.evaluator import EvaluatorBaseConfig
 from aiq.data_models.function import FunctionBaseConfig
 from aiq.data_models.function_dependencies import FunctionDependencies
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
 from aiq.data_models.llm import LLMBaseConfig
 from aiq.data_models.memory import MemoryBaseConfig
+from aiq.data_models.object_store import ObjectStoreBaseConfig
 from aiq.data_models.retriever import RetrieverBaseConfig
+from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
 from aiq.memory.interfaces import MemoryEditor
+from aiq.object_store.interfaces import ObjectStore
 from aiq.retriever.interface import AIQRetriever
 
 
@@ -91,6 +101,10 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
     async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
         pass
 
+    @abstractmethod
+    async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
+        pass
+
     async def get_llms(self, llm_names: Sequence[str | LLMRef],
                        wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
 
@@ -101,11 +115,41 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         return list(llms)
 
     @abstractmethod
-    async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
+    def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
         pass
 
     @abstractmethod
-    def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
+    async def add_auth_provider(self, name: str | AuthenticationRef, config: AuthProviderBaseConfig):
+        pass
+
+    @abstractmethod
+    async def get_auth_provider(self, auth_provider_name: str | AuthenticationRef) -> AuthProviderBase:
+        pass
+
+    async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]):
+
+        coros = [self.get_auth_provider(auth_provider_name=n) for n in auth_provider_names]
+
+        auth_providers = await asyncio.gather(*coros, return_exceptions=False)
+
+        return list(auth_providers)
+
+    @abstractmethod
+    async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig):
+        pass
+
+    async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
+        """
+        Return a list of all object store clients.
+        """
+        return list(await asyncio.gather(*[self.get_object_store_client(name) for name in object_store_names]))
+
+    @abstractmethod
+    async def get_object_store_client(self, object_store_name: str | ObjectStoreRef) -> ObjectStore:
+        pass
+
+    @abstractmethod
+    def get_object_store_config(self, object_store_name: str | ObjectStoreRef) -> ObjectStoreBaseConfig:
         pass
 
     @abstractmethod
@@ -187,6 +231,24 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
     async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
         pass
 
+    @abstractmethod
+    async def add_its_strategy(self, name: str | str, config: ITSStrategyBaseConfig):
+        pass
+
+    @abstractmethod
+    async def get_its_strategy(self,
+                               strategy_name: str | ITSStrategyRef,
+                               pipeline_type: PipelineTypeEnum,
+                               stage_type: StageTypeEnum):
+        pass
+
+    @abstractmethod
+    async def get_its_strategy_config(self,
+                                      strategy_name: str | ITSStrategyRef,
+                                      pipeline_type: PipelineTypeEnum,
+                                      stage_type: StageTypeEnum) -> ITSStrategyBaseConfig:
+        pass
+
     @abstractmethod
     def get_user_manager(self) -> UserManagerHolder:
         pass

aiq/builder/component_utils.py

@@ -21,6 +21,7 @@ from collections.abc import Iterable
 import networkx as nx
 from pydantic import BaseModel
 
+from aiq.data_models.authentication import AuthProviderBaseConfig
 from aiq.data_models.common import TypedBaseModel
 from aiq.data_models.component import ComponentGroup
 from aiq.data_models.component_ref import ComponentRef
@@ -29,8 +30,10 @@ from aiq.data_models.component_ref import generate_instance_id
 from aiq.data_models.config import AIQConfig
 from aiq.data_models.embedder import EmbedderBaseConfig
 from aiq.data_models.function import FunctionBaseConfig
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
 from aiq.data_models.llm import LLMBaseConfig
 from aiq.data_models.memory import MemoryBaseConfig
+from aiq.data_models.object_store import ObjectStoreBaseConfig
 from aiq.data_models.retriever import RetrieverBaseConfig
 from aiq.utils.type_utils import DecomposedType
 
@@ -38,11 +41,14 @@ logger = logging.getLogger(__name__)
 
 # Order in which we want to process the component groups
 _component_group_order = [
+    ComponentGroup.AUTHENTICATION,
     ComponentGroup.EMBEDDERS,
     ComponentGroup.LLMS,
     ComponentGroup.MEMORY,
+    ComponentGroup.OBJECT_STORES,
     ComponentGroup.RETRIEVERS,
-    ComponentGroup.FUNCTIONS
+    ComponentGroup.ITS_STRATEGIES,
+    ComponentGroup.FUNCTIONS,
 ]
 
 
@@ -95,6 +101,8 @@ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
             component is not a valid runtime instance, None is returned.
     """
 
+    if (isinstance(component, AuthProviderBaseConfig)):
+        return ComponentGroup.AUTHENTICATION
     if (isinstance(component, EmbedderBaseConfig)):
         return ComponentGroup.EMBEDDERS
     if (isinstance(component, FunctionBaseConfig)):
@@ -103,8 +111,12 @@ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
         return ComponentGroup.LLMS
     if (isinstance(component, MemoryBaseConfig)):
         return ComponentGroup.MEMORY
+    if (isinstance(component, ObjectStoreBaseConfig)):
+        return ComponentGroup.OBJECT_STORES
     if (isinstance(component, RetrieverBaseConfig)):
         return ComponentGroup.RETRIEVERS
+    if (isinstance(component, ITSStrategyBaseConfig)):
+        return ComponentGroup.ITS_STRATEGIES
 
     return None
 
@@ -142,7 +154,7 @@ def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
             yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
     if (decomposed_type.is_union):
         for arg in decomposed_type.args:
-            if (isinstance(value, DecomposedType(arg).root)):
+            if arg is typing.Any or (isinstance(value, DecomposedType(arg).root)):
                 yield from recursive_componentref_discovery(cls, value, arg)
     else:
         for arg in decomposed_type.args:
@@ -243,7 +255,8 @@ def build_dependency_sequence(config: "AIQConfig") -> list[ComponentInstanceData
     """
 
     total_node_count = len(config.embedders) + len(config.functions) + len(config.llms) + len(config.memory) + len(
-        config.retrievers) + 1  # +1 for the workflow
+        config.object_stores) + len(config.retrievers) + len(config.its_strategies) + len(
+            config.authentication) + 1  # +1 for the workflow
 
     dependency_map: dict
     dependency_graph: nx.DiGraph

aiq/builder/context.py

@@ -22,6 +22,9 @@ from contextvars import ContextVar
 
 from aiq.builder.intermediate_step_manager import IntermediateStepManager
 from aiq.builder.user_interaction_manager import AIQUserInteractionManager
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.authentication import AuthProviderBaseConfig
 from aiq.data_models.interactive import HumanResponse
 from aiq.data_models.interactive import InteractionPrompt
 from aiq.data_models.intermediate_step import IntermediateStep
@@ -61,6 +64,7 @@ class ActiveFunctionContextManager:
 class AIQContextState(metaclass=Singleton):
 
     def __init__(self):
+        self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
         self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
@@ -75,6 +79,9 @@ class AIQContextState(metaclass=Singleton):
                                              | None] = ContextVar(
                                                  "user_input_callback",
                                                  default=AIQUserInteractionManager.default_callback_handler)
+        self.user_auth_callback: ContextVar[Callable[[AuthProviderBaseConfig, AuthFlowType],
+                                                     Awaitable[AuthenticatedContext]]
+                                            | None] = ContextVar("user_auth_callback", default=None)
 
     @staticmethod
     def get() -> "AIQContextState":
@@ -148,6 +155,16 @@ class AIQContext:
         """
         return IntermediateStepManager(self._context_state)
 
+    @property
+    def conversation_id(self) -> str | None:
+        """
+        This property retrieves the conversation ID which is the unique identifier for the current chat conversation.
+
+        Returns:
+            str | None
+        """
+        return self._context_state.conversation_id.get()
+
     @contextmanager
     def push_active_function(self, function_name: str, input_data: typing.Any | None):
         """
@@ -213,6 +230,26 @@ class AIQContext:
         """
         return self._context_state.active_span_id_stack.get()[-1]
 
+    @property
+    def user_auth_callback(self) -> Callable[[AuthProviderBaseConfig, AuthFlowType], Awaitable[AuthenticatedContext]]:
+        """
+        Retrieves the user authentication callback function from the context state.
+
+        This property provides access to the user authentication callback function stored in the context state.
+        The callback function is responsible for handling user authentication based on the provided configuration.
+
+        Returns:
+            Callable[[AuthenticationBaseConfig], Awaitable[AuthenticatedContext]]: The user authentication
+            callback function.
+
+        Raises:
+            RuntimeError: If the user authentication callback is not set in the context.
+        """
+        callback = self._context_state.user_auth_callback.get()
+        if callback is None:
+            raise RuntimeError("User authentication callback is not set in the context.")
+        return callback
+
     @staticmethod
     def get() -> "AIQContext":
         """

aiq/builder/eval_builder.py

@@ -102,14 +102,55 @@ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
 
         return tools
 
+    def _log_build_failure_evaluator(self,
+                                     failing_evaluator_name: str,
+                                     completed_evaluators: list[str],
+                                     remaining_evaluators: list[str],
+                                     original_error: Exception) -> None:
+        """
+        Log comprehensive evaluator build failure information.
+
+        Args:
+            failing_evaluator_name (str): The name of the evaluator that failed to build
+            completed_evaluators (list[str]): List of evaluator names that were successfully built
+            remaining_evaluators (list[str]): List of evaluator names still to be built
+            original_error (Exception): The original exception that caused the failure
+        """
+        # Convert evaluator names to (name, type) tuples for consistent logging
+        completed_components = [(name, "evaluator") for name in completed_evaluators]
+        remaining_components = [(name, "evaluator") for name in remaining_evaluators]
+
+        # Use the inherited common logging method from WorkflowBuilder
+        self._log_build_failure(failing_evaluator_name,
+                                "evaluator",
+                                completed_components,
+                                remaining_components,
+                                original_error)
+
     async def populate_builder(self, config: AIQConfig):
         # Skip setting workflow if workflow config is EmptyFunctionConfig
         skip_workflow = isinstance(config.workflow, EmptyFunctionConfig)
 
         await super().populate_builder(config, skip_workflow)
-        # Instantiate the evaluators
+
+        # Initialize progress tracking for evaluators
+        completed_evaluators = []
+        remaining_evaluators = list(config.eval.evaluators.keys())
+
+        # Instantiate the evaluators with enhanced error logging
         for name, evaluator_config in config.eval.evaluators.items():
-            await self.add_evaluator(name, evaluator_config)
+            try:
+                # Remove from remaining as we start building
+                remaining_evaluators.remove(name)
+
+                await self.add_evaluator(name, evaluator_config)
+
+                # Add to completed after successful build
+                completed_evaluators.append(name)
+
+            except Exception as e:
+                self._log_build_failure_evaluator(name, completed_evaluators, remaining_evaluators, e)
+                raise
 
     @classmethod
     @asynccontextmanager

aiq/builder/function.py

@@ -48,7 +48,8 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
                  input_schema: type[BaseModel] | None = None,
                  streaming_output_schema: type[BaseModel] | type[None] | None = None,
                  single_output_schema: type[BaseModel] | type[None] | None = None,
-                 converters: list[Callable[[typing.Any], typing.Any]] | None = None):
+                 converters: list[Callable[[typing.Any], typing.Any]] | None = None,
+                 instance_name: str | None = None):
 
         super().__init__(input_schema=input_schema,
                          streaming_output_schema=streaming_output_schema,
@@ -57,6 +58,7 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
 
         self.config = config
         self.description = description
+        self.instance_name = instance_name or config.type
         self._context = AIQContext.get()
 
     def convert(self, value: typing.Any, to_type: type[_T]) -> _T:
@@ -78,6 +80,25 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
 
         return self._converter.convert(value, to_type=to_type)
 
+    def try_convert(self, value: typing.Any, to_type: type[_T]) -> _T:
+        """
+        Converts the given value to the specified type using graceful error handling.
+        If conversion fails, returns the original value and continues processing.
+
+        Parameters
+        ----------
+        value : typing.Any
+            The value to convert.
+        to_type : type
+            The type to convert the value to.
+
+        Returns
+        -------
+        _T
+            The converted value, or original value if conversion fails.
+        """
+        return self._converter.try_convert(value, to_type=to_type)
+
     @abstractmethod
     async def _ainvoke(self, value: InputT) -> SingleOutputT:
         pass
@@ -110,7 +131,7 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
             The output of the function optionally converted to the specified type.
         """
 
-        with self._context.push_active_function(self.config.type,
+        with self._context.push_active_function(self.instance_name,
                                                 input_data=value) as manager:  # Set the current invocation context
             try:
                 converted_input: InputT = self._convert_input(value)  # type: ignore
@@ -118,7 +139,7 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
                 result = await self._ainvoke(converted_input)
 
                 if to_type is not None and not isinstance(result, to_type):
-                    result = self._converter.convert(result, to_type=to_type)
+                    result = self._converter.try_convert(result, to_type=to_type)
 
                 manager.set_output(result)
 
@@ -196,16 +217,25 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
             The output of the function optionally converted to the specified type.
         """
 
-        with self._context.push_active_function(self.config.type, input_data=value):
+        with self._context.push_active_function(self.instance_name, input_data=value) as manager:
             try:
                 converted_input: InputT = self._convert_input(value)  # type: ignore
 
-                async for data in self._astream(converted_input):
+                # Collect streaming outputs to capture the final result
+                final_output: list[typing.Any] = []
 
+                async for data in self._astream(converted_input):
                     if to_type is not None and not isinstance(data, to_type):
-                        yield self._converter.convert(data, to_type=to_type)
+                        converted_data = self._converter.try_convert(data, to_type=to_type)
+                        final_output.append(converted_data)
+                        yield converted_data
                     else:
+                        final_output.append(data)
                         yield data
+
+                # Set the final output for intermediate step tracking
+                manager.set_output(final_output)
+
             except Exception as e:
                 logger.error("Error with astream in function with input: %s.", value, exc_info=True)
                 raise e
@@ -254,17 +284,17 @@ class Function(FunctionBase[InputT, StreamingOutputT, SingleOutputT], ABC):
 
 class LambdaFunction(Function[InputT, StreamingOutputT, SingleOutputT]):
 
-    def __init__(self, *, config: FunctionBaseConfig, info: FunctionInfo):
+    def __init__(self, *, config: FunctionBaseConfig, info: FunctionInfo, instance_name: str | None = None):
 
         super().__init__(config=config,
                          description=info.description,
                          input_schema=info.input_schema,
                          streaming_output_schema=info.stream_output_schema,
                          single_output_schema=info.single_output_schema,
-                         converters=info.converters)
+                         converters=info.converters,
+                         instance_name=instance_name)
 
         self._info = info
-
         self._ainvoke_fn: _InvokeFnT = info.single_fn
         self._astream_fn: _StreamFnT = info.stream_fn
 
@@ -284,8 +314,10 @@ class LambdaFunction(Function[InputT, StreamingOutputT, SingleOutputT]):
             yield x
 
     @staticmethod
-    def from_info(*, config: FunctionBaseConfig,
-                  info: FunctionInfo) -> 'LambdaFunction[InputT, StreamingOutputT, SingleOutputT]':
+    def from_info(*,
+                  config: FunctionBaseConfig,
+                  info: FunctionInfo,
+                  instance_name: str | None = None) -> 'LambdaFunction[InputT, StreamingOutputT, SingleOutputT]':
 
         input_type: type = info.input_type
         streaming_output_type = info.stream_output_type
@@ -294,4 +326,4 @@ class LambdaFunction(Function[InputT, StreamingOutputT, SingleOutputT]):
         class FunctionImpl(LambdaFunction[input_type, streaming_output_type, single_output_type]):
             pass
 
-        return FunctionImpl(config=config, info=info)
+        return FunctionImpl(config=config, info=info, instance_name=instance_name)

aiq/builder/function_base.py

@@ -373,4 +373,4 @@ class FunctionBase(typing.Generic[InputT, StreamingOutputT, SingleOutputT], ABC)
             return value
 
         # Fallback to the converter
-        return self._converter.convert(value, to_type=self.input_class)
+        return self._converter.try_convert(value, to_type=self.input_class)

aiq/builder/intermediate_step_manager.py

@@ -20,7 +20,6 @@ import typing
 from aiq.data_models.intermediate_step import IntermediateStep
 from aiq.data_models.intermediate_step import IntermediateStepPayload
 from aiq.data_models.intermediate_step import IntermediateStepState
-from aiq.data_models.invocation_node import InvocationNode
 from aiq.utils.reactive.observable import OnComplete
 from aiq.utils.reactive.observable import OnError
 from aiq.utils.reactive.observable import OnNext
@@ -37,7 +36,7 @@ class OpenStep:
     step_id: str
     step_name: str
     step_type: str
-    step_parent_id: str | None
+    step_parent_id: str
     prev_stack: list[str]
     active_stack: list[str]
 
@@ -153,15 +152,14 @@ class IntermediateStepManager:
                 return
 
             parent_step_id = open_step.step_parent_id
+        else:
+            assert False, "Invalid event state"
 
         active_function = self._context_state.active_function.get()
 
-        function_ancestry = InvocationNode(function_name=active_function.function_name,
-                                           function_id=active_function.function_id,
-                                           parent_id=parent_step_id,
-                                           parent_name=active_function.parent_name)
-
-        intermediate_step = IntermediateStep(function_ancestry=function_ancestry, payload=payload)
+        intermediate_step = IntermediateStep(parent_id=parent_step_id,
+                                             function_ancestry=active_function,
+                                             payload=payload)
 
         self._context_state.event_stream.get().on_next(intermediate_step)
 

aiq/builder/user_interaction_manager.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import time
 import uuid
 
@@ -22,6 +23,8 @@ from aiq.data_models.interactive import InteractionPrompt
 from aiq.data_models.interactive import InteractionResponse
 from aiq.data_models.interactive import InteractionStatus
 
+logger = logging.getLogger(__name__)
+
 
 class AIQUserInteractionManager:
     """

aiq/builder/workflow.py

@@ -27,19 +27,12 @@ from aiq.builder.function_base import StreamingOutputT
 from aiq.builder.llm import LLMProviderInfo
 from aiq.builder.retriever import RetrieverProviderInfo
 from aiq.data_models.config import AIQConfig
+from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
 from aiq.memory.interfaces import MemoryEditor
+from aiq.object_store.interfaces import ObjectStore
+from aiq.observability.exporter.base_exporter import BaseExporter
+from aiq.observability.exporter_manager import ExporterManager
 from aiq.runtime.runner import AIQRunner
-from aiq.utils.optional_imports import TelemetryOptionalImportError
-from aiq.utils.optional_imports import try_import_opentelemetry
-
-# Try to import OpenTelemetry modules
-# If the dependencies are not installed, use a dummy span exporter here
-try:
-    opentelemetry = try_import_opentelemetry()
-    from opentelemetry.sdk.trace.export import SpanExporter
-except TelemetryOptionalImportError:
-    from aiq.utils.optional_imports import DummySpanExporter  # pylint: disable=ungrouped-imports
-    SpanExporter = DummySpanExporter
 
 callback_handler_var: ContextVar[Any | None] = ContextVar("callback_handler_var", default=None)
 
@@ -54,8 +47,10 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
                  llms: dict[str, LLMProviderInfo] | None = None,
                  embeddings: dict[str, EmbedderProviderInfo] | None = None,
                  memory: dict[str, MemoryEditor] | None = None,
-                 exporters: dict[str, SpanExporter] | None = None,
+                 object_stores: dict[str, ObjectStore] | None = None,
+                 telemetry_exporters: dict[str, BaseExporter] | None = None,
                  retrievers: dict[str | None, RetrieverProviderInfo] | None = None,
+                 its_strategies: dict[str, StrategyBase] | None = None,
                  context_state: AIQContextState):
 
         super().__init__(input_schema=entry_fn.input_schema,
@@ -67,14 +62,17 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
         self.llms = llms or {}
         self.embeddings = embeddings or {}
         self.memory = memory or {}
+        self.telemetry_exporters = telemetry_exporters or {}
+        self.object_stores = object_stores or {}
         self.retrievers = retrievers or {}
 
+        self._exporter_manager = ExporterManager.from_exporters(self.telemetry_exporters)
+        self.its_strategies = its_strategies or {}
+
         self._entry_fn = entry_fn
 
         self._context_state = context_state
 
-        self._exporters = exporters or {}
-
     @property
     def has_streaming_output(self) -> bool:
 
@@ -91,8 +89,11 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
         Called each time we start a new workflow run. We'll create
         a new top-level workflow span here.
         """
-        async with AIQRunner(input_message=message, entry_fn=self._entry_fn,
-                             context_state=self._context_state) as runner:
+
+        async with AIQRunner(input_message=message,
+                             entry_fn=self._entry_fn,
+                             context_state=self._context_state,
+                             exporter_manager=self._exporter_manager.get()) as runner:
 
             # The caller can `yield runner` so they can do `runner.result()` or `runner.result_stream()`
             yield runner
@@ -121,8 +122,10 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
                       llms: dict[str, LLMProviderInfo] | None = None,
                       embeddings: dict[str, EmbedderProviderInfo] | None = None,
                       memory: dict[str, MemoryEditor] | None = None,
-                      exporters: dict[str, SpanExporter] | None = None,
+                      object_stores: dict[str, ObjectStore] | None = None,
+                      telemetry_exporters: dict[str, BaseExporter] | None = None,
                       retrievers: dict[str | None, RetrieverProviderInfo] | None = None,
+                      its_strategies: dict[str, StrategyBase] | None = None,
                       context_state: AIQContextState) -> 'Workflow[InputT, StreamingOutputT, SingleOutputT]':
 
         input_type: type = entry_fn.input_type
@@ -138,6 +141,8 @@ class Workflow(FunctionBase[InputT, StreamingOutputT, SingleOutputT]):
                             llms=llms,
                             embeddings=embeddings,
                             memory=memory,
-                            exporters=exporters,
+                            object_stores=object_stores,
+                            telemetry_exporters=telemetry_exporters,
                             retrievers=retrievers,
+                            its_strategies=its_strategies,
                             context_state=context_state)