diff options
| author | tanaka takashi <isazaka@gmail.com> | 2026-06-26 23:54:17 +0900 |
|---|---|---|
| committer | tanaka takashi <isazaka@gmail.com> | 2026-06-26 23:54:17 +0900 |
| commit | 3a9e8e18df4c596360da54000f65e1bebc6039c2 (patch) | |
| tree | 0c6dc5fd75cc8f751ca6d7dca01c7e7b5fb3c9f5 /src/ollama_web/routes | |
| parent | a296c2c3a294648c04286a5c027deddbbdc8d5ae (diff) | |
MCP機能追加
Diffstat (limited to 'src/ollama_web/routes')
| -rw-r--r-- | src/ollama_web/routes/chat.py | 34 | ||||
| -rw-r--r-- | src/ollama_web/routes/mcp.py | 119 |
2 files changed, 149 insertions, 4 deletions
diff --git a/src/ollama_web/routes/chat.py b/src/ollama_web/routes/chat.py index 8b5c3f4..75aabd7 100644 --- a/src/ollama_web/routes/chat.py +++ b/src/ollama_web/routes/chat.py @@ -16,6 +16,7 @@ from starlette.requests import Request from starlette.responses import StreamingResponse from ..config import settings +from ..mcp import collect_mcp_tools, make_mcp_executor from ..sessions import SessionStore, is_valid_id from ..tools.fetch import pop_fetched_files from ..tools.helper.image import resize_image @@ -57,8 +58,8 @@ def _parse_messages(raw: list[dict[str, Any]]) -> list[Message]: return out -def _session_aware_registry(session_id: str | None) -> ToolRegistry: - """Return a registry whose search_and_fetch receives the session id.""" +async def _build_registry(session_id: str | None) -> ToolRegistry: + """Return a registry with built-in tools and any configured MCP tools.""" reg = default_registry() original = reg.get("search_and_fetch") @@ -73,6 +74,24 @@ def _session_aware_registry(session_id: str | None) -> ToolRegistry: reg._tools["search_and_fetch"] = wrapped # noqa: SLF001 + try: + mcp_tools = await collect_mcp_tools() + except Exception as exc: # noqa: BLE001 + logger.warning("failed to collect MCP tools: %s", exc) + mcp_tools = [] + + tool_names = [t.get("function", {}).get("name") for t in mcp_tools] + logger.info("collected %d MCP tool(s): %s", len(mcp_tools), tool_names) + + if mcp_tools: + executor = make_mcp_executor() + for tool_def in mcp_tools: + fn = tool_def.get("function", {}) + name = fn.get("name") + if not name: + continue + reg.register_mcp_tool(name, tool_def, executor) + return reg @@ -188,7 +207,7 @@ async def _chat_event_stream( from ollama_web import llm # local import to avoid circular at module load - registry = _session_aware_registry(session_id) + registry = await _build_registry(session_id) total_image_bytes = sum( len(img.value) for img in (getattr(messages[-1], "images", None) or []) @@ -209,6 +228,13 @@ async def _chat_event_stream( sorted(capabilities), ) + # Fallback: some cloud/proxy models support tools even though ollama + # does not advertise the capability. Enable tools when the model name + # hints at a cloud endpoint and tools are not already reported. + if "tools" not in capabilities and ":cloud" in model: + capabilities.add("tools") + logger.info("enabled tools capability fallback for cloud model %s", model) + async for event in llm.astream_chat_with_tools( messages, model, think=think, host=host, registry=registry, capabilities=capabilities, @@ -304,4 +330,4 @@ async def chat(request: Request) -> StreamingResponse: async def _single_event(data: str) -> AsyncIterator[bytes]: """Yield a single SSE event and finish.""" - yield f"data: {data}\n\n".encode() + yield f"data: {data}\n\n".encode()
\ No newline at end of file diff --git a/src/ollama_web/routes/mcp.py b/src/ollama_web/routes/mcp.py new file mode 100644 index 0000000..85d4bb3 --- /dev/null +++ b/src/ollama_web/routes/mcp.py @@ -0,0 +1,119 @@ +"""MCP server configuration API endpoints.""" + +from __future__ import annotations + +import logging +from typing import Any + +from starlette.requests import Request +from starlette.responses import JSONResponse + +from ..mcp import load_mcp_config, save_mcp_config + +logger = logging.getLogger("ollama_web.mcp_route") + + +_ALLOWED_STDIO_FIELDS = {"command", "args", "env", "cwd", "encoding", "encoding_error_handler"} +_ALLOWED_HTTP_FIELDS = {"url", "headers", "timeout", "sse_read_timeout", "terminate_on_close"} + + +async def get_servers(request: Request) -> JSONResponse: + """GET /api/mcp/servers: return the current MCP server configuration.""" + data_dir = request.app.state.settings.data_dir + config = load_mcp_config(data_dir) + return JSONResponse(config) + + +async def put_servers(request: Request) -> JSONResponse: + """PUT /api/mcp/servers: persist a new MCP server configuration. + + The body must be a JSON object with a ``mcpServers`` key. Each server is + validated to contain only recognized fields and a single transport. + """ + try: + body = await request.json() + except Exception as exc: # noqa: BLE001 + logger.warning("invalid MCP config JSON: %s", exc) + return JSONResponse({"error": "Invalid JSON body"}, status_code=400) + + if not isinstance(body, dict) or "mcpServers" not in body: + return JSONResponse( + {"error": "Body must contain a 'mcpServers' object"}, + status_code=400, + ) + + raw_servers = body["mcpServers"] + if not isinstance(raw_servers, dict): + return JSONResponse( + {"error": "'mcpServers' must be an object"}, + status_code=400, + ) + + cleaned: dict[str, Any] = {"mcpServers": {}} + errors: list[str] = [] + for name, server in raw_servers.items(): + if not isinstance(server, dict): + errors.append(f"Server '{name}' must be an object") + continue + + if "command" in server: + valid = _allowed_fields(server, _ALLOWED_STDIO_FIELDS, name, errors) + if valid: + cleaned["mcpServers"][name] = _clean_stdio_server(server) + elif "url" in server: + valid = _allowed_fields(server, _ALLOWED_HTTP_FIELDS, name, errors) + if valid: + cleaned["mcpServers"][name] = _clean_http_server(server) + else: + errors.append( + f"Server '{name}' must define either 'command' (stdio) or 'url' (http)" + ) + + if errors: + return JSONResponse({"error": errors}, status_code=400) + + data_dir = request.app.state.settings.data_dir + save_mcp_config(data_dir, cleaned) + return JSONResponse(cleaned) + + +def _allowed_fields( + server: dict[str, Any], + allowed: set[str], + name: str, + errors: list[str], +) -> bool: + unknown = set(server.keys()) - allowed + if unknown: + errors.append(f"Server '{name}' has unknown fields: {', '.join(sorted(unknown))}") + return False + return True + + +def _clean_stdio_server(server: dict[str, Any]) -> dict[str, Any]: + out: dict[str, Any] = {"command": str(server["command"])} + if "args" in server: + out["args"] = [str(a) for a in server["args"]] + if "env" in server: + out["env"] = {str(k): str(v) for k, v in server["env"].items()} + if "cwd" in server and server["cwd"]: + out["cwd"] = str(server["cwd"]) + for key in ("encoding", "encoding_error_handler"): + if key in server: + out[key] = str(server[key]) + return out + + +def _clean_http_server(server: dict[str, Any]) -> dict[str, Any]: + out: dict[str, Any] = {"url": str(server["url"])} + if "headers" in server: + out["headers"] = {str(k): str(v) for k, v in server["headers"].items()} + for key in ("timeout", "sse_read_timeout"): + if key in server: + try: + out[key] = float(server[key]) + except (TypeError, ValueError): + pass + if "terminate_on_close" in server: + out["terminate_on_close"] = bool(server["terminate_on_close"]) + return out
\ No newline at end of file |
