aiqtoolkit 1.2.0a20250617__py3-none-any.whl → 1.2.0a20250618__py3-none-any.whl

aiq/builder/context.py

@@ -61,6 +61,7 @@ class ActiveFunctionContextManager:
 class AIQContextState(metaclass=Singleton):
 
     def __init__(self):
+        self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
         self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
@@ -148,6 +149,16 @@ class AIQContext:
         """
         return IntermediateStepManager(self._context_state)
 
+    @property
+    def conversation_id(self) -> str | None:
+        """
+        This property retrieves the conversation ID which is the unique identifier for the current chat conversation.
+
+        Returns:
+            str | None
+        """
+        return self._context_state.conversation_id.get()
+
     @contextmanager
     def push_active_function(self, function_name: str, input_data: typing.Any | None):
         """

aiq/data_models/api_server.py

@@ -392,7 +392,7 @@ class WebSocketUserMessage(BaseModel):
     type: typing.Literal[WebSocketMessageType.USER_MESSAGE]
     schema_type: WorkflowSchemaType
     id: str = "default"
-    thread_id: str = "default"
+    conversation_id: str | None = None
     content: UserMessageContent
     user: User = User()
     security: Security = Security()

aiq/front_ends/fastapi/message_handler.py

@@ -130,28 +130,29 @@ class MessageHandler:
 
         return None
 
-    async def process_user_message(self, message_as_validated_type: WebSocketUserMessage) -> None:
+    async def process_user_message(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
         """
         Process user messages and routes them appropriately.
 
-        :param message_as_validated_type: A WebSocketUserMessage Data Model instance.
+        :param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
         """
 
         try:
-            self._message_parent_id = message_as_validated_type.id
-            self._workflow_schema_type = message_as_validated_type.schema_type
+            self._message_parent_id = user_message_as_validated_type.id
+            self._workflow_schema_type = user_message_as_validated_type.schema_type
+            conversation_id: str = user_message_as_validated_type.conversation_id
 
-            content: BaseModel | None = await self.process_user_message_content(message_as_validated_type)
+            content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
 
             if content is None:
-                raise ValueError(f"User message content could not be found: {message_as_validated_type}")
+                raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
 
             if isinstance(content, TextContent) and (self._background_task is None):
 
                 await self._process_response()
                 self._background_task = asyncio.create_task(
                     self._websocket_reference.workflow_schema_type.get(self._workflow_schema_type)(
-                        content.text)).add_done_callback(
+                        content.text, conversation_id)).add_done_callback(
                             lambda task: asyncio.create_task(self._on_process_stream_task_done(task)))
 
         except ValueError as e:

aiq/front_ends/fastapi/websocket.py

@@ -110,10 +110,12 @@ class AIQWebSocket(WebSocketEndpoint):
 
     async def _process_message(self,
                                payload: typing.Any,
+                               conversation_id: str | None = None,
                                result_type: type | None = None,
                                output_type: type | None = None) -> None:
 
         async with self._session_manager.session(
+                conversation_id=conversation_id,
                 user_input_callback=self._message_handler.human_interaction) as session:
 
             async for value in generate_streaming_response(payload,
@@ -131,18 +133,21 @@ class AIQWebSocket(WebSocketEndpoint):
                 await self._message_handler.create_websocket_message(data_model=value,
                                                                      status=WebSocketMessageStatus.IN_PROGRESS)
 
-    async def process_generate_stream(self, payload: str):
+    async def process_generate_stream(self, payload: str, conversation_id: str) -> None:
 
-        return await self._process_message(payload, result_type=None, output_type=None)
+        return await self._process_message(payload, conversation_id=conversation_id, result_type=None, output_type=None)
 
-    async def process_chat_stream(self, payload: AIQChatRequest):
+    async def process_chat_stream(self, payload: AIQChatRequest, conversation_id: str):
 
-        return await self._process_message(payload, result_type=AIQChatResponse, output_type=AIQChatResponseChunk)
+        return await self._process_message(payload,
+                                           conversation_id=conversation_id,
+                                           result_type=AIQChatResponse,
+                                           output_type=AIQChatResponseChunk)
 
-    async def process_generate(self, payload: typing.Any):
+    async def process_generate(self, payload: typing.Any, conversation_id: str):
 
-        return await self._process_message(payload)
+        return await self._process_message(payload, conversation_id=conversation_id)
 
-    async def process_chat(self, payload: AIQChatRequest):
+    async def process_chat(self, payload: AIQChatRequest, conversation_id: str):
 
-        return await self._process_message(payload, result_type=AIQChatResponse)
+        return await self._process_message(payload, conversation_id=conversation_id, result_type=AIQChatResponse)

aiq/runtime/session.py

@@ -86,7 +86,8 @@ class AIQSessionManager:
     @asynccontextmanager
     async def session(self,
                       user_manager=None,
-                      request: Request = None,
+                      request: Request | None = None,
+                      conversation_id: str | None = None,
                       user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None):
 
         token_user_input = None
@@ -97,7 +98,10 @@ class AIQSessionManager:
         if user_manager is not None:
             token_user_manager = self._context_state.user_manager.set(user_manager)
 
-        self.set_request_attributes(request)
+        if conversation_id is not None and request is None:
+            self._context_state.conversation_id.set(conversation_id)
+
+        self.set_metadata_from_http_request(request)
 
         try:
             yield self
@@ -120,9 +124,9 @@ class AIQSessionManager:
             async with self._workflow.run(message) as runner:
                 yield runner
 
-    def set_request_attributes(self, request: Request) -> None:
+    def set_metadata_from_http_request(self, request: Request | None) -> None:
         """
-        Extracts and sets request attributes from an HTTP request.
+        Extracts and sets user metadata request attributes from a HTTP request.
         If request is None, no attributes are set.
         """
         if request is None:
@@ -138,3 +142,6 @@ class AIQSessionManager:
         self._context.metadata._request.client_host = request.client.host
         self._context.metadata._request.client_port = request.client.port
         self._context.metadata._request.cookies = request.cookies
+
+        if request.headers.get("conversation-id"):
+            self._context_state.conversation_id.set(request.headers["conversation-id"])

aiq/runtime/user_metadata.py

@@ -21,9 +21,8 @@ from aiq.data_models.api_server import Request
 
 class RequestAttributes:
     """
-    The RequestAttributes class is responsible for managing user-defined
-    metadata and attributes. It provides a way to store and
-    expose user-defined attributes to workflow tools.
+    The RequestAttributes class is responsible for managing user http and webscoket session
+    metadata. It provides a way to store and expose session attributes to workflow tools.
     """
 
     def __init__(self) -> None:

aiq/tool/server_tools.py

@@ -39,6 +39,7 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
 
         from aiq.builder.context import AIQContext
         aiq_context = AIQContext.get()
+
         method: str | None = aiq_context.metadata.method
         url_path: str | None = aiq_context.metadata.url_path
         url_scheme: str | None = aiq_context.metadata.url_scheme
@@ -48,6 +49,7 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
         client_host: str | None = aiq_context.metadata.client_host
         client_port: int | None = aiq_context.metadata.client_port
         cookies: dict[str, str] | None = aiq_context.metadata.cookies
+        conversation_id: str | None = aiq_context.conversation_id
 
         return (f"Method: {method}, "
                 f"URL Path: {url_path}, "
@@ -57,7 +59,8 @@ async def current_request_attributes(config: RequestAttributesTool, builder: Bui
                 f"Path Params: {path_params}, "
                 f"Client Host: {client_host}, "
                 f"Client Port: {client_port}, "
-                f"Cookies: {cookies}")
+                f"Cookies: {cookies}, "
+                f"Conversation Id: {conversation_id}")
 
     yield FunctionInfo.from_fn(_get_request_attributes,
                                description="Returns the acquired user defined request attriubutes.")

{aiqtoolkit-1.2.0a20250617.dist-info → aiqtoolkit-1.2.0a20250618.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiqtoolkit
-Version: 1.2.0a20250617
+Version: 1.2.0a20250618
 Summary: NVIDIA Agent Intelligence toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation

{aiqtoolkit-1.2.0a20250617.dist-info → aiqtoolkit-1.2.0a20250618.dist-info}/RECORD

@@ -19,7 +19,7 @@ aiq/agent/tool_calling_agent/register.py,sha256=kqcN2uovVBQxrIx5MszBS65opbhBrCRl
 aiq/builder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/builder/builder.py,sha256=JE3uCw6qTmWX2akUfS-Ga5WSiJVTzBjDQPniiONMjy4,7448
 aiq/builder/component_utils.py,sha256=2jIXWSLIlKxDKAO7kdXz_4BHqQNWw4t9GmcQfw0ER4g,12923
-aiq/builder/context.py,sha256=f-kw7JfsIFW83V3ZJFv2OU8LcvNVAg-vRVZ9qtbVESo,9173
+aiq/builder/context.py,sha256=34laVLvpsHxvj5sm_3zpXVA7IgxZE1ZDNXjXAi_74KE,9572
 aiq/builder/embedder.py,sha256=M-n2oXSlwE7u-nmpzRbPq6UVWtFtMvrEdKqqsbKQZLk,965
 aiq/builder/eval_builder.py,sha256=UnNgtQiDAUfT3yuwjZQVerenI09-4q0Cse9uwLjk3Fg,4657
 aiq/builder/evaluator.py,sha256=O6Gu0cUwQkrPxPX29Vf_-RopgijxPnhy7mhg_j-9A84,1162
@@ -74,7 +74,7 @@ aiq/cli/commands/workflow/templates/pyproject.toml.j2,sha256=tDV7-vbt8Of82OEdSOi
 aiq/cli/commands/workflow/templates/register.py.j2,sha256=SlOFmIZakPDu_E6DbIhUZ3yP8KhTrAQCFGBuhy9Fyg4,170
 aiq/cli/commands/workflow/templates/workflow.py.j2,sha256=NRp0MP8GtZByk7lrHp2Y5_6iEopRK2Wyrt0v0_2qQeo,1226
 aiq/data_models/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
-aiq/data_models/api_server.py,sha256=B7PS9u3k2gD23tq5Hrmtji6LhL_hPsz8O-zVHbe0IEQ,18950
+aiq/data_models/api_server.py,sha256=EWeFo4TvPt3VC6s21VjcE5Pc13XnwL2_YEhA5LLSg_k,18958
 aiq/data_models/common.py,sha256=G63rUXvDAtK6p1SrRyH0VlHGqrDgCZVVjbnzgGSl2Ic,4213
 aiq/data_models/component.py,sha256=x6jm1Fhn1k1hGu-5AjM0ywuyvs6ztaZfapD8bLUXSqc,1469
 aiq/data_models/component_ref.py,sha256=GyyIf4k80aUIn6LV9r84m5imbiVhpdaY7uKMMpYpbzU,3872
@@ -148,12 +148,12 @@ aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py,sha256=ZbAM3qZFUQJDp2s
 aiq/front_ends/fastapi/intermediate_steps_subscriber.py,sha256=2Y3ZXfiu8d85qJRWbT3-ex6iFDuXO6TnNtjQ6Cjl-tc,3131
 aiq/front_ends/fastapi/job_store.py,sha256=AZC0a8miKqSI9YbW19-AA_GPITVY9bYHXn9H8saNrfQ,6389
 aiq/front_ends/fastapi/main.py,sha256=ftLYy8YEyiY38kxReRwAVKwQsTa7QrKjnBU8V1omXaU,2789
-aiq/front_ends/fastapi/message_handler.py,sha256=3rFDXG633Upom1taW3ab_sC3KKxN8Y_WoM_kXtJ3K6o,12640
+aiq/front_ends/fastapi/message_handler.py,sha256=6511eexUf53h5YKi9Zz6ERP2kCpl9NSAN2ZmDbKNFb4,12769
 aiq/front_ends/fastapi/message_validator.py,sha256=NqjeIG0InGAS6yooEnTaYwjfy3qtQHNgmdJ4zZlxgSQ,17407
 aiq/front_ends/fastapi/register.py,sha256=-Vr6HEDy7L1IIBQZy6VFKZWWKYNtSI-cxk3l4ZPPLko,1159
 aiq/front_ends/fastapi/response_helpers.py,sha256=JnOOvurlkhh6akpzkKDNh1xCi5ClQTePyW_ScEY1pLo,9111
 aiq/front_ends/fastapi/step_adaptor.py,sha256=2ps136GzZ7EFgNrN0rPCUSJllvO-GNZPtADrGYseCNc,12551
-aiq/front_ends/fastapi/websocket.py,sha256=upRuFSceWRZcjBcjWG7l6Tl6fxY4NnMroSq372bdkpw,6569
+aiq/front_ends/fastapi/websocket.py,sha256=a7at8v7Yg_c-KVrdahU920tLYxZIWpSLeSjFdU0VSVI,7042
 aiq/front_ends/mcp/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
 aiq/front_ends/mcp/mcp_front_end_config.py,sha256=Mp2mgCh1pJkegih5DEJ3t0OKEiCiCg8-rM5vnLpLS6U,1470
 aiq/front_ends/mcp/mcp_front_end_plugin.py,sha256=XXyu5A5zWLH-eZgLDGcdCELXkwC1YPb94z6wddadH8I,3677
@@ -248,8 +248,8 @@ aiq/retriever/nemo_retriever/retriever.py,sha256=IvScUr9XuDLiMR__I3QsboLaM52N5D5
 aiq/runtime/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
 aiq/runtime/loader.py,sha256=4K96mko9BfSEdYg2cLZ6a5G8j-6PAkl6JSigSOshg1A,6823
 aiq/runtime/runner.py,sha256=WUAiliqI5Se9OgRmimpeFdrl38d9gRTJQ8uf59lvk7U,5836
-aiq/runtime/session.py,sha256=Hd92_MjYkPNdjbuoxQVV5EOJ0d_X2re449O6XHDlJZI,5039
-aiq/runtime/user_metadata.py,sha256=d7K5CFdOvaXpP2NCUIFBY2lHxB1FnjqV4yOlT4C7fuU,3697
+aiq/runtime/session.py,sha256=A7o3uPRznRQMfVX3LID6R30D3COTkMIQQOX6UGfk4BA,5410
+aiq/runtime/user_metadata.py,sha256=9EiBc-EEJzOdpf3Q1obHqAdY_kRlJ1T0TVvY0Jonk6o,3692
 aiq/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/settings/global_settings.py,sha256=Utm7GEiZTYTBZsxJt2TLp4bNWWMscU4b47MAEJehIrU,12036
 aiq/test/.namespace,sha256=Gace0pOC3ETEJf-TBVuNw0TQV6J_KtOPpEiSzMH-odo,215
@@ -259,7 +259,7 @@ aiq/tool/document_search.py,sha256=w3D3r5ZBS2jYmVAZZ7lC7xCoi25bA1RePoFjjlV1Zog,6
 aiq/tool/nvidia_rag.py,sha256=9mS3igONo1RywxXNj_ITh2-qD91x1R0f7uhOWMZQX3o,4178
 aiq/tool/register.py,sha256=l9TbCf9T7IXvvXCidzjYPDygBx_E5as5okN2kIZtieE,1461
 aiq/tool/retriever.py,sha256=DnuU4khpJkd4epDBGQsowDOqDBKFiLQrnyKXgU6IRW8,3724
-aiq/tool/server_tools.py,sha256=e_HCYleB64KqAzT7HZ-i__mZ3QsNXEZzBTGtDOYMLvw,3063
+aiq/tool/server_tools.py,sha256=AiopRoJx0IztsHPGfsDywfecKYAMGSKPYJCYNYS1RtU,3186
 aiq/tool/code_execution/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/tool/code_execution/code_sandbox.py,sha256=ha0IVHmWsEp0U5A0XywG6hp42pFdj5esM19CRS94_30,5986
 aiq/tool/code_execution/register.py,sha256=H-4eZwu6f5jANEBnyONcDdeBrivkHjw-aZU-V7CLMRs,3143
@@ -309,10 +309,10 @@ aiq/utils/reactive/base/observer_base.py,sha256=UAlyAY_ky4q2t0P81RVFo2Bs_R7z5Nde
 aiq/utils/reactive/base/subject_base.py,sha256=Ed-AC6P7cT3qkW1EXjzbd5M9WpVoeN_9KCe3OM3FLU4,2521
 aiq/utils/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/utils/settings/global_settings.py,sha256=U9TCLdoZsKq5qOVGjREipGVv9e-FlStzqy5zv82_VYk,7454
-aiqtoolkit-1.2.0a20250617.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
-aiqtoolkit-1.2.0a20250617.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-aiqtoolkit-1.2.0a20250617.dist-info/METADATA,sha256=7U1QTdteJlIPEhjypLhDj_VsNOp3dcGC5GyK4sPecD8,20274
-aiqtoolkit-1.2.0a20250617.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aiqtoolkit-1.2.0a20250617.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
-aiqtoolkit-1.2.0a20250617.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
-aiqtoolkit-1.2.0a20250617.dist-info/RECORD,,
+aiqtoolkit-1.2.0a20250618.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
+aiqtoolkit-1.2.0a20250618.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+aiqtoolkit-1.2.0a20250618.dist-info/METADATA,sha256=jVXfiZS96SmSPFu4lK1gpcF3ZdTCuujLw08gmCtGGSw,20274
+aiqtoolkit-1.2.0a20250618.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aiqtoolkit-1.2.0a20250618.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
+aiqtoolkit-1.2.0a20250618.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+aiqtoolkit-1.2.0a20250618.dist-info/RECORD,,