el-primor 3.0.0

package/turbo-vec-mcp.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+"""
+turbo-vec-mcp.py — MCP Server para TurboVec (El Primor v3.0.0)
+
+Servidor MCP (Model Context Protocol) que expone búsqueda semántica,
+almacenamiento de documentos y estadísticas usando TurboVec + SQLite.
+
+Tools:
+  - turbo_vec_search(query, k=5, collection=null)
+  - turbo_vec_store(title, content, collection="default", metadata_json="{}")
+  - turbo_vec_stats()
+"""
+
+import json
+import sys
+import sqlite3
+from pathlib import Path
+
+# ── Configuración ───────────────────────────────────────────────────────────
+PRIMOR_DIR = Path.home() / ".primor"
+DB_PATH = PRIMOR_DIR / "primor.db"
+MODEL_NAME = "sentence-transformers/multilingual-e5-small"
+MODEL_DIR = PRIMOR_DIR / "models" / "multilingual-e5-small"
+
+# Estado global lazy
+_model = None
+_index = None
+_db: sqlite3.Connection | None = None
+
+# ── Inicialización lazy ────────────────────────────────────────────────────
+def _ensure_model():
+    global _model, _index
+    if _model is None:
+        from sentence_transformers import SentenceTransformer
+        _model = SentenceTransformer(
+            str(MODEL_DIR) if MODEL_DIR.exists() else MODEL_NAME,
+            cache_folder=str(PRIMOR_DIR / "models"),
+            device="cpu",
+        )
+    if _index is None:
+        from turbo_vec import IdMapIndex
+        _index = IdMapIndex(dim=_model.get_sentence_embedding_dimension())
+    return _model, _index
+
+
+def _ensure_db():
+    global _db
+    if _db is None:
+        _db = sqlite3.connect(str(DB_PATH))
+        _db.row_factory = sqlite3.Row
+        _db.execute("PRAGMA journal_mode=WAL")
+    return _db
+
+
+def _embed(text: str) -> list[float]:
+    model, _ = _ensure_model()
+    return model.encode(text, normalize_embeddings=True).tolist()
+
+
+# ── Tools ───────────────────────────────────────────────────────────────────
+def turbo_vec_search(query: str, k: int = 5, collection: str | None = None):
+    """Búsqueda semántica en la base de conocimiento."""
+    _, index = _ensure_model()
+    db = _ensure_db()
+
+    query_vec = _embed(query)
+    results = index.search(query_vec, top_k=k)
+
+    output = []
+    for vec_id, score in results:
+        if score < 0.3:
+            continue
+        if collection:
+            row = db.execute(
+                "SELECT * FROM docs WHERE vec_id = ? AND collection = ?",
+                (vec_id, collection),
+            ).fetchone()
+        else:
+            row = db.execute(
+                "SELECT * FROM docs WHERE vec_id = ?", (vec_id,)
+            ).fetchone()
+        if row:
+            output.append({
+                "id": row["id"],
+                "vec_id": row["vec_id"],
+                "title": row["title"],
+                "content": row["content"][:500],
+                "collection": row["collection"],
+                "score": round(float(score), 4),
+                "created_at": row["created_at"],
+            })
+
+    return {"query": query, "k": k, "results": output, "total": len(output)}
+
+
+def turbo_vec_store(
+    title: str,
+    content: str,
+    collection: str = "default",
+    metadata_json: str = "{}",
+):
+    """Almacena un documento con embedding semántico."""
+    model, index = _ensure_model()
+    db = _ensure_db()
+
+    vec = _embed(f"{title}\n{content}")
+    vec_id = index.add(vec)
+
+    try:
+        metadata = json.loads(metadata_json)
+    except json.JSONDecodeError:
+        metadata = {}
+
+    cursor = db.execute(
+        """INSERT INTO docs (title, content, collection, metadata, vec_id)
+           VALUES (?, ?, ?, ?, ?)""",
+        (title, content, collection, json.dumps(metadata), vec_id),
+    )
+    db.commit()
+
+    return {
+        "id": cursor.lastrowid,
+        "vec_id": vec_id,
+        "title": title,
+        "collection": collection,
+        "content_length": len(content),
+    }
+
+
+def turbo_vec_stats():
+    """Estadísticas de la base de conocimiento."""
+    db = _ensure_db()
+    _, index = _ensure_model()
+
+    total_docs = db.execute("SELECT COUNT(*) FROM docs").fetchone()[0]
+    collections = db.execute(
+        "SELECT collection, COUNT(*) as cnt FROM docs GROUP BY collection ORDER BY cnt DESC"
+    ).fetchall()
+    index_size = index.size() if index else 0
+
+    return {
+        "total_docs": total_docs,
+        "index_size": index_size,
+        "collections": [
+            {"name": r["collection"], "count": r["cnt"]} for r in collections
+        ],
+        "model": MODEL_NAME,
+        "db_path": str(DB_PATH),
+    }
+
+
+# ── MCP JSON-RPC sobre stdio ────────────────────────────────────────────────
+TOOLS = {
+    "turbo_vec_search": turbo_vec_search,
+    "turbo_vec_store": turbo_vec_store,
+    "turbo_vec_stats": turbo_vec_stats,
+}
+
+TOOL_SCHEMAS = [
+    {
+        "name": "turbo_vec_search",
+        "description": "Búsqueda semántica en la base de conocimiento TurboVec. Encuentra documentos similares por significado.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "query": {"type": "string", "description": "Texto de búsqueda"},
+                "k": {"type": "integer", "description": "Número de resultados (default: 5)", "default": 5},
+                "collection": {"type": "string", "description": "Filtrar por colección (opcional)"},
+            },
+            "required": ["query"],
+        },
+    },
+    {
+        "name": "turbo_vec_store",
+        "description": "Almacena un documento con embedding semántico en TurboVec.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "title": {"type": "string", "description": "Título del documento"},
+                "content": {"type": "string", "description": "Contenido del documento"},
+                "collection": {"type": "string", "description": "Colección (default: 'default')", "default": "default"},
+                "metadata_json": {"type": "string", "description": "Metadatos en JSON", "default": "{}"},
+            },
+            "required": ["title", "content"],
+        },
+    },
+    {
+        "name": "turbo_vec_stats",
+        "description": "Estadísticas de la base de conocimiento: total de documentos, colecciones, tamaño del índice.",
+        "inputSchema": {"type": "object", "properties": {}},
+    },
+]
+
+
+def handle_request(req: dict) -> dict:
+    """Procesa una petición JSON-RPC."""
+    method = req.get("method", "")
+    req_id = req.get("id")
+
+    if method == "initialize":
+        return {
+            "jsonrpc": "2.0",
+            "id": req_id,
+            "result": {
+                "protocolVersion": "2024-11-05",
+                "serverInfo": {
+                    "name": "turbo-vec-mcp",
+                    "version": "3.0.0",
+                },
+                "capabilities": {"tools": {}},
+            },
+        }
+
+    if method == "tools/list":
+        return {
+            "jsonrpc": "2.0",
+            "id": req_id,
+            "result": {"tools": TOOL_SCHEMAS},
+        }
+
+    if method == "tools/call":
+        tool_name = req.get("params", {}).get("name", "")
+        arguments = req.get("params", {}).get("arguments", {})
+        if tool_name not in TOOLS:
+            return {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "error": {"code": -32601, "message": f"Tool not found: {tool_name}"},
+            }
+        try:
+            result = TOOLS[tool_name](**arguments)
+            return {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {"content": [{"type": "text", "text": json.dumps(result, ensure_ascii=False)}]},
+            }
+        except Exception as e:
+            return {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "error": {"code": -32000, "message": str(e)},
+            }
+
+    return {
+        "jsonrpc": "2.0",
+        "id": req_id,
+        "error": {"code": -32601, "message": f"Unknown method: {method}"},
+    }
+
+
+def main():
+    """Loop principal MCP sobre stdio."""
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            req = json.loads(line)
+            resp = handle_request(req)
+            sys.stdout.write(json.dumps(resp, ensure_ascii=False) + "\n")
+            sys.stdout.flush()
+        except json.JSONDecodeError:
+            continue
+
+
+if __name__ == "__main__":
+    main()