aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/memory/interfaces.py

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+
+from .models import MemoryItem
+
+
+class MemoryEditor(ABC):
+    """
+    Abstract interface for editing and
+    retrieving memory items.
+
+    A MemoryEditor is responsible for adding, searching, and
+    removing MemoryItems.
+
+    Implementations may integrate with
+    vector stores or other indexing backends.
+    """
+
+    @abstractmethod
+    async def add_items(self, items: list[MemoryItem]) -> None:
+        """
+        Insert multiple MemoryItems into the memory.
+
+        Args:
+            items (list[MemoryItem]): The items to be added.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:
+        """
+        Retrieve items relevant to the given query.
+        Relevance criteria depend on implementation.
+
+        Args:
+            query (str): The query string to match.
+            top_k (int): Maximum number of items to return.
+            kwargs (dict): Keyword arguments to pass to the search method.
+
+        Returns:
+            list[MemoryItem]: The most relevant MemoryItems.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def remove_items(self, **kwargs) -> None:
+        """
+        Remove items. Additional parameters
+        needed for deletion can be specified in keyword arguments.
+
+        Args:
+            kwargs (dict): Keyword arguments to pass to the remove-items method.
+        """
+        raise NotImplementedError
+
+
+class MemoryIOBase(ABC):
+    """
+    Base abstract class for I/O operations
+    on memory, providing a common interface for
+
+    MemoryReader and MemoryWriter to interact
+    with a MemoryEditor.
+
+    Concrete subclasses should hold a
+    reference to a MemoryEditor instance.
+    """
+
+    def __init__(self, editor: MemoryEditor) -> None:
+        self._editor = editor
+
+
+class MemoryReader(MemoryIOBase):
+    """
+    Responsible for retrieving MemoryItems
+    from the MemoryEditor based on context or queries.
+    """
+
+    @abstractmethod
+    async def retrieve(self, context: str, top_k: int = 5) -> list[MemoryItem]:
+        """
+        Retrieve a subset of
+        MemoryItems relevant to the provided context.
+
+        Args:
+            context (str): A string representing
+            the current user context or query.
+            top_k (int): Maximum number of items to return.
+
+        Returns:
+            list[MemoryItem]: Relevant MemoryItems.
+        """
+        raise NotImplementedError
+
+
+class MemoryWriter(MemoryIOBase):
+    """
+    Responsible for converting new observations
+    (textual inputs) into MemoryItems andstoring
+    them via the MemoryEditor.
+    """
+
+    @abstractmethod
+    async def write(self, observation: str, context: str | None = None) -> list[MemoryItem]:
+        """
+        Process the given observation and store the resulting MemoryItems.
+
+        Args:
+            observation (str): The new textual input to record.
+            context (Optional[str]): Additional
+            context that might influence how the observation is stored.
+
+        Returns:
+            list[MemoryItem]: The newly created MemoryItems.
+        """
+        raise NotImplementedError
+
+
+class MemoryManager(ABC):
+    """
+    Manages the lifecycle of the stored
+    memory by applying policies such as summarization,
+    reflection, forgetting, and mergingn
+    to ensure long-term coherence and relevance.
+    """
+
+    @abstractmethod
+    async def summarize(self) -> None:
+        """
+        Summarize long or numerous MemoryItems into a more compact form.
+        This may remove the original items and store a new summary item.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def reflect(self) -> None:
+        """
+        Generate higher-level insights or
+        abstractions from existing MemoryItems.
+        This may call out to an LLM or other
+        logic to produce conceptual memory.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def forget(self, criteria: Callable[[MemoryItem], bool]) -> None:
+        """
+        Remove MemoryItems that are no
+        longer relevant or have low importance.
+
+        Args:
+            criteria (Callable[[MemoryItem], bool]): A function that
+            returns True for items to forget.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def merge(self, criteria: Callable[[MemoryItem, MemoryItem], bool]) -> None:
+        """
+        Merge similar or redundant MemoryItems
+        into a smaller set of more concise items.
+
+        Args:
+            criteria (Callable[[MemoryItem, MemoryItem], bool]): A function
+            that determines which items can be merged.
+        """
+        raise NotImplementedError

aiq/memory/models.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+
+
+class MemoryItem(BaseModel):
+    """
+    Represents a single memory item consisting of structured content and associated metadata.
+
+    Attributes
+    ----------
+    conversation : list[dict[str, str]]
+        A list of dictionaries, each containing string key-value pairs.
+    user_id : str
+        Unique identifier for this MemoryItem's user.
+    tags : list[str]
+        A list of strings representing tags attached to the item.
+    metadata : dict[str, typing.Any]
+        Metadata providing context and utility for management operations.
+    memory : str or None
+        Optional memory string. Helpful when returning a memory.
+    """
+    model_config = ConfigDict(
+        # Add an example for the schema
+        json_schema_extra={
+            "examples": [{
+                "conversation": [
+                    {
+                        "role": "user",
+                        "content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts.",
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy.",
+                    },
+                ],
+                "user_id": "user_abc",
+                "metadata": {
+                    "key_value_pairs": {
+                        "type": "technology", "relevance": "high"
+                    }
+                }
+            }]
+        },
+        # Allow population of models from arbitrary types (e.g., ORM objects)
+        arbitrary_types_allowed=True,
+        # Enable aliasing if needed
+        populate_by_name=True)
+
+    conversation: list[dict[str, str]] = Field(description="List of conversation messages. "
+                                               "Each message must have a \"role\" "
+                                               "key (user or assistant. It must "
+                                               "als have a \"content\" key.")
+    tags: list[str] = Field(default_factory=list, description="List of tags applied to the item.")
+    metadata: dict[str, typing.Any] = Field(description="Metadata about the memory item.")
+    user_id: str = Field(description="The user's ID.")
+    memory: str | None = Field(default=None)
+
+
+class SearchMemoryInput(BaseModel):
+    """
+    Represents a search memory input structure.
+    """
+    model_config = ConfigDict(json_schema_extra={
+        "example": {
+            "query": "What is the user's preferred programming language?",
+            "top_k": 1,
+            "user_id": "user_abc",
+        }
+    })
+
+    query: str = Field(description="Search query for which to retrieve memory.")  # noqa: E501
+    top_k: int = Field(description="Maximum number of memories to return")
+    user_id: str = Field(description="ID of the user to search for.")
+
+
+class DeleteMemoryInput(BaseModel):
+    """
+    Represents a delete memory input structure.
+    """
+    model_config = ConfigDict(json_schema_extra={"example": {"user_id": "user_abc", }})
+
+    user_id: str = Field(description="ID of the user to delete memory for. Careful when using "
+                         "this tool; make sure you use the "
+                         "username present in the conversation.")

aiq/meta/module_to_distro.json

@@ -0,0 +1,3 @@
+{
+    "aiq": "aiqtoolkit"
+}

aiq/meta/pypi.md

@@ -0,0 +1,59 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2024-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
+
+AgentIQ is a flexible library designed to seamlessly integrate your enterprise agents—regardless of framework—with various data sources and tools. By treating agents, tools, and agentic workflows as simple function calls, AgentIQ enables true composability: build once and reuse anywhere.
+
+## Key Features
+
+- [**Framework Agnostic:**](https://docs.nvidia.com/agentiq/latest/concepts/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
+- [**Reusability:**](https://docs.nvidia.com/agentiq/latest/guides/sharing-workflows-and-tools.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
+- [**Rapid Development:**](https://docs.nvidia.com/agentiq/latest/guides/create-customize-workflows.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
+- [**Profiling:**](https://docs.nvidia.com/agentiq/latest/guides/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
+- [**Observability:**](https://docs.nvidia.com/agentiq/latest/guides/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool.
+- [**Evaluation System:**](https://docs.nvidia.com/agentiq/latest/guides/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
+- [**User Interface:**](https://docs.nvidia.com/agentiq/latest/guides/using-agentiq-ui-and-server.html) Use the AgentIQ UI chat interface to interact with your agents, visualize output, and debug workflows.
+- [**MCP Compatibility**](https://docs.nvidia.com/agentiq/latest/components/mcp.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as AgentIQ functions.
+
+With AgentIQ, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
+
+## Links
+ * [Documentation](https://docs.nvidia.com/agentiq/latest/index.html): Explore the full documentation for AgentIQ.
+ * [About AgentIQ](https://docs.nvidia.com/agentiq/latest/intro/why-agentiq.html): Learn more about the benefits of using AgentIQ.
+
+## First time user?
+ If this is your first time using AgentIQ, it is recommended to install the latest version from the [source repository](https://github.com/NVIDIA/AgentIQ?tab=readme-ov-file#get-started) on GitHub. This package is intended for users who are familiar with AgentIQ applications and need to add AgentIQ as a dependency to their project.
+
+## Feedback
+
+We would love to hear from you! Please file an issue on [GitHub](https://github.com/NVIDIA/AgentIQ/issues) if you have any feedback or feature requests.
+
+## Acknowledgements
+
+We would like to thank the following open source projects that made AgentIQ possible:
+
+- [CrewAI](https://github.com/crewAIInc/crewAI)
+- [FastAPI](https://github.com/tiangolo/fastapi)
+- [LangChain](https://github.com/langchain-ai/langchain)
+- [Llama-Index](https://github.com/run-llama/llama_index)
+- [Mem0ai](https://github.com/mem0ai/mem0)
+- [Ragas](https://github.com/explodinggradients/ragas)
+- [Semantic Kernel](https://github.com/microsoft/semantic-kernel)
+- [uv](https://github.com/astral-sh/uv)

aiq/observability/async_otel_listener.py

@@ -0,0 +1,270 @@
+# 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 logging
+import re
+from contextlib import asynccontextmanager
+from typing import Any
+
+from openinference.semconv.trace import OpenInferenceSpanKindValues
+from openinference.semconv.trace import SpanAttributes
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace import Span
+from opentelemetry.trace.propagation import set_span_in_context
+from pydantic import TypeAdapter
+
+from aiq.builder.context import AIQContextState
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.data_models.intermediate_step import IntermediateStepState
+
+logger = logging.getLogger(__name__)
+
+OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
+
+
+def _ns_timestamp(seconds_float: float) -> int:
+    """
+    Convert AgentIQ’s float `event_timestamp` (in seconds) into an integer number
+    of nanoseconds, as OpenTelemetry expects.
+    """
+    return int(seconds_float * 1e9)
+
+
+class AsyncOtelSpanListener:
+    """
+    A separate, async class that listens to the AgentIQ intermediate step
+    event stream and creates proper Otel spans:
+
+    - On FUNCTION_START => open a new top-level span
+    - On any other intermediate step => open a child subspan (immediate open/close)
+    - On FUNCTION_END => close the function’s top-level span
+
+    This runs fully independently from the normal AgentIQ workflow, so that
+    the workflow is not blocking or entangled by OTel calls.
+    """
+
+    def __init__(self, context_state: AIQContextState | None = None):
+        """
+        :param context_state: Optionally supply a specific AIQContextState.
+                              If None, uses the global singleton.
+        """
+        self._context_state = context_state or AIQContextState.get()
+
+        # Maintain a subscription so we can unsubscribe on shutdown
+        self._subscription = None
+
+        # Outstanding spans which have been opened but not yet closed
+        self._outstanding_spans: dict[str, Span] = {}
+
+        # Stack of spans, for when we need to create a child span
+        self._span_stack: list[Span] = []
+
+        self._running = False
+
+        # Prepare the tracer (optionally you might already have done this)
+        if trace.get_tracer_provider() is None or not isinstance(trace.get_tracer_provider(), TracerProvider):
+            tracer_provider = TracerProvider()
+            trace.set_tracer_provider(tracer_provider)
+
+        # We’ll optionally attach exporters if you want (out of scope to do it here).
+        # Example: tracer_provider.add_span_processor(BatchSpanProcessor(your_exporter))
+
+        self._tracer = trace.get_tracer("aiq-async-otel-listener")
+
+    def _on_next(self, step: IntermediateStep) -> None:
+        """
+        The main logic that reacts to each IntermediateStep.
+        """
+        if (step.event_state == IntermediateStepState.START):
+
+            self._process_start_event(step)
+
+        elif (step.event_state == IntermediateStepState.END):
+
+            self._process_end_event(step)
+
+    def _on_error(self, exc: Exception) -> None:
+        logger.error("Error in intermediate step subscription: %s", exc, exc_info=True)
+
+    def _on_complete(self) -> None:
+        logger.debug("Intermediate step stream completed. No more events will arrive.")
+
+    @asynccontextmanager
+    async def start(self):
+        """
+        Usage::
+
+            otel_listener = AsyncOtelSpanListener()
+            async with otel_listener.start():
+                # run your AgentIQ workflow
+                ...
+            # cleans up
+
+        This sets up the subscription to the AgentIQ event stream and starts the background loop.
+        """
+        try:
+            # Subscribe to the event stream
+            subject = self._context_state.event_stream.get()
+            self._subscription = subject.subscribe(
+                on_next=self._on_next,
+                on_error=self._on_error,
+                on_complete=self._on_complete,
+            )
+
+            self._running = True
+
+            yield  # let the caller do their workflow
+
+        finally:
+            # Cleanup
+            self._running = False
+            # Close out any running spans
+            await self._cleanup()
+
+            if self._subscription:
+                self._subscription.unsubscribe()
+            self._subscription = None
+
+    async def _cleanup(self):
+        """
+        Close any remaining open spans.
+        """
+        if self._outstanding_spans:
+            logger.warning(
+                "Not all spans were closed. Ensure all start events have a corresponding end event. Remaining: %s",
+                self._outstanding_spans)
+
+        for span_info in self._outstanding_spans.values():
+            span_info.end()
+
+        self._outstanding_spans.clear()
+
+        if self._span_stack:
+            logger.error(
+                "Not all spans were closed. Ensure all start events have a corresponding end event. Remaining: %s",
+                self._span_stack)
+
+        self._span_stack.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
+        serialization is JSON or a string
+        """
+        try:
+            return TypeAdapter(type(input_value)).dump_json(input_value).decode('utf-8'), True
+        except Exception:
+            # Fallback to string representation if we can't serialize using pydantic
+            return str(input_value), False
+
+    def _process_start_event(self, step: IntermediateStep):
+
+        parent_ctx = None
+
+        if (len(self._span_stack) > 0):
+            parent_span = self._span_stack[-1]
+
+            parent_ctx = set_span_in_context(parent_span)
+
+        # Extract start/end times from the step
+        # By convention, `span_event_timestamp` is the time we started, `event_timestamp` is the time we ended.
+        # If span_event_timestamp is missing, we default to event_timestamp (meaning zero-length).
+        s_ts = step.payload.span_event_timestamp or step.payload.event_timestamp
+        start_ns = _ns_timestamp(s_ts)
+
+        # Optional: embed the LLM/tool name if present
+        if step.payload.name:
+            sub_span_name = f"{step.payload.name}"
+        else:
+            sub_span_name = f"{step.payload.event_type}"
+
+        # Start the subspan
+        sub_span = self._tracer.start_span(
+            name=sub_span_name,
+            context=parent_ctx,
+            attributes={
+                "aiq.event_type": step.payload.event_type.value,
+                "aiq.function.id": step.function_ancestry.function_id,
+                "aiq.function.name": step.function_ancestry.function_name,
+                "aiq.subspan.name": step.payload.name or "",
+                "aiq.event_timestamp": step.event_timestamp,
+                "aiq.framework": step.payload.framework.value if step.payload.framework else "unknown",
+            },
+            start_time=start_ns,
+        )
+
+        event_type_to_span_kind = {
+            "LLM_START": OpenInferenceSpanKindValues.LLM,
+            "LLM_END": OpenInferenceSpanKindValues.LLM,
+            "LLM_NEW_TOKEN": OpenInferenceSpanKindValues.LLM,
+            "TOOL_START": OpenInferenceSpanKindValues.TOOL,
+            "TOOL_END": OpenInferenceSpanKindValues.TOOL,
+            "FUNCTION_START": OpenInferenceSpanKindValues.CHAIN,
+            "FUNCTION_END": OpenInferenceSpanKindValues.CHAIN,
+        }
+
+        span_kind = event_type_to_span_kind.get(step.event_type, OpenInferenceSpanKindValues.UNKNOWN)
+        sub_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind.value)
+
+        if step.payload.data and step.payload.data.input:
+            # optional parse
+            match = re.search(r"Human:\s*Question:\s*(.*)", str(step.payload.data.input))
+            if match:
+                human_question = match.group(1).strip()
+                sub_span.set_attribute(SpanAttributes.INPUT_VALUE, human_question)
+            else:
+                serialized_input, is_json = self._serialize_payload(step.payload.data.input)
+                sub_span.set_attribute(SpanAttributes.INPUT_VALUE, serialized_input)
+                sub_span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json" if is_json else "text/plain")
+
+        self._span_stack.append(sub_span)
+
+        self._outstanding_spans[step.UUID] = sub_span
+
+    def _process_end_event(self, step: IntermediateStep):
+
+        # Find the subspan that was created in the start event
+        sub_span = self._outstanding_spans.pop(step.UUID, None)
+
+        if sub_span is None:
+            logger.warning("No subspan found for step %s", step.UUID)
+            return
+
+        self._span_stack.pop()
+
+        # Optionally add more attributes from usage_info or data
+        usage_info = step.payload.usage_info
+        if usage_info:
+            sub_span.set_attribute("aiq.usage.num_llm_calls",
+                                   usage_info.num_llm_calls if usage_info.num_llm_calls else 0)
+            sub_span.set_attribute("aiq.usage.seconds_between_calls",
+                                   usage_info.seconds_between_calls if usage_info.seconds_between_calls else 0)
+            sub_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT,
+                                   usage_info.token_usage.prompt_tokens if usage_info.token_usage else 0)
+            sub_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION,
+                                   usage_info.token_usage.completion_tokens if usage_info.token_usage else 0)
+            sub_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL,
+                                   usage_info.token_usage.total_tokens if usage_info.token_usage else 0)
+
+        if step.payload.data and step.payload.data.output is not None:
+            serialized_output, is_json = self._serialize_payload(step.payload.data.output)
+            sub_span.set_attribute(SpanAttributes.OUTPUT_VALUE, serialized_output)
+            sub_span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json" if is_json else "text/plain")
+
+        end_ns = _ns_timestamp(step.payload.event_timestamp)
+
+        # End the subspan
+        sub_span.end(end_time=end_ns)