diff --git a/.vscode/launch.json b/.vscode/launch.json index a24fcc73..3765647c 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -44,6 +44,25 @@ ], "envFile": "${workspaceFolder}/.env" }, + { + "name": "Rogue CLI MCP", + "type": "debugpy", + "request": "launch", + "python": "./.venv/bin/python", + "module": "rogue", + "args": [ + "cli", + "--evaluated-agent-url", + "http://localhost:10001/mcp", + "--protocol", + "mcp", + "--judge-llm", + "openai/o4-mini", + "--workdir", + "./examples/tshirt_store_agent/.rogue" + ], + "envFile": "${workspaceFolder}/.env" + }, { "name": "Rogue AIO", "type": "debugpy", diff --git a/rogue/models/cli_input.py b/rogue/models/cli_input.py index 0c395c14..bedcb23f 100644 --- a/rogue/models/cli_input.py +++ b/rogue/models/cli_input.py @@ -1,7 +1,7 @@ from pathlib import Path from pydantic import BaseModel, Field, HttpUrl, SecretStr, model_validator -from rogue_sdk.types import AuthType, Scenarios +from rogue_sdk.types import AuthType, Protocol, Scenarios, Transport class CLIInput(BaseModel): @@ -10,6 +10,8 @@ class CLIInput(BaseModel): """ workdir: Path = Path(".") / ".rogue" + protocol: Protocol + transport: Transport evaluated_agent_url: HttpUrl evaluated_agent_auth_type: AuthType = AuthType.NO_AUTH evaluated_agent_credentials: SecretStr | None = None @@ -52,6 +54,8 @@ class PartialCLIInput(BaseModel): """ workdir: Path = Path(".") / ".rogue" + protocol: Protocol = Field(default=Protocol.A2A) + transport: Transport = None # type: ignore # fixed in model_post_init evaluated_agent_url: HttpUrl | None = Field(default=None) evaluated_agent_auth_type: AuthType = Field(default=AuthType.NO_AUTH) evaluated_agent_credentials: SecretStr | None = Field(default=None) @@ -73,3 +77,6 @@ def model_post_init(self, __context): self.output_report_file = self.workdir / "report.md" if self.business_context_file is None: self.business_context_file = self.workdir / "business_context.md" + + if self.transport is None: + self.transport = self.protocol.get_default_transport() diff --git a/rogue/run_cli.py b/rogue/run_cli.py index ef4f32f1..fe61717e 100644 --- a/rogue/run_cli.py +++ b/rogue/run_cli.py @@ -4,12 +4,22 @@ import requests from a2a.types import AgentCard +from fastmcp import Client +from fastmcp.client import SSETransport, StreamableHttpTransport from loguru import logger from pydantic import SecretStr, ValidationError from rich.console import Console from rich.markdown import Markdown from rogue_sdk import RogueClientConfig, RogueSDK -from rogue_sdk.types import AgentConfig, AuthType, EvaluationResults, Scenarios +from rogue_sdk.types import ( + PROTOCOL_TO_TRANSPORTS, + AgentConfig, + AuthType, + EvaluationResults, + Protocol, + Scenarios, + Transport, +) from .models.cli_input import CLIInput, PartialCLIInput @@ -25,6 +35,26 @@ def set_cli_args(parser: ArgumentParser) -> None: default="http://localhost:8000", help="Rogue server URL", ) + parser.add_argument( + "--protocol", + choices=[e.value for e in Protocol], + default=Protocol.A2A.value, + type=Protocol, + help="Protocol used to communicate with the agent." + f"Valid options are: {[e.value for e in Protocol]}", + ) + transport_options = ", ".join( + f"{protocol.value}: {[t.value for t in transports]}" + for protocol, transports in PROTOCOL_TO_TRANSPORTS.items() + ) + parser.add_argument( + "--transport", + choices=[e.value for e in Transport], + type=Transport, + required=False, + help="Transport used to communicate with the agent. " + f"Valid options are based on the protocol: {transport_options}", + ) parser.add_argument( "--evaluated-agent-url", required=False, @@ -96,6 +126,8 @@ def set_cli_args(parser: ArgumentParser) -> None: async def run_scenarios( rogue_server_url: str, + protocol: Protocol, + transport: Transport, evaluated_agent_url: str, evaluated_agent_auth_type: AuthType, evaluated_agent_auth_credentials_secret: SecretStr | None, @@ -120,6 +152,8 @@ async def run_scenarios( # Use SDK for evaluation return await _run_scenarios_with_sdk( rogue_server_url=rogue_server_url, + protocol=protocol, + transport=transport, evaluated_agent_url=evaluated_agent_url, evaluated_agent_auth_type=evaluated_agent_auth_type, evaluated_agent_auth_credentials=evaluated_agent_auth_credentials, @@ -134,6 +168,8 @@ async def run_scenarios( async def _run_scenarios_with_sdk( rogue_server_url: str, + protocol: Protocol, + transport: Transport, evaluated_agent_url: str, evaluated_agent_auth_type: AuthType, evaluated_agent_auth_credentials: str | None, @@ -159,12 +195,15 @@ async def _run_scenarios_with_sdk( # Run evaluation job = await sdk.run_evaluation( + protocol=protocol, + transport=transport, agent_url=evaluated_agent_url, scenarios=scenarios, business_context=business_context, auth_type=evaluated_agent_auth_type, auth_credentials=evaluated_agent_auth_credentials, judge_model=judge_llm, + judge_llm_api_key=judge_llm_api_key, deep_test=deep_test_mode, ) @@ -323,20 +362,77 @@ def get_cli_input(cli_args: Namespace) -> CLIInput: return cli_input -def get_agent_card(agent_url: str) -> AgentCard: - try: - response = requests.get( - f"{agent_url}/.well-known/agent.json", - timeout=5, +async def get_a2a_agent_card( + transport: Transport, + agent_url: str, + headers: dict[str, str] | None = None, +) -> AgentCard: + if transport == Transport.HTTP: + try: + response = requests.get( + f"{agent_url}/.well-known/agent.json", + timeout=5, + headers=headers, + ) + return AgentCard.model_validate(response.json()) + except Exception: + logger.debug( + "Failed to connect to agent", + extra={"agent_url": agent_url}, + exc_info=True, + ) + raise + else: + raise ValueError(f"Unsupported transport: {transport} for A2A protocol") + + +async def ping_mcp_server( + transport: Transport, + agent_url: str, + headers: dict[str, str] | None = None, +) -> None: + client: Client[StreamableHttpTransport | SSETransport] + if transport == Transport.STREAMABLE_HTTP: + client = Client[StreamableHttpTransport]( + transport=StreamableHttpTransport( + url=agent_url, + headers=headers, + ), ) - return AgentCard.model_validate(response.json()) - except Exception: - logger.debug( - "Failed to connect to agent", - extra={"agent_url": agent_url}, - exc_info=True, + elif transport == Transport.SSE: + client = Client[SSETransport]( + transport=SSETransport( + url=agent_url, + headers=headers, + ), ) - raise + else: + raise ValueError(f"Unsupported transport: {transport} for MCP protocol") + + async with client: + await client.ping() + + +async def ping_agent( + protocol: Protocol, + transport: Transport, + agent_url: str, + agent_auth_type: AuthType, + agent_auth_credentials: SecretStr | None, +) -> None: + # TODO: move this to the server side + protocol_to_ping_function = { + Protocol.MCP: ping_mcp_server, + Protocol.A2A: get_a2a_agent_card, + } + if protocol not in protocol_to_ping_function: + raise ValueError(f"Unsupported protocol: {protocol}") + + await protocol_to_ping_function[protocol]( + transport=transport, + agent_url=agent_url, + headers=agent_auth_type.get_auth_header(agent_auth_credentials), + ) async def run_cli(args: Namespace) -> int: @@ -344,7 +440,13 @@ async def run_cli(args: Namespace) -> int: logger.debug("Running CLI", extra=cli_input.model_dump()) # fast fail if the agent is not reachable - get_agent_card(cli_input.evaluated_agent_url.encoded_string()) + await ping_agent( + protocol=cli_input.protocol, + transport=cli_input.transport, + agent_url=cli_input.evaluated_agent_url.encoded_string(), + agent_auth_type=cli_input.evaluated_agent_auth_type, + agent_auth_credentials=cli_input.evaluated_agent_credentials, + ) scenarios = cli_input.get_scenarios_from_file() @@ -356,6 +458,8 @@ async def run_cli(args: Namespace) -> int: ) results, job_id = await run_scenarios( rogue_server_url=args.rogue_server_url, + protocol=cli_input.protocol, + transport=cli_input.transport, evaluated_agent_url=cli_input.evaluated_agent_url.encoded_string(), evaluated_agent_auth_type=cli_input.evaluated_agent_auth_type, evaluated_agent_auth_credentials_secret=cli_input.evaluated_agent_credentials, diff --git a/rogue/tests/models/test_cli_input.py b/rogue/tests/models/test_cli_input.py index e2571b47..589408da 100644 --- a/rogue/tests/models/test_cli_input.py +++ b/rogue/tests/models/test_cli_input.py @@ -3,6 +3,7 @@ import pytest from pydantic import HttpUrl, SecretStr, ValidationError from pytest_mock import MockerFixture +from rogue_sdk.types import Protocol, Transport from rogue.models.cli_input import AuthType, CLIInput @@ -23,6 +24,8 @@ class TestCLIInput: ) def test_check_auth_credentials(self, auth_type, credentials, should_raise): input_data = { + "protocol": Protocol.A2A, + "transport": Transport.HTTP, "evaluated_agent_url": "https://example.com", "evaluated_agent_auth_type": auth_type, "evaluated_agent_credentials": credentials, @@ -80,6 +83,8 @@ class CLIInputWithMockScenarios(CLIInput): ) cli_input = CLIInputWithMockScenarios( + protocol=Protocol.A2A, + transport=Transport.HTTP, evaluated_agent_url=HttpUrl("https://example.com"), judge_llm="example-model", business_context="example-context", diff --git a/rogue/tests/test_run_cli.py b/rogue/tests/test_run_cli.py index b88b3ac9..5a6e1fd4 100644 --- a/rogue/tests/test_run_cli.py +++ b/rogue/tests/test_run_cli.py @@ -5,7 +5,7 @@ import pytest from pydantic import HttpUrl, SecretStr from pytest_mock import MockerFixture -from rogue_sdk.types import AuthType +from rogue_sdk.types import AuthType, Protocol, Transport from rogue.models.cli_input import CLIInput from rogue.run_cli import get_cli_input @@ -23,6 +23,8 @@ ), CLIInput( workdir=Path(".") / ".rogue", + protocol=Protocol.A2A, + transport=Transport.HTTP, evaluated_agent_url=HttpUrl("https://localhost:10001"), evaluated_agent_auth_type=AuthType.NO_AUTH, evaluated_agent_credentials=None, @@ -45,6 +47,8 @@ Namespace(business_context="my business"), CLIInput( workdir=Path(".") / ".rogue", + protocol=Protocol.A2A, + transport=Transport.HTTP, evaluated_agent_url=HttpUrl("https://localhost:10001"), evaluated_agent_auth_type=AuthType.API_KEY, evaluated_agent_credentials=SecretStr("abc123"), @@ -66,6 +70,8 @@ ), CLIInput( workdir=Path(".") / ".rogue", + protocol=Protocol.A2A, + transport=Transport.HTTP, evaluated_agent_url=HttpUrl("https://overriden_agent_url:10001"), evaluated_agent_auth_type=AuthType.NO_AUTH, evaluated_agent_credentials=None, diff --git a/sdks/python/rogue_sdk/sdk.py b/sdks/python/rogue_sdk/sdk.py index c6f5297b..ebd52850 100644 --- a/sdks/python/rogue_sdk/sdk.py +++ b/sdks/python/rogue_sdk/sdk.py @@ -254,7 +254,7 @@ async def run_evaluation( agent_config = AgentConfig( protocol=protocol, - transport=transport, + transport=transport or protocol.get_default_transport(), evaluated_agent_url=HttpUrl(agent_url), evaluated_agent_auth_type=auth_type, evaluated_agent_credentials=auth_credentials, diff --git a/sdks/python/rogue_sdk/types.py b/sdks/python/rogue_sdk/types.py index 2addede0..ae68d7ad 100644 --- a/sdks/python/rogue_sdk/types.py +++ b/sdks/python/rogue_sdk/types.py @@ -105,7 +105,7 @@ class AgentConfig(BaseModel): """Configuration for the agent being evaluated.""" protocol: Protocol = Protocol.A2A - transport: Transport | None = None + transport: Transport = None # type: ignore # fixed in model_post_init evaluated_agent_url: HttpUrl evaluated_agent_auth_type: AuthType = Field( default=AuthType.NO_AUTH,