Skip to content

Commit 2ec51dc

Browse files
yuval-qfyuval-qf
authored
FIRE-831 | Rogue Server | MCP transport (#103)
* WIP - add mcp example * Fix mcp agent * Fix mcp agent * typing * Add docs * Add base logic to support different transports * Use local sdk for github action tests * Add TransportType to sdk __init__/__all__ * fix gh tests * Fix cyclical import * Add missing files * Fix missing kwarg * cicd fix * Fix import * Add mcp transport * formatting * transport -> protocol * formatting * Fix session id * Improve import time * Rename example folder * revert unreleated changes * Rabbit cr * Add transport support * Improve import time * Rabbit cr * Update launch.json path * Add transport to a2a to create a clean convention * CR * Fix partial import * Run mcp example from maing * Fix typo * Rabbit cr - Add click options to mcp examples * Fix uv step hang in github workflow
1 parent 10a4883 commit 2ec51dc

25 files changed

Lines changed: 970 additions & 232 deletions

‎.github/actions/run-tests/action.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ runs:
1111
steps:
1212
- name: "Run tests"
1313
shell: bash
14-
run: source .venv/bin/activate && uv run --dev pytest --junitxml=test-results.xml --cov=rogue --cov-report=xml
14+
run: source .venv/bin/activate && pytest --junitxml=test-results.xml --cov=rogue --cov-report=xml
1515

1616
- name: "Upload Test Results"
1717
uses: actions/upload-artifact@v4

‎.github/workflows/rogue.yml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ jobs:
2525

2626
- name: Install uv
2727
uses: astral-sh/setup-uv@v5
28+
with:
29+
prune-cache: false
2830

2931
- name: Setup Python
3032
uses: actions/setup-python@v5
@@ -49,7 +51,7 @@ jobs:
4951
echo "🚀 Starting AI agent..."
5052
5153
# Not using uv because it will reinstall the sdk from pypi
52-
source .venv/bin/activate && uv run python -m examples.tshirt_store_agent --host 0.0.0.0 --port 10001 &
54+
source .venv/bin/activate && python -m examples.tshirt_store_agent --host 0.0.0.0 --port 10001 &
5355
AGENT_PID=$!
5456
echo "Agent started with PID: $AGENT_PID"
5557
trap 'echo "🛑 Stopping agent..."; kill $AGENT_PID' EXIT

‎.github/workflows/test.yml‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,13 @@ jobs:
3333
- name: Create venv
3434
run: uv venv
3535

36-
- name: Install rogue sdk
37-
run: source .venv/bin/activate && uv pip install -e sdks/python --force-reinstall
38-
39-
- name: Install rogue server
40-
run: source .venv/bin/activate && uv sync --dev
36+
- name: Install rogue server with local sdk
37+
run: |
38+
source .venv/bin/activate
39+
uv sync --dev
40+
uv pip install -e .
41+
uv pip uninstall rogue-ai-sdk
42+
uv pip install -e sdks/python --force-reinstall
4143
4244
- name: Run tests
4345
uses: ./.github/actions/run-tests

‎.vscode/launch.json‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,16 +58,16 @@
5858
"type": "debugpy",
5959
"request": "launch",
6060
"python": "./.venv/bin/python",
61-
"program": "./examples/tshirt_store_agent",
61+
"module": "examples.tshirt_store_agent",
6262
"envFile": "${workspaceFolder}/examples/tshirt_store_agent/.env"
6363
},
6464
{
6565
"name": "shirtify mcp",
6666
"type": "debugpy",
6767
"request": "launch",
6868
"python": "./.venv/bin/python",
69-
"program": "./examples/mcp/tshirt_store_mcp",
70-
"envFile": "${workspaceFolder}/examples/mcp/tshirt_store_mcp/.env"
69+
"module": "examples.mcp.tshirt_store_langgraph_mcp",
70+
"envFile": "${workspaceFolder}/examples/mcp/tshirt_store_langgraph_mcp/.env"
7171
}
7272
]
7373
}

‎examples/mcp/tshirt_store_langgraph_mcp/__main__.py‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,31 @@
1+
from typing import Literal
2+
3+
import click
14
from dotenv import load_dotenv
2-
from mcp_agent_wrapper import mcp
5+
6+
from .mcp_agent_wrapper import get_mcp_server
37

48
load_dotenv()
59

610

7-
def main() -> None:
11+
@click.command()
12+
@click.option("--host", "host", default="127.0.0.1", help="Host to run the server on")
13+
@click.option("--port", "port", default=10001, help="Port to run the server on")
14+
@click.option(
15+
"--transport",
16+
"transport",
17+
default="streamable-http",
18+
choices=["streamable-http", "sse"],
19+
help="Transport to use for the mcp server",
20+
)
21+
def main(host: str, port: int, transport: Literal["streamable-http", "sse"]) -> None:
822
print("Starting MCP server...")
23+
mcp = get_mcp_server(host=host, port=port)
924

10-
# Can also be "sse".
1125
# When using "sse", the url will be http://localhost:10001/sse
1226
# When using "streamable-http", the url will be http://localhost:10001/mcp
1327
# stdio isn't supported in this example, since rogue won't be able to connect to it.
14-
mcp.run(transport="streamable-http")
15-
# mcp.run(transport="sse")
28+
mcp.run(transport=transport)
1629

1730

1831
if __name__ == "__main__":

‎examples/mcp/tshirt_store_langgraph_mcp/mcp_agent_wrapper.py‎

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -6,38 +6,45 @@
66
you only need to implement the send_message tool.
77
"""
88

9+
from functools import lru_cache
10+
911
from loguru import logger
1012
from mcp.server.fastmcp import Context, FastMCP
11-
from shirtify_agent import ShirtifyAgent
1213
from starlette.requests import Request
1314

14-
agent = ShirtifyAgent()
15-
mcp = FastMCP(
16-
"shirtify_agent_mcp",
17-
port=10001,
18-
host="127.0.0.1",
19-
)
15+
from .shirtify_agent import ShirtifyAgent
16+
17+
18+
@lru_cache(maxsize=1)
19+
def get_mcp_server(host: str = "127.0.0.1", port: int = 10001) -> FastMCP:
20+
agent = ShirtifyAgent()
21+
mcp = FastMCP(
22+
"shirtify_agent_mcp",
23+
host=host,
24+
port=port,
25+
)
2026

27+
@mcp.tool()
28+
def send_message(message: str, context: Context) -> str:
29+
session_id: str | None = None
30+
try:
31+
request: Request = context.request_context.request # type: ignore
2132

22-
@mcp.tool()
23-
def send_message(message: str, context: Context) -> str:
24-
session_id: str | None = None
25-
try:
26-
request: Request = context.request_context.request # type: ignore
33+
# The session id should be in the headers for streamable-http transport
34+
session_id = request.headers.get("mcp-session-id")
2735

28-
# The session id should be in the headers for streamable-http transport
29-
session_id = request.headers.get("mcp-session-id")
36+
# The session id might also be in query param when using sse transport
37+
if session_id is None:
38+
session_id = request.query_params.get("session_id")
39+
except Exception:
40+
session_id = None
41+
logger.exception("Error while extracting session id")
3042

31-
# The session id might also be in query param when using sse transport
3243
if session_id is None:
33-
session_id = request.query_params.get("session_id")
34-
except Exception:
35-
session_id = None
36-
logger.exception("Error while extracting session id")
44+
logger.error("Couldn't extract session id")
3745

38-
if session_id is None:
39-
logger.error("Couldn't extract session id")
46+
# Invoking our agent
47+
response = agent.invoke(message, session_id)
48+
return response.get("content", "")
4049

41-
# Invoking our agent
42-
response = agent.invoke(message, session_id)
43-
return response.get("content", "")
50+
return mcp

‎lefthook.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ pre-commit:
3030

3131
- name: isort
3232
glob: "*.py"
33-
run: isort {staged_files}
33+
run: isort --profile black {staged_files}
3434

3535
- name: black
3636
glob: "*.py"

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ dependencies = [
1010
"click>=8.0.0",
1111
"datasets==3.6.0",
1212
"fastapi>=0.115.0",
13+
"fastmcp>=2.12.5",
1314
"google-adk==1.5.0",
1415
"gradio==5.35.0",
1516
"langchain-openai>=0.3.35",

‎rogue/__main__.py‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,9 @@ def common_parser() -> ArgumentParser:
4444
parent_parser.add_argument(
4545
"--example",
4646
type=str,
47-
choices=["tshirt_store"],
48-
help="Run with an example agent (e.g., tshirt_store)",
47+
choices=["tshirt_store", "tshirt_store_langgraph_mcp"],
48+
help="Run with an example agent "
49+
"(e.g., tshirt_store, tshirt_store_langgraph_mcp)",
4950
)
5051
parent_parser.add_argument(
5152
"--example-host",
@@ -126,6 +127,19 @@ def start_example_agent(
126127
"--port",
127128
str(port),
128129
]
130+
elif example_name == "tshirt_store_langgraph_mcp":
131+
# Use subprocess to run the example agent
132+
cmd = [
133+
sys.executable,
134+
"-m",
135+
"examples.mcp.tshirt_store_langgraph_mcp",
136+
"--host",
137+
host,
138+
"--port",
139+
str(port),
140+
"--transport",
141+
"streamable-http",
142+
]
129143
else:
130144
logger.error(f"Unknown example: {example_name}")
131145
return None

‎rogue/evaluator_agent/__init__.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,21 @@
1-
from . import evaluator_agent, policy_evaluation, run_evaluator_agent
1+
from . import (
2+
a2a,
3+
base_evaluator_agent,
4+
evaluator_agent_factory,
5+
mcp,
6+
policy_evaluation,
7+
run_evaluator_agent,
8+
)
9+
from .a2a import A2AEvaluatorAgent
10+
from .mcp import MCPEvaluatorAgent
11+
12+
__all__ = [
13+
"base_evaluator_agent",
14+
"evaluator_agent_factory",
15+
"policy_evaluation",
16+
"run_evaluator_agent",
17+
"a2a",
18+
"mcp",
19+
"A2AEvaluatorAgent",
20+
"MCPEvaluatorAgent",
21+
]

0 commit comments

Comments
 (0)