Skip to content

Commit 7103f69

Browse files
authored
feat: add system prompt cli and envar config option to prepend system message to /api/chat. Add unit test and update README entry (#32)
1 parent 33e4fea commit 7103f69

File tree

7 files changed

+87
-12
lines changed

7 files changed

+87
-12
lines changed

README.md

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@
5959
- 💻 **Typer CLI**: Clean command-line interface with configurable options
6060
- 📊 **Structured Logging**: Uses loguru for comprehensive logging
6161
- 📦 **PyPI Package**: Easily installable via pip or uv from PyPI
62-
62+
- 🗣️ **System Prompt Configuration**: Allows setting a system prompt for the assistant's behavior
6363

6464
## Requirements
6565

@@ -204,6 +204,10 @@ CORS_ORIGINS="http://localhost:3000,http://localhost:8080,https://app.example.co
204204
- Can be overridden with `--ollama-url` CLI parameter
205205
- Useful for Docker deployments and configuration management
206206
- Example: `OLLAMA_URL=http://192.168.1.100:11434 ollama-mcp-bridge`
207+
- `SYSTEM_PROMPT`: Optional system prompt to prepend to all forwarded `/api/chat` requests
208+
- Can be set via the `SYSTEM_PROMPT` environment variable or `--system-prompt` CLI flag
209+
- If provided, the bridge will prepend a system message (role: `system`) to the beginning of the `messages` array for `/api/chat` requests unless the request already starts with a system message.
210+
- Example: `SYSTEM_PROMPT="You are a concise assistant." ollama-mcp-bridge`
207211

208212
**CORS Logging:**
209213
- The bridge logs CORS configuration at startup
@@ -235,6 +239,9 @@ ollama-mcp-bridge --ollama-url http://192.168.1.100:11434
235239
# Limit tool execution rounds (prevents excessive tool calls)
236240
ollama-mcp-bridge --max-tool-rounds 5
237241

242+
# Set a system prompt to prepend to all /api/chat requests
243+
ollama-mcp-bridge --system-prompt "You are a concise assistant."
244+
238245
# Combine options
239246
ollama-mcp-bridge --config custom.json --host 0.0.0.0 --port 8080 --ollama-url http://remote-ollama:11434 --max-tool-rounds 10
240247

@@ -253,10 +260,10 @@ ollama-mcp-bridge --version
253260
- `--host`: Host to bind the server (default: `0.0.0.0`)
254261
- `--port`: Port to bind the server (default: `8000`)
255262
- `--ollama-url`: Ollama server URL (default: `http://localhost:11434`)
256-
- `--max-tool-rounds`: Maximum tool execution rounds (default: unlimited, can also be set via `MAX_TOOL_ROUNDS` environment variable)
263+
- `--max-tool-rounds`: Maximum tool execution rounds (default: unlimited)
257264
- `--reload`: Enable auto-reload during development
258265
- `--version`: Show version information, check for updates and exit
259-
266+
- `--system-prompt`: Optional system prompt to prepend to `/api/chat` requests (default: none)
260267
### API Usage
261268

262269
The API is available at `http://localhost:8000`.

src/ollama_mcp_bridge/lifecycle.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,13 @@ async def lifespan(fastapi_app: FastAPI):
2424
config_file = getattr(fastapi_app.state, 'config_file', 'mcp-config.json')
2525
ollama_url = getattr(fastapi_app.state, 'ollama_url', 'http://localhost:11434')
2626
max_tool_rounds = getattr(fastapi_app.state, 'max_tool_rounds', None)
27-
2827
logger.info(f"Starting with config file: {config_file}, Ollama URL: {ollama_url}, Max tool rounds: {max_tool_rounds if max_tool_rounds else 'unlimited'}")
2928

29+
# Get optional system prompt
30+
system_prompt = getattr(fastapi_app.state, 'system_prompt', None)
31+
3032
# Initialize manager and load servers
31-
mcp_manager = MCPManager(ollama_url=ollama_url)
33+
mcp_manager = MCPManager(ollama_url=ollama_url, system_prompt=system_prompt)
3234
mcp_manager.max_tool_rounds = max_tool_rounds
3335
await mcp_manager.load_servers(config_file)
3436

src/ollama_mcp_bridge/main.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ def cli_app(
1616
port: int = typer.Option(8000, "--port", help="Port to bind to"),
1717
ollama_url: str = typer.Option(os.getenv("OLLAMA_URL", "http://localhost:11434"), "--ollama-url", help="Ollama server URL"),
1818
max_tool_rounds: Optional[int] = typer.Option(os.getenv("MAX_TOOL_ROUNDS", None), "--max-tool-rounds", help="Maximum tool execution rounds (default: unlimited)"),
19+
system_prompt: Optional[str] = typer.Option(os.getenv("SYSTEM_PROMPT", None), "--system-prompt", help="System prompt to prepend to messages (can also be set with SYSTEM_PROMPT env var)"),
1920
reload: bool = typer.Option(False, "--reload", help="Enable auto-reload"),
2021
version: bool = typer.Option(False, "--version", help="Show version information, check for updates and exit"),
2122
):
@@ -25,12 +26,13 @@ def cli_app(
2526
# Check for updates and print if available
2627
asyncio.run(check_for_updates(__version__, print_message=True))
2728
raise typer.Exit(0)
28-
validate_cli_inputs(config, host, port, ollama_url, max_tool_rounds)
29+
validate_cli_inputs(config, host, port, ollama_url, max_tool_rounds, system_prompt)
2930

3031
# Store config in app state so lifespan can access it
3132
app.state.config_file = config
3233
app.state.ollama_url = ollama_url
3334
app.state.max_tool_rounds = max_tool_rounds
35+
app.state.system_prompt = system_prompt
3436

3537
logger.info(f"Starting MCP proxy server on {host}:{port}")
3638
logger.info(f"Using Ollama server: {ollama_url}")

src/ollama_mcp_bridge/mcp_manager.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
class MCPManager:
1313
"""Manager for MCP servers, handling tool definitions and session management."""
1414

15-
def __init__(self, ollama_url: str = "http://localhost:11434"):
15+
def __init__(self, ollama_url: str = "http://localhost:11434", system_prompt: str = None):
1616
"""Initialize MCP Manager
1717
1818
Args:
@@ -22,6 +22,8 @@ def __init__(self, ollama_url: str = "http://localhost:11434"):
2222
self.all_tools: List[dict] = []
2323
self.exit_stack = AsyncExitStack()
2424
self.ollama_url = ollama_url
25+
# Optional system prompt that can be prepended to messages
26+
self.system_prompt = system_prompt
2527
self.http_client = httpx.AsyncClient()
2628

2729
async def load_servers(self, config_path: str):

src/ollama_mcp_bridge/proxy_service.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,20 @@ def __init__(self, mcp_manager: MCPManager):
1818
self.mcp_manager = mcp_manager
1919
self.http_client = httpx.AsyncClient(timeout=None)
2020

21+
def _maybe_prepend_system_prompt(self, messages: list) -> list:
22+
"""If a system prompt is configured on the MCP manager, ensure it is the first message.
23+
24+
Does not duplicate if the first message already has role 'system'.
25+
"""
26+
system_prompt = getattr(self.mcp_manager, 'system_prompt', None)
27+
if not system_prompt:
28+
return messages
29+
30+
# If messages is empty or first message isn't a system role, prepend
31+
if not messages or messages[0].get('role') != 'system':
32+
return [{'role': 'system', 'content': system_prompt}] + messages
33+
return messages
34+
2135
async def health_check(self) -> Dict[str, Any]:
2236
"""Check the health of the Ollama server and MCP setup"""
2337
ollama_healthy = await check_ollama_health_async(self.mcp_manager.ollama_url)
@@ -83,6 +97,7 @@ async def _proxy_with_tools_non_streaming(self, endpoint: str, payload: Dict[str
8397
payload = dict(payload)
8498
payload["tools"] = self.mcp_manager.all_tools if self.mcp_manager.all_tools else None
8599
messages = payload.get("messages") or []
100+
messages = self._maybe_prepend_system_prompt(messages)
86101

87102
# Get max tool rounds from app state (None means unlimited)
88103
max_rounds = getattr(self.mcp_manager, 'max_tool_rounds', None)
@@ -124,6 +139,7 @@ async def _proxy_with_tools_streaming(self, endpoint: str, payload: Dict[str, An
124139
payload = dict(payload)
125140
payload["tools"] = self.mcp_manager.all_tools if self.mcp_manager.all_tools else None
126141
messages = list(payload.get("messages") or [])
142+
messages = self._maybe_prepend_system_prompt(messages)
127143

128144
async def stream_ollama(payload_to_send):
129145
async with httpx.AsyncClient(timeout=None) as client:

src/ollama_mcp_bridge/utils.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,12 @@ async def iter_ndjson_chunks(chunk_iterator):
8080
except json.JSONDecodeError as e:
8181
logger.debug(f"Error parsing trailing NDJSON: {e}")
8282

83-
def validate_cli_inputs(config: str, host: str, port: int, ollama_url: str, max_tool_rounds: int = None):
84-
"""Validate CLI inputs for config file, host, port, ollama_url, and max_tool_rounds."""
83+
def validate_cli_inputs(config: str, host: str, port: int, ollama_url: str, max_tool_rounds: int = None, system_prompt: str = None):
84+
"""Validate CLI inputs for config file, host, port, ollama_url, max_tool_rounds and system_prompt.
85+
86+
Args:
87+
system_prompt: optional system prompt string; if provided, must be a non-empty string and not excessively long.
88+
"""
8589
# Validate config file exists
8690
if not os.path.isfile(config):
8791
raise BadParameter(f"Config file not found: {config}")
@@ -103,6 +107,17 @@ def validate_cli_inputs(config: str, host: str, port: int, ollama_url: str, max_
103107
if max_tool_rounds is not None and max_tool_rounds < 1:
104108
raise BadParameter(f"max_tool_rounds must be at least 1, got {max_tool_rounds}")
105109

110+
# Validate system_prompt (if provided)
111+
if system_prompt is not None:
112+
if not isinstance(system_prompt, str):
113+
raise BadParameter("system_prompt must be a string")
114+
# Reject empty or whitespace-only prompts
115+
if not system_prompt.strip():
116+
raise BadParameter("system_prompt must be a non-empty string")
117+
# Limit length to a reasonable maximum to avoid excessively large payloads
118+
if len(system_prompt) > 10000:
119+
raise BadParameter("system_prompt is too long (max 10000 characters)")
120+
106121
async def check_for_updates(current_version: str, print_message: bool = False) -> str:
107122
"""
108123
Check if a newer version of ollama-mcp-bridge is available on PyPI.

tests/test_unit.py

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,19 +140,19 @@ def test_validate_cli_max_tool_rounds():
140140
from ollama_mcp_bridge.utils import validate_cli_inputs
141141

142142
# Valid case: None
143-
validate_cli_inputs("mcp-config.json", "0.0.0.0", 8000, "http://localhost:11434", None)
143+
validate_cli_inputs("mcp-config.json", "0.0.0.0", 8000, "http://localhost:11434", None, None)
144144

145145
# Invalid max_tool_rounds (zero)
146146
from typer import BadParameter
147147
try:
148-
validate_cli_inputs("mcp-config.json", "0.0.0.0", 8000, "http://localhost:11434", 0)
148+
validate_cli_inputs("mcp-config.json", "0.0.0.0", 8000, "http://localhost:11434", 0, None)
149149
assert False, "Expected BadParameter for max_tool_rounds=0"
150150
except BadParameter:
151151
pass
152152

153153
# Invalid max_tool_rounds (negative)
154154
try:
155-
validate_cli_inputs("mcp-config.json", "0.0.0.0", 8000, "http://localhost:11434", -1)
155+
validate_cli_inputs("mcp-config.json", "0.0.0.0", 8000, "http://localhost:11434", -1, None)
156156
assert False, "Expected BadParameter for max_tool_rounds=-1"
157157
except BadParameter:
158158
pass
@@ -163,3 +163,34 @@ def test_script_installed():
163163
assert result.returncode == 0
164164
except (FileNotFoundError, subprocess.SubprocessError) as e:
165165
assert False, f"Subprocess call failed. Is the script installed? {e}"
166+
167+
168+
def test_system_prompt_prepended():
169+
"""Test that the system prompt configured on MCPManager is prepended to messages."""
170+
try:
171+
from ollama_mcp_bridge.mcp_manager import MCPManager
172+
from ollama_mcp_bridge.proxy_service import ProxyService
173+
except ImportError:
174+
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
175+
from ollama_mcp_bridge.mcp_manager import MCPManager
176+
from ollama_mcp_bridge.proxy_service import ProxyService
177+
178+
mgr = MCPManager(system_prompt="You are a helpful assistant.")
179+
ps = ProxyService(mgr)
180+
181+
# Case: user message only -> system prompt should be prepended
182+
messages = [{"role": "user", "content": "Hello"}]
183+
out = ps._maybe_prepend_system_prompt(messages)
184+
assert out[0]["role"] == "system"
185+
assert out[0]["content"] == "You are a helpful assistant."
186+
187+
# Case: existing system prompt should not be duplicated or replaced
188+
messages2 = [{"role": "system", "content": "Existing"}, {"role": "user", "content": "Hi"}]
189+
out2 = ps._maybe_prepend_system_prompt(messages2)
190+
assert out2[0]["role"] == "system"
191+
assert out2[0]["content"] == "Existing"
192+
193+
# Case: empty messages -> system prompt becomes the only message
194+
out3 = ps._maybe_prepend_system_prompt([])
195+
assert out3[0]["role"] == "system"
196+
assert out3[0]["content"] == "You are a helpful assistant."

0 commit comments

Comments
 (0)