From e12e22439aa456166433b81534ef0c2afa1ca191 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Tue, 18 Aug 2026 22:06:01 -0700 Subject: [PATCH 01/13] feat: integrate with nemo relay for observability Signed-off-by: Chantal D Gama Rose --- .../skills/aiq-configure-workflow/SKILL.md | 2 +- .../assets/config-scaffold.yml | 3 + .../references/composing-config.md | 23 +- .../references/config-schema.md | 2 - .../references/env-vars.md | 8 +- .gitignore | 4 + .secrets.baseline | 4 +- README.md | 4 +- configs/config_cli_default.yml | 35 +- configs/config_domain_routing_and_skills.yml | 17 +- configs/config_frontier_models.yml | 12 +- configs/config_mcp.yml | 20 +- configs/config_openshell.yml | 24 +- configs/config_web_azure_ai_search.yml | 28 +- configs/config_web_default_guardrails.yml | 24 +- configs/config_web_default_llamaindex.yml | 46 +- configs/config_web_frag.yml | 24 +- configs/config_web_frag_mcp_auth.yml | 24 +- configs/config_web_opensearch.yml | 24 +- .../config_web_default_with_pricing.yml | 249 ++++++ configs/nemo_relay/relay_pricing_catalog.json | 39 + docs/source/architecture/agents/clarifier.md | 2 - .../architecture/agents/deep-researcher.md | 2 - .../architecture/agents/intent-classifier.md | 2 - docs/source/contributing/testing.md | 2 +- .../customization/configuration-reference.md | 47 +- docs/source/deployment/index.md | 2 +- docs/source/deployment/kubernetes.md | 3 - docs/source/deployment/observability.md | 537 +++++++----- docs/source/deployment/production.md | 6 +- .../benchmarks/deep-research-bench.md | 34 +- docs/source/examples/cli-with-local-nims.md | 7 - .../examples/full-pipeline-llamaindex.md | 5 - docs/source/examples/full-pipeline-web.md | 7 - docs/source/integration/index.md | 1 + docs/source/profiling/index.md | 62 +- docs/source/resources/troubleshooting.md | 33 +- .../aiq_api/src/aiq_api/auth/request_trace.py | 8 +- frontends/aiq_api/src/aiq_api/auth/utils.py | 3 + frontends/aiq_api/src/aiq_api/jobs/runner.py | 144 ++- frontends/aiq_api/src/aiq_api/jobs/submit.py | 72 +- .../aiq_api/src/aiq_api/jobs/telemetry.py | 140 --- .../benchmarks/deepresearch_bench/README.md | 33 +- .../configs/config_deep_research_bench.yml | 11 +- .../config_deep_research_bench_profiling.yml | 11 +- .../configs/config_tokenomics_pricing.yml | 29 +- .../configs/config_deepsearch_qa.yml | 7 +- .../freshqa/configs/config_full_workflow.yml | 5 +- .../configs/config_shallow_research_only.yml | 7 +- frontends/cli/cli.py | 8 - mcp/uv.lock | 73 +- pyproject.toml | 1 + scripts/start_cli.sh | 6 +- src/aiq_agent/agents/chat_researcher/agent.py | 19 +- .../nodes/intent_classifier.py | 16 +- .../agents/chat_researcher/register.py | 60 +- src/aiq_agent/agents/clarifier/agent.py | 19 +- src/aiq_agent/agents/clarifier/register.py | 13 +- src/aiq_agent/agents/deep_researcher/agent.py | 15 +- .../deep_researcher/custom_middleware.py | 21 +- .../agents/deep_researcher/factory.py | 74 +- .../agents/deep_researcher/register.py | 12 +- .../agents/deep_researcher/tools/research.py | 67 +- .../tools/source_tool_batching.py | 5 +- src/aiq_agent/agents/report_rewriter/agent.py | 36 +- .../agents/shallow_researcher/agent.py | 35 +- .../agents/shallow_researcher/register.py | 10 +- src/aiq_agent/common/__init__.py | 17 - src/aiq_agent/common/callbacks.py | 9 +- src/aiq_agent/relay/__init__.py | 32 + src/aiq_agent/relay/bootstrap.py | 62 ++ src/aiq_agent/relay/config.py | 177 ++++ src/aiq_agent/relay/logging.py | 171 ++++ src/aiq_agent/relay/privacy.py | 97 +++ src/aiq_agent/relay/runtime.py | 384 ++++++++ src/aiq_agent/tokenomics/README.md | 45 +- src/aiq_agent/tokenomics/__init__.py | 2 +- src/aiq_agent/tokenomics/atof_adapter.py | 273 ++++++ src/aiq_agent/tokenomics/nat_adapter.py | 319 ------- src/aiq_agent/tokenomics/pricing.py | 4 +- src/aiq_agent/tokenomics/profile.py | 4 +- src/aiq_agent/tokenomics/report/__init__.py | 12 +- src/aiq_agent/tokenomics/report/__main__.py | 4 +- .../nodes/test_intent_classifier.py | 3 +- .../aiq_agent/agents/clarifier/test_agent.py | 46 +- .../agents/clarifier/test_register.py | 3 - .../agents/deep_researcher/test_agent.py | 9 +- .../deep_researcher/test_custom_middleware.py | 24 +- .../agents/deep_researcher/test_factory.py | 1 + .../agents/test_config_observability.py | 24 + tests/aiq_agent/common/test_callbacks.py | 49 +- tests/aiq_agent/common/test_common_init.py | 68 -- tests/aiq_agent/jobs/test_runner.py | 84 +- tests/aiq_agent/jobs/test_telemetry.py | 241 ----- tests/conftest.py | 22 + tests/test_relay_runtime.py | 820 ++++++++++++++++++ tests/tokenomics/test_atof_adapter.py | 162 ++++ tests/tokenomics/test_nat_adapter.py | 205 ----- uv.lock | 73 +- 99 files changed, 3686 insertions(+), 2088 deletions(-) create mode 100644 configs/nemo_relay/config_web_default_with_pricing.yml create mode 100644 configs/nemo_relay/relay_pricing_catalog.json delete mode 100644 frontends/aiq_api/src/aiq_api/jobs/telemetry.py create mode 100644 src/aiq_agent/relay/__init__.py create mode 100644 src/aiq_agent/relay/bootstrap.py create mode 100644 src/aiq_agent/relay/config.py create mode 100644 src/aiq_agent/relay/logging.py create mode 100644 src/aiq_agent/relay/privacy.py create mode 100644 src/aiq_agent/relay/runtime.py create mode 100644 src/aiq_agent/tokenomics/atof_adapter.py delete mode 100644 src/aiq_agent/tokenomics/nat_adapter.py create mode 100644 tests/aiq_agent/agents/test_config_observability.py delete mode 100644 tests/aiq_agent/jobs/test_telemetry.py create mode 100644 tests/conftest.py create mode 100644 tests/test_relay_runtime.py create mode 100644 tests/tokenomics/test_atof_adapter.py delete mode 100644 tests/tokenomics/test_nat_adapter.py diff --git a/.agents/skills/aiq-configure-workflow/SKILL.md b/.agents/skills/aiq-configure-workflow/SKILL.md index 1eaf986af..880eb69cd 100644 --- a/.agents/skills/aiq-configure-workflow/SKILL.md +++ b/.agents/skills/aiq-configure-workflow/SKILL.md @@ -1,6 +1,6 @@ --- name: aiq-configure-workflow -description: Use when composing, adapting, or validating an AI-Q workflow YAML under configs/ — selecting a shipped profile, enabling tools and data_source_registry sources, wiring agents and the chat_deepresearcher_agent workflow, configuring general.telemetry (Phoenix, LangSmith, Weave, OTEL) and general.front_end aiq_api settings, and pre-flighting cross-references before deploy or serve. Hand off deploy to aiq-deploy, live research to aiq-research, prompt/model edits to aiq-customize-prompts-models, and new source code to aiq-add-tool or aiq-add-data-source. +description: Use when composing, adapting, or validating an AI-Q workflow YAML under configs/ — selecting a shipped profile, enabling tools and data_source_registry sources, wiring agents and the chat_deepresearcher_agent workflow, configuring NeMo Relay observability and general.front_end aiq_api settings, and pre-flighting cross-references before deploy or serve. Hand off deploy to aiq-deploy, live research to aiq-research, prompt/model edits to aiq-customize-prompts-models, and new source code to aiq-add-tool or aiq-add-data-source. license: Apache-2.0 compatibility: Claude Code, Codex, Cursor, OpenCode, and Agent Skills-compatible tools. metadata: diff --git a/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml b/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml index a987d763c..be8b93381 100644 --- a/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml +++ b/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml @@ -91,3 +91,6 @@ workflow: enable_escalation: true enable_clarifier: false # set true and add clarifier_agent under functions: to enable checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] diff --git a/.agents/skills/aiq-configure-workflow/references/composing-config.md b/.agents/skills/aiq-configure-workflow/references/composing-config.md index 6d03c5bbd..e49bda3cf 100644 --- a/.agents/skills/aiq-configure-workflow/references/composing-config.md +++ b/.agents/skills/aiq-configure-workflow/references/composing-config.md @@ -54,18 +54,14 @@ general: level: INFO # DEBUG | INFO | WARNING | ERROR ``` -**Tracing** — enable under `general.telemetry.tracing` (exporters can coexist). -Uncomment the matching block in any `config_web_*.yml` or copy from -`docs/source/deployment/observability.md`: +**Observability** — configure NeMo Relay under `workflow.relay`. Relay logging, +ATOF, and redaction are enabled by default. OTEL is opt-in; uncomment the Relay +OpenInference endpoint in a default config to send traces to Phoenix. Keep +`workflow.relay.pricing.sources: []` unless the workflow intentionally loads an +audited catalog. See `docs/source/deployment/observability.md`. -| Backend | YAML `_type` | Setup notes | -|---------|--------------|-------------| -| Phoenix | `phoenix` | `phoenix serve`; set `endpoint`, `project` | -| LangSmith | `langsmith` or env-only | `LANGCHAIN_TRACING_V2`, `LANGCHAIN_API_KEY`, `LANGCHAIN_PROJECT` | -| Weave | `weave` | `WANDB_API_KEY`; `project`, optional `redact_pii` | -| OpenTelemetry | `otelcollector_redaction` | `endpoint`; redaction + batch fields | - -`verbose: true` on `workflow:` or agents adds console detail without a tracer. +`workflow.relay.logging` controls the console subscriber; agent and workflow +configs do not have separate verbose switches. ### `front_end` (`aiq_api`) @@ -165,8 +161,8 @@ and feature guides under `docs/source/customization/`. | `_type` | Key options to tune | Doc anchor | |---------|---------------------|------------| -| `intent_classifier` | `llm`, `tools`, `llm_timeout`, `verbose` | `configuration-reference.md` § `intent_classifier` | -| `clarifier_agent` | `llm`, `max_turns`, `exclude_tools`, `verbose` | § `clarifier_agent` | +| `intent_classifier` | `llm`, `tools`, `llm_timeout` | `configuration-reference.md` § `intent_classifier` | +| `clarifier_agent` | `llm`, `max_turns`, `exclude_tools` | § `clarifier_agent` | | `shallow_research_agent` | `llm`, `max_llm_turns`, `max_tool_iterations`, `exclude_tools` | § `shallow_research_agent` | | `deep_research_agent` | role LLMs, `exclude_tools`, `enable_source_router`, `domain_catalog_path`, `enable_citation_verification`, `skills`, `sandbox`, concurrency caps | § `deep_research_agent` | @@ -191,7 +187,6 @@ workflow: use_async_deep_research: true # needs general.front_end max_history: 20 checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - verbose: true ``` Full defaults table: `configuration-reference.md` § `workflow`. diff --git a/.agents/skills/aiq-configure-workflow/references/config-schema.md b/.agents/skills/aiq-configure-workflow/references/config-schema.md index 76fe5231a..8697bd0da 100644 --- a/.agents/skills/aiq-configure-workflow/references/config-schema.md +++ b/.agents/skills/aiq-configure-workflow/references/config-schema.md @@ -26,8 +26,6 @@ uv run python .agents/skills/aiq-configure-workflow/scripts/validate_config.py < - No `data_source_registry` - `requires_auth: true` on a source (confirm MCP/OAuth wiring) - `use_async_deep_research: true` without `general.front_end` -- LangSmith tracing without `LANGCHAIN_API_KEY` -- Weave tracing without `WANDB_API_KEY` ## Env checklist diff --git a/.agents/skills/aiq-configure-workflow/references/env-vars.md b/.agents/skills/aiq-configure-workflow/references/env-vars.md index 48b25e123..2b966e7c3 100644 --- a/.agents/skills/aiq-configure-workflow/references/env-vars.md +++ b/.agents/skills/aiq-configure-workflow/references/env-vars.md @@ -31,15 +31,13 @@ Canonical references: | `SEARCHAPI_API_KEY` | SearchAPI paper search | | `RAG_SERVER_URL`, `RAG_INGEST_URL` | Foundational RAG profiles | -## Web API, auth, and tracing +## Web API, auth, and Relay correlation | Variable | When needed | |----------|-------------| | `REQUIRE_AUTH` | Enforce API authentication. Requires validator registration. | -| `AIQ_TRACE_USER_IDENTITY_MODE`, `AIQ_TRACE_USER_IDENTITY_HMAC_SECRET` | User identity tagging for NAT spans. | -| `AIQ_TRACE_CLIENT_ID_MODE`, `AIQ_TRACE_CLIENT_ID_HMAC_SECRET`, `AIQ_TRACE_CLIENT_IP_HEADERS` | Client tagging for NAT spans. | -| `LANGCHAIN_TRACING_V2`, `LANGCHAIN_API_KEY`, `LANGCHAIN_PROJECT` | LangSmith tracing. | -| `WANDB_API_KEY` | Weave tracing. | +| `AIQ_TRACE_USER_IDENTITY_MODE`, `AIQ_TRACE_USER_IDENTITY_HMAC_SECRET` | User identity tagging for Relay-exported spans. | +| `AIQ_TRACE_CLIENT_ID_MODE`, `AIQ_TRACE_CLIENT_ID_HMAC_SECRET`, `AIQ_TRACE_CLIENT_IP_HEADERS` | Client tagging for Relay-exported spans. | ## Sandbox and artifact storage diff --git a/.gitignore b/.gitignore index da134e0bc..963acd54d 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,10 @@ env/ *.log logs/ +# NeMo Relay local observability output (regeneratable; never commit payloads) +/.nemo-relay/ +/relay/ + # Testing .pytest_cache/ .coverage diff --git a/.secrets.baseline b/.secrets.baseline index a0f4964c9..2627a5858 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -324,7 +324,7 @@ "filename": "tests/aiq_agent/common/test_common_init.py", "hashed_secret": "9d4e1e23bd5b727046a9e3b4b7db57bd8d6ee684", "is_verified": false, - "line_number": 230 + "line_number": 162 } ], "tests/knowledge_layer_tests/run_llamaindex.py": [ @@ -355,5 +355,5 @@ } ] }, - "generated_at": "2026-08-17T17:33:50Z" + "generated_at": "2026-08-19T05:02:34Z" } diff --git a/README.md b/README.md index 1571cc9c7..1b8300e3e 100644 --- a/README.md +++ b/README.md @@ -72,7 +72,7 @@ The NVIDIA AI-Q Blueprint is a deployable research backend built on the [NVIDIA - **Expanded sources** — Paper search supports Serper, SerpAPI, and SearchAPI; You.com adds web, contents, general-research, and finance-research tools; Nimble adds configurable web search; focused profiles demonstrate DuckDuckGo news, Polymarket, OpenSearch, and Azure AI Search knowledge retrieval. - **Production API and auth** — REST endpoints, async job ownership, per-user OAuth-protected MCP sources, token validator entry points, and provider lifecycle hooks support authenticated deployments; a separate public MCP server exposes stateless research tools for trusted networks. - **Opt-in policy controls** — NeMo Guardrails middleware covers selected workflow and agent boundaries, and narrow application-level encryption can protect final async output plus selected artifact-event content. -- **Observability, profiling, and cost analysis** — NAT-exported async traces preserve task, named-agent, and model/tool hierarchy across concurrent researchers. Tokenomics reports combine profiler traces with pricing configuration for cost, latency, and cache analysis. +- **Observability, profiling, and cost analysis** — NeMo Relay preserves task, named-agent, LLM, and tool hierarchy across interactive turns and async researchers. ATOF and OTEL exports feed debugging and tokenomics reports for cost, latency, and cache analysis. - **Evaluation harnesses** — Built-in benchmarks (for example, FreshQA, DeepResearch) and evaluation scripts to measure quality and iterate on prompts and agent architecture. - **Frontend options** — Run through CLI, web UI, or async jobs. Refer to [Getting started](#getting-started) and [Ways to run the agents](#ways-to-run-the-agents). - **Deployment options** - Deployment assets for [Docker Compose](deploy/compose/) and [Helm](deploy/helm/deployment-k8s/); the repository source chart honors the Helm release namespace for every namespaced resource. @@ -478,7 +478,7 @@ For development, contribution, and documentation, refer to: - **[Knowledge Layer Setup](sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md)**: RAG backends and document ingestion - **[Agent Skills](docs/source/integration/agent-skills.md)**: Install the portable AI-Q research skill in compatible coding harnesses - **[Skills and Sandbox Example](docs/source/examples/skills-sandbox/index.md)**: Run deep research with built-in skills and Modal sandbox execution -- **[Profiling and Cost Analysis](docs/source/profiling/index.md)**: Generate tokenomics and latency reports from NAT profiler traces +- **[Profiling and Cost Analysis](docs/source/profiling/index.md)**: Generate tokenomics and latency reports from Relay ATOF traces - **[Docs index](docs/README.md)**: Full documentation list and component docs - **[Changelog](docs/source/resources/changelog.md)**: Version history and changes diff --git a/configs/config_cli_default.yml b/configs/config_cli_default.yml index 5bd6b89eb..0c6c7f771 100644 --- a/configs/config_cli_default.yml +++ b/configs/config_cli_default.yml @@ -10,14 +10,10 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -26,13 +22,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -41,11 +35,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -53,11 +45,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -65,8 +55,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false functions: # ========================================================================= @@ -133,7 +121,6 @@ functions: # exclude_tools: [] max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -161,3 +148,17 @@ workflow: enable_escalation: true enable_clarifier: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + # Uncomment to export Relay traces to a local Phoenix instance. + # observability: + # opentelemetry: + # enabled: true + # endpoints: + # - type: openinference + # endpoint: http://localhost:6006/v1/traces + # service_name: aiq-relay + # resource_attributes: + # openinference.project.name: aiq-relay + # deployment.environment: development + pricing: + sources: [] diff --git a/configs/config_domain_routing_and_skills.yml b/configs/config_domain_routing_and_skills.yml index 19f468c47..1579fa76f 100644 --- a/configs/config_domain_routing_and_skills.yml +++ b/configs/config_domain_routing_and_skills.yml @@ -10,10 +10,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -48,7 +44,7 @@ general: llms: nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -56,11 +52,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -68,8 +62,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: @@ -220,12 +212,13 @@ functions: # tools: omitted -> inherits all from data_source_registry exclude_tools: # Remove basic variant; deep uses advanced_web_search_tool - web_search_tool - verbose: true domain_catalog_path: configs/domain_catalogs/deep_research_domain_catalog.yml skills: deep_research_skills sandbox: deep_research_sandbox workflow: _type: deep_research_workflow - verbose: true use_async_deep_research: true + relay: + pricing: + sources: [] diff --git a/configs/config_frontier_models.yml b/configs/config_frontier_models.yml index 15144bafd..131a60776 100644 --- a/configs/config_frontier_models.yml +++ b/configs/config_frontier_models.yml @@ -14,10 +14,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -132,7 +128,6 @@ functions: intent_classifier: _type: intent_classifier llm: gpt_luna_intent_llm - verbose: true tools: - web_search_tool # - paper_search_tool # Uncomment if SERPER_API_KEY is set @@ -146,12 +141,10 @@ functions: - knowledge_search max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent llm: gpt_luna_shallow_llm - verbose: true tools: - web_search_tool - knowledge_search @@ -166,7 +159,6 @@ functions: researcher_llm: gpt_luna_agent_llm planner_llm: gpt_sol_agent_llm writer_llm: gpt_sol_writer_llm - verbose: true tools: # - paper_search_tool # Uncomment if SERPER_API_KEY is set - advanced_web_search_tool @@ -175,8 +167,10 @@ functions: workflow: _type: chat_deepresearcher_agent - verbose: true enable_escalation: true enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] diff --git a/configs/config_mcp.yml b/configs/config_mcp.yml index 71ff96e2d..fec4d41c4 100644 --- a/configs/config_mcp.yml +++ b/configs/config_mcp.yml @@ -27,7 +27,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -36,13 +36,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -51,11 +49,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -63,11 +59,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -75,8 +69,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false functions: data_sources: @@ -109,7 +101,6 @@ functions: llm: nemotron_ultra_llm max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -136,3 +127,6 @@ workflow: enable_escalation: true use_async_deep_research: false checkpoint_db: ${AIQ_CHECKPOINT_DB} + relay: + pricing: + sources: [] diff --git a/configs/config_openshell.yml b/configs/config_openshell.yml index d026d9ac1..4b059877f 100644 --- a/configs/config_openshell.yml +++ b/configs/config_openshell.yml @@ -29,7 +29,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -38,13 +38,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: NVIDIA API Catalog-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -53,11 +51,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -65,8 +61,6 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: _type: nim @@ -77,8 +71,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false summary_llm: _type: nim @@ -127,21 +119,18 @@ functions: intent_classifier: _type: intent_classifier llm: nemotron_lightning_intent_llm - verbose: true clarifier_agent: _type: clarifier_agent llm: nemotron_ultra_llm max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent llm: nemotron_lightning_agent_llm exclude_tools: - advanced_web_search_tool - verbose: true max_llm_turns: 10 max_tool_iterations: 5 @@ -198,14 +187,17 @@ functions: writer_llm: nemotron_ultra_writer_llm exclude_tools: - web_search_tool - verbose: true skills: deep_research_skills sandbox: deep_research_sandbox workflow: _type: chat_deepresearcher_agent - verbose: true enable_escalation: true enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + # Add file or inline catalog sources to emit cost estimates. Token usage + # is still captured when no catalog matches the configured model. + sources: [] diff --git a/configs/config_web_azure_ai_search.yml b/configs/config_web_azure_ai_search.yml index 14248537d..709728507 100644 --- a/configs/config_web_azure_ai_search.yml +++ b/configs/config_web_azure_ai_search.yml @@ -10,10 +10,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -48,7 +44,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -57,13 +53,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -72,11 +66,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -84,11 +76,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -96,8 +86,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: @@ -196,7 +184,6 @@ functions: llm: nemotron_lightning_intent_llm # tools: omitted -> inherits all from data_source_registry # exclude_tools: [] - verbose: true clarifier_agent: _type: clarifier_agent @@ -207,7 +194,6 @@ functions: max_turns: 3 enable_plan_approval: true log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -215,7 +201,6 @@ functions: # tools: omitted -> inherits all from data_source_registry exclude_tools: # Remove advanced variant; shallow uses web_search_tool - advanced_web_search_tool - verbose: true max_llm_turns: 10 max_tool_iterations: 5 @@ -230,12 +215,13 @@ functions: # tools: omitted -> inherits all from data_source_registry exclude_tools: # Remove basic variant; deep uses advanced_web_search_tool - web_search_tool - verbose: true workflow: _type: chat_deepresearcher_agent - verbose: true enable_escalation: true enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] diff --git a/configs/config_web_default_guardrails.yml b/configs/config_web_default_guardrails.yml index 781602bf6..b5daa37b3 100644 --- a/configs/config_web_default_guardrails.yml +++ b/configs/config_web_default_guardrails.yml @@ -29,7 +29,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -38,13 +38,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -53,11 +51,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -65,11 +61,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -77,8 +71,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false summary_llm: @@ -232,21 +224,18 @@ functions: intent_classifier: _type: intent_classifier llm: nemotron_lightning_intent_llm - verbose: true clarifier_agent: _type: clarifier_agent llm: nemotron_ultra_llm max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent llm: nemotron_lightning_agent_llm exclude_tools: - advanced_web_search_tool - verbose: true max_llm_turns: 10 max_tool_iterations: 5 @@ -259,14 +248,15 @@ functions: writer_llm: nemotron_ultra_writer_llm exclude_tools: - web_search_tool - verbose: true workflow: _type: chat_deepresearcher_agent - verbose: true enable_escalation: true enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] middleware: - workflow_guardrails diff --git a/configs/config_web_default_llamaindex.yml b/configs/config_web_default_llamaindex.yml index 2d6770293..ea144f68c 100644 --- a/configs/config_web_default_llamaindex.yml +++ b/configs/config_web_default_llamaindex.yml @@ -10,10 +10,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -48,7 +44,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -57,14 +53,12 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim - model_name: nvidia/nemotron-3.5-lightning-30b-a3b + _type: openai + model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} temperature: 0.2 @@ -72,11 +66,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -84,11 +76,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -96,8 +86,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: @@ -197,7 +185,6 @@ functions: llm: nemotron_lightning_intent_llm # tools: omitted -> inherits all from data_source_registry # exclude_tools: [] - verbose: true clarifier_agent: _type: clarifier_agent @@ -206,7 +193,6 @@ functions: # exclude_tools: [] max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -214,7 +200,6 @@ functions: # tools: omitted -> inherits all from data_source_registry exclude_tools: # Remove advanced variant; shallow uses web_search_tool - advanced_web_search_tool - verbose: true max_llm_turns: 10 max_tool_iterations: 5 @@ -229,12 +214,29 @@ functions: # tools: omitted -> inherits all from data_source_registry exclude_tools: # Remove basic variant; deep uses advanced_web_search_tool - web_search_tool - verbose: true workflow: _type: chat_deepresearcher_agent - verbose: true enable_escalation: true enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + observability: + atof: + enabled: true + output_directory: ./relay + filename: aiq-relay.atof.jsonl + mode: append + # Uncomment to export Relay traces to a local Phoenix instance. + # opentelemetry: + # enabled: true + # endpoints: + # - type: openinference + # endpoint: http://localhost:6006/v1/traces + # service_name: aiq-relay + # resource_attributes: + # openinference.project.name: aiq-relay + # deployment.environment: development + pricing: + sources: [] diff --git a/configs/config_web_frag.yml b/configs/config_web_frag.yml index b54f07cfc..b291836d5 100644 --- a/configs/config_web_frag.yml +++ b/configs/config_web_frag.yml @@ -12,10 +12,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -50,7 +46,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -59,13 +55,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -74,11 +68,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -86,11 +78,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -98,8 +88,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false functions: # ========================================================================= @@ -175,7 +163,6 @@ functions: # exclude_tools: [] max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -204,3 +191,6 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] diff --git a/configs/config_web_frag_mcp_auth.yml b/configs/config_web_frag_mcp_auth.yml index e48f67f83..034fd6261 100644 --- a/configs/config_web_frag_mcp_auth.yml +++ b/configs/config_web_frag_mcp_auth.yml @@ -19,10 +19,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -57,7 +53,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -66,13 +62,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -81,11 +75,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -93,11 +85,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -105,8 +95,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false functions: # ========================================================================= @@ -224,7 +212,6 @@ functions: # exclude_tools: [] max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -323,3 +310,6 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] diff --git a/configs/config_web_opensearch.yml b/configs/config_web_opensearch.yml index 81af40aac..40d275251 100644 --- a/configs/config_web_opensearch.yml +++ b/configs/config_web_opensearch.yml @@ -11,10 +11,6 @@ general: console: _type: console level: INFO - # tracing: - # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` - # _type: langsmith - # project: nvidia-aiq front_end: _type: aiq_api @@ -49,7 +45,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -58,13 +54,11 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -73,11 +67,9 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - chat_template_kwargs: - enable_thinking: true nemotron_ultra_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -85,11 +77,9 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - chat_template_kwargs: - enable_thinking: false nemotron_ultra_writer_llm: - _type: nim + _type: openai model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -97,8 +87,6 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - chat_template_kwargs: - enable_thinking: false functions: # ========================================================================= @@ -182,7 +170,6 @@ functions: # exclude_tools: [] max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent @@ -211,3 +198,6 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + pricing: + sources: [] diff --git a/configs/nemo_relay/config_web_default_with_pricing.yml b/configs/nemo_relay/config_web_default_with_pricing.yml new file mode 100644 index 000000000..4198cdfb4 --- /dev/null +++ b/configs/nemo_relay/config_web_default_with_pricing.yml @@ -0,0 +1,249 @@ +# This is the default configuration for the Web mode. +# It has the following features: +# - Knowledge retrieval using LlamaIndex +# - Web search and Paper search tools by default + +general: + use_uvloop: true + telemetry: + logging: + console: + _type: console + level: INFO + + front_end: + _type: aiq_api + runner_class: aiq_api.plugin.AIQAPIWorker + # ========================================================================= + # Knowledge API is automatically enabled when knowledge_retrieval function + # is configured + # ========================================================================= + # Async Job API Settings + # ========================================================================= + # Async job infrastructure database (NAT JobStore + EventStore) + # Used by: /v1/jobs/async routes, SSE streaming, job status persistence + # Requires async driver for SQLite (aiosqlite) or PostgreSQL (asyncpg) + # Environment overrides: + # - NAT_JOB_STORE_DB_URL (direct override) + # - NAT_JOB_STORE_DB_URL_DEV / NAT_JOB_STORE_DB_URL_PROD (via NAT_ENV) + db_url: ${NAT_JOB_STORE_DB_URL:-sqlite+aiosqlite:///./jobs.db} + # Job expiry - how long completed jobs stay in database before cleanup + expiry_seconds: 86400 # 24 hours (min: 600, max: 604800/7 days) + cors: + allow_origin_regex: 'http://localhost(:\d+)?|http://127.0.0.1(:\d+)?' + allow_methods: + - GET + - POST + - DELETE + - OPTIONS + allow_headers: + - "*" + allow_credentials: true + expose_headers: + - "*" + +llms: + nemotron_lightning_intent_llm: + _type: nim + model_name: nvidia/nemotron-3.5-lightning-30b-a3b + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.1 + top_p: 0.9 + max_tokens: 1024 + num_retries: 5 + parallel_tool_calls: false + + # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. + # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. + nemotron_lightning_agent_llm: + _type: nim + model_name: nvidia/nemotron-3-ultra-550b-a55b + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.2 + top_p: 0.7 + max_tokens: 8192 + num_retries: 5 + parallel_tool_calls: false + # chat_template_kwargs: + # enable_thinking: true + + nemotron_ultra_llm: + _type: nim + model_name: nvidia/nemotron-3-ultra-550b-a55b + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.2 + top_p: 0.7 + max_tokens: 16384 + num_retries: 5 + # chat_template_kwargs: + # enable_thinking: false + + nemotron_ultra_writer_llm: + _type: nim + model_name: nvidia/nemotron-3-ultra-550b-a55b + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.2 + top_p: 0.7 + max_tokens: 32768 + num_retries: 5 + # chat_template_kwargs: + # enable_thinking: false + + # LLM for document summaries (required when generate_summary: true) + summary_llm: + _type: nim + model_name: google/gemma-4-31b-it + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + temperature: 0.1 + max_tokens: 100 + +functions: + # ========================================================================= + # Data Source Registry + # ========================================================================= + # Central registry that controls: + # 1. UI toggles — each source appears as an on/off switch in the frontend + # 2. Per-message filtering — users can select active sources per request + # 3. Tool auto-inheritance — agents with no explicit `tools` list receive + # every tool listed here (use `exclude_tools` on agents to specialize) + # + # Adding a new tool or MCP function group? Just add it here — all agents + # pick it up automatically. No per-agent config changes needed. + # + # Source entry fields: + # id — Unique key used in API payloads and filtering + # name — Display name shown in the UI + # description — Human-readable description shown in the UI + # tools — List of NAT function names or function group names + # requires_auth — (default: false) If true, the UI greys out this source + # until the user signs in. Use for sources that need + # user-level OAuth tokens (e.g., enterprise SSO). + # Sources using backend API keys (Tavily, Serper) should + # leave this false. + # default_enabled — (default: true) Whether enabled by default + # + # See docs/source/customization/tools-and-sources.md for full details. + # ========================================================================= + data_sources: + _type: data_source_registry + sources: + - id: web_search + name: "Web Search" + description: "Search the web for real-time information." + tools: + - web_search_tool + - advanced_web_search_tool + - id: knowledge_layer + name: "Knowledge Base" + description: "Search uploaded documents and files." + tools: + - knowledge_search + # Uncomment when paper_search_tool is enabled (requires SERPER_API_KEY) + # - id: paper_search + # name: "Academic Papers" + # description: "Search academic papers and scientific publications." + # tools: + # - paper_search_tool + + web_search_tool: + _type: tavily_web_search + max_results: 5 + max_content_length: 1000 + + advanced_web_search_tool: + _type: tavily_web_search + max_results: 2 + advanced_search: true + + # Knowledge Retrieval (see sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md) + knowledge_search: + _type: knowledge_retrieval + backend: llamaindex + collection_name: ${COLLECTION_NAME:-test_collection} + generate_summary: true + summary_model: summary_llm # Required when generate_summary: true + summary_db: ${AIQ_SUMMARY_DB:-sqlite+aiosqlite:///./summaries.db} + top_k: 5 + chroma_dir: ${AIQ_CHROMA_DIR:-/tmp/chroma_data} + + # Paper Search (optional - requires SERPER_API_KEY) + # Uncomment the block below and set SERPER_API_KEY to enable academic paper search. + # paper_search_tool: + # _type: paper_search + # max_results: 5 + # serper_api_key: ${SERPER_API_KEY} + + # ========================================================================= + # Agents + # ========================================================================= + # Tool inheritance: agents with no `tools` list inherit ALL tools from the + # data_source_registry above. Use `exclude_tools` to remove specific tools + # from an agent (e.g., give shallow the basic search, deep the advanced). + # To bypass auto-inherit entirely, set an explicit `tools` list. + # ========================================================================= + intent_classifier: + _type: intent_classifier + llm: nemotron_lightning_intent_llm + # tools: omitted -> inherits all from data_source_registry + # exclude_tools: [] + + clarifier_agent: + _type: clarifier_agent + llm: nemotron_ultra_llm + # tools: omitted -> inherits all from data_source_registry + # exclude_tools: [] + max_turns: 3 + log_response_max_chars: 2000 + + shallow_research_agent: + _type: shallow_research_agent + llm: nemotron_lightning_agent_llm + # tools: omitted -> inherits all from data_source_registry + exclude_tools: # Remove advanced variant; shallow uses web_search_tool + - advanced_web_search_tool + max_llm_turns: 10 + max_tool_iterations: 5 + + deep_research_agent: + _type: deep_research_agent + enable_citation_verification: true + orchestrator_llm: nemotron_ultra_llm + source_router_llm: nemotron_ultra_llm + researcher_llm: nemotron_ultra_llm + planner_llm: nemotron_ultra_llm + writer_llm: nemotron_ultra_writer_llm + # tools: omitted -> inherits all from data_source_registry + exclude_tools: # Remove basic variant; deep uses advanced_web_search_tool + - web_search_tool + +workflow: + _type: chat_deepresearcher_agent + enable_escalation: true + enable_clarifier: true + use_async_deep_research: true + checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + observability: + atof: + enabled: true + output_directory: ./relay + filename: aiq-relay.atof.jsonl + mode: append + opentelemetry: + enabled: true + endpoints: + - type: openinference + endpoint: http://localhost:6006/v1/traces + service_name: aiq-relay + resource_attributes: + openinference.project.name: aiq-relay + deployment.environment: development + pricing: + sources: + - type: file + path: configs/nemo_relay/relay_pricing_catalog.json diff --git a/configs/nemo_relay/relay_pricing_catalog.json b/configs/nemo_relay/relay_pricing_catalog.json new file mode 100644 index 000000000..ef3f33f41 --- /dev/null +++ b/configs/nemo_relay/relay_pricing_catalog.json @@ -0,0 +1,39 @@ +{ + "version": 1, + "entries": [ + { + "provider": "nvidia", + "model_id": "nvidia/nemotron-3-ultra-550b-a55b", + "aliases": [], + "currency": "USD", + "unit": "per_token", + "rates": { + "input_per_million": 0.0, + "output_per_million": 0.0, + "cache_read_per_million": 0.0 + }, + "prompt_cache": { + "read_accounting": "included_in_prompt_tokens" + }, + "pricing_as_of": "2026-08-18", + "pricing_source": "NVIDIA-hosted API access used by this AI-Q profile; zero is access-path-specific, not a self-hosting cost estimate" + }, + { + "provider": "nvidia", + "model_id": "nvidia/nemotron-3.5-lightning-30b-a3b", + "aliases": [], + "currency": "USD", + "unit": "per_token", + "rates": { + "input_per_million": 0.0, + "output_per_million": 0.0, + "cache_read_per_million": 0.0 + }, + "prompt_cache": { + "read_accounting": "included_in_prompt_tokens" + }, + "pricing_as_of": "2026-08-18", + "pricing_source": "NVIDIA-hosted API access used by this AI-Q profile; zero is access-path-specific, not a self-hosting cost estimate" + } + ] +} diff --git a/docs/source/architecture/agents/clarifier.md b/docs/source/architecture/agents/clarifier.md index 7bc33c608..a1b685d36 100644 --- a/docs/source/architecture/agents/clarifier.md +++ b/docs/source/architecture/agents/clarifier.md @@ -96,7 +96,6 @@ Configured through `ClarifierConfig` (NeMo Agent Toolkit type name: `clarifier_a | `tools` | `list[FunctionRef \| FunctionGroupRef]` | `[]` | Tools for context gathering (for example, web search) | | `max_turns` | `int` | `3` | Maximum clarification Q&A turns before auto-completing | | `log_response_max_chars` | `int` | `2000` | Maximum characters to log from LLM responses | -| `verbose` | `bool` | `false` | Enable verbose logging with `VerboseTraceCallback` | **Example YAML:** @@ -108,7 +107,6 @@ functions: tools: - web_search_tool max_turns: 3 - verbose: true ``` ## Prompt Templates diff --git a/docs/source/architecture/agents/deep-researcher.md b/docs/source/architecture/agents/deep-researcher.md index 2093c00d7..f45ca31f0 100644 --- a/docs/source/architecture/agents/deep-researcher.md +++ b/docs/source/architecture/agents/deep-researcher.md @@ -229,7 +229,6 @@ for configuration details. | `sandbox` | `FunctionRef`, inline `deep_research_sandbox`, or `None` | `None` | Optional sandbox profile for DeepAgents `execute` support | | `enable_citation_verification` | `bool` | `true` | Verify generated citations against captured sources after final report extraction | | `resource_limits` | `DeepResearchResourceLimits` | hard ceilings | Per-job request, graph-time, plan, report, shared-state, note, todo, query, and source-call budgets; configurable downward only | -| `verbose` | `bool` | `true` | Enable detailed logging | **Example YAML:** @@ -248,7 +247,6 @@ functions: resource_limits: max_research_queries: 20 max_source_tool_calls: 100 - verbose: true tools: - web_search_tool ``` diff --git a/docs/source/architecture/agents/intent-classifier.md b/docs/source/architecture/agents/intent-classifier.md index 8a8982792..ab6a28f3a 100644 --- a/docs/source/architecture/agents/intent-classifier.md +++ b/docs/source/architecture/agents/intent-classifier.md @@ -104,7 +104,6 @@ Configured through `IntentClassifierConfig` (NeMo Agent Toolkit type name: `inte | --------- | ---- | ------- | ----------- | | `llm` | `LLMRef` | required | LLM to use for classification | | `tools` | `list[FunctionRef \| FunctionGroupRef]` | `[]` | Tool references; their names and descriptions are shown to the LLM so it can assess query complexity | -| `verbose` | `bool` | `false` | Enable verbose logging using `VerboseTraceCallback` | | `llm_timeout` | `float` | `90` | Timeout in seconds for the LLM call | **Example YAML:** @@ -116,7 +115,6 @@ functions: llm: nemotron_llm tools: - web_search_tool - verbose: true llm_timeout: 90 ``` diff --git a/docs/source/contributing/testing.md b/docs/source/contributing/testing.md index 9097e59b7..19e295e20 100644 --- a/docs/source/contributing/testing.md +++ b/docs/source/contributing/testing.md @@ -23,6 +23,6 @@ Refer to each benchmark's README for details. The [Customization guide](../custo ## Debugging -- **Verbose logging:** `./scripts/start_cli.sh --verbose` or set `verbose: true` in workflow config. +- **Relay logging:** enabled by `workflow.relay.logging`; inspect the console subscriber and configured ATOF/OTEL destinations. - **Phoenix tracing:** Start `uvx --from arize-phoenix phoenix serve`, run the agent with Phoenix tracing enabled in config, then open `http://localhost:6006`. - **Common issues:** Import errors -- ensure `uv pip install -e .`; auth -- check env vars; tool not found -- check config; pre-commit cache -- `pre-commit clean`. diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index e631dad1c..b854ccb2c 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -56,11 +56,6 @@ general: console: _type: console level: INFO # DEBUG, INFO, WARNING, ERROR - tracing: - phoenix: # Optional: Phoenix observability - _type: phoenix - endpoint: http://localhost:6006/v1/traces - project: dev front_end: # Only for web/API mode _type: aiq_api runner_class: aiq_api.plugin.AIQAPIWorker @@ -79,13 +74,13 @@ general: | `use_uvloop` | `bool` | `false` | Enable uvloop for improved async I/O performance. Recommended for web mode. | | `telemetry.logging.console._type` | `str` | `console` | Logging backend type. | | `telemetry.logging.console.level` | `str` | `INFO` | Log level: `DEBUG`, `INFO`, `WARNING`, `ERROR`. | -| `telemetry.tracing` | `object` | -- | Optional tracing configuration (Phoenix, OpenTelemetry). | | `front_end._type` | `str` | -- | Front-end type. Use `aiq_api` for the web API server. Omit for CLI mode. | | `front_end.db_url` | `str` | `sqlite+aiosqlite:///./jobs.db` | Database URL for async job persistence. | | `front_end.expiry_seconds` | `int` | `86400` | How long completed jobs remain in the database (seconds). | | `front_end.cors` | `object` | -- | CORS settings for the API server. | -For `aiq_api`, request tag enrichment for NAT-exported spans is configured via +Tracing is configured through `workflow.relay`, not `general.telemetry`. +For `aiq_api`, request tag enrichment for Relay-exported spans is configured via environment variables rather than YAML fields. Refer to `frontends/aiq_api/README.md` and the [Observability](../deployment/observability.md) guide for: @@ -390,7 +385,6 @@ functions: tools: - web_search_tool - paper_search_tool - verbose: true llm_timeout: 90 ``` @@ -398,7 +392,6 @@ functions: |-----------|------|---------|-------------| | `llm` | `str` | **required** | Reference to an LLM defined in `llms` section. | | `tools` | `list[str]` | `[]` | Tool references passed to the intent prompt for tool-awareness. | -| `verbose` | `bool` | `false` | Enable verbose logging with trace callbacks. | | `llm_timeout` | `float` | `90` | Timeout in seconds for the intent classification LLM call. | ### `clarifier_agent` @@ -414,7 +407,6 @@ functions: - web_search_tool max_turns: 3 log_response_max_chars: 2000 - verbose: true ``` | Parameter | Type | Default | Description | @@ -424,7 +416,6 @@ functions: | `exclude_tools` | `list[str]` | `[]` | Tool names to exclude when inheriting from the data source registry. | | `max_turns` | `int` | `3` | Maximum number of clarification Q&A turns before auto-completing. | | `log_response_max_chars` | `int` | `2000` | Maximum characters to log from LLM responses. | -| `verbose` | `bool` | `false` | Enable verbose logging. | ### `shallow_research_agent` @@ -497,7 +488,6 @@ functions: max_todo_items: 20 max_todo_item_chars: 2048 max_total_todo_chars: 10000 - verbose: true ``` | Parameter | Type | Default | Description | @@ -519,7 +509,6 @@ functions: | `max_concurrent_source_tool_calls` | `int` | `5` | Shared cap on concurrent source-tool calls across all researcher workers in the run. | | `max_source_tool_batch_size` | `int` | `4` | Maximum concrete inputs accepted by a batch-capable source-tool wrapper in one call. | | `resource_limits` | object | See below | Non-disableable per-job request, graph, state, and provider-call ceilings. Values may be reduced but cannot exceed the defaults. | -| `verbose` | `bool` | `true` | Enable verbose logging. | `resource_limits` is enforced in both synchronous and async-job construction: @@ -574,10 +563,34 @@ workflow: enable_clarifier: true use_async_deep_research: true max_history: 20 - verbose: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} + relay: + logging: true + observability: + enable_full_payloads: true + atof: {enabled: true, output_directory: ./relay, filename: aiq-relay.atof.jsonl, mode: append} + opentelemetry: + enabled: false + endpoints: + - type: openinference + endpoint: "http://localhost:6006/v1/traces" + service_name: aiq-relay + resource_attributes: {openinference.project.name: aiq-relay} + redaction: + enabled: true + request_privacy_attributes: [data, category_profile] + pricing: + enabled: true + sources: [] ``` +Default configs do not load a pricing catalog. The dedicated +`configs/nemo_relay/config_web_default_with_pricing.yml` example loads +deployment-specific rates from `configs/nemo_relay/relay_pricing_catalog.json`. +Its zero-dollar Nemotron entries describe the NVIDIA-hosted access path used by +the example; they are not estimates for self-hosted infrastructure. Review the +catalog when the provider offer or deployment changes. + | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `_type` | `str` | **required** | Workflow type. Use `chat_deepresearcher_agent` for the full pipeline. | @@ -585,8 +598,11 @@ workflow: | `enable_clarifier` | `bool` | `true` | Run the clarifier agent before deep research to gather user requirements. | | `use_async_deep_research` | `bool` | `false` | Submit deep research as an async background job (requires [Dask](https://www.dask.org/) scheduler). | | `max_history` | `int` | `20` | Maximum number of messages to keep in conversation history before trimming. | -| `verbose` | `bool` | `false` | Enable verbose logging. | | `checkpoint_db` | `str` | `./checkpoints.db` | SQLite path or PostgreSQL DSN for persistent conversation checkpoints. | +| `relay` | `object` | enabled defaults | NeMo Relay logging, Observability v3 ATOF/OTEL destinations, PII redaction, and pricing sources. Relay instrumentation itself has no workflow disable switch. See [Observability with NeMo Relay](../deployment/observability.md). | + +Relay configuration is strict: unknown nested fields and invalid OTLP endpoint +URLs fail workflow validation instead of being silently ignored. > **Note:** `interactive_auth` is a YAML-level field consumed by the CLI entry point (`start_cli.sh` / `aiq-research`), not a Pydantic field on `ChatDeepResearcherConfig`. It can be set in YAML config files but is not part of the workflow config class. @@ -689,7 +705,6 @@ functions: tools: - web_search_tool max_turns: 3 - verbose: true shallow_research_agent: # Fast single-pass research _type: shallow_research_agent diff --git a/docs/source/deployment/index.md b/docs/source/deployment/index.md index 825792dc6..fb7cc5222 100644 --- a/docs/source/deployment/index.md +++ b/docs/source/deployment/index.md @@ -45,7 +45,7 @@ All containerized deployments run the same three services: - **[Async Job Content Encryption](./content-encryption.md)** -- Configure encryption at rest for async final reports and selected artifact event content, including Vault Transit and static-key modes. -- **[Observability](./observability.md)** -- Tracing and monitoring with Phoenix, LangSmith, Weave, and OpenTelemetry. +- **[Observability](./observability.md)** -- NeMo Relay logging, ATOF traces, Phoenix OTEL export, redaction, and cost data. - **[Production Considerations](./production.md)** -- Guidance on managed databases, horizontal scaling, security hardening, monitoring, and resource requirements. diff --git a/docs/source/deployment/kubernetes.md b/docs/source/deployment/kubernetes.md index d67deebd4..425b77ae9 100644 --- a/docs/source/deployment/kubernetes.md +++ b/docs/source/deployment/kubernetes.md @@ -279,9 +279,6 @@ For complete examples with NGC-specific flags, refer to `deploy/helm/README.md` | `NIMBLE_API_KEY` | Nimble API key for web search | | `SERPER_API_KEY` | Serper API key for Google search | | `JINA_API_KEY` | Jina API key | -| `WANDB_API_KEY` | Weights & Biases API key | -| `NVIDIA_INFERENCE_API_KEY` | Alternative inference key (defaults to `NVIDIA_API_KEY`) | -| `INFERENCE_NVIDIA_API_KEY` | Alternative inference key (defaults to `NVIDIA_API_KEY`) | ### Updating secrets diff --git a/docs/source/deployment/observability.md b/docs/source/deployment/observability.md index 9f590c927..1be0ad684 100644 --- a/docs/source/deployment/observability.md +++ b/docs/source/deployment/observability.md @@ -3,320 +3,383 @@ SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES SPDX-License-Identifier: Apache-2.0 --> -# Observability +# Observability with NeMo Relay -The AI-Q blueprint supports multiple observability backends for tracing agent execution, LLM calls, tool invocations, and token usage. Choose the backend that best fits your workflow. For more details on available backends, refer to the [NVIDIA Agent Toolkit observability documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/run-workflows/observe/observe.html). +AI-Q uses [NeMo Relay](https://docs.nvidia.com/nemo/relay) as its observability runtime. -| Backend | Best For | Setup | -|---------|----------|-------| -| [Phoenix](#phoenix) | Local development, trace visualization | Run Phoenix server, add YAML config | -| [LangSmith](#langsmith) | LLM evaluation, prompt optimization, team collaboration | Set environment variables | -| [Weights & Biases Weave](#weights--biases-weave) | Experiment tracking, model monitoring | Set environment variables | -| [OpenTelemetry Collector](#opentelemetry-collector) | Production infrastructure, enterprise redaction | YAML config with OTEL endpoint | -| [Verbose Logging](#verbose-logging) | Quick debugging, no external services | CLI flag or YAML config | +Relay gives AI-Q users four complementary views: -## Async Deep Research Trace Hierarchy +- **Developer logs** show agent, LLM, and tool activity in the terminal. +- **ATOF JSONL** provides a durable, machine-readable event stream for debugging + and post-processing. +- **OpenTelemetry (OTEL)** exports the same scope tree to backends such as + Phoenix. +- **Pricing enrichment** attaches model cost data when a configured catalog + matches the observed model. -NAT-exported traces from the async job runner preserve the DeepAgents execution -hierarchy instead of flattening named agents beside their task and model spans. The root -workflow span uses the configured function name. A `task` tool is labeled with its -subagent type, and each outer DeepAgents chain receives a distinct named-agent span: +Logging, ATOF export, full observability payloads, and redaction are enabled by +default. OTEL export is opt-in so a default AI-Q installation does not attempt +to contact an observability server. Pricing is enabled with no catalog sources; +token usage is still recorded, but monetary cost is not estimated until you +configure prices. -```text -deep_research_agent -├── task: planner-agent -│ └── planner-agent -│ └── model -└── run_research_batch - ├── researcher-agent - │ └── model - └── researcher-agent - └── model -``` - -Parallel researchers remain separate children of the shared batch span. Structural -agent spans include `agent_id`, `agent_name`, and `span_role=agent`; start metadata also -records the LangChain parent run ID, and an error close records only the exception class -as `error_type`. These structural spans deliberately omit LangGraph input/output state so -they do not duplicate prompts or results. LLM and tool spans can still contain application -content, so configure the selected exporter's redaction controls for the deployment's -privacy requirements. +## Installation -This hierarchy describes NAT-exported async-job telemetry. Third-party tracing SDKs that -instrument LangChain directly can present a different tree. +NeMo Relay and the LangChain, LangGraph, and Deep Agents integrations are +installed with AI-Q: -## Phoenix - -[Phoenix](https://docs.arize.com/phoenix) provides a local UI for visualizing traces, inspecting LLM calls, and analyzing token usage and latency. It is the recommended backend for local development. - -### Setup +```bash +./scripts/setup.sh +``` -1. Start Phoenix in an isolated `uvx` environment. This installs `arize-phoenix` outside the AI-Q project environment - on the first run: +For an existing development checkout, synchronize the environment: - ```bash - uvx --from arize-phoenix phoenix serve - ``` +```bash +uv sync +``` - This launches the Phoenix UI at [http://localhost:6006](http://localhost:6006). +Verify the installed version: -2. Enable Phoenix tracing in your YAML config: +```bash +uv run python -c 'from importlib.metadata import version; print(version("nemo-relay"))' +``` - ```yaml - general: - telemetry: - tracing: - phoenix: - _type: phoenix - endpoint: http://localhost:6006/v1/traces - project: dev - ``` +AI-Q supports the Relay version range pinned in `pyproject.toml`. - The `project` field groups traces under a named project in the Phoenix UI. +## Configuration walkthrough -### What You Can Inspect +Relay is configured under the top-level AI-Q workflow: -- **Traces** -- Full agent execution trees showing orchestrator routing, tool calls, and LLM interactions. -- **Token usage** -- Per-call input/output token counts and costs. -- **Latency** -- Time spent in each step of the agent pipeline. -- **Tool calls** -- Arguments passed to and results returned from search tools, RAG retrieval, and other data sources. +```yaml +workflow: + _type: chat_deepresearcher_agent + relay: + logging: true + observability: + enable_full_payloads: true + atof: + enabled: true + output_directory: ./relay + filename: aiq-relay.atof.jsonl + mode: append + opentelemetry: + enabled: false + redaction: + enabled: true + pricing: + enabled: true + sources: [] +``` -## LangSmith +Most users can omit this block and use the defaults. Add only the settings that +you intend to change. -[LangSmith](https://smith.langchain.com/) provides cloud-hosted tracing, evaluation datasets, and prompt optimization for LangChain-based applications. It works automatically through the LangChain integration -- no YAML config changes are needed. +| Setting | Default | Purpose | +|---|---:|---| +| `logging` | `true` | Register AI-Q's Relay console subscriber. | +| `observability.enable_full_payloads` | `true` | Preserve supported inputs, outputs, metadata, and annotated usage for sanitization and export. | +| `observability.atof.enabled` | `true` | Write Relay events to ATOF JSONL. | +| `observability.atof.mode` | `append` | Preserve events across turns and async jobs. Use `overwrite` only for a single isolated run. | +| `observability.opentelemetry.enabled` | `false` | Export Relay scopes over OTEL when explicitly enabled. | +| `redaction.enabled` | `true` | Sanitize supported sensitive values before logs and exporters receive them. | +| `pricing.sources` | `[]` | Model pricing catalogs used to enrich Relay usage. | -### Setup +Relay configuration is strict. Unknown fields, invalid endpoint URLs, and an +invalid source shape fail workflow validation instead of being ignored. -1. Create an account at [smith.langchain.com](https://smith.langchain.com/) and generate an API key. +### Change the ATOF output path -2. Set the following environment variables in `deploy/.env`: +Set `output_directory` and `filename` independently: - ```bash - LANGCHAIN_TRACING_V2=true - LANGCHAIN_API_KEY=lsv2-... - LANGCHAIN_PROJECT=aiq-research - ``` +```yaml +workflow: + relay: + observability: + atof: + enabled: true + output_directory: ./observability/traces + filename: aiq-development.atof.jsonl + mode: append +``` - The `LANGCHAIN_PROJECT` variable groups traces under a named project. If omitted, traces go to the `default` project. +Relative directories are resolved from the working directory where AI-Q is +started. Use an absolute directory for containers, services, or async workers +when their working directories might differ: -3. Start the application as usual. All LangChain and LangGraph operations are traced automatically. No YAML config changes are required -- the LangChain SDK detects these environment variables at startup. +```yaml +output_directory: /var/lib/aiq/relay +``` -### What You Can Inspect +The resulting file is +`/var/lib/aiq/relay/aiq-development.atof.jsonl`. Ensure every worker can write +to the directory. Keep `mode: append` when multiple user turns or async jobs +share a file; choose a unique filename instead of `overwrite` when you need +per-run isolation. -- **Trace trees** -- Visualize the full agent execution including orchestrator decisions, tool calls, and LLM interactions. -- **LLM calls** -- Input prompts, output completions, token counts, and latency for every model call. -- **Evaluation** -- Build datasets from traced runs and evaluate agent quality over time. +## Inspect ATOF traces -## Weights & Biases Weave +By default, AI-Q appends events to: -[Weave](https://wandb.ai/site/weave) provides experiment tracking and trace -logging integrated with the Weights & Biases platform. Weave support is an -optional NAT extra and is not installed by default. +```text +relay/aiq-relay.atof.jsonl +``` -### Setup +Each line is one JSON event. Use `jq` to inspect it: -1. Install the exporter into your local environment: +```bash +# Follow new events while AI-Q runs. +tail -f relay/aiq-relay.atof.jsonl | jq -c . + +# Count scope starts by category. +jq -s ' + [.[] | select(.kind == "scope" and .scope_category == "start")] + | group_by(.category) + | map({category: .[0].category, count: length}) +' relay/aiq-relay.atof.jsonl + +# Show LLM usage recorded on completed LLM scopes. +jq -c ' + select(.category == "llm" and .scope_category == "end") + | { + model: .category_profile.annotated_response.model, + usage: .category_profile.annotated_response.usage, + status: .metadata["otel.status_code"] + } +' relay/aiq-relay.atof.jsonl + +# Find scopes that do not have exactly one start and one end. +jq -s ' + [.[] | select(.kind == "scope")] + | sort_by(.uuid) + | group_by(.uuid) + | map({ + uuid: .[0].uuid, + name: .[0].name, + starts: map(select(.scope_category == "start")) | length, + ends: map(select(.scope_category == "end")) | length + }) + | map(select(.starts != 1 or .ends != 1)) +' relay/aiq-relay.atof.jsonl +``` - ```bash - uv pip install "nvidia-nat[weave]==1.8.0" - ``` +An empty final result from the balance check means every recorded scope closed +exactly once. ATOF `mode: append` is important for web and async-job testing: +multiple worker processes can initialize exporters, and `overwrite` can replace +events written by an earlier job. - For production or container deployments, add this exact pinned dependency - to the image build and rebuild the image. Installing it into a running - container is not a durable deployment. +## Relay logging subscriber -2. Create a [Weights & Biases](https://wandb.ai/) account if you do not have one. +AI-Q registers a process-wide Relay subscriber when `workflow.relay.logging` is +enabled. It reads the sanitized Relay lifecycle stream and renders developer logs: -3. Set the API key in `deploy/.env`: +```text +[Chain Start] shallow_research_agent +[AGENT] model-name +[Tool Calls] 1 tool(s) requested + → web_search_tool +[Tokens] prompt=1882, completion=42, model=model-name +[Tool Start] web_search_tool +[Tool Result] chars=8472 ref=sha256:... +[Chain End] shallow_research_agent +``` - ```bash - WANDB_API_KEY=your-wandb-api-key - ``` +The subscriber does not instrument the workflow itself. Relay's maintained +framework integrations and AI-Q's semantic agent/tool scopes produce events; +the subscriber only formats those events. This keeps console logging aligned +with ATOF and OTEL rather than maintaining a second callback-based trace. - Alternatively, authenticate interactively: +Raw prompts, responses, tool arguments, and tool results are not printed. +Instead, the subscriber logs sizes and stable content references after Relay +redaction. Set the normal console log level under `general.telemetry.logging`. - ```bash - wandb login - ``` +## Export Relay traces to Phoenix -4. Enable Weave tracing in your YAML config: +[Phoenix](https://docs.arize.com/phoenix) provides a local UI for inspecting the +Relay hierarchy, latency, model inputs and outputs, tool calls, token usage, and +errors. - ```yaml - general: - telemetry: - tracing: - weave: - _type: weave - project: aiq-research - ``` +Start Phoenix: -### Configuration Reference +```bash +uvx --from arize-phoenix phoenix serve +``` -The Weave exporter supports PII redaction and custom trace attributes: +Phoenix is available at [http://localhost:6006](http://localhost:6006). OTEL is +commented out in the default AI-Q configs. Uncomment or add this Relay block: ```yaml -general: - telemetry: - tracing: - weave: - _type: weave - project: aiq-research - verbose: false - redact_pii: true - redact_pii_fields: - - CREDIT_CARD - - EMAIL_ADDRESS - - PHONE_NUMBER - redact_keys: - - api_key - - authorization - attributes: - environment: development - team: research +workflow: + relay: + observability: + opentelemetry: + enabled: true + endpoints: + - type: openinference + endpoint: http://localhost:6006/v1/traces + service_name: aiq-relay + resource_attributes: + openinference.project.name: aiq-relay + deployment.environment: development ``` -| Field | Description | -|-------|-------------| -| `project` | The W&B project name. | -| `verbose` | Enable verbose logging for the Weave exporter. | -| `redact_pii` | Automatically redact PII from traces using Presidio. | -| `redact_pii_fields` | Custom PII entity types to redact (e.g., `CREDIT_CARD`, `EMAIL_ADDRESS`). Only used when `redact_pii` is `true`. | -| `redact_keys` | Additional keys to redact beyond the defaults (`api_key`, `auth_headers`, `authorization`). | -| `attributes` | Custom attributes to include in all trace spans. | +The `openinference` projection gives Phoenix semantic LLM, agent, and tool span +attributes and the corresponding UI icons. `openinference.project.name` +selects the Phoenix project. Use a distinct project name for each AI-Q +environment that you want to compare independently. -### What You Can Inspect +Relay also supports `full` and `gen_ai` OTEL projections. Use `full` when the +destination needs the richest Relay-native attributes, and `gen_ai` when the +destination expects OpenTelemetry GenAI semantic conventions. Phoenix users +should normally use `openinference`. -- **Trace timelines** -- Agent execution flows with timing breakdowns. -- **Model calls** -- Inputs, outputs, and metadata for each LLM invocation. -- **Experiment comparison** -- Compare traces across different configurations or model versions. +### Troubleshoot missing Phoenix traces -## OpenTelemetry Collector +If a trace does not appear in the expected Phoenix project: -For production environments, the AI-Q blueprint provides a custom OpenTelemetry exporter (`otelcollector_redaction`) that forwards spans to any OTEL-compatible collector (Jaeger, Grafana Tempo, Datadog, etc.) with built-in privacy redaction. +1. Check that Phoenix is listening at the configured endpoint. +2. Confirm `opentelemetry.enabled: true` and inspect the AI-Q log for export + errors. +3. Look in Phoenix's `default` project and any project configured globally on + the machine. +4. Inspect `~/.config/nemo-relay/plugins.toml`. NeMo Relay automatically + discovers user-level plugin configuration. If Relay was already configured + for another application or coding agent, that exporter can send the AI-Q + trace to its globally configured Phoenix project instead of the project you + are currently viewing. +5. Compare the Phoenix trace with the local ATOF file. If ATOF contains the + scopes, instrumentation worked and the remaining issue is OTEL destination, + project selection, export, or batching. -### Setup +Keep personal Relay configuration when it is needed by other applications. +Use an AI-Q-specific project in the workflow configuration and account for all +discovered exporters when validating where telemetry is sent. -Add the exporter to your YAML config: +## Read AI-Q traces -```yaml -general: - telemetry: - tracing: - otel: - _type: otelcollector_redaction - endpoint: http://your-otel-collector:4318/v1/traces - project: aiq-research - resource_attributes: - deployment.environment: production - service.version: "1.0.0" -``` +AI-Q creates one root trace for each user turn. Multiple turns in the same +conversation have different trace IDs and share the Phoenix `session.id`, so +the session view groups them without merging their execution trees. -### Privacy Redaction +An async deep-research job runs in a separate Relay trace because it executes +outside the request task, often in another Dask worker process. Job metadata +links the background trace to the submitted job and originating request. -The `otelcollector_redaction` exporter can automatically redact sensitive data from trace spans before they leave the application. This is useful for enterprise environments where LLM inputs and outputs may contain PII or confidential information. +A typical deep-research trace is structured as follows: -```yaml -general: - telemetry: - tracing: - otel: - _type: otelcollector_redaction - endpoint: http://your-otel-collector:4318/v1/traces - project: aiq-research - redaction_enabled: true - redaction_attributes: - - input.value - - output.value - - nat.metadata - force_redaction: false - redaction_tag: redacted +```text + +└── chat_deepresearcher_agent + ├── intent_classifier + │ └── LLM + ├── clarifier_agent + │ ├── LLM + │ └── tool + └── deep_research_agent + ├── planner-agent + │ └── LLM + ├── researcher-agent + │ ├── LLM + │ └── tool + └── writer-agent + └── LLM ``` -| Field | Description | -|-------|-------------| -| `endpoint` | The OTEL collector URL to send spans to (e.g., `http://your-otel-collector:4318/v1/traces`). | -| `project` | Logical project name attached to all exported spans. | -| `redaction_enabled` | Enable or disable redaction processing. | -| `redaction_attributes` | Span attributes to redact (defaults to `input.value`, `output.value`, `nat.metadata`). | -| `force_redaction` | Always redact, regardless of header conditions. | -| `redaction_tag` | Tag added to spans when redaction is applied. | -| `redaction_headers` | Request headers checked to determine whether to redact. | -| `resource_attributes` | Custom OTEL resource attributes attached to all spans. | +Use the tree in this order: -### Request Tags on NAT Spans +1. Start at the root and check its terminal status and duration. +2. Find the slowest agent, LLM, or tool child. +3. Inspect LLM spans for model, token usage, response status, and sanitized + input/output attributes. +4. Inspect tool spans for tool name, duration, sanitized arguments/results, and + errors. +5. For parallel researchers, compare sibling spans rather than adding their + wall-clock durations. +6. For a failed async job, search by `aiq.job.id` and confirm the root scope has + one start, one end, and an `ERROR` terminal status. -When the `aiq_api` auth middleware is enabled, NAT-exported workflow spans can -include low-cardinality request tags plus optional pseudonymous identity tags. -These tags are propagated across HTTP requests, WebSocket workflows, and async -job execution. +Internal graph-routing nodes are represented as decision metadata/events where +possible rather than noisy agent spans. Framework-generated names can still +appear when the underlying integration exposes a real execution boundary. -Always-on NAT span tags: +## Redaction and privacy -- `nat.aiq.caller.type` -- resolved caller type from auth middleware -- `nat.aiq.auth.transport` -- `bearer`, `cookie`, or `none` -- `nat.aiq.auth.verified` -- whether the request resolved to a verified principal -- `nat.aiq.access.channel` -- inferred request channel or trusted explicit access-channel header +Relay redaction runs before the AI-Q logging subscriber, ATOF sink, and OTEL +exporter. The default detectors cover common credentials and personal data. +AI-Q can also request privacy-mode sanitization for supported `data` and +`category_profile` payloads through request privacy context. -Optional pseudonymous tags: +Redaction reduces accidental disclosure; it is not a substitute for auditing +the destination's access controls, retention, and data policy. Validate every +configured exported attribute with synthetic sensitive values before enabling +full payloads in a production environment. -- `nat.enduser.id`, `nat.aiq.user.id`, `nat.aiq.auth.type` -- controlled by `AIQ_TRACE_USER_IDENTITY_MODE` -- `nat.aiq.user.email`, `nat.aiq.user.name` -- added only in `full` mode -- `nat.aiq.client.id` -- controlled by `AIQ_TRACE_CLIENT_ID_MODE=ip` +## Pricing and cost analysis -Environment variables: +Default AI-Q configs deliberately use an empty Relay pricing source list: -- `AIQ_TRACE_USER_IDENTITY_MODE=none|id|full` -- `AIQ_TRACE_USER_IDENTITY_HMAC_SECRET=` -- `AIQ_TRACE_CLIENT_ID_MODE=none|ip` -- `AIQ_TRACE_CLIENT_ID_HMAC_SECRET=` -- `AIQ_TRACE_CLIENT_IP_HEADERS=x-real-ip,x-forwarded-for` +```yaml +workflow: + relay: + pricing: + enabled: true + sources: [] +``` + +This records token usage without claiming a monetary cost. Pricing depends on +the provider, deployment, contract, region, cache policy, and date. -The `id` and `ip` modes emit HMAC-derived pseudonymous identifiers rather than -raw subjects or raw IP addresses. +Use the dedicated example when you want model cost enrichment: -### Batch Configuration +```bash +nat serve \ + --config_file configs/nemo_relay/config_web_default_with_pricing.yml \ + --port 8000 +``` -The exporter supports standard OTEL batch settings: +That config loads `configs/nemo_relay/relay_pricing_catalog.json`: ```yaml -general: - telemetry: - tracing: - otel: - _type: otelcollector_redaction - endpoint: http://your-otel-collector:4318/v1/traces - batch_size: 512 - flush_interval: 5000 - max_queue_size: 2048 - drop_on_overflow: false - shutdown_timeout: 30000 +pricing: + enabled: true + sources: + - type: file + path: configs/nemo_relay/relay_pricing_catalog.json ``` -## Verbose Logging - -For quick debugging without any external services, enable the built-in `VerboseTraceCallback` logger. This callback -records execution metadata directly to the console without printing raw prompts, tool arguments, tool results, or -model responses. This metadata-only guarantee applies only to `VerboseTraceCallback`. Phoenix and other exporters, -enabled source adapters, and external providers can still receive or retain raw prompts, tool arguments, tool results, -and model responses; configure and audit their redaction, retention, and access controls independently. +Relay matches the observed provider/model name against the catalog and adds +cost information to the annotated LLM usage. Review and date every rate before +using it operationally. A zero-dollar hosted API rate does not mean that a +self-hosted deployment has no infrastructure cost. -### Enable via CLI +Relay's pricing catalog covers model usage. AI-Q's tokenomics report also +supports per-call prices for external tools such as web search. After a run, +generate the report with: ```bash -./scripts/start_cli.sh --verbose +PYTHONPATH=src python -m aiq_agent.tokenomics.report \ + --trace relay/aiq-relay.atof.jsonl \ + --config frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml ``` -### Enable via YAML Config +The report uses Relay-attributed model cost when present and the report pricing +configuration as a fallback. It also calculates configured tool API charges +and writes a self-contained HTML report. See [Profiling and Cost +Analysis](../profiling/index.md) for report fields, phase attribution, and +pricing maintenance. -```yaml -workflow: - _type: chat_deepresearcher_agent - verbose: true +## Validate the configuration + +Validate an edited workflow before starting AI-Q: + +```bash +uv run python .agents/skills/aiq-configure-workflow/scripts/validate_config.py \ + configs/config_web_default_llamaindex.yml ``` -### What Gets Logged +Then start AI-Q, run one shallow turn and one deep-research job, and verify all +three views that you enabled: -- Chain starts and completions (orchestrator routing, agent handoffs) -- LLM invocation metadata, such as model and message counts when available -- Tool names and execution metadata -- Content lengths and redaction markers instead of raw request or response content +- console logs show agent, LLM, and tool lifecycle activity; +- ATOF contains balanced scopes and annotated LLM usage; +- Phoenix shows the expected project, per-turn traces, shared session grouping, + and an independent async-job trace. diff --git a/docs/source/deployment/production.md b/docs/source/deployment/production.md index 81de95391..18f876858 100644 --- a/docs/source/deployment/production.md +++ b/docs/source/deployment/production.md @@ -364,10 +364,12 @@ Set `LOG_LEVEL=DEBUG` for verbose output during troubleshooting. Use `LOG_LEVEL= ### Tracing -The backend supports OpenTelemetry-compatible tracing. See [Observability](./observability.md) for setup guides covering Phoenix, LangSmith, Weave, and the OTEL Collector with privacy redaction. +The backend exports NeMo Relay traces to OpenTelemetry-compatible destinations. +See [Observability](./observability.md) for ATOF, Phoenix OTEL, pricing, and +privacy-redaction guidance. If you are deploying the `aiq_api` front-end and want request correlation on -NAT-exported spans, set the relevant environment variables at deploy time rather +Relay-exported spans, set the relevant environment variables at deploy time rather than hardcoding them in code: - `AIQ_TRACE_USER_IDENTITY_MODE` diff --git a/docs/source/evaluation/benchmarks/deep-research-bench.md b/docs/source/evaluation/benchmarks/deep-research-bench.md index e480cdb4f..90cc1c445 100644 --- a/docs/source/evaluation/benchmarks/deep-research-bench.md +++ b/docs/source/evaluation/benchmarks/deep-research-bench.md @@ -59,9 +59,12 @@ python frontends/benchmarks/deepresearch_bench/scripts/export_drb_jsonl.py --inp Follow instructions in the [Deep Research Bench Github Repository](https://github.com/Ayanami0730/deep_research_bench/tree/main) to run evaluation and obtain scores. -## Optional: Phoenix Tracing +## Optional: Relay and Phoenix Tracing -If your config enables Phoenix tracing, start the Phoenix server before running `nat eval`. +AI-Q evaluation uses the same NeMo Relay observability path as interactive and +async workflows. ATOF is enabled by default. To visualize the evaluation in +Phoenix, enable the Relay OpenInference OTEL endpoint in the evaluated workflow +and start Phoenix before running `nat eval`. Start server (separate terminal): @@ -69,27 +72,26 @@ Start server (separate terminal): uvx --from arize-phoenix phoenix serve ``` -## W&B Tracking - -Evaluation runs are tracked using [Weights & Biases Weave - deep-researcher-v2 project](https://wandb.ai/nvidia-aiq/deep-researcher-v2/weave) for experiment tracking and observability. - -### Configuration - -Enable W&B tracking in your config file under `general.telemetry.tracing`: - ```yaml -general: - telemetry: - tracing: - weave: - _type: weave - project: "deep-researcher-v2" +workflow: + relay: + observability: + opentelemetry: + enabled: true + endpoints: + - type: openinference + endpoint: http://localhost:6006/v1/traces + resource_attributes: + openinference.project.name: aiq-deepresearch-bench eval: general: workflow_alias: "aiq-deepresearch-v2-baseline" ``` +See [Observability with NeMo Relay](../../deployment/observability.md) for ATOF +inspection, trace interpretation, project selection, and cost reporting. + ### workflow_alias The `workflow_alias` parameter provides a workflow-specific identifier for tracking evaluation runs: diff --git a/docs/source/examples/cli-with-local-nims.md b/docs/source/examples/cli-with-local-nims.md index 3306bc79b..1412c11d8 100644 --- a/docs/source/examples/cli-with-local-nims.md +++ b/docs/source/examples/cli-with-local-nims.md @@ -48,12 +48,6 @@ general: console: _type: console level: INFO - # Optional: trace to local Phoenix for debugging - # tracing: - # phoenix: - # _type: phoenix - # endpoint: http://localhost:6006/v1/traces - # project: local-dev # =========================================================================== # LLMs - pointing to local NIM containers @@ -118,7 +112,6 @@ functions: - web_search_tool max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent diff --git a/docs/source/examples/full-pipeline-llamaindex.md b/docs/source/examples/full-pipeline-llamaindex.md index e30db5db0..09bc5ffe3 100644 --- a/docs/source/examples/full-pipeline-llamaindex.md +++ b/docs/source/examples/full-pipeline-llamaindex.md @@ -148,7 +148,6 @@ functions: intent_classifier: _type: intent_classifier llm: nemotron_lightning_intent_llm - verbose: true tools: - web_search_tool # - paper_search_tool # Uncomment if SERPER_API_KEY is set @@ -162,12 +161,10 @@ functions: - knowledge_search max_turns: 3 log_response_max_chars: 2000 - verbose: true shallow_research_agent: _type: shallow_research_agent llm: nemotron_lightning_agent_llm - verbose: true tools: - web_search_tool - knowledge_search @@ -181,7 +178,6 @@ functions: planner_llm: nemotron_ultra_llm researcher_llm: nemotron_ultra_llm writer_llm: nemotron_ultra_writer_llm - verbose: true tools: # - paper_search_tool # Uncomment if SERPER_API_KEY is set - advanced_web_search_tool @@ -189,7 +185,6 @@ functions: workflow: _type: chat_deepresearcher_agent - verbose: true enable_escalation: true enable_clarifier: true use_async_deep_research: true diff --git a/docs/source/examples/full-pipeline-web.md b/docs/source/examples/full-pipeline-web.md index 932c1ff83..efdf8b9e7 100644 --- a/docs/source/examples/full-pipeline-web.md +++ b/docs/source/examples/full-pipeline-web.md @@ -32,12 +32,6 @@ general: console: _type: console level: INFO - # Uncomment for tracing: - # tracing: - # phoenix: - # _type: phoenix - # endpoint: http://localhost:6006/v1/traces - # project: dev # --------------------------------------------------------------------------- # Front-end: AI-Q API plugin @@ -182,7 +176,6 @@ functions: - knowledge_search max_turns: 3 # Max clarification rounds log_response_max_chars: 2000 - verbose: true # ------------------------------------------------------------------------- # Shallow research agent diff --git a/docs/source/integration/index.md b/docs/source/integration/index.md index 334c97e22..82d3433e9 100644 --- a/docs/source/integration/index.md +++ b/docs/source/integration/index.md @@ -10,3 +10,4 @@ Connect AI-Q to external systems and services. - **[Agent Skills](./agent-skills.md)** — Install the AI-Q research skill in Claude Code, OpenCode, Codex, or another Agent Skills-compatible coding harness - **[REST API](./rest-api.md)** — Async jobs API endpoints, SSE events, request/response models - **[MCP Server](./mcp-server.md)** — Expose AI-Q through a standalone, stateless, no-authentication MCP server +- **[NeMo Relay observability](../deployment/observability.md)** — Inspect AI-Q agents with developer logs, ATOF, Phoenix OTEL, redaction, and cost data diff --git a/docs/source/profiling/index.md b/docs/source/profiling/index.md index 96893756a..ab3db0978 100644 --- a/docs/source/profiling/index.md +++ b/docs/source/profiling/index.md @@ -142,31 +142,34 @@ Declare prices under `tokenomics.pricing`: tokenomics: pricing: models: - "azure/openai/gpt-5.2": - input_per_1m_tokens: 2.50 - output_per_1m_tokens: 10.00 - # Illustrative market-equivalent rates; verify current provider pricing. "nvidia/nemotron-3-ultra-550b-a55b": - input_per_1m_tokens: 0.60 - output_per_1m_tokens: 3.60 + # NVIDIA-hosted access for this example; not self-hosting cost. + input_per_1m_tokens: 0.00 + output_per_1m_tokens: 0.00 tools: - # Key "web_search" matches "advanced_web_search_tool" via substring lookup - "web_search": + # Tavily pay-as-you-go: $0.008/credit; basic uses one credit and + # advanced uses two. Use your plan's effective credit rate instead. + "web_search_tool": + cost_per_call: 0.008 + "advanced_web_search_tool": cost_per_call: 0.016 + # Default Serper Starter tier: $50 / 50,000 successful queries. + # Change this when using another tier or paper-search provider. "paper_search": - cost_per_call: 0.0003 - # Fallback for any model not listed above. - # Set to null to raise an error on unknown models instead. - default: - input_per_1m_tokens: 1.00 - output_per_1m_tokens: 4.00 + cost_per_call: 0.001 ``` You can optionally set `eval.general.output_dir` in that same file so the report’s default output path matches your eval artifacts directory (refer to `config_tokenomics_pricing.yml` in the bench configs). -**Model name lookup** uses exact match first, then substring match, then the `default`. A key of `"gpt-5.2"` matches a trace model name of `"azure/openai/gpt-5.2"` because the key is a substring of the full name. +**Model name lookup** uses exact match first, then substring match, then the +`default`. Prefer the exact provider model identifier emitted in the Relay +trace so similarly named deployments do not share prices accidentally. -**Tool name lookup** follows the same rule. A key of `"web_search"` matches `"advanced_web_search_tool"` because `"web_search"` is a substring of the tool name. Unknown tools default to $0 — no error is raised, so you only need to configure tools that have a real per-call cost. +**Tool name lookup** follows the same rule, with exact matches taking priority. +Keep wrapper and provider-facing tool names distinct so one external request is +not charged twice. Unknown and internal tools default to $0. Provider-backed +tools such as paper search must use the effective per-request price for the +configured provider and subscription plan. **`cached_input_per_1m_tokens`** is optional. When omitted, cached tokens are billed at the full input rate (no discount). Set it when your model provider charges a reduced rate for KV-cache hits. @@ -176,7 +179,7 @@ After `nat eval` completes, run: ```bash PYTHONPATH=src python -m aiq_agent.tokenomics.report \ - --trace frontends/benchmarks/deepresearch_bench/results/all_requests_profiler_traces.json \ + --trace relay/aiq-relay.atof.jsonl \ --config frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml ``` @@ -241,23 +244,12 @@ Full per-query table: cost, ISL, OSL, cached tokens, ISL:OSL ratio, LLM call cou ### Subagent Phase Attribution -The Deep Research Agent has an orchestrator, an optional source router, a planner, parallel researcher workers, and -a writer. The current adapter in `src/aiq_agent/tokenomics/nat_adapter.py` builds timing windows for `task` -invocations whose `subagent_type` it can parse. It maps `planner-agent` windows to `planner-phase` and every other -parsed task subagent to `researcher-phase`. It associates an `LLM_END` with a window using the call's completion -timestamp; calls outside task windows fall into `orchestrator-phase`. - -This does not align completely with the current runtime. The optional `source-router-agent`, `planner-agent`, and -`writer-agent` are delegated through `task()`, so source-router and writer calls are normally folded into -`researcher-phase`. Researcher workers are invoked directly by `run_research_batch` rather than through individual -`task()` calls, so their calls can instead appear in `orchestrator-phase`. The researcher bucket is therefore a -mixed task-subagent bucket, and the orchestrator bucket is partly an **unattributed/default bucket**; neither proves -which role's model performed the work. - -Phase charts are consequently best-effort diagnostics, not correct per-role cost accounting for the current runtime. -Overall token and cost totals remain useful independently of that distribution, subject to the completeness of the -trace and pricing configuration. Native role metadata on each LLM step, or adapter support for every current -execution path, is required before the phase split can be treated as authoritative. +The adapter in `src/aiq_agent/tokenomics/atof_adapter.py` reads Relay ATOF +JSONL and follows scope `parent_uuid` ancestry. Calls nested below real +`planner-agent` and `researcher-agent` scopes are attributed to those phases; +all remaining calls use the orchestrator bucket. Parallel researcher tasks use +isolated Relay asyncio contexts, so they retain correct parentage without +timing-window inference. ### Python API @@ -272,7 +264,7 @@ with open("frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pri pricing = PricingRegistry.from_dict(config["tokenomics"]["pricing"]) profiles = parse_trace( - "frontends/benchmarks/deepresearch_bench/results/all_requests_profiler_traces.json", + "relay/aiq-relay.atof.jsonl", pricing, ) diff --git a/docs/source/resources/troubleshooting.md b/docs/source/resources/troubleshooting.md index bb8a2b220..3ea2c39e7 100644 --- a/docs/source/resources/troubleshooting.md +++ b/docs/source/resources/troubleshooting.md @@ -34,7 +34,7 @@ Common issues and solutions for the AI-Q blueprint. | Issue | Cause | Fix | |-------|-------|-----| -| Agent hangs on deep research | LLM timeout or rate limit | Set `verbose: true` in config to see progress; check LLM API availability and rate limits | +| Agent hangs on deep research | LLM timeout or rate limit | Inspect Relay logs/traces and check LLM API availability and rate limits | | HTTP 429 or 503 on deep research | Nemotron hosted endpoint availability | Retry after a short delay, reduce concurrency, or follow the [self-hosting guidance](#nemotron-hosted-endpoint-availability) for consistent throughput | | Intermittent shallow-research failure with Nemotron 3.5 Lightning on NVIDIA API Catalog | The hosted serving profile can produce citation-incomplete or malformed final drafts | Use Nemotron Ultra for the shallow role, or use a validated self-hosted Lightning serving profile; see [Nemotron 3.5 Lightning on NVIDIA API Catalog](#nemotron-35-lightning-on-nvidia-api-catalog) | | Shallow research returns generic answers | Insufficient tool calls | Increase `max_tool_iterations` (default: 5) | @@ -159,34 +159,39 @@ Docker Compose deployments on the VM handle container-to-host port mapping autom ## Debugging Tips -### Enable Verbose Logging +### Inspect Relay Logging ```yaml # In your config YAML workflow: _type: chat_deepresearcher_agent - verbose: true + relay: + logging: true ``` -Or through CLI: `./scripts/start_cli.sh --verbose` +### Phoenix Tracing Through Relay -### Phoenix Tracing - -For full setup instructions covering Phoenix, LangSmith, and other tracing backends, see [Observability](../deployment/observability.md). +For full setup and trace-reading instructions, see [Observability with NeMo Relay](../deployment/observability.md). Start a Phoenix server and enable tracing in config: ```yaml -general: - telemetry: - tracing: - phoenix: - _type: phoenix - endpoint: http://localhost:6006/v1/traces - project: dev +workflow: + relay: + observability: + opentelemetry: + enabled: true + endpoints: + - type: openinference + endpoint: http://localhost:6006/v1/traces + resource_attributes: + openinference.project.name: aiq-relay ``` Then open [http://localhost:6006](http://localhost:6006) to inspect traces, token usage, and latency. +If the trace is missing, also inspect the project configured in +`~/.config/nemo-relay/plugins.toml`; Relay can discover an existing user-level +Phoenix destination. ### Check Registered Components diff --git a/frontends/aiq_api/src/aiq_api/auth/request_trace.py b/frontends/aiq_api/src/aiq_api/auth/request_trace.py index 59ca22357..631774cef 100644 --- a/frontends/aiq_api/src/aiq_api/auth/request_trace.py +++ b/frontends/aiq_api/src/aiq_api/auth/request_trace.py @@ -31,10 +31,14 @@ def get_request_trace_tags() -> dict[str, str]: @contextmanager def request_trace_tag_context(tags: dict[str, str]): - """Bind request trace tags while NAT emits spans for this request.""" + """Bind request trace tags and Relay privacy while telemetry is emitted.""" + from aiq_agent.relay.privacy import request_privacy_context + from aiq_agent.relay.privacy import request_privacy_from_tags + token = _current_request_trace_tags.set(dict(tags)) try: - yield + with request_privacy_context(request_privacy_from_tags(tags)): + yield finally: _current_request_trace_tags.reset(token) diff --git a/frontends/aiq_api/src/aiq_api/auth/utils.py b/frontends/aiq_api/src/aiq_api/auth/utils.py index 4b11f612c..6161ed2d7 100644 --- a/frontends/aiq_api/src/aiq_api/auth/utils.py +++ b/frontends/aiq_api/src/aiq_api/auth/utils.py @@ -38,6 +38,7 @@ "unknown", } ) +TRACE_REDACTION_HEADER = "x-aiq-telemetry-redact" def _load_trace_user_identity_mode() -> str: @@ -270,6 +271,8 @@ def _build_common_trace_tags( allow_explicit_override=trust_access_channel_override, ), } + if (_extract_header_text(headers, TRACE_REDACTION_HEADER) or "").lower() == "true": + tags["aiq.telemetry.redact"] = "true" if client_id_mode == "ip": client_ip = _extract_client_ip(headers, scope, client_ip_headers) diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index d279dcf85..1c7de7ccc 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -32,6 +32,8 @@ import uuid from collections.abc import Awaitable from collections.abc import Callable +from dataclasses import dataclass +from dataclasses import field from typing import TYPE_CHECKING from typing import Any @@ -67,6 +69,16 @@ _JOB_SCOPED_AGENT_KWARGS = frozenset({"job_id"}) +@dataclass(frozen=True) +class JobTraceCorrelation: + """Serializable correlation from a submitting request to an independent job trace.""" + + session_id: str | None = None + submission_trace_id: str | None = None + submission_span_id: str | None = None + request_trace_tags: dict[str, str] = field(default_factory=dict) + + def _constructor_accepts_explicit_kwargs(agent_cls: type, kwarg_names: frozenset[str]) -> bool: """Return true when a class constructor explicitly declares all requested kwargs.""" import inspect @@ -87,25 +99,6 @@ def _constructor_accepts_explicit_kwargs(agent_cls: type, kwarg_names: frozenset return kwarg_names.issubset(accepted_kwargs) -def _normalize_trace_id(trace_id: int | str | None) -> int | None: - """Convert trace ID to integer format. - - Args: - trace_id: Trace ID as int, hex string, or None. - - Returns: - Integer trace ID or None. - """ - if trace_id is None: - return None - if isinstance(trace_id, int): - return trace_id - try: - return int(trace_id, 16) - except ValueError: - return int(trace_id) - - class CancellationMonitor: """ Monitors job status for cancellation requests. @@ -601,13 +594,7 @@ async def run_agent_job( input_text: str, agent_class_path: str, agent_config_name: str, - parent_span_id: str | None = None, - parent_function_id: str | None = None, - parent_function_name: str | None = None, - parent_workflow_run_id: str | None = None, - parent_workflow_trace_id: int | str | None = None, - parent_conversation_id: str | None = None, - request_trace_tags: dict[str, str] | None = None, + trace_correlation: JobTraceCorrelation | None = None, available_documents: list[dict] | None = None, data_sources: list[str] | None = None, auth_token: str | None = None, @@ -624,7 +611,7 @@ async def run_agent_job( - Uses NAT's JobStore for status tracking - Monitors for cancellation requests and gracefully terminates the agent - Exports telemetry to Phoenix/OpenTelemetry via NAT's ExporterManager - - Propagates trace context from parent workflow for nested spans + - Starts an independent trace correlated to the submitting request and session Args: configure_logging: Whether to set up logging in the worker. @@ -636,13 +623,7 @@ async def run_agent_job( input_text: User input/query to run. agent_class_path: Full module path to agent class. agent_config_name: NAT config function name for the agent. - parent_span_id: Parent span ID for trace continuity (from caller context). - parent_function_id: Parent function ID for span hierarchy. - parent_function_name: Parent function name for span metadata. - parent_workflow_run_id: Parent workflow run ID for trace grouping. - parent_workflow_trace_id: Parent trace ID (int or hex string) for trace continuity. - parent_conversation_id: Conversation ID for session grouping in Phoenix. - request_trace_tags: Request trace tags captured at async submission time. + trace_correlation: Session and submission identifiers used to correlate this independent job trace. available_documents: Optional list of document dicts with file_name and summary. data_sources: Optional list of allowed data sources to enforce in the worker. auth_token: Optional auth token propagated from the HTTP request for @@ -657,6 +638,8 @@ async def run_agent_job( admission_token: Opaque deep-research fencing token captured at submit time. """ + trace_correlation = trace_correlation or JobTraceCorrelation() + # Propagate auth token into the current async task's context so tools # can retrieve it via get_auth_token(). Uses a ContextVar so concurrent # jobs in the same Dask worker process don't leak tokens across tasks. @@ -674,8 +657,6 @@ async def run_agent_job( install_request_trace_span_injection() - from aiq_agent.common import VerboseTraceCallback - from aiq_agent.common import is_verbose from nat.builder.framework_enum import LLMFrameworkEnum from nat.builder.workflow_builder import WorkflowBuilder from nat.front_ends.fastapi.async_jobs.job_store import JobStatus @@ -787,7 +768,7 @@ async def run_agent_job( from nat.builder.context import ContextState context_state = ContextState.get() - _conversation_id_reset = context_state.conversation_id.set(parent_conversation_id) + _conversation_id_reset = context_state.conversation_id.set(trace_correlation.session_id) # Always shadow the inherited identity, including for ownerless jobs, # so a reused worker context cannot expose a prior owner's MCP tokens. _user_id_reset = context_state.user_id.set(owner_user_id) @@ -796,6 +777,11 @@ async def run_agent_job( await _attach_middleware_to_function(builder, config, agent_config_name) fn_config = builder.get_function_config(agent_config_name) + relay_config = getattr(fn_config, "relay", None) + if relay_config is not None: + from aiq_agent.relay.bootstrap import ensure_started as ensure_relay_started + + await ensure_relay_started(relay_config) if getattr(fn_config, "type", None) == "deep_research_agent": from aiq_agent.agents.deep_researcher.register import DeepResearchAgentConfig from aiq_agent.agents.deep_researcher.register import resolve_deep_research_runtime_config @@ -835,18 +821,16 @@ async def run_agent_job( from nat.observability.exporter_manager import ExporterManager from nat.utils.reactive.subject import Subject - from .telemetry import AgentLifecycleTelemetryCallback - from .telemetry import aiq_langchain_profiler_context - telemetry_exporters = { name: configured.instance for name, configured in builder._telemetry_exporters.items() } exporter_manager = ExporterManager.from_exporters(telemetry_exporters) - # Initialize context state with trace propagation from parent + # A durable background job is an independent trace. The submitting + # request is retained only as correlation metadata. context_state.workflow_run_id.set(job_id) - workflow_trace_id = _normalize_trace_id(parent_workflow_trace_id) or uuid.uuid4().int + workflow_trace_id = uuid.uuid4().int context_state.workflow_trace_id.set(workflow_trace_id) # Event stream for exporters to subscribe to @@ -862,8 +846,6 @@ async def run_agent_job( InvocationNode( function_name=workflow_span_name, function_id=job_id, - parent_id=parent_function_id, - parent_name=parent_function_name, ) ) @@ -873,36 +855,16 @@ async def run_agent_job( provided_metadata={ "workflow_run_id": job_id, "workflow_trace_id": f"{workflow_trace_id:032x}", - "conversation_id": parent_conversation_id, + "conversation_id": trace_correlation.session_id, "agent": agent_class_path, - "parent_workflow_run_id": parent_workflow_run_id, - "parent_workflow_name": parent_function_name, + "submission_trace_id": trace_correlation.submission_trace_id, + "submission_span_id": trace_correlation.submission_span_id, } ) # Run with telemetry - exporter must start before pushing events - with request_trace_tag_context(request_trace_tags or {}): + with request_trace_tag_context(trace_correlation.request_trace_tags): async with exporter_manager.start(context_state=context_state): - # Link to parent span if provided (for nested trace continuity) - parent_metadata: TraceMetadata | None = None - if parent_span_id and parent_span_id != "root": - parent_metadata = TraceMetadata( - provided_metadata={ - "workflow_run_id": parent_workflow_run_id, - "workflow_trace_id": f"{workflow_trace_id:032x}", - "conversation_id": parent_conversation_id, - "workflow_name": parent_function_name, - } - ) - context.intermediate_step_manager.push_intermediate_step( - IntermediateStepPayload( - UUID=parent_span_id, - event_type=IntermediateStepType.SPAN_START, - name=parent_function_name or "parent_workflow", - metadata=parent_metadata, - ) - ) - # Push WORKFLOW_START first so LLM/tool events become children context.intermediate_step_manager.push_intermediate_step( IntermediateStepPayload( @@ -914,14 +876,10 @@ async def run_agent_job( ) ) - agent_telemetry_callback = AgentLifecycleTelemetryCallback(context.intermediate_step_manager) - - verbose = is_verbose(getattr(fn_config, "verbose", False)) - callbacks = [VerboseTraceCallback()] if verbose else [] + callbacks: list[Any] = [] raw_event_store = EventStore(db_url, job_id, content_cipher=job_output_cipher) event_store = BatchingEventStore(raw_event_store) - callbacks.append(agent_telemetry_callback) callbacks.append(AgentEventCallback(event_store)) # Resolve per-user MCP source tools for the job owner (Context.user_id @@ -947,7 +905,6 @@ async def run_agent_job( llm=llm, tools=agent_tools, fn_config=fn_config, - verbose=verbose, callbacks=callbacks, job_id=job_id, # Artifact harvesting rides 284's job store + event stream: the same db_url @@ -961,10 +918,10 @@ async def run_agent_job( # agents without a sandbox runtime; close()/terminate() are then no-ops. sandbox_runtime = getattr(agent, "deepagents_runtime", None) - # Replace NAT's inherited profiler for this invocation rather than adding a - # second callback with duplicate LangChain run IDs. - with aiq_langchain_profiler_context(): - result = await _run_agent( + from aiq_agent.relay import run_workflow as run_relay_workflow + + async def _execute_agent() -> Any: + return await _run_agent( agent=agent, input_text=input_text, builder=builder, @@ -978,6 +935,20 @@ async def run_agent_job( initial_files=initial_files, ) + result = await run_relay_workflow( + f"async_{agent_config_name.removesuffix('_agent')}_job", + _execute_agent, + session_id=trace_correlation.session_id, + input_value=input_text, + metadata={ + "aiq.execution.mode": "async", + "aiq.job.id": job_id, + "aiq.agent.type": agent_config_name, + "aiq.submission.trace_id": trace_correlation.submission_trace_id, + "aiq.submission.span_id": trace_correlation.submission_span_id, + }, + ) + # Emit WORKFLOW_END event for Phoenix context.intermediate_step_manager.push_intermediate_step( IntermediateStepPayload( @@ -989,16 +960,6 @@ async def run_agent_job( ) ) - if parent_metadata: - context.intermediate_step_manager.push_intermediate_step( - IntermediateStepPayload( - UUID=parent_span_id, - event_type=IntermediateStepType.SPAN_END, - name=parent_function_name or "parent_workflow", - metadata=parent_metadata, - ) - ) - # Signal event stream completion event_stream.on_complete() @@ -1230,7 +1191,6 @@ def _create_agent_instance( llm, tools: list, fn_config, - verbose: bool, callbacks: list, job_id: str | None = None, artifact_db_url: str | None = None, @@ -1254,7 +1214,6 @@ def _create_agent_instance( return agent_cls( llm_provider=llm_provider, tools=tools, - verbose=verbose, callbacks=callbacks, domain_catalog_path=fn_config.domain_catalog_path, enable_source_router=fn_config.enable_source_router, @@ -1275,7 +1234,6 @@ def _create_agent_instance( return agent_cls( llm_provider=llm_provider, tools=tools, - verbose=verbose, callbacks=callbacks, config=fn_config, job_id=job_id, @@ -1288,19 +1246,17 @@ def _create_agent_instance( return agent_cls( llm_provider=llm_provider, tools=tools, - verbose=verbose, callbacks=callbacks, job_id=job_id, ) except TypeError: pass - # Try original deep_researcher pattern (llm_provider + tools + verbose) + # Try the common llm_provider + tools pattern. try: return agent_cls( llm_provider=llm_provider, tools=tools, - verbose=verbose, callbacks=callbacks, ) except TypeError: diff --git a/frontends/aiq_api/src/aiq_api/jobs/submit.py b/frontends/aiq_api/src/aiq_api/jobs/submit.py index 59af5b750..0e550f292 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -51,6 +51,7 @@ from .admission import renew_deep_research_job_reservation from .admission import reserve_deep_research_job from .admission import validate_deep_research_input +from .runner import JobTraceCorrelation from .runner import run_agent_job logger = logging.getLogger(__name__) @@ -105,58 +106,22 @@ def _resolve_admission_principal(principal: Principal) -> Principal: return Principal(type="anonymous", sub="anonymous") -def _get_parent_trace_context() -> tuple[ - str | None, # parent_span_id - str | None, # parent_function_id - str | None, # parent_function_name - str | None, # parent_workflow_run_id - int | str | None, # parent_workflow_trace_id - str | None, # parent_conversation_id - dict[str, str], # request_trace_tags -]: - """ - Extract trace context from current workflow for propagation to async jobs. - - This enables nested spans in Phoenix - the async job will appear as a child - of the workflow that submitted it. - - Returns: - Tuple of (parent_span_id, parent_function_id, parent_function_name, - parent_workflow_run_id, parent_workflow_trace_id, parent_conversation_id, request_trace_tags) - """ +def _get_job_trace_correlation() -> JobTraceCorrelation: + """Capture identifiers that correlate an independent async-job trace to its submission.""" try: from nat.builder.context import ContextState except ImportError: - return (None, None, None, None, None, None, {}) + return JobTraceCorrelation(request_trace_tags=get_current_trace_tags()) context_state = ContextState.get() - - # Extract workflow-level context - parent_workflow_run_id = context_state.workflow_run_id.get() - parent_workflow_trace_id = context_state.workflow_trace_id.get() - parent_conversation_id = context_state.conversation_id.get() - - # Extract span hierarchy context - parent_span_id = None + workflow_trace_id = context_state.workflow_trace_id.get() active_stack = context_state.active_span_id_stack.get() - if active_stack and len(active_stack) > 1: - parent_span_id = active_stack[1] - - parent_function_id = None - parent_function_name = None - active_function = context_state.active_function.get() - if active_function and active_function.function_id != "root": - parent_function_id = active_function.function_id - parent_function_name = active_function.function_name - - return ( - parent_span_id, - parent_function_id, - parent_function_name, - parent_workflow_run_id, - parent_workflow_trace_id, - parent_conversation_id, - get_current_trace_tags(), + active_span_id = active_stack[-1] if active_stack and active_stack[-1] != "root" else None + return JobTraceCorrelation( + session_id=context_state.conversation_id.get(), + submission_trace_id=f"{workflow_trace_id:032x}" if workflow_trace_id is not None else None, + submission_span_id=active_span_id, + request_trace_tags=get_current_trace_tags(), ) @@ -326,14 +291,15 @@ async def submit_agent_job( reservation_ttl_seconds = ( admission_limits.reservation_ttl_seconds if admission_limits is not None else DEFAULT_RESERVATION_TTL_SECONDS ) - parent_trace_context = _get_parent_trace_context() + trace_correlation = _get_job_trace_correlation() if conversation_id is not None: - parent_trace_context = ( - *parent_trace_context[:5], - conversation_id, - parent_trace_context[6], + trace_correlation = JobTraceCorrelation( + session_id=conversation_id, + submission_trace_id=trace_correlation.submission_trace_id, + submission_span_id=trace_correlation.submission_span_id, + request_trace_tags=trace_correlation.request_trace_tags, ) - submission_conversation_id = parent_trace_context[5] + submission_conversation_id = trace_correlation.session_id async def _release_submission_reservations() -> None: """Conditionally release only reservations owned by this submitter.""" @@ -459,7 +425,7 @@ async def _stop_submission_lease() -> None: input_text, agent_config.class_path, agent_config.config_name, - *parent_trace_context, + trace_correlation, available_documents, data_sources, auth_token, diff --git a/frontends/aiq_api/src/aiq_api/jobs/telemetry.py b/frontends/aiq_api/src/aiq_api/jobs/telemetry.py deleted file mode 100644 index b87ea4633..000000000 --- a/frontends/aiq_api/src/aiq_api/jobs/telemetry.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""AI-Q telemetry adapters for named LangChain agent spans.""" - -from __future__ import annotations - -import ast -from collections.abc import Iterator -from contextlib import contextmanager -from typing import Any - -from langchain_core.callbacks import BaseCallbackHandler - -from nat.data_models.intermediate_step import IntermediateStepPayload -from nat.data_models.intermediate_step import IntermediateStepType -from nat.data_models.intermediate_step import TraceMetadata -from nat.plugins.langchain.callback_handler import LangchainProfilerHandler - - -def _deepagents_agent_name(kwargs: dict[str, Any]) -> str | None: - """Return the outer DeepAgents chain name, excluding internal graph nodes.""" - metadata = kwargs.get("metadata") - if not isinstance(metadata, dict): - return None - - semantic_name = metadata.get("lc_agent_name") - callback_name = kwargs.get("name") - if not isinstance(semantic_name, str) or callback_name != semantic_name: - return None - return semantic_name - - -def _task_display_name(serialized: dict[str, Any], input_str: str, inputs: dict[str, Any] | None) -> str: - name = str(serialized.get("name", "")) - if name != "task": - return name - - parsed_inputs: Any = inputs - if not isinstance(parsed_inputs, dict): - try: - parsed_inputs = ast.literal_eval(input_str) - except (SyntaxError, ValueError): - parsed_inputs = None - subagent_type = parsed_inputs.get("subagent_type") if isinstance(parsed_inputs, dict) else None - return f"task: {subagent_type}" if subagent_type else name - - -class AgentLifecycleTelemetryCallback(BaseCallbackHandler): - """Emit named NAT agent spans while preserving the active task/tool stack.""" - - run_inline = True - - def __init__(self, step_manager: Any) -> None: - super().__init__() - self._step_manager = step_manager - self._agent_names: dict[str, str] = {} - - def on_chain_start(self, serialized: dict[str, Any] | None, inputs: dict[str, Any], **kwargs: Any) -> None: - name = _deepagents_agent_name(kwargs) - run_id = str(kwargs.get("run_id", "")) - if not run_id or name is None: - return - - self._agent_names[run_id] = name - parent_run_id = kwargs.get("parent_run_id") - self._step_manager.push_intermediate_step( - IntermediateStepPayload( - UUID=run_id, - event_type=IntermediateStepType.WORKFLOW_START, - name=name, - metadata=TraceMetadata( - provided_metadata={ - "agent_id": run_id, - "agent_name": name, - "span_role": "agent", - "langchain_parent_run_id": str(parent_run_id) if parent_run_id else None, - } - ), - ) - ) - - def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None: - self._end_agent_run(outputs=outputs, error=None, **kwargs) - - def on_chain_error(self, error: BaseException, **kwargs: Any) -> None: - self._end_agent_run(outputs=None, error=error, **kwargs) - - def _end_agent_run( - self, - *, - outputs: dict[str, Any] | None, - error: BaseException | None, - **kwargs: Any, - ) -> None: - run_id = str(kwargs.get("run_id", "")) - name = self._agent_names.pop(run_id, None) - if name is None: - return - - metadata = {"agent_id": run_id, "agent_name": name, "span_role": "agent"} - if error is not None: - metadata["error_type"] = type(error).__name__ - self._step_manager.push_intermediate_step( - IntermediateStepPayload( - UUID=run_id, - event_type=IntermediateStepType.WORKFLOW_END, - name=name, - metadata=TraceMetadata(provided_metadata=metadata), - ) - ) - - -class AIQLangchainProfilerHandler(LangchainProfilerHandler): - """Preserve NAT's profiler behavior while naming DeepAgents task spans.""" - - async def on_tool_start( - self, - serialized: dict[str, Any], - input_str: str, - *, - inputs: dict[str, Any] | None = None, - **kwargs: Any, - ) -> Any: - serialized = dict(serialized) - serialized["name"] = _task_display_name(serialized, input_str, inputs) - return await super().on_tool_start(serialized, input_str, inputs=inputs, **kwargs) - - -@contextmanager -def aiq_langchain_profiler_context() -> Iterator[AIQLangchainProfilerHandler]: - """Replace NAT's inherited profiler for one AIQ job without adding a duplicate callback.""" - from nat.plugins.profiler.decorators.framework_wrapper import callback_handler_var - - profiler = AIQLangchainProfilerHandler() - token = callback_handler_var.set(profiler) - try: - yield profiler - finally: - callback_handler_var.reset(token) diff --git a/frontends/benchmarks/deepresearch_bench/README.md b/frontends/benchmarks/deepresearch_bench/README.md index 6f4f4789a..37573b635 100644 --- a/frontends/benchmarks/deepresearch_bench/README.md +++ b/frontends/benchmarks/deepresearch_bench/README.md @@ -55,9 +55,11 @@ python frontends/benchmarks/deepresearch_bench/scripts/export_drb_jsonl.py --inp Follow instructions in the [Deep Research Bench Github Repository](https://github.com/Ayanami0730/deep_research_bench/tree/main) to run evaluation and obtain scores. -## Optional: Phoenix Tracing +## Optional: Relay and Phoenix Tracing -If your config enables Phoenix tracing, start the Phoenix server before running `nat eval`. +ATOF tracing is enabled through NeMo Relay by default. To visualize an +evaluation in Phoenix, start Phoenix and enable the Relay OpenInference OTEL +endpoint in the evaluated workflow. Start server (separate terminal): @@ -65,27 +67,26 @@ Start server (separate terminal): uvx --from arize-phoenix phoenix serve ``` -## W&B Tracking - -Evaluation runs are tracked using [Weights & Biases Weave](https://wandb.ai/site/weave/) for experiment tracking and observability. - -### Configuration - -Enable W&B tracking in your config file under `general.telemetry.tracing`: - ```yaml -general: - telemetry: - tracing: - weave: - _type: weave - project: "deep-researcher-v2" +workflow: + relay: + observability: + opentelemetry: + enabled: true + endpoints: + - type: openinference + endpoint: http://localhost:6006/v1/traces + resource_attributes: + openinference.project.name: aiq-deepresearch-bench eval: general: workflow_alias: "aiq-deepresearch-v2-baseline" ``` +See the main [observability guide](../../../docs/source/deployment/observability.md) +for ATOF inspection, trace reading, project selection, and cost reporting. + ### workflow_alias The `workflow_alias` parameter provides a workflow-specific identifier for tracking evaluation runs: diff --git a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml index 13242e217..13ee1011d 100644 --- a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml +++ b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml @@ -7,14 +7,6 @@ general: console: _type: console level: INFO - # tracing: - # phoenix: - # _type: phoenix - # endpoint: http://localhost:6006/v1/traces - # project: dev - # weave: - # _type: weave - # project: "nvidia-aiq/AIQ_v2_deepresearch_bench" use_uvloop: true llms: @@ -67,6 +59,9 @@ functions: workflow: _type: deep_research_workflow + relay: + pricing: + sources: [] eval: general: diff --git a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml index 263455318..96f33e543 100644 --- a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml +++ b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml @@ -7,14 +7,6 @@ general: console: _type: console level: INFO - # tracing: - # phoenix: - # _type: phoenix - # endpoint: http://localhost:6006/v1/traces - # project: dev - # weave: - # _type: weave - # project: "nvidia-aiq/AIQ_v2_deepresearch_bench" use_uvloop: true llms: @@ -67,6 +59,9 @@ functions: workflow: _type: deep_research_workflow + relay: + pricing: + sources: [] eval: general: diff --git a/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml b/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml index ff4f4de9b..35d1196a9 100644 --- a/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml +++ b/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml @@ -15,20 +15,23 @@ eval: tokenomics: pricing: models: - "gpt-5.2": - input_per_1m_tokens: 2.50 - output_per_1m_tokens: 10.00 - # Market-equivalent examples, not NVIDIA API Catalog billing rates. - # Verify current provider pricing before using these values for accounting. "nvidia/nemotron-3-ultra-550b-a55b": - input_per_1m_tokens: 0.60 - output_per_1m_tokens: 3.60 + # NVIDIA-hosted API access used by this profile. Self-hosting is not free. + input_per_1m_tokens: 0.00 + output_per_1m_tokens: 0.00 + "nvidia/nemotron-3.5-lightning-30b-a3b": + # NVIDIA-hosted API access used by this profile. Self-hosting is not free. + input_per_1m_tokens: 0.00 + output_per_1m_tokens: 0.00 tools: - # Key "web_search" matches "advanced_web_search_tool" via substring lookup. - "web_search": + # Tavily pay-as-you-go is $0.008/credit: basic uses 1 credit, + # advanced uses 2. Monthly plans have lower effective rates. + "web_search_tool": + cost_per_call: 0.008 + "advanced_web_search_tool": cost_per_call: 0.016 + # The shipped paper-search tool defaults to Serper. This is Serper's + # Starter rate ($50 / 50,000 successful queries); change it for your tier + # or when selecting SerpAPI/SearchAPI instead. "paper_search": - cost_per_call: 0.0003 - default: - input_per_1m_tokens: 1.00 - output_per_1m_tokens: 4.00 + cost_per_call: 0.001 diff --git a/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml b/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml index 038ff6dde..dace73cf4 100644 --- a/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml +++ b/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml @@ -7,10 +7,6 @@ general: console: _type: console level: INFO - tracing: - weave: - _type: weave - project: "nvidia-aiq/AIQ_v2_deepsearch_qa" use_uvloop: true llms: @@ -64,6 +60,9 @@ functions: workflow: _type: deep_research_workflow + relay: + pricing: + sources: [] eval: general: diff --git a/frontends/benchmarks/freshqa/configs/config_full_workflow.yml b/frontends/benchmarks/freshqa/configs/config_full_workflow.yml index 85747871f..d63fb0583 100644 --- a/frontends/benchmarks/freshqa/configs/config_full_workflow.yml +++ b/frontends/benchmarks/freshqa/configs/config_full_workflow.yml @@ -93,7 +93,6 @@ functions: - web_search_tool max_llm_turns: 10 max_tool_iterations: 5 - verbose: true deep_research_agent: _type: deep_research_agent @@ -109,9 +108,11 @@ functions: workflow: _type: chat_deepresearcher_agent enable_escalation: true - verbose: true tools: - web_search_tool + relay: + pricing: + sources: [] eval: general: diff --git a/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml b/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml index e8dbbf4b8..e278c1b66 100644 --- a/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml +++ b/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml @@ -7,10 +7,6 @@ general: console: _type: console level: INFO - tracing: - weave: - _type: weave - project: "AIQ_v2_freshqa" llms: # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. @@ -57,6 +53,9 @@ functions: workflow: _type: shallow_research_workflow + relay: + pricing: + sources: [] eval: general: diff --git a/frontends/cli/cli.py b/frontends/cli/cli.py index 03ebc2b11..9a690374f 100644 --- a/frontends/cli/cli.py +++ b/frontends/cli/cli.py @@ -396,14 +396,6 @@ def main() -> None: logging.basicConfig( level=logging.INFO, format="%(levelname)s - %(name)s - %(message)s", handlers=[logging.StreamHandler()] ) - callbacks_logger = logging.getLogger("aiq_agent.callbacks") - callbacks_logger.setLevel(logging.DEBUG) - - cb_handler = logging.StreamHandler() - cb_handler.setFormatter(logging.Formatter("%(message)s")) - callbacks_logger.handlers.clear() - callbacks_logger.addHandler(cb_handler) - callbacks_logger.propagate = False else: logging.basicConfig(level=logging.WARNING, format="%(levelname)s - %(name)s - %(message)s") diff --git a/mcp/uv.lock b/mcp/uv.lock index a16634ecd..fe112e07b 100644 --- a/mcp/uv.lock +++ b/mcp/uv.lock @@ -182,6 +182,7 @@ dependencies = [ { name = "langgraph-checkpoint-postgres" }, { name = "langgraph-checkpoint-sqlite" }, { name = "mcp" }, + { name = "nemo-relay", extra = ["deepagents", "langchain", "langgraph"] }, { name = "nvidia-nat", extra = ["async-endpoints", "langchain", "mcp", "phoenix"] }, { name = "nvidia-nat-core" }, { name = "nvidia-nat-eval" }, @@ -210,6 +211,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.5.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, + { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], specifier = ">=0.7.3,<0.8" }, { name = "nvidia-nat", extras = ["langchain", "async-endpoints", "phoenix", "mcp"], specifier = "==1.8.0" }, { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, @@ -1149,7 +1151,7 @@ wheels = [ [[package]] name = "deepagents" -version = "0.6.8" +version = "0.6.12" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain" }, @@ -1159,9 +1161,9 @@ dependencies = [ { name = "langsmith" }, { name = "wcmatch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/97/77/e3b7efd9bff9cd101c085a5a3bf74180c13ab6c41a96f725cd1cb1bf53e8/deepagents-0.6.8.tar.gz", hash = "sha256:70cdd4da920cc420a8a0f729792ec559688bbbff39f7ab1508110cce9f901c06", size = 196927, upload-time = "2026-06-03T17:08:36.724Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/db/a6acdc72a9e90c3f07ed10de35c951734a02d4facb693bb59684ad368801/deepagents-0.6.12.tar.gz", hash = "sha256:1f281c0bc5a63132f62e2ee345c1dc593b23188da6e23016401f6879fbe54b5f", size = 211364, upload-time = "2026-06-25T17:26:52.775Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8e/19/1b7b76e958ac7f4e40886edc70f67aff4d7188770ab68105c9c48cbeb769/deepagents-0.6.8-py3-none-any.whl", hash = "sha256:087bdc1458202a3436854cf180f7ec059d07d2114a6c232819e9ad6533a5174a", size = 221469, upload-time = "2026-06-03T17:08:35.133Z" }, + { url = "https://files.pythonhosted.org/packages/98/49/af7219b3c13520fee047bb807cfaefba17f8e4584c551d946773589a4f08/deepagents-0.6.12-py3-none-any.whl", hash = "sha256:28b8fa0119ca0a689e3e18e288c4634e4046062acfc87a1cb34289d3af3a1c88", size = 236120, upload-time = "2026-06-25T17:26:51.736Z" }, ] [[package]] @@ -2233,16 +2235,16 @@ wheels = [ [[package]] name = "langchain" -version = "1.3.11" +version = "1.3.15" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "langgraph" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b9/a2/91a7197c604a3ce1b774b3c10dd114c3c745c6186a304fc2573b3f94d400/langchain-1.3.11.tar.gz", hash = "sha256:f3cf9cd4d2329b1a03eb8fd92b9d73e4e58a4d52570d67725fc77fbe0f104b32", size = 633374, upload-time = "2026-06-22T23:00:33.44Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/1c/b84579174a8e82ed79f4c3e0cd5a7f2323facc5ccd4d1b8390e7d175b663/langchain-1.3.15.tar.gz", hash = "sha256:ab4b775b9703f7e37babe0b325dbbaef25573bda60ecf79f7850bc875f252795", size = 665047, upload-time = "2026-08-11T19:10:52.455Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/df/a4/3a181967294f8876362cc4ba36840d50b8286fa23bb3f5e602b69eb3cb1e/langchain-1.3.11-py3-none-any.whl", hash = "sha256:7ae011f95a09b22feea1e8ae4e43f0b6164aebf4c61b8ad845b45f72ff3a90a2", size = 133639, upload-time = "2026-06-22T23:00:31.619Z" }, + { url = "https://files.pythonhosted.org/packages/09/8d/ae721f4d68ff79a17110cabc9cb39b4568b0e3f1fe0a379b926c3f81d175/langchain-1.3.15-py3-none-any.whl", hash = "sha256:c0d2d0d51ed7da249e8ab7487173872059a9dd46fb071d905957485b7334f987", size = 147001, upload-time = "2026-08-11T19:10:50.846Z" }, ] [[package]] @@ -2317,9 +2319,10 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.4.8" +version = "1.5.6" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "httpx" }, { name = "jsonpatch" }, { name = "langchain-protocol" }, { name = "langsmith" }, @@ -2330,9 +2333,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/e3/bea6d0080acf183332f24dcd74c208aee5857cf8f783c3fb0bd86027d8fb/langchain_core-1.4.8.tar.gz", hash = "sha256:5bf1f8411077c904182ad8f975943d36adcbf579c4e017b3a118b719229ebf9a", size = 957974, upload-time = "2026-06-18T19:39:23.636Z" } +sdist = { url = "https://files.pythonhosted.org/packages/27/b9/893806b89f77e1271fe6e10ce41682ff5fe43d071564e1d8e39dbb5d4d6d/langchain_core-1.5.6.tar.gz", hash = "sha256:b5f73bd9688c457b31ec73657a0ad56948f889fae27acee79286e9c285632ee6", size = 984873, upload-time = "2026-08-17T21:26:35.921Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/13/d6/bdf6f0481cc57ef300d6b1eb48cf1400c0409be715d6eb3cabadd1142a09/langchain_core-1.4.8-py3-none-any.whl", hash = "sha256:d84c28b05e3ba8d4271d0827aad5b592ccdaaf986e76768c23503f0a2045e8aa", size = 557416, upload-time = "2026-06-18T19:39:21.902Z" }, + { url = "https://files.pythonhosted.org/packages/5d/0a/890504397885c9d1ae45f2e06c3000dd9f4445439602b6a990a88b32c0ac/langchain_core-1.5.6-py3-none-any.whl", hash = "sha256:d6cf37bf695ecc22cddeb8461a684e353190b2ce430d99eb22bc11c0c7c00ea5", size = 567016, upload-time = "2026-08-17T21:26:34.595Z" }, ] [[package]] @@ -2350,7 +2353,7 @@ wheels = [ [[package]] name = "langchain-google-genai" -version = "4.2.4" +version = "4.3.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filetype" }, @@ -2358,9 +2361,9 @@ dependencies = [ { name = "langchain-core" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/54/52/de168715eb092c920531d418b8b9aafdff9e37ee80e5fc88106211ccbd47/langchain_google_genai-4.2.4.tar.gz", hash = "sha256:2f5de7a8a6552ffb64b907aca7503fd5e34d1a3240e280abcdc5f7eef480edd5", size = 270054, upload-time = "2026-05-28T21:23:00.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/ae/8ba8ee41bd20a23dee95cda109632c8b19a53141fbc81d9f87a72f0e975c/langchain_google_genai-4.3.4.tar.gz", hash = "sha256:265655baad05f799fa7b83a030eaca7cee0e32c9ab7de846b80ea7f59c26134e", size = 287396, upload-time = "2026-08-14T18:10:00.119Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/a2/5feaf21cfe6fac80eae944f3ac5348d9e5e986813256f74f8dd104617474/langchain_google_genai-4.2.4-py3-none-any.whl", hash = "sha256:0e2c1021a15c91e60b68d813bb3e793bd1d9396b3f8639b943ab4e56e5652e04", size = 68832, upload-time = "2026-05-28T21:22:59.291Z" }, + { url = "https://files.pythonhosted.org/packages/c9/f8/9fe4a28e319e9d6b20454e85fea9c243bf011dc866f2c3b9a89d64f9c1a1/langchain_google_genai-4.3.4-py3-none-any.whl", hash = "sha256:618fb0da1b9ba9def5569a8b05cb87e1389de41b8731c802958d09499490d2a5", size = 73338, upload-time = "2026-08-14T18:09:58.772Z" }, ] [[package]] @@ -2509,7 +2512,7 @@ wheels = [ [[package]] name = "langgraph" -version = "1.2.8" +version = "1.2.11" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, @@ -2519,9 +2522,9 @@ dependencies = [ { name = "pydantic" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b6/ad/583fda4c69501390b989770a465ccd0bdab1c1612eba582c012002ddf9b6/langgraph-1.2.8.tar.gz", hash = "sha256:f79d3575f45b404899358976e4fac0294eb75f8df1bfe8cd11286be7539c4548", size = 722464, upload-time = "2026-07-06T20:40:19.487Z" } +sdist = { url = "https://files.pythonhosted.org/packages/56/0d/c8e7ee98896659e1b6555db0ab115a9ca899844744645d5d894032bab1d7/langgraph-1.2.11.tar.gz", hash = "sha256:9ecfe11e50d338b34b15cf4d8a442642de103e8ae6971320efba84e4542eb363", size = 725753, upload-time = "2026-08-11T14:00:36.945Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/36/49/b958a9963606807e5a20cc75fced14aa77c5cbcc470d5bf8ae13277cd298/langgraph-1.2.8-py3-none-any.whl", hash = "sha256:aa8de1d4df44162353d117589ae0bf6930ca009b62d2d6e26cc32580794c5be6", size = 246983, upload-time = "2026-07-06T20:40:18.242Z" }, + { url = "https://files.pythonhosted.org/packages/0a/7f/c5c30e4be99ff821029c7ac872a480676bb179c9f3df85ea3f38d13f86d4/langgraph-1.2.11-py3-none-any.whl", hash = "sha256:8bab70de7b2d00b5300fb289bcf38d8b241400f3184c1e95e8ce706fb0e8686b", size = 248854, upload-time = "2026-08-11T14:00:35.494Z" }, ] [[package]] @@ -3300,6 +3303,46 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/ca/36339329c4604adbcc99c899b7eb1ce1a555c499b6a6860757dc9bfed36d/narwhals-2.22.1-py3-none-any.whl", hash = "sha256:60567d774edf77db53906f89d9fbd164e66e56d66d388e1e6990f17ac33cfb53", size = 454815, upload-time = "2026-06-05T12:34:32.289Z" }, ] +[[package]] +name = "nemo-relay" +version = "0.7.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/a259aac8df4aa78c0b3a6f3ad0fbf6305666bfaad5c29d9adc26db0f9e27/nemo_relay-0.7.3.tar.gz", hash = "sha256:ea5a1bb52e25e001dcbf6af1830616be181845e978cc848df58562556bba5604", size = 1299430, upload-time = "2026-08-14T14:40:44.832Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/9e/4eb80d2307cadcbb839dad2212a8d888667cd8c531ad0d8c0c1afc841181/nemo_relay-0.7.3-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:301dfc8334032ac52c09cc0b1421181e7c4df601abe7c82dfe51aa74ddbb3732", size = 9251369, upload-time = "2026-08-14T14:40:07.668Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a9/9fb77f7142b1381d8c3c81fdbb76782a02dcff1ee403ac6c93b13fa46b24/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:49c70c0a94cebb6cba3dd7521be11f63eac4cd9386881eb29827ac91b1bd780b", size = 8458046, upload-time = "2026-08-14T14:40:10.181Z" }, + { url = "https://files.pythonhosted.org/packages/ae/ab/b2f246f971f561a982d234d9f3ec1b29dfa4c72bf4882f3e32ce6db54dfa/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8eeaaf8c6a18440e473dd14eb1bb82e56c6514e70adee7456eddf9ce217cff89", size = 8957931, upload-time = "2026-08-14T14:40:12.547Z" }, + { url = "https://files.pythonhosted.org/packages/8f/d8/d8c25ba915467bab457d175b84a3af5c33291655867d472c49eada3b77e6/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:90ff984a89c42ebd0cfe26a0af3f180b1970e2ebcd743d19a960e547949ad2ef", size = 10325640, upload-time = "2026-08-14T14:40:15.722Z" }, + { url = "https://files.pythonhosted.org/packages/7d/fd/6ec48f47eb5ca4cca566b197ade8514baf183dddf53d6de3a296b8bc1102/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d4895a4df39a92ecbac3633fb4ec59cf1f83286c0369f433760ab28fb8d6dddf", size = 10706736, upload-time = "2026-08-14T14:40:18.377Z" }, + { url = "https://files.pythonhosted.org/packages/6b/cb/8a6f8d5f9922e75100f135c1dc4451bafc15f5d51608081f72b894e0caf4/nemo_relay-0.7.3-cp311-abi3-win_amd64.whl", hash = "sha256:f123cd45a27fca3d570559f2c26850138763411f66f387559d757e5d88bfa3c1", size = 8807604, upload-time = "2026-08-14T14:40:20.98Z" }, + { url = "https://files.pythonhosted.org/packages/f1/ac/558a05e6b9e28464b64fee8327d9e8a2aab3fb356363b8b82597b6b35b2a/nemo_relay-0.7.3-cp311-abi3-win_arm64.whl", hash = "sha256:6d5444e9a03b8b5d4bba2409105a121349580cb51481d441c6b5e699713a8763", size = 8442944, upload-time = "2026-08-14T14:40:23.433Z" }, +] + +[package.optional-dependencies] +deepagents = [ + { name = "deepagents" }, + { name = "langchain" }, + { name = "langchain-anthropic" }, + { name = "langchain-core" }, + { name = "langgraph" }, + { name = "langgraph-checkpoint" }, + { name = "langsmith" }, +] +langchain = [ + { name = "langchain" }, + { name = "langchain-core" }, + { name = "langgraph" }, + { name = "langgraph-checkpoint" }, + { name = "langsmith" }, +] +langgraph = [ + { name = "langchain" }, + { name = "langchain-core" }, + { name = "langgraph" }, + { name = "langgraph-checkpoint" }, + { name = "langsmith" }, +] + [[package]] name = "nemoguardrails" version = "0.21.0" diff --git a/pyproject.toml b/pyproject.toml index 32a0ddf72..e2066ee1f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ "nvidia-nat-profiler==1.8.0", "nvidia-nat-redis==1.8.0", "nvidia-nat-security[guardrails]==1.8.0", + "nemo-relay[deepagents,langchain,langgraph]>=0.7.3,<0.8", "deepagents>=0.6.5", "langgraph-checkpoint-postgres>=3.0.0", "langgraph-checkpoint-sqlite>=2.0.0", diff --git a/scripts/start_cli.sh b/scripts/start_cli.sh index a4e3b0eab..1d157c18b 100755 --- a/scripts/start_cli.sh +++ b/scripts/start_cli.sh @@ -37,7 +37,7 @@ while [[ $# -gt 0 ]]; do echo "" echo "Options:" echo " --config_file PATH Config file (default: configs/config_cli_default.yml)" - echo " -v, --verbose Enable verbose tracing for all agents" + echo " -v, --verbose Show detailed console logs" echo " -h, --help Show this help" echo "" echo "Available configs in configs/:" @@ -67,10 +67,6 @@ fi export AIQ_DEV_ENV=cli -if [ "$CLI_VERBOSE" = "true" ]; then - export AIQ_VERBOSE=true -fi - echo "============================================" echo " AI-Q Blueprint - CLI Mode" echo "============================================" diff --git a/src/aiq_agent/agents/chat_researcher/agent.py b/src/aiq_agent/agents/chat_researcher/agent.py index f6dcdc9cc..5aea3572a 100644 --- a/src/aiq_agent/agents/chat_researcher/agent.py +++ b/src/aiq_agent/agents/chat_researcher/agent.py @@ -48,6 +48,7 @@ from aiq_agent.common.citation_verification import EmptySourceRegistryError from aiq_agent.common.logging_utils import log_content_metadata from aiq_agent.common.logging_utils import log_identifier_ref +from aiq_agent.relay import run_agent try: from aiq_api.auth.errors import AuthError as _AuthError @@ -175,7 +176,11 @@ def _build_graph(self) -> CompiledStateGraph: async def intent_classifier_node(state: ChatResearcherState) -> dict[str, Any]: try: - return await self.intent_classifier_fn(state) + return await run_agent( + "intent_classifier", + lambda: self.intent_classifier_fn(state), + input_value=state, + ) except Exception as error: logger.warning("Intent routing failed (error_type=%s)", type(error).__name__) return { @@ -655,7 +660,17 @@ async def run( if messages: query = messages[-1].content logger.info("Query: %s", log_content_metadata(query or "")) - result = await self._graph.ainvoke(input_state, config=graph_config) + + async def _invoke_graph() -> dict[str, Any]: + effective_config = dict(graph_config or {}) + return await self._graph.ainvoke(input_state, config=effective_config) + + result = await run_agent( + "chat_deepresearcher_agent", + _invoke_graph, + session_id=thread_id, + input_value=input_state, + ) logger.info("ChatResearcherAgent: Workflow complete") diff --git a/src/aiq_agent/agents/chat_researcher/nodes/intent_classifier.py b/src/aiq_agent/agents/chat_researcher/nodes/intent_classifier.py index 3e6ea3afe..edcab04da 100644 --- a/src/aiq_agent/agents/chat_researcher/nodes/intent_classifier.py +++ b/src/aiq_agent/agents/chat_researcher/nodes/intent_classifier.py @@ -31,6 +31,7 @@ from aiq_agent.common import load_prompt from aiq_agent.common import render_prompt_template from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay import ainvoke_with_relay from ..models import RESEARCH_WORKFLOW_FAILURE_ERROR from ..models import ChatResearcherState @@ -150,9 +151,10 @@ async def run(self, state: ChatResearcherState) -> dict[str, Any]: messages: list[BaseMessage] = [SystemMessage(content=system_content)] try: - config = {"callbacks": self.callbacks} if self.callbacks else {} + # This model call crosses the NAT function boundary, so it does not + # inherit the parent graph's RunnableConfig automatically. response = await asyncio.wait_for( - self.llm.ainvoke(messages, config=config), + ainvoke_with_relay(self.llm, messages, callbacks=self.callbacks), timeout=self.llm_timeout, ) @@ -162,7 +164,7 @@ async def run(self, state: ChatResearcherState) -> dict[str, Any]: parsed = await self._repair_json_response( system_content=system_content, invalid_response=response_text, - config=config, + callbacks=self.callbacks, ) if not parsed or not isinstance(parsed, dict): @@ -253,7 +255,7 @@ async def _repair_json_response( *, system_content: str, invalid_response: str, - config: dict[str, Any], + callbacks: list[Any], ) -> dict[str, Any] | None: repair_prompt = ( f"{system_content}\n\n" @@ -264,7 +266,11 @@ async def _repair_json_response( ) try: response = await asyncio.wait_for( - self.llm.ainvoke([SystemMessage(content=repair_prompt)], config=config), + ainvoke_with_relay( + self.llm, + [SystemMessage(content=repair_prompt)], + callbacks=callbacks, + ), timeout=min(self.llm_timeout, _REPAIR_TIMEOUT_SECONDS), ) except TimeoutError: diff --git a/src/aiq_agent/agents/chat_researcher/register.py b/src/aiq_agent/agents/chat_researcher/register.py index e8b3b5256..c06e25f9a 100644 --- a/src/aiq_agent/agents/chat_researcher/register.py +++ b/src/aiq_agent/agents/chat_researcher/register.py @@ -26,11 +26,9 @@ from pydantic import Field from pydantic import ValidationError -from aiq_agent.common import VerboseTraceCallback from aiq_agent.common import _create_chat_response from aiq_agent.common import format_data_source_tools from aiq_agent.common import get_checkpointer -from aiq_agent.common import is_verbose from aiq_agent.common.citation_verification import get_or_create_session_registry from aiq_agent.common.citation_verification import reset_session_registry from aiq_agent.common.citation_verification import set_session_registry @@ -39,6 +37,10 @@ from aiq_agent.observability.otel_header_redaction_exporter import ( ensure_registered as _ensure_otel_redaction_registered, ) +from aiq_agent.relay.bootstrap import ensure_started as _ensure_relay_started +from aiq_agent.relay.config import RelayConfig +from aiq_agent.relay.runtime import ainvoke_with_relay +from aiq_agent.relay.runtime import run_workflow from nat.builder.builder import Builder from nat.builder.context import Context from nat.builder.framework_enum import LLMFrameworkEnum @@ -138,7 +140,10 @@ async def _answer_from_report_context( source_summary_markdown=source_summary_markdown, ) try: - response = await asyncio.wait_for(llm.ainvoke([HumanMessage(content=prompt)]), timeout=_REPORT_ASK_TIMEOUT_S) + response = await asyncio.wait_for( + ainvoke_with_relay(llm, [HumanMessage(content=prompt)]), + timeout=_REPORT_ASK_TIMEOUT_S, + ) except TimeoutError: logger.warning("Report ask LLM call timed out after %ss", _REPORT_ASK_TIMEOUT_S) raise @@ -200,7 +205,6 @@ class IntentClassifierConfig(FunctionBaseConfig, name="intent_classifier"): default_factory=list, description="Tool names to exclude when inheriting from registry.", ) - verbose: bool = Field(default=False) llm_timeout: float = Field( default=90, description="Timeout in seconds for the intent-classification LLM call. Default 90 if not set.", @@ -227,8 +231,7 @@ async def intent_classifier(config: IntentClassifierConfig, builder: Builder): excluded = set(config.exclude_tools) tools = [t for t in tools if getattr(t, "name", "") not in excluded] - verbose = is_verbose(config.verbose) - callbacks = [VerboseTraceCallback()] if verbose else [] + callbacks: list[Any] = [] tools_info = [{"name": getattr(t, "name", str(t)), "description": getattr(t, "description", "")} for t in tools] classifier = IntentClassifier( @@ -306,7 +309,6 @@ class ChatDeepResearcherConfig(FunctionBaseConfig, name="chat_deepresearcher_age max_history: int = Field( default=20, description="Maximum number of messages to keep in history before invoking the agent" ) - verbose: bool = Field(default=False, description="Enable verbose logging") enable_clarifier: bool = Field(default=False, description="Enable clarification of research queries") use_async_deep_research: bool = Field( default=False, @@ -317,6 +319,7 @@ class ChatDeepResearcherConfig(FunctionBaseConfig, name="chat_deepresearcher_age default="./checkpoints.db", description="SQLite database path or Postgres DSN for persistent checkpoints.", ) + relay: RelayConfig = Field(default_factory=RelayConfig, description="NeMo Relay plugins and export destinations") @register_function(config_type=ChatDeepResearcherConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) @@ -327,6 +330,8 @@ async def chat_deepresearcher_agent(config: ChatDeepResearcherConfig, builder: B Coordinates intent classification, depth routing, and research agents to produce research results based on user queries. """ + await _ensure_relay_started(config.relay) + import os import sys @@ -446,8 +451,7 @@ def validate_deep_research_tools(data_sources: list[str] | None) -> tuple[bool, return True, "" - verbose = is_verbose(config.verbose) - callbacks = [VerboseTraceCallback()] if verbose else [] + callbacks: list[Any] = [] # LLM for inline report Q&A: prefer the report writer model, fall back to the # deep researcher's orchestrator LLM (always configured). Report ask is a single @@ -531,6 +535,7 @@ async def _inline_report_edit(state: ChatResearcherState) -> str: edit_instruction=instruction, source_summary=report_context.source_summary_markdown, parent_context=report_context.model_dump_json(indent=2, exclude={"report_markdown"}), + callbacks=callbacks, ) async def _build_report_seed_files(state: ChatResearcherState) -> dict[str, str]: @@ -627,10 +632,9 @@ async def _submit_deep_job(state: ChatResearcherState) -> str: validate_deep_research_tools_fn=validate_deep_research_tools, ) - async def _run(query: object) -> ChatResearcherResponse: + async def _run_impl(query: object, nat_context_conversation_id: str) -> ChatResearcherResponse: import os import sys - import uuid # Check if API keys are missing and return graceful error response if api_key_error_response: @@ -650,24 +654,18 @@ def exit_after_error(): workflow_outcome=WorkflowFailure(error=RESEARCH_WORKFLOW_FAILURE_ERROR), ) - # For --input mode, use a fresh conversation_id to avoid loading old checkpoint state - # This ensures each run starts with a clean conversation history if "--input" in sys.argv: - nat_context_conversation_id = str(uuid.uuid4()) _log_conversation_reference( "Using fresh conversation reference for --input mode: %s", nat_context_conversation_id, ) + elif Context.get().conversation_id: + _log_conversation_reference("Thread reference for checkpointing: %s", nat_context_conversation_id) else: - nat_context_conversation_id = Context.get().conversation_id - if not nat_context_conversation_id: - nat_context_conversation_id = str(uuid.uuid4()) - _log_conversation_reference( - "No conversation-id header; generated thread reference: %s", - nat_context_conversation_id, - ) - else: - _log_conversation_reference("Thread reference for checkpointing: %s", nat_context_conversation_id) + _log_conversation_reference( + "No conversation-id header; generated thread reference: %s", + nat_context_conversation_id, + ) from aiq_agent.auth import get_current_principal @@ -791,4 +789,20 @@ def exit_after_response(): return response + async def _run(query: object) -> ChatResearcherResponse: + import sys + import uuid + + context = Context.get() + nat_context_conversation_id = ( + str(uuid.uuid4()) if "--input" in sys.argv or not context.conversation_id else context.conversation_id + ) + + return await run_workflow( + workflow_id, + lambda: _run_impl(query, nat_context_conversation_id), + session_id=nat_context_conversation_id, + input_value=query, + ) + yield FunctionInfo.from_fn(_run, description="Chat deep researcher with intent routing and escalation.") diff --git a/src/aiq_agent/agents/clarifier/agent.py b/src/aiq_agent/agents/clarifier/agent.py index 4b5ef8a2a..e37ee93af 100644 --- a/src/aiq_agent/agents/clarifier/agent.py +++ b/src/aiq_agent/agents/clarifier/agent.py @@ -56,6 +56,7 @@ from langgraph.graph import StateGraph from langgraph.graph.state import CompiledStateGraph from langgraph.prebuilt import ToolNode +from nemo_relay.integrations.langchain import NemoRelayMiddleware from aiq_agent.common import LLMProvider from aiq_agent.common import LLMRole @@ -64,6 +65,8 @@ from aiq_agent.common import load_prompt from aiq_agent.common import render_prompt_template from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay import ainvoke_with_relay +from aiq_agent.relay import run_agent from .models import ClarificationResponse from .models import ClarifierAgentState @@ -156,7 +159,6 @@ def __init__( user_prompt_callback: Callable[[str], Awaitable[str]], max_turns: int = 3, log_response_max_chars: int = 2000, - verbose: bool = False, callbacks: list[Any] | None = None, ) -> None: """ @@ -172,7 +174,6 @@ def __init__( automatically completing clarification. Defaults to 3. log_response_max_chars: Maximum characters to log from LLM responses. Used for debugging. Defaults to 2000. - verbose: Whether to enable detailed logging. Defaults to False. callbacks: Optional list of LangChain callback handlers for tracing and logging. """ @@ -181,7 +182,6 @@ def __init__( self.user_prompt_callback = user_prompt_callback self.max_turns = max_turns self.log_response_max_chars = log_response_max_chars - self.verbose = verbose self.callbacks = callbacks or [] self.system_prompt = self._load_default_prompt() @@ -520,7 +520,7 @@ async def agent_node(state: ClarifierAgentState): logger.info("Adding JSON reminder after tool results") messages.append(HumanMessage(content=JSON_REMINDER_AFTER_TOOLS)) - response = await bound_llm.ainvoke(messages) + response = await ainvoke_with_relay(bound_llm, messages, callbacks=self.callbacks) # Search-before-clarify (issue #234): on the first turn, if the model # asks for clarification without searching, nudge it once (guidance as @@ -541,7 +541,7 @@ async def agent_node(state: ClarifierAgentState): logger.info("Clarifier: model skipped search before clarifying; injecting guidance and retrying once") retry_system = SystemMessage(content=f"{rendered_system_prompt}\n\n{FORCE_SEARCH_GUIDANCE}") retry_messages = [retry_system, *messages[1:], response] - retry_response = await bound_llm.ainvoke(retry_messages) + retry_response = await ainvoke_with_relay(bound_llm, retry_messages, callbacks=self.callbacks) return {"messages": [retry_response]} return {"messages": [response]} @@ -633,7 +633,8 @@ def decide_route(state: ClarifierAgentState | dict): return "ask_for_clarification" graph.add_node("agent", agent_node) - graph.add_node("tools", ToolNode(self.tools)) + relay_middleware = NemoRelayMiddleware() + graph.add_node("tools", ToolNode(self.tools, awrap_tool_call=relay_middleware.awrap_tool_call)) graph.add_node("ask_for_clarification", ask_clarification) graph.set_entry_point("agent") @@ -666,7 +667,11 @@ async def run(self, state: ClarifierAgentState) -> ClarifierResult: logger.info("Clarifier: Starting (max %d turns)", self.max_turns) query = get_latest_user_query(state.messages) logger.info("User query: %s", log_content_metadata(query or "")) - result = await self._graph.ainvoke(state, config={"callbacks": self.callbacks}) + + async def _invoke_graph() -> dict[str, Any]: + return await self._graph.ainvoke(state, config={"callbacks": self.callbacks}) + + result = await run_agent("clarifier_agent", _invoke_graph, input_value=state) final_state = ClarifierAgentState.model_validate(result) return ClarifierResult(clarifier_log=final_state.clarifier_log) diff --git a/src/aiq_agent/agents/clarifier/register.py b/src/aiq_agent/agents/clarifier/register.py index 10ac0946f..5a5d651a0 100644 --- a/src/aiq_agent/agents/clarifier/register.py +++ b/src/aiq_agent/agents/clarifier/register.py @@ -28,7 +28,6 @@ tools: - web_search_tool max_turns: 3 - verbose: true """ import logging @@ -36,10 +35,8 @@ from pydantic import Field from aiq_agent.common import LLMProvider -from aiq_agent.common import VerboseTraceCallback from aiq_agent.common import all_mapped_tools_filtered_out from aiq_agent.common import filter_tools_by_sources -from aiq_agent.common import is_verbose from nat.builder.builder import Builder from nat.builder.context import Context from nat.builder.framework_enum import LLMFrameworkEnum @@ -68,7 +65,6 @@ class ClarifierConfig(FunctionBaseConfig, name="clarifier_agent"): tools: List of tool references for context gathering (e.g., web search). max_turns: Maximum number of clarification Q&A turns before auto-completing. log_response_max_chars: Maximum characters to log from LLM responses. - verbose: Whether to enable verbose logging with VerboseTraceCallback. """ llm: LLMRef = Field(..., description="LLM to use for generating questions") @@ -88,10 +84,6 @@ class ClarifierConfig(FunctionBaseConfig, name="clarifier_agent"): default=2000, description="Max characters to log from LLM responses", ) - verbose: bool = Field( - default=False, - description="Whether to enable verbose logging", - ) @register_function(config_type=ClarifierConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) @@ -141,8 +133,7 @@ async def clarifier_agent(config: ClarifierConfig, builder: Builder): provider = LLMProvider() provider.set_default(llm) - verbose = is_verbose(config.verbose) - callbacks = [VerboseTraceCallback(log_reasoning=True, max_chars=config.log_response_max_chars)] if verbose else [] + callbacks: list = [] async def user_prompt_callback(question: str) -> str: """ @@ -174,7 +165,6 @@ async def user_prompt_callback(question: str) -> str: user_prompt_callback=user_prompt_callback, max_turns=config.max_turns, log_response_max_chars=config.log_response_max_chars, - verbose=verbose, callbacks=callbacks, ) @@ -197,7 +187,6 @@ async def _run(state: ClarifierAgentState) -> ClarifierResult: user_prompt_callback=user_prompt_callback, max_turns=config.max_turns, log_response_max_chars=config.log_response_max_chars, - verbose=verbose, callbacks=callbacks, ) diff --git a/src/aiq_agent/agents/deep_researcher/agent.py b/src/aiq_agent/agents/deep_researcher/agent.py index c673473d3..d5403578f 100644 --- a/src/aiq_agent/agents/deep_researcher/agent.py +++ b/src/aiq_agent/agents/deep_researcher/agent.py @@ -37,6 +37,7 @@ from aiq_agent.common.citation_verification import source_entries_from_parent_context from aiq_agent.common.citation_verification import verify_citations from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay import run_agent from .custom_middleware import FinalReportCommitTracker from .custom_middleware import SourceRegistryMiddleware @@ -75,7 +76,6 @@ def __init__( llm_provider: LLMProvider, tools: Sequence[BaseTool] | None = None, *, - verbose: bool = True, callbacks: list[Any] | None = None, domain_catalog_path: str | None = None, enable_source_router: bool = True, @@ -97,7 +97,6 @@ def __init__( Args: llm_provider: LLMProvider for role-based LLM access. tools: Optional sequence of LangChain tools for research. - verbose: Enable detailed logging. callbacks: Optional list of callbacks. domain_catalog_path: Optional YAML/JSON domain catalog path for source-router-agent. enable_source_router: Enable the advisory source-router-agent before planning. @@ -114,7 +113,6 @@ def __init__( """ self.llm_provider = llm_provider self.tools = list(tools) if tools else [] - self.verbose = verbose self.callbacks = callbacks or [] self.max_research_concurrency = max_research_concurrency self.max_researcher_model_calls = max_researcher_model_calls @@ -315,9 +313,14 @@ async def run(self, state: DeepResearchAgentState) -> DeepResearchAgentState: execution_timeout = asyncio.timeout(self.resource_limits.max_execution_seconds) try: async with execution_timeout: - result = await agent.ainvoke( - state, - config={"callbacks": self.callbacks} if self.callbacks else None, + + async def _invoke_orchestrator() -> Any: + return await agent.ainvoke(state, config={"callbacks": self.callbacks}) + + result = await run_agent( + "deep_research_agent", + _invoke_orchestrator, + input_value=state, ) except TimeoutError as exc: # An inner provider/tool may raise TimeoutError for its own operation. diff --git a/src/aiq_agent/agents/deep_researcher/custom_middleware.py b/src/aiq_agent/agents/deep_researcher/custom_middleware.py index 4c44d1c47..77c12311b 100644 --- a/src/aiq_agent/agents/deep_researcher/custom_middleware.py +++ b/src/aiq_agent/agents/deep_researcher/custom_middleware.py @@ -877,10 +877,23 @@ async def awrap_model_call(self, request, handler): new_tool_calls = [] for tc in msg.tool_calls: new_tool_calls.append({**tc, "name": self._sanitize_tool_name(tc["name"])}) - new_msg = AIMessage( - content=msg.content, - tool_calls=new_tool_calls, - id=msg.id, + additional_kwargs = dict(msg.additional_kwargs) + raw_tool_calls = additional_kwargs.get("tool_calls") + if isinstance(raw_tool_calls, list): + sanitized_raw_tool_calls = [] + for raw_tool_call in raw_tool_calls: + if not isinstance(raw_tool_call, dict) or not isinstance(raw_tool_call.get("function"), dict): + sanitized_raw_tool_calls.append(raw_tool_call) + continue + function = dict(raw_tool_call["function"]) + function["name"] = self._sanitize_tool_name(str(function.get("name") or "")) + sanitized_raw_tool_calls.append({**raw_tool_call, "function": function}) + additional_kwargs["tool_calls"] = sanitized_raw_tool_calls + new_msg = msg.model_copy( + update={ + "additional_kwargs": additional_kwargs, + "tool_calls": new_tool_calls, + } ) new_result.append(new_msg) else: diff --git a/src/aiq_agent/agents/deep_researcher/factory.py b/src/aiq_agent/agents/deep_researcher/factory.py index 2ed251cf2..bf0c6252c 100644 --- a/src/aiq_agent/agents/deep_researcher/factory.py +++ b/src/aiq_agent/agents/deep_researcher/factory.py @@ -41,6 +41,8 @@ from aiq_agent.common import LLMProvider from aiq_agent.common import LLMRole from aiq_agent.common import render_prompt_template +from aiq_agent.relay import deepagents_kwargs +from aiq_agent.relay import merge_langchain_middleware from .custom_middleware import RESEARCHER_FINALIZATION_MODEL_CALLS from .custom_middleware import ArtifactHarvestMiddleware @@ -403,6 +405,7 @@ def build_researcher_runnable( *(visibility_middleware or []), ] ) + middleware = merge_langchain_middleware(middleware) return create_agent( model=researcher_model, tools=researcher_tools, @@ -647,39 +650,44 @@ def build_deep_research_graph( orchestrator_tools = [*context.tool_set.helper_tools, research_batch_tool] agent = create_deep_agent( - model=context.llm_provider.get(LLMRole.ORCHESTRATOR), - tools=orchestrator_tools, - system_prompt=context.render_prompt( - "orchestrator", - clarifier_result=context.state.clarifier_result, - # Advertise only the tools the orchestrator can actually call. Source - # tools (incl. per-user MCP tools like Google Drive) are NOT directly - # callable here — the orchestrator delegates all source access through - # run_research_batch to the researcher, which holds those tools. Listing - # them under "Available Tools" made the orchestrator call them directly - # (e.g. per_user_mcp_client__google_drive_read_file), which the runtime - # rejects since they aren't bound to this agent. - tools=[{"name": t.name, "description": t.description} for t in orchestrator_tools], - enable_source_router=context.enable_source_router, - max_research_concurrency=context.max_research_concurrency, - execution_enabled=context.runtime.execution_enabled, - parent_report_context_available=context.parent_report_context_available, - ), - subagents=build_deep_research_subagents(context), - store=InMemoryStore(), - middleware=context.middleware( - [ - *context.middleware_set.orchestrator, - FinalReportOwnershipGuardMiddleware(), - StateMutationGuardMiddleware( - writer=False, - sandbox_enabled=context.runtime.execution_enabled, + **deepagents_kwargs( + dict( + model=context.llm_provider.get(LLMRole.ORCHESTRATOR), + name="deep_research_agent", + tools=orchestrator_tools, + system_prompt=context.render_prompt( + "orchestrator", + clarifier_result=context.state.clarifier_result, + # Advertise only the tools the orchestrator can actually call. Source + # tools (incl. per-user MCP tools like Google Drive) are NOT directly + # callable here — the orchestrator delegates all source access through + # run_research_batch to the researcher, which holds those tools. Listing + # them under "Available Tools" made the orchestrator call them directly + # (e.g. per_user_mcp_client__google_drive_read_file), which the runtime + # rejects since they aren't bound to this agent. + tools=[{"name": t.name, "description": t.description} for t in orchestrator_tools], + enable_source_router=context.enable_source_router, + max_research_concurrency=context.max_research_concurrency, + execution_enabled=context.runtime.execution_enabled, + parent_report_context_available=context.parent_report_context_available, ), - TodoQuotaMiddleware(resource_limits=context.resource_limits), - RequiredWriterDelegationMiddleware(tracker=context.final_report_tracker), - ] - ), - permissions=context.permissions(ORCHESTRATOR_AGENT), - backend=context.backend, + subagents=build_deep_research_subagents(context), + store=InMemoryStore(), + middleware=context.middleware( + [ + *context.middleware_set.orchestrator, + FinalReportOwnershipGuardMiddleware(), + StateMutationGuardMiddleware( + writer=False, + sandbox_enabled=context.runtime.execution_enabled, + ), + TodoQuotaMiddleware(resource_limits=context.resource_limits), + RequiredWriterDelegationMiddleware(tracker=context.final_report_tracker), + ] + ), + permissions=context.permissions(ORCHESTRATOR_AGENT), + backend=context.backend, + ) + ) ) return agent.with_config({"recursion_limit": 2000}) diff --git a/src/aiq_agent/agents/deep_researcher/register.py b/src/aiq_agent/agents/deep_researcher/register.py index 38f368f6f..88faac7dd 100644 --- a/src/aiq_agent/agents/deep_researcher/register.py +++ b/src/aiq_agent/agents/deep_researcher/register.py @@ -27,14 +27,14 @@ from aiq_agent.common import LLMProvider from aiq_agent.common import LLMRole -from aiq_agent.common import VerboseTraceCallback from aiq_agent.common import _create_chat_response from aiq_agent.common import all_mapped_tools_filtered_out from aiq_agent.common import filter_tools_by_sources -from aiq_agent.common import is_verbose from aiq_agent.common import validate_research_source_configuration from aiq_agent.common.citation_verification import EmptySourceRegistryError from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay.bootstrap import ensure_started as _ensure_relay_started +from aiq_agent.relay.config import RelayConfig from nat.builder.builder import Builder from nat.builder.framework_enum import LLMFrameworkEnum from nat.builder.function_info import FunctionInfo @@ -79,7 +79,6 @@ class DeepResearchAgentConfig(FunctionBaseConfig, name="deep_research_agent"): default_factory=list, description="Tool names to exclude when inheriting from registry.", ) - verbose: bool = Field(default=True) domain_catalog_path: str | None = Field( default=None, description="Optional YAML/JSON domain catalog path for source-router-agent.", @@ -239,13 +238,11 @@ async def deep_research_agent(config: DeepResearchAgentConfig, builder: Builder) writer_llm = await builder.get_llm(config.writer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN) provider.configure(LLMRole.REPORT_WRITER, writer_llm) - verbose = is_verbose(config.verbose) - callbacks = [VerboseTraceCallback()] if verbose else [] + callbacks: list = [] agent = DeepResearcherAgent( llm_provider=provider, tools=tools, - verbose=verbose, callbacks=callbacks, domain_catalog_path=config.domain_catalog_path, enable_source_router=config.enable_source_router, @@ -283,7 +280,6 @@ async def _run(state: DeepResearchAgentState) -> DeepResearchAgentState: active_agent = DeepResearcherAgent( llm_provider=provider, tools=selected_tools, - verbose=verbose, callbacks=callbacks, domain_catalog_path=config.domain_catalog_path, enable_source_router=config.enable_source_router, @@ -335,11 +331,13 @@ class DeepResearchWorkflowConfig(FunctionBaseConfig, name="deep_research_workflo default=False, description="Submit deep research as an async job instead of running inline", ) + relay: RelayConfig = Field(default_factory=RelayConfig, description="NeMo Relay plugins and export destinations") @register_function(config_type=DeepResearchWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) async def deep_research_workflow(config: DeepResearchWorkflowConfig, builder: Builder): """Wrapper workflow that accepts string queries for evaluation.""" + await _ensure_relay_started(config.relay) deep_research_agent_fn = await builder.get_function("deep_research_agent") workflow_id = config.name or config.type diff --git a/src/aiq_agent/agents/deep_researcher/tools/research.py b/src/aiq_agent/agents/deep_researcher/tools/research.py index 6a43d2640..f490b5c0e 100644 --- a/src/aiq_agent/agents/deep_researcher/tools/research.py +++ b/src/aiq_agent/agents/deep_researcher/tools/research.py @@ -25,6 +25,7 @@ from typing import Any from typing import cast +import nemo_relay from langchain.agents.middleware.model_call_limit import ModelCallLimitExceededError from langchain.tools import ToolRuntime from langchain_core.messages import HumanMessage @@ -32,6 +33,7 @@ from langchain_core.tools import tool from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay import agent_scope from ..custom_middleware import ResearcherBudgetExhaustedError from ..models import EvidenceJudgment @@ -116,31 +118,33 @@ async def _run_research_query( ) -> ResearchNotes: """Run one researcher worker and return its structured notes.""" async with semaphore: - try: - result = await researcher_runnable.ainvoke( - researcher_invoke_state(query, runtime), - config=researcher_invoke_config(runtime, callbacks), - ) - except (ModelCallLimitExceededError, ResearcherBudgetExhaustedError): - logger.warning( - "Researcher worker exhausted its model-call budget (query_%s)", - log_content_metadata(query.query), - ) - return _exhausted_research_notes(query) - except Exception as exc: # noqa: BLE001 - captured as per-item failure - raise RuntimeError(f"researcher worker failed for query {query.query!r}: {exc}") from exc - - try: - structured = result.get("structured_response") if isinstance(result, dict) else None - if structured is None: - raise ValueError("researcher worker did not return structured ResearchNotes") - note = ResearchNotes.model_validate(structured) - except Exception as exc: # noqa: BLE001 - captured as per-item failure - raise ValueError( - f"researcher worker returned invalid ResearchNotes for query {query.query!r}: {exc}" - ) from exc + with agent_scope(RESEARCHER_AGENT_NAME, input_value=query) as lifecycle: + try: + result = await researcher_runnable.ainvoke( + researcher_invoke_state(query, runtime), + config=researcher_invoke_config(runtime, callbacks), + ) + except (ModelCallLimitExceededError, ResearcherBudgetExhaustedError): + logger.warning( + "Researcher worker exhausted its model-call budget (query_%s)", + log_content_metadata(query.query), + ) + return _exhausted_research_notes(query) + except Exception as exc: # noqa: BLE001 - captured as per-item failure + raise RuntimeError(f"researcher worker failed for query {query.query!r}: {exc}") from exc + + try: + structured = result.get("structured_response") if isinstance(result, dict) else None + if structured is None: + raise ValueError("researcher worker did not return structured ResearchNotes") + note = ResearchNotes.model_validate(structured) + except Exception as exc: # noqa: BLE001 - captured as per-item failure + raise ValueError( + f"researcher worker returned invalid ResearchNotes for query {query.query!r}: {exc}" + ) from exc - return note + lifecycle.output = note + return note def _research_note_slug(text: str) -> str: @@ -201,19 +205,20 @@ async def _run_research_queries( ) -> tuple[list[ResearchQuery], list[ResearchNotes], list[str]]: """Run researcher workers concurrently and collect successful query/note pairs plus surfaced errors.""" semaphore = asyncio.Semaphore(min(max_concurrency, len(queries))) - raw_results = await asyncio.gather( - *( + tasks = [ + asyncio.create_task( _run_research_query( query=query, researcher_runnable=researcher_runnable, runtime=runtime, callbacks=callbacks, semaphore=semaphore, - ) - for query in queries - ), - return_exceptions=True, - ) + ), + context=nemo_relay.fork_asyncio_context(), + ) + for query in queries + ] + raw_results = await asyncio.gather(*tasks, return_exceptions=True) successful_queries: list[ResearchQuery] = [] notes: list[ResearchNotes] = [] diff --git a/src/aiq_agent/agents/deep_researcher/tools/source_tool_batching.py b/src/aiq_agent/agents/deep_researcher/tools/source_tool_batching.py index b45498072..84ac91227 100644 --- a/src/aiq_agent/agents/deep_researcher/tools/source_tool_batching.py +++ b/src/aiq_agent/agents/deep_researcher/tools/source_tool_batching.py @@ -34,6 +34,7 @@ from pydantic import Field from aiq_agent.common.citation_verification import is_non_citable_status_output +from aiq_agent.relay import ainvoke_tool_with_relay from ..resource_limits import DEFAULT_MAX_CONSECUTIVE_SOURCE_TOOL_FAILURES @@ -273,7 +274,7 @@ async def _call_one(query: str) -> tuple[str, str | None, str | None]: async with limiter.limit(): await _ensure_source_tool_circuit_closed() try: - result = await original_tool.ainvoke({input_field_name: query}) + result = await ainvoke_tool_with_relay(original_tool, {input_field_name: query}) except SourceToolCircuitOpen: raise except Exception: # noqa: BLE001 - represented as per-item failure for the LLM @@ -315,7 +316,7 @@ async def _run_throttled(**kwargs) -> object: async with limiter.limit(): await _ensure_source_tool_circuit_closed() try: - result = await original_tool.ainvoke(kwargs) + result = await ainvoke_tool_with_relay(original_tool, kwargs) except SourceToolCircuitOpen: raise except Exception: diff --git a/src/aiq_agent/agents/report_rewriter/agent.py b/src/aiq_agent/agents/report_rewriter/agent.py index 8997f78e5..c00ed426d 100644 --- a/src/aiq_agent/agents/report_rewriter/agent.py +++ b/src/aiq_agent/agents/report_rewriter/agent.py @@ -26,6 +26,8 @@ from aiq_agent.common.citation_verification import source_entries_from_parent_context from aiq_agent.common.citation_verification import verify_citations from aiq_agent.common.logging_utils import log_identifier_ref +from aiq_agent.relay import ainvoke_with_relay +from aiq_agent.relay import run_agent from .models import ReportRewriterAgentState @@ -102,6 +104,7 @@ async def rewrite_report( source_summary: str = _DEFAULT_SOURCE_SUMMARY, parent_context: str = "{}", system_prompt: str | None = None, + callbacks: list[Any] | None = None, ) -> str: """Rewrite a report per an edit instruction with one bounded LLM call. @@ -120,12 +123,11 @@ async def rewrite_report( parent_context=parent_context, edit_instruction=instruction, ) - response = await llm.ainvoke( - [ - SystemMessage(content=rendered_prompt), - HumanMessage(content=instruction), - ] - ) + messages = [ + SystemMessage(content=rendered_prompt), + HumanMessage(content=instruction), + ] + response = await ainvoke_with_relay(llm, messages, callbacks=callbacks) revised_report = response.content if hasattr(response, "content") else str(response) revised_report = revised_report if isinstance(revised_report, str) else str(revised_report) revised_report = revised_report.strip() @@ -145,12 +147,10 @@ def __init__( llm_provider: LLMProvider, tools: Sequence[Any] | None = None, *, - verbose: bool = False, callbacks: list[Any] | None = None, job_id: str | None = None, ) -> None: self.llm_provider = llm_provider - self.verbose = verbose self.callbacks = callbacks or [] self.job_id = job_id self.system_prompt = load_prompt(AGENT_DIR / "prompts", "edit") @@ -186,14 +186,18 @@ async def run(self, state: ReportRewriterAgentState) -> ReportRewriterAgentState source_summary = self._read_text_file(state.files, SOURCE_SUMMARY_PATH) or _DEFAULT_SOURCE_SUMMARY parent_context = self._read_text_file(state.files, PARENT_CONTEXT_PATH) or "{}" - revised_report = await rewrite_report( - llm=self.llm_provider.get(LLMRole.REPORT_WRITER), - original_report=original_report, - edit_instruction=instruction, - source_summary=source_summary, - parent_context=parent_context, - system_prompt=self.system_prompt, - ) + async def _rewrite() -> str: + return await rewrite_report( + llm=self.llm_provider.get(LLMRole.REPORT_WRITER), + original_report=original_report, + edit_instruction=instruction, + source_summary=source_summary, + parent_context=parent_context, + system_prompt=self.system_prompt, + callbacks=self.callbacks, + ) + + revised_report = await run_agent("report_rewriter_agent", _rewrite, input_value=state) cited_urls = _verified_cited_urls( revised_report, _effective_parent_sources(original_report, parent_context), diff --git a/src/aiq_agent/agents/shallow_researcher/agent.py b/src/aiq_agent/agents/shallow_researcher/agent.py index c76576883..06d2c70c7 100644 --- a/src/aiq_agent/agents/shallow_researcher/agent.py +++ b/src/aiq_agent/agents/shallow_researcher/agent.py @@ -35,6 +35,7 @@ from langgraph.graph.state import CompiledStateGraph from langgraph.prebuilt import ToolNode from langgraph.prebuilt import tools_condition +from nemo_relay.integrations.langchain import NemoRelayMiddleware from aiq_agent.common import get_source_id_for_tool from aiq_agent.common import load_prompt @@ -49,6 +50,8 @@ from aiq_agent.common.citation_verification import sanitize_report from aiq_agent.common.citation_verification import verify_citations from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay import ainvoke_with_relay +from aiq_agent.relay import run_agent from ...common import LLMProvider from ...common import LLMRole @@ -360,8 +363,10 @@ async def _repair_missing_citations( try: response = await asyncio.wait_for( - self._get_llm().ainvoke( + ainvoke_with_relay( + self._get_llm(), [repair_system, *messages, repair_request], + callbacks=self.callbacks, config=repair_config, ), timeout=self.citation_repair_timeout, @@ -435,13 +440,23 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: ) full_messages = [system_message] + processed_history + [synthesis_anchor] - response = await self._get_llm().ainvoke(full_messages, config=draft_config) + response = await ainvoke_with_relay( + self._get_llm(), + full_messages, + callbacks=self.callbacks, + config=draft_config, + ) return {"messages": [response], "tool_iterations": iterations} llm = self._get_llm() llm_with_tools = llm.bind_tools(self.tools) if self.tools else llm full_messages = [system_message] + processed_history - response = await llm_with_tools.ainvoke(full_messages, config=draft_config) + response = await ainvoke_with_relay( + llm_with_tools, + full_messages, + callbacks=self.callbacks, + config=draft_config, + ) if self.tools and iterations == 0 and not getattr(response, "tool_calls", None): logger.warning("Shallow researcher returned an answer before collecting evidence; retrying once") @@ -452,8 +467,10 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: ) ) retry_llm = llm.bind_tools(self.tools, parallel_tool_calls=False) - response = await retry_llm.ainvoke( + response = await ainvoke_with_relay( + retry_llm, full_messages + [response, tool_required], + callbacks=self.callbacks, config=draft_config, ) retry_tool_calls = getattr(response, "tool_calls", None) or [] @@ -483,7 +500,8 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: builder.set_entry_point("agent") - tool_node = ToolNode(self.tools) + relay_middleware = NemoRelayMiddleware() + tool_node = ToolNode(self.tools, awrap_tool_call=relay_middleware.awrap_tool_call) # Per-agent allowlist mirrors the deep researcher: only tools this # agent was loaded with are candidates for source capture. The @@ -572,9 +590,12 @@ async def run(self, state: ShallowResearchAgentState) -> ShallowResearchAgentSta recursion_limit = (self.max_llm_turns * 2) + 10 config = {"recursion_limit": recursion_limit} - if self.callbacks: + + async def _invoke_graph() -> dict[str, Any]: config["callbacks"] = self.callbacks - result = await self._graph.ainvoke(state, config=config) + return await self._graph.ainvoke(state, config=config) + + result = await run_agent("shallow_research_agent", _invoke_graph, input_value=state) # Post-process: verify citations against source registry validated_result = dict(result) diff --git a/src/aiq_agent/agents/shallow_researcher/register.py b/src/aiq_agent/agents/shallow_researcher/register.py index 6332470bb..e0116b409 100644 --- a/src/aiq_agent/agents/shallow_researcher/register.py +++ b/src/aiq_agent/agents/shallow_researcher/register.py @@ -21,14 +21,14 @@ from pydantic import Field from aiq_agent.common import LLMProvider -from aiq_agent.common import VerboseTraceCallback from aiq_agent.common import _create_chat_response from aiq_agent.common import all_mapped_tools_filtered_out from aiq_agent.common import filter_tools_by_sources -from aiq_agent.common import is_verbose from aiq_agent.common import validate_research_source_configuration from aiq_agent.common.citation_verification import EmptySourceRegistryError from aiq_agent.common.logging_utils import log_content_metadata +from aiq_agent.relay.bootstrap import ensure_started as _ensure_relay_started +from aiq_agent.relay.config import RelayConfig from nat.builder.builder import Builder from nat.builder.framework_enum import LLMFrameworkEnum from nat.builder.function_info import FunctionInfo @@ -99,8 +99,7 @@ async def shallow_research_agent(config: ShallowResearchAgentConfig, builder: Bu provider = LLMProvider() provider.set_default(llm) - verbose = is_verbose(config.verbose) - callbacks = [VerboseTraceCallback()] if verbose else [] + callbacks: list = [] # No shared agent is built here: it is (re)built per request inside _run, since # the active tool set depends on the request's data_sources and the per-user MCP @@ -210,12 +209,13 @@ class ShallowResearchWorkflowConfig(FunctionBaseConfig, name="shallow_research_w for the shallow_research_agent. Use this as the workflow for evaluation. """ - pass + relay: RelayConfig = Field(default_factory=RelayConfig, description="NeMo Relay plugins and export destinations") @register_function(config_type=ShallowResearchWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) async def shallow_research_workflow(config: ShallowResearchWorkflowConfig, builder: Builder): """Wrapper workflow that accepts string queries for evaluation.""" + await _ensure_relay_started(config.relay) shallow_research_agent_fn = await builder.get_function("shallow_research_agent") workflow_id = config.name or config.type diff --git a/src/aiq_agent/common/__init__.py b/src/aiq_agent/common/__init__.py index 442249211..e84bd78bd 100644 --- a/src/aiq_agent/common/__init__.py +++ b/src/aiq_agent/common/__init__.py @@ -19,7 +19,6 @@ import datetime import logging -import os import aiosqlite from langgraph.checkpoint.base import BaseCheckpointSaver @@ -101,22 +100,6 @@ ] -# @environment_variable AIQ_VERBOSE -# @category Debug -# @type bool -# @default false -# @required false -# Enable verbose logging output. Accepts true/1/yes or false/0/no. -def is_verbose(config_verbose: bool) -> bool: - """Check if verbose mode is enabled via env var or config.""" - env_verbose = os.getenv("AIQ_VERBOSE", "").lower() - if env_verbose in ("true", "1", "yes"): - return True - if env_verbose in ("false", "0", "no"): - return False - return config_verbose - - def _create_chat_response( content: str, response_id: str = "conversational_response", diff --git a/src/aiq_agent/common/callbacks.py b/src/aiq_agent/common/callbacks.py index f68e90abb..77add2170 100644 --- a/src/aiq_agent/common/callbacks.py +++ b/src/aiq_agent/common/callbacks.py @@ -14,7 +14,6 @@ # limitations under the License. import logging -import os from typing import Any from langchain_core.callbacks import BaseCallbackHandler @@ -40,16 +39,12 @@ RESET_ALL = "\033[0m" -def is_verbose_enabled() -> bool: - return os.environ.get("AIQ_VERBOSE", "").lower() in ("1", "true", "yes") - - class ResearchLogger: """Colored logging utilities for research agents.""" - def __init__(self, logger_instance: logging.Logger, verbose: bool | None = None): + def __init__(self, logger_instance: logging.Logger, verbose: bool = False): self.logger = logger_instance - self.verbose = verbose if verbose is not None else is_verbose_enabled() + self.verbose = verbose def section(self, label: str, message: str, *args): self.logger.info(f"{BOLD}[{label}]{RESET_ALL} {message}", *args) diff --git a/src/aiq_agent/relay/__init__.py b/src/aiq_agent/relay/__init__.py new file mode 100644 index 000000000..ce35049f2 --- /dev/null +++ b/src/aiq_agent/relay/__init__.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AI-Q's NeMo Relay integration boundary. + +Relay is initialized for every AI-Q process. Agent-specific integrations live +behind this package so observability wiring does not leak into research logic. +""" + +from .config import RelayConfig +from .logging import register_logging_subscriber +from .runtime import agent_scope +from .runtime import ainvoke_tool_with_relay +from .runtime import ainvoke_with_relay +from .runtime import deepagents_kwargs +from .runtime import merge_langchain_middleware +from .runtime import run_agent +from .runtime import run_workflow +from .runtime import workflow_scope + +__all__ = [ + "agent_scope", + "ainvoke_tool_with_relay", + "ainvoke_with_relay", + "deepagents_kwargs", + "merge_langchain_middleware", + "register_logging_subscriber", + "RelayConfig", + "run_agent", + "run_workflow", + "workflow_scope", +] diff --git a/src/aiq_agent/relay/bootstrap.py b/src/aiq_agent/relay/bootstrap.py new file mode 100644 index 000000000..aaf303811 --- /dev/null +++ b/src/aiq_agent/relay/bootstrap.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AI-Q-owned lifecycle for NeMo Relay's plugin host.""" + +from __future__ import annotations + +import asyncio +import atexit +import json + +import nemo_relay +from nemo_relay import plugin + +from .config import RelayConfig +from .logging import register_logging_subscriber +from .privacy import deregister_privacy_sanitizers +from .privacy import register_privacy_sanitizers + +_lock = asyncio.Lock() +_active_config: str | None = None + + +async def ensure_started(config: RelayConfig | None = None) -> None: + """Initialize Relay's supported plugin host once for the effective AI-Q config.""" + global _active_config + relay_config = config or RelayConfig() + plugin_config = relay_config.to_plugin_config() + serialized = json.dumps(plugin_config, sort_keys=True) + async with _lock: + if _active_config == serialized: + return + plugin.validate(plugin_config) + await plugin.initialize(plugin_config) + register_privacy_sanitizers(relay_config.redaction) + if relay_config.logging: + register_logging_subscriber() + else: + nemo_relay.subscribers.deregister("aiq-relay-logging") + _active_config = serialized + + +def _reset_state() -> None: + global _active_config + nemo_relay.subscribers.deregister("aiq-relay-logging") + deregister_privacy_sanitizers() + _active_config = None + + +async def shutdown_async() -> None: + """Flush Relay exporters from an asynchronous application lifecycle.""" + await plugin.clear_async() + _reset_state() + + +def shutdown() -> None: + """Flush Relay exporters after the application event loop has stopped.""" + plugin.clear() + _reset_state() + + +atexit.register(shutdown) diff --git a/src/aiq_agent/relay/config.py b/src/aiq_agent/relay/config.py new file mode 100644 index 000000000..8e14a48d5 --- /dev/null +++ b/src/aiq_agent/relay/config.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed AI-Q configuration for NeMo Relay plugins.""" + +from __future__ import annotations + +from typing import Any +from typing import Literal + +from pydantic import AnyHttpUrl +from pydantic import BaseModel +from pydantic import ConfigDict +from pydantic import Field + + +class _RelayBaseConfig(BaseModel): + """Strict base for operator-facing Relay YAML.""" + + model_config = ConfigDict(extra="forbid") + + +class RelayAtofConfig(_RelayBaseConfig): + """Raw Relay event export configuration.""" + + enabled: bool = True + output_directory: str = "./relay" + filename: str = "aiq-relay.atof.jsonl" + mode: Literal["append", "overwrite"] = "append" + + +class RelayOpenTelemetryEndpointConfig(_RelayBaseConfig): + """One Relay 0.7 typed OTLP trace destination.""" + + type: Literal["full", "gen_ai", "openinference"] = "openinference" + endpoint: AnyHttpUrl = AnyHttpUrl("http://localhost:6006/v1/traces") + transport: Literal["http_binary", "grpc"] = "http_binary" + service_name: str = "aiq-relay" + service_namespace: str | None = None + service_version: str | None = None + instrumentation_scope: str = "nemo-relay" + timeout_millis: int = Field(default=3000, gt=0) + header_env: dict[str, str] = Field(default_factory=dict) + resource_attributes: dict[str, str] = Field(default_factory=lambda: {"openinference.project.name": "aiq-relay"}) + + +class RelayOpenTelemetryConfig(_RelayBaseConfig): + """Relay OpenTelemetry fan-out configuration.""" + + enabled: bool = False + endpoints: list[RelayOpenTelemetryEndpointConfig] = Field( + default_factory=lambda: [RelayOpenTelemetryEndpointConfig()] + ) + + +class RelayObservabilityConfig(_RelayBaseConfig): + """Relay Observability plugin version 3 configuration.""" + + enable_full_payloads: bool = True + atof: RelayAtofConfig = Field(default_factory=RelayAtofConfig) + opentelemetry: RelayOpenTelemetryConfig = Field(default_factory=RelayOpenTelemetryConfig) + + +class RelayRedactionConfig(_RelayBaseConfig): + """Config-driven sanitization applied before subscribers and exporters.""" + + enabled: bool = True + request_privacy_attributes: list[Literal["data", "category_profile"]] = Field( + default_factory=lambda: ["data", "category_profile"] + ) + detectors: list[ + Literal[ + "email", + "phone", + "api_key", + "bearer_token", + "jwt", + "credit_card", + "aws_access_key_id", + "aws_secret_access_key", + "gcp_api_key", + "azure_storage_account_key", + ] + ] = Field( + default_factory=lambda: [ + "email", + "phone", + "api_key", + "bearer_token", + "jwt", + "credit_card", + "aws_access_key_id", + "aws_secret_access_key", + "gcp_api_key", + "azure_storage_account_key", + ] + ) + + +class RelayPricingConfig(_RelayBaseConfig): + """Model-pricing enrichment for managed LLM responses.""" + + enabled: bool = True + sources: list[dict[str, Any]] = Field(default_factory=list) + + +class RelayConfig(_RelayBaseConfig): + """AI-Q-owned Relay configuration; Relay itself is always instrumented.""" + + logging: bool = True + observability: RelayObservabilityConfig = Field(default_factory=RelayObservabilityConfig) + redaction: RelayRedactionConfig = Field(default_factory=RelayRedactionConfig) + pricing: RelayPricingConfig = Field(default_factory=RelayPricingConfig) + + def to_plugin_config(self) -> dict[str, Any]: + """Translate AI-Q YAML into Relay's supported plugin document shape.""" + observability = self.observability + atof = observability.atof + otel = observability.opentelemetry + components: list[dict[str, Any]] = [ + { + "kind": "observability", + "enabled": True, + "config": { + "version": 3, + "enable_full_payloads": observability.enable_full_payloads, + "atof": { + "enabled": atof.enabled, + "sinks": [ + { + "type": "file", + "output_directory": atof.output_directory, + "filename": atof.filename, + "mode": atof.mode, + } + ] + if atof.enabled + else [], + }, + "opentelemetry": { + "enabled": otel.enabled, + "endpoints": [ + endpoint.model_dump(mode="json", exclude_none=True) for endpoint in otel.endpoints + ] + if otel.enabled + else [], + }, + }, + } + ] + if self.redaction.enabled: + components.append( + { + "kind": "pii_redaction", + "enabled": True, + "config": { + "version": 1, + "profiles": [ + { + "mode": "builtin", + "priority": 80 + index, + "builtin": {"action": "redact", "detector": detector}, + } + for index, detector in enumerate(self.redaction.detectors) + ], + }, + } + ) + if self.pricing.enabled: + components.append( + { + "kind": "pricing", + "enabled": True, + "config": {"sources": self.pricing.sources}, + } + ) + return {"version": 1, "components": components} diff --git a/src/aiq_agent/relay/logging.py b/src/aiq_agent/relay/logging.py new file mode 100644 index 000000000..0dfeb1731 --- /dev/null +++ b/src/aiq_agent/relay/logging.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Developer-safe console logging for NeMo Relay lifecycle events.""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +import nemo_relay + +from aiq_agent.common.callbacks import BOLD +from aiq_agent.common.callbacks import CYAN +from aiq_agent.common.callbacks import DIM +from aiq_agent.common.callbacks import GREEN +from aiq_agent.common.callbacks import MAGENTA +from aiq_agent.common.callbacks import RED +from aiq_agent.common.callbacks import RESET +from aiq_agent.common.callbacks import RESET_ALL +from aiq_agent.common.callbacks import YELLOW +from aiq_agent.common.logging_utils import log_content_metadata + +logger = logging.getLogger(__name__) + +SUBSCRIBER_NAME = "aiq-relay-logging" +_depths: dict[str, int] = {} +_depth_lock = threading.RLock() + + +def _event_value(event: Any, name: str, default: Any = None) -> Any: + value = getattr(event, name, default) + return value if value is not None else default + + +def _nested(value: Any, *path: str | int) -> Any: + for part in path: + if isinstance(part, int) and isinstance(value, list) and len(value) > part: + value = value[part] + elif isinstance(part, str) and isinstance(value, dict): + value = value.get(part) + else: + return None + return value + + +def _event_depth(event: Any, phase: str) -> int: + uuid = str(_event_value(event, "uuid", "")) + parent_uuid = str(_event_value(event, "parent_uuid", "")) + with _depth_lock: + if phase == "start": + depth = _depths.get(parent_uuid, -1) + 1 + _depths[uuid] = depth + return depth + depth = _depths.pop(uuid, _depths.get(parent_uuid, -1) + 1) + return depth + + +def _log_agent(name: str, phase: str, indent: str, status: str | None) -> None: + if "agent" not in name.lower() and "subagent" not in name.lower(): + return + if phase == "start": + logger.info("%s%s[Chain Start] %s%s", indent, CYAN, name, RESET) + elif status == "ERROR": + logger.error("%s%s[Chain Error] %s%s", indent, RED, name, RESET) + else: + logger.info("%s%s[Chain End] %s%s", indent, CYAN, name, RESET) + + +def _log_llm(event: Any, name: str, phase: str, status: str | None) -> None: + data = _event_value(event, "data") + profile = _event_value(event, "category_profile", {}) or {} + annotated = profile.get("annotated_response", {}) if isinstance(profile, dict) else {} + if phase == "start": + logger.info("-" * 30) + logger.info("%s[AGENT]%s %s", BOLD, RESET_ALL, name) + if data is not None: + logger.info("%sAgent input: %s%s", YELLOW, log_content_metadata(data), RESET) + return + if status == "ERROR": + logger.error("%s[LLM Error] %s%s", RED, name, RESET) + return + + reasoning = _nested(data, "generations", 0, 0, "message", "additional_kwargs", "reasoning_content") + response = annotated.get("message") if isinstance(annotated, dict) else None + response = response or _nested(data, "generations", 0, 0, "message", "content") + tool_calls = annotated.get("tool_calls") if isinstance(annotated, dict) else None + tool_calls = ( + tool_calls + or _nested(data, "generations", 0, 0, "message", "tool_calls") + or _nested(data, "generations", 0, 0, "message", "additional_kwargs", "tool_calls") + ) + response_metadata = _nested(data, "generations", 0, 0, "message", "response_metadata") or {} + usage = annotated.get("usage", {}) if isinstance(annotated, dict) else {} + if not usage and isinstance(response_metadata, dict): + usage = response_metadata.get("token_usage", {}) or {} + model = annotated.get("model") if isinstance(annotated, dict) else None + model = model or (profile.get("model_name") if isinstance(profile, dict) else None) + if not model and isinstance(response_metadata, dict): + model = next( + (response_metadata[key] for key in ("model_name", "model", "model_id") if response_metadata.get(key)), + None, + ) + if reasoning: + logger.info("%s[Reasoning] %s%s", MAGENTA, log_content_metadata(reasoning), RESET_ALL) + if response: + logger.info("%s[Agent Response] %s%s", CYAN, log_content_metadata(response), RESET) + if isinstance(tool_calls, list) and tool_calls: + logger.info("%s[Tool Calls] %d tool(s) requested%s", GREEN, len(tool_calls), RESET) + for tool_call in tool_calls: + function = tool_call.get("function", {}) if isinstance(tool_call, dict) else {} + tool_name = tool_call.get("name") if isinstance(tool_call, dict) else None + tool_name = tool_name or (function.get("name") if isinstance(function, dict) else None) or "unknown" + tool_args = tool_call.get("args") if isinstance(tool_call, dict) else None + tool_args = tool_args if tool_args is not None else tool_call.get("arguments") + tool_args = tool_args if tool_args is not None else function.get("arguments", {}) + logger.info("%s → %s%s", GREEN, tool_name, RESET) + logger.info("%s Args: %s%s", DIM, log_content_metadata(tool_args), RESET_ALL) + if usage: + logger.info( + "%s[Tokens] prompt=%s, completion=%s, model=%s%s", + DIM, + usage.get("prompt_tokens", "N/A"), + usage.get("completion_tokens", "N/A"), + model or "unknown", + RESET_ALL, + ) + logger.info("-" * 30) + + +def _log_tool(event: Any, name: str, phase: str, status: str | None) -> None: + data = _event_value(event, "data") + if phase == "start": + logger.info("%s[Tool Start] %s%s", GREEN, name, RESET) + if data is not None: + logger.info("%s Input: %s%s", DIM, log_content_metadata(data), RESET_ALL) + elif status == "ERROR": + logger.error("%s[Tool Error] %s%s", RED, name, RESET) + else: + logger.info("%s[Tool Result] %s%s", GREEN, log_content_metadata(data), RESET) + + +def log_event(event: Any) -> None: + """Render sanitized Relay events in AI-Q's established callback format.""" + category = str(_event_value(event, "category", "")) + kind = str(_event_value(event, "kind", "")) + name = str(_event_value(event, "name", "unknown")) + phase = str(_event_value(event, "scope_category", "event")) + metadata = _event_value(event, "metadata", {}) or {} + status = metadata.get("otel.status_code") if isinstance(metadata, dict) else None + + if kind != "scope": + logger.debug("[Relay Event] %s category=%s", name, category) + return + + indent = " " * _event_depth(event, phase) + if category == "llm": + _log_llm(event, name, phase, status) + elif category == "tool": + _log_tool(event, name, phase, status) + elif category == "agent": + _log_agent(name, phase, indent, status) + + +def register_logging_subscriber() -> None: + """Register the process-global AI-Q Relay logger exactly once.""" + nemo_relay.subscribers.deregister(SUBSCRIBER_NAME) + with _depth_lock: + _depths.clear() + nemo_relay.subscribers.register(SUBSCRIBER_NAME, log_event) diff --git a/src/aiq_agent/relay/privacy.py b/src/aiq_agent/relay/privacy.py new file mode 100644 index 000000000..0bd70780f --- /dev/null +++ b/src/aiq_agent/relay/privacy.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request-scoped privacy controls for Relay observability payloads.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any + +import nemo_relay + +from .config import RelayRedactionConfig + +_PRIVACY_SANITIZER = "aiq-request-privacy" +_request_privacy_enabled: ContextVar[bool] = ContextVar("aiq_relay_request_privacy", default=False) + + +@contextmanager +def request_privacy_context(enabled: bool) -> Iterator[None]: + """Apply payload redaction to Relay events emitted by the current request.""" + token = _request_privacy_enabled.set(enabled) + try: + yield + finally: + _request_privacy_enabled.reset(token) + + +def request_privacy_from_tags(tags: dict[str, str], tag: str = "aiq.telemetry.redact") -> bool: + """Resolve the trusted request privacy decision carried with trace tags.""" + return str(tags.get(tag, "")).strip().lower() == "true" + + +def _redact_event_fields( + _event: Any, + fields: nemo_relay.EventSanitizeFields, + *, + attributes: tuple[str, ...], +) -> nemo_relay.EventSanitizeFields: + """Replace only configured fields that Relay supplied for this event.""" + if not _request_privacy_enabled.get(): + return fields + sanitized = dict(fields) + for attribute in attributes: + if attribute in sanitized: + sanitized[attribute] = None + metadata = sanitized.get("metadata") + if isinstance(metadata, dict): + sanitized["metadata"] = {**metadata, "aiq.telemetry.redacted": True} + return nemo_relay.EventSanitizeFields(**sanitized) + + +def _redact_llm_request(request: Any, _context: Any) -> Any: + return None if _request_privacy_enabled.get() else request + + +def _redact_llm_response(response: Any, _context: Any) -> Any: + return None if _request_privacy_enabled.get() else response + + +def _redact_tool_payload(_tool_name: str, payload: Any) -> Any: + return "[REDACTED]" if _request_privacy_enabled.get() else payload + + +def deregister_privacy_sanitizers() -> None: + """Remove AI-Q's process-global Relay privacy sanitizers.""" + guardrails = nemo_relay.guardrails + guardrails.deregister_scope_sanitize_start(_PRIVACY_SANITIZER) + guardrails.deregister_scope_sanitize_end(_PRIVACY_SANITIZER) + guardrails.deregister_mark_sanitize(_PRIVACY_SANITIZER) + guardrails.deregister_llm_sanitize_request(_PRIVACY_SANITIZER) + guardrails.deregister_llm_sanitize_response(_PRIVACY_SANITIZER) + guardrails.deregister_tool_sanitize_request(_PRIVACY_SANITIZER) + guardrails.deregister_tool_sanitize_response(_PRIVACY_SANITIZER) + + +def register_privacy_sanitizers(config: RelayRedactionConfig) -> None: + """Register observation-only sanitizers for request privacy decisions.""" + deregister_privacy_sanitizers() + if not config.enabled: + return + + attributes = tuple(config.request_privacy_attributes) + + def sanitize_event(event: Any, fields: nemo_relay.EventSanitizeFields) -> nemo_relay.EventSanitizeFields: + return _redact_event_fields(event, fields, attributes=attributes) + + guardrails = nemo_relay.guardrails + guardrails.register_scope_sanitize_start(_PRIVACY_SANITIZER, 10, sanitize_event) + guardrails.register_scope_sanitize_end(_PRIVACY_SANITIZER, 10, sanitize_event) + guardrails.register_mark_sanitize(_PRIVACY_SANITIZER, 10, sanitize_event) + guardrails.register_llm_sanitize_request(_PRIVACY_SANITIZER, 10, _redact_llm_request) + guardrails.register_llm_sanitize_response(_PRIVACY_SANITIZER, 10, _redact_llm_response) + guardrails.register_tool_sanitize_request(_PRIVACY_SANITIZER, 10, _redact_tool_payload) + guardrails.register_tool_sanitize_response(_PRIVACY_SANITIZER, 10, _redact_tool_payload) diff --git a/src/aiq_agent/relay/runtime.py b/src/aiq_agent/relay/runtime.py new file mode 100644 index 000000000..7f2dcd91e --- /dev/null +++ b/src/aiq_agent/relay/runtime.py @@ -0,0 +1,384 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NeMo Relay framework integration helpers.""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +from collections.abc import Awaitable +from collections.abc import Callable +from collections.abc import Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Any +from typing import TypeVar +from uuid import uuid4 + +import nemo_relay +from langchain.agents.middleware import AgentMiddleware +from langchain.agents.middleware import ModelRequest +from langchain.agents.middleware import ModelResponse +from langchain.agents.middleware import ToolCallRequest +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage +from langchain_core.messages import BaseMessage +from langchain_core.messages import messages_to_dict +from langchain_core.runnables import RunnableBinding +from langchain_core.runnables.config import merge_configs +from nemo_relay.integrations.langchain import NemoRelayMiddleware +from pydantic import BaseModel + +_T = TypeVar("_T") +_aiq_scope_active: ContextVar[bool] = ContextVar("aiq_relay_scope_active", default=False) +logger = logging.getLogger(__name__) + + +@dataclass +class _AgentScopeLifecycle: + handle: Any + output: Any = None + + +def _log_capture_failure(operation: str, error: Exception) -> None: + """Report Relay capture failures without exposing payloads or changing execution.""" + logger.warning("NeMo Relay: %s failed (error_type=%s)", operation, type(error).__name__) + + +@dataclass +class _NamedModelAdapter: + """Give otherwise valid chat-model implementations a Relay call name.""" + + wrapped: Any + model_name: str + + def __getattr__(self, name: str) -> Any: + return getattr(self.wrapped, name) + + +def _normalize_chat_nvidia_binding( + runnable: Any, + config: dict[str, Any], +) -> tuple[Any, dict[str, Any], dict[str, Any]]: + """Expose a bound ChatNVIDIA model to Relay without losing bound tools. + + Relay 0.7.3 handles propagation headers through ``ChatNVIDIA.default_headers`` + only when the model is a direct ChatNVIDIA instance. LangChain's + ``bind_tools()`` returns a RunnableBinding, which otherwise makes Relay fall + back to the unsupported ``extra_headers`` model parameter. + """ + try: + from langchain_nvidia_ai_endpoints import ChatNVIDIA + except ImportError: + return runnable, {}, config + + if ( + not isinstance(runnable, RunnableBinding) + or not isinstance(runnable.bound, ChatNVIDIA) + or runnable.config_factories + ): + return runnable, {}, config + + return runnable.bound, dict(runnable.kwargs), merge_configs(runnable.config, config) + + +# Work around NVIDIA/NeMo-Relay#805 until DeepAgents emits nested local-subagent Agent scopes. +class _DelegatedAgentScopeMiddleware(AgentMiddleware): + """Create a semantic Agent scope for the subagent selected by DeepAgents.""" + + @staticmethod + def _agent_name(request: Any) -> str | None: + tool_call = getattr(request, "tool_call", None) + if not isinstance(tool_call, dict) or tool_call.get("name") != "task": + return None + arguments = tool_call.get("args") + if not isinstance(arguments, dict): + return None + name = arguments.get("subagent_type") + return name if isinstance(name, str) and name else None + + def wrap_tool_call(self, request: Any, handler: Callable[[Any], Any]) -> Any: + name = self._agent_name(request) + if name is None: + return handler(request) + with agent_scope(name, input_value=getattr(request, "tool_call", None)) as lifecycle: + result = handler(request) + lifecycle.output = result + return result + + async def awrap_tool_call(self, request: Any, handler: Callable[[Any], Awaitable[Any]]) -> Any: + name = self._agent_name(request) + if name is None: + return await handler(request) + with agent_scope(name, input_value=getattr(request, "tool_call", None)) as lifecycle: + result = await handler(request) + lifecycle.output = result + return result + + +def deepagents_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + """Attach Relay's supported DeepAgents middleware.""" + + from nemo_relay.integrations.deepagents import add_nemo_relay_integration + + observed = add_nemo_relay_integration(kwargs) + middleware = list(observed.get("middleware") or ()) + if not any(isinstance(item, _DelegatedAgentScopeMiddleware) for item in middleware): + middleware.append(_DelegatedAgentScopeMiddleware()) + observed["middleware"] = middleware + return observed + + +def merge_langchain_middleware(middleware: Sequence[Any] | None) -> list[Any]: + """Attach Relay managed execution to an application-owned LangChain agent.""" + merged = list(middleware or ()) + if not any(isinstance(item, NemoRelayMiddleware) for item in merged): + merged.insert(0, NemoRelayMiddleware()) + return merged + + +async def ainvoke_with_relay( + runnable: Any, + input_value: Any, + *, + callbacks: Sequence[Any] | None = None, + config: dict[str, Any] | None = None, +) -> Any: + """Run a direct LangChain model call through Relay's maintained middleware.""" + effective_config = dict(config or {}) + configured_callbacks = callbacks if callbacks is not None else effective_config.get("callbacks") + if configured_callbacks: + effective_config["callbacks"] = list(configured_callbacks) + else: + effective_config.pop("callbacks", None) + messages = list(input_value) + system_message = ( + messages.pop(0) if messages and isinstance(messages[0], BaseMessage) and messages[0].type == "system" else None + ) + model, model_settings, effective_config = _normalize_chat_nvidia_binding(runnable, effective_config) + if not any( + isinstance(getattr(model, attribute, None), str) and getattr(model, attribute) + for attribute in ("model", "model_name", "model_id", "deployment_name") + ): + model = _NamedModelAdapter(model, type(model).__name__) + request = ModelRequest( + model=model, + messages=messages, + system_message=system_message, + model_settings=model_settings, + ) + + async def invoke(next_request: ModelRequest[Any]) -> ModelResponse[Any]: + next_messages = list(next_request.messages) + if next_request.system_message is not None: + next_messages.insert(0, next_request.system_message) + parameters = inspect.signature(next_request.model.ainvoke).parameters + accepts_config = "config" in parameters or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() + ) + kwargs = dict(next_request.model_settings) + if isinstance(next_request.model, _NamedModelAdapter) and not isinstance( + next_request.model.wrapped, + BaseChatModel, + ): + kwargs = {} + if not any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()): + kwargs = {name: value for name, value in kwargs.items() if name in parameters} + if accepts_config: + response = await next_request.model.ainvoke(next_messages, config=effective_config, **kwargs) + else: + response = await next_request.model.ainvoke(next_messages, **kwargs) + if not isinstance(response, BaseMessage): + content = getattr(response, "content", None) + if not isinstance(content, str | list): + message = f"Relay-managed LangChain model returned {type(response).__name__}, expected BaseMessage" + raise TypeError(message) + response = AIMessage(content=content) + return ModelResponse(result=[response]) + + response = await NemoRelayMiddleware().awrap_model_call(request, invoke) + if not response.result: + raise RuntimeError("Relay-managed LangChain model returned no messages") + return response.result[-1] + + +async def ainvoke_tool_with_relay(tool: Any, args: dict[str, Any]) -> Any: + """Run a direct LangChain tool call through Relay's maintained middleware.""" + request = ToolCallRequest( + tool_call={"name": tool.name, "args": args, "id": f"aiq-{uuid4()}"}, + tool=tool, + state={}, + runtime=None, + ) + + async def invoke(next_request: ToolCallRequest) -> Any: + if next_request.tool is None: + raise RuntimeError(f"Relay-managed tool {next_request.tool_call['name']!r} is unavailable") + return await next_request.tool.ainvoke(next_request.tool_call.get("args") or {}) + + return await NemoRelayMiddleware().awrap_tool_call(request, invoke) + + +@contextmanager +def _semantic_scope( + name: str, + scope_type: Any, + component_type: str, + *, + session_id: str | None = None, + input_value: Any = None, + metadata: dict[str, Any] | None = None, +): + """Create a semantic scope and mark nested AI-Q scope execution.""" + scope_token = None + if not _aiq_scope_active.get(): + scope_token = _aiq_scope_active.set(True) + scope_metadata = { + "aiq.component.name": name, + "aiq.component.type": component_type, + "aiq.framework": "nemo-agent-toolkit", + } + scope_metadata.update(metadata or {}) + if session_id: + scope_metadata["session_id"] = session_id + lifecycle = _AgentScopeLifecycle(None) + status_metadata: dict[str, Any] = {"otel.status_code": "UNSET"} + try: + try: + lifecycle.handle = nemo_relay.scope.push( + name, + scope_type, + metadata=_safe_value(scope_metadata), + input=_safe_value(input_value) if input_value is not None else None, + ) + except Exception as capture_error: + _log_capture_failure("semantic scope start", capture_error) + try: + yield lifecycle + except BaseException as error: + status_metadata = { + "error_type": type(error).__name__, + "otel.status_code": "ERROR", + "otel.status_description": str(error), + } + raise + else: + status_metadata = {"otel.status_code": "OK"} + finally: + try: + if lifecycle.handle is not None: + try: + output = _safe_value(lifecycle.output) if lifecycle.output is not None else None + nemo_relay.scope.pop(lifecycle.handle, output=output, metadata=status_metadata) + except Exception as capture_error: + _log_capture_failure("semantic scope end", capture_error) + finally: + if scope_token is not None: + _aiq_scope_active.reset(scope_token) + + +@contextmanager +def agent_scope(name: str, *, session_id: str | None = None, input_value: Any = None): + """Create an Agent scope around an application-owned agent boundary.""" + with _semantic_scope( + name, + nemo_relay.ScopeType.Agent, + "agent", + session_id=session_id, + input_value=input_value, + ) as lifecycle: + yield lifecycle + + +@contextmanager +def workflow_scope( + name: str, + *, + session_id: str | None = None, + input_value: Any = None, + metadata: dict[str, Any] | None = None, +): + """Create a NAT workflow scope above application-owned agent scopes.""" + with _semantic_scope( + name, + nemo_relay.ScopeType.Function, + "workflow", + session_id=session_id, + input_value=input_value, + metadata=metadata, + ) as lifecycle: + yield lifecycle + + +async def run_agent( + name: str, + operation: Callable[[], Awaitable[_T]], + *, + session_id: str | None = None, + input_value: Any = None, +) -> _T: + """Run an agent with a fresh stack at a request boundary and shared stack when nested.""" + + async def _run() -> _T: + with agent_scope(name, session_id=session_id, input_value=input_value) as lifecycle: + result = await operation() + lifecycle.output = result + return result + + if _aiq_scope_active.get(): + return await _run() + + async def _run_isolated() -> _T: + with nemo_relay.use_scope_stack(nemo_relay.create_scope_stack()): + return await _run() + + return await asyncio.create_task(_run_isolated()) + + +async def run_workflow( + name: str, + operation: Callable[[], Awaitable[_T]], + *, + session_id: str | None = None, + input_value: Any = None, + metadata: dict[str, Any] | None = None, +) -> _T: + """Run one NAT request as a Relay workflow root with semantic input and output.""" + + async def _run() -> _T: + with workflow_scope( + name, + session_id=session_id, + input_value=input_value, + metadata=metadata, + ) as lifecycle: + result = await operation() + lifecycle.output = result + return result + + if _aiq_scope_active.get(): + return await _run() + + async def _run_isolated() -> _T: + with nemo_relay.use_scope_stack(nemo_relay.create_scope_stack()): + return await _run() + + return await asyncio.create_task(_run_isolated()) + + +def _safe_value(value: Any) -> Any: + """Project framework state to JSON-compatible Relay event values.""" + if value is None or isinstance(value, str | int | float | bool): + return value + if isinstance(value, dict): + return {str(key): _safe_value(item) for key, item in value.items()} + if isinstance(value, list | tuple | set): + return [_safe_value(item) for item in value] + if isinstance(value, BaseMessage): + return messages_to_dict([value])[0] + if isinstance(value, BaseModel): + return value.model_dump(mode="json") + return {"type": type(value).__name__} diff --git a/src/aiq_agent/tokenomics/README.md b/src/aiq_agent/tokenomics/README.md index c0e95c34b..c5cfc18fe 100644 --- a/src/aiq_agent/tokenomics/README.md +++ b/src/aiq_agent/tokenomics/README.md @@ -1,16 +1,17 @@ # AIQ Tokenomics -Post-eval analysis module for the Deep Research Agent. Parses a NAT profiler trace, attributes costs and token counts to workflow phases (Orchestrator / Planner / Researcher), and renders a self-contained interactive HTML report. +Post-eval analysis module for the Deep Research Agent. Parses NeMo Relay ATOF JSONL, attributes costs and token counts to workflow phases (Orchestrator / Planner / Researcher), and renders a self-contained interactive HTML report. --- ## Background -### The subagent attribution problem +### Subagent attribution -The workflow is registered as `deep_research_agent`, and NAT still emits `FUNCTION_START` / `FUNCTION_END` for **tools** (e.g. search). Planner and Researcher subagents are inline LangGraph graphs inside the **`task`** tool: they do not appear as their own `FUNCTION_*` scopes, and traces from this stack usually have no per-step metadata (such as `function_ancestry`) that identifies subagent phase. - -This module uses **timing-window attribution**: every `task` TOOL_START/END pair brackets one subagent run and carries `subagent_type` in the tool input. Each `LLM_END` is classified using its **`event_timestamp`** (completion time): if it falls inside a task window, that phase applies; otherwise orchestrator. Overlapping researcher windows (parallel invocations) all yield `researcher-phase` — correct phase even when the specific instance is ambiguous. +The adapter follows Relay `parent_uuid` ancestry. LLM calls below real +`planner-agent` and `researcher-agent` scopes are attributed directly to those +phases; other calls are attributed to the orchestrator. No timing-window +inference or NAT callback renaming is required. --- @@ -20,7 +21,7 @@ This module uses **timing-window attribution**: every `task` TOOL_START/END pair src/aiq_agent/tokenomics/ ├── pricing.py # PricingRegistry — maps model names to per-token prices ├── profile.py # RequestProfile, PhaseStats — structured data classes -├── nat_adapter.py # parse_trace() — NAT JSON → list[RequestProfile] +├── atof_adapter.py # parse_trace() — Relay ATOF JSONL → list[RequestProfile] └── report.py # generate_report() — builds and renders HTML dashboard ``` @@ -37,37 +38,37 @@ Pricing lives in a YAML file under `tokenomics.pricing`. Prices are in **USD per tokenomics: pricing: models: - # Illustrative market-equivalent rates; verify current provider pricing. "nvidia/nemotron-3-ultra-550b-a55b": - input_per_1m_tokens: 0.60 - output_per_1m_tokens: 3.60 + # NVIDIA-hosted access for this example; not self-hosting cost. + input_per_1m_tokens: 0.00 + output_per_1m_tokens: 0.00 tools: - # Tool name lookup is substring-based: "web_search" matches "advanced_web_search_tool" - # and "tavily_search" because the key is a substring of those names. - "web_search": + # Tavily pay-as-you-go rates. Monthly plans have a lower effective rate. + "web_search_tool": + cost_per_call: 0.008 + "advanced_web_search_tool": cost_per_call: 0.016 + # Serper Starter; change for the selected tier/provider. "paper_search": - cost_per_call: 0.0003 - # Fallback for any model not explicitly listed. - # Set to null to raise an error on unknown models instead. - default: - input_per_1m_tokens: 1.00 - output_per_1m_tokens: 4.00 + cost_per_call: 0.001 ``` Model name lookup is: exact match → substring match → default. This means a key of `"nemotron-3-ultra"` will match a trace model name of `"nvidia/nemotron-3-ultra-550b-a55b"`. -Tool name lookup follows the same substring rule. Unknown tools default to $0/call — no error is raised, so you can configure only the costly tools and omit free internal ones. +Tool name lookup tries exact names before substring matching. Unknown tools +default to $0/call, so free internal tools can be omitted. Configure only the +provider-facing tool scope to avoid charging both a wrapper and its underlying +API call. --- ## Generating a report -Run after `nat eval` completes. The trace file is written to the `output_dir` configured in the eval config. +Run after a Relay-instrumented workflow completes. The default ATOF sink writes `relay/aiq-relay.atof.jsonl`. ```bash PYTHONPATH=src python -m aiq_agent.tokenomics.report \ - --trace frontends/benchmarks/deepresearch_bench/results/all_requests_profiler_traces.json \ + --trace relay/aiq-relay.atof.jsonl \ --config frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml \ [--output path/to/report.html] ``` @@ -137,7 +138,7 @@ with open("frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pri pricing = PricingRegistry.from_dict(config["tokenomics"]["pricing"]) # Parse trace → one RequestProfile per query -profiles = parse_trace("results/all_requests_profiler_traces.json", pricing) +profiles = parse_trace("relay/aiq-relay.atof.jsonl", pricing) for prof in profiles: print(f"Query {prof.request_index}: ${prof.total_cost_usd:.4f}, " diff --git a/src/aiq_agent/tokenomics/__init__.py b/src/aiq_agent/tokenomics/__init__.py index b1d7e86eb..eb87cd8b9 100644 --- a/src/aiq_agent/tokenomics/__init__.py +++ b/src/aiq_agent/tokenomics/__init__.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .nat_adapter import parse_trace +from .atof_adapter import parse_trace from .pricing import ModelPrice from .pricing import ModelPriceConfig from .pricing import PricingRegistry diff --git a/src/aiq_agent/tokenomics/atof_adapter.py b/src/aiq_agent/tokenomics/atof_adapter.py new file mode 100644 index 000000000..138df5e3c --- /dev/null +++ b/src/aiq_agent/tokenomics/atof_adapter.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Convert NeMo Relay ATOF JSONL events into tokenomics request profiles.""" + +from __future__ import annotations + +import json +import logging +from datetime import datetime +from pathlib import Path +from typing import Any + +from .pricing import PricingRegistry +from .profile import PHASE_ORCHESTRATOR +from .profile import PHASE_PLANNER +from .profile import PHASE_RESEARCHER +from .profile import PhaseStats +from .profile import RequestProfile + +logger = logging.getLogger(__name__) + + +def _timestamp(value: Any) -> float: + if isinstance(value, int | float): + return float(value) + if not isinstance(value, str): + return 0.0 + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + except ValueError: + return 0.0 + + +def _nested(value: Any, *path: str) -> Any: + for key in path: + if not isinstance(value, dict): + return None + value = value.get(key) + return value + + +def _integer(*values: Any) -> int: + for value in values: + if isinstance(value, int | float): + return int(value) + return 0 + + +def _number(*values: Any) -> float | None: + for value in values: + if isinstance(value, int | float): + return float(value) + return None + + +def _question(data: Any) -> str: + if isinstance(data, str): + return data + if isinstance(data, dict): + for key in ("input", "query", "question"): + if key in data: + return _question(data[key]) + messages = data.get("messages") + if isinstance(messages, list): + for message in reversed(messages): + if isinstance(message, dict): + content = message.get("content") or _nested(message, "data", "content") + if isinstance(content, str): + return content + return "" if data is None else json.dumps(data, ensure_ascii=False, default=str) + + +def _load_events(path: str) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + with Path(path).open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + if not line.strip(): + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + logger.warning("Skipping invalid ATOF JSON on line %d", line_number) + continue + if isinstance(event, dict): + events.append(event) + return events + + +def _phase_for(event: dict[str, Any], starts: dict[str, dict[str, Any]]) -> str: + parent_uuid = event.get("parent_uuid") + visited: set[str] = set() + while isinstance(parent_uuid, str) and parent_uuid not in visited: + visited.add(parent_uuid) + parent = starts.get(parent_uuid) + if parent is None: + break + name = str(parent.get("name") or "").lower() + if name == PHASE_PLANNER: + return PHASE_PLANNER + if name == "researcher-agent" or (parent.get("category") == "agent" and "researcher" in name): + return PHASE_RESEARCHER + parent_uuid = parent.get("parent_uuid") + return PHASE_ORCHESTRATOR + + +def _root_uuid(event: dict[str, Any], starts: dict[str, dict[str, Any]]) -> str | None: + event_uuid = event.get("uuid") + current = event_uuid if event_uuid in starts else event.get("parent_uuid") + if not isinstance(current, str): + return None + visited: set[str] = set() + while current not in visited: + visited.add(current) + parent = starts.get(current, {}).get("parent_uuid") + if not isinstance(parent, str) or parent not in starts: + return current + current = parent + return None + + +def _usage(event: dict[str, Any]) -> tuple[int, int, int, int, float | None]: + profile = event.get("category_profile") or {} + annotated = profile.get("annotated_response") if isinstance(profile, dict) else {} + annotated = annotated if isinstance(annotated, dict) else {} + usage = annotated.get("usage") or profile.get("usage") or {} + usage = usage if isinstance(usage, dict) else {} + input_details = usage.get("input_tokens_details") or {} + output_details = usage.get("output_tokens_details") or {} + cost = usage.get("cost") or annotated.get("cost") or profile.get("cost") + cost_total = ( + _number(cost.get("total"), cost.get("total_cost"), cost.get("usd")) if isinstance(cost, dict) else _number(cost) + ) + return ( + _integer(usage.get("prompt_tokens"), usage.get("input_tokens")), + _integer(usage.get("cached_tokens"), _nested(input_details, "cached_tokens")), + _integer(usage.get("completion_tokens"), usage.get("output_tokens")), + _integer(usage.get("reasoning_tokens"), _nested(output_details, "reasoning_tokens")), + cost_total, + ) + + +def _model(event: dict[str, Any]) -> str: + profile = event.get("category_profile") or {} + annotated = profile.get("annotated_response") if isinstance(profile, dict) else {} + for value in ( + annotated.get("model") if isinstance(annotated, dict) else None, + annotated.get("model_name") if isinstance(annotated, dict) else None, + profile.get("model_name") if isinstance(profile, dict) else None, + event.get("name"), + ): + if isinstance(value, str) and value: + return value + return "unknown" + + +def _parse_request( + request_index: int, + root: dict[str, Any], + events: list[dict[str, Any]], + starts: dict[str, dict[str, Any]], + pricing: PricingRegistry, +) -> RequestProfile: + ends = { + str(event.get("uuid")): event + for event in events + if event.get("kind") == "scope" and event.get("scope_category") == "end" + } + root_end = ends.get(str(root.get("uuid")), {}) + duration_s = max(0.0, _timestamp(root_end.get("timestamp")) - _timestamp(root.get("timestamp"))) + phase_model_stats: dict[tuple[str, str], PhaseStats] = {} + model_call_counters: dict[str, int] = {} + llm_call_events: list[dict[str, Any]] = [] + tool_call_events: list[dict[str, Any]] = [] + tool_calls: dict[str, int] = {} + + for start in events: + if start.get("kind") != "scope" or start.get("scope_category") != "start": + continue + end = ends.get(str(start.get("uuid"))) + if end is None: + continue + category = start.get("category") + dur_s = max(0.0, _timestamp(end.get("timestamp")) - _timestamp(start.get("timestamp"))) + if category == "tool": + name = str(start.get("name") or "unknown") + tool_calls[name] = tool_calls.get(name, 0) + 1 + tool_call_events.append( + {"tool": name, "dur_s": round(dur_s, 3), "cost_usd": pricing.get_tool(name).cost_per_call} + ) + continue + if category != "llm": + continue + + model = _model(end) + prompt_tokens, cached_tokens, completion_tokens, reasoning_tokens, relay_cost = _usage(end) + phase = _phase_for(start, starts) + key = (phase, model) + stats = phase_model_stats.setdefault(key, PhaseStats(phase=phase, model=model)) + try: + price = pricing.get(model) + calculated_cost = price.cost(prompt_tokens, cached_tokens, completion_tokens) + savings = price.cache_savings(cached_tokens) + except KeyError: + calculated_cost = savings = 0.0 + if relay_cost is None: + logger.warning("No price for model %r and no Relay cost; cost will be 0", model) + cost = relay_cost if relay_cost is not None else calculated_cost + stats.llm_calls += 1 + stats.prompt_tokens += prompt_tokens + stats.cached_tokens += cached_tokens + stats.completion_tokens += completion_tokens + stats.cost_usd += cost + stats.cache_savings_usd += savings + call_index = model_call_counters.get(model, 0) + model_call_counters[model] = call_index + 1 + llm_call_events.append( + { + "uuid": start.get("uuid"), + "isl": prompt_tokens, + "osl": completion_tokens, + "cached": cached_tokens, + "reasoning": reasoning_tokens, + "dur_s": round(dur_s, 3), + "tps": round(completion_tokens / dur_s, 2) if dur_s else 0.0, + "model": model, + "phase": phase, + "call_idx": call_index, + } + ) + + phases = list(phase_model_stats.values()) + return RequestProfile( + request_index=request_index, + question=_question(root.get("data")), + duration_s=duration_s, + phases=phases, + tool_calls=tool_calls, + llm_call_events=llm_call_events, + tool_call_events=tool_call_events, + total_llm_calls=sum(phase.llm_calls for phase in phases), + total_prompt_tokens=sum(phase.prompt_tokens for phase in phases), + total_cached_tokens=sum(phase.cached_tokens for phase in phases), + total_completion_tokens=sum(phase.completion_tokens for phase in phases), + total_cost_usd=sum(phase.cost_usd for phase in phases), + total_tool_cost_usd=sum(event["cost_usd"] for event in tool_call_events), + total_cache_savings_usd=sum(phase.cache_savings_usd for phase in phases), + ) + + +def parse_trace(path: str, pricing: PricingRegistry) -> list[RequestProfile]: + """Parse Relay ATOF JSONL into one profile per workflow root scope.""" + events = _load_events(path) + starts = { + str(event.get("uuid")): event + for event in events + if event.get("kind") == "scope" and event.get("scope_category") == "start" and event.get("uuid") + } + explicit_roots = [ + event for event in starts.values() if _nested(event, "metadata", "aiq.component.type") == "workflow" + ] + roots = explicit_roots or [event for event in starts.values() if event.get("parent_uuid") not in starts] + roots.sort(key=lambda event: _timestamp(event.get("timestamp"))) + + profiles: list[RequestProfile] = [] + for root in roots: + root_uuid = str(root["uuid"]) + request_events = [event for event in events if _root_uuid(event, starts) == root_uuid] + try: + profiles.append(_parse_request(len(profiles), root, request_events, starts, pricing)) + except Exception: + logger.exception("Failed to parse Relay request root %s; skipping", root_uuid) + return profiles diff --git a/src/aiq_agent/tokenomics/nat_adapter.py b/src/aiq_agent/tokenomics/nat_adapter.py deleted file mode 100644 index 26a2a375b..000000000 --- a/src/aiq_agent/tokenomics/nat_adapter.py +++ /dev/null @@ -1,319 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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. - -""" -NAT trace → list[RequestProfile] -================================= - -Converts a NAT profiler trace JSON file (produced by ``nat eval``) into -structured :class:`~aiq_agent.tokenomics.profile.RequestProfile` objects -ready for the tokenomics HTML report. - -Architecture note ------------------ -The workflow is registered as ``deep_research_agent``. NAT 1.5.0 traces still -emit ``FUNCTION_START`` / ``FUNCTION_END`` for **tools** (e.g. search helpers), -but **planner-agent** and **researcher-agent** runs live inside the ``task`` -tool: they do not get distinct ``FUNCTION_*`` names. Traces from this stack -typically have no per-step ``function_ancestry`` (or equivalent) carrying -subagent identity — calling ``subagent.ainvoke()`` does not surface as separate -NAT function scopes for Planner vs Researcher. - -Subagent attribution is therefore inferred post-hoc via timing windows: every -``task`` TOOL_START/END pair brackets one subagent invocation and carries -``subagent_type`` in its input. For each ``LLM_END`` we use that step's -``event_timestamp`` (completion time, not ``span_event_timestamp``): if it -lies inside a task window, the call is attributed to that phase; otherwise -**orchestrator-phase**. - -``_build_task_windows`` appends windows in ``task`` TOOL_END order. -``_infer_phase`` returns the **first** window in that list whose bounds contain -``ts``. Overlapping researcher windows share the same phase label, so order is -unimportant in the common parallel-researcher case. - -If NAT later attaches subagent phase directly on each step (e.g. -``function_ancestry`` or explicit ``FUNCTION_*`` scopes for subagents), -``_infer_phase`` can be replaced with a field read and the rest of this module -can stay the same. -""" - -from __future__ import annotations - -import ast -import json -import logging -from dataclasses import dataclass -from dataclasses import field -from typing import Any - -from .pricing import PricingRegistry -from .profile import PHASE_ORCHESTRATOR -from .profile import PHASE_PLANNER -from .profile import PHASE_RESEARCHER -from .profile import PhaseStats -from .profile import RequestProfile - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Internal helpers -# --------------------------------------------------------------------------- - - -@dataclass -class _TaskWindow: - """Time span of a single subagent (task tool) invocation.""" - - uuid: str - subagent_type: str # "planner-agent" | "researcher-agent" - start_ts: float - end_ts: float = field(default=0.0) - - @property - def phase(self) -> str: - if self.subagent_type == "planner-agent": - return PHASE_PLANNER - return PHASE_RESEARCHER # any other subagent_type → researcher-phase - - -def _extract_subagent_type(raw_input: Any) -> str | None: - """Pull subagent_type out of the task tool's input field.""" - if isinstance(raw_input, dict): - return raw_input.get("subagent_type") - if isinstance(raw_input, str): - # NAT stores tool inputs as Python-repr strings, not JSON - try: - parsed = ast.literal_eval(raw_input) - if isinstance(parsed, dict): - return parsed.get("subagent_type") - except Exception: - pass - # Last resort: substring scan (handles malformed reprs) - for candidate in ("planner-agent", "researcher-agent"): - if candidate in raw_input: - return candidate - return None - - -def _build_task_windows(steps: list[dict]) -> list[_TaskWindow]: - """Build a list of completed task-tool windows from a request's steps.""" - open_windows: dict[str, _TaskWindow] = {} - closed: list[_TaskWindow] = [] - - for step in steps: - payload = step["payload"] - event_type = payload["event_type"] - name = payload.get("name", "") - uuid = payload["UUID"] - ts = payload["event_timestamp"] - - if event_type == "TOOL_START" and name == "task": - raw_input = (payload.get("data") or {}).get("input") - subagent_type = _extract_subagent_type(raw_input) - if subagent_type: - open_windows[uuid] = _TaskWindow(uuid=uuid, subagent_type=subagent_type, start_ts=ts) - else: - logger.debug("task TOOL_START missing subagent_type, uuid=%s", uuid) - - elif event_type == "TOOL_END" and name == "task": - win = open_windows.pop(uuid, None) - if win is not None: - win.end_ts = ts - closed.append(win) - - if open_windows: - logger.warning("%d task windows never closed (truncated trace?)", len(open_windows)) - - return closed - - -def _infer_phase(ts: float, windows: list[_TaskWindow]) -> str: - """ - Return the phase label for an LLM call from its ``LLM_END`` time ``ts``. - - ``windows`` is ordered by ``task`` TOOL_END (see ``_build_task_windows``). - The first window with ``start_ts <= ts <= end_ts`` wins. Overlapping - researcher windows all map to ``researcher-phase`` anyway. - """ - for win in windows: - if win.start_ts <= ts <= win.end_ts: - return win.phase - return PHASE_ORCHESTRATOR - - -def _parse_request(request_index: int, steps: list[dict], pricing: PricingRegistry) -> RequestProfile: - """Convert one request's step list into a RequestProfile.""" - - # --- Workflow timing and question --- - wf_start_ts = wf_end_ts = 0.0 - question = "" - for step in steps: - payload = step["payload"] - et = payload["event_type"] - if et == "WORKFLOW_START": - wf_start_ts = payload["event_timestamp"] - question = (payload.get("data") or {}).get("input") or "" - elif et == "WORKFLOW_END": - wf_end_ts = payload["event_timestamp"] - - duration_s = max(0.0, wf_end_ts - wf_start_ts) - - # --- Subagent phase windows --- - task_windows = _build_task_windows(steps) - - # --- Single forward pass: accumulate all events --- - phase_model_stats: dict[tuple[str, str], PhaseStats] = {} - model_call_counters: dict[str, int] = {} - llm_call_events: list[dict] = [] - tool_call_events: list[dict] = [] - tool_calls: dict[str, int] = {} - tool_start_times: dict[str, tuple[str, float]] = {} # uuid -> (name, start_ts) - - for step in steps: - payload = step["payload"] - et = payload["event_type"] - uuid = payload["UUID"] - ts = payload["event_timestamp"] - - if et == "TOOL_START": - name = payload.get("name") or "unknown" - tool_start_times[uuid] = (name, ts) - - elif et == "TOOL_END": - name = payload.get("name") or "unknown" - tool_calls[name] = tool_calls.get(name, 0) + 1 - dur_s = 0.0 - if uuid in tool_start_times: - _, start_ts = tool_start_times.pop(uuid) - dur_s = max(0.0, ts - start_ts) - tool_price = pricing.get_tool(name) - tool_call_events.append( - { - "tool": name, - "dur_s": round(dur_s, 3), - "cost_usd": tool_price.cost_per_call, - } - ) - - elif et == "LLM_END": - # span_event_timestamp is set by LangchainProfilerHandler at LLM_START - span_ts = payload.get("span_event_timestamp", ts) - model = payload.get("name") or "unknown" - usage = (payload.get("usage_info") or {}).get("token_usage") or {} - - prompt_tokens = usage.get("prompt_tokens", 0) - cached_tokens = usage.get("cached_tokens", 0) - completion_tokens = usage.get("completion_tokens", 0) - reasoning_tokens = usage.get("reasoning_tokens", 0) - - dur_s = max(0.0, ts - span_ts) - tps = completion_tokens / dur_s if dur_s > 0 else 0.0 - - # Window match uses LLM_END event_timestamp (completion), not span_event_timestamp. - phase = _infer_phase(ts, task_windows) - key = (phase, model) - - if key not in phase_model_stats: - phase_model_stats[key] = PhaseStats(phase=phase, model=model) - - try: - price = pricing.get(model) - cost = price.cost(prompt_tokens, cached_tokens, completion_tokens) - savings = price.cache_savings(cached_tokens) - except KeyError: - logger.warning("No price for model %r — cost will be 0", model) - cost = savings = 0.0 - - ps = phase_model_stats[key] - ps.llm_calls += 1 - ps.prompt_tokens += prompt_tokens - ps.cached_tokens += cached_tokens - ps.completion_tokens += completion_tokens - ps.cost_usd += cost - ps.cache_savings_usd += savings - - # Per-call observation (for distribution charts) - call_idx = model_call_counters.get(model, 0) - model_call_counters[model] = call_idx + 1 - - llm_call_events.append( - { - "uuid": uuid, - "isl": prompt_tokens, - "osl": completion_tokens, - "cached": cached_tokens, - "reasoning": reasoning_tokens, - "dur_s": round(dur_s, 3), - "tps": round(tps, 2), - "model": model, - "phase": phase, - "call_idx": call_idx, - } - ) - - # --- Roll up to request-level totals --- - phases = list(phase_model_stats.values()) - total_tool_cost_usd = sum(ev["cost_usd"] for ev in tool_call_events) - return RequestProfile( - request_index=request_index, - question=question, - duration_s=duration_s, - phases=phases, - tool_calls=tool_calls, - llm_call_events=llm_call_events, - tool_call_events=tool_call_events, - total_llm_calls=sum(p.llm_calls for p in phases), - total_prompt_tokens=sum(p.prompt_tokens for p in phases), - total_cached_tokens=sum(p.cached_tokens for p in phases), - total_completion_tokens=sum(p.completion_tokens for p in phases), - total_cost_usd=sum(p.cost_usd for p in phases), - total_tool_cost_usd=total_tool_cost_usd, - total_cache_savings_usd=sum(p.cache_savings_usd for p in phases), - ) - - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - - -def parse_trace(path: str, pricing: PricingRegistry) -> list[RequestProfile]: - """ - Parse a NAT profiler trace JSON file and return one - :class:`~aiq_agent.tokenomics.profile.RequestProfile` per request. - - Parameters - ---------- - path: - Path to the ``all_requests_profiler_traces.json`` file produced by - ``nat eval``. - pricing: - A :class:`~aiq_agent.tokenomics.pricing.PricingRegistry` built from - the ``tokenomics.pricing`` section of the eval config YAML. - """ - with open(path) as f: - data = json.load(f) - - profiles = [] - for item in data: - idx = item.get("request_number", len(profiles)) - steps = item.get("intermediate_steps", []) - try: - profiles.append(_parse_request(idx, steps, pricing)) - except Exception: - logger.exception("Failed to parse request %d — skipping", idx) - - return profiles diff --git a/src/aiq_agent/tokenomics/pricing.py b/src/aiq_agent/tokenomics/pricing.py index 109bf64c3..30df41e59 100644 --- a/src/aiq_agent/tokenomics/pricing.py +++ b/src/aiq_agent/tokenomics/pricing.py @@ -46,7 +46,7 @@ class PricingRegistryConfig(BaseModel): """ Pricing table read from the ``tokenomics.pricing`` section of the eval config YAML. ``models`` is keyed by the exact model name that appears in - NAT traces (e.g. ``"azure/openai/gpt-5.2"``). ``default`` is used as a + Relay traces (e.g. ``"nvidia/nemotron-3-ultra-550b-a55b"``). ``default`` is used as a fallback when no model key matches. ``tools`` is keyed by tool name as it appears in the trace. """ @@ -98,7 +98,7 @@ class PricingRegistry: Model lookup order: 1. Exact match on ``model_name``. 2. Substring match — useful for versioned or provider-prefixed names - (e.g. ``"azure/openai/gpt-5.2"`` matches key ``"gpt-5.2"``). + (e.g. ``"nvidia/nemotron-3-ultra-550b-a55b"`` matches key ``"nemotron-3-ultra"``). 3. ``default`` price, if configured. 4. :class:`KeyError`. diff --git a/src/aiq_agent/tokenomics/profile.py b/src/aiq_agent/tokenomics/profile.py index e5587e013..6442cdd23 100644 --- a/src/aiq_agent/tokenomics/profile.py +++ b/src/aiq_agent/tokenomics/profile.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from dataclasses import field -# Canonical phase names produced by nat_adapter. +# Canonical phase names produced by atof_adapter. PHASE_ORCHESTRATOR = "orchestrator" PHASE_PLANNER = "planner-agent" PHASE_RESEARCHER = "researcher-phase" @@ -77,7 +77,7 @@ class RequestProfile: total_cache_savings_usd: float = 0.0 total_llm_calls: int = 0 - # One entry per (phase, model) pair — populated by nat_adapter + # One entry per (phase, model) pair — populated by atof_adapter phases: list[PhaseStats] = field(default_factory=list) # tool_name → invocation count diff --git a/src/aiq_agent/tokenomics/report/__init__.py b/src/aiq_agent/tokenomics/report/__init__.py index 7f7df012a..45ed5f31f 100644 --- a/src/aiq_agent/tokenomics/report/__init__.py +++ b/src/aiq_agent/tokenomics/report/__init__.py @@ -14,20 +14,20 @@ # limitations under the License. """ -Generate a self-contained tokenomics HTML report from a NAT profiler trace. +Generate a self-contained tokenomics HTML report from Relay ATOF JSONL. Single-run ---------- python -m aiq_agent.tokenomics.report \\ - --trace results/all_requests_profiler_traces.json \\ + --trace relay/aiq-relay.atof.jsonl \\ --config configs/config_tokenomics_pricing.yml \\ [--output results/tokenomics_report.html] Comparison (two or more runs) ------------------------------ python -m aiq_agent.tokenomics.report \\ - --trace results/run_a/all_requests_profiler_traces.json \\ - --trace results/run_b/all_requests_profiler_traces.json \\ + --trace results/run_a/aiq-relay.atof.jsonl \\ + --trace results/run_b/aiq-relay.atof.jsonl \\ --config configs/config_tokenomics_pricing.yml Passing ``--trace`` more than once activates comparison mode: every tab @@ -42,7 +42,7 @@ import yaml -from ..nat_adapter import parse_trace +from ..atof_adapter import parse_trace from ..pricing import PricingRegistry from ._report_builders import _build_comparison_data from ._report_builders import _build_report_data @@ -61,7 +61,7 @@ def generate_report( Parameters ---------- trace_path: - Path to a single ``all_requests_profiler_traces.json``, or a list of + Path to a single Relay ATOF JSONL file, or a list of paths for comparison mode. When more than one path is provided every tab (Overview, Cost, Latency, Tokens, Efficiency, Per-Query) shows A-vs-B comparison charts instead of single-run visualisations. diff --git a/src/aiq_agent/tokenomics/report/__main__.py b/src/aiq_agent/tokenomics/report/__main__.py index f21df5dcc..cf99bb143 100644 --- a/src/aiq_agent/tokenomics/report/__main__.py +++ b/src/aiq_agent/tokenomics/report/__main__.py @@ -22,14 +22,14 @@ from . import generate_report if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Generate a tokenomics HTML report from a NAT profiler trace.") + parser = argparse.ArgumentParser(description="Generate a tokenomics HTML report from Relay ATOF JSONL.") parser.add_argument( "--trace", required=True, action="append", metavar="TRACE", help=( - "Path to all_requests_profiler_traces.json. " + "Path to a Relay ATOF JSONL file. " "Repeat the flag to compare multiple runs (e.g. --trace run_a/traces.json --trace run_b/traces.json)." ), ) diff --git a/tests/aiq_agent/agents/chat_researcher/nodes/test_intent_classifier.py b/tests/aiq_agent/agents/chat_researcher/nodes/test_intent_classifier.py index 772ab1eda..3beca2fdb 100644 --- a/tests/aiq_agent/agents/chat_researcher/nodes/test_intent_classifier.py +++ b/tests/aiq_agent/agents/chat_researcher/nodes/test_intent_classifier.py @@ -401,7 +401,8 @@ async def test_run_with_callbacks(self, mock_llm): # ainvoke(rendered_prompt, config=config) assert call_args[0][0] # first positional arg is the prompt string config = call_args[1].get("config", {}) - assert config.get("callbacks") == [mock_callback] + callbacks = config.get("callbacks") + assert callbacks == [mock_callback] @pytest.mark.asyncio async def test_run_does_not_pass_prior_report_content_to_classifier_llm(self, mock_llm): diff --git a/tests/aiq_agent/agents/clarifier/test_agent.py b/tests/aiq_agent/agents/clarifier/test_agent.py index 5260d3691..2c3552fa7 100644 --- a/tests/aiq_agent/agents/clarifier/test_agent.py +++ b/tests/aiq_agent/agents/clarifier/test_agent.py @@ -27,6 +27,7 @@ from langchain_core.messages import SystemMessage from langchain_core.messages import ToolMessage from langchain_core.tools import tool +from nemo_relay.integrations.langchain import NemoRelayMiddleware from aiq_agent.agents.clarifier.agent import DEFAULT_CLARIFICATION_PROMPT from aiq_agent.agents.clarifier.agent import FORCE_SEARCH_GUIDANCE @@ -88,7 +89,6 @@ def test_init_with_defaults(self, mock_llm_provider, mock_user_callback): assert agent.user_prompt_callback == mock_user_callback assert agent.max_turns == 3 assert agent.log_response_max_chars == 2000 - assert agent.verbose is False assert agent.callbacks == [] assert agent.system_prompt is not None @@ -123,16 +123,6 @@ def test_init_with_callbacks(self, mock_llm_provider, mock_user_callback): assert agent.callbacks == [mock_callback] - def test_init_with_verbose(self, mock_llm_provider, mock_user_callback): - """Test initialization with verbose mode.""" - agent = ClarifierAgent( - llm_provider=mock_llm_provider, - user_prompt_callback=mock_user_callback, - verbose=True, - ) - - assert agent.verbose is True - def test_graph_property(self, mock_llm_provider, mock_user_callback): """Test graph property returns compiled graph.""" agent = ClarifierAgent( @@ -693,6 +683,40 @@ async def test_force_search_not_triggered_when_llm_searches_first(self, mock_llm for call in mock_llm.ainvoke.call_args_list: assert not any(FORCE_SEARCH_GUIDANCE in str(m.content) for m in call.args[0]) + @pytest.mark.asyncio + async def test_tool_node_uses_relay_middleware(self, mock_llm_provider, mock_llm, monkeypatch): + """Application-owned ToolNode executions pass through Relay's maintained hook.""" + observed_tool_names: list[str] = [] + + async def observe_tool_call(self, request, handler): + observed_tool_names.append(request.tool_call["name"]) + return await handler(request) + + monkeypatch.setattr(NemoRelayMiddleware, "awrap_tool_call", observe_tool_call) + mock_llm.ainvoke = AsyncMock( + side_effect=[ + AIMessage( + content="", + tool_calls=[{"name": "web_search_tool", "args": {"query": "AI"}, "id": "call_1"}], + ), + AIMessage( + content=ClarificationResponse( + needs_clarification=False, + clarification_question=None, + ).model_dump_json() + ), + ] + ) + agent = ClarifierAgent( + llm_provider=mock_llm_provider, + tools=[web_search_tool], + user_prompt_callback=AsyncMock(), + ) + + await agent.run(ClarifierAgentState(messages=[HumanMessage(content="Research AI")])) + + assert observed_tool_names == ["web_search_tool"] + @pytest.mark.asyncio async def test_force_search_guidance_not_in_state_messages(self, mock_llm_provider, mock_llm): """The force_search guidance must be injected ephemerally only; it must diff --git a/tests/aiq_agent/agents/clarifier/test_register.py b/tests/aiq_agent/agents/clarifier/test_register.py index e61eaca31..ebed92672 100644 --- a/tests/aiq_agent/agents/clarifier/test_register.py +++ b/tests/aiq_agent/agents/clarifier/test_register.py @@ -34,7 +34,6 @@ def test_config_with_required_fields(self): assert config.tools == [] assert config.max_turns == 3 assert config.log_response_max_chars == 2000 - assert config.verbose is False def test_config_with_all_fields(self): """Test config with all fields specified.""" @@ -43,14 +42,12 @@ def test_config_with_all_fields(self): tools=["tool1", "tool2"], max_turns=5, log_response_max_chars=1000, - verbose=True, ) assert config.llm == "test_llm" assert config.tools == ["tool1", "tool2"] assert config.max_turns == 5 assert config.log_response_max_chars == 1000 - assert config.verbose is True def test_config_tools_default_factory(self): """Test tools default to empty list.""" diff --git a/tests/aiq_agent/agents/deep_researcher/test_agent.py b/tests/aiq_agent/agents/deep_researcher/test_agent.py index 21051e6e7..2e949441a 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_agent.py +++ b/tests/aiq_agent/agents/deep_researcher/test_agent.py @@ -208,7 +208,6 @@ def test_init_with_defaults(self, mock_llm_provider, real_tool, mock_create_deep assert agent.llm_provider == mock_llm_provider assert len(agent.tools) == 1 - assert agent.verbose is True assert agent.callbacks == [] assert agent.deepagents_runtime.skill_sources_for("orchestrator") is None assert agent.enable_source_router is True @@ -232,7 +231,6 @@ def test_init_with_custom_settings(self, mock_llm_provider, real_tool, mock_crea agent = DeepResearcherAgent( llm_provider=mock_llm_provider, tools=[real_tool], - verbose=False, callbacks=callbacks, enable_citation_verification=False, skills=DeepResearchSkillsConfig(agents={"researcher-agent": ("research",)}), @@ -245,7 +243,6 @@ def test_init_with_custom_settings(self, mock_llm_provider, real_tool, mock_crea max_source_tool_batch_size=4, ) - assert agent.verbose is False assert agent.callbacks == callbacks assert agent.max_research_concurrency == 2 assert agent.max_researcher_model_calls == 12 @@ -373,7 +370,6 @@ async def test_registered_run_cancellation_finalizes_only_request_owned_agent(se config = DeepResearchAgentConfig( orchestrator_llm="llm", tools=["web_search_tool"], - verbose=False, sandbox=DeepResearchSandboxConfig() if owns_active_agent else None, ) state = DeepResearchAgentState(messages=[HumanMessage(content="cancel this request")]) @@ -423,7 +419,6 @@ async def test_registered_run_only_forces_cleanup_for_job_resource_timeout(self, config = DeepResearchAgentConfig( orchestrator_llm="llm", tools=["web_search_tool"], - verbose=False, sandbox=DeepResearchSandboxConfig(), ) state = DeepResearchAgentState(messages=[HumanMessage(content="bounded request")]) @@ -463,7 +458,7 @@ async def test_registered_run_raises_typed_source_configuration_failure( builder = MagicMock() builder.get_tools = AsyncMock(return_value=[web_search_tool]) builder.get_llm = AsyncMock(return_value=MagicMock()) - config = DeepResearchAgentConfig(orchestrator_llm="llm", tools=["web_search_tool"], verbose=False) + config = DeepResearchAgentConfig(orchestrator_llm="llm", tools=["web_search_tool"]) state = DeepResearchAgentState(messages=[HumanMessage(content="research this")], data_sources=data_sources) original_description = web_search_tool.description web_search_tool.description = tool_description @@ -490,7 +485,7 @@ async def test_explicit_empty_selection_skips_request_agent_setup(self): builder = MagicMock() builder.get_tools = AsyncMock(return_value=[web_search_tool]) builder.get_llm = AsyncMock(return_value=MagicMock()) - config = DeepResearchAgentConfig(orchestrator_llm="llm", tools=["web_search_tool"], verbose=False) + config = DeepResearchAgentConfig(orchestrator_llm="llm", tools=["web_search_tool"]) state = DeepResearchAgentState(messages=[HumanMessage(content="research this")], data_sources=[]) with patch.object(deep_register, "filter_tools_by_sources") as filter_tools: diff --git a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py index 617ae19ca..b252d8119 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py +++ b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py @@ -1026,14 +1026,29 @@ def test_dot_suffix_with_invalid_base_passes_through(self, middleware): @pytest.mark.asyncio async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): - """Integration: middleware sanitizes tool_calls in AIMessage.""" + """Sanitize tool names without dropping provider, usage, or response metadata.""" from langchain.agents.middleware.types import ModelResponse ai_msg = AIMessage( content="", + additional_kwargs={ + "tool_calls": [ + { + "id": "tc1", + "type": "function", + "function": { + "name": "advanced_web_search_tool<|channel|>commentary", + "arguments": '{"question":"test"}', + }, + } + ], + "provider_field": "preserve-me", + }, + response_metadata={"model_name": "nvidia/nemotron-3-ultra-550b-a55b", "finish_reason": "tool_calls"}, tool_calls=[ {"name": "advanced_web_search_tool<|channel|>commentary", "args": {"question": "test"}, "id": "tc1"}, ], + usage_metadata={"input_tokens": 100, "output_tokens": 20, "total_tokens": 120}, ) mock_response = ModelResponse(result=[ai_msg]) mock_handler = AsyncMock(return_value=mock_response) @@ -1041,7 +1056,12 @@ async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): result = await middleware.awrap_model_call(mock_request, mock_handler) - assert result.result[0].tool_calls[0]["name"] == "advanced_web_search_tool" + message = result.result[0] + assert message.tool_calls[0]["name"] == "advanced_web_search_tool" + assert message.additional_kwargs["tool_calls"][0]["function"]["name"] == "advanced_web_search_tool" + assert message.additional_kwargs["provider_field"] == "preserve-me" + assert message.response_metadata == ai_msg.response_metadata + assert message.usage_metadata == ai_msg.usage_metadata @pytest.mark.asyncio async def test_awrap_model_call_no_tool_calls_passthrough(self, middleware): diff --git a/tests/aiq_agent/agents/deep_researcher/test_factory.py b/tests/aiq_agent/agents/deep_researcher/test_factory.py index 5d29868d8..5118c93e2 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_factory.py +++ b/tests/aiq_agent/agents/deep_researcher/test_factory.py @@ -537,3 +537,4 @@ class FakeSummarizationMiddleware(AgentMiddleware): assert "StructuredResponseTextFallbackMiddleware" in middleware_names assert "ToolVisibilityMiddleware" in middleware_names assert kwargs["middleware"][-2] is shared_middleware[0] + assert middleware_names[0] == "NemoRelayMiddleware" diff --git a/tests/aiq_agent/agents/test_config_observability.py b/tests/aiq_agent/agents/test_config_observability.py new file mode 100644 index 000000000..7e3857471 --- /dev/null +++ b/tests/aiq_agent/agents/test_config_observability.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from aiq_agent.agents.chat_researcher.register import ChatDeepResearcherConfig +from aiq_agent.agents.chat_researcher.register import IntentClassifierConfig +from aiq_agent.agents.clarifier.register import ClarifierConfig +from aiq_agent.agents.deep_researcher.register import DeepResearchAgentConfig +from aiq_agent.agents.shallow_researcher.register import ShallowResearchAgentConfig + + +@pytest.mark.parametrize( + "config_type", + [ + IntentClassifierConfig, + ChatDeepResearcherConfig, + ClarifierConfig, + ShallowResearchAgentConfig, + DeepResearchAgentConfig, + ], +) +def test_agent_configs_do_not_expose_legacy_verbose_switch(config_type: type) -> None: + assert "verbose" not in config_type.model_fields diff --git a/tests/aiq_agent/common/test_callbacks.py b/tests/aiq_agent/common/test_callbacks.py index 9f98d020a..67ff51ddd 100644 --- a/tests/aiq_agent/common/test_callbacks.py +++ b/tests/aiq_agent/common/test_callbacks.py @@ -16,56 +16,15 @@ """Tests for callback handlers and logging utilities.""" import logging -import os from unittest.mock import MagicMock -from unittest.mock import patch import pytest from aiq_agent.common.callbacks import ResearchLogger from aiq_agent.common.callbacks import VerboseTraceCallback -from aiq_agent.common.callbacks import is_verbose_enabled from aiq_agent.common.logging_utils import log_content_metadata -class TestIsVerboseEnabled: - """Tests for the is_verbose_enabled function.""" - - def test_verbose_enabled_true(self): - """Test is_verbose_enabled returns True for enabled values.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "1"}): - assert is_verbose_enabled() is True - - with patch.dict(os.environ, {"AIQ_VERBOSE": "true"}): - assert is_verbose_enabled() is True - - with patch.dict(os.environ, {"AIQ_VERBOSE": "yes"}): - assert is_verbose_enabled() is True - - with patch.dict(os.environ, {"AIQ_VERBOSE": "TRUE"}): - assert is_verbose_enabled() is True - - def test_verbose_disabled(self): - """Test is_verbose_enabled returns False for disabled/empty values.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "0"}): - assert is_verbose_enabled() is False - - with patch.dict(os.environ, {"AIQ_VERBOSE": "false"}): - assert is_verbose_enabled() is False - - with patch.dict(os.environ, {"AIQ_VERBOSE": "no"}): - assert is_verbose_enabled() is False - - with patch.dict(os.environ, {"AIQ_VERBOSE": ""}): - assert is_verbose_enabled() is False - - def test_verbose_unset(self): - """Test is_verbose_enabled returns False when env var is not set.""" - with patch.dict(os.environ, clear=True): - os.environ.pop("AIQ_VERBOSE", None) - assert is_verbose_enabled() is False - - class TestResearchLogger: """Tests for the ResearchLogger class.""" @@ -82,11 +41,9 @@ def test_research_logger_init_with_verbose_param(self, mock_logger): logger_non_verbose = ResearchLogger(mock_logger, verbose=False) assert logger_non_verbose.verbose is False - def test_research_logger_init_from_env(self, mock_logger): - """Test ResearchLogger initialization from environment variable.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "true"}): - logger = ResearchLogger(mock_logger) - assert logger.verbose is True + def test_research_logger_defaults_to_non_verbose(self, mock_logger): + """Research logging does not consult a process-global verbosity switch.""" + assert ResearchLogger(mock_logger).verbose is False def test_section_logs_info(self, mock_logger): """Test section method logs at info level.""" diff --git a/tests/aiq_agent/common/test_common_init.py b/tests/aiq_agent/common/test_common_init.py index fde907207..28fd4c72a 100644 --- a/tests/aiq_agent/common/test_common_init.py +++ b/tests/aiq_agent/common/test_common_init.py @@ -30,77 +30,9 @@ from aiq_agent.common import format_data_source_tools from aiq_agent.common import get_checkpointer from aiq_agent.common import is_postgres_dsn -from aiq_agent.common import is_verbose from aiq_agent.common import parse_data_sources -class TestIsVerbose: - """Tests for the is_verbose function.""" - - def test_verbose_env_true_overrides_config_false(self): - """Test that env var 'true' overrides config False.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "true"}): - assert is_verbose(config_verbose=False) is True - - def test_verbose_env_1_overrides_config_false(self): - """Test that env var '1' overrides config False.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "1"}): - assert is_verbose(config_verbose=False) is True - - def test_verbose_env_yes_overrides_config_false(self): - """Test that env var 'yes' overrides config False.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "yes"}): - assert is_verbose(config_verbose=False) is True - - def test_verbose_env_false_overrides_config_true(self): - """Test that env var 'false' overrides config True.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "false"}): - assert is_verbose(config_verbose=True) is False - - def test_verbose_env_0_overrides_config_true(self): - """Test that env var '0' overrides config True.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "0"}): - assert is_verbose(config_verbose=True) is False - - def test_verbose_env_no_overrides_config_true(self): - """Test that env var 'no' overrides config True.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "no"}): - assert is_verbose(config_verbose=True) is False - - def test_verbose_env_empty_uses_config_true(self): - """Test that empty env var falls back to config.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": ""}): - assert is_verbose(config_verbose=True) is True - - def test_verbose_env_empty_uses_config_false(self): - """Test that empty env var falls back to config False.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": ""}): - assert is_verbose(config_verbose=False) is False - - def test_verbose_env_unset_uses_config_true(self): - """Test that unset env var falls back to config True.""" - with patch.dict(os.environ, clear=True): - os.environ.pop("AIQ_VERBOSE", None) - assert is_verbose(config_verbose=True) is True - - def test_verbose_env_unset_uses_config_false(self): - """Test that unset env var falls back to config False.""" - with patch.dict(os.environ, clear=True): - os.environ.pop("AIQ_VERBOSE", None) - assert is_verbose(config_verbose=False) is False - - def test_verbose_env_uppercase_true(self): - """Test that uppercase 'TRUE' is recognized.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "TRUE"}): - assert is_verbose(config_verbose=False) is True - - def test_verbose_env_random_value_uses_config(self): - """Test that random env value falls back to config.""" - with patch.dict(os.environ, {"AIQ_VERBOSE": "maybe"}): - assert is_verbose(config_verbose=True) is True - assert is_verbose(config_verbose=False) is False - - class TestCreateChatResponse: """Tests for the _create_chat_response function.""" diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 5aa12f548..10bc498c8 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -393,6 +393,28 @@ class TestSubmitDeepResearchJob: principal = Principal(type="test", sub="user-1", email="test@example.com", name="Test User") + def test_job_trace_correlation_keeps_session_and_submission_ids(self): + from aiq_api.auth.request_trace import request_trace_tag_context + from aiq_api.jobs.submit import _get_job_trace_correlation + from nat.builder.context import ContextState + + state = ContextState.get() + conversation_token = state.conversation_id.set("conversation-1") + trace_token = state.workflow_trace_id.set(0x1234) + span_token = state.active_span_id_stack.set(["root", "submission-span"]) + try: + with request_trace_tag_context({"tenant": "test"}): + correlation = _get_job_trace_correlation() + finally: + state.active_span_id_stack.reset(span_token) + state.workflow_trace_id.reset(trace_token) + state.conversation_id.reset(conversation_token) + + assert correlation.session_id == "conversation-1" + assert correlation.submission_trace_id == f"{0x1234:032x}" + assert correlation.submission_span_id == "submission-span" + assert correlation.request_trace_tags == {"tenant": "test"} + @pytest.fixture(autouse=True) def _isolate_admission_store(self): """Keep legacy submit tests scoped to submission wiring, not admission-store integration.""" @@ -510,7 +532,7 @@ async def test_submit_agent_job_passes_explicit_conversation_id_to_worker(self): from aiq_api.jobs.runner import run_agent_job worker_args = inspect.signature(run_agent_job).bind(*job_args).arguments - assert worker_args["parent_conversation_id"] == "customer-collection" + assert worker_args["trace_correlation"].session_id == "customer-collection" @pytest.mark.asyncio async def test_submit_agent_job_passes_initial_files_and_output_metadata(self): @@ -691,6 +713,7 @@ async def test_worker_binds_conversation_before_tool_construction_and_invocation from types import SimpleNamespace from aiq_api.jobs.crypto import ContentEncryptionConfig + from aiq_api.jobs.runner import JobTraceCorrelation from aiq_api.jobs.runner import run_agent_job from nat.builder.context import Context from nat.builder.context import ContextState @@ -703,6 +726,8 @@ async def __aexit__(self, exc_type, exc, tb): return False observed: dict[str, str | None] = {} + relay_observed: dict[str, object] = {} + nat_events = [] class ContextAwareKnowledgeTool: def __init__(self, construction_conversation_id): @@ -724,7 +749,7 @@ async def __aexit__(self, exc_type, exc, tb): return False def get_function_config(self, _name): - return SimpleNamespace(tools=["knowledge_retrieval"], exclude_tools=[], verbose=False) + return SimpleNamespace(tools=["knowledge_retrieval"], exclude_tools=[]) async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API async def build_tool(): @@ -738,7 +763,9 @@ async def build_tool(): return list(await asyncio.gather(build_tool())) class FakeExporterManager: - def start(self, *, context_state): # noqa: ARG002 - mirrors NAT API + def start(self, *, context_state): + observed["worker_trace_id"] = f"{context_state.workflow_trace_id.get():032x}" + context_state.event_stream.get().subscribe(nat_events.append) return AsyncContext() def create_agent(*, tools, **_kwargs): @@ -748,6 +775,10 @@ async def run_agent(*, agent, **_kwargs): observed["resolved_collection"] = await agent.tools[0].ainvoke() raise RuntimeError("stop after context assertion") + async def run_relay_workflow(name, operation, **kwargs): + relay_observed.update({"name": name, **kwargs}) + return await operation() + mock_job_store = MagicMock(update_status=AsyncMock()) config = SimpleNamespace(functions={}, middleware={}) db_url = f"sqlite:///{tmp_path / 'conversation-context.db'}" @@ -771,6 +802,7 @@ async def run_agent(*, agent, **_kwargs): patch("aiq_api.jobs.runner._create_agent_instance", side_effect=create_agent), patch("aiq_api.jobs.runner._run_agent", side_effect=run_agent), patch("aiq_api.jobs.runner._run_lease_refresher"), + patch("aiq_agent.relay.run_workflow", side_effect=run_relay_workflow), patch("aiq_api.mcp_auth.runtime_tools.open_per_user_mcp_tools", AsyncMock(return_value=[])), ): await run_agent_job( @@ -783,7 +815,11 @@ async def run_agent(*, agent, **_kwargs): "input", "aiq_agent.agents.shallow_researcher.agent.ShallowResearcherAgent", "shallow_research_agent", - parent_conversation_id=conversation_id, + trace_correlation=JobTraceCorrelation( + session_id=conversation_id, + submission_trace_id="1" * 32, + submission_span_id="submission-span", + ), content_encryption_policy=ContentEncryptionConfig(mode="off").policy_identity, owner_user_id=owner_user_id, ) @@ -794,7 +830,22 @@ async def run_agent(*, agent, **_kwargs): "invocation_conversation_id": conversation_id, "invocation_user_id": owner_user_id, "resolved_collection": expected_collection, + "worker_trace_id": observed["worker_trace_id"], + } + assert observed["worker_trace_id"] != "1" * 32 + assert relay_observed == { + "name": "async_shallow_research_job", + "session_id": conversation_id, + "input_value": "input", + "metadata": { + "aiq.execution.mode": "async", + "aiq.job.id": "job-1", + "aiq.agent.type": "shallow_research_agent", + "aiq.submission.trace_id": "1" * 32, + "aiq.submission.span_id": "submission-span", + }, } + assert all(step.payload.UUID != "submission-span" for step in nat_events) assert outer_context.conversation_id.get() == "stale-parent-context" assert outer_context.user_id.get() == "jwt:stale-owner" finally: @@ -948,7 +999,7 @@ async def __aexit__(self, exc_type, exc, tb): return False def get_function_config(self, _name): - return SimpleNamespace(tools=[], exclude_tools=[], verbose=False) + return SimpleNamespace(tools=[], exclude_tools=[]) async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API return [] @@ -1064,7 +1115,7 @@ async def __aexit__(self, exc_type, exc, tb): return False def get_function_config(self, _name): - return SimpleNamespace(tools=[], exclude_tools=[], verbose=False) + return SimpleNamespace(tools=[], exclude_tools=[]) async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API return [] @@ -1191,7 +1242,7 @@ async def __aexit__(self, exc_type, exc, tb): return False def get_function_config(self, _name): - return SimpleNamespace(tools=[], exclude_tools=[], verbose=False) + return SimpleNamespace(tools=[], exclude_tools=[]) async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API return [] @@ -1512,7 +1563,7 @@ async def __aexit__(self, exc_type, exc, tb): return False def get_function_config(self, _name): - return SimpleNamespace(tools=[], exclude_tools=[], verbose=False) + return SimpleNamespace(tools=[], exclude_tools=[]) async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API return [] @@ -2843,7 +2894,6 @@ def __init__( *, llm_provider, tools, - verbose, callbacks, domain_catalog_path=None, enable_source_router=True, @@ -2860,7 +2910,6 @@ def __init__( ): self.llm_provider = llm_provider self.tools = tools - self.verbose = verbose self.callbacks = callbacks self.domain_catalog_path = domain_catalog_path self.enable_source_router = enable_source_router @@ -2903,7 +2952,6 @@ def __init__( llm="llm", tools=["tool"], fn_config=fn_config, - verbose=True, callbacks=["callback"], job_id="job-123", ) @@ -2978,14 +3026,12 @@ def __init__( llm_provider, tools=None, *, - verbose=False, callbacks=None, config=None, job_id=None, ): self.llm_provider = llm_provider self.tools = tools - self.verbose = verbose self.callbacks = callbacks self.config = config self.job_id = job_id @@ -3003,14 +3049,12 @@ def __init__( llm="llm", tools=["tool"], fn_config=fn_config, - verbose=True, callbacks=["callback"], job_id="job-123", ) assert agent.llm_provider == "provider" assert agent.tools == ["tool"] - assert agent.verbose is True assert agent.callbacks == ["callback"] assert agent.config is fn_config assert agent.job_id == "job-123" @@ -3026,13 +3070,11 @@ def __init__( llm_provider, tools=None, *, - verbose=False, callbacks=None, job_id=None, ): self.llm_provider = llm_provider self.tools = tools - self.verbose = verbose self.callbacks = callbacks self.job_id = job_id @@ -3042,14 +3084,12 @@ def __init__( llm="llm", tools=["tool"], fn_config=DeepResearchAgentConfig(orchestrator_llm="llm"), - verbose=True, callbacks=["callback"], job_id="job-123", ) assert agent.llm_provider == "provider" assert agent.tools == ["tool"] - assert agent.verbose is True assert agent.callbacks == ["callback"] assert agent.job_id == "job-123" @@ -3082,7 +3122,6 @@ def test_async_deep_researcher_constructor_applies_config_tuning(self): llm=mock_llm, tools=[], fn_config=fn_config, - verbose=False, callbacks=[], job_id="async-job-123", ) @@ -3127,7 +3166,6 @@ def stop(self): orchestrator_llm="llm", enable_citation_verification=False, ), - verbose=False, callbacks=[], job_id="async-job-123", ) @@ -3303,7 +3341,6 @@ def async_test_search(query: str) -> str: llm=mock_llm, tools=[async_test_search], fn_config=fn_config, - verbose=False, callbacks=[], job_id="async-job-123", ) @@ -3370,7 +3407,6 @@ def test_async_deep_researcher_empty_data_sources_keeps_internal_tools(self): llm=mock_llm, tools=[], fn_config=DeepResearchAgentConfig(orchestrator_llm="llm"), - verbose=False, callbacks=[], job_id="async-job-123", ) @@ -3400,7 +3436,6 @@ def __init__( *, llm_provider, tools, - verbose, callbacks, domain_catalog_path=None, enable_source_router=True, @@ -3430,7 +3465,6 @@ def __init__( llm="llm", tools=["tool"], fn_config=fn_config, - verbose=True, callbacks=["callback"], job_id="job-123", ) diff --git a/tests/aiq_agent/jobs/test_telemetry.py b/tests/aiq_agent/jobs/test_telemetry.py deleted file mode 100644 index 0380382fc..000000000 --- a/tests/aiq_agent/jobs/test_telemetry.py +++ /dev/null @@ -1,241 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import asyncio -from uuid import UUID - -import pytest -from langchain_core.messages import AIMessage -from langchain_core.messages import HumanMessage -from langchain_core.outputs import ChatGeneration -from langchain_core.outputs import LLMResult - -from nat.builder.context import Context -from nat.builder.context import ContextState -from nat.data_models.intermediate_step import IntermediateStepPayload -from nat.data_models.intermediate_step import IntermediateStepType -from nat.data_models.invocation_node import InvocationNode -from nat.utils.reactive.subject import Subject - - -@pytest.fixture -def telemetry_context(): - """Provide an isolated NAT event stream and span stack for callback tests.""" - state = ContextState.get() - event_stream = Subject() - tokens = [ - (state.event_stream, state.event_stream.set(event_stream)), - (state.active_span_id_stack, state.active_span_id_stack.set(["root"])), - ( - state.active_function, - state.active_function.set(InvocationNode(function_name="deep_research_agent", function_id="job-1")), - ), - (state.workflow_run_id, state.workflow_run_id.set("job-1")), - (state.workflow_trace_id, state.workflow_trace_id.set(0x1234)), - ] - manager = Context(state).intermediate_step_manager - events = [] - event_stream.subscribe(events.append) - - yield state, manager, events - - for context_var, token in reversed(tokens): - context_var.reset(token) - - -@pytest.mark.asyncio -async def test_agent_telemetry_nests_named_subagent_and_model_under_task(telemetry_context): - """DeepAgents task -> named agent -> model must be an explicit NAT span hierarchy.""" - from aiq_api.jobs.telemetry import AgentLifecycleTelemetryCallback - from aiq_api.jobs.telemetry import AIQLangchainProfilerHandler - - _, manager, events = telemetry_context - task_id = UUID("00000000-0000-0000-0000-000000000001") - agent_id = UUID("00000000-0000-0000-0000-000000000002") - model_id = UUID("00000000-0000-0000-0000-000000000003") - - manager.push_intermediate_step( - IntermediateStepPayload( - UUID="job-1", - event_type=IntermediateStepType.WORKFLOW_START, - name="deep_research_agent", - ) - ) - agent_callback = AgentLifecycleTelemetryCallback(manager) - profiler_callback = AIQLangchainProfilerHandler() - - await profiler_callback.on_tool_start( - {"name": "task"}, - "{'subagent_type': 'planner-agent'}", - run_id=task_id, - inputs={"subagent_type": "planner-agent"}, - ) - agent_callback.on_chain_start( - None, - inputs={"messages": [HumanMessage(content="plan the research")]}, - run_id=agent_id, - parent_run_id=task_id, - name="planner-agent", - metadata={"lc_agent_name": "planner-agent"}, - ) - await profiler_callback.on_chat_model_start( - {}, - [[HumanMessage(content="plan the research")]], - run_id=model_id, - parent_run_id=agent_id, - metadata={"ls_model_name": "test-model"}, - invocation_params={}, - ) - - starts = {event.UUID: event for event in events if event.event_state.value == "START"} - assert starts[str(task_id)].payload.name == "task: planner-agent" - assert starts[str(task_id)].parent_id == "job-1" - assert starts[str(agent_id)].event_type == IntermediateStepType.WORKFLOW_START - assert starts[str(agent_id)].payload.name == "planner-agent" - assert starts[str(agent_id)].payload.metadata.provided_metadata["span_role"] == "agent" - assert starts[str(agent_id)].parent_id == str(task_id) - assert starts[str(model_id)].parent_id == str(agent_id) - - result = LLMResult( - generations=[[ChatGeneration(message=AIMessage(content="plan complete"))]], - llm_output={"model_name": "test-model"}, - ) - await profiler_callback.on_llm_end(result, run_id=model_id) - agent_callback.on_chain_end({"messages": [AIMessage(content="plan complete")]}, run_id=agent_id) - await profiler_callback.on_tool_end("plan complete", run_id=task_id, name="task") - manager.push_intermediate_step( - IntermediateStepPayload( - UUID="job-1", - event_type=IntermediateStepType.WORKFLOW_END, - name="deep_research_agent", - ) - ) - - started_ids = {event.UUID for event in events if event.event_state.value == "START"} - assert all(event.parent_id == "root" or event.parent_id in started_ids for event in events) - assert manager.get_outstanding_step_count() == 0 - assert profiler_callback.step_manager.get_outstanding_step_count() == 0 - - -@pytest.mark.asyncio -async def test_parallel_researcher_spans_share_batch_parent_without_sharing_identity(telemetry_context): - """Concurrent researcher runs remain distinct children of one batch tool span.""" - from aiq_api.jobs.telemetry import AgentLifecycleTelemetryCallback - from aiq_api.jobs.telemetry import AIQLangchainProfilerHandler - - _, manager, events = telemetry_context - batch_id = UUID("00000000-0000-0000-0000-000000000010") - researcher_ids = [ - UUID("00000000-0000-0000-0000-000000000011"), - UUID("00000000-0000-0000-0000-000000000012"), - ] - - manager.push_intermediate_step( - IntermediateStepPayload( - UUID="job-1", - event_type=IntermediateStepType.WORKFLOW_START, - name="deep_research_agent", - ) - ) - agent_callback = AgentLifecycleTelemetryCallback(manager) - profiler_callback = AIQLangchainProfilerHandler() - await profiler_callback.on_tool_start( - {"name": "run_research_batch"}, - "{}", - run_id=batch_id, - inputs={}, - ) - - async def run_researcher(run_id: UUID) -> None: - agent_callback.on_chain_start( - None, - inputs={}, - run_id=run_id, - parent_run_id=batch_id, - name="researcher-agent", - metadata={"lc_agent_name": "researcher-agent"}, - ) - await asyncio.sleep(0) - agent_callback.on_chain_end({}, run_id=run_id) - - await asyncio.gather(*(run_researcher(run_id) for run_id in researcher_ids)) - - starts = { - event.UUID: event - for event in events - if event.event_state.value == "START" and event.payload.name == "researcher-agent" - } - assert set(starts) == {str(run_id) for run_id in researcher_ids} - assert {event.parent_id for event in starts.values()} == {str(batch_id)} - - -def test_agent_lifecycle_spans_do_not_capture_graph_state(telemetry_context): - """Structural agent spans must not duplicate LangGraph state into telemetry.""" - from aiq_api.jobs.telemetry import AgentLifecycleTelemetryCallback - - _, manager, events = telemetry_context - callback = AgentLifecycleTelemetryCallback(manager) - run_id = UUID("00000000-0000-0000-0000-000000000020") - - callback.on_chain_start( - None, - inputs={"messages": [HumanMessage(content="sensitive input")]}, - run_id=run_id, - name="researcher-agent", - metadata={"lc_agent_name": "researcher-agent"}, - ) - callback.on_chain_end( - {"messages": [AIMessage(content="sensitive output")]}, - run_id=run_id, - ) - - assert len(events) == 2 - assert all(event.payload.data is None for event in events) - assert all("sensitive input" not in str(event.payload.metadata) for event in events) - assert all("sensitive output" not in str(event.payload.metadata) for event in events) - - -@pytest.mark.parametrize( - ("name", "metadata", "expected"), - [ - ("general-purpose", {"lc_agent_name": "general-purpose"}, True), - ("reviewer", {"lc_agent_name": "reviewer"}, True), - ("agent", {"langgraph_node": "agent"}, False), - ("model", {"lc_agent_name": "reviewer"}, False), - ], -) -def test_agent_lifecycle_requires_matching_deepagents_identity(telemetry_context, name, metadata, expected): - """Only the outer chain identified by DeepAgents metadata is an agent boundary.""" - from aiq_api.jobs.telemetry import AgentLifecycleTelemetryCallback - - _, manager, events = telemetry_context - callback = AgentLifecycleTelemetryCallback(manager) - run_id = UUID("00000000-0000-0000-0000-000000000030") - - callback.on_chain_start(None, inputs={}, run_id=run_id, name=name, metadata=metadata) - callback.on_chain_end({}, run_id=run_id) - - assert [event.payload.name for event in events] == ([name, name] if expected else []) - assert manager.get_outstanding_step_count() == 0 - - -def test_aiq_profiler_context_replaces_and_restores_nat_profiler(telemetry_context): - """AIQ must customize NAT's inherited profiler instead of installing a duplicate callback.""" - from aiq_api.jobs.telemetry import AIQLangchainProfilerHandler - from aiq_api.jobs.telemetry import aiq_langchain_profiler_context - from nat.plugins.langchain.callback_handler import LangchainProfilerHandler - from nat.plugins.profiler.decorators.framework_wrapper import callback_handler_var - - nat_profiler = LangchainProfilerHandler() - token = callback_handler_var.set(nat_profiler) - try: - with aiq_langchain_profiler_context() as aiq_profiler: - assert isinstance(aiq_profiler, AIQLangchainProfilerHandler) - assert callback_handler_var.get() is aiq_profiler - assert callback_handler_var.get() is not nat_profiler - - assert callback_handler_var.get() is nat_profiler - finally: - callback_handler_var.reset(token) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..f4d5f25a6 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository-wide pytest isolation for process-discovered configuration.""" + +import os +import shutil +import tempfile +from pathlib import Path + +# NeMo Relay discovers ``$XDG_CONFIG_HOME/nemo-relay/plugins.toml`` during +# import/initialization. Unit tests create real Relay scopes, so inheriting a +# developer's XDG directory would export fixture traffic to their configured +# observability destinations. Set this during conftest import, before pytest +# imports test modules that import Relay. +_TEST_XDG_CONFIG_HOME = Path(tempfile.mkdtemp(prefix="aiq-pytest-xdg-")) +os.environ["XDG_CONFIG_HOME"] = str(_TEST_XDG_CONFIG_HOME) + + +def pytest_unconfigure() -> None: + """Remove the process-local Relay discovery directory after the test run.""" + shutil.rmtree(_TEST_XDG_CONFIG_HOME, ignore_errors=True) diff --git a/tests/test_relay_runtime.py b/tests/test_relay_runtime.py new file mode 100644 index 000000000..bf3fbd996 --- /dev/null +++ b/tests/test_relay_runtime.py @@ -0,0 +1,820 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import json +from collections import Counter +from http.server import BaseHTTPRequestHandler +from http.server import ThreadingHTTPServer +from pathlib import Path +from threading import Thread +from types import SimpleNamespace +from uuid import uuid4 + +import nemo_relay +import pytest +from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel +from langchain_core.messages import AIMessage +from langchain_core.messages import HumanMessage +from langchain_core.tools import tool +from nemo_relay import plugin +from nemo_relay.integrations.langchain._serialization import payload_to_model_request +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest +from pydantic import ValidationError + +from aiq_agent.agents.deep_researcher.models import ResearchQuery +from aiq_agent.agents.deep_researcher.tools.research import _run_research_queries +from aiq_agent.relay.bootstrap import ensure_started +from aiq_agent.relay.bootstrap import shutdown_async +from aiq_agent.relay.config import RelayConfig +from aiq_agent.relay.config import RelayOpenTelemetryEndpointConfig +from aiq_agent.relay.logging import log_event +from aiq_agent.relay.privacy import request_privacy_context +from aiq_agent.relay.runtime import _normalize_chat_nvidia_binding +from aiq_agent.relay.runtime import _safe_value +from aiq_agent.relay.runtime import ainvoke_tool_with_relay +from aiq_agent.relay.runtime import ainvoke_with_relay +from aiq_agent.relay.runtime import deepagents_kwargs +from aiq_agent.relay.runtime import merge_langchain_middleware +from aiq_agent.relay.runtime import run_agent +from aiq_agent.relay.runtime import run_workflow + + +def test_deepagents_integration_and_delegated_agent_scope_are_enabled() -> None: + kwargs = deepagents_kwargs( + { + "model": "test", + "tools": [], + "name": "test-agent", + "subagents": [{"name": "runtime-agent", "description": "test", "model": "test", "tools": []}], + } + ) + assert [type(middleware).__name__ for middleware in kwargs["middleware"][-2:]] == [ + "NemoRelayDeepAgentsMiddleware", + "_DelegatedAgentScopeMiddleware", + ] + assert [type(middleware).__name__ for middleware in kwargs["subagents"][0]["middleware"][-1:]] == [ + "NemoRelayDeepAgentsMiddleware" + ] + + +def test_langchain_managed_execution_middleware_is_enabled() -> None: + middleware = merge_langchain_middleware([]) + + assert [type(item).__name__ for item in middleware] == ["NemoRelayMiddleware"] + + +def test_bound_chat_nvidia_uses_supported_relay_header_path() -> None: + from langchain.agents.middleware import ModelRequest + from langchain_nvidia_ai_endpoints import ChatNVIDIA + + @tool + def search(query: str) -> str: + """Search for a query.""" + return query + + model = ChatNVIDIA( + model="nvidia/nemotron-3-ultra-550b-a55b", + api_key="test-key", # pragma: allowlist secret + base_url="https://example.invalid/v1", + ) + bound_model = model.bind_tools([search], parallel_tool_calls=False) + normalized, model_settings, config = _normalize_chat_nvidia_binding( + bound_model, + {"tags": ["request"]}, + ) + + assert normalized is model + assert model_settings["tools"][0]["function"]["name"] == "search" + assert model_settings["parallel_tool_calls"] is False + assert config["tags"] == ["request"] + + request = ModelRequest( + model=normalized, + messages=[HumanMessage(content="test")], + model_settings=model_settings, + ) + relay_request = nemo_relay.LLMRequest( + {"traceparent": "00-test-trace-test-span-01"}, + {"model_settings": model_settings}, + ) + converted = payload_to_model_request(request, relay_request) + + assert isinstance(converted.model, ChatNVIDIA) + assert converted.model.default_headers["traceparent"] == "00-test-trace-test-span-01" + assert "extra_headers" not in converted.model_settings + assert converted.model_settings["tools"][0]["function"]["name"] == "search" + + +def test_relay_logging_subscriber_matches_verbose_trace_labels(caplog) -> None: + caplog.set_level("INFO") + root_uuid = str(uuid4()) + llm_uuid = str(uuid4()) + tool_uuid = str(uuid4()) + researcher_uuid = str(uuid4()) + events = [ + SimpleNamespace( + kind="scope", + category="agent", + name="test-agent", + scope_category="start", + metadata={}, + uuid=root_uuid, + parent_uuid=None, + ), + SimpleNamespace( + kind="scope", + category="llm", + name="ChatNVIDIA", + scope_category="start", + metadata={}, + uuid=llm_uuid, + parent_uuid=root_uuid, + data={"messages": "redacted Relay input"}, + category_profile={}, + ), + SimpleNamespace( + kind="scope", + category="llm", + name="ChatNVIDIA", + scope_category="end", + metadata={"otel.status_code": "OK"}, + uuid=llm_uuid, + parent_uuid=root_uuid, + data={ + "generations": [ + [ + { + "message": { + "content": "redacted Relay response", + "additional_kwargs": {"reasoning_content": "redacted Relay reasoning"}, + "response_metadata": { + "model_name": "test-model", + "token_usage": {"prompt_tokens": 10, "completion_tokens": 4}, + }, + } + } + ] + ] + }, + category_profile={ + "model_name": "test-model", + "annotated_response": { + "tool_calls": [ + { + "name": "web_search", + "arguments": {"query": "redacted Relay query"}, + } + ] + }, + }, + ), + SimpleNamespace( + kind="scope", + category="tool", + name="web_search", + scope_category="start", + metadata={}, + uuid=tool_uuid, + parent_uuid=root_uuid, + data={"query": "redacted Relay query"}, + ), + SimpleNamespace( + kind="scope", + category="tool", + name="web_search", + scope_category="end", + metadata={"otel.status_code": "OK"}, + uuid=tool_uuid, + parent_uuid=root_uuid, + data={"result": "redacted Relay result"}, + ), + SimpleNamespace( + kind="scope", + category="agent", + name="researcher-agent", + scope_category="start", + uuid=researcher_uuid, + parent_uuid=root_uuid, + metadata={}, + ), + SimpleNamespace( + kind="scope", + category="agent", + name="researcher-agent", + scope_category="end", + uuid=researcher_uuid, + parent_uuid=root_uuid, + metadata={"otel.status_code": "OK"}, + ), + SimpleNamespace( + kind="scope", + category="agent", + name="test-agent", + scope_category="end", + metadata={"otel.status_code": "OK"}, + uuid=root_uuid, + parent_uuid=None, + ), + ] + for event in events: + log_event(event) + + for label in ( + "[Chain Start] test-agent", + "[AGENT]", + "[Reasoning]", + "[Agent Response]", + "[Tool Calls] 1 tool(s) requested", + "→ web_search", + "Args: chars=33", + "[Tokens] prompt=10, completion=4, model=test-model", + "[Tool Start] web_search", + "[Tool Result]", + "[Chain Start] researcher-agent", + "[Chain End] researcher-agent", + "[Chain End] test-agent", + ): + assert label in caplog.text + assert "[Relay Scope]" not in caplog.text + assert "[Relay LLM]" not in caplog.text + assert "[Relay Tool]" not in caplog.text + + +def test_safe_value_projects_pydantic_like_state() -> None: + class State: + pass + + assert _safe_value({"state": State()}) == {"state": {"type": "State"}} + + +@pytest.mark.asyncio +async def test_callback_does_not_duplicate_middleware_managed_llm_and_tool_scopes(tmp_path: Path) -> None: + class TestChatModel(FakeMessagesListChatModel): + model_name: str = "managed-model" + + @tool + def managed_tool(value: str) -> str: + """Return the supplied value.""" + return value + + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "managed.jsonl" + config.observability.opentelemetry.enabled = False + + async def operation() -> None: + model = TestChatModel( + responses=[ + AIMessage( + content="response", + response_metadata={"model_name": "managed-model"}, + usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}, + ) + ] + ) + response = await ainvoke_with_relay( + model, + [HumanMessage(content="request")], + ) + assert response.content == "response" + assert await ainvoke_tool_with_relay(managed_tool, {"value": "result"}) == "result" + + await ensure_started(config) + try: + await run_agent("managed-agent", operation) + finally: + await shutdown_async() + + events = [json.loads(line) for line in (tmp_path / "managed.jsonl").read_text().splitlines()] + starts = [event for event in events if event["kind"] == "scope" and event["scope_category"] == "start"] + assert [(event["category"], event["name"]) for event in starts] == [ + ("agent", "managed-agent"), + ("llm", "managed-model"), + ("tool", "managed_tool"), + ] + llm_end = next(event for event in events if event["category"] == "llm" and event["scope_category"] == "end") + assert llm_end["category_profile"]["model_name"] == "managed-model" + assert llm_end["category_profile"]["annotated_response"]["usage"] == { + "completion_tokens": 4, + "prompt_tokens": 10, + "total_tokens": 14, + } + + +@pytest.mark.asyncio +async def test_deepagents_task_uses_runtime_subagent_name_for_nested_scope(tmp_path: Path) -> None: + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "delegation.jsonl" + config.observability.opentelemetry.enabled = False + delegation_middleware = deepagents_kwargs({"model": "test", "tools": [], "name": "parent"})["middleware"][-1] + request = SimpleNamespace( + tool_call={ + "name": "task", + "args": {"subagent_type": "runtime-selected-agent", "description": "research this"}, + } + ) + + async def delegated_agent(_: object) -> str: + async def model_call(_: nemo_relay.LLMRequest) -> dict[str, str]: + return {"response": "done"} + + await nemo_relay.llm.execute( + "managed-model", + nemo_relay.LLMRequest({}, {"messages": []}), + model_call, + ) + return "done" + + async def task_call(_: object) -> str: + return await delegation_middleware.awrap_tool_call(request, delegated_agent) + + async def operation() -> None: + await nemo_relay.tools.execute("task", request.tool_call["args"], task_call) + + await ensure_started(config) + try: + await run_agent("deep_research_agent", operation) + finally: + await shutdown_async() + + events = [json.loads(line) for line in (tmp_path / "delegation.jsonl").read_text().splitlines()] + starts = { + event["name"]: event for event in events if event["kind"] == "scope" and event["scope_category"] == "start" + } + assert starts["task"]["parent_uuid"] == starts["deep_research_agent"]["uuid"] + assert starts["runtime-selected-agent"]["parent_uuid"] == starts["deep_research_agent"]["uuid"] + assert starts["managed-model"]["parent_uuid"] == starts["runtime-selected-agent"]["uuid"] + + +@pytest.mark.asyncio +async def test_concurrent_researchers_do_not_share_mutable_relay_agent_scopes(tmp_path: Path) -> None: + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "concurrent-researchers.jsonl" + config.observability.opentelemetry.enabled = False + + class Researcher: + async def ainvoke(self, state, config=None): # noqa: ARG002 + await asyncio.sleep(0.01 if "slow" in state["messages"][0].content else 0) + return { + "structured_response": { + "query_topic": "test", + "target_components": ["test"], + "summary": "test", + "findings": [], + "gaps": [], + "sources": [], + "narrative_notes": "test", + "language": "English", + } + } + + queries = [ + ResearchQuery( + query=query, + preferred_tools=["test_tool"], + target_components=["test"], + rationale="test", + ) + for query in ("slow query", "fast query") + ] + + async def operation() -> None: + successful, notes, errors = await _run_research_queries( + queries=queries, + researcher_runnable=Researcher(), + runtime=None, + callbacks=[], + max_concurrency=2, + ) + assert successful == queries + assert len(notes) == 2 + assert errors == [] + + await ensure_started(config) + try: + await run_agent("deep_research_agent", operation) + finally: + await shutdown_async() + + events = [json.loads(line) for line in (tmp_path / "concurrent-researchers.jsonl").read_text().splitlines()] + scope_events = [event for event in events if event["kind"] == "scope"] + starts = [event for event in scope_events if event["scope_category"] == "start"] + assert [event["name"] for event in starts] == [ + "deep_research_agent", + "researcher-agent", + "researcher-agent", + ] + researcher_scopes = [event for event in scope_events if event["name"] == "researcher-agent"] + assert Counter(event["scope_category"] for event in researcher_scopes) == {"start": 2, "end": 2} + assert len({event["uuid"] for event in researcher_scopes}) == 2 + assert {event["parent_uuid"] for event in researcher_scopes} == {starts[0]["uuid"]} + + +@pytest.mark.asyncio +async def test_semantic_scope_capture_failures_do_not_change_agent_execution(monkeypatch) -> None: + operation_calls = 0 + + async def operation() -> str: + nonlocal operation_calls + operation_calls += 1 + return "result" + + def fail_start(*args, **kwargs): # noqa: ARG001 + raise RuntimeError("synthetic Relay start failure") + + monkeypatch.setattr(nemo_relay.scope, "push", fail_start) + assert await run_agent("test-agent", operation) == "result" + + monkeypatch.setattr(nemo_relay.scope, "push", lambda *args, **kwargs: object()) + + def fail_end(*args, **kwargs): # noqa: ARG001 + raise RuntimeError("synthetic Relay end failure") + + monkeypatch.setattr(nemo_relay.scope, "pop", fail_end) + assert await run_agent("test-agent", operation) == "result" + assert operation_calls == 2 + + +def test_relay_config_is_accepted_by_plugin_validator() -> None: + report = plugin.validate(RelayConfig().to_plugin_config()) + + assert not [diagnostic for diagnostic in report["diagnostics"] if diagnostic["level"] == "error"] + + +@pytest.mark.parametrize( + "value", + [ + {"unknown": True}, + {"observability": {"unknown": True}}, + {"observability": {"opentelemetry": {"endpoints": [{"endpoint": "not-a-url"}]}}}, + {"redaction": {"request_privacy_attributes": ["input.value"]}}, + ], +) +def test_relay_config_rejects_unknown_or_invalid_values(value: dict) -> None: + with pytest.raises(ValidationError): + RelayConfig.model_validate(value) + + +@pytest.mark.asyncio +async def test_plugin_managed_atof_redacts_before_export(tmp_path: Path) -> None: + config = RelayConfig() + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "events.jsonl" + config.observability.opentelemetry.enabled = False + + await ensure_started(config) + try: + with nemo_relay.scope.scope( + "redaction-test", + nemo_relay.ScopeType.Agent, + input={"email": "person@example.com"}, + ): + nemo_relay.scope.event("secret", data={"api_key": "sk-1234567890abcdef"}) # pragma: allowlist secret + finally: + await shutdown_async() + + exported = (tmp_path / "events.jsonl").read_text() + assert "redaction-test" in exported + assert "person@example.com" not in exported + assert "sk-1234567890abcdef" not in exported + + +@pytest.mark.asyncio +async def test_request_privacy_sanitizes_relay_without_changing_execution(tmp_path: Path) -> None: + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "private-events.jsonl" + config.observability.opentelemetry.enabled = False + + @tool + def echo_private(value: str) -> str: + """Echo a value for privacy testing.""" + return f"tool-result:{value}" + + class TestChatModel(FakeMessagesListChatModel): + model_name: str = "test-model" + + model = TestChatModel(responses=[AIMessage(content="private-model-output")]) + + async def operation() -> dict[str, str]: + message = await ainvoke_with_relay(model, [HumanMessage(content="private-model-input")]) + tool_result = await ainvoke_tool_with_relay(echo_private, {"value": "private-tool-input"}) + return {"model": str(message.content), "tool": tool_result} + + await ensure_started(config) + try: + with request_privacy_context(True): + result = await run_workflow("private-workflow", operation, input_value="private-workflow-input") + finally: + await shutdown_async() + + assert result == {"model": "private-model-output", "tool": "tool-result:private-tool-input"} + exported = (tmp_path / "private-events.jsonl").read_text() + for private_value in ( + "private-workflow-input", + "private-model-input", + "private-model-output", + "private-tool-input", + "tool-result:private-tool-input", + ): + assert private_value not in exported + + +@pytest.mark.asyncio +async def test_two_turn_parity_has_two_traces_one_session_no_duplicates_and_balanced_scopes(tmp_path: Path) -> None: + @tool + def echo(text: str) -> str: + """Return the supplied text.""" + return text + + received: list[bytes] = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 - stdlib callback name + received.append(self.rfile.read(int(self.headers["content-length"]))) + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + server_thread = Thread(target=server.serve_forever, daemon=True) + server_thread.start() + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "two-turns.jsonl" + config.observability.opentelemetry.enabled = True + config.observability.opentelemetry.endpoints = [ + RelayOpenTelemetryEndpointConfig( + endpoint=f"http://127.0.0.1:{server.server_port}/v1/traces", + timeout_millis=1000, + ) + ] + + async def turn(turn_number: int) -> dict[str, int]: + class TestChatModel(FakeMessagesListChatModel): + model_name: str = "ChatNVIDIA" + + async def classify_intent() -> dict[str, str]: + model = TestChatModel( + responses=[ + AIMessage( + content=f"answer-{turn_number}", + response_metadata={ + "model_name": "test-model", + "token_usage": {"prompt_tokens": 10, "completion_tokens": 4}, + }, + ) + ] + ) + await ainvoke_with_relay(model, [HumanMessage(content=f"question-{turn_number}")]) + return {"intent": "research"} + + await run_agent("intent_classifier", classify_intent, input_value={"turn": turn_number}) + + await ainvoke_tool_with_relay(echo, {"text": f"turn-{turn_number}"}) + + async def nested_agent() -> None: + return None + + await run_agent("shallow_research_agent", nested_agent) + return {"turn": turn_number} + + try: + await ensure_started(config) + await run_workflow( + "", + lambda: run_agent( + "chat_deepresearcher_agent", + lambda: turn(1), + input_value={"question": "question-1"}, + ), + session_id="same-session", + input_value={"question": "question-1"}, + ) + await run_workflow( + "", + lambda: run_agent( + "chat_deepresearcher_agent", + lambda: turn(2), + input_value={"question": "question-2"}, + ), + session_id="same-session", + input_value={"question": "question-2"}, + ) + finally: + await shutdown_async() + server.shutdown() + server.server_close() + server_thread.join() + + events = [json.loads(line) for line in (tmp_path / "two-turns.jsonl").read_text().splitlines()] + scope_events = [event for event in events if event["kind"] == "scope"] + lifecycle_counts = Counter((event["uuid"], event["scope_category"]) for event in scope_events) + scope_uuids = {event["uuid"] for event in scope_events} + assert all(lifecycle_counts[(scope_uuid, phase)] == 1 for scope_uuid in scope_uuids for phase in ("start", "end")) + + root_starts = [ + event for event in scope_events if event["name"] == "" and event["scope_category"] == "start" + ] + assert len(root_starts) == 2 + assert len({event["uuid"] for event in root_starts}) == 2 + assert {event["metadata"]["session_id"] for event in root_starts} == {"same-session"} + assert {event["metadata"]["aiq.framework"] for event in root_starts} == {"nemo-agent-toolkit"} + assert all(event["data"] is not None for event in root_starts) + root_ends = [event for event in scope_events if event["name"] == "" and event["scope_category"] == "end"] + assert all(event["data"] is not None for event in root_ends) + assert sum(event["category"] == "llm" and event["scope_category"] == "start" for event in scope_events) == 2 + assert sum(event["category"] == "tool" and event["scope_category"] == "start" for event in scope_events) == 2 + assert not {"LangGraph", "tools_condition", "should_escalate", "agent", "tools"}.intersection( + event["name"] for event in scope_events + ) + classifier_starts = [ + event + for event in scope_events + if event["category"] == "agent" and event["name"] == "intent_classifier" and event["scope_category"] == "start" + ] + assert len(classifier_starts) == 2 + assert all(event["data"] is not None for event in classifier_starts) + classifier_ends = [ + event + for event in scope_events + if event["category"] == "agent" and event["name"] == "intent_classifier" and event["scope_category"] == "end" + ] + assert all(event["data"] is not None for event in classifier_ends) + root_uuids = {event["uuid"] for event in root_starts} + agent_starts = [ + event + for event in scope_events + if event["category"] == "agent" + and event["name"] == "chat_deepresearcher_agent" + and event["scope_category"] == "start" + ] + assert len(agent_starts) == 2 + assert {event["parent_uuid"] for event in agent_starts} == root_uuids + agent_uuids = {event["uuid"] for event in agent_starts} + classifier_uuids = {event["uuid"] for event in classifier_starts} + assert {event["parent_uuid"] for event in classifier_starts} == agent_uuids + llm_starts = [event for event in scope_events if event["category"] == "llm" and event["scope_category"] == "start"] + assert {event["parent_uuid"] for event in llm_starts} == classifier_uuids + + spans = [] + for body in received: + request = ExportTraceServiceRequest() + request.ParseFromString(body) + spans.extend( + span + for resource_spans in request.resource_spans + for scope_spans in resource_spans.scope_spans + for span in scope_spans.spans + ) + assert len({span.trace_id for span in spans}) == 2 + root_spans = [span for span in spans if span.name == ""] + assert len(root_spans) == 2 + assert all( + {"input.value", "output.value"}.issubset({attribute.key for attribute in span.attributes}) + for span in root_spans + ) + assert all( + any( + attribute.key == "openinference.span.kind" and attribute.value.string_value == "CHAIN" + for attribute in span.attributes + ) + for span in root_spans + ) + llm_spans = [span for span in spans if span.name == "ChatNVIDIA"] + tool_spans = [span for span in spans if span.name == "echo"] + assert len(llm_spans) == 2 + assert len(tool_spans) == 2 + assert all( + {"input.value", "output.value"}.issubset({attribute.key for attribute in span.attributes}) + for span in [*llm_spans, *tool_spans] + ) + session_ids = { + attribute.value.string_value + for span in root_spans + for attribute in span.attributes + if attribute.key == "session.id" + } + assert session_ids == {"same-session"} + + +@pytest.mark.asyncio +async def test_request_scope_closes_on_error_timeout_and_cancellation(tmp_path: Path) -> None: + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "interruptions.jsonl" + config.observability.opentelemetry.enabled = False + + async def fail() -> None: + raise RuntimeError("synthetic failure") + + async def wait_forever(started: asyncio.Event | None = None) -> None: + if started is not None: + started.set() + await asyncio.Event().wait() + + await ensure_started(config) + try: + with pytest.raises(RuntimeError, match="synthetic failure"): + await run_agent("error-agent", fail, session_id="same-session") + with pytest.raises(TimeoutError): + await asyncio.wait_for( + run_agent("timeout-agent", wait_forever, session_id="same-session"), + timeout=0.01, + ) + + started = asyncio.Event() + cancelled = asyncio.create_task( + run_agent("cancelled-agent", lambda: wait_forever(started), session_id="same-session") + ) + await started.wait() + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + + await run_agent("recovery-agent", _completed_operation, session_id="same-session") + finally: + await shutdown_async() + + events = [json.loads(line) for line in (tmp_path / "interruptions.jsonl").read_text().splitlines()] + roots = [event for event in events if event["kind"] == "scope" and event["category"] == "agent"] + counts = Counter((event["uuid"], event["scope_category"]) for event in roots) + assert {event["name"] for event in roots} == { + "error-agent", + "timeout-agent", + "cancelled-agent", + "recovery-agent", + } + assert all(counts[(scope_uuid, phase)] == 1 for scope_uuid, _ in counts for phase in ("start", "end")) + end_status = { + event["name"]: event["metadata"]["otel.status_code"] for event in roots if event["scope_category"] == "end" + } + assert end_status == { + "error-agent": "ERROR", + "timeout-agent": "ERROR", + "cancelled-agent": "ERROR", + "recovery-agent": "OK", + } + + +async def _completed_operation() -> None: + return None + + +@pytest.mark.asyncio +async def test_plugin_managed_otel_exports_protobuf_trace(tmp_path: Path) -> None: + received: list[tuple[str, str | None, bytes]] = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 - stdlib callback name + body = self.rfile.read(int(self.headers["content-length"])) + received.append((self.path, self.headers.get("content-type"), body)) + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + server_thread = Thread(target=server.serve_forever, daemon=True) + server_thread.start() + config = RelayConfig() + config.observability.atof.output_directory = str(tmp_path) + config.observability.opentelemetry.enabled = True + config.observability.opentelemetry.endpoints = [ + RelayOpenTelemetryEndpointConfig( + type=projection, + endpoint=f"http://127.0.0.1:{server.server_port}/v1/traces?projection={projection}", + timeout_millis=1000, + ) + for projection in ("openinference", "full", "gen_ai") + ] + try: + await ensure_started(config) + with nemo_relay.scope.scope("otel-test", nemo_relay.ScopeType.Agent): + pass + await shutdown_async() + finally: + server.shutdown() + server.server_close() + server_thread.join() + + assert len(received) == 3 + assert {path for path, _, _ in received} == { + "/v1/traces?projection=openinference", + "/v1/traces?projection=full", + "/v1/traces?projection=gen_ai", + } + assert all(content_type == "application/x-protobuf" for _, content_type, _ in received) + assert all(body for _, _, body in received) diff --git a/tests/tokenomics/test_atof_adapter.py b/tests/tokenomics/test_atof_adapter.py new file mode 100644 index 000000000..cea05cdfd --- /dev/null +++ b/tests/tokenomics/test_atof_adapter.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Relay ATOF tokenomics post-processing.""" + +import json +from pathlib import Path + +import pytest + +from aiq_agent.tokenomics.atof_adapter import parse_trace +from aiq_agent.tokenomics.pricing import PricingRegistry +from aiq_agent.tokenomics.profile import PHASE_ORCHESTRATOR +from aiq_agent.tokenomics.profile import PHASE_PLANNER +from aiq_agent.tokenomics.profile import PHASE_RESEARCHER + + +def _pricing() -> PricingRegistry: + return PricingRegistry.from_dict( + { + "models": { + "test-model": { + "input_per_1m_tokens": 1.0, + "cached_input_per_1m_tokens": 0.5, + "output_per_1m_tokens": 2.0, + } + }, + "tools": {"search": {"cost_per_call": 0.01}}, + } + ) + + +def _scope( + uuid: str, + category: str, + name: str, + phase: str, + timestamp: str, + *, + parent_uuid: str = "ambient", + data: object = None, + metadata: dict | None = None, + category_profile: dict | None = None, +) -> dict: + return { + "atof_version": "0.1", + "kind": "scope", + "uuid": uuid, + "parent_uuid": parent_uuid, + "category": category, + "name": name, + "scope_category": phase, + "timestamp": timestamp, + "data": data, + "metadata": metadata, + "category_profile": category_profile, + } + + +def _write(path: Path, events: list[dict]) -> None: + path.write_text("".join(f"{json.dumps(event)}\n" for event in events), encoding="utf-8") + + +def test_parse_trace_uses_real_agent_ancestry_and_relay_cost(tmp_path: Path) -> None: + root_metadata = {"aiq.component.type": "workflow", "session_id": "session-1"} + usage = { + "annotated_response": { + "model": "test-model", + "usage": { + "prompt_tokens": 100, + "cached_tokens": 20, + "completion_tokens": 40, + "reasoning_tokens": 7, + "cost": {"total": 0.25}, + }, + } + } + events = [ + _scope( + "root", + "function", + "workflow", + "start", + "2026-01-01T00:00:00Z", + data={"query": "why?"}, + metadata=root_metadata, + ), + _scope("orch", "llm", "test-model", "start", "2026-01-01T00:00:01Z", parent_uuid="root"), + _scope("orch", "llm", "test-model", "end", "2026-01-01T00:00:02Z", parent_uuid="root", category_profile=usage), + _scope("planner", "agent", "planner-agent", "start", "2026-01-01T00:00:02Z", parent_uuid="root"), + _scope("plan-llm", "llm", "test-model", "start", "2026-01-01T00:00:03Z", parent_uuid="planner"), + _scope( + "plan-llm", + "llm", + "test-model", + "end", + "2026-01-01T00:00:04Z", + parent_uuid="planner", + category_profile=usage, + ), + _scope("planner", "agent", "planner-agent", "end", "2026-01-01T00:00:04Z", parent_uuid="root"), + _scope("researcher", "agent", "researcher-agent", "start", "2026-01-01T00:00:04Z", parent_uuid="root"), + _scope("research-llm", "llm", "test-model", "start", "2026-01-01T00:00:05Z", parent_uuid="researcher"), + _scope( + "research-llm", + "llm", + "test-model", + "end", + "2026-01-01T00:00:06Z", + parent_uuid="researcher", + category_profile=usage, + ), + _scope("search", "tool", "search", "start", "2026-01-01T00:00:06Z", parent_uuid="researcher"), + _scope("search", "tool", "search", "end", "2026-01-01T00:00:07Z", parent_uuid="researcher"), + _scope("researcher", "agent", "researcher-agent", "end", "2026-01-01T00:00:07Z", parent_uuid="root"), + _scope("root", "function", "workflow", "end", "2026-01-01T00:00:08Z", metadata=root_metadata), + ] + path = tmp_path / "relay.atof.jsonl" + _write(path, events) + + profiles = parse_trace(str(path), _pricing()) + + assert len(profiles) == 1 + profile = profiles[0] + assert profile.question == "why?" + assert profile.duration_s == pytest.approx(8.0) + assert profile.total_llm_calls == 3 + assert profile.total_cost_usd == pytest.approx(0.75) + assert profile.total_tool_cost_usd == pytest.approx(0.01) + assert {event["phase"] for event in profile.llm_call_events} == { + PHASE_ORCHESTRATOR, + PHASE_PLANNER, + PHASE_RESEARCHER, + } + assert profile.llm_call_events[0]["reasoning"] == 7 + + +def test_parse_trace_skips_invalid_json_and_uses_catalog_fallback(tmp_path: Path) -> None: + usage = {"annotated_response": {"model": "test-model", "usage": {"input_tokens": 100, "output_tokens": 50}}} + events = [ + _scope( + "root", + "function", + "workflow", + "start", + "2026-01-01T00:00:00Z", + metadata={"aiq.component.type": "workflow"}, + ), + _scope("llm", "llm", "test-model", "start", "2026-01-01T00:00:01Z", parent_uuid="root"), + _scope("llm", "llm", "test-model", "end", "2026-01-01T00:00:02Z", parent_uuid="root", category_profile=usage), + _scope("root", "function", "workflow", "end", "2026-01-01T00:00:03Z"), + ] + path = tmp_path / "relay.atof.jsonl" + _write(path, events) + with path.open("a", encoding="utf-8") as stream: + stream.write("not-json\n") + + profile = parse_trace(str(path), _pricing())[0] + + assert profile.total_prompt_tokens == 100 + assert profile.total_completion_tokens == 50 + assert profile.total_cost_usd == pytest.approx(0.0002) diff --git a/tests/tokenomics/test_nat_adapter.py b/tests/tokenomics/test_nat_adapter.py deleted file mode 100644 index 467467213..000000000 --- a/tests/tokenomics/test_nat_adapter.py +++ /dev/null @@ -1,205 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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. - -"""Tests for NAT trace parsing and phase inference. - -Module under test: src/aiq_agent/tokenomics/nat_adapter.py -""" - -import json -from pathlib import Path - -import pytest - -from aiq_agent.tokenomics.nat_adapter import _build_task_windows -from aiq_agent.tokenomics.nat_adapter import _extract_subagent_type -from aiq_agent.tokenomics.nat_adapter import _infer_phase -from aiq_agent.tokenomics.nat_adapter import _TaskWindow -from aiq_agent.tokenomics.nat_adapter import parse_trace -from aiq_agent.tokenomics.pricing import PricingRegistry -from aiq_agent.tokenomics.profile import PHASE_ORCHESTRATOR -from aiq_agent.tokenomics.profile import PHASE_PLANNER -from aiq_agent.tokenomics.profile import PHASE_RESEARCHER - - -def _payload( - event_type: str, - ts: float, - uuid: str = "step-uuid", - **kwargs: object, -) -> dict: - p: dict = {"event_type": event_type, "event_timestamp": ts, "UUID": uuid} - p.update(kwargs) - return {"payload": p} - - -@pytest.mark.parametrize( - "raw,expected", - [ - ({"subagent_type": "planner-agent"}, "planner-agent"), - ({"subagent_type": "researcher-agent"}, "researcher-agent"), - ("{'subagent_type': 'planner-agent', 'description': 'x'}", "planner-agent"), - ("malformed but researcher-agent string", "researcher-agent"), - ], -) -def test_extract_subagent_type(raw, expected): - assert _extract_subagent_type(raw) == expected - - -def test_extract_subagent_type_none(): - assert _extract_subagent_type(None) is None - assert _extract_subagent_type({}) is None - assert _extract_subagent_type("no marker here") is None - - -def test_build_task_windows_closes_pairs(): - steps = [ - _payload( - "TOOL_START", - 10.0, - uuid="t1", - name="task", - data={"input": {"subagent_type": "planner-agent"}}, - ), - _payload("TOOL_END", 20.0, uuid="t1", name="task"), - ] - wins = _build_task_windows(steps) - assert len(wins) == 1 - assert wins[0].subagent_type == "planner-agent" - assert wins[0].start_ts == 10.0 - assert wins[0].end_ts == 20.0 - - -def test_build_task_windows_string_input(): - steps = [ - _payload( - "TOOL_START", - 1.0, - uuid="u", - name="task", - data={"input": "{'subagent_type': 'researcher-agent'}"}, - ), - _payload("TOOL_END", 2.0, uuid="u", name="task"), - ] - wins = _build_task_windows(steps) - assert len(wins) == 1 - assert wins[0].phase == PHASE_RESEARCHER - - -def test_infer_phase_orchestrator_outside_windows(): - wins = [_TaskWindow(uuid="a", subagent_type="planner-agent", start_ts=10.0, end_ts=20.0)] - assert _infer_phase(5.0, wins) == PHASE_ORCHESTRATOR - assert _infer_phase(25.0, wins) == PHASE_ORCHESTRATOR - - -def test_infer_phase_inside_window(): - wins = [_TaskWindow(uuid="a", subagent_type="planner-agent", start_ts=10.0, end_ts=20.0)] - assert _infer_phase(15.0, wins) == PHASE_PLANNER - - -def test_infer_phase_first_match_wins_on_overlap(): - planner = _TaskWindow(uuid="p", subagent_type="planner-agent", start_ts=10.0, end_ts=25.0) - researcher = _TaskWindow(uuid="r", subagent_type="researcher-agent", start_ts=15.0, end_ts=30.0) - ts = 18.0 - assert _infer_phase(ts, [planner, researcher]) == PHASE_PLANNER - assert _infer_phase(ts, [researcher, planner]) == PHASE_RESEARCHER - - -def _minimal_pricing() -> PricingRegistry: - return PricingRegistry.from_dict( - { - "models": { - "test-model": { - "input_per_1m_tokens": 1.0, - "output_per_1m_tokens": 2.0, - }, - }, - "default": {"input_per_1m_tokens": 1.0, "output_per_1m_tokens": 2.0}, - "tools": {}, - } - ) - - -def _llm_end(ts: float, uuid: str, span_ts: float | None = None) -> dict: - body = { - "event_type": "LLM_END", - "event_timestamp": ts, - "UUID": uuid, - "name": "test-model", - "usage_info": { - "token_usage": { - "prompt_tokens": 1000, - "cached_tokens": 0, - "completion_tokens": 500, - }, - }, - } - if span_ts is not None: - body["span_event_timestamp"] = span_ts - return {"payload": body} - - -def test_parse_trace_end_to_end(tmp_path: Path): - """Orchestrator LLM outside task; planner LLM inside task window; one tool call.""" - steps = [ - _payload("WORKFLOW_START", 100.0, uuid="w0", data={"input": "my question?"}), - _llm_end(101.0, "l0", span_ts=100.5), - _payload( - "TOOL_START", - 102.0, - uuid="task1", - name="task", - data={"input": {"subagent_type": "planner-agent"}}, - ), - _llm_end(103.0, "l1", span_ts=102.5), - _payload("TOOL_START", 103.5, uuid="tool-a", name="search_tool"), - _payload("TOOL_END", 104.0, uuid="tool-a", name="search_tool"), - _payload("TOOL_END", 105.0, uuid="task1", name="task"), - _payload("WORKFLOW_END", 106.0, uuid="w1"), - ] - trace_path = tmp_path / "trace.json" - trace_path.write_text(json.dumps([{"request_number": 0, "intermediate_steps": steps}]), encoding="utf-8") - - profiles = parse_trace(str(trace_path), _minimal_pricing()) - assert len(profiles) == 1 - prof = profiles[0] - assert prof.request_index == 0 - assert prof.question == "my question?" - assert prof.duration_s == pytest.approx(6.0) - assert prof.total_llm_calls == 2 - assert prof.tool_calls.get("search_tool") == 1 - - orch = [p for p in prof.phases if p.phase == PHASE_ORCHESTRATOR] - plan = [p for p in prof.phases if p.phase == PHASE_PLANNER] - assert len(orch) == 1 and orch[0].llm_calls == 1 - assert len(plan) == 1 and plan[0].llm_calls == 1 - - assert prof.llm_call_events[0]["phase"] == PHASE_ORCHESTRATOR - assert prof.llm_call_events[1]["phase"] == PHASE_PLANNER - - -def test_parse_trace_skips_broken_request(tmp_path: Path): - bad = [{"request_number": 0, "intermediate_steps": "not-a-list"}] - good_steps = [ - _payload("WORKFLOW_START", 1.0, data={"input": ""}), - _payload("WORKFLOW_END", 2.0), - ] - good = [{"request_number": 1, "intermediate_steps": good_steps}] - trace_path = tmp_path / "trace.json" - trace_path.write_text(json.dumps(bad + good), encoding="utf-8") - - profiles = parse_trace(str(trace_path), _minimal_pricing()) - assert len(profiles) == 1 - assert profiles[0].request_index == 1 diff --git a/uv.lock b/uv.lock index 51eb8db87..6e953798a 100644 --- a/uv.lock +++ b/uv.lock @@ -211,6 +211,7 @@ dependencies = [ { name = "langgraph-checkpoint-postgres" }, { name = "langgraph-checkpoint-sqlite" }, { name = "mcp" }, + { name = "nemo-relay", extra = ["deepagents", "langchain", "langgraph"] }, { name = "nvidia-nat", extra = ["async-endpoints", "langchain", "mcp", "phoenix"] }, { name = "nvidia-nat-core" }, { name = "nvidia-nat-eval" }, @@ -298,6 +299,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.5.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, + { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], specifier = ">=0.7.3,<0.8" }, { name = "nvidia-nat", extras = ["langchain", "async-endpoints", "phoenix", "mcp"], specifier = "==1.8.0" }, { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, @@ -1528,7 +1530,7 @@ wheels = [ [[package]] name = "deepagents" -version = "0.6.8" +version = "0.6.12" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain" }, @@ -1538,9 +1540,9 @@ dependencies = [ { name = "langsmith" }, { name = "wcmatch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/97/77/e3b7efd9bff9cd101c085a5a3bf74180c13ab6c41a96f725cd1cb1bf53e8/deepagents-0.6.8.tar.gz", hash = "sha256:70cdd4da920cc420a8a0f729792ec559688bbbff39f7ab1508110cce9f901c06", size = 196927, upload-time = "2026-06-03T17:08:36.724Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/db/a6acdc72a9e90c3f07ed10de35c951734a02d4facb693bb59684ad368801/deepagents-0.6.12.tar.gz", hash = "sha256:1f281c0bc5a63132f62e2ee345c1dc593b23188da6e23016401f6879fbe54b5f", size = 211364, upload-time = "2026-06-25T17:26:52.775Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8e/19/1b7b76e958ac7f4e40886edc70f67aff4d7188770ab68105c9c48cbeb769/deepagents-0.6.8-py3-none-any.whl", hash = "sha256:087bdc1458202a3436854cf180f7ec059d07d2114a6c232819e9ad6533a5174a", size = 221469, upload-time = "2026-06-03T17:08:35.133Z" }, + { url = "https://files.pythonhosted.org/packages/98/49/af7219b3c13520fee047bb807cfaefba17f8e4584c551d946773589a4f08/deepagents-0.6.12-py3-none-any.whl", hash = "sha256:28b8fa0119ca0a689e3e18e288c4634e4046062acfc87a1cb34289d3af3a1c88", size = 236120, upload-time = "2026-06-25T17:26:51.736Z" }, ] [[package]] @@ -2803,16 +2805,16 @@ wheels = [ [[package]] name = "langchain" -version = "1.3.11" +version = "1.3.15" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "langgraph" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b9/a2/91a7197c604a3ce1b774b3c10dd114c3c745c6186a304fc2573b3f94d400/langchain-1.3.11.tar.gz", hash = "sha256:f3cf9cd4d2329b1a03eb8fd92b9d73e4e58a4d52570d67725fc77fbe0f104b32", size = 633374, upload-time = "2026-06-22T23:00:33.44Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/1c/b84579174a8e82ed79f4c3e0cd5a7f2323facc5ccd4d1b8390e7d175b663/langchain-1.3.15.tar.gz", hash = "sha256:ab4b775b9703f7e37babe0b325dbbaef25573bda60ecf79f7850bc875f252795", size = 665047, upload-time = "2026-08-11T19:10:52.455Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/df/a4/3a181967294f8876362cc4ba36840d50b8286fa23bb3f5e602b69eb3cb1e/langchain-1.3.11-py3-none-any.whl", hash = "sha256:7ae011f95a09b22feea1e8ae4e43f0b6164aebf4c61b8ad845b45f72ff3a90a2", size = 133639, upload-time = "2026-06-22T23:00:31.619Z" }, + { url = "https://files.pythonhosted.org/packages/09/8d/ae721f4d68ff79a17110cabc9cb39b4568b0e3f1fe0a379b926c3f81d175/langchain-1.3.15-py3-none-any.whl", hash = "sha256:c0d2d0d51ed7da249e8ab7487173872059a9dd46fb071d905957485b7334f987", size = 147001, upload-time = "2026-08-11T19:10:50.846Z" }, ] [[package]] @@ -2887,9 +2889,10 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.4.8" +version = "1.5.6" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "httpx" }, { name = "jsonpatch" }, { name = "langchain-protocol" }, { name = "langsmith" }, @@ -2900,9 +2903,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/e3/bea6d0080acf183332f24dcd74c208aee5857cf8f783c3fb0bd86027d8fb/langchain_core-1.4.8.tar.gz", hash = "sha256:5bf1f8411077c904182ad8f975943d36adcbf579c4e017b3a118b719229ebf9a", size = 957974, upload-time = "2026-06-18T19:39:23.636Z" } +sdist = { url = "https://files.pythonhosted.org/packages/27/b9/893806b89f77e1271fe6e10ce41682ff5fe43d071564e1d8e39dbb5d4d6d/langchain_core-1.5.6.tar.gz", hash = "sha256:b5f73bd9688c457b31ec73657a0ad56948f889fae27acee79286e9c285632ee6", size = 984873, upload-time = "2026-08-17T21:26:35.921Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/13/d6/bdf6f0481cc57ef300d6b1eb48cf1400c0409be715d6eb3cabadd1142a09/langchain_core-1.4.8-py3-none-any.whl", hash = "sha256:d84c28b05e3ba8d4271d0827aad5b592ccdaaf986e76768c23503f0a2045e8aa", size = 557416, upload-time = "2026-06-18T19:39:21.902Z" }, + { url = "https://files.pythonhosted.org/packages/5d/0a/890504397885c9d1ae45f2e06c3000dd9f4445439602b6a990a88b32c0ac/langchain_core-1.5.6-py3-none-any.whl", hash = "sha256:d6cf37bf695ecc22cddeb8461a684e353190b2ce430d99eb22bc11c0c7c00ea5", size = 567016, upload-time = "2026-08-17T21:26:34.595Z" }, ] [[package]] @@ -2920,7 +2923,7 @@ wheels = [ [[package]] name = "langchain-google-genai" -version = "4.2.4" +version = "4.3.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filetype" }, @@ -2928,9 +2931,9 @@ dependencies = [ { name = "langchain-core" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/54/52/de168715eb092c920531d418b8b9aafdff9e37ee80e5fc88106211ccbd47/langchain_google_genai-4.2.4.tar.gz", hash = "sha256:2f5de7a8a6552ffb64b907aca7503fd5e34d1a3240e280abcdc5f7eef480edd5", size = 270054, upload-time = "2026-05-28T21:23:00.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/ae/8ba8ee41bd20a23dee95cda109632c8b19a53141fbc81d9f87a72f0e975c/langchain_google_genai-4.3.4.tar.gz", hash = "sha256:265655baad05f799fa7b83a030eaca7cee0e32c9ab7de846b80ea7f59c26134e", size = 287396, upload-time = "2026-08-14T18:10:00.119Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/a2/5feaf21cfe6fac80eae944f3ac5348d9e5e986813256f74f8dd104617474/langchain_google_genai-4.2.4-py3-none-any.whl", hash = "sha256:0e2c1021a15c91e60b68d813bb3e793bd1d9396b3f8639b943ab4e56e5652e04", size = 68832, upload-time = "2026-05-28T21:22:59.291Z" }, + { url = "https://files.pythonhosted.org/packages/c9/f8/9fe4a28e319e9d6b20454e85fea9c243bf011dc866f2c3b9a89d64f9c1a1/langchain_google_genai-4.3.4-py3-none-any.whl", hash = "sha256:618fb0da1b9ba9def5569a8b05cb87e1389de41b8731c802958d09499490d2a5", size = 73338, upload-time = "2026-08-14T18:09:58.772Z" }, ] [[package]] @@ -3105,7 +3108,7 @@ wheels = [ [[package]] name = "langgraph" -version = "1.2.8" +version = "1.2.11" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, @@ -3115,9 +3118,9 @@ dependencies = [ { name = "pydantic" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b6/ad/583fda4c69501390b989770a465ccd0bdab1c1612eba582c012002ddf9b6/langgraph-1.2.8.tar.gz", hash = "sha256:f79d3575f45b404899358976e4fac0294eb75f8df1bfe8cd11286be7539c4548", size = 722464, upload-time = "2026-07-06T20:40:19.487Z" } +sdist = { url = "https://files.pythonhosted.org/packages/56/0d/c8e7ee98896659e1b6555db0ab115a9ca899844744645d5d894032bab1d7/langgraph-1.2.11.tar.gz", hash = "sha256:9ecfe11e50d338b34b15cf4d8a442642de103e8ae6971320efba84e4542eb363", size = 725753, upload-time = "2026-08-11T14:00:36.945Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/36/49/b958a9963606807e5a20cc75fced14aa77c5cbcc470d5bf8ae13277cd298/langgraph-1.2.8-py3-none-any.whl", hash = "sha256:aa8de1d4df44162353d117589ae0bf6930ca009b62d2d6e26cc32580794c5be6", size = 246983, upload-time = "2026-07-06T20:40:18.242Z" }, + { url = "https://files.pythonhosted.org/packages/0a/7f/c5c30e4be99ff821029c7ac872a480676bb179c9f3df85ea3f38d13f86d4/langgraph-1.2.11-py3-none-any.whl", hash = "sha256:8bab70de7b2d00b5300fb289bcf38d8b241400f3184c1e95e8ce706fb0e8686b", size = 248854, upload-time = "2026-08-11T14:00:35.494Z" }, ] [[package]] @@ -4050,6 +4053,46 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/ca/36339329c4604adbcc99c899b7eb1ce1a555c499b6a6860757dc9bfed36d/narwhals-2.22.1-py3-none-any.whl", hash = "sha256:60567d774edf77db53906f89d9fbd164e66e56d66d388e1e6990f17ac33cfb53", size = 454815, upload-time = "2026-06-05T12:34:32.289Z" }, ] +[[package]] +name = "nemo-relay" +version = "0.7.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/a259aac8df4aa78c0b3a6f3ad0fbf6305666bfaad5c29d9adc26db0f9e27/nemo_relay-0.7.3.tar.gz", hash = "sha256:ea5a1bb52e25e001dcbf6af1830616be181845e978cc848df58562556bba5604", size = 1299430, upload-time = "2026-08-14T14:40:44.832Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/9e/4eb80d2307cadcbb839dad2212a8d888667cd8c531ad0d8c0c1afc841181/nemo_relay-0.7.3-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:301dfc8334032ac52c09cc0b1421181e7c4df601abe7c82dfe51aa74ddbb3732", size = 9251369, upload-time = "2026-08-14T14:40:07.668Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a9/9fb77f7142b1381d8c3c81fdbb76782a02dcff1ee403ac6c93b13fa46b24/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:49c70c0a94cebb6cba3dd7521be11f63eac4cd9386881eb29827ac91b1bd780b", size = 8458046, upload-time = "2026-08-14T14:40:10.181Z" }, + { url = "https://files.pythonhosted.org/packages/ae/ab/b2f246f971f561a982d234d9f3ec1b29dfa4c72bf4882f3e32ce6db54dfa/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8eeaaf8c6a18440e473dd14eb1bb82e56c6514e70adee7456eddf9ce217cff89", size = 8957931, upload-time = "2026-08-14T14:40:12.547Z" }, + { url = "https://files.pythonhosted.org/packages/8f/d8/d8c25ba915467bab457d175b84a3af5c33291655867d472c49eada3b77e6/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:90ff984a89c42ebd0cfe26a0af3f180b1970e2ebcd743d19a960e547949ad2ef", size = 10325640, upload-time = "2026-08-14T14:40:15.722Z" }, + { url = "https://files.pythonhosted.org/packages/7d/fd/6ec48f47eb5ca4cca566b197ade8514baf183dddf53d6de3a296b8bc1102/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d4895a4df39a92ecbac3633fb4ec59cf1f83286c0369f433760ab28fb8d6dddf", size = 10706736, upload-time = "2026-08-14T14:40:18.377Z" }, + { url = "https://files.pythonhosted.org/packages/6b/cb/8a6f8d5f9922e75100f135c1dc4451bafc15f5d51608081f72b894e0caf4/nemo_relay-0.7.3-cp311-abi3-win_amd64.whl", hash = "sha256:f123cd45a27fca3d570559f2c26850138763411f66f387559d757e5d88bfa3c1", size = 8807604, upload-time = "2026-08-14T14:40:20.98Z" }, + { url = "https://files.pythonhosted.org/packages/f1/ac/558a05e6b9e28464b64fee8327d9e8a2aab3fb356363b8b82597b6b35b2a/nemo_relay-0.7.3-cp311-abi3-win_arm64.whl", hash = "sha256:6d5444e9a03b8b5d4bba2409105a121349580cb51481d441c6b5e699713a8763", size = 8442944, upload-time = "2026-08-14T14:40:23.433Z" }, +] + +[package.optional-dependencies] +deepagents = [ + { name = "deepagents" }, + { name = "langchain" }, + { name = "langchain-anthropic" }, + { name = "langchain-core" }, + { name = "langgraph" }, + { name = "langgraph-checkpoint" }, + { name = "langsmith" }, +] +langchain = [ + { name = "langchain" }, + { name = "langchain-core" }, + { name = "langgraph" }, + { name = "langgraph-checkpoint" }, + { name = "langsmith" }, +] +langgraph = [ + { name = "langchain" }, + { name = "langchain-core" }, + { name = "langgraph" }, + { name = "langgraph-checkpoint" }, + { name = "langsmith" }, +] + [[package]] name = "nemoguardrails" version = "0.21.0" From 63564d0eeb0f448b43200ba00f72a5227c4e2b59 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Tue, 18 Aug 2026 22:19:34 -0700 Subject: [PATCH 02/13] remove empty relay pricing from configs Signed-off-by: Chantal D Gama Rose --- .../assets/config-scaffold.yml | 3 --- .../references/composing-config.md | 6 ++--- configs/config_cli_default.yml | 26 +++++++++---------- configs/config_domain_routing_and_skills.yml | 3 --- configs/config_frontier_models.yml | 3 --- configs/config_mcp.yml | 3 --- configs/config_openshell.yml | 5 ---- configs/config_web_azure_ai_search.yml | 3 --- configs/config_web_default_guardrails.yml | 3 --- configs/config_web_default_llamaindex.yml | 2 -- configs/config_web_frag.yml | 3 --- configs/config_web_frag_mcp_auth.yml | 3 --- configs/config_web_opensearch.yml | 3 --- .../customization/configuration-reference.md | 6 ++--- docs/source/deployment/observability.md | 16 ++---------- .../configs/config_deep_research_bench.yml | 4 --- .../config_deep_research_bench_profiling.yml | 4 --- .../configs/config_deepsearch_qa.yml | 4 --- .../freshqa/configs/config_full_workflow.yml | 4 --- .../configs/config_shallow_research_only.yml | 4 --- 20 files changed, 19 insertions(+), 89 deletions(-) diff --git a/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml b/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml index be8b93381..a987d763c 100644 --- a/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml +++ b/.agents/skills/aiq-configure-workflow/assets/config-scaffold.yml @@ -91,6 +91,3 @@ workflow: enable_escalation: true enable_clarifier: false # set true and add clarifier_agent under functions: to enable checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] diff --git a/.agents/skills/aiq-configure-workflow/references/composing-config.md b/.agents/skills/aiq-configure-workflow/references/composing-config.md index e49bda3cf..f10ee1958 100644 --- a/.agents/skills/aiq-configure-workflow/references/composing-config.md +++ b/.agents/skills/aiq-configure-workflow/references/composing-config.md @@ -56,9 +56,9 @@ general: **Observability** — configure NeMo Relay under `workflow.relay`. Relay logging, ATOF, and redaction are enabled by default. OTEL is opt-in; uncomment the Relay -OpenInference endpoint in a default config to send traces to Phoenix. Keep -`workflow.relay.pricing.sources: []` unless the workflow intentionally loads an -audited catalog. See `docs/source/deployment/observability.md`. +OpenInference endpoint in a default config to send traces to Phoenix. Omit +`workflow.relay.pricing` unless the workflow intentionally loads an audited +catalog. See `docs/source/deployment/observability.md`. `workflow.relay.logging` controls the console subscriber; agent and workflow configs do not have separate verbose switches. diff --git a/configs/config_cli_default.yml b/configs/config_cli_default.yml index 0c6c7f771..c30aafd47 100644 --- a/configs/config_cli_default.yml +++ b/configs/config_cli_default.yml @@ -148,17 +148,15 @@ workflow: enable_escalation: true enable_clarifier: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - # Uncomment to export Relay traces to a local Phoenix instance. - # observability: - # opentelemetry: - # enabled: true - # endpoints: - # - type: openinference - # endpoint: http://localhost:6006/v1/traces - # service_name: aiq-relay - # resource_attributes: - # openinference.project.name: aiq-relay - # deployment.environment: development - pricing: - sources: [] + # Uncomment to export Relay traces to a local Phoenix instance. + # relay: + # observability: + # opentelemetry: + # enabled: true + # endpoints: + # - type: openinference + # endpoint: http://localhost:6006/v1/traces + # service_name: aiq-relay + # resource_attributes: + # openinference.project.name: aiq-relay + # deployment.environment: development diff --git a/configs/config_domain_routing_and_skills.yml b/configs/config_domain_routing_and_skills.yml index 1579fa76f..54d118c02 100644 --- a/configs/config_domain_routing_and_skills.yml +++ b/configs/config_domain_routing_and_skills.yml @@ -219,6 +219,3 @@ functions: workflow: _type: deep_research_workflow use_async_deep_research: true - relay: - pricing: - sources: [] diff --git a/configs/config_frontier_models.yml b/configs/config_frontier_models.yml index 131a60776..2135f2a7a 100644 --- a/configs/config_frontier_models.yml +++ b/configs/config_frontier_models.yml @@ -171,6 +171,3 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] diff --git a/configs/config_mcp.yml b/configs/config_mcp.yml index fec4d41c4..269bd176c 100644 --- a/configs/config_mcp.yml +++ b/configs/config_mcp.yml @@ -127,6 +127,3 @@ workflow: enable_escalation: true use_async_deep_research: false checkpoint_db: ${AIQ_CHECKPOINT_DB} - relay: - pricing: - sources: [] diff --git a/configs/config_openshell.yml b/configs/config_openshell.yml index 4b059877f..1ed09f543 100644 --- a/configs/config_openshell.yml +++ b/configs/config_openshell.yml @@ -196,8 +196,3 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - # Add file or inline catalog sources to emit cost estimates. Token usage - # is still captured when no catalog matches the configured model. - sources: [] diff --git a/configs/config_web_azure_ai_search.yml b/configs/config_web_azure_ai_search.yml index 709728507..b4f92032d 100644 --- a/configs/config_web_azure_ai_search.yml +++ b/configs/config_web_azure_ai_search.yml @@ -222,6 +222,3 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] diff --git a/configs/config_web_default_guardrails.yml b/configs/config_web_default_guardrails.yml index b5daa37b3..f478196d8 100644 --- a/configs/config_web_default_guardrails.yml +++ b/configs/config_web_default_guardrails.yml @@ -255,8 +255,5 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] middleware: - workflow_guardrails diff --git a/configs/config_web_default_llamaindex.yml b/configs/config_web_default_llamaindex.yml index ea144f68c..f3d184718 100644 --- a/configs/config_web_default_llamaindex.yml +++ b/configs/config_web_default_llamaindex.yml @@ -238,5 +238,3 @@ workflow: # resource_attributes: # openinference.project.name: aiq-relay # deployment.environment: development - pricing: - sources: [] diff --git a/configs/config_web_frag.yml b/configs/config_web_frag.yml index b291836d5..9843168fc 100644 --- a/configs/config_web_frag.yml +++ b/configs/config_web_frag.yml @@ -191,6 +191,3 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] diff --git a/configs/config_web_frag_mcp_auth.yml b/configs/config_web_frag_mcp_auth.yml index 034fd6261..a64cc6451 100644 --- a/configs/config_web_frag_mcp_auth.yml +++ b/configs/config_web_frag_mcp_auth.yml @@ -310,6 +310,3 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] diff --git a/configs/config_web_opensearch.yml b/configs/config_web_opensearch.yml index 40d275251..3735fd637 100644 --- a/configs/config_web_opensearch.yml +++ b/configs/config_web_opensearch.yml @@ -198,6 +198,3 @@ workflow: enable_clarifier: true use_async_deep_research: true checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} - relay: - pricing: - sources: [] diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index b854ccb2c..8bb49c43c 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -579,12 +579,10 @@ workflow: redaction: enabled: true request_privacy_attributes: [data, category_profile] - pricing: - enabled: true - sources: [] ``` -Default configs do not load a pricing catalog. The dedicated +The default pricing source list is empty, so default configs omit the pricing +block and do not load a catalog. The dedicated `configs/nemo_relay/config_web_default_with_pricing.yml` example loads deployment-specific rates from `configs/nemo_relay/relay_pricing_catalog.json`. Its zero-dollar Nemotron entries describe the NVIDIA-hosted access path used by diff --git a/docs/source/deployment/observability.md b/docs/source/deployment/observability.md index 1be0ad684..bb8b60740 100644 --- a/docs/source/deployment/observability.md +++ b/docs/source/deployment/observability.md @@ -66,9 +66,6 @@ workflow: enabled: false redaction: enabled: true - pricing: - enabled: true - sources: [] ``` Most users can omit this block and use the defaults. Add only the settings that @@ -315,17 +312,8 @@ full payloads in a production environment. ## Pricing and cost analysis -Default AI-Q configs deliberately use an empty Relay pricing source list: - -```yaml -workflow: - relay: - pricing: - enabled: true - sources: [] -``` - -This records token usage without claiming a monetary cost. Pricing depends on +Default AI-Q configs omit the Relay pricing block. The resulting empty source +list records token usage without claiming a monetary cost. Pricing depends on the provider, deployment, contract, region, cache policy, and date. Use the dedicated example when you want model cost enrichment: diff --git a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml index 13ee1011d..795fababd 100644 --- a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml +++ b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench.yml @@ -59,10 +59,6 @@ functions: workflow: _type: deep_research_workflow - relay: - pricing: - sources: [] - eval: general: workflow_alias: "aiq-deepresearcher" diff --git a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml index 96f33e543..eb2ca79c4 100644 --- a/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml +++ b/frontends/benchmarks/deepresearch_bench/configs/config_deep_research_bench_profiling.yml @@ -59,10 +59,6 @@ functions: workflow: _type: deep_research_workflow - relay: - pricing: - sources: [] - eval: general: workflow_alias: "aiq-deepresearcher" diff --git a/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml b/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml index dace73cf4..bded59a8a 100644 --- a/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml +++ b/frontends/benchmarks/deepsearch_qa/configs/config_deepsearch_qa.yml @@ -60,10 +60,6 @@ functions: workflow: _type: deep_research_workflow - relay: - pricing: - sources: [] - eval: general: output_dir: frontends/benchmarks/deepsearch_qa/results diff --git a/frontends/benchmarks/freshqa/configs/config_full_workflow.yml b/frontends/benchmarks/freshqa/configs/config_full_workflow.yml index d63fb0583..270358913 100644 --- a/frontends/benchmarks/freshqa/configs/config_full_workflow.yml +++ b/frontends/benchmarks/freshqa/configs/config_full_workflow.yml @@ -110,10 +110,6 @@ workflow: enable_escalation: true tools: - web_search_tool - relay: - pricing: - sources: [] - eval: general: workflow_alias: "freshqa_eval_full_workflow_with_intent_classifier_and_depth_router" diff --git a/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml b/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml index e278c1b66..4d31e1054 100644 --- a/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml +++ b/frontends/benchmarks/freshqa/configs/config_shallow_research_only.yml @@ -53,10 +53,6 @@ functions: workflow: _type: shallow_research_workflow - relay: - pricing: - sources: [] - eval: general: workflow_alias: "freshqa_eval_shallow_research_only" From bb93ed832f6ab84ad71763fefe4d973bf8fdddaf Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Tue, 18 Aug 2026 22:26:14 -0700 Subject: [PATCH 03/13] fix mcp tests Signed-off-by: Chantal D Gama Rose --- mcp/tests/test_config_and_packaging.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/mcp/tests/test_config_and_packaging.py b/mcp/tests/test_config_and_packaging.py index 1f788a86e..4136ca95c 100644 --- a/mcp/tests/test_config_and_packaging.py +++ b/mcp/tests/test_config_and_packaging.py @@ -80,13 +80,12 @@ def test_public_mcp_config_preserves_reference_orchestration_choices() -> None: "web_search_tool", } assert functions["intent_classifier"]["_type"] == "intent_classifier" - assert functions["clarifier_agent"] == { - "_type": "clarifier_agent", - "llm": "nemotron_ultra_llm", - "max_turns": 3, - "log_response_max_chars": 2000, - "verbose": True, - } + clarifier = functions["clarifier_agent"] + assert clarifier["_type"] == "clarifier_agent" + assert clarifier["llm"] == "nemotron_ultra_llm" + assert clarifier["max_turns"] == 3 + assert clarifier["log_response_max_chars"] == 2000 + assert "verbose" not in clarifier assert functions["shallow_research_agent"]["exclude_tools"] == ["advanced_web_search_tool"] assert functions["deep_research_agent"]["exclude_tools"] == ["web_search_tool"] @@ -105,8 +104,10 @@ def test_public_mcp_config_uses_only_public_models_sources_and_environment_names config = yaml.safe_load(_CONFIG_PATH.read_text()) text = _CONFIG_PATH.read_text().lower() - assert {entry["_type"] for entry in config["llms"].values()} == {"nim"} + llms = config["llms"].values() + assert {entry["_type"] for entry in llms} == {"openai"} assert {entry["base_url"] for entry in config["llms"].values()} == {"https://integrate.api.nvidia.com/v1"} + assert all(entry["model_name"].startswith("nvidia/") for entry in config["llms"].values()) assert config["functions"]["web_search_tool"]["_type"] == "tavily_web_search" assert config["functions"]["advanced_web_search_tool"]["_type"] == "tavily_web_search" assert "nvidia_api_key" in text From ac2d18574102db6e321697170c75bb1b7aa2fa5b Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Aug 2026 09:48:36 -0700 Subject: [PATCH 04/13] fix: CI failures and tests Signed-off-by: Chantal D Gama Rose --- README.md | 2 +- configs/config_cli_default.yml | 2 +- .../config_web_default_with_pricing.yml | 2 +- .../customization/configuration-reference.md | 2 +- docs/source/resources/troubleshooting.md | 3 +- frontends/aiq_api/pyproject.toml | 1 + frontends/aiq_api/src/aiq_api/jobs/runner.py | 15 +- frontends/aiq_api/src/aiq_api/jobs/submit.py | 8 +- .../benchmarks/deepresearch_bench/README.md | 2 +- mcp/Dockerfile | 3 +- .../agents/deep_researcher/tools/research.py | 21 +- src/aiq_agent/relay/runtime.py | 84 +++++-- src/aiq_agent/tokenomics/atof_adapter.py | 40 +++- .../agents/deep_researcher/test_agent.py | 3 +- .../deep_researcher/test_custom_middleware.py | 10 +- .../test_deepagents_runtime.py | 119 ---------- tests/aiq_agent/jobs/test_runner.py | 22 +- .../aiq_agent/test_default_model_profiles.py | 211 ------------------ tests/conftest.py | 22 -- tests/test_relay_runtime.py | 120 +++++++++- tests/tokenomics/test_atof_adapter.py | 46 ++++ 21 files changed, 324 insertions(+), 414 deletions(-) delete mode 100644 tests/aiq_agent/test_default_model_profiles.py delete mode 100644 tests/conftest.py diff --git a/README.md b/README.md index 1b8300e3e..6ac29531e 100644 --- a/README.md +++ b/README.md @@ -72,7 +72,7 @@ The NVIDIA AI-Q Blueprint is a deployable research backend built on the [NVIDIA - **Expanded sources** — Paper search supports Serper, SerpAPI, and SearchAPI; You.com adds web, contents, general-research, and finance-research tools; Nimble adds configurable web search; focused profiles demonstrate DuckDuckGo news, Polymarket, OpenSearch, and Azure AI Search knowledge retrieval. - **Production API and auth** — REST endpoints, async job ownership, per-user OAuth-protected MCP sources, token validator entry points, and provider lifecycle hooks support authenticated deployments; a separate public MCP server exposes stateless research tools for trusted networks. - **Opt-in policy controls** — NeMo Guardrails middleware covers selected workflow and agent boundaries, and narrow application-level encryption can protect final async output plus selected artifact-event content. -- **Observability, profiling, and cost analysis** — NeMo Relay preserves task, named-agent, LLM, and tool hierarchy across interactive turns and async researchers. ATOF and OTEL exports feed debugging and tokenomics reports for cost, latency, and cache analysis. +- **Observability, profiling, and cost analysis** — NeMo Relay preserves task, named-agent, LLM, and tool hierarchy across interactive turns and async researchers. ATOF feeds local debugging and tokenomics reports; OTEL exports traces to external observability backends. - **Evaluation harnesses** — Built-in benchmarks (for example, FreshQA, DeepResearch) and evaluation scripts to measure quality and iterate on prompts and agent architecture. - **Frontend options** — Run through CLI, web UI, or async jobs. Refer to [Getting started](#getting-started) and [Ways to run the agents](#ways-to-run-the-agents). - **Deployment options** - Deployment assets for [Docker Compose](deploy/compose/) and [Helm](deploy/helm/deployment-k8s/); the repository source chart honors the Helm release namespace for every namespaced resource. diff --git a/configs/config_cli_default.yml b/configs/config_cli_default.yml index c30aafd47..1a6609362 100644 --- a/configs/config_cli_default.yml +++ b/configs/config_cli_default.yml @@ -155,7 +155,7 @@ workflow: # enabled: true # endpoints: # - type: openinference - # endpoint: http://localhost:6006/v1/traces + # endpoint: ${RELAY_OTEL_ENDPOINT:-http://localhost:6006/v1/traces} # service_name: aiq-relay # resource_attributes: # openinference.project.name: aiq-relay diff --git a/configs/nemo_relay/config_web_default_with_pricing.yml b/configs/nemo_relay/config_web_default_with_pricing.yml index 4198cdfb4..64c1f4606 100644 --- a/configs/nemo_relay/config_web_default_with_pricing.yml +++ b/configs/nemo_relay/config_web_default_with_pricing.yml @@ -238,7 +238,7 @@ workflow: enabled: true endpoints: - type: openinference - endpoint: http://localhost:6006/v1/traces + endpoint: ${RELAY_OTEL_ENDPOINT:-http://localhost:6006/v1/traces} service_name: aiq-relay resource_attributes: openinference.project.name: aiq-relay diff --git a/docs/source/customization/configuration-reference.md b/docs/source/customization/configuration-reference.md index 8bb49c43c..9fc421bfe 100644 --- a/docs/source/customization/configuration-reference.md +++ b/docs/source/customization/configuration-reference.md @@ -573,7 +573,7 @@ workflow: enabled: false endpoints: - type: openinference - endpoint: "http://localhost:6006/v1/traces" + endpoint: "${RELAY_OTEL_ENDPOINT:-http://localhost:6006/v1/traces}" service_name: aiq-relay resource_attributes: {openinference.project.name: aiq-relay} redaction: diff --git a/docs/source/resources/troubleshooting.md b/docs/source/resources/troubleshooting.md index 3ea2c39e7..0edec8f92 100644 --- a/docs/source/resources/troubleshooting.md +++ b/docs/source/resources/troubleshooting.md @@ -183,12 +183,13 @@ workflow: enabled: true endpoints: - type: openinference - endpoint: http://localhost:6006/v1/traces + endpoint: ${RELAY_OTEL_ENDPOINT:-http://localhost:6006/v1/traces} resource_attributes: openinference.project.name: aiq-relay ``` Then open [http://localhost:6006](http://localhost:6006) to inspect traces, token usage, and latency. +Set `RELAY_OTEL_ENDPOINT` to use a remote Phoenix or collector endpoint; local Phoenix is the default. If the trace is missing, also inspect the project configured in `~/.config/nemo-relay/plugins.toml`; Relay can discover an existing user-level Phoenix destination. diff --git a/frontends/aiq_api/pyproject.toml b/frontends/aiq_api/pyproject.toml index 883d67e57..882ab61f3 100644 --- a/frontends/aiq_api/pyproject.toml +++ b/frontends/aiq_api/pyproject.toml @@ -37,6 +37,7 @@ classifiers = [ "Programming Language :: Python :: 3.13", ] dependencies = [ + "aiq-agent", "fastapi>=0.100.0", "dask[distributed]>=2024.1.0", "sqlalchemy>=2.0.0", diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 1c7de7ccc..07e81efeb 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -185,6 +185,17 @@ def check(self) -> None: # well under GHOST_JOB_TIMEOUT_SECONDS so a live worker refreshes several times # before the reaper's timeout. LEASE_REFRESH_INTERVAL_SECONDS = 60 +RELAY_STARTUP_TIMEOUT_SECONDS = 30 + + +async def _ensure_relay_started_for_job(relay_config: Any, job_id: str) -> None: + """Start Relay without allowing observability initialization to stall a job.""" + from aiq_agent.relay.bootstrap import ensure_started + + try: + await asyncio.wait_for(ensure_started(relay_config), timeout=RELAY_STARTUP_TIMEOUT_SECONDS) + except Exception as exc: + logger.warning("Relay startup failed for job %s (error_type=%s)", job_id, type(exc).__name__) def _db_now_expr(db_url: str) -> str: @@ -779,9 +790,7 @@ async def run_agent_job( fn_config = builder.get_function_config(agent_config_name) relay_config = getattr(fn_config, "relay", None) if relay_config is not None: - from aiq_agent.relay.bootstrap import ensure_started as ensure_relay_started - - await ensure_relay_started(relay_config) + await _ensure_relay_started_for_job(relay_config, job_id) if getattr(fn_config, "type", None) == "deep_research_agent": from aiq_agent.agents.deep_researcher.register import DeepResearchAgentConfig from aiq_agent.agents.deep_researcher.register import resolve_deep_research_runtime_config diff --git a/frontends/aiq_api/src/aiq_api/jobs/submit.py b/frontends/aiq_api/src/aiq_api/jobs/submit.py index 0e550f292..2bfa691df 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -27,6 +27,7 @@ import secrets import time from contextlib import suppress +from dataclasses import replace from functools import partial from typing import Any @@ -293,12 +294,7 @@ async def submit_agent_job( ) trace_correlation = _get_job_trace_correlation() if conversation_id is not None: - trace_correlation = JobTraceCorrelation( - session_id=conversation_id, - submission_trace_id=trace_correlation.submission_trace_id, - submission_span_id=trace_correlation.submission_span_id, - request_trace_tags=trace_correlation.request_trace_tags, - ) + trace_correlation = replace(trace_correlation, session_id=conversation_id) submission_conversation_id = trace_correlation.session_id async def _release_submission_reservations() -> None: diff --git a/frontends/benchmarks/deepresearch_bench/README.md b/frontends/benchmarks/deepresearch_bench/README.md index 37573b635..2c324feca 100644 --- a/frontends/benchmarks/deepresearch_bench/README.md +++ b/frontends/benchmarks/deepresearch_bench/README.md @@ -75,7 +75,7 @@ workflow: enabled: true endpoints: - type: openinference - endpoint: http://localhost:6006/v1/traces + endpoint: ${RELAY_OTEL_ENDPOINT:-http://localhost:6006/v1/traces} resource_attributes: openinference.project.name: aiq-deepresearch-bench diff --git a/mcp/Dockerfile b/mcp/Dockerfile index fd46488c8..faa5353fc 100644 --- a/mcp/Dockerfile +++ b/mcp/Dockerfile @@ -65,7 +65,8 @@ RUN apt-get update \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates \ && rm -rf /var/lib/apt/lists/* \ && groupadd --gid 10001 aiq \ - && useradd --uid 10001 --gid 10001 --create-home --home-dir /home/aiq --shell /usr/sbin/nologin aiq + && useradd --uid 10001 --gid 10001 --create-home --home-dir /home/aiq --shell /usr/sbin/nologin aiq \ + && install -d --owner=10001 --group=10001 /app/relay WORKDIR /app diff --git a/src/aiq_agent/agents/deep_researcher/tools/research.py b/src/aiq_agent/agents/deep_researcher/tools/research.py index f490b5c0e..36a18b509 100644 --- a/src/aiq_agent/agents/deep_researcher/tools/research.py +++ b/src/aiq_agent/agents/deep_researcher/tools/research.py @@ -131,7 +131,12 @@ async def _run_research_query( ) return _exhausted_research_notes(query) except Exception as exc: # noqa: BLE001 - captured as per-item failure - raise RuntimeError(f"researcher worker failed for query {query.query!r}: {exc}") from exc + logger.warning( + "Researcher worker failed (error_type=%s, query_%s)", + type(exc).__name__, + log_content_metadata(query.query), + ) + raise RuntimeError("researcher worker failed") from exc try: structured = result.get("structured_response") if isinstance(result, dict) else None @@ -139,9 +144,15 @@ async def _run_research_query( raise ValueError("researcher worker did not return structured ResearchNotes") note = ResearchNotes.model_validate(structured) except Exception as exc: # noqa: BLE001 - captured as per-item failure - raise ValueError( - f"researcher worker returned invalid ResearchNotes for query {query.query!r}: {exc}" - ) from exc + missing_response = "researcher worker did not return structured ResearchNotes" + if isinstance(exc, ValueError) and str(exc) == missing_response: + raise + logger.warning( + "Researcher worker returned invalid ResearchNotes (error_type=%s, query_%s)", + type(exc).__name__, + log_content_metadata(query.query), + ) + raise ValueError("researcher worker returned invalid ResearchNotes") from exc lifecycle.output = note return note @@ -226,7 +237,7 @@ async def _run_research_queries( for query, raw_result in zip(queries, raw_results, strict=False): if isinstance(raw_result, BaseException): error = str(raw_result) or raw_result.__class__.__name__ - errors.append(f"{query.query}: {error}") + errors.append(error) else: successful_queries.append(query) notes.append(raw_result) diff --git a/src/aiq_agent/relay/runtime.py b/src/aiq_agent/relay/runtime.py index 7f2dcd91e..412431b99 100644 --- a/src/aiq_agent/relay/runtime.py +++ b/src/aiq_agent/relay/runtime.py @@ -14,6 +14,7 @@ from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass +from itertools import islice from typing import Any from typing import TypeVar from uuid import uuid4 @@ -35,6 +36,9 @@ _T = TypeVar("_T") _aiq_scope_active: ContextVar[bool] = ContextVar("aiq_relay_scope_active", default=False) logger = logging.getLogger(__name__) +_SAFE_VALUE_MAX_DEPTH = 12 +_SAFE_VALUE_MAX_ITEMS = 100 +_SAFE_VALUE_MAX_STRING_LENGTH = 16_384 @dataclass @@ -154,7 +158,10 @@ async def ainvoke_with_relay( effective_config["callbacks"] = list(configured_callbacks) else: effective_config.pop("callbacks", None) - messages = list(input_value) + if isinstance(input_value, str | BaseMessage): + messages = [input_value] + else: + messages = list(input_value) system_message = ( messages.pop(0) if messages and isinstance(messages[0], BaseMessage) and messages[0].type == "system" else None ) @@ -171,7 +178,7 @@ async def ainvoke_with_relay( model_settings=model_settings, ) - async def invoke(next_request: ModelRequest[Any]) -> ModelResponse[Any]: + async def invoke_call(next_request: ModelRequest[Any]) -> ModelResponse[Any]: next_messages = list(next_request.messages) if next_request.system_message is not None: next_messages.insert(0, next_request.system_message) @@ -199,7 +206,28 @@ async def invoke(next_request: ModelRequest[Any]) -> ModelResponse[Any]: response = AIMessage(content=content) return ModelResponse(result=[response]) - response = await NemoRelayMiddleware().awrap_model_call(request, invoke) + invocation_started = False + invocation_error: BaseException | None = None + + async def invoke(next_request: ModelRequest[Any]) -> ModelResponse[Any]: + nonlocal invocation_error + nonlocal invocation_started + invocation_started = True + try: + return await invoke_call(next_request) + except BaseException as error: + invocation_error = error + raise + + try: + response = await NemoRelayMiddleware().awrap_model_call(request, invoke) + except Exception as error: + if invocation_started: + if invocation_error is not None: + raise invocation_error + raise + _log_capture_failure("model middleware", error) + response = await invoke(request) if not response.result: raise RuntimeError("Relay-managed LangChain model returned no messages") return response.result[-1] @@ -214,12 +242,33 @@ async def ainvoke_tool_with_relay(tool: Any, args: dict[str, Any]) -> Any: runtime=None, ) - async def invoke(next_request: ToolCallRequest) -> Any: + async def invoke_call(next_request: ToolCallRequest) -> Any: if next_request.tool is None: raise RuntimeError(f"Relay-managed tool {next_request.tool_call['name']!r} is unavailable") return await next_request.tool.ainvoke(next_request.tool_call.get("args") or {}) - return await NemoRelayMiddleware().awrap_tool_call(request, invoke) + invocation_started = False + invocation_error: BaseException | None = None + + async def invoke(next_request: ToolCallRequest) -> Any: + nonlocal invocation_error + nonlocal invocation_started + invocation_started = True + try: + return await invoke_call(next_request) + except BaseException as error: + invocation_error = error + raise + + try: + return await NemoRelayMiddleware().awrap_tool_call(request, invoke) + except Exception as error: + if invocation_started: + if invocation_error is not None: + raise invocation_error + raise + _log_capture_failure("tool middleware", error) + return await invoke(request) @contextmanager @@ -262,7 +311,7 @@ def _semantic_scope( status_metadata = { "error_type": type(error).__name__, "otel.status_code": "ERROR", - "otel.status_description": str(error), + "otel.status_description": type(error).__name__, } raise else: @@ -369,16 +418,23 @@ async def _run_isolated() -> _T: return await asyncio.create_task(_run_isolated()) -def _safe_value(value: Any) -> Any: +def _safe_value(value: Any, *, _depth: int = 0) -> Any: """Project framework state to JSON-compatible Relay event values.""" - if value is None or isinstance(value, str | int | float | bool): + if _depth >= _SAFE_VALUE_MAX_DEPTH: + return {"type": type(value).__name__, "truncated": True} + if isinstance(value, str): + return value[:_SAFE_VALUE_MAX_STRING_LENGTH] + if value is None or isinstance(value, int | float | bool): return value - if isinstance(value, dict): - return {str(key): _safe_value(item) for key, item in value.items()} - if isinstance(value, list | tuple | set): - return [_safe_value(item) for item in value] if isinstance(value, BaseMessage): - return messages_to_dict([value])[0] + return _safe_value(messages_to_dict([value])[0], _depth=_depth + 1) if isinstance(value, BaseModel): - return value.model_dump(mode="json") + return _safe_value(value.model_dump(mode="json"), _depth=_depth + 1) + if isinstance(value, dict): + return { + str(key)[:_SAFE_VALUE_MAX_STRING_LENGTH]: _safe_value(item, _depth=_depth + 1) + for key, item in islice(value.items(), _SAFE_VALUE_MAX_ITEMS) + } + if isinstance(value, list | tuple | set): + return [_safe_value(item, _depth=_depth + 1) for item in islice(value, _SAFE_VALUE_MAX_ITEMS)] return {"type": type(value).__name__} diff --git a/src/aiq_agent/tokenomics/atof_adapter.py b/src/aiq_agent/tokenomics/atof_adapter.py index 138df5e3c..04382ac6b 100644 --- a/src/aiq_agent/tokenomics/atof_adapter.py +++ b/src/aiq_agent/tokenomics/atof_adapter.py @@ -21,6 +21,11 @@ logger = logging.getLogger(__name__) +def _event_uuid(event: dict[str, Any]) -> str | None: + value = event.get("uuid") + return value if isinstance(value, str) and value else None + + def _timestamp(value: Any) -> float: if isinstance(value, int | float): return float(value) @@ -105,9 +110,10 @@ def _phase_for(event: dict[str, Any], starts: dict[str, dict[str, Any]]) -> str: def _root_uuid(event: dict[str, Any], starts: dict[str, dict[str, Any]]) -> str | None: - event_uuid = event.get("uuid") - current = event_uuid if event_uuid in starts else event.get("parent_uuid") - if not isinstance(current, str): + event_uuid = _event_uuid(event) + parent_uuid = event.get("parent_uuid") + current = event_uuid if event_uuid in starts else parent_uuid + if not isinstance(current, str) or not current: return None visited: set[str] = set() while current not in visited: @@ -162,8 +168,9 @@ def _parse_request( pricing: PricingRegistry, ) -> RequestProfile: ends = { - str(event.get("uuid")): event + event_uuid: event for event in events + if (event_uuid := _event_uuid(event)) is not None if event.get("kind") == "scope" and event.get("scope_category") == "end" } root_end = ends.get(str(root.get("uuid")), {}) @@ -252,22 +259,31 @@ def parse_trace(path: str, pricing: PricingRegistry) -> list[RequestProfile]: """Parse Relay ATOF JSONL into one profile per workflow root scope.""" events = _load_events(path) starts = { - str(event.get("uuid")): event + event_uuid: event for event in events - if event.get("kind") == "scope" and event.get("scope_category") == "start" and event.get("uuid") + if (event_uuid := _event_uuid(event)) is not None + if event.get("kind") == "scope" and event.get("scope_category") == "start" } explicit_roots = [ event for event in starts.values() if _nested(event, "metadata", "aiq.component.type") == "workflow" ] - roots = explicit_roots or [event for event in starts.values() if event.get("parent_uuid") not in starts] + roots = explicit_roots or [ + event + for event in starts.values() + if not isinstance(event.get("parent_uuid"), str) or event.get("parent_uuid") not in starts + ] roots.sort(key=lambda event: _timestamp(event.get("timestamp"))) profiles: list[RequestProfile] = [] - for root in roots: - root_uuid = str(root["uuid"]) + for request_index, root in enumerate(roots): + root_uuid = root["uuid"] request_events = [event for event in events if _root_uuid(event, starts) == root_uuid] try: - profiles.append(_parse_request(len(profiles), root, request_events, starts, pricing)) - except Exception: - logger.exception("Failed to parse Relay request root %s; skipping", root_uuid) + profiles.append(_parse_request(request_index, root, request_events, starts, pricing)) + except Exception as exc: + logger.warning( + "Failed to parse Relay request; skipping (request_index=%d, error_type=%s)", + request_index, + type(exc).__name__, + ) return profiles diff --git a/tests/aiq_agent/agents/deep_researcher/test_agent.py b/tests/aiq_agent/agents/deep_researcher/test_agent.py index 2e949441a..91a814fc2 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_agent.py +++ b/tests/aiq_agent/agents/deep_researcher/test_agent.py @@ -1360,7 +1360,8 @@ async def ainvoke(self, state, config=None): await batch_tool.ainvoke({"queries": query_payloads}) assert "run_research_batch failed for 1 of 3 researcher worker" in str(exc_info.value) - assert "search backend exploded" in str(exc_info.value) + assert "researcher worker failed" in str(exc_info.value) + assert "search backend exploded" not in str(exc_info.value) assert "timed out" not in str(exc_info.value) assert "2 successful researcher worker(s) were registered and persisted under /shared/" in str(exc_info.value) assert "resubmit only the failed queries" in str(exc_info.value) diff --git a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py index b252d8119..d96f861ac 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py +++ b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py @@ -1029,6 +1029,8 @@ async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): """Sanitize tool names without dropping provider, usage, or response metadata.""" from langchain.agents.middleware.types import ModelResponse + response_metadata = {"model_name": "nvidia/nemotron-3-ultra-550b-a55b", "finish_reason": "tool_calls"} + usage_metadata = {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120} ai_msg = AIMessage( content="", additional_kwargs={ @@ -1044,11 +1046,11 @@ async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): ], "provider_field": "preserve-me", }, - response_metadata={"model_name": "nvidia/nemotron-3-ultra-550b-a55b", "finish_reason": "tool_calls"}, + response_metadata=response_metadata, tool_calls=[ {"name": "advanced_web_search_tool<|channel|>commentary", "args": {"question": "test"}, "id": "tc1"}, ], - usage_metadata={"input_tokens": 100, "output_tokens": 20, "total_tokens": 120}, + usage_metadata=usage_metadata, ) mock_response = ModelResponse(result=[ai_msg]) mock_handler = AsyncMock(return_value=mock_response) @@ -1060,8 +1062,8 @@ async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): assert message.tool_calls[0]["name"] == "advanced_web_search_tool" assert message.additional_kwargs["tool_calls"][0]["function"]["name"] == "advanced_web_search_tool" assert message.additional_kwargs["provider_field"] == "preserve-me" - assert message.response_metadata == ai_msg.response_metadata - assert message.usage_metadata == ai_msg.usage_metadata + assert message.response_metadata == response_metadata + assert message.usage_metadata == usage_metadata @pytest.mark.asyncio async def test_awrap_model_call_no_tool_calls_passthrough(self, middleware): diff --git a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py index 079613681..75b16c978 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py +++ b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py @@ -19,7 +19,6 @@ import logging from contextlib import nullcontext -from pathlib import Path from threading import Event from threading import Thread from typing import Any @@ -27,7 +26,6 @@ from unittest.mock import patch import pytest -import yaml from deepagents.backends import CompositeBackend from deepagents.backends import FilesystemBackend from deepagents.backends import StateBackend @@ -45,123 +43,6 @@ SYNTHESIS_SKILL_SOURCE = f"{BUILTIN_SKILL_SOURCE}synthesis/" -def test_frontier_profile_uses_validated_gpt_role_split() -> None: - """Keep the shipped GPT profile aligned with the validated Sol/Luna topology.""" - config = yaml.safe_load(Path("configs/config_frontier_models.yml").read_text(encoding="utf-8")) - llms = config["llms"] - - assert llms["gpt_luna_intent_llm"] == { - "_type": "openai", - "model_name": "gpt-5.6-luna", - "api_key": "${OPENAI_API_KEY}", - "max_tokens": 1024, - "num_retries": 2, - "parallel_tool_calls": False, - } - assert llms["gpt_luna_shallow_llm"] == { - **llms["gpt_luna_intent_llm"], - "max_tokens": 8192, - } - assert llms["gpt_sol_agent_llm"] == { - "_type": "openai", - "model_name": "gpt-5.6-sol", - "api_key": "${OPENAI_API_KEY}", - "max_tokens": 16384, - "num_retries": 2, - "parallel_tool_calls": False, - } - assert llms["gpt_sol_writer_llm"] == { - **llms["gpt_sol_agent_llm"], - "max_tokens": 32768, - } - assert llms["gpt_luna_agent_llm"] == { - **llms["gpt_sol_agent_llm"], - "model_name": "gpt-5.6-luna", - } - assert config["functions"]["intent_classifier"]["llm"] == "gpt_luna_intent_llm" - assert config["functions"]["shallow_research_agent"]["llm"] == "gpt_luna_shallow_llm" - deep_research = config["functions"]["deep_research_agent"] - assert { - role: deep_research[f"{role}_llm"] - for role in ("orchestrator", "source_router", "researcher", "planner", "writer") - } == { - "orchestrator": "gpt_sol_agent_llm", - "source_router": "gpt_luna_agent_llm", - "researcher": "gpt_luna_agent_llm", - "planner": "gpt_sol_agent_llm", - "writer": "gpt_sol_writer_llm", - } - - -def test_openshell_workflow_only_diverges_for_skills_and_sandbox_wiring() -> None: - """Keep the OpenShell workflow aligned with the standard web config. - - The visualization chart skill now ships only in the skills and sandbox - example configs, so the standard web config wires no deep_research_skills - at all and renders chart-worthy data as Markdown tables. OpenShell layers - the sandbox-gated research and synthesis collections plus the on-demand - visualization skill on top of a sandbox, so the skills function, the - sandbox function, and their two agent refs are the only divergence from - the standard config. - """ - - def load(path: str) -> dict[str, Any]: - text = Path(path).read_text(encoding="utf-8") - text = text.replace("${AIQ_OPENSHELL_REQUIRE_HARD_LANDLOCK:-true}", "true") - return yaml.safe_load(text) - - standard = load("configs/config_web_default_llamaindex.yml") - openshell = load("configs/config_openshell.yml") - - standard_functions = standard["functions"].copy() - openshell_functions = openshell["functions"].copy() - - assert "deep_research_skills" not in standard_functions - assert "deep_research_sandbox" not in standard_functions - assert "skills" not in standard_functions["deep_research_agent"] - assert "sandbox" not in standard_functions["deep_research_agent"] - - openshell_skills = openshell_functions.pop("deep_research_skills") - assert openshell_skills["_type"] == "deep_research_skills" - assert openshell_skills["agents"]["researcher-agent"] == ["research"] - assert "visualization" in openshell_skills["agents"]["writer-agent"] - assert "research" in openshell_skills["require_sandbox"] - openshell_functions.pop("deep_research_sandbox") - - openshell_agent = openshell_functions["deep_research_agent"] = openshell_functions["deep_research_agent"].copy() - assert openshell_agent.pop("skills") == "deep_research_skills" - assert openshell_agent.pop("sandbox") == "deep_research_sandbox" - - assert openshell["general"] == standard["general"] - assert openshell["llms"] == standard["llms"] - assert openshell_functions == standard_functions - assert openshell["workflow"] == standard["workflow"] - - -def test_modal_reference_profile_enables_bounded_artifact_capture() -> None: - """Keep the shipped Modal profile's capture policy validated.""" - config = yaml.safe_load(Path("configs/config_domain_routing_and_skills.yml").read_text(encoding="utf-8")) - sandbox_data = config["functions"]["deep_research_sandbox"].copy() - - assert sandbox_data.pop("_type") == "deep_research_sandbox" - sandbox = DeepResearchSandboxConfig.model_validate(sandbox_data) - - assert sandbox.provider == "modal" - assert sandbox.artifact_capture.enabled is True - assert sandbox.artifact_capture.max_file_bytes == 50_000_000 - assert sandbox.artifact_capture.allow_extensions == ( - ".png", - ".jpg", - ".jpeg", - ".webp", - ".csv", - ".json", - ".md", - ".ipynb", - ".pdf", - ) - - class TestSkillCollections: """Public skill config uses collection names, not DeepAgents virtual paths.""" diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 10bc498c8..021440044 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -63,6 +63,7 @@ - Error message filtering for CancelledError """ +import asyncio import inspect from unittest.mock import AsyncMock from unittest.mock import MagicMock @@ -92,6 +93,22 @@ def fixture_event_store_cache_guard(): EventStore.dispose_all_engines() +@pytest.mark.asyncio +async def test_relay_startup_timeout_does_not_block_job(monkeypatch, caplog): + from aiq_api.jobs import runner + + async def never_starts(_config): + await asyncio.Event().wait() + + monkeypatch.setattr("aiq_agent.relay.bootstrap.ensure_started", never_starts) + monkeypatch.setattr(runner, "RELAY_STARTUP_TIMEOUT_SECONDS", 0.001) + + await runner._ensure_relay_started_for_job(object(), "job-1") + + assert "job-1" in caplog.text + assert "TimeoutError" in caplog.text + + @pytest.fixture(name="content_encryption_manager_guard") def fixture_content_encryption_manager_guard(): """Reset content-encryption globals even when a test assertion fails.""" @@ -824,15 +841,15 @@ async def run_relay_workflow(name, operation, **kwargs): owner_user_id=owner_user_id, ) + worker_trace_id = observed.pop("worker_trace_id") assert observed == { "construction_conversation_id": conversation_id, "construction_user_id": owner_user_id, "invocation_conversation_id": conversation_id, "invocation_user_id": owner_user_id, "resolved_collection": expected_collection, - "worker_trace_id": observed["worker_trace_id"], } - assert observed["worker_trace_id"] != "1" * 32 + assert worker_trace_id != "1" * 32 assert relay_observed == { "name": "async_shallow_research_job", "session_id": conversation_id, @@ -845,6 +862,7 @@ async def run_relay_workflow(name, operation, **kwargs): "aiq.submission.span_id": "submission-span", }, } + assert nat_events assert all(step.payload.UUID != "submission-span" for step in nat_events) assert outer_context.conversation_id.get() == "stale-parent-context" assert outer_context.user_id.get() == "jwt:stale-owner" diff --git a/tests/aiq_agent/test_default_model_profiles.py b/tests/aiq_agent/test_default_model_profiles.py deleted file mode 100644 index 1cc204c91..000000000 --- a/tests/aiq_agent/test_default_model_profiles.py +++ /dev/null @@ -1,211 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import json -from pathlib import Path -from unittest.mock import MagicMock - -import pytest -import yaml - -from aiq_agent.agents.shallow_researcher.agent import ShallowResearcherAgent -from aiq_agent.common import LLMProvider - -REPO_ROOT = Path(__file__).resolve().parents[2] -SHARED_SHALLOW_PROMPT = REPO_ROOT / "src/aiq_agent/agents/shallow_researcher/prompts/researcher.j2" -BREV_GETTING_STARTED_NOTEBOOK = REPO_ROOT / "docs/notebooks/0_Getting_Started_with_AIQ.ipynb" - -ULTRA_MODEL = "nvidia/nemotron-3-ultra-550b-a55b" -LIGHTNING_MODEL = "nvidia/nemotron-3.5-lightning-30b-a3b" -BUILD_BASE_URL = "https://integrate.api.nvidia.com/v1" - -CONFIG_GLOBS = ( - ".agents/skills/aiq-configure-workflow/assets/config-scaffold.yml", - "configs/config_*.yml", - "frontends/benchmarks/**/configs/*.yml", -) -CONFIG_PATHS = tuple(sorted(path for pattern in CONFIG_GLOBS for path in REPO_ROOT.glob(pattern))) -FRESHQA_CONFIG_PATHS = tuple(sorted(REPO_ROOT.glob("frontends/benchmarks/freshqa/configs/*.yml"))) -SHALLOW_PROFILE_PATHS = ( - REPO_ROOT / "configs/config_web_default_guardrails.yml", - REPO_ROOT / "configs/config_frontier_models.yml", -) - -DEPRECATED_REFERENCES = ( - "/".join(("nvidia", "nemotron-3-super-120b-a12b")), - "/".join(("nvidia", "nemotron-3-nano-30b-a3b")), - "/".join(("nvidia", "nemotron-mini-4b-instruct")), - "/".join(("nvidia", "llama-nemotron-embed-vl-1b-v2")), - "/".join(("nvidia", "nemotron-nano-12b-v2-vl")), - "/".join(("openai", "gpt-oss-120b")), - ".".join(("inference-api", "nvidia", "com")), -) -SCANNED_SUFFIXES = { - ".baseline", - ".example", - ".ipynb", - ".json", - ".md", - ".py", - ".sh", - ".ts", - ".tsx", - ".yaml", - ".yml", -} -IGNORED_PARTS = { - ".git", - ".mypy_cache", - ".pytest_cache", - ".ruff_cache", - ".venv", - "build", - "dist", - "node_modules", - "results", -} - - -def _load_config(path: Path) -> dict: - return yaml.safe_load(path.read_text(encoding="utf-8")) - - -def _model_for_alias(config: dict, alias: str) -> str: - return config["llms"][alias]["model_name"] - - -def _thinking_enabled(config: dict, alias: str) -> bool: - return bool(config["llms"][alias].get("chat_template_kwargs", {}).get("enable_thinking", False)) - - -def _registered_source_tools(config: dict) -> set[str]: - functions = config.get("functions", {}) - registries = ( - function - for function in functions.values() - if isinstance(function, dict) and function.get("_type") == "data_source_registry" - ) - return { - tool for registry in registries for source in registry.get("sources", []) for tool in source.get("tools", []) - } - - -@pytest.mark.parametrize("config_path", CONFIG_PATHS, ids=lambda path: str(path.relative_to(REPO_ROOT))) -def test_default_profiles_use_role_appropriate_models(config_path: Path): - config = _load_config(config_path) - functions = config.get("functions", {}) - is_frontier_profile = config_path.name == "config_frontier_models.yml" - - for function in functions.values(): - if not isinstance(function, dict): - continue - - function_type = function.get("_type") - if function_type == "intent_classifier": - if is_frontier_profile: - continue - alias = function["llm"] - assert alias == "nemotron_lightning_intent_llm" - assert _model_for_alias(config, alias) == LIGHTNING_MODEL - assert config["llms"][alias]["base_url"] == BUILD_BASE_URL - assert config["llms"][alias]["api_key"] == "${NVIDIA_API_KEY}" - assert config["llms"][alias]["temperature"] == 0.1 - assert config["llms"][alias]["top_p"] == 0.9 - assert config["llms"][alias]["max_tokens"] == 1024 - assert not config["llms"][alias]["parallel_tool_calls"] - assert not _thinking_enabled(config, alias) - elif function_type == "shallow_research_agent": - if is_frontier_profile: - continue - alias = function["llm"] - assert alias == "nemotron_lightning_agent_llm" - assert _model_for_alias(config, alias) == LIGHTNING_MODEL - assert config["llms"][alias]["base_url"] == BUILD_BASE_URL - assert config["llms"][alias]["api_key"] == "${NVIDIA_API_KEY}" - assert config["llms"][alias]["temperature"] == 0.2 - assert config["llms"][alias]["top_p"] == 0.7 - assert config["llms"][alias]["max_tokens"] == 8192 - assert not config["llms"][alias]["parallel_tool_calls"] - assert _thinking_enabled(config, alias) - elif not is_frontier_profile and function_type == "clarifier_agent": - assert _model_for_alias(config, function["llm"]) == ULTRA_MODEL - elif not is_frontier_profile and function_type == "deep_research_agent": - assert function["writer_llm"] == "nemotron_ultra_writer_llm" - for role in ( - "orchestrator_llm", - "source_router_llm", - "researcher_llm", - "planner_llm", - "writer_llm", - ): - assert _model_for_alias(config, function[role]) == ULTRA_MODEL - - -@pytest.mark.parametrize("config_path", FRESHQA_CONFIG_PATHS, ids=lambda path: path.name) -def test_freshqa_research_tools_are_registered_data_sources(config_path: Path): - config = _load_config(config_path) - source_tools = _registered_source_tools(config) - - for function in config.get("functions", {}).values(): - if isinstance(function, dict) and function.get("_type") in { - "shallow_research_agent", - "deep_research_agent", - }: - assert set(function.get("tools", [])) <= source_tools - - -@pytest.mark.parametrize("config_path", SHALLOW_PROFILE_PATHS, ids=lambda path: path.name) -def test_shallow_profiles_use_the_shared_citation_prompt(config_path: Path): - """Default Lightning and frontier Luna must share the hardened prompt and runtime path.""" - config = _load_config(config_path) - shallow = config["functions"]["shallow_research_agent"] - agent = ShallowResearcherAgent(llm_provider=MagicMock(spec=LLMProvider), tools=[]) - - assert shallow["_type"] == "shallow_research_agent" - assert "system_prompt" not in shallow - assert agent.system_prompt == SHARED_SHALLOW_PROMPT.read_text(encoding="utf-8") - - -def test_brev_getting_started_uses_ultra_for_shallow_research(): - """The Brev launchable avoids the hosted Lightning shallow-serving limitation.""" - notebook = json.loads(BREV_GETTING_STARTED_NOTEBOOK.read_text(encoding="utf-8")) - config_cells = [ - cell - for cell in notebook["cells"] - if cell.get("cell_type") == "code" - and cell.get("source", [""])[0].startswith("%%writefile config_simple_researcher.yml") - ] - - assert len(config_cells) == 1 - config = yaml.safe_load("".join(config_cells[0]["source"][2:])) - shallow = config["functions"]["shallow_research_agent"] - shallow_alias = shallow["llm"] - web_search = config["functions"]["web_search_tool"] - - assert config["functions"]["intent_classifier"]["llm"] == "nemotron_lightning_intent_llm" - assert shallow_alias == "nemotron_ultra_shallow_llm" - assert _model_for_alias(config, shallow_alias) == ULTRA_MODEL - assert config["llms"][shallow_alias]["max_tokens"] == 8192 - assert not config["llms"][shallow_alias]["parallel_tool_calls"] - assert _thinking_enabled(config, shallow_alias) - assert "nemotron_lightning_agent_llm" not in config["llms"] - assert shallow["max_llm_turns"] == 20 - assert shallow["max_tool_iterations"] == 5 - assert web_search["max_results"] == 5 - assert web_search["max_retries"] == 3 - assert not web_search["advanced_search"] - - -def test_deprecated_model_and_endpoint_references_are_absent(): - violations: list[str] = [] - - for path in REPO_ROOT.rglob("*"): - if not path.is_file() or path.suffix not in SCANNED_SUFFIXES or IGNORED_PARTS.intersection(path.parts): - continue - - text = path.read_text(encoding="utf-8") - for reference in DEPRECATED_REFERENCES: - if reference in text: - violations.append(f"{path.relative_to(REPO_ROOT)}: {reference}") - - assert not violations, "Deprecated references remain:\n" + "\n".join(violations) diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index f4d5f25a6..000000000 --- a/tests/conftest.py +++ /dev/null @@ -1,22 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Repository-wide pytest isolation for process-discovered configuration.""" - -import os -import shutil -import tempfile -from pathlib import Path - -# NeMo Relay discovers ``$XDG_CONFIG_HOME/nemo-relay/plugins.toml`` during -# import/initialization. Unit tests create real Relay scopes, so inheriting a -# developer's XDG directory would export fixture traffic to their configured -# observability destinations. Set this during conftest import, before pytest -# imports test modules that import Relay. -_TEST_XDG_CONFIG_HOME = Path(tempfile.mkdtemp(prefix="aiq-pytest-xdg-")) -os.environ["XDG_CONFIG_HOME"] = str(_TEST_XDG_CONFIG_HOME) - - -def pytest_unconfigure() -> None: - """Remove the process-local Relay discovery directory after the test run.""" - shutil.rmtree(_TEST_XDG_CONFIG_HOME, ignore_errors=True) diff --git a/tests/test_relay_runtime.py b/tests/test_relay_runtime.py index bf3fbd996..ccf20bbb5 100644 --- a/tests/test_relay_runtime.py +++ b/tests/test_relay_runtime.py @@ -248,6 +248,104 @@ class State: assert _safe_value({"state": State()}) == {"state": {"type": "State"}} +def test_safe_value_bounds_nested_and_large_values() -> None: + nested: object = "leaf" + for _ in range(20): + nested = {"nested": nested} + + projected = _safe_value({"nested": nested, "items": list(range(101)), "text": "x" * 20_000}) + + assert len(projected["items"]) < 101 + assert len(projected["text"]) < 20_000 + current = projected["nested"] + while "nested" in current: + current = current["nested"] + assert current["truncated"] is True + + +@pytest.mark.asyncio +async def test_relay_model_call_accepts_scalar_inputs(monkeypatch) -> None: + calls = [] + + class Model: + model_name = "test-model" + + async def ainvoke(self, messages, config=None): # noqa: ARG002 + calls.append(messages) + return AIMessage(content="done") + + async def passthrough(_self, request, handler): + return await handler(request) + + monkeypatch.setattr("aiq_agent.relay.runtime.NemoRelayMiddleware.awrap_model_call", passthrough) + message = HumanMessage(content="question") + + await ainvoke_with_relay(Model(), "question") + await ainvoke_with_relay(Model(), message) + + assert calls == [["question"], [message]] + + +@pytest.mark.asyncio +async def test_relay_model_middleware_fallback_does_not_retry_started_calls(monkeypatch, caplog) -> None: + calls = 0 + + class Model: + model_name = "test-model" + + async def ainvoke(self, messages, config=None): # noqa: ARG002 + nonlocal calls + calls += 1 + return AIMessage(content="done") + + async def fail_before(_self, request, handler): # noqa: ARG001 + raise RuntimeError("private middleware detail") + + monkeypatch.setattr("aiq_agent.relay.runtime.NemoRelayMiddleware.awrap_model_call", fail_before) + assert (await ainvoke_with_relay(Model(), [])).content == "done" + assert calls == 1 + assert "private middleware detail" not in caplog.text + + async def fail_after(_self, request, handler): + await handler(request) + raise RuntimeError("post-invocation failure") + + monkeypatch.setattr("aiq_agent.relay.runtime.NemoRelayMiddleware.awrap_model_call", fail_after) + with pytest.raises(RuntimeError, match="post-invocation failure"): + await ainvoke_with_relay(Model(), []) + assert calls == 2 + + +@pytest.mark.asyncio +async def test_relay_tool_middleware_fallback_does_not_retry_started_calls(monkeypatch, caplog) -> None: + calls = 0 + + class Tool: + name = "test_tool" + + async def ainvoke(self, args): + nonlocal calls + calls += 1 + return args["value"] + + async def fail_before(_self, request, handler): # noqa: ARG001 + raise RuntimeError("private middleware detail") + + monkeypatch.setattr("aiq_agent.relay.runtime.NemoRelayMiddleware.awrap_tool_call", fail_before) + assert await ainvoke_tool_with_relay(Tool(), {"value": "done"}) == "done" + assert calls == 1 + assert "private middleware detail" not in caplog.text + + async def fail_after(_self, request, handler): + await handler(request) + raise RuntimeError("post-invocation failure") + + monkeypatch.setattr("aiq_agent.relay.runtime.NemoRelayMiddleware.awrap_tool_call", fail_after) + with pytest.raises(RuntimeError, match="post-invocation failure"): + await ainvoke_tool_with_relay(Tool(), {"value": "done"}) + assert calls == 2 + + @pytest.mark.asyncio async def test_callback_does_not_duplicate_middleware_managed_llm_and_tool_scopes(tmp_path: Path) -> None: class TestChatModel(FakeMessagesListChatModel): @@ -612,10 +710,13 @@ async def nested_agent() -> None: input_value={"question": "question-2"}, ) finally: - await shutdown_async() - server.shutdown() - server.server_close() - server_thread.join() + try: + await shutdown_async() + finally: + server.shutdown() + server.server_close() + server_thread.join(timeout=5) + assert not server_thread.is_alive() events = [json.loads(line) for line in (tmp_path / "two-turns.jsonl").read_text().splitlines()] scope_events = [event for event in events if event["kind"] == "scope"] @@ -804,11 +905,14 @@ def log_message(self, format: str, *args: object) -> None: await ensure_started(config) with nemo_relay.scope.scope("otel-test", nemo_relay.ScopeType.Agent): pass - await shutdown_async() finally: - server.shutdown() - server.server_close() - server_thread.join() + try: + await shutdown_async() + finally: + server.shutdown() + server.server_close() + server_thread.join(timeout=5) + assert not server_thread.is_alive() assert len(received) == 3 assert {path for path, _, _ in received} == { diff --git a/tests/tokenomics/test_atof_adapter.py b/tests/tokenomics/test_atof_adapter.py index cea05cdfd..0d7ce0406 100644 --- a/tests/tokenomics/test_atof_adapter.py +++ b/tests/tokenomics/test_atof_adapter.py @@ -8,6 +8,7 @@ import pytest +from aiq_agent.tokenomics import atof_adapter from aiq_agent.tokenomics.atof_adapter import parse_trace from aiq_agent.tokenomics.pricing import PricingRegistry from aiq_agent.tokenomics.profile import PHASE_ORCHESTRATOR @@ -160,3 +161,48 @@ def test_parse_trace_skips_invalid_json_and_uses_catalog_fallback(tmp_path: Path assert profile.total_prompt_tokens == 100 assert profile.total_completion_tokens == 50 assert profile.total_cost_usd == pytest.approx(0.0002) + + +def test_parse_trace_ignores_non_string_identifiers(tmp_path: Path) -> None: + events = [ + _scope( + "root", + "function", + "workflow", + "start", + "2026-01-01T00:00:00Z", + metadata={"aiq.component.type": "workflow"}, + ), + _scope("root", "function", "workflow", "end", "2026-01-01T00:00:01Z"), + {"kind": "scope", "scope_category": "start", "uuid": ["invalid"], "parent_uuid": {"invalid": True}}, + ] + path = tmp_path / "relay.atof.jsonl" + _write(path, events) + + profiles = parse_trace(str(path), _pricing()) + + assert len(profiles) == 1 + + +def test_parse_trace_failure_log_excludes_exception_content(tmp_path: Path, monkeypatch, caplog) -> None: + events = [ + _scope( + "root", + "function", + "workflow", + "start", + "2026-01-01T00:00:00Z", + metadata={"aiq.component.type": "workflow"}, + ) + ] + path = tmp_path / "relay.atof.jsonl" + _write(path, events) + + def fail(*args, **kwargs): # noqa: ARG001 + raise RuntimeError("customer-secret") + + monkeypatch.setattr(atof_adapter, "_parse_request", fail) + + assert parse_trace(str(path), _pricing()) == [] + assert "RuntimeError" in caplog.text + assert "customer-secret" not in caplog.text From 33b005d5eeebcb84e05d021763b08fb27a6d1307 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Aug 2026 11:02:45 -0700 Subject: [PATCH 05/13] use relay's callbacks for deepagents Signed-off-by: Chantal D Gama Rose --- mcp/uv.lock | 63 ++++++++---------- pyproject.toml | 7 +- src/aiq_agent/agents/deep_researcher/agent.py | 4 +- src/aiq_agent/relay/runtime.py | 42 +----------- .../agents/deep_researcher/test_agent.py | 7 +- tests/test_relay_runtime.py | 54 +-------------- uv.lock | 65 +++++++++---------- 7 files changed, 70 insertions(+), 172 deletions(-) diff --git a/mcp/uv.lock b/mcp/uv.lock index fe112e07b..0e382edee 100644 --- a/mcp/uv.lock +++ b/mcp/uv.lock @@ -202,16 +202,16 @@ requires-dist = [ { name = "aiosqlite", specifier = ">=0.19.0" }, { name = "asyncpg", specifier = ">=0.29.0" }, { name = "boto3", marker = "extra == 's3'", specifier = ">=1.35.0,<2" }, - { name = "deepagents", specifier = ">=0.6.5" }, + { name = "deepagents", specifier = ">=0.7.4,<0.8" }, { name = "en-core-web-lg", marker = "extra == 'pii'", url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.8.0/en_core_web_lg-3.8.0-py3-none-any.whl" }, { name = "knowledge-layer", extras = ["all"], editable = "../sources/knowledge_layer" }, - { name = "langchain-modal", specifier = "==0.0.5" }, + { name = "langchain-modal", specifier = "==0.0.6" }, { name = "langgraph-checkpoint-postgres", specifier = ">=3.0.0" }, { name = "langgraph-checkpoint-sqlite", specifier = ">=2.0.0" }, { name = "mcp", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.5.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, - { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], specifier = ">=0.7.3,<0.8" }, + { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0" }, { name = "nvidia-nat", extras = ["langchain", "async-endpoints", "phoenix", "mcp"], specifier = "==1.8.0" }, { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, @@ -340,7 +340,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/07/38/e321b0e05d8cc068a [[package]] name = "anthropic" -version = "0.108.0" +version = "0.124.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -352,9 +352,9 @@ dependencies = [ { name = "sniffio" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/c7/d7f6d2e3975893958081f0282751217757333a3830d0d95859023d7006d0/anthropic-0.108.0.tar.gz", hash = "sha256:91b70253debb477a99f7ca43dac3f71e52207db79d4b06f104080b8dd1693e3b", size = 909409, upload-time = "2026-06-09T16:37:43.584Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/44/2afd86c9aac387b41f1b133cd3f595e0b110a2b980625f7987e897cee543/anthropic-0.124.0.tar.gz", hash = "sha256:b5c855a2912b157829a667ce21e10561f2b23bca98bb03af6abee2aa7f4a4cf3", size = 1054437, upload-time = "2026-08-19T16:51:31.493Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ab/40/75a937ddd8f230ec129d27de60df69ce8afcab1d0b15f7d651a5a95fac8a/anthropic-0.108.0-py3-none-any.whl", hash = "sha256:bdee7b14c13cf5a60b2c8ae0cf195720e0ea7fd8ab90df5a3899c50f1c91c4be", size = 870079, upload-time = "2026-06-09T16:37:44.895Z" }, + { url = "https://files.pythonhosted.org/packages/55/6a/3ab35f7bb9a2c512c80638eb108fbbcee21c433a00baca8ab0ded485e007/anthropic-0.124.0-py3-none-any.whl", hash = "sha256:e7f86e3182c4173edf09fd8f42aa328d1aa9246355a90a300287c355228bb356", size = 1147647, upload-time = "2026-08-19T16:51:33.459Z" }, ] [[package]] @@ -622,11 +622,11 @@ wheels = [ [[package]] name = "bracex" -version = "2.6" +version = "3.0.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/63/9a/fec38644694abfaaeca2798b58e276a8e61de49e2e37494ace423395febc/bracex-2.6.tar.gz", hash = "sha256:98f1347cd77e22ee8d967a30ad4e310b233f7754dbf31ff3fceb76145ba47dc7", size = 26642, upload-time = "2025-06-22T19:12:31.254Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ac/01/5f394b8bcd6e5b92f73130990960423bbb19711f906bd9fe9ea5557c667c/bracex-3.0.1.tar.gz", hash = "sha256:4e38e32392e4a4780fe15d644bfc7c8514057cfc3861e060b11814ce829c25e4", size = 44019, upload-time = "2026-07-20T13:43:00.335Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/2a/9186535ce58db529927f6cf5990a849aa9e052eea3e2cfefe20b9e1802da/bracex-2.6-py3-none-any.whl", hash = "sha256:0b0049264e7340b3ec782b5cb99beb325f36c3782a32e36e876452fd49a09952", size = 11508, upload-time = "2025-06-22T19:12:29.781Z" }, + { url = "https://files.pythonhosted.org/packages/b8/8f/6f7273a7adb8d73fc8d21ede4376a3e475e52f98435c6007f69100dec8ca/bracex-3.0.1-py3-none-any.whl", hash = "sha256:6523ad83aeb5098a4ee597cff0f964442ff74e460bd3fafaffab6a013ff2288c", size = 11940, upload-time = "2026-07-20T13:42:59.268Z" }, ] [[package]] @@ -1151,7 +1151,7 @@ wheels = [ [[package]] name = "deepagents" -version = "0.6.12" +version = "0.7.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain" }, @@ -1159,11 +1159,12 @@ dependencies = [ { name = "langchain-core" }, { name = "langchain-google-genai" }, { name = "langsmith" }, + { name = "packaging" }, { name = "wcmatch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e0/db/a6acdc72a9e90c3f07ed10de35c951734a02d4facb693bb59684ad368801/deepagents-0.6.12.tar.gz", hash = "sha256:1f281c0bc5a63132f62e2ee345c1dc593b23188da6e23016401f6879fbe54b5f", size = 211364, upload-time = "2026-06-25T17:26:52.775Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/20/6a2fff4d42984f6302f211c164d4cf1b1894086b576736d81a935a67186e/deepagents-0.7.7.tar.gz", hash = "sha256:68b8f0c861a95065b3eda39b5c3748943e09500898bc39e7e689f0b7deba9f72", size = 282001, upload-time = "2026-08-18T15:06:55.05Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/98/49/af7219b3c13520fee047bb807cfaefba17f8e4584c551d946773589a4f08/deepagents-0.6.12-py3-none-any.whl", hash = "sha256:28b8fa0119ca0a689e3e18e288c4634e4046062acfc87a1cb34289d3af3a1c88", size = 236120, upload-time = "2026-06-25T17:26:51.736Z" }, + { url = "https://files.pythonhosted.org/packages/e2/15/0e029e7f119c4c7a9cb17702b604153af71571d36c646c2837e5f4d3074d/deepagents-0.7.7-py3-none-any.whl", hash = "sha256:ddce2604799205ef51dc2c17fae613c9e162bbbc90206fcff2860c4e6f48bdb1", size = 308616, upload-time = "2026-08-18T15:06:53.532Z" }, ] [[package]] @@ -2249,16 +2250,16 @@ wheels = [ [[package]] name = "langchain-anthropic" -version = "1.4.8" +version = "1.5.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anthropic" }, { name = "langchain-core" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/98/22/40ab129b08329ca295b391aa1d48267692b42594757084c6918e22b655ac/langchain_anthropic-1.4.8.tar.gz", hash = "sha256:c76891b2044d56105ff13c106ed12650637b53bd598a4bdf15b4796eefa2a4ec", size = 708524, upload-time = "2026-06-26T21:28:46.916Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/c6/97c439282c13225beb56ba96ed2a5b1cb2c32eacb84c518fe475ce43711d/langchain_anthropic-1.5.6.tar.gz", hash = "sha256:648fdab25573fc9d29543c4b4af1d682b7b6e548d452bb46e67ebc586e195629", size = 720743, upload-time = "2026-08-13T02:30:48.262Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/14/746235c4da89d9bc6a608c5f489f628e03feb8f697195c146e452c8f23c8/langchain_anthropic-1.4.8-py3-none-any.whl", hash = "sha256:778e9301b6fd517824f76ec1776975ce8add97a1f6a36c50ae3c2f4b03a66f7f", size = 52366, upload-time = "2026-06-26T21:28:45.535Z" }, + { url = "https://files.pythonhosted.org/packages/de/5e/61b288074742179041cb5390430e52429a0fc41f4a94efec20042bda009f/langchain_anthropic-1.5.6-py3-none-any.whl", hash = "sha256:c358ed2ca90ef75254bb73a96b0a8a8ef0efc1deabe529ea06db6ff403001a27", size = 56547, upload-time = "2026-08-13T02:30:47.118Z" }, ] [[package]] @@ -2410,15 +2411,15 @@ wheels = [ [[package]] name = "langchain-modal" -version = "0.0.5" +version = "0.0.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "deepagents" }, { name = "modal" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/cc/a7/4c65b5204837dd1d26b93aeb01768d053bf6f92988f476842b7b9a5bf61c/langchain_modal-0.0.5.tar.gz", hash = "sha256:acd3ee16264ad97d0040f20b7c711948cf90bbf1fd791e9effdf2be3c82f38bf", size = 190600, upload-time = "2026-06-03T21:26:52.553Z" } +sdist = { url = "https://files.pythonhosted.org/packages/74/e8/c364ac18bd1d348f80bc334818ba89e1e0544b87cacf42d0b78cd5a49916/langchain_modal-0.0.6.tar.gz", hash = "sha256:7e1fe4dc0cae0937cfb9bd82946787e351fc9d4972be2963cdba8d43a194c2f7", size = 191018, upload-time = "2026-07-29T18:02:01.618Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/94/ce0d8423c85d61422f00022852185296bcd43c98e7bc2ba9075a1dcf6adc/langchain_modal-0.0.5-py3-none-any.whl", hash = "sha256:123d9b3ce781a3cf59d87e5297d1ea8608a1f2042327593321cb8ab6257d8599", size = 4633, upload-time = "2026-06-03T21:26:51.591Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a9/500b528c5a837b2960a3f87d3f558f5dde4299a0ba2ab916fd103bf16ff8/langchain_modal-0.0.6-py3-none-any.whl", hash = "sha256:6d43caa94e6846151592f1828a9c9380b7ceda90ace45fce1c3cad74c4edfa8c", size = 4754, upload-time = "2026-07-29T18:02:00.62Z" }, ] [[package]] @@ -2600,7 +2601,7 @@ wheels = [ [[package]] name = "langsmith" -version = "0.9.8" +version = "0.11.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -2618,9 +2619,9 @@ dependencies = [ { name = "xxhash" }, { name = "zstandard" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0c/68/8d8471233ee0cd82c2af946d76f80a01aeb8bb04160c392c1229fddf5d3d/langsmith-0.9.8.tar.gz", hash = "sha256:8c3d6a6d5246a3ea6d439b726d59edefba31dfb251de9eedb256119bbea4439e", size = 4710812, upload-time = "2026-07-06T19:06:10.866Z" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/57/7b6c11080c9e082ebf1456a2e2372fae8f23a85a5ae2869bdd5ab9a6507d/langsmith-0.11.1.tar.gz", hash = "sha256:47998977366acb3ba3093881fd465cbf11a5f8c2f4e87e40a17dda203f6dedf4", size = 4814724, upload-time = "2026-08-19T15:47:44.116Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4a/67/a85caaa99117bbc988a0df7faa39e7f68344361854638d86bcce0ffe3619/langsmith-0.9.8-py3-none-any.whl", hash = "sha256:098da9fc6c184284f17913cb813a41e28c5ab1508e90bd50db40c28166681017", size = 671148, upload-time = "2026-07-06T19:06:08.911Z" }, + { url = "https://files.pythonhosted.org/packages/f1/85/0ad6df25588122760b2b40ec182eafc55532c9e66015c83e16675bd34a29/langsmith-0.11.1-py3-none-any.whl", hash = "sha256:cfc3437a9cf27440cd0095c24df945edbceb6df10b579bc8b3980b4ad367835f", size = 744589, upload-time = "2026-08-19T15:47:42.043Z" }, ] [[package]] @@ -3305,18 +3306,8 @@ wheels = [ [[package]] name = "nemo-relay" -version = "0.7.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/a259aac8df4aa78c0b3a6f3ad0fbf6305666bfaad5c29d9adc26db0f9e27/nemo_relay-0.7.3.tar.gz", hash = "sha256:ea5a1bb52e25e001dcbf6af1830616be181845e978cc848df58562556bba5604", size = 1299430, upload-time = "2026-08-14T14:40:44.832Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ed/9e/4eb80d2307cadcbb839dad2212a8d888667cd8c531ad0d8c0c1afc841181/nemo_relay-0.7.3-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:301dfc8334032ac52c09cc0b1421181e7c4df601abe7c82dfe51aa74ddbb3732", size = 9251369, upload-time = "2026-08-14T14:40:07.668Z" }, - { url = "https://files.pythonhosted.org/packages/c7/a9/9fb77f7142b1381d8c3c81fdbb76782a02dcff1ee403ac6c93b13fa46b24/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:49c70c0a94cebb6cba3dd7521be11f63eac4cd9386881eb29827ac91b1bd780b", size = 8458046, upload-time = "2026-08-14T14:40:10.181Z" }, - { url = "https://files.pythonhosted.org/packages/ae/ab/b2f246f971f561a982d234d9f3ec1b29dfa4c72bf4882f3e32ce6db54dfa/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8eeaaf8c6a18440e473dd14eb1bb82e56c6514e70adee7456eddf9ce217cff89", size = 8957931, upload-time = "2026-08-14T14:40:12.547Z" }, - { url = "https://files.pythonhosted.org/packages/8f/d8/d8c25ba915467bab457d175b84a3af5c33291655867d472c49eada3b77e6/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:90ff984a89c42ebd0cfe26a0af3f180b1970e2ebcd743d19a960e547949ad2ef", size = 10325640, upload-time = "2026-08-14T14:40:15.722Z" }, - { url = "https://files.pythonhosted.org/packages/7d/fd/6ec48f47eb5ca4cca566b197ade8514baf183dddf53d6de3a296b8bc1102/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d4895a4df39a92ecbac3633fb4ec59cf1f83286c0369f433760ab28fb8d6dddf", size = 10706736, upload-time = "2026-08-14T14:40:18.377Z" }, - { url = "https://files.pythonhosted.org/packages/6b/cb/8a6f8d5f9922e75100f135c1dc4451bafc15f5d51608081f72b894e0caf4/nemo_relay-0.7.3-cp311-abi3-win_amd64.whl", hash = "sha256:f123cd45a27fca3d570559f2c26850138763411f66f387559d757e5d88bfa3c1", size = 8807604, upload-time = "2026-08-14T14:40:20.98Z" }, - { url = "https://files.pythonhosted.org/packages/f1/ac/558a05e6b9e28464b64fee8327d9e8a2aab3fb356363b8b82597b6b35b2a/nemo_relay-0.7.3-cp311-abi3-win_arm64.whl", hash = "sha256:6d5444e9a03b8b5d4bba2409105a121349580cb51481d441c6b5e699713a8763", size = 8442944, upload-time = "2026-08-14T14:40:23.433Z" }, -] +version = "0.8.0" +source = { git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0#2557abb5ee87d61fe914b9bc9b8442210920f7a0" } [package.optional-dependencies] deepagents = [ @@ -6022,14 +6013,14 @@ wheels = [ [[package]] name = "wcmatch" -version = "10.1" +version = "11.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "bracex" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/79/3e/c0bdc27cf06f4e47680bd5803a07cb3dfd17de84cde92dd217dcb9e05253/wcmatch-10.1.tar.gz", hash = "sha256:f11f94208c8c8484a16f4f48638a85d771d9513f4ab3f37595978801cb9465af", size = 117421, upload-time = "2025-06-22T19:14:02.49Z" } +sdist = { url = "https://files.pythonhosted.org/packages/57/43/30e407989e313677dbb9d5f045f966549a7254834571e342eaa4b55cc67b/wcmatch-11.0.1.tar.gz", hash = "sha256:1ea2b4fa678b8ca268253798d5963935df39132d47c3e241c0a0732224005e7d", size = 144662, upload-time = "2026-08-14T15:20:40.477Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/eb/d8/0d1d2e9d3fabcf5d6840362adcf05f8cf3cd06a73358140c3a97189238ae/wcmatch-10.1-py3-none-any.whl", hash = "sha256:5848ace7dbb0476e5e55ab63c6bbd529745089343427caa5537f230cc01beb8a", size = 39854, upload-time = "2025-06-22T19:14:00.978Z" }, + { url = "https://files.pythonhosted.org/packages/ce/77/7a02b0f05b3ffcdbef9719ce3ee0b508d6a29b58e95299f1580055671db3/wcmatch-11.0.1-py3-none-any.whl", hash = "sha256:fd149ecddb9f0a88ea780017d6dde17c994e494e7f7303d4e3c9d6251f978f4b", size = 43449, upload-time = "2026-08-14T15:20:39.379Z" }, ] [[package]] diff --git a/pyproject.toml b/pyproject.toml index e2066ee1f..d98b56266 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,11 +41,11 @@ dependencies = [ "nvidia-nat-profiler==1.8.0", "nvidia-nat-redis==1.8.0", "nvidia-nat-security[guardrails]==1.8.0", - "nemo-relay[deepagents,langchain,langgraph]>=0.7.3,<0.8", - "deepagents>=0.6.5", + "nemo-relay[deepagents,langchain,langgraph]>=0.8,<0.9", + "deepagents>=0.7.4,<0.8", "langgraph-checkpoint-postgres>=3.0.0", "langgraph-checkpoint-sqlite>=2.0.0", - "langchain-modal==0.0.5", + "langchain-modal==0.0.6", "psycopg[binary]>=3.0.0", "asyncpg>=0.29.0", "aiosqlite>=0.19.0", @@ -274,6 +274,7 @@ you-com = { workspace = true } duckduckgo-news-search = { workspace = true } polymarket-prediction-market = { workspace = true } knowledge-layer = { workspace = true } +nemo-relay = { git = "https://github.com/NVIDIA/NeMo-Relay.git", rev = "2557abb5ee87d61fe914b9bc9b8442210920f7a0" } # pragma: allowlist secret aiq-api = { workspace = true } aiq-research-cli = { workspace = true } aiq-debug = { workspace = true } diff --git a/src/aiq_agent/agents/deep_researcher/agent.py b/src/aiq_agent/agents/deep_researcher/agent.py index d5403578f..a8c0613c0 100644 --- a/src/aiq_agent/agents/deep_researcher/agent.py +++ b/src/aiq_agent/agents/deep_researcher/agent.py @@ -26,6 +26,7 @@ from uuid import uuid4 from langchain_core.tools import BaseTool +from nemo_relay.integrations.deepagents import NemoRelayDeepAgentsCallbackHandler from aiq_agent.common import LLMProvider from aiq_agent.common import load_prompt @@ -315,7 +316,8 @@ async def run(self, state: DeepResearchAgentState) -> DeepResearchAgentState: async with execution_timeout: async def _invoke_orchestrator() -> Any: - return await agent.ainvoke(state, config={"callbacks": self.callbacks}) + callbacks = [*self.callbacks, NemoRelayDeepAgentsCallbackHandler()] + return await agent.ainvoke(state, config={"callbacks": callbacks}) result = await run_agent( "deep_research_agent", diff --git a/src/aiq_agent/relay/runtime.py b/src/aiq_agent/relay/runtime.py index 412431b99..5465cae4b 100644 --- a/src/aiq_agent/relay/runtime.py +++ b/src/aiq_agent/relay/runtime.py @@ -20,7 +20,6 @@ from uuid import uuid4 import nemo_relay -from langchain.agents.middleware import AgentMiddleware from langchain.agents.middleware import ModelRequest from langchain.agents.middleware import ModelResponse from langchain.agents.middleware import ToolCallRequest @@ -89,51 +88,12 @@ def _normalize_chat_nvidia_binding( return runnable.bound, dict(runnable.kwargs), merge_configs(runnable.config, config) -# Work around NVIDIA/NeMo-Relay#805 until DeepAgents emits nested local-subagent Agent scopes. -class _DelegatedAgentScopeMiddleware(AgentMiddleware): - """Create a semantic Agent scope for the subagent selected by DeepAgents.""" - - @staticmethod - def _agent_name(request: Any) -> str | None: - tool_call = getattr(request, "tool_call", None) - if not isinstance(tool_call, dict) or tool_call.get("name") != "task": - return None - arguments = tool_call.get("args") - if not isinstance(arguments, dict): - return None - name = arguments.get("subagent_type") - return name if isinstance(name, str) and name else None - - def wrap_tool_call(self, request: Any, handler: Callable[[Any], Any]) -> Any: - name = self._agent_name(request) - if name is None: - return handler(request) - with agent_scope(name, input_value=getattr(request, "tool_call", None)) as lifecycle: - result = handler(request) - lifecycle.output = result - return result - - async def awrap_tool_call(self, request: Any, handler: Callable[[Any], Awaitable[Any]]) -> Any: - name = self._agent_name(request) - if name is None: - return await handler(request) - with agent_scope(name, input_value=getattr(request, "tool_call", None)) as lifecycle: - result = await handler(request) - lifecycle.output = result - return result - - def deepagents_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: """Attach Relay's supported DeepAgents middleware.""" from nemo_relay.integrations.deepagents import add_nemo_relay_integration - observed = add_nemo_relay_integration(kwargs) - middleware = list(observed.get("middleware") or ()) - if not any(isinstance(item, _DelegatedAgentScopeMiddleware) for item in middleware): - middleware.append(_DelegatedAgentScopeMiddleware()) - observed["middleware"] = middleware - return observed + return add_nemo_relay_integration(kwargs) def merge_langchain_middleware(middleware: Sequence[Any] | None) -> list[Any]: diff --git a/tests/aiq_agent/agents/deep_researcher/test_agent.py b/tests/aiq_agent/agents/deep_researcher/test_agent.py index 91a814fc2..447a78efd 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_agent.py +++ b/tests/aiq_agent/agents/deep_researcher/test_agent.py @@ -32,6 +32,7 @@ from langchain_core.messages import ToolMessage from langchain_core.runnables import RunnableLambda from langchain_core.tools import tool +from nemo_relay.integrations.deepagents import NemoRelayDeepAgentsCallbackHandler from aiq_agent.agents.deep_researcher.custom_middleware import FinalReportCommitTracker from aiq_agent.agents.deep_researcher.models import DeepResearchAgentState @@ -1772,9 +1773,9 @@ async def test_run_with_callbacks(self, mock_llm_provider, real_tool, mock_creat await agent.run(state) - # Callbacks should have been passed to ainvoke - call_kwargs = mock_create_deep_agent.ainvoke.call_args - assert call_kwargs is not None + callbacks = mock_create_deep_agent.ainvoke.call_args.kwargs["config"]["callbacks"] + assert callbacks[0] is mock_callback + assert isinstance(callbacks[1], NemoRelayDeepAgentsCallbackHandler) @pytest.mark.asyncio async def test_run_handles_error(self, mock_llm_provider, real_tool, caplog): diff --git a/tests/test_relay_runtime.py b/tests/test_relay_runtime.py index ccf20bbb5..fb2cc55d3 100644 --- a/tests/test_relay_runtime.py +++ b/tests/test_relay_runtime.py @@ -40,7 +40,7 @@ from aiq_agent.relay.runtime import run_workflow -def test_deepagents_integration_and_delegated_agent_scope_are_enabled() -> None: +def test_deepagents_integration_is_enabled() -> None: kwargs = deepagents_kwargs( { "model": "test", @@ -49,10 +49,7 @@ def test_deepagents_integration_and_delegated_agent_scope_are_enabled() -> None: "subagents": [{"name": "runtime-agent", "description": "test", "model": "test", "tools": []}], } ) - assert [type(middleware).__name__ for middleware in kwargs["middleware"][-2:]] == [ - "NemoRelayDeepAgentsMiddleware", - "_DelegatedAgentScopeMiddleware", - ] + assert [type(middleware).__name__ for middleware in kwargs["middleware"][-1:]] == ["NemoRelayDeepAgentsMiddleware"] assert [type(middleware).__name__ for middleware in kwargs["subagents"][0]["middleware"][-1:]] == [ "NemoRelayDeepAgentsMiddleware" ] @@ -401,53 +398,6 @@ async def operation() -> None: } -@pytest.mark.asyncio -async def test_deepagents_task_uses_runtime_subagent_name_for_nested_scope(tmp_path: Path) -> None: - config = RelayConfig() - config.logging = False - config.observability.atof.output_directory = str(tmp_path) - config.observability.atof.filename = "delegation.jsonl" - config.observability.opentelemetry.enabled = False - delegation_middleware = deepagents_kwargs({"model": "test", "tools": [], "name": "parent"})["middleware"][-1] - request = SimpleNamespace( - tool_call={ - "name": "task", - "args": {"subagent_type": "runtime-selected-agent", "description": "research this"}, - } - ) - - async def delegated_agent(_: object) -> str: - async def model_call(_: nemo_relay.LLMRequest) -> dict[str, str]: - return {"response": "done"} - - await nemo_relay.llm.execute( - "managed-model", - nemo_relay.LLMRequest({}, {"messages": []}), - model_call, - ) - return "done" - - async def task_call(_: object) -> str: - return await delegation_middleware.awrap_tool_call(request, delegated_agent) - - async def operation() -> None: - await nemo_relay.tools.execute("task", request.tool_call["args"], task_call) - - await ensure_started(config) - try: - await run_agent("deep_research_agent", operation) - finally: - await shutdown_async() - - events = [json.loads(line) for line in (tmp_path / "delegation.jsonl").read_text().splitlines()] - starts = { - event["name"]: event for event in events if event["kind"] == "scope" and event["scope_category"] == "start" - } - assert starts["task"]["parent_uuid"] == starts["deep_research_agent"]["uuid"] - assert starts["runtime-selected-agent"]["parent_uuid"] == starts["deep_research_agent"]["uuid"] - assert starts["managed-model"]["parent_uuid"] == starts["runtime-selected-agent"]["uuid"] - - @pytest.mark.asyncio async def test_concurrent_researchers_do_not_share_mutable_relay_agent_scopes(tmp_path: Path) -> None: config = RelayConfig() diff --git a/uv.lock b/uv.lock index 6e953798a..a72e0141a 100644 --- a/uv.lock +++ b/uv.lock @@ -290,16 +290,16 @@ requires-dist = [ { name = "aiosqlite", specifier = ">=0.19.0" }, { name = "asyncpg", specifier = ">=0.29.0" }, { name = "boto3", marker = "extra == 's3'", specifier = ">=1.35.0,<2" }, - { name = "deepagents", specifier = ">=0.6.5" }, + { name = "deepagents", specifier = ">=0.7.4,<0.8" }, { name = "en-core-web-lg", marker = "extra == 'pii'", url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.8.0/en_core_web_lg-3.8.0-py3-none-any.whl" }, { name = "knowledge-layer", extras = ["all"], editable = "sources/knowledge_layer" }, - { name = "langchain-modal", specifier = "==0.0.5" }, + { name = "langchain-modal", specifier = "==0.0.6" }, { name = "langgraph-checkpoint-postgres", specifier = ">=3.0.0" }, { name = "langgraph-checkpoint-sqlite", specifier = ">=2.0.0" }, { name = "mcp", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.5.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, - { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], specifier = ">=0.7.3,<0.8" }, + { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0" }, { name = "nvidia-nat", extras = ["langchain", "async-endpoints", "phoenix", "mcp"], specifier = "==1.8.0" }, { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, @@ -362,6 +362,7 @@ version = "0.1.0" source = { editable = "frontends/aiq_api" } dependencies = [ { name = "aiosqlite" }, + { name = "aiq-agent" }, { name = "asyncpg" }, { name = "dask", extra = ["distributed"] }, { name = "fastapi" }, @@ -386,6 +387,7 @@ postgres = [ [package.metadata] requires-dist = [ { name = "aiosqlite", specifier = ">=0.19.0" }, + { name = "aiq-agent", editable = "." }, { name = "asyncpg", specifier = ">=0.29.0" }, { name = "dask", extras = ["distributed"], specifier = ">=2024.1.0" }, { name = "fastapi", specifier = ">=0.100.0" }, @@ -497,7 +499,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/07/38/e321b0e05d8cc068a [[package]] name = "anthropic" -version = "0.108.0" +version = "0.124.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -509,9 +511,9 @@ dependencies = [ { name = "sniffio" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/c7/d7f6d2e3975893958081f0282751217757333a3830d0d95859023d7006d0/anthropic-0.108.0.tar.gz", hash = "sha256:91b70253debb477a99f7ca43dac3f71e52207db79d4b06f104080b8dd1693e3b", size = 909409, upload-time = "2026-06-09T16:37:43.584Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/44/2afd86c9aac387b41f1b133cd3f595e0b110a2b980625f7987e897cee543/anthropic-0.124.0.tar.gz", hash = "sha256:b5c855a2912b157829a667ce21e10561f2b23bca98bb03af6abee2aa7f4a4cf3", size = 1054437, upload-time = "2026-08-19T16:51:31.493Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ab/40/75a937ddd8f230ec129d27de60df69ce8afcab1d0b15f7d651a5a95fac8a/anthropic-0.108.0-py3-none-any.whl", hash = "sha256:bdee7b14c13cf5a60b2c8ae0cf195720e0ea7fd8ab90df5a3899c50f1c91c4be", size = 870079, upload-time = "2026-06-09T16:37:44.895Z" }, + { url = "https://files.pythonhosted.org/packages/55/6a/3ab35f7bb9a2c512c80638eb108fbbcee21c433a00baca8ab0ded485e007/anthropic-0.124.0-py3-none-any.whl", hash = "sha256:e7f86e3182c4173edf09fd8f42aa328d1aa9246355a90a300287c355228bb356", size = 1147647, upload-time = "2026-08-19T16:51:33.459Z" }, ] [[package]] @@ -844,11 +846,11 @@ wheels = [ [[package]] name = "bracex" -version = "2.6" +version = "3.0.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/63/9a/fec38644694abfaaeca2798b58e276a8e61de49e2e37494ace423395febc/bracex-2.6.tar.gz", hash = "sha256:98f1347cd77e22ee8d967a30ad4e310b233f7754dbf31ff3fceb76145ba47dc7", size = 26642, upload-time = "2025-06-22T19:12:31.254Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ac/01/5f394b8bcd6e5b92f73130990960423bbb19711f906bd9fe9ea5557c667c/bracex-3.0.1.tar.gz", hash = "sha256:4e38e32392e4a4780fe15d644bfc7c8514057cfc3861e060b11814ce829c25e4", size = 44019, upload-time = "2026-07-20T13:43:00.335Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/2a/9186535ce58db529927f6cf5990a849aa9e052eea3e2cfefe20b9e1802da/bracex-2.6-py3-none-any.whl", hash = "sha256:0b0049264e7340b3ec782b5cb99beb325f36c3782a32e36e876452fd49a09952", size = 11508, upload-time = "2025-06-22T19:12:29.781Z" }, + { url = "https://files.pythonhosted.org/packages/b8/8f/6f7273a7adb8d73fc8d21ede4376a3e475e52f98435c6007f69100dec8ca/bracex-3.0.1-py3-none-any.whl", hash = "sha256:6523ad83aeb5098a4ee597cff0f964442ff74e460bd3fafaffab6a013ff2288c", size = 11940, upload-time = "2026-07-20T13:42:59.268Z" }, ] [[package]] @@ -1530,7 +1532,7 @@ wheels = [ [[package]] name = "deepagents" -version = "0.6.12" +version = "0.7.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain" }, @@ -1538,11 +1540,12 @@ dependencies = [ { name = "langchain-core" }, { name = "langchain-google-genai" }, { name = "langsmith" }, + { name = "packaging" }, { name = "wcmatch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e0/db/a6acdc72a9e90c3f07ed10de35c951734a02d4facb693bb59684ad368801/deepagents-0.6.12.tar.gz", hash = "sha256:1f281c0bc5a63132f62e2ee345c1dc593b23188da6e23016401f6879fbe54b5f", size = 211364, upload-time = "2026-06-25T17:26:52.775Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/20/6a2fff4d42984f6302f211c164d4cf1b1894086b576736d81a935a67186e/deepagents-0.7.7.tar.gz", hash = "sha256:68b8f0c861a95065b3eda39b5c3748943e09500898bc39e7e689f0b7deba9f72", size = 282001, upload-time = "2026-08-18T15:06:55.05Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/98/49/af7219b3c13520fee047bb807cfaefba17f8e4584c551d946773589a4f08/deepagents-0.6.12-py3-none-any.whl", hash = "sha256:28b8fa0119ca0a689e3e18e288c4634e4046062acfc87a1cb34289d3af3a1c88", size = 236120, upload-time = "2026-06-25T17:26:51.736Z" }, + { url = "https://files.pythonhosted.org/packages/e2/15/0e029e7f119c4c7a9cb17702b604153af71571d36c646c2837e5f4d3074d/deepagents-0.7.7-py3-none-any.whl", hash = "sha256:ddce2604799205ef51dc2c17fae613c9e162bbbc90206fcff2860c4e6f48bdb1", size = 308616, upload-time = "2026-08-18T15:06:53.532Z" }, ] [[package]] @@ -2819,16 +2822,16 @@ wheels = [ [[package]] name = "langchain-anthropic" -version = "1.4.8" +version = "1.5.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anthropic" }, { name = "langchain-core" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/98/22/40ab129b08329ca295b391aa1d48267692b42594757084c6918e22b655ac/langchain_anthropic-1.4.8.tar.gz", hash = "sha256:c76891b2044d56105ff13c106ed12650637b53bd598a4bdf15b4796eefa2a4ec", size = 708524, upload-time = "2026-06-26T21:28:46.916Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/c6/97c439282c13225beb56ba96ed2a5b1cb2c32eacb84c518fe475ce43711d/langchain_anthropic-1.5.6.tar.gz", hash = "sha256:648fdab25573fc9d29543c4b4af1d682b7b6e548d452bb46e67ebc586e195629", size = 720743, upload-time = "2026-08-13T02:30:48.262Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/14/746235c4da89d9bc6a608c5f489f628e03feb8f697195c146e452c8f23c8/langchain_anthropic-1.4.8-py3-none-any.whl", hash = "sha256:778e9301b6fd517824f76ec1776975ce8add97a1f6a36c50ae3c2f4b03a66f7f", size = 52366, upload-time = "2026-06-26T21:28:45.535Z" }, + { url = "https://files.pythonhosted.org/packages/de/5e/61b288074742179041cb5390430e52429a0fc41f4a94efec20042bda009f/langchain_anthropic-1.5.6-py3-none-any.whl", hash = "sha256:c358ed2ca90ef75254bb73a96b0a8a8ef0efc1deabe529ea06db6ff403001a27", size = 56547, upload-time = "2026-08-13T02:30:47.118Z" }, ] [[package]] @@ -2980,15 +2983,15 @@ wheels = [ [[package]] name = "langchain-modal" -version = "0.0.5" +version = "0.0.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "deepagents" }, { name = "modal" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/cc/a7/4c65b5204837dd1d26b93aeb01768d053bf6f92988f476842b7b9a5bf61c/langchain_modal-0.0.5.tar.gz", hash = "sha256:acd3ee16264ad97d0040f20b7c711948cf90bbf1fd791e9effdf2be3c82f38bf", size = 190600, upload-time = "2026-06-03T21:26:52.553Z" } +sdist = { url = "https://files.pythonhosted.org/packages/74/e8/c364ac18bd1d348f80bc334818ba89e1e0544b87cacf42d0b78cd5a49916/langchain_modal-0.0.6.tar.gz", hash = "sha256:7e1fe4dc0cae0937cfb9bd82946787e351fc9d4972be2963cdba8d43a194c2f7", size = 191018, upload-time = "2026-07-29T18:02:01.618Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/94/ce0d8423c85d61422f00022852185296bcd43c98e7bc2ba9075a1dcf6adc/langchain_modal-0.0.5-py3-none-any.whl", hash = "sha256:123d9b3ce781a3cf59d87e5297d1ea8608a1f2042327593321cb8ab6257d8599", size = 4633, upload-time = "2026-06-03T21:26:51.591Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a9/500b528c5a837b2960a3f87d3f558f5dde4299a0ba2ab916fd103bf16ff8/langchain_modal-0.0.6-py3-none-any.whl", hash = "sha256:6d43caa94e6846151592f1828a9c9380b7ceda90ace45fce1c3cad74c4edfa8c", size = 4754, upload-time = "2026-07-29T18:02:00.62Z" }, ] [[package]] @@ -3196,7 +3199,7 @@ wheels = [ [[package]] name = "langsmith" -version = "0.9.8" +version = "0.11.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -3214,9 +3217,9 @@ dependencies = [ { name = "xxhash" }, { name = "zstandard" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0c/68/8d8471233ee0cd82c2af946d76f80a01aeb8bb04160c392c1229fddf5d3d/langsmith-0.9.8.tar.gz", hash = "sha256:8c3d6a6d5246a3ea6d439b726d59edefba31dfb251de9eedb256119bbea4439e", size = 4710812, upload-time = "2026-07-06T19:06:10.866Z" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/57/7b6c11080c9e082ebf1456a2e2372fae8f23a85a5ae2869bdd5ab9a6507d/langsmith-0.11.1.tar.gz", hash = "sha256:47998977366acb3ba3093881fd465cbf11a5f8c2f4e87e40a17dda203f6dedf4", size = 4814724, upload-time = "2026-08-19T15:47:44.116Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4a/67/a85caaa99117bbc988a0df7faa39e7f68344361854638d86bcce0ffe3619/langsmith-0.9.8-py3-none-any.whl", hash = "sha256:098da9fc6c184284f17913cb813a41e28c5ab1508e90bd50db40c28166681017", size = 671148, upload-time = "2026-07-06T19:06:08.911Z" }, + { url = "https://files.pythonhosted.org/packages/f1/85/0ad6df25588122760b2b40ec182eafc55532c9e66015c83e16675bd34a29/langsmith-0.11.1-py3-none-any.whl", hash = "sha256:cfc3437a9cf27440cd0095c24df945edbceb6df10b579bc8b3980b4ad367835f", size = 744589, upload-time = "2026-08-19T15:47:42.043Z" }, ] [[package]] @@ -4055,18 +4058,8 @@ wheels = [ [[package]] name = "nemo-relay" -version = "0.7.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/a259aac8df4aa78c0b3a6f3ad0fbf6305666bfaad5c29d9adc26db0f9e27/nemo_relay-0.7.3.tar.gz", hash = "sha256:ea5a1bb52e25e001dcbf6af1830616be181845e978cc848df58562556bba5604", size = 1299430, upload-time = "2026-08-14T14:40:44.832Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ed/9e/4eb80d2307cadcbb839dad2212a8d888667cd8c531ad0d8c0c1afc841181/nemo_relay-0.7.3-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:301dfc8334032ac52c09cc0b1421181e7c4df601abe7c82dfe51aa74ddbb3732", size = 9251369, upload-time = "2026-08-14T14:40:07.668Z" }, - { url = "https://files.pythonhosted.org/packages/c7/a9/9fb77f7142b1381d8c3c81fdbb76782a02dcff1ee403ac6c93b13fa46b24/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:49c70c0a94cebb6cba3dd7521be11f63eac4cd9386881eb29827ac91b1bd780b", size = 8458046, upload-time = "2026-08-14T14:40:10.181Z" }, - { url = "https://files.pythonhosted.org/packages/ae/ab/b2f246f971f561a982d234d9f3ec1b29dfa4c72bf4882f3e32ce6db54dfa/nemo_relay-0.7.3-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8eeaaf8c6a18440e473dd14eb1bb82e56c6514e70adee7456eddf9ce217cff89", size = 8957931, upload-time = "2026-08-14T14:40:12.547Z" }, - { url = "https://files.pythonhosted.org/packages/8f/d8/d8c25ba915467bab457d175b84a3af5c33291655867d472c49eada3b77e6/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:90ff984a89c42ebd0cfe26a0af3f180b1970e2ebcd743d19a960e547949ad2ef", size = 10325640, upload-time = "2026-08-14T14:40:15.722Z" }, - { url = "https://files.pythonhosted.org/packages/7d/fd/6ec48f47eb5ca4cca566b197ade8514baf183dddf53d6de3a296b8bc1102/nemo_relay-0.7.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d4895a4df39a92ecbac3633fb4ec59cf1f83286c0369f433760ab28fb8d6dddf", size = 10706736, upload-time = "2026-08-14T14:40:18.377Z" }, - { url = "https://files.pythonhosted.org/packages/6b/cb/8a6f8d5f9922e75100f135c1dc4451bafc15f5d51608081f72b894e0caf4/nemo_relay-0.7.3-cp311-abi3-win_amd64.whl", hash = "sha256:f123cd45a27fca3d570559f2c26850138763411f66f387559d757e5d88bfa3c1", size = 8807604, upload-time = "2026-08-14T14:40:20.98Z" }, - { url = "https://files.pythonhosted.org/packages/f1/ac/558a05e6b9e28464b64fee8327d9e8a2aab3fb356363b8b82597b6b35b2a/nemo_relay-0.7.3-cp311-abi3-win_arm64.whl", hash = "sha256:6d5444e9a03b8b5d4bba2409105a121349580cb51481d441c6b5e699713a8763", size = 8442944, upload-time = "2026-08-14T14:40:23.433Z" }, -] +version = "0.8.0" +source = { git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0#2557abb5ee87d61fe914b9bc9b8442210920f7a0" } [package.optional-dependencies] deepagents = [ @@ -7508,14 +7501,14 @@ wheels = [ [[package]] name = "wcmatch" -version = "10.1" +version = "11.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "bracex" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/79/3e/c0bdc27cf06f4e47680bd5803a07cb3dfd17de84cde92dd217dcb9e05253/wcmatch-10.1.tar.gz", hash = "sha256:f11f94208c8c8484a16f4f48638a85d771d9513f4ab3f37595978801cb9465af", size = 117421, upload-time = "2025-06-22T19:14:02.49Z" } +sdist = { url = "https://files.pythonhosted.org/packages/57/43/30e407989e313677dbb9d5f045f966549a7254834571e342eaa4b55cc67b/wcmatch-11.0.1.tar.gz", hash = "sha256:1ea2b4fa678b8ca268253798d5963935df39132d47c3e241c0a0732224005e7d", size = 144662, upload-time = "2026-08-14T15:20:40.477Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/eb/d8/0d1d2e9d3fabcf5d6840362adcf05f8cf3cd06a73358140c3a97189238ae/wcmatch-10.1-py3-none-any.whl", hash = "sha256:5848ace7dbb0476e5e55ab63c6bbd529745089343427caa5537f230cc01beb8a", size = 39854, upload-time = "2025-06-22T19:14:00.978Z" }, + { url = "https://files.pythonhosted.org/packages/ce/77/7a02b0f05b3ffcdbef9719ce3ee0b508d6a29b58e95299f1580055671db3/wcmatch-11.0.1-py3-none-any.whl", hash = "sha256:fd149ecddb9f0a88ea780017d6dde17c994e494e7f7303d4e3c9d6251f978f4b", size = 43449, upload-time = "2026-08-14T15:20:39.379Z" }, ] [[package]] From 4cbe83ee83189f9139a1be940a0e5d3f190dc188 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Aug 2026 22:44:52 -0700 Subject: [PATCH 06/13] revert to nim type in configs Signed-off-by: Chantal D Gama Rose --- configs/config_cli_default.yml | 16 ++++-- configs/config_domain_routing_and_skills.yml | 8 ++- configs/config_mcp.yml | 16 ++++-- configs/config_openshell.yml | 14 +++-- configs/config_web_azure_ai_search.yml | 16 ++++-- configs/config_web_default_guardrails.yml | 16 ++++-- configs/config_web_default_llamaindex.yml | 16 ++++-- configs/config_web_frag.yml | 16 ++++-- configs/config_web_frag_mcp_auth.yml | 16 ++++-- configs/config_web_opensearch.yml | 16 ++++-- .../config_web_default_with_pricing.yml | 16 +++--- frontends/cli/cli.py | 18 ++++--- mcp/uv.lock | 4 +- pyproject.toml | 2 +- tests/frontends/test_cli.py | 51 +++++++++++++++++++ uv.lock | 4 +- 16 files changed, 190 insertions(+), 55 deletions(-) create mode 100644 tests/frontends/test_cli.py diff --git a/configs/config_cli_default.yml b/configs/config_cli_default.yml index 1a6609362..871282ef8 100644 --- a/configs/config_cli_default.yml +++ b/configs/config_cli_default.yml @@ -13,7 +13,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -22,11 +22,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -35,9 +37,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -45,9 +49,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -55,6 +61,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false functions: # ========================================================================= diff --git a/configs/config_domain_routing_and_skills.yml b/configs/config_domain_routing_and_skills.yml index 54d118c02..f728e3cc7 100644 --- a/configs/config_domain_routing_and_skills.yml +++ b/configs/config_domain_routing_and_skills.yml @@ -44,7 +44,7 @@ general: llms: nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -52,9 +52,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -62,6 +64,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: diff --git a/configs/config_mcp.yml b/configs/config_mcp.yml index 269bd176c..c6908a37f 100644 --- a/configs/config_mcp.yml +++ b/configs/config_mcp.yml @@ -27,7 +27,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -36,11 +36,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -49,9 +51,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -59,9 +63,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -69,6 +75,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false functions: data_sources: diff --git a/configs/config_openshell.yml b/configs/config_openshell.yml index 1ed09f543..b107e95cb 100644 --- a/configs/config_openshell.yml +++ b/configs/config_openshell.yml @@ -29,7 +29,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -38,11 +38,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: NVIDIA API Catalog-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -51,9 +53,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -61,6 +65,8 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: _type: nim @@ -71,6 +77,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false summary_llm: _type: nim diff --git a/configs/config_web_azure_ai_search.yml b/configs/config_web_azure_ai_search.yml index b4f92032d..1b7a77f5e 100644 --- a/configs/config_web_azure_ai_search.yml +++ b/configs/config_web_azure_ai_search.yml @@ -44,7 +44,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -53,11 +53,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -66,9 +68,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -76,9 +80,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -86,6 +92,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: diff --git a/configs/config_web_default_guardrails.yml b/configs/config_web_default_guardrails.yml index f478196d8..3ec8027f0 100644 --- a/configs/config_web_default_guardrails.yml +++ b/configs/config_web_default_guardrails.yml @@ -29,7 +29,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -38,11 +38,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -51,9 +53,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -61,9 +65,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -71,6 +77,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false summary_llm: diff --git a/configs/config_web_default_llamaindex.yml b/configs/config_web_default_llamaindex.yml index f3d184718..dcb458aae 100644 --- a/configs/config_web_default_llamaindex.yml +++ b/configs/config_web_default_llamaindex.yml @@ -44,7 +44,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -53,11 +53,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -66,9 +68,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -76,9 +80,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -86,6 +92,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: diff --git a/configs/config_web_frag.yml b/configs/config_web_frag.yml index 9843168fc..d803e7082 100644 --- a/configs/config_web_frag.yml +++ b/configs/config_web_frag.yml @@ -46,7 +46,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -55,11 +55,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -68,9 +70,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -78,9 +82,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -88,6 +94,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false functions: # ========================================================================= diff --git a/configs/config_web_frag_mcp_auth.yml b/configs/config_web_frag_mcp_auth.yml index a64cc6451..d2929ab60 100644 --- a/configs/config_web_frag_mcp_auth.yml +++ b/configs/config_web_frag_mcp_auth.yml @@ -53,7 +53,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -62,11 +62,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -75,9 +77,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -85,9 +89,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -95,6 +101,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false functions: # ========================================================================= diff --git a/configs/config_web_opensearch.yml b/configs/config_web_opensearch.yml index 3735fd637..8a649f0db 100644 --- a/configs/config_web_opensearch.yml +++ b/configs/config_web_opensearch.yml @@ -45,7 +45,7 @@ general: llms: nemotron_lightning_intent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -54,11 +54,13 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3.5-lightning-30b-a3b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -67,9 +69,11 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -77,9 +81,11 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: - _type: openai + _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b base_url: "https://integrate.api.nvidia.com/v1" api_key: ${NVIDIA_API_KEY} @@ -87,6 +93,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 + chat_template_kwargs: + enable_thinking: false functions: # ========================================================================= diff --git a/configs/nemo_relay/config_web_default_with_pricing.yml b/configs/nemo_relay/config_web_default_with_pricing.yml index 64c1f4606..f3b337fe5 100644 --- a/configs/nemo_relay/config_web_default_with_pricing.yml +++ b/configs/nemo_relay/config_web_default_with_pricing.yml @@ -53,6 +53,8 @@ llms: max_tokens: 1024 num_retries: 5 parallel_tool_calls: false + chat_template_kwargs: + enable_thinking: false # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. @@ -66,8 +68,8 @@ llms: max_tokens: 8192 num_retries: 5 parallel_tool_calls: false - # chat_template_kwargs: - # enable_thinking: true + chat_template_kwargs: + enable_thinking: true nemotron_ultra_llm: _type: nim @@ -78,8 +80,8 @@ llms: top_p: 0.7 max_tokens: 16384 num_retries: 5 - # chat_template_kwargs: - # enable_thinking: false + chat_template_kwargs: + enable_thinking: false nemotron_ultra_writer_llm: _type: nim @@ -90,8 +92,8 @@ llms: top_p: 0.7 max_tokens: 32768 num_retries: 5 - # chat_template_kwargs: - # enable_thinking: false + chat_template_kwargs: + enable_thinking: false # LLM for document summaries (required when generate_summary: true) summary_llm: @@ -238,7 +240,7 @@ workflow: enabled: true endpoints: - type: openinference - endpoint: ${RELAY_OTEL_ENDPOINT:-http://localhost:6006/v1/traces} + endpoint: http://localhost:6006/v1/traces service_name: aiq-relay resource_attributes: openinference.project.name: aiq-relay diff --git a/frontends/cli/cli.py b/frontends/cli/cli.py index 9a690374f..97efafb11 100644 --- a/frontends/cli/cli.py +++ b/frontends/cli/cli.py @@ -23,6 +23,7 @@ import warnings from pathlib import Path +import nemo_relay import yaml from prompt_toolkit import PromptSession from prompt_toolkit.formatted_text import HTML @@ -372,13 +373,18 @@ def _on_step(step: IntermediateStep) -> None: else: result = await runner.result(to_type=str) - parse_and_display_response(result, verbose=verbose) + # Relay subscriber delivery is asynchronous. Wait until the run context has + # closed its outer scopes before displaying the answer and opening the next + # prompt, otherwise late lifecycle logs can be painted after ``You:``. + await nemo_relay.subscribers.flush_async() - # Check if the response indicates a critical error (e.g., missing API key) - # This is a fallback in case validation didn't catch it earlier - if "Missing Required API Keys" in result or "Missing keys:" in result: - console.print("[bold red]Cannot continue without required API keys. Exiting.[/bold red]") - break + parse_and_display_response(result, verbose=verbose) + + # Check if the response indicates a critical error (e.g., missing API key) + # This is a fallback in case validation didn't catch it earlier + if "Missing Required API Keys" in result or "Missing keys:" in result: + console.print("[bold red]Cannot continue without required API keys. Exiting.[/bold red]") + break except (EOFError, KeyboardInterrupt): console.print("\n\n[bold green]Goodbye! Happy researching![/bold green]") diff --git a/mcp/uv.lock b/mcp/uv.lock index 0e382edee..40f103743 100644 --- a/mcp/uv.lock +++ b/mcp/uv.lock @@ -211,7 +211,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.5.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, - { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0" }, + { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=ffb24817442bac99212da0971b13bdad5bc4d84d" }, { name = "nvidia-nat", extras = ["langchain", "async-endpoints", "phoenix", "mcp"], specifier = "==1.8.0" }, { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, @@ -3307,7 +3307,7 @@ wheels = [ [[package]] name = "nemo-relay" version = "0.8.0" -source = { git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0#2557abb5ee87d61fe914b9bc9b8442210920f7a0" } +source = { git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=ffb24817442bac99212da0971b13bdad5bc4d84d#ffb24817442bac99212da0971b13bdad5bc4d84d" } [package.optional-dependencies] deepagents = [ diff --git a/pyproject.toml b/pyproject.toml index d98b56266..c33b8f8f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -274,7 +274,7 @@ you-com = { workspace = true } duckduckgo-news-search = { workspace = true } polymarket-prediction-market = { workspace = true } knowledge-layer = { workspace = true } -nemo-relay = { git = "https://github.com/NVIDIA/NeMo-Relay.git", rev = "2557abb5ee87d61fe914b9bc9b8442210920f7a0" } # pragma: allowlist secret +nemo-relay = { git = "https://github.com/NVIDIA/NeMo-Relay.git", rev = "ffb24817442bac99212da0971b13bdad5bc4d84d" } # pragma: allowlist secret aiq-api = { workspace = true } aiq-research-cli = { workspace = true } aiq-debug = { workspace = true } diff --git a/tests/frontends/test_cli.py b/tests/frontends/test_cli.py new file mode 100644 index 000000000..493a61b9b --- /dev/null +++ b/tests/frontends/test_cli.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import asynccontextmanager +from types import SimpleNamespace + +import pytest +from aiq_research_cli import cli + + +@pytest.mark.asyncio +async def test_interactive_loop_flushes_relay_before_display_and_next_prompt(monkeypatch) -> None: + events: list[str] = [] + responses = iter(["research this", "q"]) + + async def prompt_async(*args, **kwargs): # noqa: ARG001 + events.append("prompt") + return next(responses) + + async def flush_async() -> None: + events.append("flush") + + class Runner: + async def result(self, *, to_type): # noqa: ARG002 + events.append("result") + return "answer" + + class Session: + @asynccontextmanager + async def run(self, user_input): # noqa: ARG002 + yield Runner() + events.append("run-exit") + + class SessionManager: + @asynccontextmanager + async def session(self, *, user_input_callback): # noqa: ARG002 + yield Session() + + monkeypatch.setattr(cli.prompt_session, "prompt_async", prompt_async) + monkeypatch.setattr(cli.nemo_relay.subscribers, "flush_async", flush_async) + monkeypatch.setattr(cli.console, "print", lambda *args, **kwargs: None) + monkeypatch.setattr(cli, "parse_and_display_response", lambda *args, **kwargs: events.append("display")) + monkeypatch.setattr( + cli.ContextState, + "get", + lambda: SimpleNamespace(conversation_id=SimpleNamespace(set=lambda value: None)), + ) + + await cli.interactive_loop(SessionManager(), verbose=True) + + assert events == ["prompt", "result", "run-exit", "flush", "display", "prompt"] diff --git a/uv.lock b/uv.lock index a72e0141a..2c1f1ce17 100644 --- a/uv.lock +++ b/uv.lock @@ -299,7 +299,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.5.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, - { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0" }, + { name = "nemo-relay", extras = ["deepagents", "langchain", "langgraph"], git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=ffb24817442bac99212da0971b13bdad5bc4d84d" }, { name = "nvidia-nat", extras = ["langchain", "async-endpoints", "phoenix", "mcp"], specifier = "==1.8.0" }, { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, @@ -4059,7 +4059,7 @@ wheels = [ [[package]] name = "nemo-relay" version = "0.8.0" -source = { git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=2557abb5ee87d61fe914b9bc9b8442210920f7a0#2557abb5ee87d61fe914b9bc9b8442210920f7a0" } +source = { git = "https://github.com/NVIDIA/NeMo-Relay.git?rev=ffb24817442bac99212da0971b13bdad5bc4d84d#ffb24817442bac99212da0971b13bdad5bc4d84d" } [package.optional-dependencies] deepagents = [ From 0d61f96967610307b3732d0da31bfa04aa7643cd Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Aug 2026 23:32:53 -0700 Subject: [PATCH 07/13] fix: address Relay review feedback Signed-off-by: Chantal D Gama Rose --- configs/config_web_default_llamaindex.yml | 2 - .../config_web_default_with_pricing.yml | 2 +- frontends/aiq_api/src/aiq_api/jobs/runner.py | 8 +- .../configs/config_tokenomics_pricing.yml | 37 --------- frontends/cli/cli.py | 5 +- mcp/scripts/check_license_inventory.py | 10 +++ mcp/tests/test_config_and_packaging.py | 5 +- src/aiq_agent/agents/chat_researcher/agent.py | 26 ++++-- .../agents/chat_researcher/register.py | 3 +- src/aiq_agent/agents/deep_researcher/agent.py | 6 +- .../agents/deep_researcher/tools/research.py | 13 ++- .../agents/shallow_researcher/agent.py | 5 +- src/aiq_agent/relay/__init__.py | 12 +++ src/aiq_agent/relay/bootstrap.py | 12 +++ src/aiq_agent/relay/config.py | 12 +++ src/aiq_agent/relay/logging.py | 12 +++ src/aiq_agent/relay/privacy.py | 12 +++ src/aiq_agent/relay/runtime.py | 80 ++++++++++++------- src/aiq_agent/tokenomics/atof_adapter.py | 24 +++++- .../deep_researcher/test_custom_middleware.py | 6 +- .../agents/deep_researcher/test_factory.py | 3 +- .../agents/test_config_observability.py | 24 ------ tests/aiq_agent/common/test_callbacks.py | 3 +- tests/aiq_agent/jobs/test_runner.py | 14 ++++ tests/frontends/test_cli.py | 42 +++++++++- tests/scripts/test_start_cli.py | 33 ++++++++ tests/test_relay_runtime.py | 38 ++++++++- tests/tokenomics/test_atof_adapter.py | 36 +++++++++ 28 files changed, 367 insertions(+), 118 deletions(-) delete mode 100644 frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml delete mode 100644 tests/aiq_agent/agents/test_config_observability.py create mode 100644 tests/scripts/test_start_cli.py diff --git a/configs/config_web_default_llamaindex.yml b/configs/config_web_default_llamaindex.yml index dcb458aae..ff4d235d1 100644 --- a/configs/config_web_default_llamaindex.yml +++ b/configs/config_web_default_llamaindex.yml @@ -56,8 +56,6 @@ llms: chat_template_kwargs: enable_thinking: false - # Known limitation: Build-hosted Lightning can intermittently produce citation-incomplete shallow drafts. - # AI-Q fails closed; see docs/source/resources/troubleshooting.md#nemotron-35-lightning-on-nvidia-api-catalog. nemotron_lightning_agent_llm: _type: nim model_name: nvidia/nemotron-3-ultra-550b-a55b diff --git a/configs/nemo_relay/config_web_default_with_pricing.yml b/configs/nemo_relay/config_web_default_with_pricing.yml index f3b337fe5..7e6ef633e 100644 --- a/configs/nemo_relay/config_web_default_with_pricing.yml +++ b/configs/nemo_relay/config_web_default_with_pricing.yml @@ -248,4 +248,4 @@ workflow: pricing: sources: - type: file - path: configs/nemo_relay/relay_pricing_catalog.json + path: ${AIQ_RELAY_PRICING_CATALOG:-/app/configs/nemo_relay/relay_pricing_catalog.json} diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 07e81efeb..2c34749cf 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -198,6 +198,12 @@ async def _ensure_relay_started_for_job(relay_config: Any, job_id: str) -> None: logger.warning("Relay startup failed for job %s (error_type=%s)", job_id, type(exc).__name__) +def _resolve_job_relay_config(config: Any, function_config: Any) -> Any: + """Prefer workflow Relay settings for a separately executed async agent.""" + workflow_config = getattr(config, "workflow", None) + return getattr(workflow_config, "relay", None) or getattr(function_config, "relay", None) + + def _db_now_expr(db_url: str) -> str: """Return the DB current-time SQL expression for this backend. @@ -788,7 +794,7 @@ async def run_agent_job( await _attach_middleware_to_function(builder, config, agent_config_name) fn_config = builder.get_function_config(agent_config_name) - relay_config = getattr(fn_config, "relay", None) + relay_config = _resolve_job_relay_config(config, fn_config) if relay_config is not None: await _ensure_relay_started_for_job(relay_config, job_id) if getattr(fn_config, "type", None) == "deep_research_agent": diff --git a/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml b/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml deleted file mode 100644 index 35d1196a9..000000000 --- a/frontends/benchmarks/deepresearch_bench/configs/config_tokenomics_pricing.yml +++ /dev/null @@ -1,37 +0,0 @@ -# Tokenomics pricing only (not loaded by `nat eval`). -# -# NAT's top-level config schema rejects unknown keys; keep `tokenomics` here and pass -# this file to `python -m aiq_agent.tokenomics.report --config ...`. -# -# Prices are USD per 1 million tokens. Tool costs are USD per invocation. -# -# Optional: mirror `eval.general.output_dir` from your profiling eval config so the -# report defaults to the same folder as the trace when you omit `--output`. - -eval: - general: - output_dir: frontends/benchmarks/deepresearch_bench/results - -tokenomics: - pricing: - models: - "nvidia/nemotron-3-ultra-550b-a55b": - # NVIDIA-hosted API access used by this profile. Self-hosting is not free. - input_per_1m_tokens: 0.00 - output_per_1m_tokens: 0.00 - "nvidia/nemotron-3.5-lightning-30b-a3b": - # NVIDIA-hosted API access used by this profile. Self-hosting is not free. - input_per_1m_tokens: 0.00 - output_per_1m_tokens: 0.00 - tools: - # Tavily pay-as-you-go is $0.008/credit: basic uses 1 credit, - # advanced uses 2. Monthly plans have lower effective rates. - "web_search_tool": - cost_per_call: 0.008 - "advanced_web_search_tool": - cost_per_call: 0.016 - # The shipped paper-search tool defaults to Serper. This is Serper's - # Starter rate ($50 / 50,000 successful queries); change it for your tier - # or when selecting SerpAPI/SearchAPI instead. - "paper_search": - cost_per_call: 0.001 diff --git a/frontends/cli/cli.py b/frontends/cli/cli.py index 97efafb11..8b74b5f10 100644 --- a/frontends/cli/cli.py +++ b/frontends/cli/cli.py @@ -376,7 +376,10 @@ def _on_step(step: IntermediateStep) -> None: # Relay subscriber delivery is asynchronous. Wait until the run context has # closed its outer scopes before displaying the answer and opening the next # prompt, otherwise late lifecycle logs can be painted after ``You:``. - await nemo_relay.subscribers.flush_async() + try: + await nemo_relay.subscribers.flush_async() + except Exception as error: # noqa: BLE001 - telemetry must not suppress the answer + logger.warning("NeMo Relay flush failed (error_type=%s)", type(error).__name__) parse_and_display_response(result, verbose=verbose) diff --git a/mcp/scripts/check_license_inventory.py b/mcp/scripts/check_license_inventory.py index 0a301c84b..68365e07a 100644 --- a/mcp/scripts/check_license_inventory.py +++ b/mcp/scripts/check_license_inventory.py @@ -34,6 +34,13 @@ ("knowledge-layer", "1.0.0"): "../sources/knowledge_layer", ("tavily-web-search", "1.0.0"): "../sources/tavily_web_search", } +_APPROVED_GIT_SOURCES = { + ("nemo-relay", "0.8.0"): ( + "https://github.com/NVIDIA/NeMo-Relay.git" + "?rev=ffb24817442bac99212da0971b13bdad5bc4d84d" + "#ffb24817442bac99212da0971b13bdad5bc4d84d" + ) +} _MCP_LOCK_PATH = Path(__file__).resolve().parents[1] / "uv.lock" # These distributions omit license metadata but bundle a license file. A version @@ -227,6 +234,9 @@ def validate_lock_sources(lock_path: Path = _MCP_LOCK_PATH) -> None: if source != {"registry": "https://pypi.org/simple"}: raise ValueError(f"dependency uses an unapproved registry source: {name}=={version}") continue + if source.get("git") is not None: + if source == {"git": _APPROVED_GIT_SOURCES.get((name, version))}: + continue editable = source.get("editable") if not isinstance(editable, str): raise ValueError(f"dependency uses an unapproved lock source: {name}=={version}") diff --git a/mcp/tests/test_config_and_packaging.py b/mcp/tests/test_config_and_packaging.py index 4136ca95c..f1be3e389 100644 --- a/mcp/tests/test_config_and_packaging.py +++ b/mcp/tests/test_config_and_packaging.py @@ -104,8 +104,6 @@ def test_public_mcp_config_uses_only_public_models_sources_and_environment_names config = yaml.safe_load(_CONFIG_PATH.read_text()) text = _CONFIG_PATH.read_text().lower() - llms = config["llms"].values() - assert {entry["_type"] for entry in llms} == {"openai"} assert {entry["base_url"] for entry in config["llms"].values()} == {"https://integrate.api.nvidia.com/v1"} assert all(entry["model_name"].startswith("nvidia/") for entry in config["llms"].values()) assert config["functions"]["web_search_tool"]["_type"] == "tavily_web_search" @@ -169,6 +167,9 @@ def test_root_workspace_excludes_the_independent_mcp_project() -> None: ) } continue + if "git" in source: + assert package["name"] == "nemo-relay" + continue editable = source.get("editable") assert isinstance(editable, str) assert not Path(editable).is_absolute() diff --git a/src/aiq_agent/agents/chat_researcher/agent.py b/src/aiq_agent/agents/chat_researcher/agent.py index 5aea3572a..51cc29559 100644 --- a/src/aiq_agent/agents/chat_researcher/agent.py +++ b/src/aiq_agent/agents/chat_researcher/agent.py @@ -177,10 +177,14 @@ def _build_graph(self) -> CompiledStateGraph: async def intent_classifier_node(state: ChatResearcherState) -> dict[str, Any]: try: return await run_agent( - "intent_classifier", - lambda: self.intent_classifier_fn(state), - input_value=state, - ) + "intent_classifier", + lambda: self.intent_classifier_fn(state), + input_value={ + "message_count": len(state.messages), + "data_source_count": len(state.data_sources or []), + "has_active_report": bool(state.active_report_job_id), + }, + ) except Exception as error: logger.warning("Intent routing failed (error_type=%s)", type(error).__name__) return { @@ -665,11 +669,23 @@ async def _invoke_graph() -> dict[str, Any]: effective_config = dict(graph_config or {}) return await self._graph.ainvoke(input_state, config=effective_config) + input_data_sources = ( + input_state.get("data_sources") if isinstance(input_state, dict) else input_state.data_sources + ) + relay_input_metadata = { + "message_count": len(messages), + "data_source_count": len(input_data_sources or []), + "has_active_report": bool( + input_state.get("active_report_job_id") + if isinstance(input_state, dict) + else input_state.active_report_job_id + ), + } result = await run_agent( "chat_deepresearcher_agent", _invoke_graph, session_id=thread_id, - input_value=input_state, + input_value=relay_input_metadata, ) logger.info("ChatResearcherAgent: Workflow complete") diff --git a/src/aiq_agent/agents/chat_researcher/register.py b/src/aiq_agent/agents/chat_researcher/register.py index c06e25f9a..8b6ddf693 100644 --- a/src/aiq_agent/agents/chat_researcher/register.py +++ b/src/aiq_agent/agents/chat_researcher/register.py @@ -26,6 +26,7 @@ from pydantic import Field from pydantic import ValidationError +from aiq_agent.common import VerboseTraceCallback from aiq_agent.common import _create_chat_response from aiq_agent.common import format_data_source_tools from aiq_agent.common import get_checkpointer @@ -284,7 +285,7 @@ async def context_aware_intent_router(config: ContextAwareIntentRouterConfig, bu raise ValueError("context_aware_intent_router requires exactly one catalog tool") prompt = load_prompt(Path(__file__).parent / "prompts", "context_aware_intent_router.j2") - callbacks = [VerboseTraceCallback()] if is_verbose(config.verbose) else [] + callbacks = [VerboseTraceCallback()] if config.verbose else [] router = ContextAwareIntentRouter( llm, tools[0], diff --git a/src/aiq_agent/agents/deep_researcher/agent.py b/src/aiq_agent/agents/deep_researcher/agent.py index a8c0613c0..4f79cbb8e 100644 --- a/src/aiq_agent/agents/deep_researcher/agent.py +++ b/src/aiq_agent/agents/deep_researcher/agent.py @@ -322,7 +322,11 @@ async def _invoke_orchestrator() -> Any: result = await run_agent( "deep_research_agent", _invoke_orchestrator, - input_value=state, + input_value={ + "message_count": len(messages), + "query_character_count": len(query), + "file_count": len(state.files), + }, ) except TimeoutError as exc: # An inner provider/tool may raise TimeoutError for its own operation. diff --git a/src/aiq_agent/agents/deep_researcher/tools/research.py b/src/aiq_agent/agents/deep_researcher/tools/research.py index 36a18b509..79b39259f 100644 --- a/src/aiq_agent/agents/deep_researcher/tools/research.py +++ b/src/aiq_agent/agents/deep_researcher/tools/research.py @@ -45,6 +45,12 @@ _NO_TOOL_RUNTIME = cast(ToolRuntime, None) logger = logging.getLogger(__name__) + + +class _MissingStructuredResponseError(ValueError): + """Raised when a researcher worker returns no structured response.""" + + _NOTE_SLUG_MAX_LENGTH = 64 RESEARCHER_AGENT_NAME = "researcher-agent" @@ -141,12 +147,11 @@ async def _run_research_query( try: structured = result.get("structured_response") if isinstance(result, dict) else None if structured is None: - raise ValueError("researcher worker did not return structured ResearchNotes") + raise _MissingStructuredResponseError("researcher worker did not return structured ResearchNotes") note = ResearchNotes.model_validate(structured) + except _MissingStructuredResponseError: + raise except Exception as exc: # noqa: BLE001 - captured as per-item failure - missing_response = "researcher worker did not return structured ResearchNotes" - if isinstance(exc, ValueError) and str(exc) == missing_response: - raise logger.warning( "Researcher worker returned invalid ResearchNotes (error_type=%s, query_%s)", type(exc).__name__, diff --git a/src/aiq_agent/agents/shallow_researcher/agent.py b/src/aiq_agent/agents/shallow_researcher/agent.py index 06d2c70c7..6036de980 100644 --- a/src/aiq_agent/agents/shallow_researcher/agent.py +++ b/src/aiq_agent/agents/shallow_researcher/agent.py @@ -35,7 +35,6 @@ from langgraph.graph.state import CompiledStateGraph from langgraph.prebuilt import ToolNode from langgraph.prebuilt import tools_condition -from nemo_relay.integrations.langchain import NemoRelayMiddleware from aiq_agent.common import get_source_id_for_tool from aiq_agent.common import load_prompt @@ -52,6 +51,7 @@ from aiq_agent.common.logging_utils import log_content_metadata from aiq_agent.relay import ainvoke_with_relay from aiq_agent.relay import run_agent +from aiq_agent.relay.runtime import awrap_tool_call_with_relay from ...common import LLMProvider from ...common import LLMRole @@ -500,8 +500,7 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: builder.set_entry_point("agent") - relay_middleware = NemoRelayMiddleware() - tool_node = ToolNode(self.tools, awrap_tool_call=relay_middleware.awrap_tool_call) + tool_node = ToolNode(self.tools, awrap_tool_call=awrap_tool_call_with_relay) # Per-agent allowlist mirrors the deep researcher: only tools this # agent was loaded with are candidates for source capture. The diff --git a/src/aiq_agent/relay/__init__.py b/src/aiq_agent/relay/__init__.py index ce35049f2..b8cea7d82 100644 --- a/src/aiq_agent/relay/__init__.py +++ b/src/aiq_agent/relay/__init__.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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. """AI-Q's NeMo Relay integration boundary. diff --git a/src/aiq_agent/relay/bootstrap.py b/src/aiq_agent/relay/bootstrap.py index aaf303811..46f3d5148 100644 --- a/src/aiq_agent/relay/bootstrap.py +++ b/src/aiq_agent/relay/bootstrap.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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. """AI-Q-owned lifecycle for NeMo Relay's plugin host.""" diff --git a/src/aiq_agent/relay/config.py b/src/aiq_agent/relay/config.py index 8e14a48d5..08a9bb6f5 100644 --- a/src/aiq_agent/relay/config.py +++ b/src/aiq_agent/relay/config.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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. """Typed AI-Q configuration for NeMo Relay plugins.""" diff --git a/src/aiq_agent/relay/logging.py b/src/aiq_agent/relay/logging.py index 0dfeb1731..518a4f9b8 100644 --- a/src/aiq_agent/relay/logging.py +++ b/src/aiq_agent/relay/logging.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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. """Developer-safe console logging for NeMo Relay lifecycle events.""" diff --git a/src/aiq_agent/relay/privacy.py b/src/aiq_agent/relay/privacy.py index 0bd70780f..df0f6181d 100644 --- a/src/aiq_agent/relay/privacy.py +++ b/src/aiq_agent/relay/privacy.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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. """Request-scoped privacy controls for Relay observability payloads.""" diff --git a/src/aiq_agent/relay/runtime.py b/src/aiq_agent/relay/runtime.py index 5465cae4b..61c04443e 100644 --- a/src/aiq_agent/relay/runtime.py +++ b/src/aiq_agent/relay/runtime.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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. """NeMo Relay framework integration helpers.""" @@ -126,11 +138,20 @@ async def ainvoke_with_relay( messages.pop(0) if messages and isinstance(messages[0], BaseMessage) and messages[0].type == "system" else None ) model, model_settings, effective_config = _normalize_chat_nvidia_binding(runnable, effective_config) + named_source = getattr(model, "bound", model) + resolved_name = next( + ( + value + for attribute in ("model", "model_name", "model_id", "deployment_name") + if isinstance(value := getattr(named_source, attribute, None), str) and value + ), + None, + ) if not any( isinstance(getattr(model, attribute, None), str) and getattr(model, attribute) for attribute in ("model", "model_name", "model_id", "deployment_name") ): - model = _NamedModelAdapter(model, type(model).__name__) + model = _NamedModelAdapter(model, resolved_name or type(model).__name__) request = ModelRequest( model=model, messages=messages, @@ -193,29 +214,19 @@ async def invoke(next_request: ModelRequest[Any]) -> ModelResponse[Any]: return response.result[-1] -async def ainvoke_tool_with_relay(tool: Any, args: dict[str, Any]) -> Any: - """Run a direct LangChain tool call through Relay's maintained middleware.""" - request = ToolCallRequest( - tool_call={"name": tool.name, "args": args, "id": f"aiq-{uuid4()}"}, - tool=tool, - state={}, - runtime=None, - ) - - async def invoke_call(next_request: ToolCallRequest) -> Any: - if next_request.tool is None: - raise RuntimeError(f"Relay-managed tool {next_request.tool_call['name']!r} is unavailable") - return await next_request.tool.ainvoke(next_request.tool_call.get("args") or {}) - +async def awrap_tool_call_with_relay( + request: ToolCallRequest, handler: Callable[[ToolCallRequest], Awaitable[_T]] +) -> _T: + """Capture a LangChain tool call while preserving execution if Relay setup fails.""" invocation_started = False invocation_error: BaseException | None = None - async def invoke(next_request: ToolCallRequest) -> Any: + async def invoke(next_request: ToolCallRequest) -> _T: nonlocal invocation_error nonlocal invocation_started invocation_started = True try: - return await invoke_call(next_request) + return await handler(next_request) except BaseException as error: invocation_error = error raise @@ -231,6 +242,23 @@ async def invoke(next_request: ToolCallRequest) -> Any: return await invoke(request) +async def ainvoke_tool_with_relay(tool: Any, args: dict[str, Any]) -> Any: + """Run a direct LangChain tool call through Relay's maintained middleware.""" + request = ToolCallRequest( + tool_call={"name": tool.name, "args": args, "id": f"aiq-{uuid4()}"}, + tool=tool, + state={}, + runtime=None, + ) + + async def invoke_call(next_request: ToolCallRequest) -> Any: + if next_request.tool is None: + raise RuntimeError(f"Relay-managed tool {next_request.tool_call['name']!r} is unavailable") + return await next_request.tool.ainvoke(next_request.tool_call.get("args") or {}) + + return await awrap_tool_call_with_relay(request, invoke_call) + + @contextmanager def _semantic_scope( name: str, @@ -337,14 +365,7 @@ async def _run() -> _T: lifecycle.output = result return result - if _aiq_scope_active.get(): - return await _run() - - async def _run_isolated() -> _T: - with nemo_relay.use_scope_stack(nemo_relay.create_scope_stack()): - return await _run() - - return await asyncio.create_task(_run_isolated()) + return await _run_at_request_boundary(_run) async def run_workflow( @@ -368,12 +389,17 @@ async def _run() -> _T: lifecycle.output = result return result + return await _run_at_request_boundary(_run) + + +async def _run_at_request_boundary(operation: Callable[[], Awaitable[_T]]) -> _T: + """Reuse a nested scope or isolate a new request on its own task and stack.""" if _aiq_scope_active.get(): - return await _run() + return await operation() async def _run_isolated() -> _T: with nemo_relay.use_scope_stack(nemo_relay.create_scope_stack()): - return await _run() + return await operation() return await asyncio.create_task(_run_isolated()) diff --git a/src/aiq_agent/tokenomics/atof_adapter.py b/src/aiq_agent/tokenomics/atof_adapter.py index 04382ac6b..bfba754e0 100644 --- a/src/aiq_agent/tokenomics/atof_adapter.py +++ b/src/aiq_agent/tokenomics/atof_adapter.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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. """Convert NeMo Relay ATOF JSONL events into tokenomics request profiles.""" @@ -126,8 +138,9 @@ def _root_uuid(event: dict[str, Any], starts: dict[str, dict[str, Any]]) -> str def _usage(event: dict[str, Any]) -> tuple[int, int, int, int, float | None]: - profile = event.get("category_profile") or {} - annotated = profile.get("annotated_response") if isinstance(profile, dict) else {} + profile = event.get("category_profile") + profile = profile if isinstance(profile, dict) else {} + annotated = profile.get("annotated_response") annotated = annotated if isinstance(annotated, dict) else {} usage = annotated.get("usage") or profile.get("usage") or {} usage = usage if isinstance(usage, dict) else {} @@ -274,10 +287,15 @@ def parse_trace(path: str, pricing: PricingRegistry) -> list[RequestProfile]: ] roots.sort(key=lambda event: _timestamp(event.get("timestamp"))) + events_by_root: dict[str, list[dict[str, Any]]] = {} + for event in events: + if (root_uuid := _root_uuid(event, starts)) is not None: + events_by_root.setdefault(root_uuid, []).append(event) + profiles: list[RequestProfile] = [] for request_index, root in enumerate(roots): root_uuid = root["uuid"] - request_events = [event for event in events if _root_uuid(event, starts) == root_uuid] + request_events = events_by_root.get(root_uuid, []) try: profiles.append(_parse_request(request_index, root, request_events, starts, pricing)) except Exception as exc: diff --git a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py index d96f861ac..933cdf955 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py +++ b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py @@ -1031,6 +1031,8 @@ async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): response_metadata = {"model_name": "nvidia/nemotron-3-ultra-550b-a55b", "finish_reason": "tool_calls"} usage_metadata = {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120} + expected_response_metadata = dict(response_metadata) + expected_usage_metadata = dict(usage_metadata) ai_msg = AIMessage( content="", additional_kwargs={ @@ -1062,8 +1064,8 @@ async def test_awrap_model_call_sanitizes_tool_calls(self, middleware): assert message.tool_calls[0]["name"] == "advanced_web_search_tool" assert message.additional_kwargs["tool_calls"][0]["function"]["name"] == "advanced_web_search_tool" assert message.additional_kwargs["provider_field"] == "preserve-me" - assert message.response_metadata == response_metadata - assert message.usage_metadata == usage_metadata + assert message.response_metadata == expected_response_metadata + assert message.usage_metadata == expected_usage_metadata @pytest.mark.asyncio async def test_awrap_model_call_no_tool_calls_passthrough(self, middleware): diff --git a/tests/aiq_agent/agents/deep_researcher/test_factory.py b/tests/aiq_agent/agents/deep_researcher/test_factory.py index 5118c93e2..805cd887f 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_factory.py +++ b/tests/aiq_agent/agents/deep_researcher/test_factory.py @@ -16,6 +16,7 @@ """Tests for deep researcher graph and middleware factory helpers.""" from unittest.mock import MagicMock +from unittest.mock import NonCallableMagicMock from unittest.mock import patch from deepagents.middleware.filesystem import _apply_permissions_to_ls_results @@ -500,7 +501,7 @@ class FakeSummarizationMiddleware(AgentMiddleware): researcher_agent = MagicMock() researcher_model = MagicMock() shared_middleware = [MagicMock(name="shared_middleware")] - backend = MagicMock() + backend = NonCallableMagicMock() with ( patch( diff --git a/tests/aiq_agent/agents/test_config_observability.py b/tests/aiq_agent/agents/test_config_observability.py deleted file mode 100644 index 7e3857471..000000000 --- a/tests/aiq_agent/agents/test_config_observability.py +++ /dev/null @@ -1,24 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import pytest - -from aiq_agent.agents.chat_researcher.register import ChatDeepResearcherConfig -from aiq_agent.agents.chat_researcher.register import IntentClassifierConfig -from aiq_agent.agents.clarifier.register import ClarifierConfig -from aiq_agent.agents.deep_researcher.register import DeepResearchAgentConfig -from aiq_agent.agents.shallow_researcher.register import ShallowResearchAgentConfig - - -@pytest.mark.parametrize( - "config_type", - [ - IntentClassifierConfig, - ChatDeepResearcherConfig, - ClarifierConfig, - ShallowResearchAgentConfig, - DeepResearchAgentConfig, - ], -) -def test_agent_configs_do_not_expose_legacy_verbose_switch(config_type: type) -> None: - assert "verbose" not in config_type.model_fields diff --git a/tests/aiq_agent/common/test_callbacks.py b/tests/aiq_agent/common/test_callbacks.py index 67ff51ddd..c43248572 100644 --- a/tests/aiq_agent/common/test_callbacks.py +++ b/tests/aiq_agent/common/test_callbacks.py @@ -41,8 +41,9 @@ def test_research_logger_init_with_verbose_param(self, mock_logger): logger_non_verbose = ResearchLogger(mock_logger, verbose=False) assert logger_non_verbose.verbose is False - def test_research_logger_defaults_to_non_verbose(self, mock_logger): + def test_research_logger_defaults_to_non_verbose(self, mock_logger, monkeypatch): """Research logging does not consult a process-global verbosity switch.""" + monkeypatch.setenv("AIQ_VERBOSE", "true") assert ResearchLogger(mock_logger).verbose is False def test_section_logs_info(self, mock_logger): diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 021440044..9fdd52ee6 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -109,6 +109,20 @@ async def never_starts(_config): assert "TimeoutError" in caplog.text +def test_async_job_uses_workflow_relay_config() -> None: + """Async agents inherit the outer workflow's effective Relay settings.""" + from types import SimpleNamespace + + from aiq_api.jobs.runner import _resolve_job_relay_config + + workflow_relay = object() + agent_relay = object() + config = SimpleNamespace(workflow=SimpleNamespace(relay=workflow_relay)) + + assert _resolve_job_relay_config(config, SimpleNamespace(relay=agent_relay)) is workflow_relay + assert _resolve_job_relay_config(SimpleNamespace(), SimpleNamespace(relay=agent_relay)) is agent_relay + + @pytest.fixture(name="content_encryption_manager_guard") def fixture_content_encryption_manager_guard(): """Reset content-encryption globals even when a test assertion fails.""" diff --git a/tests/frontends/test_cli.py b/tests/frontends/test_cli.py index 493a61b9b..496cfdafb 100644 --- a/tests/frontends/test_cli.py +++ b/tests/frontends/test_cli.py @@ -1,6 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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 sys from contextlib import asynccontextmanager from types import SimpleNamespace @@ -8,8 +22,32 @@ from aiq_research_cli import cli +@pytest.mark.parametrize( + ("arguments", "expected_level"), + [([], logging.WARNING), (["--verbose"], logging.INFO)], +) +def test_main_sets_log_level_from_verbose_flag(monkeypatch, arguments: list[str], expected_level: int) -> None: + observed: dict[str, int] = {} + + class LoggingConfigured(Exception): + pass + + def configure_logging(**kwargs) -> None: + observed["level"] = kwargs["level"] + raise LoggingConfigured + + monkeypatch.setattr(sys, "argv", ["aiq-research", *arguments]) + monkeypatch.setattr(cli.logging, "basicConfig", configure_logging) + + with pytest.raises(LoggingConfigured): + cli.main() + + assert observed == {"level": expected_level} + + +@pytest.mark.parametrize("flush_fails", [False, True]) @pytest.mark.asyncio -async def test_interactive_loop_flushes_relay_before_display_and_next_prompt(monkeypatch) -> None: +async def test_interactive_loop_flushes_relay_before_display_and_next_prompt(monkeypatch, flush_fails: bool) -> None: events: list[str] = [] responses = iter(["research this", "q"]) @@ -19,6 +57,8 @@ async def prompt_async(*args, **kwargs): # noqa: ARG001 async def flush_async() -> None: events.append("flush") + if flush_fails: + raise RuntimeError("private relay failure") class Runner: async def result(self, *, to_type): # noqa: ARG002 diff --git a/tests/scripts/test_start_cli.py b/tests/scripts/test_start_cli.py new file mode 100644 index 000000000..4de6ca532 --- /dev/null +++ b/tests/scripts/test_start_cli.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the local CLI launcher.""" + +import os +import shutil +import subprocess +from pathlib import Path + + +def test_start_cli_forwards_verbose_flag(tmp_path: Path) -> None: + """The shell launcher forwards verbose mode to the Python CLI.""" + scripts_dir = tmp_path / "scripts" + bin_dir = tmp_path / ".venv" / "bin" + scripts_dir.mkdir() + bin_dir.mkdir(parents=True) + launcher = scripts_dir / "start_cli.sh" + shutil.copy2(Path("scripts/start_cli.sh"), launcher) + (bin_dir / "activate").write_text("", encoding="utf-8") + arguments_path = tmp_path / "arguments.txt" + fake_cli = bin_dir / "aiq-research" + fake_cli.write_text('#!/bin/bash\nprintf "%s\\n" "$@" > "$AIQ_TEST_ARGUMENTS_PATH"\n', encoding="utf-8") + fake_cli.chmod(0o755) + env = {**os.environ, "AIQ_TEST_ARGUMENTS_PATH": str(arguments_path)} + + subprocess.run([launcher, "--verbose"], cwd=tmp_path, env=env, check=True, capture_output=True, text=True) + + assert arguments_path.read_text(encoding="utf-8").splitlines() == [ + "--config_file", + "configs/config_cli_default.yml", + "--verbose", + ] diff --git a/tests/test_relay_runtime.py b/tests/test_relay_runtime.py index fb2cc55d3..b208a2b7e 100644 --- a/tests/test_relay_runtime.py +++ b/tests/test_relay_runtime.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, 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 asyncio import json @@ -283,6 +295,30 @@ async def passthrough(_self, request, handler): assert calls == [["question"], [message]] +@pytest.mark.asyncio +async def test_relay_model_name_comes_from_bound_model(monkeypatch) -> None: + observed_names: list[str] = [] + + class BoundModel: + model_name = "provider-model" + + class Binding: + bound = BoundModel() + + async def ainvoke(self, messages, config=None): # noqa: ARG002 + return AIMessage(content="done") + + async def capture_name(_self, request, handler): + observed_names.append(request.model.model_name) + return await handler(request) + + monkeypatch.setattr("aiq_agent.relay.runtime.NemoRelayMiddleware.awrap_model_call", capture_name) + + await ainvoke_with_relay(Binding(), []) + + assert observed_names == ["provider-model"] + + @pytest.mark.asyncio async def test_relay_model_middleware_fallback_does_not_retry_started_calls(monkeypatch, caplog) -> None: calls = 0 @@ -864,7 +900,7 @@ def log_message(self, format: str, *args: object) -> None: server_thread.join(timeout=5) assert not server_thread.is_alive() - assert len(received) == 3 + assert len(received) >= 3 assert {path for path, _, _ in received} == { "/v1/traces?projection=openinference", "/v1/traces?projection=full", diff --git a/tests/tokenomics/test_atof_adapter.py b/tests/tokenomics/test_atof_adapter.py index 0d7ce0406..d95beef8b 100644 --- a/tests/tokenomics/test_atof_adapter.py +++ b/tests/tokenomics/test_atof_adapter.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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. """Tests for Relay ATOF tokenomics post-processing.""" @@ -184,6 +196,30 @@ def test_parse_trace_ignores_non_string_identifiers(tmp_path: Path) -> None: assert len(profiles) == 1 +def test_parse_trace_tolerates_non_mapping_category_profile(tmp_path: Path) -> None: + events = [ + _scope( + "root", + "function", + "workflow", + "start", + "2026-01-01T00:00:00Z", + metadata={"aiq.component.type": "workflow"}, + ), + _scope("llm", "llm", "test-model", "start", "2026-01-01T00:00:01Z", parent_uuid="root"), + _scope("llm", "llm", "test-model", "end", "2026-01-01T00:00:02Z", parent_uuid="root"), + _scope("root", "function", "workflow", "end", "2026-01-01T00:00:03Z"), + ] + events[2]["category_profile"] = ["malformed"] + path = tmp_path / "relay.atof.jsonl" + _write(path, events) + + profiles = parse_trace(str(path), _pricing()) + + assert len(profiles) == 1 + assert profiles[0].total_llm_calls == 1 + + def test_parse_trace_failure_log_excludes_exception_content(tmp_path: Path, monkeypatch, caplog) -> None: events = [ _scope( From da1581ffd5576ca35c211d4e97ec90606b0e986d Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Aug 2026 23:46:59 -0700 Subject: [PATCH 08/13] fix: allow pinned Relay source in MCP SBOM Signed-off-by: Chantal D Gama Rose --- mcp/scripts/check_license_inventory.py | 13 +++++++++++-- mcp/tests/test_release_checks.py | 8 +++++++- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/mcp/scripts/check_license_inventory.py b/mcp/scripts/check_license_inventory.py index 68365e07a..0d5a63692 100644 --- a/mcp/scripts/check_license_inventory.py +++ b/mcp/scripts/check_license_inventory.py @@ -15,6 +15,7 @@ from importlib.metadata import distribution from pathlib import Path from typing import Any +from urllib.parse import quote _DIRECT_RUNTIME_DEPENDENCIES = { "aiq-agent", @@ -104,6 +105,14 @@ def _canonicalize(name: str) -> str: return re.sub(r"[-_.]+", "-", name).lower() +def _expected_component_purl(name: str, version: str) -> str: + purl = f"pkg:pypi/{name}@{version}" + git_source = _APPROVED_GIT_SOURCES.get((name, version)) + if git_source is not None: + return f"{purl}?vcs_url={quote(git_source, safe=':/')}" + return purl + + def _license_files(dist: Any) -> list[dict[str, str]]: evidence: dict[str, dict[str, str]] = {} for relative in dist.files or (): @@ -211,8 +220,8 @@ def validate_sbom(sbom: dict[str, Any]) -> None: observed_local.add(record) continue - if component.get("purl") != f"pkg:pypi/{name}@{version}": - raise ValueError(f"dependency is not from the public PyPI source contract: {name}=={version}") + if component.get("purl") != _expected_component_purl(name, version): + raise ValueError(f"dependency is not from the approved source contract: {name}=={version}") if observed_local != set(_LOCAL_SOURCE_COMPONENTS): raise ValueError("SBOM local component set differs from the approved public source contract") diff --git a/mcp/tests/test_release_checks.py b/mcp/tests/test_release_checks.py index 753d5f2d2..66d1eb376 100644 --- a/mcp/tests/test_release_checks.py +++ b/mcp/tests/test_release_checks.py @@ -175,7 +175,7 @@ def test_license_inventory_requires_every_direct_runtime_dependency(monkeypatch: validate_inventory(inventory) -def test_sbom_contract_accepts_the_exact_approved_local_component_set() -> None: +def test_sbom_contract_accepts_approved_sources() -> None: sbom = { "bomFormat": "CycloneDX", "specVersion": "1.5", @@ -189,6 +189,12 @@ def test_sbom_contract_accepts_the_exact_approved_local_component_set() -> None: "version": "0.31.0", "purl": "pkg:pypi/asyncpg@0.31.0", }, + { + "name": "nemo-relay", + "version": "0.8.0", + "purl": "pkg:pypi/nemo-relay@0.8.0?vcs_url=https://github.com/NVIDIA/NeMo-Relay.git%3Frev%3D" + "ffb24817442bac99212da0971b13bdad5bc4d84d%23ffb24817442bac99212da0971b13bdad5bc4d84d", + }, ], } From 37976c81c5fccc4189327b980ab32b1f76c3003a Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Wed, 19 Aug 2026 23:56:36 -0700 Subject: [PATCH 09/13] revert: remove temporary Relay SBOM exception Signed-off-by: Chantal D Gama Rose --- mcp/scripts/check_license_inventory.py | 13 ++----------- mcp/tests/test_release_checks.py | 8 +------- 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/mcp/scripts/check_license_inventory.py b/mcp/scripts/check_license_inventory.py index 0d5a63692..68365e07a 100644 --- a/mcp/scripts/check_license_inventory.py +++ b/mcp/scripts/check_license_inventory.py @@ -15,7 +15,6 @@ from importlib.metadata import distribution from pathlib import Path from typing import Any -from urllib.parse import quote _DIRECT_RUNTIME_DEPENDENCIES = { "aiq-agent", @@ -105,14 +104,6 @@ def _canonicalize(name: str) -> str: return re.sub(r"[-_.]+", "-", name).lower() -def _expected_component_purl(name: str, version: str) -> str: - purl = f"pkg:pypi/{name}@{version}" - git_source = _APPROVED_GIT_SOURCES.get((name, version)) - if git_source is not None: - return f"{purl}?vcs_url={quote(git_source, safe=':/')}" - return purl - - def _license_files(dist: Any) -> list[dict[str, str]]: evidence: dict[str, dict[str, str]] = {} for relative in dist.files or (): @@ -220,8 +211,8 @@ def validate_sbom(sbom: dict[str, Any]) -> None: observed_local.add(record) continue - if component.get("purl") != _expected_component_purl(name, version): - raise ValueError(f"dependency is not from the approved source contract: {name}=={version}") + if component.get("purl") != f"pkg:pypi/{name}@{version}": + raise ValueError(f"dependency is not from the public PyPI source contract: {name}=={version}") if observed_local != set(_LOCAL_SOURCE_COMPONENTS): raise ValueError("SBOM local component set differs from the approved public source contract") diff --git a/mcp/tests/test_release_checks.py b/mcp/tests/test_release_checks.py index 66d1eb376..753d5f2d2 100644 --- a/mcp/tests/test_release_checks.py +++ b/mcp/tests/test_release_checks.py @@ -175,7 +175,7 @@ def test_license_inventory_requires_every_direct_runtime_dependency(monkeypatch: validate_inventory(inventory) -def test_sbom_contract_accepts_approved_sources() -> None: +def test_sbom_contract_accepts_the_exact_approved_local_component_set() -> None: sbom = { "bomFormat": "CycloneDX", "specVersion": "1.5", @@ -189,12 +189,6 @@ def test_sbom_contract_accepts_approved_sources() -> None: "version": "0.31.0", "purl": "pkg:pypi/asyncpg@0.31.0", }, - { - "name": "nemo-relay", - "version": "0.8.0", - "purl": "pkg:pypi/nemo-relay@0.8.0?vcs_url=https://github.com/NVIDIA/NeMo-Relay.git%3Frev%3D" - "ffb24817442bac99212da0971b13bdad5bc4d84d%23ffb24817442bac99212da0971b13bdad5bc4d84d", - }, ], } From b7136df291379fec9948f7f938a1c806f82c11a5 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Thu, 20 Aug 2026 11:35:57 -0700 Subject: [PATCH 10/13] fix: redact exception messages for privacy Signed-off-by: Chantal D Gama Rose --- src/aiq_agent/relay/privacy.py | 5 ++++- tests/test_relay_runtime.py | 38 ++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/aiq_agent/relay/privacy.py b/src/aiq_agent/relay/privacy.py index df0f6181d..d015f4a63 100644 --- a/src/aiq_agent/relay/privacy.py +++ b/src/aiq_agent/relay/privacy.py @@ -60,7 +60,10 @@ def _redact_event_fields( sanitized[attribute] = None metadata = sanitized.get("metadata") if isinstance(metadata, dict): - sanitized["metadata"] = {**metadata, "aiq.telemetry.redacted": True} + metadata = {**metadata, "aiq.telemetry.redacted": True} + if "otel.status_description" in metadata: + metadata["otel.status_description"] = "[REDACTED]" + sanitized["metadata"] = metadata return nemo_relay.EventSanitizeFields(**sanitized) diff --git a/tests/test_relay_runtime.py b/tests/test_relay_runtime.py index b208a2b7e..b5ed1dc85 100644 --- a/tests/test_relay_runtime.py +++ b/tests/test_relay_runtime.py @@ -30,6 +30,7 @@ from langchain_core.messages import HumanMessage from langchain_core.tools import tool from nemo_relay import plugin +from nemo_relay.integrations.deepagents import NemoRelayDeepAgentsCallbackHandler from nemo_relay.integrations.langchain._serialization import payload_to_model_request from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from pydantic import ValidationError @@ -611,6 +612,43 @@ async def operation() -> dict[str, str]: assert private_value not in exported +@pytest.mark.asyncio +async def test_request_privacy_redacts_deepagents_error_description(tmp_path: Path) -> None: + config = RelayConfig() + config.logging = False + config.observability.atof.output_directory = str(tmp_path) + config.observability.atof.filename = "private-error.jsonl" + config.observability.opentelemetry.enabled = False + private_error = "proprietary-error-canary" + + await ensure_started(config) + try: + run_id = uuid4() + with request_privacy_context(True), nemo_relay.use_scope_stack(nemo_relay.create_scope_stack()): + callback = NemoRelayDeepAgentsCallbackHandler() + callback.on_chain_start( + {}, + {}, + run_id=run_id, + name="DeepAgent", + metadata={"lc_versions": {"deepagents": "test"}, "ls_integration": "deepagents"}, + ) + callback.on_chain_error(RuntimeError(private_error), run_id=run_id) + finally: + await shutdown_async() + + exported = (tmp_path / "private-error.jsonl").read_text() + assert private_error not in exported + events = [json.loads(line) for line in exported.splitlines()] + assert [(event["name"], event["scope_category"]) for event in events] == [ + ("DeepAgent", "start"), + ("DeepAgent", "end"), + ] + assert len({event["uuid"] for event in events}) == 1 + assert events[-1]["metadata"]["otel.status_code"] == "ERROR" + assert events[-1]["metadata"]["otel.status_description"] == "[REDACTED]" + + @pytest.mark.asyncio async def test_two_turn_parity_has_two_traces_one_session_no_duplicates_and_balanced_scopes(tmp_path: Path) -> None: @tool From 71e20d535c4a533b075f2fd49f9783e83977b226 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Thu, 20 Aug 2026 15:57:00 -0700 Subject: [PATCH 11/13] update release checks to allow non pypi urls Signed-off-by: Chantal D Gama Rose --- mcp/scripts/check_license_inventory.py | 22 ++++++++- mcp/tests/test_release_checks.py | 66 ++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/mcp/scripts/check_license_inventory.py b/mcp/scripts/check_license_inventory.py index 68365e07a..120d6ff90 100644 --- a/mcp/scripts/check_license_inventory.py +++ b/mcp/scripts/check_license_inventory.py @@ -15,6 +15,8 @@ from importlib.metadata import distribution from pathlib import Path from typing import Any +from urllib.parse import parse_qs +from urllib.parse import urlsplit _DIRECT_RUNTIME_DEPENDENCIES = { "aiq-agent", @@ -211,8 +213,24 @@ def validate_sbom(sbom: dict[str, Any]) -> None: observed_local.add(record) continue - if component.get("purl") != f"pkg:pypi/{name}@{version}": - raise ValueError(f"dependency is not from the public PyPI source contract: {name}=={version}") + purl = str(component.get("purl")) + expected_purl = f"pkg:pypi/{name}@{version}" + if purl == expected_purl: + continue + + parsed_purl = urlsplit(purl) + approved_git_source = _APPROVED_GIT_SOURCES.get(record) + qualifiers = parse_qs(parsed_purl.query, keep_blank_values=True, strict_parsing=True) + base_purl = parsed_purl._replace(query="", fragment="").geturl() + if ( + approved_git_source is not None + and base_purl == expected_purl + and qualifiers == {"vcs_url": [approved_git_source]} + and not parsed_purl.fragment + ): + continue + + raise ValueError(f"dependency is not from the public PyPI source contract: {name}=={version}") if observed_local != set(_LOCAL_SOURCE_COMPONENTS): raise ValueError("SBOM local component set differs from the approved public source contract") diff --git a/mcp/tests/test_release_checks.py b/mcp/tests/test_release_checks.py index 753d5f2d2..2df1f5093 100644 --- a/mcp/tests/test_release_checks.py +++ b/mcp/tests/test_release_checks.py @@ -18,12 +18,17 @@ _DIRECT_RUNTIME_DEPENDENCIES = _NAMESPACE["_DIRECT_RUNTIME_DEPENDENCIES"] _PLATFORM_EXCLUDED = _NAMESPACE["_PLATFORM_EXCLUDED"] _WEAK_COPYLEFT = _NAMESPACE["_WEAK_COPYLEFT"] +_APPROVED_GIT_SOURCES = _NAMESPACE["_APPROVED_GIT_SOURCES"] _evidence_fingerprint = _NAMESPACE["_evidence_fingerprint"] validate_inventory = _NAMESPACE["validate_inventory"] validate_sbom = _NAMESPACE["validate_sbom"] validate_lock_sources = _NAMESPACE["validate_lock_sources"] _MCP_LOCK_PATH = _NAMESPACE["_MCP_LOCK_PATH"] +_TEST_RELAY_GIT_SOURCE = "https://github.com/NVIDIA/NeMo-Relay.git?rev=test-revision#test-revision" +_TEST_RELAY_VCS_QUALIFIER = "vcs_url=https://github.com/NVIDIA/NeMo-Relay.git%3Frev%3Dtest-revision%23test-revision" +_TEST_RELAY_PURL = f"pkg:pypi/nemo-relay@0.8.0?{_TEST_RELAY_VCS_QUALIFIER}" + def _row(name: str, version: str = "1.0", **overrides: Any) -> dict[str, Any]: row: dict[str, Any] = { @@ -195,6 +200,67 @@ def test_sbom_contract_accepts_the_exact_approved_local_component_set() -> None: validate_sbom(sbom) +def test_sbom_contract_accepts_exact_approved_vcs_qualified_purl(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(_APPROVED_GIT_SOURCES, ("nemo-relay", "0.8.0"), _TEST_RELAY_GIT_SOURCE) + sbom = { + "bomFormat": "CycloneDX", + "specVersion": "1.5", + "metadata": {"component": {"name": "aiq-mcp-server", "version": "0.1.0"}}, + "components": [ + {"name": "aiq-agent", "version": "2.2.0"}, + {"name": "knowledge-layer", "version": "1.0.0"}, + {"name": "tavily-web-search", "version": "1.0.0"}, + { + "name": "nemo-relay", + "version": "0.8.0", + "purl": _TEST_RELAY_PURL, + }, + ], + } + + validate_sbom(sbom) + + +@pytest.mark.parametrize( + "purl", + [ + pytest.param( + "pkg:pypi/nemo-relay@0.8.0?vcs_url=https://github.com/NVIDIA/NeMo-Relay.git%3Frev%3Dwrong%23wrong", + id="different-revision", + ), + pytest.param( + "pkg:pypi/asyncpg@0.31.0?vcs_url=https://github.com/MagicStack/asyncpg.git%3Frev%3Dabc%23abc", + id="unapproved-package", + ), + pytest.param( + f"{_TEST_RELAY_PURL}&subdirectory=python", + id="extra-qualifier", + ), + pytest.param( + f"{_TEST_RELAY_PURL}&{_TEST_RELAY_VCS_QUALIFIER}", + id="duplicate-vcs-qualifier", + ), + ], +) +def test_sbom_contract_rejects_unapproved_vcs_qualified_purl(purl: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(_APPROVED_GIT_SOURCES, ("nemo-relay", "0.8.0"), _TEST_RELAY_GIT_SOURCE) + name, version = ("asyncpg", "0.31.0") if "asyncpg" in purl else ("nemo-relay", "0.8.0") + sbom = { + "bomFormat": "CycloneDX", + "specVersion": "1.5", + "metadata": {"component": {"name": "aiq-mcp-server", "version": "0.1.0"}}, + "components": [ + {"name": "aiq-agent", "version": "2.2.0"}, + {"name": "knowledge-layer", "version": "1.0.0"}, + {"name": "tavily-web-search", "version": "1.0.0"}, + {"name": name, "version": version, "purl": purl}, + ], + } + + with pytest.raises(ValueError, match="dependency is not from the public PyPI source contract"): + validate_sbom(sbom) + + def test_sbom_contract_requires_every_approved_local_component() -> None: sbom = { "bomFormat": "CycloneDX", From 832116128725414e98b0ae557bf927eaacec57d9 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Thu, 20 Aug 2026 16:18:53 -0700 Subject: [PATCH 12/13] resolve mcp failures in ci Signed-off-by: Chantal D Gama Rose --- mcp/Dockerfile | 1 + mcp/tests/test_deployment_assets.py | 1 + 2 files changed, 2 insertions(+) diff --git a/mcp/Dockerfile b/mcp/Dockerfile index faa5353fc..f6cfd2626 100644 --- a/mcp/Dockerfile +++ b/mcp/Dockerfile @@ -18,6 +18,7 @@ RUN rm -f /etc/apt/apt.conf.d/docker-clean \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ build-essential \ ca-certificates \ + git \ && rm -rf /var/lib/apt/lists/* RUN python -m pip install --no-cache-dir "uv==${UV_VERSION}" diff --git a/mcp/tests/test_deployment_assets.py b/mcp/tests/test_deployment_assets.py index 5a4688bbd..9c9b2c81a 100644 --- a/mcp/tests/test_deployment_assets.py +++ b/mcp/tests/test_deployment_assets.py @@ -219,6 +219,7 @@ def test_release_dockerfile_is_public_reproducible_and_non_root() -> None: assert text.count("FROM ${PYTHON_IMAGE}") == 2 assert "AS builder" in text assert "AS release" in text + assert " git \\\n" in _dockerfile_stage(text, "builder") assert "uv sync" in text assert "--project /app/mcp" in text assert "--frozen" in text From d61dc1d483139a456e908a278da9dee6fefc5ab0 Mon Sep 17 00:00:00 2001 From: Chantal D Gama Rose Date: Thu, 20 Aug 2026 17:12:01 -0700 Subject: [PATCH 13/13] remove obsolete verbose arg Signed-off-by: Chantal D Gama Rose --- frontends/aiq_api/src/aiq_api/jobs/runner.py | 20 +++++++++++--------- tests/aiq_agent/jobs/test_runner.py | 1 - 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 2c34749cf..6904a965e 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -67,6 +67,7 @@ ) _CONFIGURABLE_AGENT_KWARGS = frozenset({"config", "job_id"}) _JOB_SCOPED_AGENT_KWARGS = frozenset({"job_id"}) +_SHALLOW_RESEARCH_AGENT_KWARGS = frozenset({"max_tool_iterations", "enforce_citations"}) @dataclass(frozen=True) @@ -1218,10 +1219,12 @@ def _create_agent_instance( 1. DeepResearcherAgent explicit config pattern 2. llm_provider + tools + config/job_id pattern 3. llm_provider + tools + job_id pattern - 4. llm_provider + tools pattern - 5. llm + tools pattern (simpler agents) + 4. ShallowResearcherAgent config pattern + 5. llm_provider + tools pattern + 6. llm + tools pattern (simpler agents) """ from aiq_agent.agents.deep_researcher.register import DeepResearchAgentConfig + from aiq_agent.agents.shallow_researcher.register import ShallowResearchAgentConfig if isinstance(fn_config, DeepResearchAgentConfig) and _constructor_accepts_explicit_kwargs( agent_cls, _DEEP_RESEARCH_AGENT_KWARGS @@ -1267,23 +1270,22 @@ def _create_agent_instance( except TypeError: pass - # Try the common llm_provider + tools pattern. - try: + if isinstance(fn_config, ShallowResearchAgentConfig) and _constructor_accepts_explicit_kwargs( + agent_cls, _SHALLOW_RESEARCH_AGENT_KWARGS + ): return agent_cls( llm_provider=llm_provider, tools=tools, + max_tool_iterations=fn_config.max_tool_iterations, + enforce_citations=fn_config.enforce_citations, callbacks=callbacks, ) - except TypeError: - pass - # Try llm_provider + tools pattern (ShallowResearcherAgent style) + # Try the common llm_provider + tools pattern. try: return agent_cls( llm_provider=llm_provider, tools=tools, - max_tool_iterations=getattr(fn_config, "max_tool_iterations", 5), - enforce_citations=getattr(fn_config, "enforce_citations", False), callbacks=callbacks, ) except TypeError: diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 9fdd52ee6..cdc793e44 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -3037,7 +3037,6 @@ def __init__( llm="llm", tools=["tool"], fn_config=fn_config, - verbose=False, callbacks=["callback"], )