Merge branch 'develop' of https://github.com/NVIDIA/NeMo-Agent-Toolkit into david-notebooks-e2e

Signed-off-by: David Gardner <[email protected]>

Author: dagardner-nv

ci/vale/styles/config/vocabularies/nat/accept.txt

@@ -143,7 +143,7 @@ Tavily
 [Tt]okenization
 [Tt]okenizer(s?)
 triages
-[Uu]ncomment
+[Uu]ncomment(ed)?
 [Uu]nencrypted
 [Uu]nittest(s?)
 [Uu]nprocessable

examples/README.md

@@ -121,6 +121,7 @@ To run the examples, install the NeMo Agent toolkit from source, if you haven't
 3. [Adding Tools and Agents](notebooks/3_adding_tools_to_agents.ipynb) - Adding tools to your agentic workflow
 4. [Multi-Agent Orchestration](notebooks/4_multi_agent_orchestration.ipynb) - Setting up a multi-agent orchestration workflow
 5. [Observability, Evaluation, and Profiling](notebooks/5_observability_evaluation_and_profiling.ipynb) - Instrumenting with observability, evaluation and profiling tools
+6. [Optimizing Model Selection, Parameters, and Prompts](notebooks/6_optimize_model_selection.ipynb) - Use NAT Optimize to compare models, parameters, and prompt variations
 
 #### Brev Launchables
 

examples/advanced_agents/alert_triage_agent/src/nat_alert_triage_agent/register.py

@@ -25,6 +25,7 @@
 from nat.data_models.component_ref import LLMRef
 from nat.data_models.function import FunctionBaseConfig
 from nat.profiler.decorators.function_tracking import track_function
+from nat.data_models.optimizable import OptimizableMixin
 
 # flake8: noqa
 # Import any tools which need to be automatically registered here
@@ -43,7 +44,7 @@
 from .prompts import ALERT_TRIAGE_AGENT_PROMPT
 
 
-class AlertTriageAgentWorkflowConfig(FunctionBaseConfig, name="alert_triage_agent"):
+class AlertTriageAgentWorkflowConfig(FunctionBaseConfig, OptimizableMixin, name="alert_triage_agent"):
     """
     Configuration for the Alert Triage Agent workflow. This agent orchestrates multiple diagnostic tools
     to analyze and triage alerts by:
@@ -54,7 +55,7 @@ class AlertTriageAgentWorkflowConfig(FunctionBaseConfig, name="alert_triage_agen
     """
     tool_names: list[str] = []
     llm_name: LLMRef
-    offline_mode: bool = Field(default=True, description="Whether to run in offline model")
+    offline_mode: bool = Field(default=True, description="Whether to run in offline mode")
     offline_data_path: str | None = Field(
         default="examples/advanced_agents/alert_triage_agent/data/offline_data.csv",
         description="Path to the main offline dataset in CSV format containing alerts and their simulated environments")

examples/notebooks/6_optimize_model_selection.ipynb

@@ -26,6 +26,7 @@ We showcase the building blocks that make up the agentic system, including tools
 3. [Adding Tools and Agents](3_adding_tools_to_agents.ipynb) - Adding tools to your agentic workflow
 4. [Multi-Agent Orchestration](4_multi_agent_orchestration.ipynb) - Setting up a multi-agent orchestration workflow
 5. [Observability, Evaluation, and Profiling](5_observability_evaluation_and_profiling.ipynb) - Instrumenting with observability, evaluation and profiling tools
+6. [Optimizing Model Selection, Parameters, and Prompts](6_optimize_model_selection.ipynb) - Use NAT Optimize to compare models, parameters, and prompt variations
 
 We recommend opening these notebooks in a Jupyter Lab environment or Google Colab environment.
 

examples/notebooks/README.md

@@ -21,7 +21,7 @@ dependencies = [
   # version when adding a new package. If unsure, default to using `~=` instead of `==`. Does not apply to nvidia-nat packages.
   # Keep sorted!!!
   "nvidia-nat~=1.4",
-  "zep-cloud~=2.2.0",
+  "zep-cloud~=3.0",
 ]
 requires-python = ">=3.11,<3.14"
 description = "Subpackage for Zep integration in NeMo Agent toolkit"

packages/nvidia_nat_zep_cloud/pyproject.toml

@@ -16,90 +16,254 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 
+from zep_cloud import ApiError
+from zep_cloud import NotFoundError
 from zep_cloud.client import AsyncZep
 from zep_cloud.types import Message
 
+from nat.builder.context import Context
 from nat.memory.interfaces import MemoryEditor
 from nat.memory.models import MemoryItem
 
+logger = logging.getLogger(__name__)
+
 
 class ZepEditor(MemoryEditor):
     """
-    Wrapper class that implements NAT interfaces for Zep Integrations Async.
+    Wrapper class that implements NAT interfaces for Zep v3 Integrations Async.
+    Uses thread-based memory management with automatic user creation.
     """
 
-    def __init__(self, zep_client: AsyncZep):
+    def __init__(self, zep_client: AsyncZep) -> None:
         """
-        Initialize class with Predefined Mem0 Client.
+        Initialize class with Zep v3 AsyncZep Client.
 
         Args:
-        zep_client (AsyncZep): Async client instance.
+            zep_client (AsyncZep): Async client instance.
         """
         self._client = zep_client
 
-    async def add_items(self, items: list[MemoryItem]) -> None:
+    async def _ensure_user_exists(self, user_id: str) -> None:
+        """
+        Ensure a user exists in Zep v3, creating if necessary.
+
+        Args:
+            user_id (str): The user ID to check/create.
         """
-        Insert Multiple MemoryItems into the memory. Each MemoryItem is translated and uploaded.
+        logger.debug("Checking if Zep user exists")
+        try:
+            await self._client.user.get(user_id=user_id)
+            logger.debug("Zep user already exists")
+        except NotFoundError:
+            # User doesn't exist, create with basic info
+            logger.info("Zep user not found, creating...")
+            try:
+                # Set defaults only for default_user, otherwise use just user_id
+                if user_id == "default_user":
+                    email = "[email protected]"
+                    first_name = "Jane"
+                    last_name = "Doe"
+                    await self._client.user.add(user_id=user_id,
+                                                email=email,
+                                                first_name=first_name,
+                                                last_name=last_name)
+                else:
+                    # For non-default users, just use user_id (email/names not required)
+                    await self._client.user.add(user_id=user_id)
+
+                logger.info("Created Zep user")
+            except ApiError as e:
+                # Check if user was created by another request (409 Conflict)
+                if e.response_data and e.response_data.get("status_code") == 409:
+                    logger.info("Zep user already exists (409), continuing")
+                else:
+                    logger.error("Failed creating Zep user: %s", str(e))  # noqa: TRY400
+                    raise
+        except ApiError as e:
+            logger.error("Failed fetching Zep user: %s", str(e))  # noqa: TRY400
+            raise
+
+    async def add_items(self, items: list[MemoryItem], **kwargs) -> None:
         """
+        Insert Multiple MemoryItems into the memory using Zep v3 thread API.
+        Each MemoryItem is translated and uploaded to a thread.
+        Uses conversation_id from NAT context as thread_id for multi-thread support.
+
+        Args:
+            items (list[MemoryItem]): The items to be added.
+            kwargs (dict): Provider-specific keyword arguments.
+
+                - ignore_roles (list[str], optional): List of role types to ignore when adding
+                  messages to graph memory. Available roles: system, assistant, user,
+                  function, tool.
+        """
+        # Extract Zep-specific parameters
+        ignore_roles = kwargs.get("ignore_roles", None)
 
         coroutines = []
+        created_threads: set[str] = set()
+        ensured_users: set[str] = set()
 
-        # Iteratively insert memories into Mem0
+        # Iteratively insert memories into Zep using threads
         for memory_item in items:
             conversation = memory_item.conversation
-            session_id = memory_item.user_id
+            user_id = memory_item.user_id or "default_user"  # Validate user_id
+
+            # Get thread_id from NAT context (unique per UI conversation)
+            thread_id = Context.get().conversation_id
+
+            # Fallback to default thread ID if no conversation_id available
+            if not thread_id:
+                thread_id = "default_zep_thread"
+
             messages = []
+
+            # Ensure user exists before creating thread (only once per user)
+            if user_id not in ensured_users:
+                await self._ensure_user_exists(user_id)
+                ensured_users.add(user_id)
+
+            # Skip if no conversation data
+            if not conversation:
+                continue
+
             for msg in conversation:
-                messages.append(Message(content=msg["content"], role_type=msg["role"]))
+                # Create Message - role field instead of role_type in V3
+                message = Message(content=msg["content"], role=msg["role"])
+                messages.append(message)
+
+            # Ensure thread exists once per thread_id
+            thread_ready = True
+            if thread_id not in created_threads:
+                logger.info("Ensuring Zep thread exists (thread_id=%s)", thread_id)
+                try:
+                    await self._client.thread.create(thread_id=thread_id, user_id=user_id)
+                    logger.info("Created Zep thread (thread_id=%s)", thread_id)
+                    created_threads.add(thread_id)
+                except ApiError as create_error:
+                    if create_error.response_data and create_error.response_data.get("status_code") == 409:
+                        logger.info("Zep thread already exists (thread_id=%s)", thread_id)
+                        created_threads.add(thread_id)
+                    else:
+                        logger.exception("Thread create failed (thread_id=%s)", thread_id)
+                        thread_ready = False
+
+            # Skip this item if thread creation failed unexpectedly
+            if not thread_ready:
+                continue
+
+            # Add messages to thread using Zep v3 API
+            logger.info("Queueing add_messages (thread_id=%s, count=%d)", thread_id, len(messages))
 
-            coroutines.append(self._client.memory.add(session_id=session_id, messages=messages))
+            # Build add_messages parameters
+            add_messages_params = {"thread_id": thread_id, "messages": messages}
+            if ignore_roles is not None:
+                add_messages_params["ignore_roles"] = ignore_roles
+
+            coroutines.append(self._client.thread.add_messages(**add_messages_params))
 
         await asyncio.gather(*coroutines)
 
-    async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:
+    async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:  # noqa: ARG002
         """
-        Retrieve items relevant to the given query.
+        Retrieve memory from Zep v3 using the high-level get_user_context API.
+        Uses conversation_id from NAT context as thread_id for multi-thread support.
+
+        Zep returns pre-formatted memory optimized for LLM consumption, including
+        relevant facts, timestamps, and structured information from its knowledge graph.
 
         Args:
-            query (str): The query string to match.
-            top_k (int): Maximum number of items to return.
-            kwargs: Other keyword arguments for search.
+            query (str): The query string (not used by Zep's high-level API, included for interface compatibility).
+            top_k (int): Maximum number of items to return (not used by Zep's context API).
+            kwargs: Zep-specific keyword arguments.
+
+                - user_id (str, required for response construction): Used only to construct the
+                  returned MemoryItem. Zep v3's thread.get_user_context() only requires thread_id.
+                - mode (str, optional): Retrieval mode. Zep server default is "summary". This
+                  implementation uses mode="basic" (NAT's default) for performance (P95 < 200ms).
+                  "summary" provides more comprehensive memory at the cost of latency.
 
         Returns:
-            list[MemoryItem]: The most relevant MemoryItems for the given query.
+            list[MemoryItem]: A single MemoryItem containing the formatted context from Zep.
         """
+        # Validate required kwargs
+        if "user_id" not in kwargs or not kwargs["user_id"]:
+            raise ValueError("user_id is required.")
+        user_id = kwargs.pop("user_id")
+        mode = kwargs.pop("mode", "basic")  # Get mode, default to "basic" for fast retrieval
+
+        # Get thread_id from NAT context
+        thread_id = Context.get().conversation_id
 
-        session_id = kwargs.pop("user_id")  # Ensure user ID is in keyword arguments
-        limit = top_k
+        # Fallback to default thread ID if no conversation_id available
+        if not thread_id:
+            thread_id = "default_zep_thread"
 
-        search_result = await self._client.memory.search_sessions(session_ids=[session_id],
-                                                                  text=query,
-                                                                  limit=limit,
-                                                                  search_scope="messages",
-                                                                  **kwargs)
+        try:
+            # Use Zep v3 thread.get_user_context - returns pre-formatted context
+            memory_response = await self._client.thread.get_user_context(thread_id=thread_id, mode=mode)
+            context_string = memory_response.context or ""
 
-        # Construct MemoryItem instances
-        memories = []
+            # Return as a single MemoryItem with the formatted context
+            if context_string:
+                return [
+                    MemoryItem(conversation=[],
+                               user_id=user_id,
+                               memory=context_string,
+                               metadata={
+                                   "mode": mode, "thread_id": thread_id
+                               })
+                ]
+            else:
+                return []
 
-        for res in search_result.results:
-            memories.append(
-                MemoryItem(conversation=[],
-                           user_id=session_id,
-                           memory=res.message.content,
-                           metadata={
-                               "relevance_score": res.score,
-                               "created_at": res.message.created_at,
-                               "updated_at": res.message.updated_at
-                           }))
+        except NotFoundError:
+            # Thread doesn't exist or no context available
+            return []
+        except ApiError as e:
+            logger.error("get_user_context failed (thread_id=%s): %s", thread_id, str(e))  # noqa: TRY400
+            raise
 
-        return memories
+    async def remove_items(self, **kwargs) -> None:
+        """
+        Remove memory items based on provided criteria.
+
+        Supports two deletion modes:
+
+        1. Delete a specific thread by thread_id
+        2. Delete all threads for a user by user_id
+
+        Args:
+            kwargs: Additional parameters.
+
+                - thread_id (str, optional): Thread ID to delete a specific thread.
+                - user_id (str, optional): User ID to delete all threads for that user.
+        """
+        if "thread_id" in kwargs:
+            # Delete specific thread
+            thread_id = kwargs.pop("thread_id")
+            logger.info("Deleting thread (thread_id=%s)", thread_id)
+            await self._client.thread.delete(thread_id=thread_id)
+        elif "user_id" in kwargs:
+            # Delete all threads for a user
+            user_id = kwargs.pop("user_id")
+            logger.debug("Deleting all threads for user (user_id=%s)", user_id)
 
-    async def remove_items(self, **kwargs):
+            # Get all threads for this user
+            threads = await self._client.user.get_threads(user_id=user_id)
+            logger.debug("Found %d threads for user (user_id=%s)", len(threads), user_id)
 
-        if "session_id" in kwargs:
-            session_id = kwargs.pop("session_id")
-            await self._client.memory.delete(session_id)
+            # Delete each thread
+            delete_coroutines = []
+            for thread in threads:
+                if thread.thread_id:
+                    logger.debug("Queueing deletion of thread (thread_id=%s)", thread.thread_id)
+                    delete_coroutines.append(self._client.thread.delete(thread_id=thread.thread_id))
 
+            if delete_coroutines:
+                await asyncio.gather(*delete_coroutines)
+                logger.info("Deleted %d threads for user", len(delete_coroutines))
         else:
-            raise ValueError("session_id not provided as part of the tool call. ")
+            raise ValueError("Either thread_id or user_id is required.")

packages/nvidia_nat_zep_cloud/src/nat/plugins/zep_cloud/zep_editor.py

@@ -196,6 +196,14 @@ def _process_start_event(self, event: IntermediateStep):
         span_kind = event_type_to_span_kind(event.event_type)
         sub_span.set_attribute(f"{self._span_prefix}.span.kind", span_kind.value)
 
+        # Enable session grouping by setting session.id from conversation_id
+        try:
+            conversation_id = self._context_state.conversation_id.get()
+            if conversation_id:
+                sub_span.set_attribute("session.id", conversation_id)
+        except (AttributeError, LookupError):
+            pass
+
         if event.payload.data and event.payload.data.input:
             match = re.search(r"Human:\s*Question:\s*(.*)", str(event.payload.data.input))
             if match: