Update README ; applying reviewers suggestions

Signed-off-by: Michele Pangrazzi <[email protected]>

Author: mpangrazzi

examples/basic/frameworks/haystack_deep_research_agent/README.md

@@ -195,7 +195,7 @@ You can customize the workflow by:
 
 1. **Modifying the system prompt** in `config.yml` to change the agent's behavior
 2. **Adding more document types** by extending the RAG tool to support other file formats
-3. **Changing the LLM model** by updating the NVIDIA model names in the configuration (`agent_model` and `rag_model` in `config.yml`). See Haystack's NvidiaChatGenerator docs: [NvidiaChatGenerator](https://docs.haystack.deepset.ai/docs/nvidiachatgenerator)
+3. **Changing the LLM model** by updating the top-level `llms` section in `config.yml`. This example defines `agent_llm` and `rag_llm` using the `nim` provider so they can leverage common parameters like `temperature`, `top_p`, and `max_tokens`. The workflow references them via the builder. See Haystack's NvidiaChatGenerator docs: [NvidiaChatGenerator](https://docs.haystack.deepset.ai/docs/nvidiachatgenerator)
 4. **Adjusting search parameters** to optimize for your use case
 
 ## Troubleshooting
@@ -213,9 +213,9 @@ You can customize the workflow by:
 
 The workflow demonstrates several key NeMo-Agent-Toolkit patterns:
 
-- **Function Registration**: Each tool is registered as a function with its own configuration
-- **Builder Pattern**: The NeMo-Agent-Toolkit Builder is used to create and manage tools and LLMs
-- **Component Integration**: Haystack components are wrapped as NeMo-Agent-Toolkit functions
+- **Workflow Registration**: The agent is exposed as a workflow function with a pydantic config
+- **Builder LLM Integration**: LLMs are defined under top-level `llms:` and accessed via `builder.get_llm_config(...)`
+- **Component Integration**: Haystack components are composed into tools within the workflow
 - **Error Handling**: Robust error handling with fallback behaviors
 - **Async Operations**: All operations are asynchronous for better performance
 

examples/basic/frameworks/haystack_deep_research_agent/src/aiq_haystack_deep_research_agent/configs/config.yml

@@ -16,12 +16,23 @@
 general:
   use_uvloop: true
 
+llms:
+  rag_llm:
+    _type: nim
+    model: nvidia/llama-3.3-nemotron-super-49b-v1
+    base_url: https://integrate.api.nvidia.com/v1
+    api_key: ${NVIDIA_API_KEY}
+  agent_llm:
+    _type: nim
+    model: nvidia/llama-3.3-nemotron-super-49b-v1
+    base_url: https://integrate.api.nvidia.com/v1
+    api_key: ${NVIDIA_API_KEY}
+
 workflow:
   _type: haystack_deep_research_agent
-  agent_model: nvidia/llama-3.3-nemotron-super-49b-v1
-  rag_model: nvidia/llama-3.3-nemotron-super-49b-v1
-  nvidia_api_url: https://integrate.api.nvidia.com/v1
   max_agent_steps: 20
+  search_top_k: 10
+  rag_top_k: 15
   opensearch_url: http://localhost:9200
   index_on_startup: true
   data_dir: /data

examples/basic/frameworks/haystack_deep_research_agent/src/aiq_haystack_deep_research_agent/pipelines/indexing.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from pathlib import Path
-from typing import Tuple
 
 from haystack.core.pipeline import Pipeline
 from haystack.components.converters.pypdf import PyPDFToDocument
@@ -12,7 +11,7 @@
 from haystack.document_stores.types import DuplicatePolicy
 
 
-def _gather_sources(base_dir: Path) -> Tuple[list[Path], list[Path]]:
+def _gather_sources(base_dir: Path) -> tuple[list[Path], list[Path]]:
     pdfs = list(base_dir.glob("**/*.pdf"))
     texts = list(base_dir.glob("**/*.txt")) + list(base_dir.glob("**/*.md"))
     return pdfs, texts

examples/basic/frameworks/haystack_deep_research_agent/src/aiq_haystack_deep_research_agent/pipelines/rag.py

@@ -1,8 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Tuple
-
 from haystack.components.builders import ChatPromptBuilder
 from haystack.core.pipeline import Pipeline
 from haystack.core.super_component import SuperComponent
@@ -16,18 +14,29 @@
 
 def create_rag_tool(
     document_store,
-    rag_model: str,
-    nvidia_api_url: str,
-    secret_provider,
     *,
     top_k: int = 15,
-) -> Tuple[ComponentTool, Pipeline]:
+    generator: NvidiaChatGenerator | None = None,
+) -> tuple[ComponentTool, Pipeline]:
+    """
+    Build a RAG tool composed of OpenSearch retriever and NvidiaChatGenerator.
+
+    Args:
+        document_store: OpenSearch document store instance.
+        top_k: Number of documents to retrieve for RAG.
+        generator: Pre-configured NvidiaChatGenerator created from builder LLM config.
+
+    Returns:
+        (ComponentTool, Pipeline): The tool and underlying pipeline.
+
+    Raises:
+        ValueError: If a generator is not provided.
+    """
     retriever = OpenSearchBM25Retriever(document_store=document_store, top_k=top_k)
-    generator = NvidiaChatGenerator(
-        model=rag_model,
-        api_base_url=nvidia_api_url,
-        api_key=secret_provider,
-    )
+    if generator is None:
+        raise ValueError(
+            "NvidiaChatGenerator instance must be provided via builder-configured LLM."
+        )
 
     template = """
 	{% for document in documents %}

examples/basic/frameworks/haystack_deep_research_agent/src/aiq_haystack_deep_research_agent/pipelines/search.py

@@ -10,9 +10,18 @@
 from haystack.tools import ComponentTool
 
 
-def create_search_tool() -> ComponentTool:
+def create_search_tool(top_k: int = 10) -> ComponentTool:
+    """
+    Build a Haystack web search tool pipeline.
+
+    Args:
+        top_k: Number of search results to retrieve from Serper.
+
+    Returns:
+        ComponentTool: A Haystack tool that executes web search and returns formatted text.
+    """
     search_pipeline = Pipeline()
-    search_pipeline.add_component("search", SerperDevWebSearch(top_k=10))
+    search_pipeline.add_component("search", SerperDevWebSearch(top_k=top_k))
     search_pipeline.add_component(
         "fetcher",
         LinkContentFetcher(timeout=3, raise_on_failure=False, retry_attempts=2),

examples/basic/frameworks/haystack_deep_research_agent/src/aiq_haystack_deep_research_agent/register.py

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
+# pyright: reportMissingTypeStubs=false, reportMissingImports=false
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,6 +19,7 @@
 from aiq.builder.builder import Builder
 from aiq.cli.register_workflow import register_function
 from aiq.data_models.function import FunctionBaseConfig
+from aiq.llm.nim_llm import NIMModelConfig
 
 logger = logging.getLogger(__name__)
 
@@ -37,11 +39,10 @@ class HaystackDeepResearchWorkflowConfig(
     When you use information from the document database, cite the text used from the source document.
     It is important that you cite accurately.
     """
-    agent_model: str = "meta/llama-3.1-8b-instruct"
-    rag_model: str = "meta/llama-3.1-8b-instruct"
     max_agent_steps: int = 20
+    search_top_k: int = 10
+    rag_top_k: int = 15
     opensearch_url: str = "http://localhost:9200"
-    nvidia_api_url: str = "https://integrate.api.nvidia.com/v1"
     # Indexing configuration
     index_on_startup: bool = True
     # Default to "/data" so users can mount a volume or place files at repo_root/data.
@@ -54,7 +55,10 @@ async def haystack_deep_research_agent_workflow(
     config: HaystackDeepResearchWorkflowConfig, builder: Builder
 ):
     """
-    Main workflow that creates and returns the deep research agent
+    Main workflow that creates and returns the deep research agent.
+
+    Uses top-level `llms` configuration via builder to instantiate Haystack NvidiaChatGenerator
+    for both the agent and RAG tool, per review suggestions.
     """
     from haystack.components.agents import Agent
     from haystack.utils import Secret
@@ -71,7 +75,7 @@ async def haystack_deep_research_agent_workflow(
     logger.info(f"Starting Haystack Deep Research Agent workflow with config: {config}")
 
     # Create search tool
-    search_tool = create_search_tool()
+    search_tool = create_search_tool(top_k=config.search_top_k)
 
     # Create document store
     document_store = OpenSearchDocumentStore(
@@ -85,22 +89,34 @@ async def haystack_deep_research_agent_workflow(
             document_store=document_store, data_dir=config.data_dir, logger=logger
         )
 
-    # Create RAG tool
+    def _nim_to_haystack_generator(cfg: NIMModelConfig) -> NvidiaChatGenerator:
+        return NvidiaChatGenerator(
+            model=cfg.model_name,
+            api_base_url=cfg.base_url,
+            api_key=Secret.from_env_var("NVIDIA_API_KEY"),
+        )
+
+    # Instantiate LLMs via builder configs (expecting NIM)
+    rag_llm_cfg = builder.get_llm_config("rag_llm")
+    agent_llm_cfg = builder.get_llm_config("agent_llm")
+
+    if not isinstance(rag_llm_cfg, NIMModelConfig):
+        raise TypeError("llms.rag_llm must be of type 'nim'.")
+    if not isinstance(agent_llm_cfg, NIMModelConfig):
+        raise TypeError("llms.agent_llm must be of type 'nim'.")
+
+    rag_generator = _nim_to_haystack_generator(rag_llm_cfg)
     rag_tool, _ = create_rag_tool(
         document_store=document_store,
-        rag_model=config.rag_model,
-        nvidia_api_url=config.nvidia_api_url,
-        secret_provider=Secret.from_env_var("NVIDIA_API_KEY"),
-        top_k=15,
+        top_k=config.rag_top_k,
+        generator=rag_generator,
     )
 
     # Create the agent
+    agent_generator = _nim_to_haystack_generator(agent_llm_cfg)
+
     agent = Agent(
-        chat_generator=NvidiaChatGenerator(
-            model=config.agent_model,
-            api_base_url=config.nvidia_api_url,
-            api_key=Secret.from_env_var("NVIDIA_API_KEY"),
-        ),
+        chat_generator=agent_generator,
         tools=Toolset(tools=[search_tool, rag_tool]),
         system_prompt=config.system_prompt,
         exit_conditions=["text"],
@@ -113,14 +129,13 @@ async def haystack_deep_research_agent_workflow(
 
     async def _response_fn(input_message: str) -> str:
         """
-        Process the input message and generate a research response
-        Implements the exact logic from the notebook
+        Process the input message and generate a research response.
 
         Args:
             input_message: The user's research question
 
         Returns:
-            Comprehensive research report
+            Comprehensive research report.
         """
         try:
             logger.info(f"Processing research query: {input_message}")

examples/basic/frameworks/haystack_deep_research_agent/tests/test_haystack_deep_research_agent.py

@@ -82,10 +82,13 @@ def test_config_yaml_loads_and_has_keys() -> None:
     assert "_type: haystack_deep_research_agent" in text
     # key fields expected
     for key in [
-        "agent_model:",
-        "rag_model:",
-        "nvidia_api_url:",
+        "llms:",
+        "rag_llm:",
+        "agent_llm:",
+        "workflow:",
         "max_agent_steps:",
+        "search_top_k:",
+        "rag_top_k:",
         "opensearch_url:",
         "index_on_startup:",
         "data_dir:",