aiqtoolkit 1.2.0a20250707__py3-none-any.whl → 1.2.0rc1__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 (197) hide show
  1. aiq/agent/base.py +170 -8
  2. aiq/agent/dual_node.py +1 -1
  3. aiq/agent/react_agent/agent.py +112 -111
  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 +123 -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 +11 -2
  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/llm/aws_bedrock_llm.py +2 -1
  126. aiq/llm/nim_llm.py +2 -1
  127. aiq/llm/openai_llm.py +2 -1
  128. aiq/object_store/__init__.py +20 -0
  129. aiq/object_store/in_memory_object_store.py +74 -0
  130. aiq/object_store/interfaces.py +84 -0
  131. aiq/object_store/models.py +36 -0
  132. aiq/object_store/register.py +20 -0
  133. aiq/observability/__init__.py +14 -0
  134. aiq/observability/exporter/__init__.py +14 -0
  135. aiq/observability/exporter/base_exporter.py +449 -0
  136. aiq/observability/exporter/exporter.py +78 -0
  137. aiq/observability/exporter/file_exporter.py +33 -0
  138. aiq/observability/exporter/processing_exporter.py +269 -0
  139. aiq/observability/exporter/raw_exporter.py +52 -0
  140. aiq/observability/exporter/span_exporter.py +264 -0
  141. aiq/observability/exporter_manager.py +335 -0
  142. aiq/observability/mixin/__init__.py +14 -0
  143. aiq/observability/mixin/batch_config_mixin.py +26 -0
  144. aiq/observability/mixin/collector_config_mixin.py +23 -0
  145. aiq/observability/mixin/file_mixin.py +288 -0
  146. aiq/observability/mixin/file_mode.py +23 -0
  147. aiq/observability/mixin/resource_conflict_mixin.py +134 -0
  148. aiq/observability/mixin/serialize_mixin.py +61 -0
  149. aiq/observability/mixin/type_introspection_mixin.py +183 -0
  150. aiq/observability/processor/__init__.py +14 -0
  151. aiq/observability/processor/batching_processor.py +316 -0
  152. aiq/observability/processor/intermediate_step_serializer.py +28 -0
  153. aiq/observability/processor/processor.py +68 -0
  154. aiq/observability/register.py +32 -116
  155. aiq/observability/utils/__init__.py +14 -0
  156. aiq/observability/utils/dict_utils.py +236 -0
  157. aiq/observability/utils/time_utils.py +31 -0
  158. aiq/profiler/calc/__init__.py +14 -0
  159. aiq/profiler/calc/calc_runner.py +623 -0
  160. aiq/profiler/calc/calculations.py +288 -0
  161. aiq/profiler/calc/data_models.py +176 -0
  162. aiq/profiler/calc/plot.py +345 -0
  163. aiq/profiler/data_models.py +2 -0
  164. aiq/profiler/profile_runner.py +16 -13
  165. aiq/runtime/loader.py +8 -2
  166. aiq/runtime/runner.py +23 -9
  167. aiq/runtime/session.py +16 -5
  168. aiq/tool/chat_completion.py +74 -0
  169. aiq/tool/code_execution/README.md +152 -0
  170. aiq/tool/code_execution/code_sandbox.py +151 -72
  171. aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
  172. aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +139 -24
  173. aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +3 -1
  174. aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +27 -2
  175. aiq/tool/code_execution/register.py +7 -3
  176. aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
  177. aiq/tool/mcp/exceptions.py +142 -0
  178. aiq/tool/mcp/mcp_client.py +17 -3
  179. aiq/tool/mcp/mcp_tool.py +1 -1
  180. aiq/tool/register.py +1 -0
  181. aiq/tool/server_tools.py +2 -2
  182. aiq/utils/exception_handlers/automatic_retries.py +289 -0
  183. aiq/utils/exception_handlers/mcp.py +211 -0
  184. aiq/utils/io/model_processing.py +28 -0
  185. aiq/utils/log_utils.py +37 -0
  186. aiq/utils/string_utils.py +38 -0
  187. aiq/utils/type_converter.py +18 -2
  188. aiq/utils/type_utils.py +87 -0
  189. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/METADATA +37 -9
  190. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/RECORD +195 -80
  191. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/entry_points.txt +3 -0
  192. aiq/front_ends/fastapi/websocket.py +0 -153
  193. aiq/observability/async_otel_listener.py +0 -470
  194. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/WHEEL +0 -0
  195. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  196. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/licenses/LICENSE.md +0 -0
  197. {aiqtoolkit-1.2.0a20250707.dist-info → aiqtoolkit-1.2.0rc1.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,316 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import asyncio
17
+ import logging
18
+ import time
19
+ from collections import deque
20
+ from collections.abc import Awaitable
21
+ from collections.abc import Callable
22
+ from typing import Any
23
+ from typing import Generic
24
+ from typing import TypeVar
25
+
26
+ from aiq.observability.processor.processor import Processor
27
+
28
+ logger = logging.getLogger(__name__)
29
+
30
+ T = TypeVar('T')
31
+
32
+
33
+ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
34
+ """Pass-through batching processor that accumulates items and outputs batched lists.
35
+
36
+ This processor fits properly into the generics design by implementing Processor[T, List[T]].
37
+ It accumulates individual items and outputs them as batches when size or time thresholds
38
+ are met. The batched output continues through the processing pipeline.
39
+
40
+ CRITICAL: Implements proper cleanup to ensure NO ITEMS ARE LOST during shutdown.
41
+ The ProcessingExporter._cleanup() method calls shutdown() on all processors.
42
+
43
+ Key Features:
44
+ - Pass-through design: Processor[T, List[T]]
45
+ - Size-based and time-based batching
46
+ - Fits into generics processing pipeline design
47
+ - GUARANTEED: No items lost during cleanup
48
+ - Comprehensive statistics and monitoring
49
+ - Proper cleanup and shutdown handling
50
+ - High-performance async implementation
51
+ - Back-pressure handling with queue limits
52
+
53
+ Cleanup Guarantee:
54
+ When ProcessingExporter._cleanup() calls shutdown(), this processor:
55
+ 1. Stops accepting new items
56
+ 2. Processes all queued items as final batch
57
+ 3. Returns final batch to continue through pipeline
58
+ 4. Ensures zero data loss during shutdown
59
+
60
+ Usage in Pipeline:
61
+ ```python
62
+ # Individual spans → Batched spans → Continue processing
63
+ exporter.add_processor(BatchingProcessor[Span](batch_size=100))
64
+ exporter.add_processor(BatchedSpanProcessor()) # Processes List[Span]
65
+ ```
66
+
67
+ Args:
68
+ batch_size: Maximum items per batch (default: 100)
69
+ flush_interval: Max seconds to wait before flushing (default: 5.0)
70
+ max_queue_size: Maximum items to queue before blocking (default: 1000)
71
+ drop_on_overflow: If True, drop items when queue is full (default: False)
72
+ shutdown_timeout: Max seconds to wait for final batch processing (default: 10.0)
73
+ """
74
+
75
+ def __init__(self,
76
+ batch_size: int = 100,
77
+ flush_interval: float = 5.0,
78
+ max_queue_size: int = 1000,
79
+ drop_on_overflow: bool = False,
80
+ shutdown_timeout: float = 10.0,
81
+ done_callback: Callable[[list[T]], Awaitable[None]] | None = None):
82
+ self._batch_size = batch_size
83
+ self._flush_interval = flush_interval
84
+ self._max_queue_size = max_queue_size
85
+ self._drop_on_overflow = drop_on_overflow
86
+ self._shutdown_timeout = shutdown_timeout
87
+ self._done_callback = done_callback
88
+
89
+ # Batching state
90
+ self._batch_queue: deque[T] = deque()
91
+ self._last_flush_time = time.time()
92
+ self._flush_task: asyncio.Task | None = None
93
+ self._batch_lock = asyncio.Lock()
94
+ self._shutdown_requested = False
95
+ self._shutdown_complete = False
96
+ self._shutdown_complete_event: asyncio.Event | None = None
97
+
98
+ # Final batch handling for cleanup
99
+ self._final_batch: list[T] | None = None
100
+ self._final_batch_processed = False
101
+
102
+ # Callback for immediate export of scheduled batches
103
+ self._done = None
104
+
105
+ # Statistics
106
+ self._batches_created = 0
107
+ self._items_processed = 0
108
+ self._items_dropped = 0
109
+ self._queue_overflows = 0
110
+ self._shutdown_batches = 0
111
+
112
+ async def process(self, item: T) -> list[T]:
113
+ """Process an item by adding it to the batch queue.
114
+
115
+ Returns a batch when batching conditions are met, otherwise returns empty list.
116
+ This maintains the Processor[T, List[T]] contract while handling batching logic.
117
+
118
+ During shutdown, immediately returns items as single-item batches to ensure
119
+ no data loss.
120
+
121
+ Args:
122
+ item: The item to add to the current batch
123
+
124
+ Returns:
125
+ List[T]: A batch of items when ready, empty list otherwise
126
+ """
127
+ if self._shutdown_requested:
128
+ # During shutdown, return item immediately as single-item batch
129
+ # This ensures no items are lost even if shutdown is in progress
130
+ self._items_processed += 1
131
+ self._shutdown_batches += 1
132
+ logger.debug("Shutdown mode: returning single-item batch for item %s", item)
133
+ return [item]
134
+
135
+ async with self._batch_lock:
136
+ # Handle queue overflow
137
+ if len(self._batch_queue) >= self._max_queue_size:
138
+ self._queue_overflows += 1
139
+
140
+ if self._drop_on_overflow:
141
+ # Drop the item and return empty
142
+ self._items_dropped += 1
143
+ logger.warning("Dropping item due to queue overflow (dropped: %d)", self._items_dropped)
144
+ return []
145
+ # Force flush to make space, then add item
146
+ logger.warning("Queue overflow, forcing flush of %d items", len(self._batch_queue))
147
+ forced_batch = await self._create_batch()
148
+ if forced_batch:
149
+ # Add current item to queue and return the forced batch
150
+ self._batch_queue.append(item)
151
+ self._items_processed += 1
152
+ return forced_batch
153
+
154
+ # Add item to batch queue
155
+ self._batch_queue.append(item)
156
+ self._items_processed += 1
157
+
158
+ # Check flush conditions
159
+ should_flush = (len(self._batch_queue) >= self._batch_size
160
+ or (time.time() - self._last_flush_time) >= self._flush_interval)
161
+
162
+ if should_flush:
163
+ return await self._create_batch()
164
+ # Schedule a time-based flush if not already scheduled
165
+ if self._flush_task is None or self._flush_task.done():
166
+ self._flush_task = asyncio.create_task(self._schedule_flush())
167
+ return []
168
+
169
+ def set_done_callback(self, callback: Callable[[list[T]], Awaitable[None]]):
170
+ """Set callback function for immediate export of scheduled batches."""
171
+ self._done_callback = callback
172
+
173
+ async def _schedule_flush(self):
174
+ """Schedule a flush after the flush interval."""
175
+ try:
176
+ await asyncio.sleep(self._flush_interval)
177
+ async with self._batch_lock:
178
+ if not self._shutdown_requested and self._batch_queue:
179
+ batch = await self._create_batch()
180
+ if batch:
181
+ # Immediately export scheduled batches via callback
182
+ if self._done_callback is not None:
183
+ try:
184
+ await self._done_callback(batch)
185
+ logger.debug("Scheduled flush exported batch of %d items", len(batch))
186
+ except Exception as e:
187
+ logger.error("Error exporting scheduled batch: %s", e, exc_info=True)
188
+ else:
189
+ logger.warning("Scheduled flush created batch of %d items but no export callback set",
190
+ len(batch))
191
+ except asyncio.CancelledError:
192
+ pass
193
+ except Exception as e:
194
+ logger.error("Error in scheduled flush: %s", e, exc_info=True)
195
+
196
+ async def _create_batch(self) -> list[T]:
197
+ """Create a batch from the current queue."""
198
+ if not self._batch_queue:
199
+ return []
200
+
201
+ batch = list(self._batch_queue)
202
+ self._batch_queue.clear()
203
+ self._last_flush_time = time.time()
204
+ self._batches_created += 1
205
+
206
+ logger.debug("Created batch of %d items (total: %d items in %d batches)",
207
+ len(batch),
208
+ self._items_processed,
209
+ self._batches_created)
210
+
211
+ return batch
212
+
213
+ async def force_flush(self) -> list[T]:
214
+ """Force an immediate flush of all queued items.
215
+
216
+ Returns:
217
+ List[T]: The current batch, empty list if no items queued
218
+ """
219
+ async with self._batch_lock:
220
+ return await self._create_batch()
221
+
222
+ async def shutdown(self) -> None:
223
+ """Shutdown the processor and ensure all items are processed.
224
+
225
+ CRITICAL: This method is called by ProcessingExporter._cleanup() to ensure
226
+ no items are lost during shutdown. It creates a final batch from any
227
+ remaining items and stores it for processing.
228
+
229
+ The final batch will be processed by the next process() call or can be
230
+ retrieved via get_final_batch().
231
+ """
232
+ if self._shutdown_requested:
233
+ logger.debug("Shutdown already requested, waiting for completion")
234
+ # Wait for shutdown to complete using event instead of polling
235
+ try:
236
+ await asyncio.wait_for(self._shutdown_complete_event.wait(), timeout=self._shutdown_timeout)
237
+ logger.debug("Shutdown completion detected via event")
238
+ except asyncio.TimeoutError:
239
+ logger.warning("Shutdown completion timeout exceeded (%s seconds)", self._shutdown_timeout)
240
+ return
241
+
242
+ logger.info("Starting shutdown of BatchingProcessor (queue size: %d)", len(self._batch_queue))
243
+ self._shutdown_requested = True
244
+
245
+ try:
246
+ # Cancel scheduled flush task
247
+ if self._flush_task and not self._flush_task.done():
248
+ self._flush_task.cancel()
249
+ try:
250
+ await self._flush_task
251
+ except asyncio.CancelledError:
252
+ pass
253
+
254
+ # Create final batch from remaining items
255
+ async with self._batch_lock:
256
+ if self._batch_queue:
257
+ self._final_batch = await self._create_batch()
258
+ logger.info("Created final batch of %d items during shutdown", len(self._final_batch))
259
+ else:
260
+ self._final_batch = []
261
+ logger.info("No items remaining during shutdown")
262
+
263
+ self._shutdown_complete = True
264
+ self._shutdown_complete_event.set()
265
+ logger.info("BatchingProcessor shutdown completed successfully")
266
+
267
+ except Exception as e:
268
+ logger.error("Error during BatchingProcessor shutdown: %s", e, exc_info=True)
269
+ self._shutdown_complete = True
270
+ self._shutdown_complete_event.set()
271
+
272
+ def get_final_batch(self) -> list[T]:
273
+ """Get the final batch created during shutdown.
274
+
275
+ This method allows the exporter to retrieve and process any items
276
+ that were queued when shutdown was called.
277
+
278
+ Returns:
279
+ List[T]: Final batch of items, empty list if none
280
+ """
281
+ if self._final_batch is not None:
282
+ final_batch = self._final_batch
283
+ self._final_batch = None # Clear to avoid double processing
284
+ self._final_batch_processed = True
285
+ return final_batch
286
+ return []
287
+
288
+ def has_final_batch(self) -> bool:
289
+ """Check if there's a final batch waiting to be processed.
290
+
291
+ Returns:
292
+ bool: True if final batch exists and hasn't been processed
293
+ """
294
+ return self._final_batch is not None and not self._final_batch_processed
295
+
296
+ def get_stats(self) -> dict[str, Any]:
297
+ """Get comprehensive batching statistics."""
298
+ return {
299
+ "current_queue_size": len(self._batch_queue),
300
+ "batch_size_limit": self._batch_size,
301
+ "flush_interval": self._flush_interval,
302
+ "max_queue_size": self._max_queue_size,
303
+ "drop_on_overflow": self._drop_on_overflow,
304
+ "shutdown_timeout": self._shutdown_timeout,
305
+ "batches_created": self._batches_created,
306
+ "items_processed": self._items_processed,
307
+ "items_dropped": self._items_dropped,
308
+ "queue_overflows": self._queue_overflows,
309
+ "shutdown_batches": self._shutdown_batches,
310
+ "shutdown_requested": self._shutdown_requested,
311
+ "shutdown_complete": self._shutdown_complete,
312
+ "final_batch_size": len(self._final_batch) if self._final_batch else 0,
313
+ "final_batch_processed": self._final_batch_processed,
314
+ "avg_items_per_batch": self._items_processed / max(1, self._batches_created),
315
+ "drop_rate": self._items_dropped / max(1, self._items_processed) * 100 if self._items_processed > 0 else 0
316
+ }
@@ -0,0 +1,28 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from aiq.data_models.intermediate_step import IntermediateStep
17
+ from aiq.observability.mixin.serialize_mixin import SerializeMixin
18
+ from aiq.observability.processor.processor import Processor
19
+ from aiq.utils.type_utils import override
20
+
21
+
22
+ class IntermediateStepSerializer(SerializeMixin, Processor[IntermediateStep, str]):
23
+ """A File processor that exports telemetry traces to a local file."""
24
+
25
+ @override
26
+ async def process(self, item: IntermediateStep) -> str:
27
+ serialized_payload, _ = self._serialize_payload(item)
28
+ return serialized_payload
@@ -0,0 +1,68 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from abc import ABC
17
+ from abc import abstractmethod
18
+ from typing import Generic
19
+ from typing import TypeVar
20
+
21
+ from aiq.observability.mixin.type_introspection_mixin import TypeIntrospectionMixin
22
+
23
+ InputT = TypeVar('InputT')
24
+ OutputT = TypeVar('OutputT')
25
+
26
+
27
+ class Processor(Generic[InputT, OutputT], TypeIntrospectionMixin, ABC):
28
+ """Generic protocol for processors that can convert between types in export pipelines.
29
+
30
+ Processors are the building blocks of processing pipelines in exporters. They can
31
+ transform data from one type to another, enabling flexible data processing chains.
32
+
33
+ The generic types work as follows:
34
+ - InputT: The type of items that this processor accepts
35
+ - OutputT: The type of items that this processor produces
36
+
37
+ Key Features:
38
+ - Type-safe transformations through generics
39
+ - Type introspection capabilities via TypeIntrospectionMixin
40
+ - Async processing support
41
+ - Chainable in processing pipelines
42
+
43
+ Inheritance Structure:
44
+ - Inherits from TypeIntrospectionMixin for type introspection capabilities
45
+ - Implements Generic[InputT, OutputT] for type safety
46
+ - Abstract base class requiring implementation of process()
47
+
48
+ Example:
49
+ .. code-block:: python
50
+
51
+ class SpanToOtelProcessor(Processor[Span, OtelSpan]):
52
+ async def process(self, item: Span) -> OtelSpan:
53
+ return convert_span_to_otel(item)
54
+
55
+ Note:
56
+ Processors are typically added to ProcessingExporter instances to create
57
+ transformation pipelines. The exporter validates type compatibility between
58
+ chained processors.
59
+ """
60
+
61
+ @abstractmethod
62
+ async def process(self, item: InputT) -> OutputT:
63
+ """Process an item and return a potentially different type.
64
+
65
+ Args:
66
+ item: The item to process
67
+ """
68
+ pass
@@ -14,7 +14,6 @@
14
14
  # limitations under the License.
15
15
 
16
16
  import logging
17
- import os
18
17
 
19
18
  from pydantic import Field
20
19
 
@@ -23,112 +22,53 @@ from aiq.cli.register_workflow import register_logging_method
23
22
  from aiq.cli.register_workflow import register_telemetry_exporter
24
23
  from aiq.data_models.logging import LoggingBaseConfig
25
24
  from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
26
- from aiq.utils.optional_imports import telemetry_optional_import
27
- from aiq.utils.optional_imports import try_import_phoenix
25
+ from aiq.observability.mixin.file_mode import FileMode
28
26
 
29
27
  logger = logging.getLogger(__name__)
30
28
 
31
29
 
32
- class PhoenixTelemetryExporter(TelemetryExporterBaseConfig, name="phoenix"):
33
- """A telemetry exporter to transmit traces to externally hosted phoenix service."""
30
+ class FileTelemetryExporterConfig(TelemetryExporterBaseConfig, name="file"):
31
+ """A telemetry exporter that writes runtime traces to local files with optional rolling."""
34
32
 
35
- endpoint: str = Field(description="The phoenix endpoint to export telemetry traces.")
36
- project: str = Field(description="The project name to group the telemetry traces.")
33
+ output_path: str = Field(description="Output path for logs. When rolling is disabled: exact file path. "
34
+ "When rolling is enabled: directory path or file path (directory + base name).")
35
+ project: str = Field(description="Name to affiliate with this application.")
36
+ mode: FileMode = Field(
37
+ default=FileMode.APPEND,
38
+ description="File write mode: 'append' to add to existing file or 'overwrite' to start fresh.")
39
+ enable_rolling: bool = Field(default=False, description="Enable rolling log files based on size limits.")
40
+ max_file_size: int = Field(
41
+ default=10 * 1024 * 1024, # 10MB
42
+ description="Maximum file size in bytes before rolling to a new file.")
43
+ max_files: int = Field(default=5, description="Maximum number of rolled files to keep.")
44
+ cleanup_on_init: bool = Field(default=False, description="Clean up old files during initialization.")
37
45
 
38
46
 
39
- @register_telemetry_exporter(config_type=PhoenixTelemetryExporter)
40
- async def phoenix_telemetry_exporter(config: PhoenixTelemetryExporter, builder: Builder):
41
- """Create a Phoenix telemetry exporter."""
42
- try:
43
- # If the dependencies are not installed, a TelemetryOptionalImportError will be raised
44
- phoenix = try_import_phoenix() # noqa: F841
45
- from phoenix.otel import HTTPSpanExporter
46
-
47
- yield HTTPSpanExporter(config.endpoint)
48
- except ConnectionError as ex:
49
- logger.warning(
50
- "Unable to connect to Phoenix at port 6006. Are you sure Phoenix is running?\n %s",
51
- ex,
52
- exc_info=True,
53
- )
54
-
55
-
56
- class LangfuseTelemetryExporter(TelemetryExporterBaseConfig, name="langfuse"):
57
- """A telemetry exporter to transmit traces to externally hosted langfuse service."""
58
-
59
- endpoint: str = Field(description="The langfuse OTEL endpoint (/api/public/otel/v1/traces)")
60
- public_key: str = Field(description="The Langfuse public key", default="")
61
- secret_key: str = Field(description="The Langfuse secret key", default="")
62
-
63
-
64
- @register_telemetry_exporter(config_type=LangfuseTelemetryExporter)
65
- async def langfuse_telemetry_exporter(config: LangfuseTelemetryExporter, builder: Builder):
66
- """Create a Langfuse telemetry exporter."""
67
-
68
- import base64
69
-
70
- trace_exporter = telemetry_optional_import("opentelemetry.exporter.otlp.proto.http.trace_exporter")
71
-
72
- secret_key = config.secret_key or os.environ.get("LANGFUSE_SECRET_KEY")
73
- public_key = config.public_key or os.environ.get("LANGFUSE_PUBLIC_KEY")
74
- if not secret_key or not public_key:
75
- raise ValueError("secret and public keys are required for langfuse")
76
-
77
- credentials = f"{public_key}:{secret_key}".encode("utf-8")
78
- auth_header = base64.b64encode(credentials).decode("utf-8")
79
- headers = {"Authorization": f"Basic {auth_header}"}
80
-
81
- yield trace_exporter.OTLPSpanExporter(endpoint=config.endpoint, headers=headers)
82
-
83
-
84
- class LangsmithTelemetryExporter(TelemetryExporterBaseConfig, name="langsmith"):
85
- """A telemetry exporter to transmit traces to externally hosted langsmith service."""
86
-
87
- endpoint: str = Field(
88
- description="The langsmith OTEL endpoint",
89
- default="https://api.smith.langchain.com/otel/v1/traces",
90
- )
91
- api_key: str = Field(description="The Langsmith API key", default="")
92
- project: str = Field(description="The project name to group the telemetry traces.")
93
-
94
-
95
- @register_telemetry_exporter(config_type=LangsmithTelemetryExporter)
96
- async def langsmith_telemetry_exporter(config: LangsmithTelemetryExporter, builder: Builder):
97
- """Create a Langsmith telemetry exporter."""
98
-
99
- trace_exporter = telemetry_optional_import("opentelemetry.exporter.otlp.proto.http.trace_exporter")
100
-
101
- api_key = config.api_key or os.environ.get("LANGSMITH_API_KEY")
102
- if not api_key:
103
- raise ValueError("API key is required for langsmith")
104
-
105
- headers = {"x-api-key": api_key, "Langsmith-Project": config.project}
106
- yield trace_exporter.OTLPSpanExporter(endpoint=config.endpoint, headers=headers)
107
-
108
-
109
- class OtelCollectorTelemetryExporter(TelemetryExporterBaseConfig, name="otelcollector"):
110
- """A telemetry exporter to transmit traces to externally hosted otel collector service."""
111
-
112
- endpoint: str = Field(description="The otel endpoint to export telemetry traces.")
113
- project: str = Field(description="The project name to group the telemetry traces.")
114
-
47
+ @register_telemetry_exporter(config_type=FileTelemetryExporterConfig)
48
+ async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder): # pylint: disable=W0613
49
+ """
50
+ Build and return a FileExporter for file-based telemetry export with optional rolling.
51
+ """
115
52
 
116
- @register_telemetry_exporter(config_type=OtelCollectorTelemetryExporter)
117
- async def otel_telemetry_exporter(config: OtelCollectorTelemetryExporter, builder: Builder):
118
- """Create an OpenTelemetry telemetry exporter."""
53
+ from aiq.observability.exporter.file_exporter import FileExporter
119
54
 
120
- trace_exporter = telemetry_optional_import("opentelemetry.exporter.otlp.proto.http.trace_exporter")
121
- yield trace_exporter.OTLPSpanExporter(endpoint=config.endpoint)
55
+ yield FileExporter(output_path=config.output_path,
56
+ project=config.project,
57
+ mode=config.mode,
58
+ enable_rolling=config.enable_rolling,
59
+ max_file_size=config.max_file_size,
60
+ max_files=config.max_files,
61
+ cleanup_on_init=config.cleanup_on_init)
122
62
 
123
63
 
124
- class ConsoleLoggingMethod(LoggingBaseConfig, name="console"):
64
+ class ConsoleLoggingMethodConfig(LoggingBaseConfig, name="console"):
125
65
  """A logger to write runtime logs to the console."""
126
66
 
127
67
  level: str = Field(description="The logging level of console logger.")
128
68
 
129
69
 
130
- @register_logging_method(config_type=ConsoleLoggingMethod)
131
- async def console_logging_method(config: ConsoleLoggingMethod, builder: Builder):
70
+ @register_logging_method(config_type=ConsoleLoggingMethodConfig)
71
+ async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder): # pylint: disable=W0613
132
72
  """
133
73
  Build and return a StreamHandler for console-based logging.
134
74
  """
@@ -146,7 +86,7 @@ class FileLoggingMethod(LoggingBaseConfig, name="file"):
146
86
 
147
87
 
148
88
  @register_logging_method(config_type=FileLoggingMethod)
149
- async def file_logging_method(config: FileLoggingMethod, builder: Builder):
89
+ async def file_logging_method(config: FileLoggingMethod, builder: Builder): # pylint: disable=W0613
150
90
  """
151
91
  Build and return a FileHandler for file-based logging.
152
92
  """
@@ -154,27 +94,3 @@ async def file_logging_method(config: FileLoggingMethod, builder: Builder):
154
94
  handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
155
95
  handler.setLevel(level)
156
96
  yield handler
157
-
158
-
159
- class PatronusTelemetryExporter(TelemetryExporterBaseConfig, name="patronus"):
160
- """A telemetry exporter to transmit traces to Patronus service."""
161
-
162
- endpoint: str = Field(description="The Patronus OTEL endpoint")
163
- api_key: str = Field(description="The Patronus API key", default="")
164
- project: str = Field(description="The project name to group the telemetry traces.")
165
-
166
-
167
- @register_telemetry_exporter(config_type=PatronusTelemetryExporter)
168
- async def patronus_telemetry_exporter(config: PatronusTelemetryExporter, builder: Builder):
169
- """Create a Patronus telemetry exporter."""
170
- trace_exporter = telemetry_optional_import("opentelemetry.exporter.otlp.proto.grpc.trace_exporter")
171
-
172
- api_key = config.api_key or os.environ.get("PATRONUS_API_KEY")
173
- if not api_key:
174
- raise ValueError("API key is required for Patronus")
175
-
176
- headers = {
177
- "x-api-key": api_key,
178
- "pat-project-name": config.project,
179
- }
180
- yield trace_exporter.OTLPSpanExporter(endpoint=config.endpoint, headers=headers)
@@ -0,0 +1,14 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.