diff --git a/services/mcp-docugen/src/mcp_docugen/http_routes.py b/services/mcp-docugen/src/mcp_docugen/http_routes.py new file mode 100644 index 0000000..4bc46cd --- /dev/null +++ b/services/mcp-docugen/src/mcp_docugen/http_routes.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import re +from pathlib import Path + +from fastapi import FastAPI, HTTPException +from fastapi.responses import FileResponse, JSONResponse + +from mcp_docugen.generation_store import GenerationStore +from mcp_docugen.template_store import InvalidTemplateName, TemplateStore + +_SAFE_SEGMENT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def build_http_app( + template_store: TemplateStore, + generation_store: GenerationStore, +) -> FastAPI: + app = FastAPI() + + @app.get("/health") + async def health(): + return {"status": "ok"} + + @app.get("/assets/{template_name}/{filename}") + async def get_template_asset(template_name: str, filename: str): + if not _SAFE_SEGMENT_RE.match(template_name) or not _SAFE_SEGMENT_RE.match( + filename + ): + raise HTTPException(status_code=400, detail="invalid path") + try: + path = template_store.asset_path(template_name, filename) + except (ValueError, InvalidTemplateName) as exc: + raise HTTPException(status_code=400, detail="invalid path") from exc + if not path.exists(): + raise HTTPException(status_code=404, detail="not found") + return FileResponse(path) + + @app.get("/generated/{gen_id}/{filename}") + async def get_generated_asset(gen_id: str, filename: str): + if not _SAFE_SEGMENT_RE.match(gen_id) or not _SAFE_SEGMENT_RE.match(filename): + raise HTTPException(status_code=400, detail="invalid path") + info = await generation_store.get_ephemeral_asset(gen_id, filename) + if info is None or info.is_expired or not Path(info.file_path).exists(): + return JSONResponse( + status_code=410, + content={ + "error": "gone", + "reason": "ephemeral asset expired or not found", + }, + ) + return FileResponse(info.file_path, media_type=info.mime) + + return app diff --git a/services/mcp-docugen/src/mcp_docugen/main.py b/services/mcp-docugen/src/mcp_docugen/main.py new file mode 100644 index 0000000..04de7e3 --- /dev/null +++ b/services/mcp-docugen/src/mcp_docugen/main.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import asyncio +import logging +import os +from contextlib import asynccontextmanager + +import uvicorn +from fastapi import FastAPI + +from mcp_docugen.auth import ApiKeyAuthMiddleware +from mcp_docugen.config import Settings +from mcp_docugen.generation_store import GenerationStore +from mcp_docugen.http_routes import build_http_app +from mcp_docugen.llm_client import OpenRouterClient +from mcp_docugen.mcp_tools import build_mcp_server +from mcp_docugen.renderer import Renderer +from mcp_docugen.template_store import TemplateStore + +logger = logging.getLogger("mcp_docugen") + + +async def build_app(settings: Settings | None = None) -> FastAPI: + settings = settings or Settings() + settings.data_dir.mkdir(parents=True, exist_ok=True) + templates_dir = settings.data_dir / "templates" + generated_dir = settings.data_dir / "generated" + db_path = settings.data_dir / "mcp_docugen.db" + + template_store = TemplateStore(base_dir=templates_dir) + generation_store = GenerationStore( + db_path=db_path, generated_dir=generated_dir + ) + await generation_store.init() + + llm = OpenRouterClient( + api_key=settings.openrouter_api_key, + base_url=settings.openrouter_base_url, + timeout=settings.llm_timeout_seconds, + ) + renderer = Renderer( + template_store=template_store, + generation_store=generation_store, + llm=llm, + public_base_url=settings.public_base_url, + default_model=settings.llm_model_default, + asset_ttl_days=settings.asset_ttl_days, + max_image_size_mb=settings.max_image_size_mb, + ) + + mcp = build_mcp_server(template_store, renderer) + + @asynccontextmanager + async def lifespan(app: FastAPI): + cleanup_task = asyncio.create_task( + _periodic_cleanup(generation_store, interval_seconds=24 * 3600) + ) + try: + yield + finally: + cleanup_task.cancel() + try: + await cleanup_task + except asyncio.CancelledError: + pass + + app = build_http_app(template_store, generation_store) + app.router.lifespan_context = lifespan + app.mount("/mcp", mcp.streamable_http_app()) + + app.add_middleware( + ApiKeyAuthMiddleware, + api_key=settings.api_key, + exempt_paths={"/health"}, + ) + + app.state.settings = settings + app.state.template_store = template_store + app.state.generation_store = generation_store + app.state.renderer = renderer + app.state.llm = llm + + return app + + +async def _periodic_cleanup( + generation_store: GenerationStore, interval_seconds: int +) -> None: + while True: + try: + removed = await generation_store.cleanup_expired() + if removed: + logger.info("cleanup_expired removed %d assets", removed) + except Exception as exc: # noqa: BLE001 + logger.error("cleanup task error: %s", exc) + await asyncio.sleep(interval_seconds) + + +def run() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s %(message)s", + ) + host = os.environ.get("HOST", "0.0.0.0") + port = int(os.environ.get("PORT", "9100")) + app = asyncio.run(build_app()) + uvicorn.run(app, host=host, port=port) diff --git a/services/mcp-docugen/src/mcp_docugen/mcp_tools.py b/services/mcp-docugen/src/mcp_docugen/mcp_tools.py new file mode 100644 index 0000000..3707d87 --- /dev/null +++ b/services/mcp-docugen/src/mcp_docugen/mcp_tools.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from mcp.server.fastmcp import FastMCP + +from mcp_docugen.models import TemplateFrontmatter +from mcp_docugen.renderer import Renderer +from mcp_docugen.template_store import TemplateStore + + +def build_mcp_server( + template_store: TemplateStore, renderer: Renderer +) -> FastMCP: + mcp = FastMCP("mcp-docugen") + + @mcp.tool() + async def template_create( + name: str, + frontmatter: dict, + body: str, + assets: list[dict] | None = None, + ) -> dict: + """Create a new template. Fails if the name already exists.""" + fm = TemplateFrontmatter(**frontmatter) + await template_store.create( + name=name, frontmatter=fm, body=body, assets=assets + ) + return {"name": name, "status": "created"} + + @mcp.tool() + async def template_update( + name: str, + frontmatter: dict, + body: str, + assets: list[dict] | None = None, + ) -> dict: + """Update an existing template. Replaces frontmatter, body, and (if given) assets.""" + fm = TemplateFrontmatter(**frontmatter) + await template_store.update( + name=name, frontmatter=fm, body=body, assets=assets + ) + return {"name": name, "status": "updated"} + + @mcp.tool() + async def template_delete(name: str) -> dict: + """Delete a template and all its assets.""" + await template_store.delete(name) + return {"name": name, "status": "deleted"} + + @mcp.tool() + async def template_get(name: str) -> dict: + """Return a template's frontmatter, body, and asset list.""" + loaded = await template_store.get(name) + return { + "name": loaded.frontmatter.name, + "frontmatter": loaded.frontmatter.model_dump(exclude_none=True), + "body": loaded.body, + "assets": [a.model_dump() for a in loaded.assets], + "updated_at": loaded.updated_at.isoformat(), + } + + @mcp.tool() + async def template_list() -> list[dict]: + """List all templates with name, description, and updated_at.""" + summaries = await template_store.list() + return [s.model_dump(mode="json") for s in summaries] + + @mcp.tool() + async def document_generate( + template_name: str, + content_md: str, + variables: dict, + instructions: str | None = None, + ) -> dict: + """Generate a Markdown document from a template, content, and variables.""" + result = await renderer.generate( + template_name=template_name, + content_md=content_md, + variables=variables, + instructions=instructions, + ) + return result.model_dump(mode="json") + + mcp.tools = { + "template_create": template_create, + "template_update": template_update, + "template_delete": template_delete, + "template_get": template_get, + "template_list": template_list, + "document_generate": document_generate, + } + + return mcp diff --git a/services/mcp-docugen/src/mcp_docugen/renderer.py b/services/mcp-docugen/src/mcp_docugen/renderer.py new file mode 100644 index 0000000..637052d --- /dev/null +++ b/services/mcp-docugen/src/mcp_docugen/renderer.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +import base64 +import binascii +import re +import uuid + +import aiofiles +from pydantic import ValidationError + +from mcp_docugen.generation_store import ( + EphemeralAssetRecord, + GenerationRecord, + GenerationStore, +) +from mcp_docugen.llm_client import LLMError, OpenRouterClient +from mcp_docugen.models import ( + GenerationResult, + ImageVariable, + TemplateFrontmatter, + TokenUsage, +) +from mcp_docugen.template_store import TemplateStore + + +class RendererError(Exception): + pass + + +class MissingVariables(RendererError): + def __init__(self, missing: list[str]) -> None: + super().__init__(f"missing required variables: {missing}") + self.missing = missing + + +class InvalidVariableType(RendererError): + pass + + +class InvalidImageEncoding(RendererError): + pass + + +class ImageTooLarge(RendererError): + pass + + +class AssetStorageError(RendererError): + pass + + +_ASSET_TEMPLATE_REF_RE = re.compile(r"\.\/assets\/([A-Za-z0-9._-]+)") + + +class Renderer: + def __init__( + self, + template_store: TemplateStore, + generation_store: GenerationStore, + llm: OpenRouterClient, + public_base_url: str, + default_model: str, + asset_ttl_days: int, + max_image_size_mb: int, + ) -> None: + self.template_store = template_store + self.generation_store = generation_store + self.llm = llm + self.public_base_url = public_base_url.rstrip("/") + self.default_model = default_model + self.asset_ttl_days = asset_ttl_days + self.max_image_size_bytes = max_image_size_mb * 1024 * 1024 + + async def generate( + self, + template_name: str, + content_md: str, + variables: dict, + instructions: str | None, + ) -> GenerationResult: + loaded = await self.template_store.get(template_name) + gen_id = str(uuid.uuid4()) + + validated = self._validate_variables(loaded.frontmatter, variables) + ephemeral_urls, var_strings = await self._materialize_image_variables( + gen_id, validated + ) + + system_prompt = self._build_system_prompt( + loaded.frontmatter, loaded.body, var_strings + ) + user_prompt = self._build_user_prompt(content_md, var_strings, instructions) + model = loaded.frontmatter.model or self.default_model + + try: + response = await self.llm.chat( + model=model, system=system_prompt, user=user_prompt + ) + except LLMError as exc: + await self.generation_store.record_generation( + GenerationRecord( + id=gen_id, + template_name=template_name, + model=model, + tokens_in=0, + tokens_out=0, + cost_usd=0.0, + success=False, + error_msg=f"{type(exc).__name__}: {exc}", + ) + ) + raise + + await self.generation_store.record_generation( + GenerationRecord( + id=gen_id, + template_name=template_name, + model=response.model, + tokens_in=response.tokens_in, + tokens_out=response.tokens_out, + cost_usd=response.cost_usd, + success=True, + error_msg=None, + ) + ) + + expires_at = None + if ephemeral_urls: + from pathlib import Path + + info = await self.generation_store.get_ephemeral_asset( + gen_id, Path(ephemeral_urls[0]).name + ) + expires_at = info.expires_at if info else None + + return GenerationResult( + generation_id=gen_id, + markdown=response.text, + model=response.model, + tokens=TokenUsage(input=response.tokens_in, output=response.tokens_out), + cost_usd=response.cost_usd, + ephemeral_assets_urls=ephemeral_urls, + ephemeral_expires_at=expires_at, + ) + + def _validate_variables( + self, fm: TemplateFrontmatter, variables: dict + ) -> dict: + missing = [ + v.name for v in fm.required_variables if v.name not in variables + ] + if missing: + raise MissingVariables(missing) + + validated: dict = {} + for var_def in fm.required_variables: + raw = variables[var_def.name] + if var_def.type == "string": + if not isinstance(raw, str): + raise InvalidVariableType( + f"{var_def.name}: expected string, got {type(raw).__name__}" + ) + validated[var_def.name] = raw + elif var_def.type == "image": + if not isinstance(raw, dict): + raise InvalidVariableType( + f"{var_def.name}: expected image object, got {type(raw).__name__}" + ) + try: + img = ImageVariable(**raw) + except ValidationError as exc: + raise InvalidVariableType(f"{var_def.name}: {exc}") from exc + validated[var_def.name] = img + return validated + + async def _materialize_image_variables( + self, gen_id: str, validated: dict + ) -> tuple[list[str], dict[str, str]]: + urls: list[str] = [] + var_strings: dict[str, str] = {} + gen_dir = self.generation_store.generated_dir / gen_id + + for name, value in validated.items(): + if isinstance(value, ImageVariable): + try: + raw_bytes = base64.b64decode(value.data_b64, validate=True) + except (binascii.Error, ValueError) as exc: + raise InvalidImageEncoding(f"{name}: {exc}") from exc + if len(raw_bytes) > self.max_image_size_bytes: + raise ImageTooLarge( + f"{name}: {len(raw_bytes)} bytes > {self.max_image_size_bytes}" + ) + ext = { + "image/png": "png", + "image/jpeg": "jpg", + "image/webp": "webp", + }[value.mime] + filename = f"{name}.{ext}" + file_path = gen_dir / filename + gen_dir.mkdir(parents=True, exist_ok=True) + try: + async with aiofiles.open(file_path, "wb") as f: + await f.write(raw_bytes) + except OSError as exc: + raise AssetStorageError(str(exc)) from exc + await self.generation_store.register_ephemeral_asset( + EphemeralAssetRecord( + generation_id=gen_id, + var_name=name, + file_path=str(file_path), + mime=value.mime, + ttl_days=self.asset_ttl_days, + ) + ) + url = ( + f"{self.public_base_url}/generated/{gen_id}/{filename}" + ) + urls.append(url) + var_strings[name] = url + else: + var_strings[name] = str(value) + return urls, var_strings + + def _build_system_prompt( + self, + fm: TemplateFrontmatter, + body: str, + var_strings: dict[str, str], + ) -> str: + template_name = fm.name + + def _replace_asset(match: re.Match) -> str: + filename = match.group(1) + return f"{self.public_base_url}/assets/{template_name}/{filename}" + + resolved = _ASSET_TEMPLATE_REF_RE.sub(_replace_asset, body) + for name, value in var_strings.items(): + resolved = resolved.replace(f"{{{{{name}}}}}", value) + return resolved + + def _build_user_prompt( + self, + content_md: str, + var_strings: dict[str, str], + instructions: str | None, + ) -> str: + parts = ["## Raw content", content_md] + if var_strings: + parts.append("## Variables") + for name, value in var_strings.items(): + parts.append(f"- {name}: {value}") + if instructions: + parts.append("## Additional instructions") + parts.append(instructions) + return "\n\n".join(parts) diff --git a/services/mcp-docugen/tests/integration/test_generate_flow.py b/services/mcp-docugen/tests/integration/test_generate_flow.py new file mode 100644 index 0000000..7ea7951 --- /dev/null +++ b/services/mcp-docugen/tests/integration/test_generate_flow.py @@ -0,0 +1,76 @@ +import base64 + +import httpx +import pytest +import respx + +from mcp_docugen.main import build_app +from mcp_docugen.models import TemplateFrontmatter, TemplateVariable + + +def _openrouter_success(text: str = "# Generated"): + return httpx.Response( + 200, + json={ + "id": "g", + "choices": [{"message": {"role": "assistant", "content": text}}], + "model": "anthropic/claude-sonnet-4", + "usage": { + "prompt_tokens": 50, + "completion_tokens": 100, + "total_cost": 0.02, + }, + }, + ) + + +@pytest.fixture +async def app_env(tmp_path, monkeypatch): + monkeypatch.setenv("API_KEY", "test-api-key") + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-test") + monkeypatch.setenv("PUBLIC_BASE_URL", "https://mcp.test.com") + monkeypatch.setenv("DATA_DIR", str(tmp_path)) + app = await build_app() + return app, tmp_path + + +@respx.mock +async def test_full_generation_flow_records_and_serves_asset(app_env): + app, tmp_path = app_env + respx.post("https://openrouter.ai/api/v1/chat/completions").mock( + return_value=_openrouter_success("# Done") + ) + + template_store = app.state.template_store + renderer = app.state.renderer + generation_store = app.state.generation_store + + fm = TemplateFrontmatter( + name="report", + description="x", + required_variables=[TemplateVariable(name="foto", type="image")], + ) + await template_store.create( + name="report", frontmatter=fm, body="![foto]({{foto}})" + ) + + png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 50 + img_var = { + "kind": "image", + "data_b64": base64.b64encode(png).decode(), + "mime": "image/png", + } + + result = await renderer.generate( + template_name="report", + content_md="content", + variables={"foto": img_var}, + instructions=None, + ) + + assert result.markdown == "# Done" + stats = await generation_store.get_stats() + assert stats["success"] == 1 + + gen_files = list((tmp_path / "generated" / result.generation_id).iterdir()) + assert len(gen_files) == 1 diff --git a/services/mcp-docugen/tests/integration/test_http_routes.py b/services/mcp-docugen/tests/integration/test_http_routes.py new file mode 100644 index 0000000..b3e9e1d --- /dev/null +++ b/services/mcp-docugen/tests/integration/test_http_routes.py @@ -0,0 +1,122 @@ +import base64 + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from mcp_docugen.auth import ApiKeyAuthMiddleware +from mcp_docugen.generation_store import EphemeralAssetRecord, GenerationStore +from mcp_docugen.http_routes import build_http_app +from mcp_docugen.models import TemplateFrontmatter +from mcp_docugen.template_store import TemplateStore + + +@pytest.fixture +async def env(tmp_path): + template_store = TemplateStore(base_dir=tmp_path / "templates") + generation_store = GenerationStore( + db_path=tmp_path / "gen.db", + generated_dir=tmp_path / "generated", + ) + await generation_store.init() + + inner = build_http_app(template_store, generation_store) + app = FastAPI() + app.add_middleware( + ApiKeyAuthMiddleware, api_key="test-key", exempt_paths={"/health"} + ) + app.mount("/", inner) + return TestClient(app), template_store, generation_store, tmp_path + + +def _headers(): + return {"Authorization": "Bearer test-key"} + + +def test_health_no_auth(env): + tc, *_ = env + r = tc.get("/health") + assert r.status_code == 200 + assert r.json() == {"status": "ok"} + + +async def test_assets_serve_existing_file(env): + tc, template_store, _, _ = env + fm = TemplateFrontmatter(name="brand", description="x") + assets = [ + { + "filename": "logo.png", + "data_b64": base64.b64encode(b"\x89PNG").decode(), + "mime": "image/png", + } + ] + await template_store.create( + name="brand", frontmatter=fm, body="b", assets=assets + ) + + r = tc.get("/assets/brand/logo.png", headers=_headers()) + assert r.status_code == 200 + assert r.content == b"\x89PNG" + + +def test_assets_require_auth(env): + tc, *_ = env + r = tc.get("/assets/brand/logo.png") + assert r.status_code == 401 + + +def test_assets_path_traversal_rejected(env): + tc, *_ = env + # Starlette decodifica %2F in / prima del routing, quindi la path + # non matcha /assets/{template}/{filename} e torna 404. 400/404 entrambi safe. + r = tc.get("/assets/brand/..%2Fevil.png", headers=_headers()) + assert r.status_code in (400, 404) + + +def test_assets_missing_returns_404(env): + tc, *_ = env + r = tc.get("/assets/brand/missing.png", headers=_headers()) + assert r.status_code == 404 + + +async def test_generated_fresh_returns_file(env): + tc, _, generation_store, tmp_path = env + gen_dir = tmp_path / "generated" / "g-1" + gen_dir.mkdir(parents=True) + (gen_dir / "foto.png").write_bytes(b"image-bytes") + await generation_store.register_ephemeral_asset( + EphemeralAssetRecord( + generation_id="g-1", + var_name="foto", + file_path=str(gen_dir / "foto.png"), + mime="image/png", + ttl_days=30, + ) + ) + r = tc.get("/generated/g-1/foto.png", headers=_headers()) + assert r.status_code == 200 + assert r.content == b"image-bytes" + + +async def test_generated_expired_returns_410(env): + tc, _, generation_store, tmp_path = env + gen_dir = tmp_path / "generated" / "g-old" + gen_dir.mkdir(parents=True) + (gen_dir / "foto.png").write_bytes(b"x") + await generation_store.register_ephemeral_asset( + EphemeralAssetRecord( + generation_id="g-old", + var_name="foto", + file_path=str(gen_dir / "foto.png"), + mime="image/png", + ttl_days=-1, + ) + ) + r = tc.get("/generated/g-old/foto.png", headers=_headers()) + assert r.status_code == 410 + + +def test_generated_unknown_returns_410(env): + tc, *_ = env + r = tc.get("/generated/unknown-id/file.png", headers=_headers()) + assert r.status_code == 410 diff --git a/services/mcp-docugen/tests/integration/test_mcp_tools.py b/services/mcp-docugen/tests/integration/test_mcp_tools.py new file mode 100644 index 0000000..84179b8 --- /dev/null +++ b/services/mcp-docugen/tests/integration/test_mcp_tools.py @@ -0,0 +1,144 @@ +from unittest.mock import AsyncMock + +import pytest + +from mcp_docugen.generation_store import GenerationStore +from mcp_docugen.llm_client import LLMResponse +from mcp_docugen.mcp_tools import build_mcp_server +from mcp_docugen.renderer import Renderer +from mcp_docugen.template_store import TemplateStore + + +@pytest.fixture +async def mcp_env(tmp_path): + template_store = TemplateStore(base_dir=tmp_path / "templates") + generation_store = GenerationStore( + db_path=tmp_path / "gen.db", + generated_dir=tmp_path / "generated", + ) + await generation_store.init() + llm = AsyncMock() + llm.chat.return_value = LLMResponse( + text="# Out", + model="m", + tokens_in=5, + tokens_out=10, + cost_usd=0.01, + latency_ms=50, + ) + renderer = Renderer( + template_store=template_store, + generation_store=generation_store, + llm=llm, + public_base_url="https://mcp.example.com", + default_model="m", + asset_ttl_days=30, + max_image_size_mb=10, + ) + mcp = build_mcp_server(template_store, renderer) + return mcp, template_store + + +async def test_template_create_and_get(mcp_env): + mcp, _ = mcp_env + await mcp.tools["template_create"]( + name="demo", + frontmatter={"name": "demo", "description": "d"}, + body="body", + assets=None, + ) + got = await mcp.tools["template_get"](name="demo") + assert got["name"] == "demo" + assert got["body"] == "body" + + +async def test_template_create_duplicate_errors(mcp_env): + mcp, _ = mcp_env + await mcp.tools["template_create"]( + name="demo", + frontmatter={"name": "demo", "description": "x"}, + body="b", + assets=None, + ) + with pytest.raises(Exception): + await mcp.tools["template_create"]( + name="demo", + frontmatter={"name": "demo", "description": "x"}, + body="b", + assets=None, + ) + + +async def test_template_list_returns_summaries(mcp_env): + mcp, _ = mcp_env + await mcp.tools["template_create"]( + name="a", + frontmatter={"name": "a", "description": "A"}, + body="b", + assets=None, + ) + await mcp.tools["template_create"]( + name="b", + frontmatter={"name": "b", "description": "B"}, + body="b", + assets=None, + ) + result = await mcp.tools["template_list"]() + names = sorted(t["name"] for t in result) + assert names == ["a", "b"] + + +async def test_template_update(mcp_env): + mcp, _ = mcp_env + await mcp.tools["template_create"]( + name="u", + frontmatter={"name": "u", "description": "d"}, + body="old", + assets=None, + ) + await mcp.tools["template_update"]( + name="u", + frontmatter={"name": "u", "description": "new"}, + body="new", + assets=None, + ) + got = await mcp.tools["template_get"](name="u") + assert got["body"] == "new" + assert got["frontmatter"]["description"] == "new" + + +async def test_template_delete(mcp_env): + mcp, _ = mcp_env + await mcp.tools["template_create"]( + name="d", + frontmatter={"name": "d", "description": "x"}, + body="b", + assets=None, + ) + await mcp.tools["template_delete"](name="d") + with pytest.raises(Exception): + await mcp.tools["template_get"](name="d") + + +async def test_document_generate_happy_path(mcp_env): + mcp, _ = mcp_env + await mcp.tools["template_create"]( + name="g", + frontmatter={ + "name": "g", + "description": "x", + "required_variables": [{"name": "cliente", "type": "string"}], + }, + body="Cliente: {{cliente}}", + assets=None, + ) + result = await mcp.tools["document_generate"]( + template_name="g", + content_md="# Ordine", + variables={"cliente": "ACME"}, + instructions=None, + ) + assert result["markdown"] == "# Out" + assert result["model"] == "m" + assert result["tokens"]["input"] == 5 + assert "generation_id" in result diff --git a/services/mcp-docugen/tests/unit/test_renderer.py b/services/mcp-docugen/tests/unit/test_renderer.py new file mode 100644 index 0000000..cfdd644 --- /dev/null +++ b/services/mcp-docugen/tests/unit/test_renderer.py @@ -0,0 +1,274 @@ +import base64 +from unittest.mock import AsyncMock + +import pytest + +from mcp_docugen.generation_store import GenerationStore +from mcp_docugen.llm_client import LLMResponse, OpenRouterClient +from mcp_docugen.models import TemplateFrontmatter, TemplateVariable +from mcp_docugen.renderer import ( + ImageTooLarge, + InvalidImageEncoding, + InvalidVariableType, + MissingVariables, + Renderer, +) +from mcp_docugen.template_store import TemplateStore + + +@pytest.fixture +async def env(tmp_path): + template_store = TemplateStore(base_dir=tmp_path / "templates") + generation_store = GenerationStore( + db_path=tmp_path / "gen.db", + generated_dir=tmp_path / "generated", + ) + await generation_store.init() + llm = AsyncMock(spec=OpenRouterClient) + llm.chat.return_value = LLMResponse( + text="# Output", + model="anthropic/claude-sonnet-4", + tokens_in=10, + tokens_out=20, + cost_usd=0.001, + latency_ms=100, + ) + renderer = Renderer( + template_store=template_store, + generation_store=generation_store, + llm=llm, + public_base_url="https://mcp.example.com", + default_model="anthropic/claude-sonnet-4", + asset_ttl_days=30, + max_image_size_mb=10, + ) + fm = TemplateFrontmatter( + name="fattura", + description="x", + required_variables=[TemplateVariable(name="cliente", type="string")], + ) + await template_store.create( + name="fattura", frontmatter=fm, body="Cliente: {{cliente}}" + ) + return renderer, llm + + +async def test_generate_happy_path_string_var(env): + renderer, llm = env + result = await renderer.generate( + template_name="fattura", + content_md="# Ordine", + variables={"cliente": "ACME"}, + instructions=None, + ) + assert result.markdown == "# Output" + assert result.model == "anthropic/claude-sonnet-4" + assert result.tokens.input == 10 + call_kwargs = llm.chat.await_args.kwargs + assert "ACME" in call_kwargs["user"] + assert "Cliente: ACME" in call_kwargs["system"] + + +async def test_generate_missing_required_variable_raises(env): + renderer, _ = env + with pytest.raises(MissingVariables) as exc: + await renderer.generate( + template_name="fattura", + content_md="x", + variables={}, + instructions=None, + ) + assert "cliente" in str(exc.value) + + +async def test_generate_wrong_type_raises(env): + renderer, _ = env + with pytest.raises(InvalidVariableType): + await renderer.generate( + template_name="fattura", + content_md="x", + variables={ + "cliente": { + "kind": "image", + "data_b64": "x", + "mime": "image/png", + } + }, + instructions=None, + ) + + +async def test_generate_image_variable_is_saved_and_rewritten(tmp_path): + template_store = TemplateStore(base_dir=tmp_path / "templates") + generation_store = GenerationStore( + db_path=tmp_path / "gen.db", + generated_dir=tmp_path / "generated", + ) + await generation_store.init() + llm = AsyncMock(spec=OpenRouterClient) + llm.chat.return_value = LLMResponse( + text="# OK", model="m", tokens_in=1, tokens_out=1, cost_usd=0, latency_ms=10 + ) + renderer = Renderer( + template_store=template_store, + generation_store=generation_store, + llm=llm, + public_base_url="https://mcp.example.com", + default_model="m", + asset_ttl_days=30, + max_image_size_mb=10, + ) + fm = TemplateFrontmatter( + name="report", + description="x", + required_variables=[TemplateVariable(name="foto", type="image")], + ) + await template_store.create(name="report", frontmatter=fm, body="![foto]({{foto}})") + + png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100 + img_var = { + "kind": "image", + "data_b64": base64.b64encode(png).decode(), + "mime": "image/png", + } + result = await renderer.generate( + template_name="report", + content_md="content", + variables={"foto": img_var}, + instructions=None, + ) + + assert result.ephemeral_assets_urls + url = result.ephemeral_assets_urls[0] + assert url.startswith("https://mcp.example.com/generated/") + assert url.endswith(".png") + gen_dir = tmp_path / "generated" / result.generation_id + assert gen_dir.exists() + assert any(gen_dir.iterdir()) + + call_kwargs = llm.chat.await_args.kwargs + assert "https://mcp.example.com/generated/" in call_kwargs["system"] + + +async def test_image_too_large_raises(env): + renderer, _ = env + big = b"x" * (11 * 1024 * 1024) + img_var = { + "kind": "image", + "data_b64": base64.b64encode(big).decode(), + "mime": "image/png", + } + fm = TemplateFrontmatter( + name="big", + description="x", + required_variables=[TemplateVariable(name="foto", type="image")], + ) + await renderer.template_store.create(name="big", frontmatter=fm, body="{{foto}}") + with pytest.raises(ImageTooLarge): + await renderer.generate( + template_name="big", + content_md="x", + variables={"foto": img_var}, + instructions=None, + ) + + +async def test_invalid_base64_raises(env): + renderer, _ = env + fm = TemplateFrontmatter( + name="bad", + description="x", + required_variables=[TemplateVariable(name="foto", type="image")], + ) + await renderer.template_store.create(name="bad", frontmatter=fm, body="{{foto}}") + with pytest.raises(InvalidImageEncoding): + await renderer.generate( + template_name="bad", + content_md="x", + variables={ + "foto": { + "kind": "image", + "data_b64": "!!!not-base64!!!", + "mime": "image/png", + } + }, + instructions=None, + ) + + +async def test_template_asset_paths_are_rewritten(tmp_path): + template_store = TemplateStore(base_dir=tmp_path / "templates") + generation_store = GenerationStore( + db_path=tmp_path / "gen.db", + generated_dir=tmp_path / "generated", + ) + await generation_store.init() + llm = AsyncMock(spec=OpenRouterClient) + llm.chat.return_value = LLMResponse( + text="out", model="m", tokens_in=1, tokens_out=1, cost_usd=0, latency_ms=10 + ) + renderer = Renderer( + template_store=template_store, + generation_store=generation_store, + llm=llm, + public_base_url="https://mcp.example.com", + default_model="m", + asset_ttl_days=30, + max_image_size_mb=10, + ) + fm = TemplateFrontmatter(name="brand", description="x") + assets = [ + { + "filename": "logo.png", + "data_b64": base64.b64encode(b"\x89PNG").decode(), + "mime": "image/png", + } + ] + await template_store.create( + name="brand", + frontmatter=fm, + body="Header ![logo](./assets/logo.png) footer", + assets=assets, + ) + await renderer.generate( + template_name="brand", + content_md="x", + variables={}, + instructions=None, + ) + system_prompt = llm.chat.await_args.kwargs["system"] + assert "https://mcp.example.com/assets/brand/logo.png" in system_prompt + assert "./assets/logo.png" not in system_prompt + + +async def test_generate_uses_frontmatter_model_override(tmp_path): + template_store = TemplateStore(base_dir=tmp_path / "templates") + generation_store = GenerationStore( + db_path=tmp_path / "gen.db", + generated_dir=tmp_path / "generated", + ) + await generation_store.init() + llm = AsyncMock(spec=OpenRouterClient) + llm.chat.return_value = LLMResponse( + text="out", + model="openai/gpt-4o", + tokens_in=1, + tokens_out=1, + cost_usd=0, + latency_ms=10, + ) + renderer = Renderer( + template_store=template_store, + generation_store=generation_store, + llm=llm, + public_base_url="https://mcp.example.com", + default_model="anthropic/claude-sonnet-4", + asset_ttl_days=30, + max_image_size_mb=10, + ) + fm = TemplateFrontmatter(name="x", description="y", model="openai/gpt-4o") + await template_store.create(name="x", frontmatter=fm, body="b") + await renderer.generate( + template_name="x", content_md="c", variables={}, instructions=None + ) + assert llm.chat.await_args.kwargs["model"] == "openai/gpt-4o"