aiqtoolkit 1.2.0a20250707__py3-none-any.whl → 1.2.0a20250731__py3-none-any.whl

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
@@ -76,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":
@@ -224,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

@@ -80,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
@@ -120,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)
 
@@ -198,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

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)