aiqtoolkit 1.1.0a20250501__py3-none-any.whl → 1.1.0a20250502__py3-none-any.whl

aiq/data_models/api_server.py

@@ -32,8 +32,57 @@ from aiq.data_models.interactive import HumanPrompt
 from aiq.utils.type_converter import GlobalTypeConverter
 
 
+class ChatContentType(str, Enum):
+    """
+    ChatContentType is an Enum that represents the type of Chat content.
+    """
+    TEXT = "text"
+    IMAGE_URL = "image_url"
+    INPUT_AUDIO = "input_audio"
+
+
+class InputAudio(BaseModel):
+    data: str = "default"
+    format: str = "default"
+
+
+class AudioContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: typing.Literal[ChatContentType.INPUT_AUDIO] = ChatContentType.INPUT_AUDIO
+    input_audio: InputAudio = InputAudio()
+
+
+class ImageUrl(BaseModel):
+    url: HttpUrl = HttpUrl(url="http://default.com")
+
+
+class ImageContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: typing.Literal[ChatContentType.IMAGE_URL] = ChatContentType.IMAGE_URL
+    image_url: ImageUrl = ImageUrl()
+
+
+class TextContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: typing.Literal[ChatContentType.TEXT] = ChatContentType.TEXT
+    text: str = "default"
+
+
+class Security(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    api_key: str = "default"
+    token: str = "default"
+
+
+UserContent = typing.Annotated[TextContent | ImageContent | AudioContent, Discriminator("type")]
+
+
 class Message(BaseModel):
-    content: str
+    content: str | list[UserContent]
     role: str
 
 
@@ -65,6 +114,20 @@ class AIQChatRequest(BaseModel):
                               max_tokens=max_tokens,
                               top_p=top_p)
 
+    @staticmethod
+    def from_content(content: list[UserContent],
+                     *,
+                     model: str | None = None,
+                     temperature: float | None = None,
+                     max_tokens: int | None = None,
+                     top_p: float | None = None) -> "AIQChatRequest":
+
+        return AIQChatRequest(messages=[Message(content=content, role="user")],
+                              model=model,
+                              temperature=temperature,
+                              max_tokens=max_tokens,
+                              top_p=top_p)
+
 
 class AIQChoiceMessage(BaseModel):
     content: str | None = None
@@ -231,15 +294,6 @@ class UserMessageContentRoleType(str, Enum):
     ASSISTANT = "assistant"
 
 
-class ChatContentType(str, Enum):
-    """
-    ChatContentType is an Enum that represents the type of Chat content.
-    """
-    TEXT = "text"
-    IMAGE_URL = "image_url"
-    INPUT_AUDIO = "input_audio"
-
-
 class WebSocketMessageType(str, Enum):
     """
     WebSocketMessageType is an Enum that represents WebSocket Message types.
@@ -270,46 +324,6 @@ class WebSocketMessageStatus(str, Enum):
     COMPLETE = "complete"
 
 
-class InputAudio(BaseModel):
-    data: str = "default"
-    format: str = "default"
-
-
-class AudioContent(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    type: typing.Literal[ChatContentType.INPUT_AUDIO] = ChatContentType.INPUT_AUDIO
-    input_audio: InputAudio = InputAudio()
-
-
-class ImageUrl(BaseModel):
-    url: HttpUrl = HttpUrl(url="http://default.com")
-
-
-class ImageContent(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    type: typing.Literal[ChatContentType.IMAGE_URL] = ChatContentType.IMAGE_URL
-    image_url: ImageUrl = ImageUrl()
-
-
-class TextContent(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    type: typing.Literal[ChatContentType.TEXT] = ChatContentType.TEXT
-    text: str = "default"
-
-
-class Security(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    api_key: str = "default"
-    token: str = "default"
-
-
-UserContent = typing.Annotated[TextContent | ImageContent | AudioContent, Discriminator("type")]
-
-
 class UserMessages(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -487,7 +501,9 @@ GlobalTypeConverter.register_converter(_generate_response_to_chat_response)
 
 # ======== AIQChatRequest Converters ========
 def _aiq_chat_request_to_string(data: AIQChatRequest) -> str:
-    return data.messages[-1].content
+    if isinstance(data.messages[-1].content, str):
+        return data.messages[-1].content
+    return str(data.messages[-1].content)
 
 
 GlobalTypeConverter.register_converter(_aiq_chat_request_to_string)

aiq/observability/async_otel_listener.py

@@ -16,6 +16,7 @@
 import logging
 import re
 from contextlib import asynccontextmanager
+from contextlib import contextmanager
 from typing import Any
 
 from openinference.semconv.trace import OpenInferenceSpanKindValues
@@ -30,6 +31,17 @@ from aiq.builder.context import AIQContextState
 from aiq.data_models.intermediate_step import IntermediateStep
 from aiq.data_models.intermediate_step import IntermediateStepState
 
+try:
+    from weave.trace.context import weave_client_context
+    from weave.trace.context.call_context import get_current_call
+    from weave.trace.context.call_context import set_call_stack
+    from weave.trace.weave_client import Call
+    WEAVE_AVAILABLE = True
+except ImportError:
+    WEAVE_AVAILABLE = False
+    # we simply don't do anything if weave is not available
+    pass
+
 logger = logging.getLogger(__name__)
 
 OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
@@ -84,6 +96,17 @@ class AsyncOtelSpanListener:
 
         self._tracer = trace.get_tracer("aiq-async-otel-listener")
 
+        # Initialize Weave-specific components if available
+        self.gc = None
+        self._weave_calls = {}
+        if WEAVE_AVAILABLE:
+            try:
+                # Try to get the weave client, but don't fail if Weave isn't initialized
+                self.gc = weave_client_context.require_weave_client()
+            except Exception:
+                # Weave is not initialized, so we don't do anything
+                pass
+
     def _on_next(self, step: IntermediateStep) -> None:
         """
         The main logic that reacts to each IntermediateStep.
@@ -159,6 +182,12 @@ class AsyncOtelSpanListener:
 
         self._span_stack.clear()
 
+        # Clean up any lingering Weave calls if Weave is available and initialized
+        if self.gc is not None and self._weave_calls:
+            for _, call in list(self._weave_calls.items()):
+                self.gc.finish_call(call, {"status": "incomplete"})
+            self._weave_calls.clear()
+
     def _serialize_payload(self, input_value: Any) -> tuple[str, bool]:
         """
         Serialize the input value to a string. Returns a tuple with the serialized value and a boolean indicating if the
@@ -237,6 +266,10 @@ class AsyncOtelSpanListener:
 
         self._outstanding_spans[step.UUID] = sub_span
 
+        # Create corresponding Weave call if Weave is available and initialized
+        if self.gc is not None:
+            self._create_weave_call(step, sub_span)
+
     def _process_end_event(self, step: IntermediateStep):
 
         # Find the subspan that was created in the start event
@@ -271,3 +304,114 @@ class AsyncOtelSpanListener:
 
         # End the subspan
         sub_span.end(end_time=end_ns)
+
+        # Finish corresponding Weave call if Weave is available and initialized
+        if self.gc is not None:
+            self._finish_weave_call(step, sub_span)
+
+    @contextmanager
+    def parent_call(self, trace_id: str, parent_call_id: str):
+        """Context manager to set a parent call context for Weave.
+        This allows connecting AIQ spans to existing traces from other frameworks.
+        """
+        dummy_call = Call(trace_id=trace_id, id=parent_call_id, _op_name="", project_id="", parent_id=None, inputs={})
+        with set_call_stack([dummy_call]):
+            yield
+
+    def _create_weave_call(self, step: IntermediateStep, span: Span) -> None:
+        """
+        Create a Weave call directly from the span and step data,
+        connecting to existing framework traces if available.
+        """
+        # Check for existing Weave trace/call
+        existing_call = get_current_call()
+
+        # Extract parent call if applicable
+        parent_call = None
+
+        # If we have an existing Weave call from another framework (e.g., LangChain),
+        # use it as the parent
+        if existing_call is not None:
+            parent_call = existing_call
+            logger.debug(f"Found existing Weave call: {existing_call.id} from trace: {existing_call.trace_id}")
+        # Otherwise, check our internal stack for parent relationships
+        elif len(self._weave_calls) > 0 and len(self._span_stack) > 1:
+            # Get the parent span using stack position (one level up)
+            parent_span_id = self._span_stack[-2].get_span_context().span_id
+            # Find the corresponding weave call for this parent span
+            for uuid, call in self._weave_calls.items():
+                if getattr(call, "span_id", None) == parent_span_id:
+                    parent_call = call
+                    break
+
+        # Generate a meaningful operation name based on event type
+        event_type = step.payload.event_type.split(".")[-1]
+        if step.payload.name:
+            op_name = f"aiq.{event_type}.{step.payload.name}"
+        else:
+            op_name = f"aiq.{event_type}"
+
+        # Create input dictionary
+        inputs = {}
+        if step.payload.data and step.payload.data.input is not None:
+            try:
+                # Add the input to the Weave call
+                inputs["input"] = step.payload.data.input
+            except Exception:
+                # If serialization fails, use string representation
+                inputs["input"] = str(step.payload.data.input)
+
+        # Create the Weave call
+        call = self.gc.create_call(
+            op_name,
+            inputs=inputs,
+            parent=parent_call,
+            attributes=span.attributes,
+            display_name=op_name,
+        )
+
+        # Store the call with step UUID as key
+        self._weave_calls[step.UUID] = call
+
+        # Store span ID for parent reference
+        setattr(call, "span_id", span.get_span_context().span_id)
+
+        return call
+
+    def _finish_weave_call(self, step: IntermediateStep, span: Span) -> None:
+        """
+        Finish a previously created Weave call
+        """
+        # Find the call for this step
+        call = self._weave_calls.pop(step.UUID, None)
+
+        if call is None:
+            logger.warning("No Weave call found for step %s", step.UUID)
+            return
+
+        # Create output dictionary
+        outputs = {}
+        if step.payload.data and step.payload.data.output is not None:
+            try:
+                # Add the output to the Weave call
+                outputs["output"] = step.payload.data.output
+            except Exception:
+                # If serialization fails, use string representation
+                outputs["output"] = str(step.payload.data.output)
+
+        # Add usage information if available
+        usage_info = step.payload.usage_info
+        if usage_info:
+            if usage_info.token_usage:
+                outputs["prompt_tokens"] = usage_info.token_usage.prompt_tokens
+                outputs["completion_tokens"] = usage_info.token_usage.completion_tokens
+                outputs["total_tokens"] = usage_info.token_usage.total_tokens
+
+            if usage_info.num_llm_calls:
+                outputs["num_llm_calls"] = usage_info.num_llm_calls
+
+            if usage_info.seconds_between_calls:
+                outputs["seconds_between_calls"] = usage_info.seconds_between_calls
+
+        # Finish the call with outputs
+        self.gc.finish_call(call, outputs)

{aiqtoolkit-1.1.0a20250501.dist-info → aiqtoolkit-1.1.0a20250502.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiqtoolkit
-Version: 1.1.0a20250501
+Version: 1.1.0a20250502
 Summary: Agent Intelligence Toolkit (AIQ Toolkit)
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation
@@ -250,6 +250,8 @@ Provides-Extra: zep-cloud
 Requires-Dist: aiqtoolkit-zep-cloud; extra == "zep-cloud"
 Provides-Extra: agno
 Requires-Dist: aiqtoolkit-agno; extra == "agno"
+Provides-Extra: weave
+Requires-Dist: aiqtoolkit-weave; extra == "weave"
 Provides-Extra: examples
 Requires-Dist: aiq_email_phishing_analyzer; extra == "examples"
 Requires-Dist: aiq_multi_frameworks; extra == "examples"

{aiqtoolkit-1.1.0a20250501.dist-info → aiqtoolkit-1.1.0a20250502.dist-info}/RECORD

@@ -73,7 +73,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=XgB33ZRF4OOWNT18p6sroGd9dmBGnaRNJxXLVbFhtAA,16825
+aiq/data_models/api_server.py,sha256=7BY7meyACBhrchO2-I1-X2_DhNoezdjJ5bGihK7eUyI,17533
 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
@@ -172,7 +172,7 @@ aiq/memory/models.py,sha256=4TZW2VSroLx0Ea11F_33_Rmx1diAk1jFpz-45jyPpnc,4026
 aiq/meta/module_to_distro.json,sha256=1XV7edobFrdDKvsSoynfodXg_hczUWpDrQzGkW9qqEs,28
 aiq/meta/pypi.md,sha256=LZTol6Cn-tU0cp9DynN6xCKRPONA9PVKWk3isU15SLc,4296
 aiq/observability/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aiq/observability/async_otel_listener.py,sha256=24DPWcYpRCEShHC437vH_oKgNN1sv4dJCDEaOs30A6E,11195
+aiq/observability/async_otel_listener.py,sha256=nDhJWDrsBygt03c0vsLWlB6f3ZzOh-3VjGgnJyw9vEc,17015
 aiq/observability/register.py,sha256=eCi76I9fDPWIckKDTgUN1zeinI3mPrHvhU3XkROT6G8,3921
 aiq/plugins/.namespace,sha256=Gace0pOC3ETEJf-TBVuNw0TQV6J_KtOPpEiSzMH-odo,215
 aiq/profiler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -303,10 +303,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.1.0a20250501.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
-aiqtoolkit-1.1.0a20250501.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-aiqtoolkit-1.1.0a20250501.dist-info/METADATA,sha256=Zq-gfkBzJNe-8hBeO4eI2-zc3Sq5XeOShmiXshnLpCw,19764
-aiqtoolkit-1.1.0a20250501.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-aiqtoolkit-1.1.0a20250501.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
-aiqtoolkit-1.1.0a20250501.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
-aiqtoolkit-1.1.0a20250501.dist-info/RECORD,,
+aiqtoolkit-1.1.0a20250502.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
+aiqtoolkit-1.1.0a20250502.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+aiqtoolkit-1.1.0a20250502.dist-info/METADATA,sha256=Gjr7NeisyZd3siwKG7hoQHUDT_Z4hGDfjGqVqqQskYs,19836
+aiqtoolkit-1.1.0a20250502.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
+aiqtoolkit-1.1.0a20250502.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
+aiqtoolkit-1.1.0a20250502.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+aiqtoolkit-1.1.0a20250502.dist-info/RECORD,,