"""Mock OpenAI-compatible LLM server for E2E tests."""
import argparse
import json
import re
import time
import uuid
from aiohttp import web
CANNED_RESPONSES = [
(re.compile(r"hello|hi|hey", re.IGNORECASE), "Hello! How can I help you today?"),
(re.compile(r"2\s*\+\s*2|two plus two", re.IGNORECASE), "The answer is 4."),
(re.compile(r"skill|install", re.IGNORECASE), "I can help you with skills management."),
(re.compile(r"html.?test|injection.?test", re.IGNORECASE),
'Here is some content: and
and end of content.'),
]
DEFAULT_RESPONSE = "I understand your request."
def match_response(messages: list[dict]) -> str:
"""Find canned response for the last user message."""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content", "")
# Handle content that may be a list (multi-modal)
if isinstance(content, list):
content = " ".join(
part.get("text", "") for part in content if part.get("type") == "text"
)
for pattern, response in CANNED_RESPONSES:
if pattern.search(content):
return response
return DEFAULT_RESPONSE
return DEFAULT_RESPONSE
async def chat_completions(request: web.Request) -> web.StreamResponse:
"""Handle POST /v1/chat/completions."""
body = await request.json()
messages = body.get("messages", [])
stream = body.get("stream", False)
response_text = match_response(messages)
completion_id = f"mock-{uuid.uuid4().hex[:8]}"
if not stream:
return web.json_response({
"id": completion_id,
"object": "chat.completion",
"created": int(time.time()),
"model": "mock-model",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": response_text},
"finish_reason": "stop",
}],
"usage": {"prompt_tokens": 10, "completion_tokens": len(response_text.split()), "total_tokens": 15},
})
# Streaming response: split into word-boundary chunks
resp = web.StreamResponse(
status=200,
headers={"Content-Type": "text/event-stream", "Cache-Control": "no-cache"},
)
await resp.prepare(request)
# First chunk: role
chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": "mock-model",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
}
await resp.write(f"data: {json.dumps(chunk)}\n\n".encode())
# Content chunks: split on spaces
words = response_text.split(" ")
for i, word in enumerate(words):
text = word if i == 0 else f" {word}"
chunk["choices"][0]["delta"] = {"content": text}
await resp.write(f"data: {json.dumps(chunk)}\n\n".encode())
# Final chunk: finish_reason
chunk["choices"][0]["delta"] = {}
chunk["choices"][0]["finish_reason"] = "stop"
await resp.write(f"data: {json.dumps(chunk)}\n\n".encode())
await resp.write(b"data: [DONE]\n\n")
return resp
async def models(_request: web.Request) -> web.Response:
"""Handle GET /v1/models."""
return web.json_response({
"object": "list",
"data": [{"id": "mock-model", "object": "model", "owned_by": "test"}],
})
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=0)
args = parser.parse_args()
app = web.Application()
app.router.add_post("/v1/chat/completions", chat_completions)
app.router.add_get("/v1/models", models)
# Use aiohttp's runner to get the actual bound port
import asyncio
async def start():
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "127.0.0.1", args.port)
await site.start()
# Extract the actual port from the bound socket
port = site._server.sockets[0].getsockname()[1]
print(f"MOCK_LLM_PORT={port}", flush=True)
# Block forever
await asyncio.Event().wait()
asyncio.run(start())
if __name__ == "__main__":
main()