diff --git a/gradient_adk/cli/agent/deployment/deploy_service.py b/gradient_adk/cli/agent/deployment/deploy_service.py index 8222129..a01730f 100644 --- a/gradient_adk/cli/agent/deployment/deploy_service.py +++ b/gradient_adk/cli/agent/deployment/deploy_service.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Protocol +import gradient_adk from gradient_adk.logging import get_logger from gradient_adk.digital_ocean_api.client_async import AsyncDigitalOceanGenAI from gradient_adk.digital_ocean_api.models import ( @@ -73,6 +74,7 @@ async def deploy_agent( source_dir: Path, project_id: str, api_token: str, + description: str | None = None, ) -> str: """Deploy an agent to the platform. @@ -90,6 +92,7 @@ async def deploy_agent( source_dir: Directory containing the agent code project_id: DigitalOcean project ID api_token: DigitalOcean API token to include in .env + description: Optional description for the deployment (max 1000 chars) Returns: agent_workspace_uuid: UUID of the deployed agent workspace @@ -125,6 +128,7 @@ async def deploy_agent( agent_deployment_name=agent_deployment_name, code_artifact=code_artifact, project_id=project_id, + description=description, ) # Poll for deployment completion @@ -302,6 +306,7 @@ async def _create_or_update_deployment( agent_deployment_name: str, code_artifact: AgentDeploymentCodeArtifact, project_id: str, + description: str | None = None, ) -> str: """Create or update the deployment based on what exists. @@ -312,6 +317,7 @@ async def _create_or_update_deployment( agent_deployment_name: Name of the deployment code_artifact: Code artifact metadata project_id: Project ID + description: Optional description for the deployment Returns: UUID of the created release @@ -324,6 +330,8 @@ async def _create_or_update_deployment( agent_deployment_name=agent_deployment_name, agent_deployment_code_artifact=code_artifact, project_id=project_id, + library_version=gradient_adk.__version__, + description=description, ) workspace_output = await self.client.create_agent_workspace(workspace_input) @@ -349,6 +357,8 @@ async def _create_or_update_deployment( agent_workspace_name=agent_workspace_name, agent_deployment_name=agent_deployment_name, agent_deployment_code_artifact=code_artifact, + library_version=gradient_adk.__version__, + description=description, ) deployment_output = await self.client.create_agent_workspace_deployment( deployment_input @@ -366,6 +376,7 @@ async def _create_or_update_deployment( agent_workspace_name=agent_workspace_name, agent_deployment_name=agent_deployment_name, agent_deployment_code_artifact=code_artifact, + library_version=gradient_adk.__version__, ) release_output = await self.client.create_agent_deployment_release( release_input diff --git a/gradient_adk/cli/cli.py b/gradient_adk/cli/cli.py index cdcf253..1fb1e1e 100644 --- a/gradient_adk/cli/cli.py +++ b/gradient_adk/cli/cli.py @@ -8,7 +8,10 @@ from gradient_adk.cli.agent.deployment.deploy_service import AgentDeployService from gradient_adk.cli.agent.direct_launch_service import DirectLaunchService from gradient_adk.cli.agent.traces_service import GalileoTracesService -from gradient_adk.cli.agent.evaluation_service import EvaluationService, validate_evaluation_dataset +from gradient_adk.cli.agent.evaluation_service import ( + EvaluationService, + validate_evaluation_dataset, +) from gradient_adk.cli.agent.env_utils import get_do_api_token, EnvironmentError @@ -61,6 +64,7 @@ def _configure_agent( agent_name: Optional[str] = None, deployment_name: Optional[str] = None, entrypoint_file: Optional[str] = None, + description: Optional[str] = None, interactive: bool = True, skip_entrypoint_prompt: bool = False, # New parameter for init ) -> None: @@ -79,6 +83,7 @@ def _configure_agent( agent_name=agent_name, agent_environment=deployment_name, entrypoint_file=entrypoint_file, + description=description, interactive=False, ) else: @@ -87,6 +92,7 @@ def _configure_agent( agent_name=agent_name, agent_environment=deployment_name, entrypoint_file=entrypoint_file, + description=description, interactive=interactive, ) @@ -140,6 +146,14 @@ def _create_project_structure() -> None: env_path.write_text(env_content) +# Default description for agents created with `gradient agent init` +_DEFAULT_INIT_DESCRIPTION = ( + "Example LangGraph agent. Invoke: curl -X POST " + '-H "Authorization: Bearer $DIGITALOCEAN_API_TOKEN" ' + '-H "Content-Type: application/json" -d \'{"prompt": "hello"}\'' +) + + @agent_app.command("init") def agent_init( agent_name: Optional[str] = typer.Option( @@ -164,6 +178,7 @@ def agent_init( agent_name=agent_name, deployment_name=deployment_name, entrypoint_file=entrypoint_file, + description=_DEFAULT_INIT_DESCRIPTION, interactive=interactive, skip_entrypoint_prompt=True, # Don't prompt for entrypoint in init ) @@ -190,6 +205,11 @@ def agent_configure( "--entrypoint-file", help="Python file containing @entrypoint decorated function", ), + description: Optional[str] = typer.Option( + None, + "--description", + help="Description for the agent deployment (max 1000 characters)", + ), interactive: bool = typer.Option( True, "--interactive/--no-interactive", help="Interactive prompt mode" ), @@ -199,6 +219,7 @@ def agent_configure( agent_name=agent_name, deployment_name=deployment_name, entrypoint_file=entrypoint_file, + description=description, interactive=interactive, ) @@ -380,6 +401,9 @@ async def deploy(): ) raise typer.Exit(1) + # Get description from config (optional) + description = _agent_config_manager.get_description() + # Create deploy service with injected client deploy_service = AgentDeployService(client=client) @@ -390,6 +414,7 @@ async def deploy(): source_dir=Path.cwd(), project_id=project_id, api_token=api_token, + description=description, ) typer.echo( @@ -420,6 +445,23 @@ async def deploy(): # Get error message with fallback error_msg = str(e) if str(e) else repr(e) + # Check for "feature not enabled" error + if "feature not enabled" in error_msg.lower(): + typer.echo(f"❌ Deployment failed: {error_msg}", err=True) + typer.echo( + "\nThe Gradient ADK is currently in public preview. To access it, enable it for your team via:", + err=True, + ) + typer.echo( + " https://cloud.digitalocean.com/account/feature-preview", + err=True, + ) + typer.echo( + "\nIt may take up to 5 minutes to take effect.", + err=True, + ) + raise typer.Exit(1) + typer.echo(f"❌ Deployment failed: {error_msg}", err=True) typer.echo( @@ -647,14 +689,14 @@ def agent_evaluate( def prompt_and_validate_dataset() -> str: """Prompt for dataset file path and validate it. Re-prompts on error.""" nonlocal dataset_file - + while True: if dataset_file is None: dataset_file = typer.prompt("Dataset file path") - + dataset_path = Path(dataset_file) is_valid, errors = validate_evaluation_dataset(dataset_path) - + if is_valid: return dataset_file else: @@ -676,7 +718,7 @@ async def get_interactive_inputs(): if test_case_name is None: test_case_name = typer.prompt("Evaluation test case name") - + # Validate dataset file immediately when prompting if dataset_file is None or not Path(dataset_file).exists(): prompt_and_validate_dataset() @@ -922,4 +964,4 @@ async def run_evaluation(): def run(): - app() \ No newline at end of file + app() diff --git a/gradient_adk/cli/config/agent_config_manager.py b/gradient_adk/cli/config/agent_config_manager.py index a28c89c..ca1dc56 100644 --- a/gradient_adk/cli/config/agent_config_manager.py +++ b/gradient_adk/cli/config/agent_config_manager.py @@ -17,11 +17,15 @@ def get_agent_environment(self) -> Optional[str]: def get_entrypoint_file(self) -> Optional[str]: raise NotImplementedError + def get_description(self) -> Optional[str]: + raise NotImplementedError + def configure( self, agent_name: Optional[str] = None, agent_environment: Optional[str] = None, entrypoint_file: Optional[str] = None, + description: Optional[str] = None, interactive: bool = True, ) -> None: raise NotImplementedError diff --git a/gradient_adk/cli/config/yaml_agent_config_manager.py b/gradient_adk/cli/config/yaml_agent_config_manager.py index d511a70..79bde12 100644 --- a/gradient_adk/cli/config/yaml_agent_config_manager.py +++ b/gradient_adk/cli/config/yaml_agent_config_manager.py @@ -41,11 +41,16 @@ def get_entrypoint_file(self) -> Optional[str]: config = self.load_config() return config.get("entrypoint_file") if config else None + def get_description(self) -> Optional[str]: + config = self.load_config() + return config.get("description") if config else None + def configure( self, agent_name: Optional[str] = None, agent_environment: Optional[str] = None, entrypoint_file: Optional[str] = None, + description: Optional[str] = None, interactive: bool = True, ) -> None: """Configure agent settings and save to YAML file.""" @@ -67,6 +72,7 @@ def configure( entrypoint_file = typer.prompt( "Entrypoint file (e.g., main.py, agent.py)", default="main.py" ) + # Note: description is optional and not prompted for in interactive mode else: if not all([agent_name, agent_environment, entrypoint_file]): typer.echo( @@ -92,8 +98,16 @@ def configure( ) raise typer.Exit(1) + # Validate description length if provided + if description is not None and len(description) > 1000: + typer.echo( + f"Error: Description exceeds maximum length of 1000 characters (got {len(description)}).", + err=True, + ) + raise typer.Exit(1) + self._validate_entrypoint_file(entrypoint_file) - self._save_config(agent_name, agent_environment, entrypoint_file) + self._save_config(agent_name, agent_environment, entrypoint_file, description) def _validate_name(self, name: str) -> bool: """Validate that a name only contains alphanumeric characters, hyphens, and underscores.""" @@ -159,7 +173,11 @@ def _show_entrypoint_example(self) -> None: ) def _save_config( - self, agent_name: str, agent_environment: str, entrypoint_file: str + self, + agent_name: str, + agent_environment: str, + entrypoint_file: str, + description: Optional[str] = None, ) -> None: """Save configuration to YAML file.""" config = { @@ -168,6 +186,10 @@ def _save_config( "entrypoint_file": entrypoint_file, } + # Only include description if provided + if description is not None: + config["description"] = description + try: with open(self.config_file, "w") as f: yaml.safe_dump(config, f, default_flow_style=False) @@ -175,6 +197,8 @@ def _save_config( typer.echo(f" Agent workspace name: {agent_name}") typer.echo(f" Agent deployment name: {agent_environment}") typer.echo(f" Entrypoint: {entrypoint_file}") + if description: + typer.echo(f" Description: {description[:50]}{'...' if len(description) > 50 else ''}") except Exception as e: typer.echo(f"Error writing configuration file: {e}", err=True) raise typer.Exit(1) diff --git a/gradient_adk/digital_ocean_api/models.py b/gradient_adk/digital_ocean_api/models.py index 3f963a4..55d84f1 100644 --- a/gradient_adk/digital_ocean_api/models.py +++ b/gradient_adk/digital_ocean_api/models.py @@ -235,6 +235,9 @@ class AgentDeploymentRelease(BaseModel): created_by_user_email: Optional[str] = Field( None, description="Email of user that created the agent deployment release" ) + library_version: Optional[str] = Field( + None, description="Version of the ADK library used to create this release" + ) class AgentLoggingConfig(BaseModel): @@ -438,6 +441,14 @@ class CreateAgentWorkspaceDeploymentInput(BaseModel): agent_deployment_code_artifact: AgentDeploymentCodeArtifact = Field( ..., description="The agent deployment code artifact" ) + library_version: Optional[str] = Field( + None, description="Version of the ADK library used to create this deployment" + ) + description: Optional[str] = Field( + None, + description="Description of the agent deployment (max 1000 characters)", + max_length=1000, + ) class CreateAgentWorkspaceDeploymentOutput(BaseModel): @@ -464,6 +475,9 @@ class CreateAgentDeploymentReleaseInput(BaseModel): agent_deployment_code_artifact: AgentDeploymentCodeArtifact = Field( ..., description="The agent deployment code artifact" ) + library_version: Optional[str] = Field( + None, description="Version of the ADK library used to create this release" + ) class CreateAgentDeploymentReleaseOutput(BaseModel): @@ -513,6 +527,14 @@ class CreateAgentWorkspaceInput(BaseModel): ..., description="The agent deployment code artifact" ) project_id: str = Field(..., description="The project id") + library_version: Optional[str] = Field( + None, description="Version of the ADK library used to create this workspace" + ) + description: Optional[str] = Field( + None, + description="Description of the agent workspace deployment (max 1000 characters)", + max_length=1000, + ) class CreateAgentWorkspaceOutput(BaseModel): diff --git a/gradient_adk/runtime/langgraph/langgraph_instrumentor.py b/gradient_adk/runtime/langgraph/langgraph_instrumentor.py index 7a2ba29..edc28fc 100644 --- a/gradient_adk/runtime/langgraph/langgraph_instrumentor.py +++ b/gradient_adk/runtime/langgraph/langgraph_instrumentor.py @@ -241,13 +241,33 @@ def _had_hits_since(intr, token) -> bool: def _get_captured_payloads_with_type(intr, token) -> tuple: """Get captured API request/response payloads and classify the call type. + When using httpx, both request() and send() are intercepted, which can result + in two CapturedRequest records for a single HTTP call. The request() interception + captures the request payload but not the response (since it delegates to send()), + while the send() interception captures both request and response payloads. + + This function searches for a captured request that has a response payload, + preferring the last one (which is typically from send()). + Returns: (request_payload, response_payload, is_llm, is_retriever) """ try: captured = intr.get_captured_requests_since(token) if captured: - # Use the first captured request (most common case) + # Search in reverse order to find a captured request with a response. + # This handles the case where both request() and send() are intercepted: + # - captured[0] from request(): has request_payload, no response_payload + # - captured[1] from send(): has both request_payload and response_payload + for call in reversed(captured): + if call.response_payload is not None: + url = call.url + is_llm = is_inference_url(url) + is_retriever = is_kbaas_url(url) + return call.request_payload, call.response_payload, is_llm, is_retriever + + # Fallback to the first captured request if none have a response + # (e.g., if the HTTP request failed or response wasn't captured) call = captured[0] url = call.url is_llm = is_inference_url(url) diff --git a/integration_tests/configure/test_adk_agents_configure.py b/integration_tests/configure/test_adk_agents_configure.py index e0b0e2d..c672b3b 100644 --- a/integration_tests/configure/test_adk_agents_configure.py +++ b/integration_tests/configure/test_adk_agents_configure.py @@ -473,6 +473,316 @@ async def my_agent(query, context): logger.info("Custom entrypoint path configured correctly") + @pytest.mark.cli + def test_configure_with_description_happy_path(self): + """ + Test that configure works with a valid description. + Verifies that: + - The command exits with code 0 + - Description is saved correctly in config file + """ + logger = logging.getLogger(__name__) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create an entrypoint file + main_py = temp_path / "main.py" + main_py.write_text(""" +from gradient_adk import entrypoint + +@entrypoint +async def main(query, context): + return {"result": "test"} +""") + + workspace_name = "test-agent" + deployment_name = "main" + entrypoint_file = "main.py" + description = "This is a test agent that does amazing things." + + logger.info(f"Running gradient agent configure with description in {temp_dir}") + + result = subprocess.run( + [ + "gradient", + "agent", + "configure", + "--agent-workspace-name", + workspace_name, + "--deployment-name", + deployment_name, + "--entrypoint-file", + entrypoint_file, + "--description", + description, + "--no-interactive", + ], + capture_output=True, + text=True, + timeout=60, + cwd=temp_dir, + ) + + assert result.returncode == 0, f"Command failed with stderr: {result.stderr}" + + # Verify config file was created with description + config_file = temp_path / ".gradient" / "agent.yml" + assert config_file.exists(), ".gradient/agent.yml was not created" + + with open(config_file, "r") as f: + config = yaml.safe_load(f) + + assert config["agent_name"] == workspace_name + assert config["agent_environment"] == deployment_name + assert config["entrypoint_file"] == entrypoint_file + assert config.get("description") == description, \ + f"Expected description '{description}', got '{config.get('description')}'" + + logger.info("Description configured correctly") + + @pytest.mark.cli + def test_configure_description_too_long(self): + """ + Test that configure fails when description exceeds 1000 characters. + """ + logger = logging.getLogger(__name__) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create an entrypoint file + main_py = temp_path / "main.py" + main_py.write_text(""" +from gradient_adk import entrypoint + +@entrypoint +async def main(query, context): + return {"result": "test"} +""") + + # Create a description that exceeds 1000 characters + long_description = "x" * 1001 + + logger.info(f"Running gradient agent configure with long description in {temp_dir}") + + result = subprocess.run( + [ + "gradient", + "agent", + "configure", + "--agent-workspace-name", + "test-agent", + "--deployment-name", + "main", + "--entrypoint-file", + "main.py", + "--description", + long_description, + "--no-interactive", + ], + capture_output=True, + text=True, + timeout=60, + cwd=temp_dir, + ) + + assert result.returncode != 0, "Command should have failed with long description" + + combined_output = result.stdout + result.stderr + assert any( + term in combined_output.lower() + for term in ["description", "1000", "length", "exceeds", "error"] + ), f"Expected error about description length, got: {combined_output}" + + logger.info("Correctly failed with description too long") + + @pytest.mark.cli + def test_configure_description_at_max_length(self): + """ + Test that configure works with description at exactly 1000 characters. + """ + logger = logging.getLogger(__name__) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create an entrypoint file + main_py = temp_path / "main.py" + main_py.write_text(""" +from gradient_adk import entrypoint + +@entrypoint +async def main(query, context): + return {"result": "test"} +""") + + # Create a description that is exactly 1000 characters + max_description = "x" * 1000 + + logger.info(f"Running gradient agent configure with max length description in {temp_dir}") + + result = subprocess.run( + [ + "gradient", + "agent", + "configure", + "--agent-workspace-name", + "test-agent", + "--deployment-name", + "main", + "--entrypoint-file", + "main.py", + "--description", + max_description, + "--no-interactive", + ], + capture_output=True, + text=True, + timeout=60, + cwd=temp_dir, + ) + + assert result.returncode == 0, f"Command failed with stderr: {result.stderr}" + + # Verify config file has correct description + config_file = temp_path / ".gradient" / "agent.yml" + with open(config_file, "r") as f: + config = yaml.safe_load(f) + + assert config.get("description") == max_description, \ + f"Expected 1000-char description, got {len(config.get('description', ''))} chars" + + logger.info("Max length description configured correctly") + + @pytest.mark.cli + def test_configure_updates_existing_config_with_description(self): + """ + Test that configure can add a description to an existing config. + """ + logger = logging.getLogger(__name__) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create an entrypoint file + main_py = temp_path / "main.py" + main_py.write_text(""" +from gradient_adk import entrypoint + +@entrypoint +async def main(query, context): + return {"result": "test"} +""") + + # Create initial config without description + gradient_dir = temp_path / ".gradient" + gradient_dir.mkdir() + config_file = gradient_dir / "agent.yml" + + initial_config = { + "agent_name": "test-agent", + "agent_environment": "main", + "entrypoint_file": "main.py", + } + with open(config_file, "w") as f: + yaml.safe_dump(initial_config, f) + + # Verify initial config has no description + with open(config_file, "r") as f: + config = yaml.safe_load(f) + assert "description" not in config, "Initial config should not have description" + + # Run configure to add description + new_description = "Adding a description to existing config." + + logger.info(f"Running gradient agent configure to add description in {temp_dir}") + + result = subprocess.run( + [ + "gradient", + "agent", + "configure", + "--agent-workspace-name", + "test-agent", + "--deployment-name", + "main", + "--entrypoint-file", + "main.py", + "--description", + new_description, + "--no-interactive", + ], + capture_output=True, + text=True, + timeout=60, + cwd=temp_dir, + ) + + assert result.returncode == 0, f"Command failed with stderr: {result.stderr}" + + # Verify config was updated with description + with open(config_file, "r") as f: + config = yaml.safe_load(f) + + assert config.get("description") == new_description, \ + f"Expected description '{new_description}', got '{config.get('description')}'" + + logger.info("Description successfully added to existing config") + + @pytest.mark.cli + def test_configure_without_description_does_not_add_description(self): + """ + Test that configure without --description does not add a description field. + """ + logger = logging.getLogger(__name__) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) + + # Create an entrypoint file + main_py = temp_path / "main.py" + main_py.write_text(""" +from gradient_adk import entrypoint + +@entrypoint +async def main(query, context): + return {"result": "test"} +""") + + logger.info(f"Running gradient agent configure without description in {temp_dir}") + + result = subprocess.run( + [ + "gradient", + "agent", + "configure", + "--agent-workspace-name", + "test-agent", + "--deployment-name", + "main", + "--entrypoint-file", + "main.py", + "--no-interactive", + ], + capture_output=True, + text=True, + timeout=60, + cwd=temp_dir, + ) + + assert result.returncode == 0, f"Command failed with stderr: {result.stderr}" + + # Verify config file does NOT have description field + config_file = temp_path / ".gradient" / "agent.yml" + with open(config_file, "r") as f: + config = yaml.safe_load(f) + + assert "description" not in config, \ + f"Config should not have description field when not provided, but got: {config}" + + logger.info("Config correctly created without description field") + @pytest.mark.cli def test_configure_requires_all_options_in_non_interactive(self): """ @@ -551,7 +861,8 @@ def test_configure_help(self): f"Should show --agent-workspace-name option. Got: {combined_output}" assert "--deployment-name" in combined_output, "Should show --deployment-name option" assert "--entrypoint-file" in combined_output, "Should show --entrypoint-file option" + assert "--description" in combined_output, "Should show --description option" assert "--interactive" in combined_output or "--no-interactive" in combined_output, \ "Should show --interactive option" - logger.info("Help output is correct") \ No newline at end of file + logger.info("Help output is correct") diff --git a/tests/runtime/langgraph/langgraph_instrumentor_test.py b/tests/runtime/langgraph/langgraph_instrumentor_test.py index 9a915bd..274ac18 100644 --- a/tests/runtime/langgraph/langgraph_instrumentor_test.py +++ b/tests/runtime/langgraph/langgraph_instrumentor_test.py @@ -570,3 +570,88 @@ def test_get_captured_payloads_with_type_no_captures(): assert resp is None assert is_llm is False assert is_retriever is False + + +def test_get_captured_payloads_with_type_double_capture(): + """Test _get_captured_payloads_with_type handles double-capture from request() and send(). + + When using httpx, both request() and send() are intercepted, creating two + CapturedRequest records: + - The first (from request()) has request_payload but NO response_payload + - The second (from send()) has both request_payload AND response_payload + + The function should find and return the one with the response. + """ + mock_intr = MagicMock() + + # First capture from request() - no response + mock_captured_1 = MagicMock() + mock_captured_1.url = "https://kbaas.do-ai.run/retrieve" + mock_captured_1.request_payload = {"query": "test"} + mock_captured_1.response_payload = None # No response captured here + + # Second capture from send() - has response + mock_captured_2 = MagicMock() + mock_captured_2.url = "https://kbaas.do-ai.run/retrieve" + mock_captured_2.request_payload = {"query": "test"} + mock_captured_2.response_payload = {"results": [{"text_content": "result"}]} + + mock_intr.get_captured_requests_since.return_value = [mock_captured_1, mock_captured_2] + + req, resp, is_llm, is_retriever = _get_captured_payloads_with_type(mock_intr, 0) + + # Should return the payload from the second capture (which has the response) + assert req == {"query": "test"} + assert resp == {"results": [{"text_content": "result"}]} + assert is_llm is False + assert is_retriever is True + + +def test_get_captured_payloads_with_type_double_capture_fallback(): + """Test _get_captured_payloads_with_type falls back to first capture if none have response. + + If for some reason neither capture has a response (e.g., request failed), + the function should fall back to the first captured request. + """ + mock_intr = MagicMock() + + # Both captures have no response + mock_captured_1 = MagicMock() + mock_captured_1.url = "https://kbaas.do-ai.run/retrieve" + mock_captured_1.request_payload = {"query": "test1"} + mock_captured_1.response_payload = None + + mock_captured_2 = MagicMock() + mock_captured_2.url = "https://kbaas.do-ai.run/retrieve" + mock_captured_2.request_payload = {"query": "test2"} + mock_captured_2.response_payload = None + + mock_intr.get_captured_requests_since.return_value = [mock_captured_1, mock_captured_2] + + req, resp, is_llm, is_retriever = _get_captured_payloads_with_type(mock_intr, 0) + + # Should fall back to first capture + assert req == {"query": "test1"} + assert resp is None + assert is_llm is False + assert is_retriever is True + + +def test_get_captured_payloads_with_type_single_capture_with_response(): + """Test _get_captured_payloads_with_type works with single capture that has response.""" + mock_intr = MagicMock() + + # Single capture with response (normal case when only send() is triggered) + mock_captured = MagicMock() + mock_captured.url = "https://inference.do-ai.run/v1/chat" + mock_captured.request_payload = {"messages": [{"role": "user", "content": "Hello"}]} + mock_captured.response_payload = {"choices": [{"message": {"content": "Hi!"}}]} + + mock_intr.get_captured_requests_since.return_value = [mock_captured] + + req, resp, is_llm, is_retriever = _get_captured_payloads_with_type(mock_intr, 0) + + assert req == {"messages": [{"role": "user", "content": "Hello"}]} + assert resp == {"choices": [{"message": {"content": "Hi!"}}]} + assert is_llm is True + assert is_retriever is False