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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of aiqtoolkit might be problematic. Click here for more details.

Files changed (198) hide show
  1. aiq/agent/base.py +171 -8
  2. aiq/agent/dual_node.py +1 -1
  3. aiq/agent/react_agent/agent.py +113 -113
  4. aiq/agent/react_agent/register.py +31 -14
  5. aiq/agent/rewoo_agent/agent.py +36 -35
  6. aiq/agent/rewoo_agent/register.py +2 -2
  7. aiq/agent/tool_calling_agent/agent.py +3 -7
  8. aiq/authentication/__init__.py +14 -0
  9. aiq/authentication/api_key/__init__.py +14 -0
  10. aiq/authentication/api_key/api_key_auth_provider.py +92 -0
  11. aiq/authentication/api_key/api_key_auth_provider_config.py +124 -0
  12. aiq/authentication/api_key/register.py +26 -0
  13. aiq/authentication/exceptions/__init__.py +14 -0
  14. aiq/authentication/exceptions/api_key_exceptions.py +38 -0
  15. aiq/authentication/exceptions/auth_code_grant_exceptions.py +86 -0
  16. aiq/authentication/exceptions/call_back_exceptions.py +38 -0
  17. aiq/authentication/exceptions/request_exceptions.py +54 -0
  18. aiq/authentication/http_basic_auth/__init__.py +0 -0
  19. aiq/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  20. aiq/authentication/http_basic_auth/register.py +30 -0
  21. aiq/authentication/interfaces.py +93 -0
  22. aiq/authentication/oauth2/__init__.py +14 -0
  23. aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  24. aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  25. aiq/authentication/oauth2/register.py +25 -0
  26. aiq/authentication/register.py +21 -0
  27. aiq/builder/builder.py +64 -2
  28. aiq/builder/component_utils.py +16 -3
  29. aiq/builder/context.py +26 -0
  30. aiq/builder/eval_builder.py +43 -2
  31. aiq/builder/function.py +32 -4
  32. aiq/builder/function_base.py +1 -1
  33. aiq/builder/intermediate_step_manager.py +6 -8
  34. aiq/builder/user_interaction_manager.py +3 -0
  35. aiq/builder/workflow.py +23 -18
  36. aiq/builder/workflow_builder.py +420 -73
  37. aiq/cli/commands/info/list_mcp.py +103 -16
  38. aiq/cli/commands/sizing/__init__.py +14 -0
  39. aiq/cli/commands/sizing/calc.py +294 -0
  40. aiq/cli/commands/sizing/sizing.py +27 -0
  41. aiq/cli/commands/start.py +1 -0
  42. aiq/cli/entrypoint.py +2 -0
  43. aiq/cli/register_workflow.py +80 -0
  44. aiq/cli/type_registry.py +151 -30
  45. aiq/data_models/api_server.py +117 -11
  46. aiq/data_models/authentication.py +231 -0
  47. aiq/data_models/common.py +35 -7
  48. aiq/data_models/component.py +17 -9
  49. aiq/data_models/component_ref.py +33 -0
  50. aiq/data_models/config.py +60 -3
  51. aiq/data_models/embedder.py +1 -0
  52. aiq/data_models/function_dependencies.py +8 -0
  53. aiq/data_models/interactive.py +10 -1
  54. aiq/data_models/intermediate_step.py +15 -5
  55. aiq/data_models/its_strategy.py +30 -0
  56. aiq/data_models/llm.py +1 -0
  57. aiq/data_models/memory.py +1 -0
  58. aiq/data_models/object_store.py +44 -0
  59. aiq/data_models/retry_mixin.py +35 -0
  60. aiq/data_models/span.py +187 -0
  61. aiq/data_models/telemetry_exporter.py +2 -2
  62. aiq/embedder/nim_embedder.py +2 -1
  63. aiq/embedder/openai_embedder.py +2 -1
  64. aiq/eval/config.py +19 -1
  65. aiq/eval/dataset_handler/dataset_handler.py +75 -1
  66. aiq/eval/evaluate.py +53 -10
  67. aiq/eval/rag_evaluator/evaluate.py +23 -12
  68. aiq/eval/remote_workflow.py +7 -2
  69. aiq/eval/runners/__init__.py +14 -0
  70. aiq/eval/runners/config.py +39 -0
  71. aiq/eval/runners/multi_eval_runner.py +54 -0
  72. aiq/eval/usage_stats.py +6 -0
  73. aiq/eval/utils/weave_eval.py +5 -1
  74. aiq/experimental/__init__.py +0 -0
  75. aiq/experimental/decorators/__init__.py +0 -0
  76. aiq/experimental/decorators/experimental_warning_decorator.py +130 -0
  77. aiq/experimental/inference_time_scaling/__init__.py +0 -0
  78. aiq/experimental/inference_time_scaling/editing/__init__.py +0 -0
  79. aiq/experimental/inference_time_scaling/editing/iterative_plan_refinement_editor.py +147 -0
  80. aiq/experimental/inference_time_scaling/editing/llm_as_a_judge_editor.py +204 -0
  81. aiq/experimental/inference_time_scaling/editing/motivation_aware_summarization.py +107 -0
  82. aiq/experimental/inference_time_scaling/functions/__init__.py +0 -0
  83. aiq/experimental/inference_time_scaling/functions/execute_score_select_function.py +105 -0
  84. aiq/experimental/inference_time_scaling/functions/its_tool_orchestration_function.py +205 -0
  85. aiq/experimental/inference_time_scaling/functions/its_tool_wrapper_function.py +146 -0
  86. aiq/experimental/inference_time_scaling/functions/plan_select_execute_function.py +224 -0
  87. aiq/experimental/inference_time_scaling/models/__init__.py +0 -0
  88. aiq/experimental/inference_time_scaling/models/editor_config.py +132 -0
  89. aiq/experimental/inference_time_scaling/models/its_item.py +48 -0
  90. aiq/experimental/inference_time_scaling/models/scoring_config.py +112 -0
  91. aiq/experimental/inference_time_scaling/models/search_config.py +120 -0
  92. aiq/experimental/inference_time_scaling/models/selection_config.py +154 -0
  93. aiq/experimental/inference_time_scaling/models/stage_enums.py +43 -0
  94. aiq/experimental/inference_time_scaling/models/strategy_base.py +66 -0
  95. aiq/experimental/inference_time_scaling/models/tool_use_config.py +41 -0
  96. aiq/experimental/inference_time_scaling/register.py +36 -0
  97. aiq/experimental/inference_time_scaling/scoring/__init__.py +0 -0
  98. aiq/experimental/inference_time_scaling/scoring/llm_based_agent_scorer.py +168 -0
  99. aiq/experimental/inference_time_scaling/scoring/llm_based_plan_scorer.py +168 -0
  100. aiq/experimental/inference_time_scaling/scoring/motivation_aware_scorer.py +111 -0
  101. aiq/experimental/inference_time_scaling/search/__init__.py +0 -0
  102. aiq/experimental/inference_time_scaling/search/multi_llm_planner.py +128 -0
  103. aiq/experimental/inference_time_scaling/search/multi_query_retrieval_search.py +122 -0
  104. aiq/experimental/inference_time_scaling/search/single_shot_multi_plan_planner.py +128 -0
  105. aiq/experimental/inference_time_scaling/selection/__init__.py +0 -0
  106. aiq/experimental/inference_time_scaling/selection/best_of_n_selector.py +63 -0
  107. aiq/experimental/inference_time_scaling/selection/llm_based_agent_output_selector.py +131 -0
  108. aiq/experimental/inference_time_scaling/selection/llm_based_output_merging_selector.py +159 -0
  109. aiq/experimental/inference_time_scaling/selection/llm_based_plan_selector.py +128 -0
  110. aiq/experimental/inference_time_scaling/selection/threshold_selector.py +58 -0
  111. aiq/front_ends/console/authentication_flow_handler.py +233 -0
  112. aiq/front_ends/console/console_front_end_plugin.py +13 -24
  113. aiq/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  114. aiq/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  115. aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  116. aiq/front_ends/fastapi/fastapi_front_end_config.py +20 -0
  117. aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  118. aiq/front_ends/fastapi/fastapi_front_end_plugin.py +14 -1
  119. aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +353 -31
  120. aiq/front_ends/fastapi/html_snippets/__init__.py +14 -0
  121. aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  122. aiq/front_ends/fastapi/main.py +2 -0
  123. aiq/front_ends/fastapi/message_handler.py +102 -84
  124. aiq/front_ends/fastapi/step_adaptor.py +2 -1
  125. aiq/front_ends/simple_base/simple_front_end_plugin_base.py +4 -2
  126. aiq/llm/aws_bedrock_llm.py +2 -1
  127. aiq/llm/nim_llm.py +2 -1
  128. aiq/llm/openai_llm.py +2 -1
  129. aiq/object_store/__init__.py +20 -0
  130. aiq/object_store/in_memory_object_store.py +74 -0
  131. aiq/object_store/interfaces.py +84 -0
  132. aiq/object_store/models.py +36 -0
  133. aiq/object_store/register.py +20 -0
  134. aiq/observability/__init__.py +14 -0
  135. aiq/observability/exporter/__init__.py +14 -0
  136. aiq/observability/exporter/base_exporter.py +449 -0
  137. aiq/observability/exporter/exporter.py +78 -0
  138. aiq/observability/exporter/file_exporter.py +33 -0
  139. aiq/observability/exporter/processing_exporter.py +269 -0
  140. aiq/observability/exporter/raw_exporter.py +52 -0
  141. aiq/observability/exporter/span_exporter.py +264 -0
  142. aiq/observability/exporter_manager.py +335 -0
  143. aiq/observability/mixin/__init__.py +14 -0
  144. aiq/observability/mixin/batch_config_mixin.py +26 -0
  145. aiq/observability/mixin/collector_config_mixin.py +23 -0
  146. aiq/observability/mixin/file_mixin.py +288 -0
  147. aiq/observability/mixin/file_mode.py +23 -0
  148. aiq/observability/mixin/resource_conflict_mixin.py +134 -0
  149. aiq/observability/mixin/serialize_mixin.py +61 -0
  150. aiq/observability/mixin/type_introspection_mixin.py +183 -0
  151. aiq/observability/processor/__init__.py +14 -0
  152. aiq/observability/processor/batching_processor.py +316 -0
  153. aiq/observability/processor/intermediate_step_serializer.py +28 -0
  154. aiq/observability/processor/processor.py +68 -0
  155. aiq/observability/register.py +32 -116
  156. aiq/observability/utils/__init__.py +14 -0
  157. aiq/observability/utils/dict_utils.py +236 -0
  158. aiq/observability/utils/time_utils.py +31 -0
  159. aiq/profiler/calc/__init__.py +14 -0
  160. aiq/profiler/calc/calc_runner.py +623 -0
  161. aiq/profiler/calc/calculations.py +288 -0
  162. aiq/profiler/calc/data_models.py +176 -0
  163. aiq/profiler/calc/plot.py +345 -0
  164. aiq/profiler/data_models.py +2 -0
  165. aiq/profiler/profile_runner.py +16 -13
  166. aiq/runtime/loader.py +8 -2
  167. aiq/runtime/runner.py +23 -9
  168. aiq/runtime/session.py +16 -5
  169. aiq/tool/chat_completion.py +74 -0
  170. aiq/tool/code_execution/README.md +152 -0
  171. aiq/tool/code_execution/code_sandbox.py +151 -72
  172. aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
  173. aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +139 -24
  174. aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +3 -1
  175. aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +27 -2
  176. aiq/tool/code_execution/register.py +7 -3
  177. aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
  178. aiq/tool/mcp/exceptions.py +142 -0
  179. aiq/tool/mcp/mcp_client.py +17 -3
  180. aiq/tool/mcp/mcp_tool.py +1 -1
  181. aiq/tool/register.py +1 -0
  182. aiq/tool/server_tools.py +2 -2
  183. aiq/utils/exception_handlers/automatic_retries.py +289 -0
  184. aiq/utils/exception_handlers/mcp.py +211 -0
  185. aiq/utils/io/model_processing.py +28 -0
  186. aiq/utils/log_utils.py +37 -0
  187. aiq/utils/string_utils.py +38 -0
  188. aiq/utils/type_converter.py +18 -2
  189. aiq/utils/type_utils.py +87 -0
  190. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/METADATA +37 -9
  191. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/RECORD +196 -81
  192. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/entry_points.txt +3 -0
  193. aiq/front_ends/fastapi/websocket.py +0 -153
  194. aiq/observability/async_otel_listener.py +0 -470
  195. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/WHEEL +0 -0
  196. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  197. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/licenses/LICENSE.md +0 -0
  198. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0a20250731.dist-info}/top_level.txt +0 -0
aiq/cli/type_registry.py CHANGED
@@ -30,6 +30,7 @@ from pydantic import Tag
30
30
  from pydantic import computed_field
31
31
  from pydantic import field_validator
32
32
 
33
+ from aiq.authentication.interfaces import AuthProviderBase
33
34
  from aiq.builder.builder import Builder
34
35
  from aiq.builder.builder import EvalBuilder
35
36
  from aiq.builder.embedder import EmbedderProviderInfo
@@ -40,6 +41,8 @@ from aiq.builder.function_base import FunctionBase
40
41
  from aiq.builder.function_info import FunctionInfo
41
42
  from aiq.builder.llm import LLMProviderInfo
42
43
  from aiq.builder.retriever import RetrieverProviderInfo
44
+ from aiq.data_models.authentication import AuthProviderBaseConfig
45
+ from aiq.data_models.authentication import AuthProviderBaseConfigT
43
46
  from aiq.data_models.common import TypedBaseModelT
44
47
  from aiq.data_models.component import AIQComponentEnum
45
48
  from aiq.data_models.config import AIQConfig
@@ -52,66 +55,69 @@ from aiq.data_models.front_end import FrontEndBaseConfig
52
55
  from aiq.data_models.front_end import FrontEndConfigT
53
56
  from aiq.data_models.function import FunctionBaseConfig
54
57
  from aiq.data_models.function import FunctionConfigT
58
+ from aiq.data_models.its_strategy import ITSStrategyBaseConfig
59
+ from aiq.data_models.its_strategy import ITSStrategyBaseConfigT
55
60
  from aiq.data_models.llm import LLMBaseConfig
56
61
  from aiq.data_models.llm import LLMBaseConfigT
57
62
  from aiq.data_models.logging import LoggingBaseConfig
58
63
  from aiq.data_models.logging import LoggingMethodConfigT
59
64
  from aiq.data_models.memory import MemoryBaseConfig
60
65
  from aiq.data_models.memory import MemoryBaseConfigT
66
+ from aiq.data_models.object_store import ObjectStoreBaseConfig
67
+ from aiq.data_models.object_store import ObjectStoreBaseConfigT
61
68
  from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
62
69
  from aiq.data_models.registry_handler import RegistryHandlerBaseConfigT
63
70
  from aiq.data_models.retriever import RetrieverBaseConfig
64
71
  from aiq.data_models.retriever import RetrieverBaseConfigT
65
72
  from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
66
73
  from aiq.data_models.telemetry_exporter import TelemetryExporterConfigT
74
+ from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
67
75
  from aiq.memory.interfaces import MemoryEditor
76
+ from aiq.object_store.interfaces import ObjectStore
77
+ from aiq.observability.exporter.base_exporter import BaseExporter
68
78
  from aiq.registry_handlers.registry_handler_base import AbstractRegistryHandler
69
- from aiq.utils.optional_imports import TelemetryOptionalImportError
70
- from aiq.utils.optional_imports import try_import_opentelemetry
71
-
72
- # Try to import OpenTelemetry modules
73
- # If the dependencies are not installed, use a dummy span exporter here
74
- try:
75
- opentelemetry = try_import_opentelemetry()
76
- from opentelemetry.sdk.trace.export import SpanExporter
77
- except TelemetryOptionalImportError:
78
- from aiq.utils.optional_imports import DummySpanExporter # pylint: disable=ungrouped-imports
79
- SpanExporter = DummySpanExporter
80
79
 
81
80
  logger = logging.getLogger(__name__)
82
81
 
82
+ AuthProviderBuildCallableT = Callable[[AuthProviderBaseConfigT, Builder], AsyncIterator[AuthProviderBase]]
83
+ EmbedderClientBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[typing.Any]]
84
+ EmbedderProviderBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[EmbedderProviderInfo]]
85
+ EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIterator[EvaluatorInfo]]
83
86
  FrontEndBuildCallableT = Callable[[FrontEndConfigT, AIQConfig], AsyncIterator[FrontEndBase]]
84
- TelemetryExporterBuildCallableT = Callable[[TelemetryExporterConfigT, Builder], AsyncIterator[SpanExporter]]
85
- LoggingMethodBuildCallableT = Callable[[LoggingMethodConfigT, Builder], AsyncIterator[Handler]]
86
87
  FunctionBuildCallableT = Callable[[FunctionConfigT, Builder], AsyncIterator[FunctionInfo | Callable | FunctionBase]]
87
- LLMProviderBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[LLMProviderInfo]]
88
+ ITSStrategyBuildCallableT = Callable[[ITSStrategyBaseConfigT, Builder], AsyncIterator[StrategyBase]]
88
89
  LLMClientBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[typing.Any]]
89
- EmbedderProviderBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[EmbedderProviderInfo]]
90
- EmbedderClientBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[typing.Any]]
91
- EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIterator[EvaluatorInfo]]
90
+ LLMProviderBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[LLMProviderInfo]]
91
+ LoggingMethodBuildCallableT = Callable[[LoggingMethodConfigT, Builder], AsyncIterator[Handler]]
92
92
  MemoryBuildCallableT = Callable[[MemoryBaseConfigT, Builder], AsyncIterator[MemoryEditor]]
93
- RetrieverProviderBuildCallableT = Callable[[RetrieverBaseConfigT, Builder], AsyncIterator[RetrieverProviderInfo]]
94
- RetrieverClientBuildCallableT = Callable[[RetrieverBaseConfigT, Builder], AsyncIterator[typing.Any]]
93
+ ObjectStoreBuildCallableT = Callable[[ObjectStoreBaseConfigT, Builder], AsyncIterator[ObjectStore]]
95
94
  RegistryHandlerBuildCallableT = Callable[[RegistryHandlerBaseConfigT], AsyncIterator[AbstractRegistryHandler]]
95
+ RetrieverClientBuildCallableT = Callable[[RetrieverBaseConfigT, Builder], AsyncIterator[typing.Any]]
96
+ RetrieverProviderBuildCallableT = Callable[[RetrieverBaseConfigT, Builder], AsyncIterator[RetrieverProviderInfo]]
97
+ TelemetryExporterBuildCallableT = Callable[[TelemetryExporterConfigT, Builder], AsyncIterator[BaseExporter]]
96
98
  ToolWrapperBuildCallableT = Callable[[str, Function, Builder], typing.Any]
97
99
 
98
- TeleExporterRegisteredCallableT = Callable[[TelemetryExporterConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
99
- LoggingMethodRegisteredCallableT = Callable[[LoggingMethodConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
100
+ AuthProviderRegisteredCallableT = Callable[[AuthProviderBaseConfigT, Builder],
101
+ AbstractAsyncContextManager[AuthProviderBase]]
102
+ EmbedderClientRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
103
+ EmbedderProviderRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder],
104
+ AbstractAsyncContextManager[EmbedderProviderInfo]]
105
+ EvaluatorRegisteredCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AbstractAsyncContextManager[EvaluatorInfo]]
100
106
  FrontEndRegisteredCallableT = Callable[[FrontEndConfigT, AIQConfig], AbstractAsyncContextManager[FrontEndBase]]
101
107
  FunctionRegisteredCallableT = Callable[[FunctionConfigT, Builder],
102
108
  AbstractAsyncContextManager[FunctionInfo | Callable | FunctionBase]]
103
- LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
109
+ ITSStrategyRegisterCallableT = Callable[[ITSStrategyBaseConfigT, Builder], AbstractAsyncContextManager[StrategyBase]]
104
110
  LLMClientRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
105
- EmbedderProviderRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder],
106
- AbstractAsyncContextManager[EmbedderProviderInfo]]
107
- EmbedderClientRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
108
- EvaluatorRegisteredCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AbstractAsyncContextManager[EvaluatorInfo]]
111
+ LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
112
+ LoggingMethodRegisteredCallableT = Callable[[LoggingMethodConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
109
113
  MemoryRegisteredCallableT = Callable[[MemoryBaseConfigT, Builder], AbstractAsyncContextManager[MemoryEditor]]
110
- RetrieverProviderRegisteredCallableT = Callable[[RetrieverBaseConfigT, Builder],
111
- AbstractAsyncContextManager[RetrieverProviderInfo]]
112
- RetrieverClientRegisteredCallableT = Callable[[RetrieverBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
114
+ ObjectStoreRegisteredCallableT = Callable[[ObjectStoreBaseConfigT, Builder], AbstractAsyncContextManager[ObjectStore]]
113
115
  RegistryHandlerRegisteredCallableT = Callable[[RegistryHandlerBaseConfigT],
114
116
  AbstractAsyncContextManager[AbstractRegistryHandler]]
117
+ RetrieverClientRegisteredCallableT = Callable[[RetrieverBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
118
+ RetrieverProviderRegisteredCallableT = Callable[[RetrieverBaseConfigT, Builder],
119
+ AbstractAsyncContextManager[RetrieverProviderInfo]]
120
+ TeleExporterRegisteredCallableT = Callable[[TelemetryExporterConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
115
121
 
116
122
 
117
123
  class RegisteredInfo(BaseModel, typing.Generic[TypedBaseModelT]):
@@ -181,6 +187,14 @@ class RegisteredLLMProviderInfo(RegisteredInfo[LLMBaseConfig]):
181
187
  build_fn: LLMProviderRegisteredCallableT = Field(repr=False)
182
188
 
183
189
 
190
+ class RegisteredAuthProviderInfo(RegisteredInfo[AuthProviderBaseConfig]):
191
+ """
192
+ Represents a registered Authentication provider. Authentication providers facilitate the authentication process.
193
+ """
194
+
195
+ build_fn: AuthProviderRegisteredCallableT = Field(repr=False)
196
+
197
+
184
198
  class RegisteredLLMClientInfo(RegisteredInfo[LLMBaseConfig]):
185
199
  """
186
200
  Represents a registered LLM client. LLM Clients are the clients that interact with the LLM providers and are
@@ -226,6 +240,22 @@ class RegisteredMemoryInfo(RegisteredInfo[MemoryBaseConfig]):
226
240
  build_fn: MemoryRegisteredCallableT = Field(repr=False)
227
241
 
228
242
 
243
+ class RegisteredObjectStoreInfo(RegisteredInfo[ObjectStoreBaseConfig]):
244
+ """
245
+ Represents a registered Object Store object which adheres to the object store interface.
246
+ """
247
+
248
+ build_fn: ObjectStoreRegisteredCallableT = Field(repr=False)
249
+
250
+
251
+ class RegisteredITSStrategyInfo(RegisteredInfo[ITSStrategyBaseConfig]):
252
+ """
253
+ Represents a registered Inference Time Scaling (ITS) strategy.
254
+ """
255
+
256
+ build_fn: ITSStrategyRegisterCallableT = Field(repr=False)
257
+
258
+
229
259
  class RegisteredToolWrapper(BaseModel):
230
260
  """
231
261
  Represents a registered tool wrapper. Tool wrappers are used to wrap the functions in a particular LLM framework.
@@ -288,6 +318,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
288
318
  self._llm_client_provider_to_framework: dict[type[LLMBaseConfig], dict[str, RegisteredLLMClientInfo]] = {}
289
319
  self._llm_client_framework_to_provider: dict[str, dict[type[LLMBaseConfig], RegisteredLLMClientInfo]] = {}
290
320
 
321
+ # Authentication
322
+ self._registered_auth_provider_infos: dict[type[AuthProviderBaseConfig], RegisteredAuthProviderInfo] = {}
323
+
291
324
  # Embedders
292
325
  self._registered_embedder_provider_infos: dict[type[EmbedderBaseConfig], RegisteredEmbedderProviderInfo] = {}
293
326
  self._embedder_client_provider_to_framework: dict[type[EmbedderBaseConfig],
@@ -302,6 +335,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
302
335
  # Memory
303
336
  self._registered_memory_infos: dict[type[MemoryBaseConfig], RegisteredMemoryInfo] = {}
304
337
 
338
+ # Object Stores
339
+ self._registered_object_store_infos: dict[type[ObjectStoreBaseConfig], RegisteredObjectStoreInfo] = {}
340
+
305
341
  # Retrievers
306
342
  self._registered_retriever_provider_infos: dict[type[RetrieverBaseConfig], RegisteredRetrieverProviderInfo] = {}
307
343
  self._retriever_client_provider_to_framework: dict[type[RetrieverBaseConfig],
@@ -317,6 +353,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
317
353
  # Tool Wrappers
318
354
  self._registered_tool_wrappers: dict[str, RegisteredToolWrapper] = {}
319
355
 
356
+ # ITS Strategies
357
+ self._registered_its_strategies: dict[type[ITSStrategyBaseConfig], RegisteredITSStrategyInfo] = {}
358
+
320
359
  # Packages
321
360
  self._registered_packages: dict[str, RegisteredPackage] = {}
322
361
 
@@ -458,9 +497,29 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
458
497
  f"Registered configs: {set(self._registered_llm_provider_infos.keys())}") from err
459
498
 
460
499
  def get_registered_llm_providers(self) -> list[RegisteredInfo[LLMBaseConfig]]:
461
-
462
500
  return list(self._registered_llm_provider_infos.values())
463
501
 
502
+ def register_auth_provider(self, info: RegisteredAuthProviderInfo):
503
+
504
+ if (info.config_type in self._registered_auth_provider_infos):
505
+ raise ValueError(
506
+ f"An Authentication Provider with the same config type `{info.config_type}` has already been "
507
+ "registered.")
508
+
509
+ self._registered_auth_provider_infos[info.config_type] = info
510
+
511
+ self._registration_changed()
512
+
513
+ def get_auth_provider(self, config_type: type[AuthProviderBaseConfig]) -> RegisteredAuthProviderInfo:
514
+ try:
515
+ return self._registered_auth_provider_infos[config_type]
516
+ except KeyError as err:
517
+ raise KeyError(f"Could not find a registered Authentication Provider for config `{config_type}`. "
518
+ f"Registered configs: {set(self._registered_auth_provider_infos.keys())}") from err
519
+
520
+ def get_registered_auth_providers(self) -> list[RegisteredInfo[AuthProviderBaseConfig]]:
521
+ return list(self._registered_auth_provider_infos.values())
522
+
464
523
  def register_llm_client(self, info: RegisteredLLMClientInfo):
465
524
 
466
525
  if (info.config_type in self._llm_client_provider_to_framework
@@ -580,6 +639,28 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
580
639
 
581
640
  return list(self._registered_memory_infos.values())
582
641
 
642
+ def register_object_store(self, info: RegisteredObjectStoreInfo):
643
+
644
+ if (info.config_type in self._registered_object_store_infos):
645
+ raise ValueError(f"An Object Store with the same config type `{info.config_type}` has already been "
646
+ "registered.")
647
+
648
+ self._registered_object_store_infos[info.config_type] = info
649
+
650
+ self._registration_changed()
651
+
652
+ def get_object_store(self, config_type: type[ObjectStoreBaseConfig]) -> RegisteredObjectStoreInfo:
653
+
654
+ try:
655
+ return self._registered_object_store_infos[config_type]
656
+ except KeyError as err:
657
+ raise KeyError(f"Could not find a registered Object Store for config `{config_type}`. "
658
+ f"Registered configs: {set(self._registered_object_store_infos.keys())}") from err
659
+
660
+ def get_registered_object_stores(self) -> list[RegisteredInfo[ObjectStoreBaseConfig]]:
661
+
662
+ return list(self._registered_object_store_infos.values())
663
+
583
664
  def register_retriever_provider(self, info: RegisteredRetrieverProviderInfo):
584
665
 
585
666
  if (info.config_type in self._registered_retriever_provider_infos):
@@ -647,6 +728,25 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
647
728
  raise KeyError(f"Could not find a registered tool wrapper for LLM framework `{llm_framework}`. "
648
729
  f"Registered LLM frameworks: {set(self._registered_tool_wrappers.keys())}") from err
649
730
 
731
+ def register_its_strategy(self, info: RegisteredITSStrategyInfo):
732
+ if (info.config_type in self._registered_its_strategies):
733
+ raise ValueError(
734
+ f"An ITS strategy with the same config type `{info.config_type}` has already been registered.")
735
+
736
+ self._registered_its_strategies[info.config_type] = info
737
+
738
+ self._registration_changed()
739
+
740
+ def get_its_strategy(self, config_type: type[ITSStrategyBaseConfig]) -> RegisteredITSStrategyInfo:
741
+ try:
742
+ strategy = self._registered_its_strategies[config_type]
743
+ except Exception as e:
744
+ raise KeyError(f"Could not find a registered ITS strategy for config `{config_type}`. ") from e
745
+ return strategy
746
+
747
+ def get_registered_its_strategies(self) -> list[RegisteredInfo[ITSStrategyBaseConfig]]:
748
+ return list(self._registered_its_strategies.values())
749
+
650
750
  def register_registry_handler(self, info: RegisteredRegistryHandlerInfo):
651
751
 
652
752
  if (info.config_type in self._registered_memory_infos):
@@ -684,6 +784,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
684
784
  if component_type == AIQComponentEnum.FRONT_END:
685
785
  return self._registered_front_end_infos
686
786
 
787
+ if component_type == AIQComponentEnum.AUTHENTICATION_PROVIDER:
788
+ return self._registered_auth_provider_infos
789
+
687
790
  if component_type == AIQComponentEnum.FUNCTION:
688
791
  return self._registered_functions
689
792
 
@@ -726,6 +829,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
726
829
  if component_type == AIQComponentEnum.MEMORY:
727
830
  return self._registered_memory_infos
728
831
 
832
+ if component_type == AIQComponentEnum.OBJECT_STORE:
833
+ return self._registered_object_store_infos
834
+
729
835
  if component_type == AIQComponentEnum.REGISTRY_HANDLER:
730
836
  return self._registered_registry_handler_infos
731
837
 
@@ -738,6 +844,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
738
844
  if component_type == AIQComponentEnum.PACKAGE:
739
845
  return self._registered_packages
740
846
 
847
+ if component_type == AIQComponentEnum.ITS_STRATEGY:
848
+ return self._registered_its_strategies
849
+
741
850
  raise ValueError(f"Supplied an unsupported component type {component_type}")
742
851
 
743
852
  def get_registered_types_by_component_type( # pylint: disable=R0911
@@ -787,6 +896,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
787
896
  if component_type == AIQComponentEnum.PACKAGE:
788
897
  return list(self._registered_packages)
789
898
 
899
+ if component_type == AIQComponentEnum.ITS_STRATEGY:
900
+ return [i.static_type() for i in self._registered_its_strategies]
901
+
790
902
  raise ValueError(f"Supplied an unsupported component type {component_type}")
791
903
 
792
904
  def get_registered_channel_info_by_channel_type(self, channel_type: str) -> RegisteredRegistryHandlerInfo:
@@ -818,6 +930,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
818
930
 
819
931
  def compute_annotation(self, cls: type[TypedBaseModelT]):
820
932
 
933
+ if issubclass(cls, AuthProviderBaseConfig):
934
+ return self._do_compute_annotation(cls, self.get_registered_auth_providers())
935
+
821
936
  if issubclass(cls, EmbedderBaseConfig):
822
937
  return self._do_compute_annotation(cls, self.get_registered_embedder_providers())
823
938
 
@@ -836,6 +951,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
836
951
  if issubclass(cls, MemoryBaseConfig):
837
952
  return self._do_compute_annotation(cls, self.get_registered_memorys())
838
953
 
954
+ if issubclass(cls, ObjectStoreBaseConfig):
955
+ return self._do_compute_annotation(cls, self.get_registered_object_stores())
956
+
839
957
  if issubclass(cls, RegistryHandlerBaseConfig):
840
958
  return self._do_compute_annotation(cls, self.get_registered_registry_handlers())
841
959
 
@@ -848,6 +966,9 @@ class TypeRegistry: # pylint: disable=too-many-public-methods
848
966
  if issubclass(cls, LoggingBaseConfig):
849
967
  return self._do_compute_annotation(cls, self.get_registered_logging_method())
850
968
 
969
+ if issubclass(cls, ITSStrategyBaseConfig):
970
+ return self._do_compute_annotation(cls, self.get_registered_its_strategies())
971
+
851
972
  raise ValueError(f"Supplied an unsupported component type {cls}")
852
973
 
853
974
 
@@ -26,6 +26,7 @@ from pydantic import Discriminator
26
26
  from pydantic import Field
27
27
  from pydantic import HttpUrl
28
28
  from pydantic import conlist
29
+ from pydantic import field_serializer
29
30
  from pydantic import field_validator
30
31
  from pydantic_core.core_schema import ValidationInfo
31
32
 
@@ -111,16 +112,47 @@ class Message(BaseModel):
111
112
  class AIQChatRequest(BaseModel):
112
113
  """
113
114
  AIQChatRequest is a data model that represents a request to the AIQ Toolkit chat API.
115
+ Fully compatible with OpenAI Chat Completions API specification.
114
116
  """
115
117
 
116
- # Allow extra fields in the model_config to support derived models
117
- model_config = ConfigDict(extra="allow")
118
-
118
+ # Required fields
119
119
  messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
120
- model: str | None = None
121
- temperature: float | None = None
122
- max_tokens: int | None = None
123
- top_p: float | None = None
120
+
121
+ # Optional fields (OpenAI Chat Completions API compatible)
122
+ model: str | None = Field(default=None, description="name of the model to use")
123
+ frequency_penalty: float | None = Field(default=0.0,
124
+ description="Penalty for new tokens based on frequency in text")
125
+ logit_bias: dict[str, float] | None = Field(default=None,
126
+ description="Modify likelihood of specified tokens appearing")
127
+ logprobs: bool | None = Field(default=None, description="Whether to return log probabilities")
128
+ top_logprobs: int | None = Field(default=None, description="Number of most likely tokens to return")
129
+ max_tokens: int | None = Field(default=None, description="Maximum number of tokens to generate")
130
+ n: int | None = Field(default=1, description="Number of chat completion choices to generate")
131
+ presence_penalty: float | None = Field(default=0.0, description="Penalty for new tokens based on presence in text")
132
+ response_format: dict[str, typing.Any] | None = Field(default=None, description="Response format specification")
133
+ seed: int | None = Field(default=None, description="Random seed for deterministic sampling")
134
+ service_tier: typing.Literal["auto", "default"] | None = Field(default=None,
135
+ description="Service tier for the request")
136
+ stream: bool | None = Field(default=False, description="Whether to stream partial message deltas")
137
+ stream_options: dict[str, typing.Any] | None = Field(default=None, description="Options for streaming")
138
+ temperature: float | None = Field(default=1.0, description="Sampling temperature between 0 and 2")
139
+ top_p: float | None = Field(default=None, description="Nucleus sampling parameter")
140
+ tools: list[dict[str, typing.Any]] | None = Field(default=None, description="List of tools the model may call")
141
+ tool_choice: str | dict[str, typing.Any] | None = Field(default=None, description="Controls which tool is called")
142
+ parallel_tool_calls: bool | None = Field(default=True, description="Whether to enable parallel function calling")
143
+ user: str | None = Field(default=None, description="Unique identifier representing end-user")
144
+
145
+ model_config = ConfigDict(extra="allow",
146
+ json_schema_extra={
147
+ "example": {
148
+ "model": "nvidia/nemotron",
149
+ "messages": [{
150
+ "role": "user", "content": "who are you?"
151
+ }],
152
+ "temperature": 0.7,
153
+ "stream": False
154
+ }
155
+ })
124
156
 
125
157
  @staticmethod
126
158
  def from_string(data: str,
@@ -156,10 +188,17 @@ class AIQChoiceMessage(BaseModel):
156
188
  role: str | None = None
157
189
 
158
190
 
191
+ class AIQChoiceDelta(BaseModel):
192
+ """Delta object for streaming responses (OpenAI-compatible)"""
193
+ content: str | None = None
194
+ role: str | None = None
195
+
196
+
159
197
  class AIQChoice(BaseModel):
160
198
  model_config = ConfigDict(extra="allow")
161
199
 
162
- message: AIQChoiceMessage
200
+ message: AIQChoiceMessage | None = None
201
+ delta: AIQChoiceDelta | None = None
163
202
  finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None
164
203
  index: int
165
204
  # logprobs: AIQChoiceLogprobs | None = None
@@ -197,16 +236,24 @@ class AIQResponseBaseModelIntermediate(BaseModel, AIQResponseSerializable):
197
236
  class AIQChatResponse(AIQResponseBaseModelOutput):
198
237
  """
199
238
  AIQChatResponse is a data model that represents a response from the AIQ Toolkit chat API.
239
+ Fully compatible with OpenAI Chat Completions API specification.
200
240
  """
201
241
 
202
242
  # Allow extra fields in the model_config to support derived models
203
243
  model_config = ConfigDict(extra="allow")
204
244
  id: str
205
- object: str
245
+ object: str = "chat.completion"
206
246
  model: str = ""
207
247
  created: datetime.datetime
208
248
  choices: list[AIQChoice]
209
249
  usage: AIQUsage | None = None
250
+ system_fingerprint: str | None = None
251
+ service_tier: typing.Literal["scale", "default"] | None = None
252
+
253
+ @field_serializer('created')
254
+ def serialize_created(self, created: datetime.datetime) -> int:
255
+ """Serialize datetime to Unix timestamp for OpenAI compatibility"""
256
+ return int(created.timestamp())
210
257
 
211
258
  @staticmethod
212
259
  def from_string(data: str,
@@ -238,6 +285,7 @@ class AIQChatResponse(AIQResponseBaseModelOutput):
238
285
  class AIQChatResponseChunk(AIQResponseBaseModelOutput):
239
286
  """
240
287
  AIQChatResponseChunk is a data model that represents a response chunk from the AIQ Toolkit chat streaming API.
288
+ Fully compatible with OpenAI Chat Completions API specification.
241
289
  """
242
290
 
243
291
  # Allow extra fields in the model_config to support derived models
@@ -248,6 +296,14 @@ class AIQChatResponseChunk(AIQResponseBaseModelOutput):
248
296
  created: datetime.datetime
249
297
  model: str = ""
250
298
  object: str = "chat.completion.chunk"
299
+ system_fingerprint: str | None = None
300
+ service_tier: typing.Literal["scale", "default"] | None = None
301
+ usage: AIQUsage | None = None
302
+
303
+ @field_serializer('created')
304
+ def serialize_created(self, created: datetime.datetime) -> int:
305
+ """Serialize datetime to Unix timestamp for OpenAI compatibility"""
306
+ return int(created.timestamp())
251
307
 
252
308
  @staticmethod
253
309
  def from_string(data: str,
@@ -273,6 +329,36 @@ class AIQChatResponseChunk(AIQResponseBaseModelOutput):
273
329
  model=model,
274
330
  object=object_)
275
331
 
332
+ @staticmethod
333
+ def create_streaming_chunk(content: str,
334
+ *,
335
+ id_: str | None = None,
336
+ created: datetime.datetime | None = None,
337
+ model: str | None = None,
338
+ role: str | None = None,
339
+ finish_reason: str | None = None,
340
+ usage: AIQUsage | None = None,
341
+ system_fingerprint: str | None = None) -> "AIQChatResponseChunk":
342
+ """Create an OpenAI-compatible streaming chunk"""
343
+ if id_ is None:
344
+ id_ = str(uuid.uuid4())
345
+ if created is None:
346
+ created = datetime.datetime.now(datetime.timezone.utc)
347
+ if model is None:
348
+ model = ""
349
+
350
+ delta = AIQChoiceDelta(content=content,
351
+ role=role) if content is not None or role is not None else AIQChoiceDelta()
352
+
353
+ return AIQChatResponseChunk(
354
+ id=id_,
355
+ choices=[AIQChoice(index=0, message=None, delta=delta, finish_reason=finish_reason)],
356
+ created=created,
357
+ model=model,
358
+ object="chat.completion.chunk",
359
+ usage=usage,
360
+ system_fingerprint=system_fingerprint)
361
+
276
362
 
277
363
  class AIQResponseIntermediateStep(AIQResponseBaseModelIntermediate):
278
364
  """
@@ -563,7 +649,7 @@ GlobalTypeConverter.register_converter(_string_to_aiq_chat_response)
563
649
 
564
650
 
565
651
  def _chat_response_to_chat_response_chunk(data: AIQChatResponse) -> AIQChatResponseChunk:
566
-
652
+ # Preserve original message structure for backward compatibility
567
653
  return AIQChatResponseChunk(id=data.id, choices=data.choices, created=data.created, model=data.model)
568
654
 
569
655
 
@@ -572,7 +658,13 @@ GlobalTypeConverter.register_converter(_chat_response_to_chat_response_chunk)
572
658
 
573
659
  # ======== AIQChatResponseChunk Converters ========
574
660
  def _aiq_chat_response_chunk_to_string(data: AIQChatResponseChunk) -> str:
575
- return data.choices[0].message.content or ""
661
+ if data.choices and len(data.choices) > 0:
662
+ choice = data.choices[0]
663
+ if choice.delta and choice.delta.content:
664
+ return choice.delta.content
665
+ if choice.message and choice.message.content:
666
+ return choice.message.content
667
+ return ""
576
668
 
577
669
 
578
670
  GlobalTypeConverter.register_converter(_aiq_chat_response_chunk_to_string)
@@ -586,3 +678,17 @@ def _string_to_aiq_chat_response_chunk(data: str) -> AIQChatResponseChunk:
586
678
 
587
679
 
588
680
  GlobalTypeConverter.register_converter(_string_to_aiq_chat_response_chunk)
681
+
682
+
683
+ # ======== AINodeMessageChunk Converters ========
684
+ def _ai_message_chunk_to_aiq_chat_response_chunk(data) -> AIQChatResponseChunk:
685
+ '''Converts LangChain AINodeMessageChunk to AIQChatResponseChunk'''
686
+ content = ""
687
+ if hasattr(data, 'content') and data.content is not None:
688
+ content = str(data.content)
689
+ elif hasattr(data, 'text') and data.text is not None:
690
+ content = str(data.text)
691
+ elif hasattr(data, 'message') and data.message is not None:
692
+ content = str(data.message)
693
+
694
+ return AIQChatResponseChunk.create_streaming_chunk(content=content, role="assistant", finish_reason=None)