Add W&B Weave tracing

packages/agentiq_weave/pyproject.toml

@@ -0,0 +1,41 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools >= 64", "setuptools-scm>=8"]
+
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["aiq.*"]
+
+
+[tool.setuptools_scm]
+root = "../.."
+
+
+[project]
+name = "agentiq-weave"
+dynamic = ["version"]
+dependencies = [
+  # Keep package version constraints as open as possible to avoid conflicts with other packages. Always define a minimum
+  # version when adding a new package. If unsure, default to using `~=` instead of `==`. Does not apply to aiq packages.
+  # Keep sorted!!!
+  "agentiq",
+  "weave>=0.51.44"
+]
+requires-python = ">=3.12"
+description = "Subpackage for Weave integration in AgentIQ"
+readme = "src/aiq/meta/pypi.md"
+keywords = ["ai", "observability", "wandb"]
+classifiers = ["Programming Language :: Python"]
+
+
+[tool.uv]
+config-settings = { editable_mode = "compat" }
+
+
+[tool.uv.sources]
+agentiq = { workspace = true }
+
+
+[project.entry-points.'aiq.components']
+aiq_weave = "aiq.plugins.weave.register"

packages/agentiq_weave/src/aiq/meta/pypi.md

@@ -0,0 +1,23 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http:/www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+![NVIDIA AgentIQ](https://media.githubusercontent.com/media/NVIDIA/AgentIQ/refs/heads/main/docs/source/_static/agentiq_banner.png "AgentIQ banner image")
+
+# NVIDIA AgentIQ Subpackage
+This is a subpackage for Weights and Biases (W&B) Weave integration in AgentIQ.
+
+For more information about AgentIQ, please visit the [AgentIQ package](https://pypi.org/project/agentiq/).

packages/agentiq_weave/src/aiq/plugins/weave/register.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=unused-import
+# flake8: noqa
+# isort:skip_file
+
+# Import any providers which need to be automatically registered here
+
+from . import weave_sdk

packages/agentiq_weave/src/aiq/plugins/weave/weave_sdk.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Optional
+
+from pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_telemetry_exporter
+from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
+
+
+def set_wandb_api_key(config_api_key: Optional[str] = None) -> Optional[str]:
+    """
+    Get the W&B API key from various sources in order of priority:
+    1. Config provided key
+    2. WANDB_API_KEY environment variable
+    Returns:
+        The API key if found, None otherwise
+    """
+    if config_api_key:
+        return config_api_key
+    # Check environment variable
+    env_api_key = os.environ.get("WANDB_API_KEY")
+    if env_api_key:
+        return env_api_key
+    return None
+
+
+class WeaveTelemetryExporter(TelemetryExporterBaseConfig, name="weave"):
+    """A telemetry exporter to transmit traces to Weights & Biases Weave using OpenTelemetry."""
+    entity: str = Field(description="The W&B entity/organization.")
+    project: str = Field(description="The W&B project name.")
+    api_key: Optional[str] = Field(
+        default=None,
+        description="Your W&B API key for auth. If not provided, look for WANDB_API_KEY environment variable.")
+
+
+@register_telemetry_exporter(config_type=WeaveTelemetryExporter)
+async def weave_telemetry_exporter(config: WeaveTelemetryExporter, builder: Builder):
+    if config.api_key:
+        set_wandb_api_key(config.api_key)
+
+    import weave
+    _ = weave.init(project_name=f"{config.entity}/{config.project}")
+
+    class NoOpSpanExporter:
+
+        def export(self, spans):
+            return None
+
+        def shutdown(self):
+            return None
+
+    # just yielding None errors with 'NoneType' object has no attribute 'export'
+    yield NoOpSpanExporter()

pyproject.toml

@@ -54,14 +54,15 @@ maintainers = [{ name = "NVIDIA Corporation" }]
 
 
 [project.optional-dependencies]
-# Optional dependencies are things that users would want to install with AIQToolkit. i.e. `uv pip install aiq[langchain]`
-crewai = ["aiqtoolkit-crewai"]
-langchain = ["aiqtoolkit-langchain"]
-llama-index = ["aiqtoolkit-llama-index"]
-mem0ai = ["aiqtoolkit-mem0ai"]
-semantic-kernel = ["aiqtoolkit-semantic-kernel"]
-zep-cloud = ["aiqtoolkit-zep-cloud"]
-agno = ["aiqtoolkit-agno"]
+# Optional dependencies are things that users would want to install with AgentIQ. i.e. `uv pip install aiq[langchain]`
+crewai = ["agentiq-crewai"]
+langchain = ["agentiq-langchain"]
+llama-index = ["agentiq-llama-index"]
+mem0ai = ["agentiq-mem0ai"]
+semantic-kernel = ["agentiq-semantic-kernel"]
+zep-cloud = ["agentiq-zep-cloud"]
+agno = ["agentiq-agno"]
+weave = ["agentiq-weave"]
 
 examples = [
   "aiq_email_phishing_analyzer",
@@ -88,14 +89,15 @@ config-settings = { editable_mode = "compat" }
 
 [tool.uv.sources]
 # Workspace members
-aiqtoolkit-crewai = { workspace = true }
-aiqtoolkit-langchain = { workspace = true }
-aiqtoolkit-llama-index = { workspace = true }
-aiqtoolkit-mem0ai = { workspace = true }
-aiqtoolkit-semantic-kernel = { workspace = true }
-aiqtoolkit-test = { workspace = true }
-aiqtoolkit-zep-cloud = { workspace = true }
-aiqtoolkit-agno = { workspace = true }
+agentiq-crewai = { workspace = true }
+agentiq-langchain = { workspace = true }
+agentiq-llama-index = { workspace = true }
+agentiq-mem0ai = { workspace = true }
+agentiq-semantic-kernel = { workspace = true }
+agentiq-test = { workspace = true }
+agentiq-zep-cloud = { workspace = true }
+agentiq-agno = { workspace = true }
+agentiq-weave = { workspace = true }
 
 # All examples here
 aiq_email_phishing_analyzer = { path = "examples/email_phishing_analyzer", editable = true }
@@ -306,8 +308,7 @@ argument-naming-style = "snake_case"
 attr-naming-style = "snake_case"
 
 # Regular expression matching correct attribute names. Overrides attr-naming-
-# style. If left empty, attribute names will be checked with the set naming
-# style.
+# style. If left empty, attribute names will be checked with the set naming style.
 # attr-rgx =
 
 # Bad variable names which should always be refused, separated by a comma.

src/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.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,14 @@ def __init__(self, context_state: AIQContextState | None = None):
 
         self._tracer = trace.get_tracer("aiq-async-otel-listener")
 
+        # Initialize Weave-specific components if available
+        if WEAVE_AVAILABLE:
+            # get the weave client
+            self.gc = weave_client_context.require_weave_client()
+            self._weave_calls: dict[str, Call] = {}
+        else:
+            self._weave_calls = {}
+
     def _on_next(self, step: IntermediateStep) -> None:
         """
         The main logic that reacts to each IntermediateStep.
@@ -159,6 +179,12 @@ async def _cleanup(self):
 
         self._span_stack.clear()
 
+        # Clean up any lingering Weave calls if Weave is available
+        if WEAVE_AVAILABLE:
+            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
@@ -234,6 +260,10 @@ def _process_start_event(self, step: IntermediateStep):
 
         self._outstanding_spans[step.UUID] = sub_span
 
+        # Create corresponding Weave call if Weave is available
+        if WEAVE_AVAILABLE:
+            self._create_weave_call(step, sub_span)
+
     def _process_end_event(self, step: IntermediateStep):
 
         # Find the subspan that was created in the start event
@@ -268,3 +298,114 @@ def _process_end_event(self, step: IntermediateStep):
 
         # End the subspan
         sub_span.end(end_time=end_ns)
+
+        # Finish corresponding Weave call if Weave is available
+        if WEAVE_AVAILABLE:
+            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)