From 14cc75f4c4c7832c2bec7dbb47d18a2b29001b3c Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Mon, 3 Aug 2026 21:55:22 -0700 Subject: [PATCH 1/2] support GPT frontier structured responses Signed-off-by: Ajay Thorve --- configs/config_frontier_models.yml | 55 +++++--- .../deep_researcher/custom_middleware.py | 59 ++++++++ .../agents/deep_researcher/factory.py | 4 + .../deep_researcher/test_custom_middleware.py | 126 ++++++++++++++++++ .../test_deepagents_runtime.py | 35 +++++ .../agents/deep_researcher/test_factory.py | 9 ++ 6 files changed, 266 insertions(+), 22 deletions(-) diff --git a/configs/config_frontier_models.yml b/configs/config_frontier_models.yml index b3af9f0f1..7dcacefd8 100644 --- a/configs/config_frontier_models.yml +++ b/configs/config_frontier_models.yml @@ -1,6 +1,8 @@ -# This is an example configuration for hybrid frontier model mode. +# This is an example configuration for GPT frontier model mode. # It has the following features: -# - A combination of a open-source model for researcher and a frontier model for orchestrating +# - GPT Sol for orchestration, planning, and report writing +# - GPT Luna for source routing and research +# - NVIDIA API Catalog access through the standard NVIDIA_API_KEY # - Knowledge retrieval using LlamaIndex # - Web search and Paper search tools by default @@ -70,23 +72,32 @@ llms: chat_template_kwargs: enable_thinking: true - nemotron_ultra_llm: + gpt_sol_agent_llm: _type: nim - model_name: nvidia/nemotron-3-ultra-550b-a55b + model_name: openai/gpt-5.6-sol base_url: "https://integrate.api.nvidia.com/v1" - temperature: 0.7 - top_p: 0.7 - max_tokens: 65536 - num_retries: 5 - chat_template_kwargs: - enable_thinking: true + api_key: ${NVIDIA_API_KEY} + max_tokens: 16384 + num_retries: 2 + parallel_tool_calls: false - # Frontier model for orchestrating and planning in deep research agent - # You need to set the OPENAI_API_KEY environment variable - openai_gpt_5_2: - _type: openai - model_name: "gpt-5.2" - api_key: ${OPENAI_API_KEY} + gpt_sol_writer_llm: + _type: nim + model_name: openai/gpt-5.6-sol + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + max_tokens: 32768 + num_retries: 2 + parallel_tool_calls: false + + gpt_luna_agent_llm: + _type: nim + model_name: openai/gpt-5.6-luna + base_url: "https://integrate.api.nvidia.com/v1" + api_key: ${NVIDIA_API_KEY} + max_tokens: 16384 + num_retries: 2 + parallel_tool_calls: false # LLM for document summaries (required when generate_summary: true) summary_llm: @@ -137,7 +148,7 @@ functions: clarifier_agent: _type: clarifier_agent - llm: openai_gpt_5_2 + llm: gpt_sol_agent_llm tools: - web_search_tool - knowledge_search @@ -158,11 +169,11 @@ functions: deep_research_agent: _type: deep_research_agent enable_citation_verification: true - orchestrator_llm: openai_gpt_5_2 - source_router_llm: nemotron_ultra_llm - researcher_llm: nemotron_ultra_llm - planner_llm: openai_gpt_5_2 - writer_llm: openai_gpt_5_2 + orchestrator_llm: gpt_sol_agent_llm + source_router_llm: gpt_luna_agent_llm + 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 diff --git a/src/aiq_agent/agents/deep_researcher/custom_middleware.py b/src/aiq_agent/agents/deep_researcher/custom_middleware.py index e41034a30..fcd56083b 100644 --- a/src/aiq_agent/agents/deep_researcher/custom_middleware.py +++ b/src/aiq_agent/agents/deep_researcher/custom_middleware.py @@ -32,6 +32,8 @@ from langchain_core.messages import HumanMessage from langchain_core.messages import SystemMessage from langchain_core.messages import ToolMessage +from pydantic import BaseModel +from pydantic import ValidationError from aiq_agent.common import get_source_id_for_tool from aiq_agent.common import load_prompt @@ -98,6 +100,63 @@ def _tool_result_failed(result: object) -> bool: return status == "error" +class StructuredResponseTextFallbackMiddleware(AgentMiddleware): + """Recover one exact JSON response when a provider skips the output tool.""" + + def __init__(self, schema: type[BaseModel]) -> None: + self.schema = schema + schema_json = json.dumps(schema.model_json_schema(), separators=(",", ":"), ensure_ascii=False) + self._correction = ( + "The previous response did not produce the required structured result. Do not call tools. " + "Return exactly one JSON object matching this JSON Schema, with no Markdown fences or prose:\n" + f"{schema_json}" + ) + + def _promote(self, response: ModelResponse) -> ModelResponse: + if response.structured_response is not None or len(response.result) != 1: + return response + message = response.result[0] + if not isinstance(message, AIMessage) or message.tool_calls or not isinstance(message.content, str): + return response + try: + structured = self.schema.model_validate_json(message.content) + except ValidationError: + return response + logger.info("Recovered %s from schema-valid JSON message content", self.schema.__name__) + return ModelResponse(result=response.result, structured_response=structured) + + @staticmethod + def _needs_correction(response: ModelResponse) -> bool: + if response.structured_response is not None or len(response.result) != 1: + return False + message = response.result[0] + return isinstance(message, AIMessage) and not message.tool_calls and isinstance(message.content, str) + + def _correction_request(self, request): + return request.override( + messages=[*request.messages, HumanMessage(content=self._correction)], + tools=[], + tool_choice=None, + response_format=None, + ) + + def wrap_model_call(self, request, handler): + """Promote JSON text, with one tools-disabled corrective call when needed.""" + response = self._promote(handler(request)) + if not self._needs_correction(response): + return response + logger.warning("Retrying %s as a tools-disabled JSON response", self.schema.__name__) + return self._promote(handler(self._correction_request(request))) + + async def awrap_model_call(self, request, handler): + """Promote JSON text, with one tools-disabled corrective call when needed.""" + response = self._promote(await handler(request)) + if not self._needs_correction(response): + return response + logger.warning("Retrying %s as a tools-disabled JSON response", self.schema.__name__) + return self._promote(await handler(self._correction_request(request))) + + class FinalReportCommitTracker: """Run-local proof of the writer's most recent successful report mutation.""" diff --git a/src/aiq_agent/agents/deep_researcher/factory.py b/src/aiq_agent/agents/deep_researcher/factory.py index 83e75f92d..078c122cf 100644 --- a/src/aiq_agent/agents/deep_researcher/factory.py +++ b/src/aiq_agent/agents/deep_researcher/factory.py @@ -54,6 +54,7 @@ from .custom_middleware import SourceRoutingGuardMiddleware from .custom_middleware import SourceRoutingPersistenceMiddleware from .custom_middleware import StateMutationGuardMiddleware +from .custom_middleware import StructuredResponseTextFallbackMiddleware from .custom_middleware import TodoQuotaMiddleware from .custom_middleware import TodoSuppressionMiddleware from .custom_middleware import ToolNameSanitizationMiddleware @@ -386,6 +387,7 @@ def build_researcher_runnable( FilesystemMiddleware(backend=backend, _permissions=filesystem_permissions), create_summarization_middleware(researcher_model, backend), PatchToolCallsMiddleware(), + StructuredResponseTextFallbackMiddleware(ResearchNotes), *researcher_middleware, *(visibility_middleware or []), ] @@ -457,6 +459,7 @@ def build_deep_research_subagents(context: DeepResearchGraphContext) -> list[dic sandbox_enabled=context.runtime.execution_enabled, ), TodoSuppressionMiddleware(), + StructuredResponseTextFallbackMiddleware(SourceRoutingPlan), SourceRoutingPersistenceMiddleware( backend=context.backend, state_budget=context.state_budget, @@ -487,6 +490,7 @@ def build_deep_research_subagents(context: DeepResearchGraphContext) -> list[dic sandbox_enabled=context.runtime.execution_enabled, ), TodoSuppressionMiddleware(), + StructuredResponseTextFallbackMiddleware(ResearchPlan), PlanPersistenceMiddleware( backend=context.backend, state_budget=context.state_budget, 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 b6b6e341b..fe8005d23 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py +++ b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py @@ -26,6 +26,7 @@ from deepagents.backends import StateBackend from deepagents.middleware.filesystem import FilesystemMiddleware from langchain.agents import create_agent +from langchain.agents.middleware.types import ModelResponse from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel from langchain_core.messages import AIMessage from langchain_core.messages import HumanMessage @@ -44,10 +45,12 @@ from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingGuardMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingPersistenceMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import StateMutationGuardMiddleware +from aiq_agent.agents.deep_researcher.custom_middleware import StructuredResponseTextFallbackMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import TodoQuotaMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import TodoSuppressionMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolNameSanitizationMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolVisibilityMiddleware +from aiq_agent.agents.deep_researcher.models import SourceRoutingPlan from aiq_agent.agents.deep_researcher.resource_limits import DeepResearchResourceLimits from aiq_agent.agents.deep_researcher.resource_limits import StateBudgetLedger from aiq_agent.agents.deep_researcher.tools.source_registry import build_get_verified_sources_tool @@ -63,6 +66,129 @@ def bind_tools(self, tools, *, tool_choice=None, **kwargs): return self +class TestStructuredResponseTextFallbackMiddleware: + """Strictly recover provider responses that contain the requested contract as JSON text.""" + + @staticmethod + def _routing() -> dict[str, object]: + return { + "domain_id": "general", + "domain_name": "General", + "routing_reason": "Best fit", + "recommendations": [], + "fallback_sources": [], + "planner_guidance": "Use web search.", + } + + def test_promotes_exact_schema_valid_json_in_agent(self) -> None: + payload = self._routing() + model = _ToolBindingFakeChatModel(responses=[AIMessage(content=json.dumps(payload))]) + agent = create_agent( + model=model, + tools=[], + middleware=[StructuredResponseTextFallbackMiddleware(SourceRoutingPlan)], + response_format=SourceRoutingPlan, + ) + + result = agent.invoke({"messages": [HumanMessage(content="Route this request.")]}) + + assert result["structured_response"] == SourceRoutingPlan.model_validate(payload) + + def test_corrects_empty_response_in_agent(self) -> None: + payload = self._routing() + model = _ToolBindingFakeChatModel(responses=[AIMessage(content=""), AIMessage(content=json.dumps(payload))]) + agent = create_agent( + model=model, + tools=[], + middleware=[StructuredResponseTextFallbackMiddleware(SourceRoutingPlan)], + response_format=SourceRoutingPlan, + ) + + result = agent.invoke({"messages": [HumanMessage(content="Route this request.")]}) + + assert result["structured_response"] == SourceRoutingPlan.model_validate(payload) + + @pytest.mark.parametrize( + "content", + ["", "```json\n{}\n```", "{}\nUse this routing plan.", "[]", "{}"], + ) + def test_retries_non_exact_or_schema_invalid_content_without_tools(self, content: str) -> None: + payload = self._routing() + responses = [ + ModelResponse(result=[AIMessage(content=content)]), + ModelResponse(result=[AIMessage(content=json.dumps(payload))]), + ] + request = MagicMock() + request.messages = [HumanMessage(content="Route this request.")] + corrected_request = object() + request.override.return_value = corrected_request + handler = MagicMock(side_effect=responses) + middleware = StructuredResponseTextFallbackMiddleware(SourceRoutingPlan) + + result = middleware.wrap_model_call(request, handler) + + assert result.structured_response == SourceRoutingPlan.model_validate(payload) + assert handler.call_args_list == [((request,),), ((corrected_request,),)] + request.override.assert_called_once() + overrides = request.override.call_args.kwargs + assert overrides["tools"] == [] + assert overrides["tool_choice"] is None + assert overrides["response_format"] is None + assert overrides["messages"][:-1] == request.messages + correction = overrides["messages"][-1] + assert isinstance(correction, HumanMessage) + assert "exactly one JSON object" in str(correction.content) + assert '"domain_id"' in str(correction.content) + + @pytest.mark.asyncio + async def test_async_correction_is_bounded_to_one_retry(self) -> None: + responses = [ + ModelResponse(result=[AIMessage(content="")]), + ModelResponse(result=[AIMessage(content="still invalid")]), + ] + request = MagicMock() + request.messages = [HumanMessage(content="Route this request.")] + corrected_request = object() + request.override.return_value = corrected_request + handler = AsyncMock(side_effect=responses) + middleware = StructuredResponseTextFallbackMiddleware(SourceRoutingPlan) + + result = await middleware.awrap_model_call(request, handler) + + assert result is responses[1] + assert result.structured_response is None + assert handler.await_args_list == [((request,),), ((corrected_request,),)] + + @pytest.mark.asyncio + async def test_preserves_native_structured_response(self) -> None: + structured = SourceRoutingPlan.model_validate(self._routing()) + response = ModelResponse(result=[AIMessage(content="")], structured_response=structured) + middleware = StructuredResponseTextFallbackMiddleware(SourceRoutingPlan) + handler = AsyncMock(return_value=response) + + result = await middleware.awrap_model_call(None, handler) + + assert result is response + handler.assert_awaited_once_with(None) + + @pytest.mark.asyncio + async def test_does_not_intercept_tool_calls(self) -> None: + response = ModelResponse( + result=[ + AIMessage( + content=json.dumps(self._routing()), + tool_calls=[{"name": "lookup_source_catalog", "args": {}, "id": "lookup-1"}], + ) + ] + ) + middleware = StructuredResponseTextFallbackMiddleware(SourceRoutingPlan) + + result = await middleware.awrap_model_call(None, AsyncMock(return_value=response)) + + assert result is response + assert result.structured_response is None + + class TestSourceRoutingGuardMiddleware: """Tests for the orchestrator's required source-routing transition.""" 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 170411117..7fffbbdca 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py +++ b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py @@ -45,6 +45,41 @@ 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_sol_agent_llm"] == { + "_type": "nim", + "model_name": "openai/gpt-5.6-sol", + "base_url": "https://integrate.api.nvidia.com/v1", + "api_key": "${NVIDIA_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": "openai/gpt-5.6-luna", + } + 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_sandbox_wiring() -> None: """Keep the OpenShell workflow aligned with the standard web config.""" diff --git a/tests/aiq_agent/agents/deep_researcher/test_factory.py b/tests/aiq_agent/agents/deep_researcher/test_factory.py index 58a0e9b26..ddfa821a3 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_factory.py +++ b/tests/aiq_agent/agents/deep_researcher/test_factory.py @@ -33,6 +33,7 @@ from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingGuardMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingPersistenceMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import StateMutationGuardMiddleware +from aiq_agent.agents.deep_researcher.custom_middleware import StructuredResponseTextFallbackMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import TodoQuotaMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import TodoSuppressionMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolNameSanitizationMiddleware @@ -263,10 +264,17 @@ def test_subagents_route_tools_and_writer_skills(): assert any(isinstance(item, ToolVisibilityMiddleware) for item in by_name["planner-agent"]["middleware"]) assert any(isinstance(item, ToolVisibilityMiddleware) for item in by_name["writer-agent"]["middleware"]) assert any(isinstance(item, TodoSuppressionMiddleware) for item in by_name["source-router-agent"]["middleware"]) + assert any( + isinstance(item, StructuredResponseTextFallbackMiddleware) + for item in by_name["source-router-agent"]["middleware"] + ) assert any( isinstance(item, SourceRoutingPersistenceMiddleware) for item in by_name["source-router-agent"]["middleware"] ) assert any(isinstance(item, TodoSuppressionMiddleware) for item in by_name["planner-agent"]["middleware"]) + assert any( + isinstance(item, StructuredResponseTextFallbackMiddleware) for item in by_name["planner-agent"]["middleware"] + ) assert any(isinstance(item, TodoSuppressionMiddleware) for item in by_name["writer-agent"]["middleware"]) assert any(isinstance(item, RequiredOutputFileMiddleware) for item in by_name["writer-agent"]["middleware"]) @@ -516,5 +524,6 @@ class FakeSummarizationMiddleware(AgentMiddleware): assert "FilesystemMiddleware" in middleware_names assert "FakeSummarizationMiddleware" in middleware_names assert "PatchToolCallsMiddleware" in middleware_names + assert "StructuredResponseTextFallbackMiddleware" in middleware_names assert "ToolVisibilityMiddleware" in middleware_names assert kwargs["middleware"][-2] is shared_middleware[0] From de584362e185a7cecc7f39de8a557dd95899a3ce Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Tue, 4 Aug 2026 08:13:43 -0700 Subject: [PATCH 2/2] fix frontier OpenAI provider wiring Signed-off-by: Ajay Thorve --- configs/config_frontier_models.yml | 23 ++++++++----------- .../test_deepagents_runtime.py | 9 ++++---- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/configs/config_frontier_models.yml b/configs/config_frontier_models.yml index 7dcacefd8..f87e6dd9b 100644 --- a/configs/config_frontier_models.yml +++ b/configs/config_frontier_models.yml @@ -2,7 +2,7 @@ # It has the following features: # - GPT Sol for orchestration, planning, and report writing # - GPT Luna for source routing and research -# - NVIDIA API Catalog access through the standard NVIDIA_API_KEY +# - OpenAI API access through the standard OPENAI_API_KEY # - Knowledge retrieval using LlamaIndex # - Web search and Paper search tools by default @@ -73,28 +73,25 @@ llms: enable_thinking: true gpt_sol_agent_llm: - _type: nim - model_name: openai/gpt-5.6-sol - base_url: "https://integrate.api.nvidia.com/v1" - api_key: ${NVIDIA_API_KEY} + _type: openai + model_name: "gpt-5.6-sol" + api_key: ${OPENAI_API_KEY} max_tokens: 16384 num_retries: 2 parallel_tool_calls: false gpt_sol_writer_llm: - _type: nim - model_name: openai/gpt-5.6-sol - base_url: "https://integrate.api.nvidia.com/v1" - api_key: ${NVIDIA_API_KEY} + _type: openai + model_name: "gpt-5.6-sol" + api_key: ${OPENAI_API_KEY} max_tokens: 32768 num_retries: 2 parallel_tool_calls: false gpt_luna_agent_llm: - _type: nim - model_name: openai/gpt-5.6-luna - base_url: "https://integrate.api.nvidia.com/v1" - api_key: ${NVIDIA_API_KEY} + _type: openai + model_name: "gpt-5.6-luna" + api_key: ${OPENAI_API_KEY} max_tokens: 16384 num_retries: 2 parallel_tool_calls: false 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 7fffbbdca..3e46a6c52 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py +++ b/tests/aiq_agent/agents/deep_researcher/test_deepagents_runtime.py @@ -51,10 +51,9 @@ def test_frontier_profile_uses_validated_gpt_role_split() -> None: llms = config["llms"] assert llms["gpt_sol_agent_llm"] == { - "_type": "nim", - "model_name": "openai/gpt-5.6-sol", - "base_url": "https://integrate.api.nvidia.com/v1", - "api_key": "${NVIDIA_API_KEY}", + "_type": "openai", + "model_name": "gpt-5.6-sol", + "api_key": "${OPENAI_API_KEY}", "max_tokens": 16384, "num_retries": 2, "parallel_tool_calls": False, @@ -65,7 +64,7 @@ def test_frontier_profile_uses_validated_gpt_role_split() -> None: } assert llms["gpt_luna_agent_llm"] == { **llms["gpt_sol_agent_llm"], - "model_name": "openai/gpt-5.6-luna", + "model_name": "gpt-5.6-luna", } deep_research = config["functions"]["deep_research_agent"] assert {