knowledge-worker 0.6.0__py3-none-any.whl

ollama_proxy/extractor_adapter.py

@@ -0,0 +1,168 @@
+"""
+extractor_adapter.py — local-Gemma drop-in for mygraph's extractor (v1.5).
+
+mygraph/extractor.py uses Claude with native tool-use to emit a structured
+candidates payload. Ollama models don't speak Anthropic tool-use; they speak
+Ollama's `format` (JSON mode + optional JSON schema). This adapter:
+
+  - reuses the same prompt template + schema as mygraph/extractor.py
+  - calls Ollama with `format=<schema>` for constrained JSON output
+  - returns the same dict shape, so validator.py / review.py / merge.py work
+    without modification
+
+Usage (drop-in):
+    from ollama_proxy.extractor_adapter import extract as gemma_extract
+    payload = gemma_extract(Path("notes.md"), out_path=Path("out.json"))
+
+Or via the helper CLI:
+    python ollama_proxy/extractor_adapter.py path/to/file.md [out.json]
+
+Env:
+    OLLAMA_DEFAULT_MODEL  default model tag (gemma4:e4b)
+    OLLAMA_BASE_URL       http://127.0.0.1:11434
+    GEMMA_NUM_CTX         optional context window override (default 8192)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+
+# Make mygraph/ importable regardless of cwd.
+_HERE = Path(__file__).resolve().parent
+_MYGRAPH = _HERE.parent / "mygraph"
+if str(_MYGRAPH) not in sys.path:
+    sys.path.insert(0, str(_MYGRAPH))
+
+from mygraph import Graph, NODE_TYPES, EDGE_TYPES, slug  # noqa: E402
+from extractor import (  # noqa: E402
+    EXTRACTION_TOOL,
+    PROMPT_TEMPLATE,
+    build_source_decl,
+    ensure_provenance_edges,
+)
+
+try:
+    import httpx
+except ImportError as e:
+    raise SystemExit(
+        "extractor_adapter: missing dep. Run: pip install httpx"
+    ) from e
+
+
+OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/")
+DEFAULT_MODEL = os.environ.get("OLLAMA_DEFAULT_MODEL", "gemma4:e4b")
+NUM_CTX = int(os.environ.get("GEMMA_NUM_CTX", "8192"))
+
+# JSON schema mirrors mygraph/extractor.py's tool input_schema.
+RESPONSE_SCHEMA = EXTRACTION_TOOL["input_schema"]
+
+
+def _ollama_chat(prompt: str, model: str, schema: dict) -> dict:
+    """Call /api/chat with format=schema, return parsed JSON content."""
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "stream": False,
+        "format": schema,                 # Ollama structured-output mode
+        "options": {"num_ctx": NUM_CTX, "temperature": 0.2},
+    }
+    with httpx.Client(timeout=600.0) as c:
+        r = c.post(f"{OLLAMA_BASE_URL}/api/chat", json=payload)
+    if r.status_code != 200:
+        raise RuntimeError(
+            f"extractor_adapter: ollama returned {r.status_code}: {r.text[:500]}"
+        )
+    body = r.json()
+    content = (body.get("message") or {}).get("content", "")
+    if not content:
+        raise RuntimeError(f"extractor_adapter: empty response from {model}")
+    try:
+        return json.loads(content)
+    except json.JSONDecodeError as e:
+        # Fallback: strip code fences if the model wrapped output despite format hint.
+        stripped = content.strip()
+        if stripped.startswith("```"):
+            stripped = stripped.strip("`")
+            stripped = stripped.split("\n", 1)[1] if "\n" in stripped else stripped
+            stripped = stripped.rsplit("```", 1)[0] if stripped.endswith("```") else stripped
+            return json.loads(stripped)
+        raise RuntimeError(
+            f"extractor_adapter: model returned non-JSON: {content[:300]}"
+        ) from e
+
+
+def extract(md_path: Path, out_path: Path | None = None,
+            model: str = DEFAULT_MODEL) -> dict:
+    """End-to-end extract via local Ollama. Same return shape as
+    mygraph.extractor.extract — drop-in compatible with validate()/review()/merge()."""
+    g = Graph.load()
+    decl = build_source_decl(md_path)
+    source_text = md_path.read_text(encoding="utf-8")
+    existing_ids = sorted(g.nodes.keys())
+    prompt = PROMPT_TEMPLATE.format(
+        source_id=decl["source_id"],
+        source_label=decl["source_label"],
+        source_path=decl["source_path"],
+        node_types=", ".join(sorted(NODE_TYPES)),
+        edge_types=", ".join(sorted(EDGE_TYPES)),
+        existing_ids="\n".join(f"  - {i}" for i in existing_ids),
+        source_text=source_text,
+    )
+    payload = _ollama_chat(prompt, model=model, schema=RESPONSE_SCHEMA)
+
+    # Ensure required keys exist so validator never crashes on a missing top-level.
+    payload.setdefault("source", {
+        "id": decl["source_id"],
+        "label": decl["source_label"],
+        "body": "",
+    })
+    payload.setdefault("nodes", [])
+    payload.setdefault("edges", [])
+    injected = ensure_provenance_edges(payload)
+    if injected:
+        print(
+            "extractor_adapter: gateway returned missing provenance edges; "
+            f"synthesized {injected} MENTIONED_IN edges.",
+            file=sys.stderr,
+        )
+
+    payload.setdefault("_meta", {})
+    payload["_meta"]["source_path"] = decl["source_path"]
+    payload["_meta"]["ingested_at"] = decl["ingested_at"]
+    payload["_meta"]["model"] = model
+    payload["_meta"]["backend"] = "ollama"
+
+    if out_path:
+        out_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    return payload
+
+
+def main(argv: list[str]) -> int:
+    if len(argv) < 2:
+        print("Usage: python extractor_adapter.py <path/to/file.md> [out.json] [--model NAME]")
+        return 1
+    md = Path(argv[1]).expanduser().resolve()
+    out = None
+    model = DEFAULT_MODEL
+    rest = argv[2:]
+    if "--model" in rest:
+        i = rest.index("--model")
+        model = rest[i + 1]
+        del rest[i:i + 2]
+    if rest:
+        out = Path(rest[0]).expanduser().resolve()
+    else:
+        out = md.parent / f"{md.stem}.candidates.gemma.json"
+    payload = extract(md, out, model=model)
+    print(f"extractor_adapter: model={model}  "
+          f"nodes={len(payload.get('nodes', []))}  "
+          f"edges={len(payload.get('edges', []))}  ->  {out}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

ollama_proxy/proxy.py

@@ -0,0 +1,143 @@
+"""
+proxy.py — thin Ollama-API-compatible passthrough with logging (v1.5).
+
+Why this exists alongside server.py (the MCP server):
+  - AnythingLLM and most external LLM clients speak Ollama's REST API directly,
+    NOT MCP. So the MCP server (server.py) is for Claude Code/Cowork; this
+    proxy is for AnythingLLM and other Ollama-compatible clients.
+  - Tailscale-expose THIS port (default 11435) instead of raw 11434 so we get:
+      * per-call logging (latency, model, prompt size, error)
+      * a single chokepoint for future auth/rate-limit if we ever leave the tailnet
+  - 100% pass-through: same request/response shape as Ollama. AnythingLLM points
+    at this URL and treats it as Ollama.
+
+Usage:
+    python proxy.py                      # binds 127.0.0.1:11435 → 127.0.0.1:11434
+    python proxy.py --host 0.0.0.0       # listen on all interfaces (tailnet)
+    python proxy.py --port 8080
+    OLLAMA_BASE_URL=http://127.0.0.1:11434 python proxy.py
+
+Deps: fastapi uvicorn httpx
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+
+try:
+    import httpx
+    from fastapi import FastAPI, Request
+    from fastapi.responses import JSONResponse, Response, StreamingResponse
+    import uvicorn
+except ImportError as e:
+    print(
+        "proxy.py: missing deps. Install with:\n"
+        "    pip install fastapi 'uvicorn[standard]' httpx",
+        file=sys.stderr,
+    )
+    raise SystemExit(1) from e
+
+
+OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/")
+LOG_PATH = Path(os.environ.get(
+    "OLLAMA_PROXY_LOG",
+    str(Path(__file__).parent / "proxy.log.jsonl"),
+))
+
+app = FastAPI(title="ollama-proxy", version="1.5.0")
+_client = httpx.AsyncClient(base_url=OLLAMA_BASE_URL, timeout=300.0)
+
+
+def _log(record: dict) -> None:
+    record.setdefault("ts", datetime.now(timezone.utc).isoformat())
+    try:
+        with LOG_PATH.open("a", encoding="utf-8") as f:
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+    except OSError:
+        pass  # never let logging break a request
+
+
+@app.get("/healthz")
+async def healthz() -> dict:
+    try:
+        r = await _client.get("/api/tags")
+        r.raise_for_status()
+        return {"ok": True, "ollama": OLLAMA_BASE_URL,
+                "models": [m["name"] for m in r.json().get("models", [])]}
+    except Exception as e:
+        return JSONResponse({"ok": False, "error": str(e)}, status_code=502)
+
+
+@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
+async def passthrough(path: str, request: Request):
+    """Mirror every request to Ollama. Stream responses if the client asked
+    for stream=true; otherwise buffer + log."""
+    body = await request.body()
+    headers = {k: v for k, v in request.headers.items()
+               if k.lower() not in {"host", "content-length"}}
+    target = f"/{path}"
+    qs = request.url.query
+    if qs:
+        target = f"{target}?{qs}"
+
+    # detect streaming vs buffered
+    streaming = False
+    if body and request.method == "POST":
+        try:
+            payload = json.loads(body)
+            streaming = bool(payload.get("stream", True if path.startswith("api/") else False))
+            model = payload.get("model")
+            prompt_len = len(json.dumps(payload.get("messages") or payload.get("prompt") or ""))
+        except Exception:
+            payload, model, prompt_len = None, None, len(body)
+    else:
+        payload, model, prompt_len = None, None, 0
+
+    t0 = time.perf_counter()
+    if streaming:
+        async def stream_iter():
+            async with _client.stream(request.method, target, content=body, headers=headers) as r:
+                status = r.status_code
+                async for chunk in r.aiter_raw():
+                    yield chunk
+                _log({
+                    "kind": "request", "method": request.method, "path": target,
+                    "status": status, "model": model, "prompt_len": prompt_len,
+                    "duration_ms": int((time.perf_counter() - t0) * 1000),
+                    "stream": True,
+                })
+        return StreamingResponse(stream_iter(), media_type="application/x-ndjson")
+
+    r = await _client.request(request.method, target, content=body, headers=headers)
+    _log({
+        "kind": "request", "method": request.method, "path": target,
+        "status": r.status_code, "model": model, "prompt_len": prompt_len,
+        "duration_ms": int((time.perf_counter() - t0) * 1000),
+        "stream": False,
+    })
+    return Response(content=r.content, status_code=r.status_code,
+                    media_type=r.headers.get("content-type"))
+
+
+def main(argv: list[str]) -> int:
+    p = argparse.ArgumentParser(description="Logging passthrough proxy in front of Ollama")
+    p.add_argument("--host", default=os.environ.get("PROXY_HOST", "127.0.0.1"),
+                   help="Bind host (set 0.0.0.0 to be reachable on the tailnet)")
+    p.add_argument("--port", type=int, default=int(os.environ.get("PROXY_PORT", "11435")),
+                   help="Bind port (default 11435, Ollama's 11434 + 1)")
+    args = p.parse_args(argv)
+    print(f"proxy: forwarding http://{args.host}:{args.port}/* -> {OLLAMA_BASE_URL}/*",
+          file=sys.stderr)
+    print(f"proxy: log -> {LOG_PATH}", file=sys.stderr)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

ollama_proxy/server.py

@@ -0,0 +1,194 @@
+"""
+server.py — MCP server wrapping local Ollama (v1.5).
+
+Stateless MCP server exposing four tools:
+  - chat(model, messages, options?)        Chat completion via /api/chat
+  - generate(model, prompt, options?)      Single-shot completion via /api/generate
+  - list_models()                          Installed models via /api/tags
+  - embed(model, input)                    Embeddings via /api/embed
+
+Transports:
+  - stdio  (default)   — for Claude Code, Cowork, local MCP clients
+  - sse                — for remote MCP clients over HTTP (optional)
+
+Default model: gemma4:e4b (override with OLLAMA_DEFAULT_MODEL env var).
+Ollama base : http://127.0.0.1:11434 (override with OLLAMA_BASE_URL).
+
+Usage:
+    python server.py                  # stdio transport (default)
+    python server.py --sse --port 7421
+    OLLAMA_DEFAULT_MODEL=gemma4:latest python server.py
+
+Deps: mcp[cli] httpx
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Any
+
+import httpx
+
+try:
+    from mcp.server.fastmcp import FastMCP
+except ImportError as e:
+    print(
+        "server.py: missing dependency `mcp`. Install with:\n"
+        "    pip install 'mcp[cli]' httpx",
+        file=sys.stderr,
+    )
+    raise SystemExit(1) from e
+
+
+OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/")
+DEFAULT_MODEL = os.environ.get("OLLAMA_DEFAULT_MODEL", "gemma4:e4b")
+TIMEOUT_S = float(os.environ.get("OLLAMA_TIMEOUT_S", "300"))
+
+mcp = FastMCP("ollama-proxy")
+_client = httpx.Client(base_url=OLLAMA_BASE_URL, timeout=TIMEOUT_S)
+
+
+def _post(path: str, payload: dict) -> dict:
+    r = _client.post(path, json=payload)
+    r.raise_for_status()
+    return r.json()
+
+
+def _get(path: str) -> dict:
+    r = _client.get(path)
+    r.raise_for_status()
+    return r.json()
+
+
+@mcp.tool()
+def chat(
+    messages: list[dict],
+    model: str | None = None,
+    options: dict | None = None,
+    format: str | None = None,
+) -> dict:
+    """Chat completion via Ollama /api/chat.
+
+    Args:
+        messages: list of {"role": "user|assistant|system", "content": str}
+        model:    Ollama model tag. Defaults to OLLAMA_DEFAULT_MODEL (gemma4:e4b).
+        options:  Ollama runtime options (temperature, num_ctx, etc.).
+        format:   "json" to force JSON-mode output.
+
+    Returns the raw Ollama response: {message: {role, content}, done, ...}
+    """
+    payload: dict[str, Any] = {
+        "model": model or DEFAULT_MODEL,
+        "messages": messages,
+        "stream": False,
+    }
+    if options:
+        payload["options"] = options
+    if format:
+        payload["format"] = format
+    return _post("/api/chat", payload)
+
+
+@mcp.tool()
+def generate(
+    prompt: str,
+    model: str | None = None,
+    options: dict | None = None,
+    format: str | None = None,
+    system: str | None = None,
+) -> dict:
+    """Single-shot completion via Ollama /api/generate.
+
+    Args:
+        prompt:   the user prompt
+        model:    defaults to OLLAMA_DEFAULT_MODEL
+        options:  Ollama runtime options
+        format:   "json" to force JSON output
+        system:   optional system prompt
+    """
+    payload: dict[str, Any] = {
+        "model": model or DEFAULT_MODEL,
+        "prompt": prompt,
+        "stream": False,
+    }
+    if options:
+        payload["options"] = options
+    if format:
+        payload["format"] = format
+    if system:
+        payload["system"] = system
+    return _post("/api/generate", payload)
+
+
+@mcp.tool()
+def list_models() -> dict:
+    """List models installed in the local Ollama instance."""
+    return _get("/api/tags")
+
+
+@mcp.tool()
+def embed(input: str | list[str], model: str | None = None) -> dict:
+    """Get embeddings via Ollama /api/embed.
+
+    Args:
+        input: a string or list of strings to embed
+        model: embedding-capable model tag (e.g. "nomic-embed-text")
+    """
+    payload = {
+        "model": model or DEFAULT_MODEL,
+        "input": input,
+    }
+    return _post("/api/embed", payload)
+
+
+def main(argv: list[str]) -> int:
+    p = argparse.ArgumentParser(description="MCP server wrapping local Ollama")
+    p.add_argument("--sse", action="store_true",
+                   help="Use SSE/HTTP transport instead of stdio (for remote MCP clients)")
+    p.add_argument("--port", type=int, default=int(os.environ.get("MCP_PORT", "7421")),
+                   help="HTTP port for --sse mode (default 7421)")
+    p.add_argument("--host", default=os.environ.get("MCP_HOST", "127.0.0.1"),
+                   help="Bind host for --sse mode (default 127.0.0.1; set 0.0.0.0 for tailnet)")
+    args = p.parse_args(argv)
+
+    # Sanity check: is Ollama reachable?
+    try:
+        _get("/api/tags")
+    except Exception as e:
+        print(
+            f"server.py: cannot reach Ollama at {OLLAMA_BASE_URL} ({e}).\n"
+            "Is the Ollama daemon running? Try: ollama serve",
+            file=sys.stderr,
+        )
+        # Don't exit — let the MCP client see the error per-call. But warn loudly.
+
+    if args.sse:
+        # FastMCP's SSE transport binds host/port via settings.
+        mcp.settings.host = args.host
+        mcp.settings.port = args.port
+        print(f"server.py: SSE transport on http://{args.host}:{args.port}/sse",
+              file=sys.stderr)
+        mcp.run(transport="sse")
+    else:
+        # Friendly hint when launched interactively — MCP-stdio is for clients,
+        # not humans. If stdin is a TTY, the user almost certainly meant to use
+        # --sse or register this in an MCP client config.
+        if sys.stdin.isatty():
+            print(
+                "server.py: stdio transport ready. This server is meant to be\n"
+                "  spawned by an MCP client (Claude Code, Cowork, Claude Desktop)\n"
+                "  via subprocess pipes — running it in a terminal will produce\n"
+                "  JSON-RPC parse errors on every keystroke. To test interactively:\n"
+                "    npx @modelcontextprotocol/inspector python server.py\n"
+                "  Or run with --sse to expose an HTTP/SSE endpoint instead.\n"
+                "  Press Ctrl-C to exit.\n",
+                file=sys.stderr,
+            )
+        mcp.run(transport="stdio")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))