diff --git a/src/cmdforge/cli/__init__.py b/src/cmdforge/cli/__init__.py index 5d17ef6..67a0ff2 100644 --- a/src/cmdforge/cli/__init__.py +++ b/src/cmdforge/cli/__init__.py @@ -196,6 +196,13 @@ def main(): p_reg_update_readme.add_argument("--all", action="store_true", dest="update_all", help="Update README for all published tools that have a local README.md") p_reg_update_readme.set_defaults(func=cmd_registry) + # registry describe (semantic search) + p_reg_describe = registry_sub.add_parser("describe", help="Find tools by describing what you need") + p_reg_describe.add_argument("query", help="Natural language description (e.g., 'something that analyzes CSV files')") + p_reg_describe.add_argument("-n", "--limit", type=int, default=20, help="Max results (default: 20)") + p_reg_describe.add_argument("--json", action="store_true", help="Output as JSON") + p_reg_describe.set_defaults(func=cmd_registry) + # registry my-tools p_reg_mytools = registry_sub.add_parser("my-tools", help="List your published tools") p_reg_mytools.set_defaults(func=cmd_registry) diff --git a/src/cmdforge/cli/registry_commands.py b/src/cmdforge/cli/registry_commands.py index 1204d0e..de3071d 100644 --- a/src/cmdforge/cli/registry_commands.py +++ b/src/cmdforge/cli/registry_commands.py @@ -37,6 +37,8 @@ def cmd_registry(args): return _cmd_registry_status(args) elif args.registry_cmd == "browse": return _cmd_registry_browse(args) + elif args.registry_cmd == "describe": + return _cmd_registry_describe(args) elif args.registry_cmd == "config": return _cmd_registry_config(args) else: @@ -53,6 +55,7 @@ def cmd_registry(args): print(" my-tools List your published tools") print(" status Check moderation status of a tool") print(" browse Browse tools (GUI)") + print(" describe Find tools by describing what you need (AI)") print(" config [action] Manage registry settings (admin)") return 0 @@ -894,6 +897,67 @@ def _cmd_registry_status(args): return 0 +def _cmd_registry_describe(args): + """Find tools by describing what you need (semantic search).""" + from ..registry_client import RegistryError, get_client + + query = args.query + limit = getattr(args, 'limit', 20) + + try: + client = get_client() + result = client.semantic_search(query, limit=limit) + + # JSON output + if getattr(args, 'json', False): + print(json.dumps(result, indent=2)) + return 0 + + available = result.get("available", False) + tools = result.get("data", []) + error = result.get("error") + + if error: + print(f"Error: {error}", file=sys.stderr) + return 1 + + if not available: + print("AI search is not available on this registry.") + print("The Ollama embedding service may be offline or the feature is disabled.") + return 1 + + if not tools: + print(f'No tools found matching: "{query}"') + print("Try different phrasing or use keyword search: cmdforge registry search ") + return 0 + + print(f'Found {len(tools)} tools matching: "{query}"\n') + + for i, tool in enumerate(tools, 1): + similarity = tool.get("similarity", 0) + pct = f"{similarity * 100:.0f}%" + owner = tool.get("owner", "") + name = tool.get("name", "") + desc = tool.get("description", "") + + print(f" {i}. {owner}/{name} ({pct} match)") + if desc: + print(f" {desc[:70]}{'...' if len(desc) > 70 else ''}") + print() + + except RegistryError as e: + if e.code == "CONNECTION_ERROR": + print("Could not connect to the registry.", file=sys.stderr) + else: + print(f"Error: {e.message}", file=sys.stderr) + return 1 + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + return 1 + + return 0 + + def _cmd_registry_browse(args): """Browse tools (GUI).""" from ..gui import run_gui diff --git a/src/cmdforge/gui/pages/registry_page.py b/src/cmdforge/gui/pages/registry_page.py index 1c3f68b..fed6479 100644 --- a/src/cmdforge/gui/pages/registry_page.py +++ b/src/cmdforge/gui/pages/registry_page.py @@ -108,6 +108,24 @@ class InstallWorker(QThread): self.error.emit(str(e)) +class SemanticSearchWorker(QThread): + """Background worker for semantic search.""" + finished = Signal(dict) # {"data": [...], "available": bool} + error = Signal(str) + + def __init__(self, query: str): + super().__init__() + self.query = query + + def run(self): + try: + client = RegistryClient() + result = client.semantic_search(self.query) + self.finished.emit(result) + except Exception as e: + self.error.emit(str(e)) + + class RegistryPage(QWidget): """Registry browser page.""" @@ -115,6 +133,7 @@ class RegistryPage(QWidget): super().__init__() self.main_window = main_window self._search_worker = None + self._semantic_worker = None self._install_worker = None self._versions_worker = None self._selected_tool = None @@ -198,6 +217,36 @@ class RegistryPage(QWidget): layout.addWidget(filters_box) + # Semantic search row + semantic_box = QWidget() + semantic_layout = QHBoxLayout(semantic_box) + semantic_layout.setContentsMargins(0, 0, 0, 0) + semantic_layout.setSpacing(8) + + ai_badge = QLabel("AI") + ai_badge.setStyleSheet( + "background: #805ad5; color: white; border-radius: 4px; " + "padding: 2px 8px; font-weight: bold; font-size: 11px;" + ) + ai_badge.setFixedHeight(24) + semantic_layout.addWidget(ai_badge) + + self.semantic_input = QLineEdit() + self.semantic_input.setPlaceholderText("Describe what you need...") + self.semantic_input.setToolTip( + "Describe what you're looking for in plain language " + "(e.g., 'something that analyzes CSV files')" + ) + self.semantic_input.returnPressed.connect(self._do_semantic_search) + semantic_layout.addWidget(self.semantic_input, 1) + + self.btn_semantic = QPushButton("Find") + self.btn_semantic.setToolTip("Search by description using AI") + self.btn_semantic.clicked.connect(self._do_semantic_search) + semantic_layout.addWidget(self.btn_semantic) + + layout.addWidget(semantic_box) + # Active tags display self.tags_widget = QWidget() self.tags_layout = QHBoxLayout(self.tags_widget) @@ -638,6 +687,98 @@ class RegistryPage(QWidget): self.btn_browse.setEnabled(True) self.status_label.setText(f"Error: {error}") + def _do_semantic_search(self): + """Perform semantic search.""" + query = self.semantic_input.text().strip() + if not query: + return + + self.btn_semantic.setEnabled(False) + self.btn_semantic.setText("Searching...") + self.status_label.setText("Searching by description...") + self.results_table.setRowCount(0) + + self._semantic_worker = SemanticSearchWorker(query) + self._semantic_worker.finished.connect(self._on_semantic_results) + self._semantic_worker.error.connect(self._on_semantic_error) + self._semantic_worker.start() + + def _on_semantic_results(self, result: dict): + """Handle semantic search results.""" + self.btn_semantic.setEnabled(True) + self.btn_semantic.setText("Find") + + available = result.get("available", False) + tools = result.get("data", []) + error = result.get("error") + + if error: + self.status_label.setText(f"Search error: {error}") + return + + if not available: + self.status_label.setText("AI search not available on this registry") + return + + if not tools: + self.status_label.setText("No matching tools found. Try different phrasing.") + return + + self.status_label.setText(f"Found {len(tools)} tools by description match") + + # Disable pagination for semantic results + self.btn_prev.setEnabled(False) + self.btn_next.setEnabled(False) + + self.results_table.setRowCount(len(tools)) + for row, tool in enumerate(tools): + tool_name = tool.get("name", "") + similarity = tool.get("similarity", 0) + + # Installed indicator + installed_item = QTableWidgetItem() + if tool_name in self._installed_tools: + installed_version = self._installed_tools[tool_name] + registry_version = tool.get("version", "1.0.0") + if installed_version != registry_version: + installed_item.setText("↑") + installed_item.setToolTip(f"Update available: {installed_version} → {registry_version}") + installed_item.setForeground(QColor("#48bb78")) + else: + installed_item.setText("✓") + installed_item.setToolTip("Installed") + installed_item.setForeground(QColor("#4299e1")) + installed_item.setTextAlignment(Qt.AlignCenter) + self.results_table.setItem(row, 0, installed_item) + + # Name (with similarity %) + name_item = QTableWidgetItem(tool_name) + name_item.setData(Qt.UserRole, tool) + self.results_table.setItem(row, 1, name_item) + + # Owner + self.results_table.setItem(row, 2, QTableWidgetItem(tool.get("owner", ""))) + + # Rating column: show similarity % instead + pct = f"{similarity * 100:.0f}% match" + match_item = QTableWidgetItem(pct) + match_item.setToolTip(f"Semantic similarity: {similarity:.4f}") + self.results_table.setItem(row, 3, match_item) + + # Downloads + downloads = tool.get("downloads", 0) + downloads_str = self._format_downloads(downloads) + self.results_table.setItem(row, 4, QTableWidgetItem(downloads_str)) + + # Version + self.results_table.setItem(row, 5, QTableWidgetItem(tool.get("version", "1.0.0"))) + + def _on_semantic_error(self, error: str): + """Handle semantic search error.""" + self.btn_semantic.setEnabled(True) + self.btn_semantic.setText("Find") + self.status_label.setText(f"Search failed: {error}") + def _prev_page(self): """Go to previous page.""" if self._current_page > 1: diff --git a/src/cmdforge/registry/app.py b/src/cmdforge/registry/app.py index 8cd49f9..590cd9b 100644 --- a/src/cmdforge/registry/app.py +++ b/src/cmdforge/registry/app.py @@ -695,6 +695,55 @@ def create_app() -> Flask: return jsonify({"data": data, "meta": paginate(page, per_page, total)}) + @app.route("/api/v1/tools/semantic-search", methods=["GET"]) + def semantic_search_tools() -> Response: + """Semantic search using Ollama embeddings.""" + from .settings import get_setting + from .embeddings import semantic_search as _semantic_search + + enabled = get_setting(g.db, "embeddings.enabled") + if not enabled: + return error_response( + "SERVICE_UNAVAILABLE", + "Semantic search is disabled", + 503, + ) + + query_text = request.args.get("q", "").strip() + if not query_text: + return error_response("VALIDATION_ERROR", "Missing search query 'q'") + if len(query_text) > 500: + return error_response("VALIDATION_ERROR", "Query too long (max 500 characters)") + + limit = request.args.get("limit", 20, type=int) + if limit is None or limit < 1: + limit = 20 + limit = min(limit, 50) + + ollama_url = get_setting(g.db, "embeddings.ollama_url") + if ollama_url is None: + ollama_url = "http://192.168.0.186:11434" + model = get_setting(g.db, "embeddings.model") + if model is None: + model = "nomic-embed-text" + min_score = get_setting(g.db, "embeddings.min_score") + if min_score is None: + min_score = 0.3 + + results, available = _semantic_search( + g.db, query_text, ollama_url, model, top_k=limit, min_score=min_score, + ) + + return jsonify({ + "data": results, + "meta": { + "query": query_text, + "count": len(results), + "search_type": "semantic", + "available": available, + }, + }) + @app.route("/api/v1/tools/search", methods=["GET"]) def search_tools() -> Response: query_text = request.args.get("q", "").strip() @@ -2641,6 +2690,34 @@ def create_app() -> Flask: ) g.db.commit() + # Best-effort embedding generation (after commit, outside transaction) + # Only embed public + approved tools (matches search/backfill policy) + if visibility == "public" and moderation_status == "approved": + try: + from .embeddings import build_embed_text, generate_embedding, store_embedding + from .settings import get_setting as _get_setting + embed_enabled = _get_setting(g.db, "embeddings.enabled") + if embed_enabled: + embed_text = build_embed_text(name, description, tags) + if embed_text: + ollama_url = _get_setting(g.db, "embeddings.ollama_url") + if ollama_url is None: + ollama_url = "http://192.168.0.186:11434" + embed_model = _get_setting(g.db, "embeddings.model") + if embed_model is None: + embed_model = "nomic-embed-text" + tool_row = query_one( + g.db, + "SELECT id FROM tools WHERE owner = ? AND name = ? AND version = ?", + [owner, name, version], + ) + if tool_row: + vec = generate_embedding(embed_text, ollama_url, embed_model) + if vec: + store_embedding(g.db, tool_row["id"], vec, embed_model) + except Exception: + pass # Never block publish + response = jsonify({ "data": { "owner": owner, @@ -3989,6 +4066,48 @@ def create_app() -> Flask: } }) + # ─── Admin Embeddings ───────────────────────────────────────────────────── + + @app.route("/api/v1/admin/embeddings/backfill", methods=["POST"]) + @require_token + @require_admin + def admin_embeddings_backfill() -> Response: + """Generate embeddings for all public/approved tools missing them.""" + from .embeddings import backfill_all_embeddings + stats = backfill_all_embeddings(g.db) + return jsonify({"data": stats}) + + @app.route("/api/v1/admin/embeddings/status", methods=["GET"]) + @require_token + @require_admin + def admin_embeddings_status() -> Response: + """Get embedding coverage stats.""" + from .settings import get_setting as _get_setting + model = _get_setting(g.db, "embeddings.model") or "nomic-embed-text" + + total = query_one( + g.db, + "SELECT COUNT(*) as cnt FROM tools WHERE visibility = 'public' AND moderation_status = 'approved'", + )["cnt"] + + embedded = query_one( + g.db, + "SELECT COUNT(*) as cnt FROM tool_embeddings WHERE model = ?", + [model], + )["cnt"] + + coverage = (embedded / total * 100) if total > 0 else 0.0 + + return jsonify({ + "data": { + "total_tools": total, + "embedded": embedded, + "coverage_pct": round(coverage, 1), + "model": model, + "enabled": bool(_get_setting(g.db, "embeddings.enabled")), + } + }) + @app.route("/api/v1/admin/reports", methods=["GET"]) @require_moderator def admin_list_reports() -> Response: diff --git a/src/cmdforge/registry/db.py b/src/cmdforge/registry/db.py index 8377574..b55335f 100644 --- a/src/cmdforge/registry/db.py +++ b/src/cmdforge/registry/db.py @@ -413,6 +413,15 @@ CREATE TABLE IF NOT EXISTS password_reset_tokens ( CREATE INDEX IF NOT EXISTS idx_reset_tokens_hash ON password_reset_tokens(token_hash); CREATE INDEX IF NOT EXISTS idx_reset_tokens_publisher ON password_reset_tokens(publisher_id); + +-- Semantic search embeddings +CREATE TABLE IF NOT EXISTS tool_embeddings ( + tool_id INTEGER PRIMARY KEY REFERENCES tools(id) ON DELETE CASCADE, + embedding BLOB NOT NULL, + dimensions INTEGER NOT NULL, + model TEXT NOT NULL, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP +); """ @@ -530,6 +539,21 @@ def migrate_db(conn: sqlite3.Connection) -> None: except sqlite3.OperationalError: pass + # Ensure tool_embeddings table exists (for existing databases) + try: + conn.executescript(""" + CREATE TABLE IF NOT EXISTS tool_embeddings ( + tool_id INTEGER PRIMARY KEY REFERENCES tools(id) ON DELETE CASCADE, + embedding BLOB NOT NULL, + dimensions INTEGER NOT NULL, + model TEXT NOT NULL, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ); + """) + conn.commit() + except sqlite3.OperationalError: + pass + # Ensure indexes exist try: conn.execute("CREATE INDEX IF NOT EXISTS idx_tools_owner ON tools(owner)") diff --git a/src/cmdforge/registry/embeddings.py b/src/cmdforge/registry/embeddings.py new file mode 100644 index 0000000..16e970f --- /dev/null +++ b/src/cmdforge/registry/embeddings.py @@ -0,0 +1,266 @@ +"""Semantic search via Ollama embeddings. + +Generates embeddings for tool metadata, stores them in SQLite, +and performs cosine similarity search for natural language queries. +""" + +from __future__ import annotations + +import math +import sqlite3 +import struct +import time +from typing import Dict, List, Optional, Tuple + +import requests + +from .settings import get_setting + + +def build_embed_text( + name: Optional[str], + description: Optional[str], + tags: Optional[List[str]], +) -> Optional[str]: + """Build the text to embed from tool metadata. + + Joins name, description, and tags with ' | ' separator. + Returns None if all fields are empty/None. + """ + parts = [] + if name and name.strip(): + parts.append(name.strip()) + if description and description.strip(): + parts.append(description.strip()) + if tags: + tag_str = ", ".join(str(t) for t in tags if t) + if tag_str: + parts.append(tag_str) + return " | ".join(parts) if parts else None + + +def generate_embedding( + text: str, + ollama_url: str, + model: str, + timeout: int = 10, +) -> Optional[List[float]]: + """Generate an embedding vector via Ollama's /api/embed endpoint. + + Returns float list or None on any failure. + """ + try: + resp = requests.post( + f"{ollama_url.rstrip('/')}/api/embed", + json={"model": model, "input": text}, + timeout=timeout, + ) + resp.raise_for_status() + data = resp.json() + # Ollama returns {"embeddings": [[...]]} for /api/embed + embeddings = data.get("embeddings") + if embeddings and len(embeddings) > 0: + return embeddings[0] + return None + except Exception: + return None + + +def pack_embedding(vector: List[float]) -> bytes: + """Pack a float vector into a compact binary blob for SQLite storage.""" + return struct.pack(f"<{len(vector)}f", *vector) + + +def unpack_embedding(blob: bytes) -> Optional[List[float]]: + """Unpack a binary blob back into a float vector. + + Returns None on malformed data. + """ + if not blob or len(blob) % 4 != 0: + return None + count = len(blob) // 4 + try: + return list(struct.unpack(f"<{count}f", blob)) + except struct.error: + return None + + +def cosine_similarity(a: List[float], b: List[float]) -> float: + """Compute cosine similarity between two vectors. + + Returns 0.0 on length mismatch or zero-magnitude vectors. + """ + if len(a) != len(b): + return 0.0 + dot = sum(x * y for x, y in zip(a, b)) + norm_a = math.sqrt(sum(x * x for x in a)) + norm_b = math.sqrt(sum(x * x for x in b)) + if norm_a == 0.0 or norm_b == 0.0: + return 0.0 + return dot / (norm_a * norm_b) + + +def store_embedding( + conn: sqlite3.Connection, + tool_id: int, + embedding: List[float], + model: str, + commit: bool = True, +) -> None: + """Upsert an embedding for a tool. + + Set commit=False for batch operations where the caller manages commits. + """ + blob = pack_embedding(embedding) + conn.execute( + """ + INSERT INTO tool_embeddings (tool_id, embedding, dimensions, model, updated_at) + VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(tool_id) DO UPDATE SET + embedding = excluded.embedding, + dimensions = excluded.dimensions, + model = excluded.model, + updated_at = CURRENT_TIMESTAMP + """, + [tool_id, blob, len(embedding), model], + ) + if commit: + conn.commit() + + +def semantic_search( + conn: sqlite3.Connection, + query: str, + ollama_url: str, + model: str, + top_k: int = 20, + min_score: float = 0.3, +) -> Tuple[List[Dict], bool]: + """Perform semantic search over stored embeddings. + + Returns (results, available) where: + - results: list of dicts with tool info + similarity score, ranked + - available: True if Ollama was reachable, False otherwise + + Skips stored embeddings with mismatched model or dimensions. + """ + top_k = min(top_k, 50) + + query_vec = generate_embedding(query, ollama_url, model) + if query_vec is None: + return [], False + + query_dim = len(query_vec) + + # Load all stored embeddings with tool metadata + rows = conn.execute( + """ + SELECT te.tool_id, te.embedding, te.dimensions, te.model, + t.owner, t.name, t.version, t.description, t.category, + t.tags, t.downloads + FROM tool_embeddings te + JOIN tools t ON t.id = te.tool_id + WHERE t.visibility = 'public' AND t.moderation_status = 'approved' + """, + ).fetchall() + + scored = [] + for row in rows: + # Skip model or dimension mismatches + if row["model"] != model or row["dimensions"] != query_dim: + continue + + vec = unpack_embedding(row["embedding"]) + if vec is None: + continue + + score = cosine_similarity(query_vec, vec) + if score >= min_score: + tags = row["tags"] or "[]" + try: + import json + tags_list = json.loads(tags) + except Exception: + tags_list = [] + + scored.append({ + "owner": row["owner"], + "name": row["name"], + "version": row["version"], + "description": row["description"] or "", + "category": row["category"] or "", + "tags": tags_list, + "downloads": row["downloads"] or 0, + "similarity": round(score, 4), + }) + + # Sort by similarity descending + scored.sort(key=lambda x: x["similarity"], reverse=True) + return scored[:top_k], True + + +def backfill_all_embeddings( + conn: sqlite3.Connection, + ollama_url: Optional[str] = None, + model: Optional[str] = None, + batch_size: int = 20, + throttle: float = 0.5, +) -> Dict[str, int]: + """Generate embeddings for all public/approved tools that don't have one. + + Returns stats dict with total, success, failed counts. + """ + if ollama_url is None: + ollama_url = get_setting(conn, "embeddings.ollama_url") + if ollama_url is None: + ollama_url = "http://192.168.0.186:11434" + if model is None: + model = get_setting(conn, "embeddings.model") + if model is None: + model = "nomic-embed-text" + + # Find tools without embeddings (or with stale model) + rows = conn.execute( + """ + SELECT t.id, t.name, t.description, t.tags + FROM tools t + LEFT JOIN tool_embeddings te ON te.tool_id = t.id AND te.model = ? + WHERE t.visibility = 'public' + AND t.moderation_status = 'approved' + AND te.tool_id IS NULL + """, + [model], + ).fetchall() + + stats = {"total": len(rows), "success": 0, "failed": 0} + + for i, row in enumerate(rows): + tags = row["tags"] or "[]" + try: + import json + tags_list = json.loads(tags) + except Exception: + tags_list = [] + + text = build_embed_text(row["name"], row["description"], tags_list) + if text is None: + stats["failed"] += 1 + continue + + embedding = generate_embedding(text, ollama_url, model) + if embedding is None: + stats["failed"] += 1 + else: + store_embedding(conn, row["id"], embedding, model, commit=False) + stats["success"] += 1 + + # Commit in batches + if (i + 1) % batch_size == 0: + conn.commit() + + # Throttle to avoid hammering Ollama + if throttle > 0 and i < len(rows) - 1: + time.sleep(throttle) + + conn.commit() + return stats diff --git a/src/cmdforge/registry/settings.py b/src/cmdforge/registry/settings.py index 960b84a..e2d0b1a 100644 --- a/src/cmdforge/registry/settings.py +++ b/src/cmdforge/registry/settings.py @@ -158,6 +158,36 @@ DEFAULT_SETTINGS: List[Setting] = [ category="moderation", ), + # Embeddings (semantic search) + Setting( + key="embeddings.ollama_url", + value="http://192.168.0.186:11434", + value_type="string", + description="Ollama server URL for generating embeddings", + category="embeddings", + ), + Setting( + key="embeddings.model", + value="nomic-embed-text", + value_type="string", + description="Ollama model to use for embeddings", + category="embeddings", + ), + Setting( + key="embeddings.enabled", + value=True, + value_type="bool", + description="Enable semantic search via embeddings", + category="embeddings", + ), + Setting( + key="embeddings.min_score", + value=0.3, + value_type="float", + description="Minimum cosine similarity threshold for search results (0.0-1.0)", + category="embeddings", + ), + # Rate limits (can override defaults) Setting( key="rate_limit.publish.limit", diff --git a/src/cmdforge/registry_client.py b/src/cmdforge/registry_client.py index 78106eb..4c2f29c 100644 --- a/src/cmdforge/registry_client.py +++ b/src/cmdforge/registry_client.py @@ -1022,6 +1022,47 @@ class RegistryClient: return response.json().get("data", {}) + def semantic_search(self, query: str, limit: int = 20) -> Dict: + """Semantic search by natural language description. + + Returns {"data": [...], "available": bool, "error": str|None} + + available=True, data=[] -> no matching tools + available=False, data=[] -> Ollama unreachable or feature disabled (503) + available=True, error="..." -> validation error (bad query) + """ + try: + response = self._request( + "GET", + "/tools/semantic-search", + params={"q": query, "limit": min(limit, 50)}, + ) + except RegistryError: + return {"data": [], "available": False, "error": None} + + if response.status_code == 503: + return {"data": [], "available": False, "error": None} + + if response.status_code == 400: + # Validation error — service is available but query was bad + try: + err = response.json().get("error", {}) + msg = err.get("message", "Invalid query") + except Exception: + msg = "Invalid query" + return {"data": [], "available": True, "error": msg} + + if response.status_code != 200: + return {"data": [], "available": False, "error": None} + + data = response.json() + meta = data.get("meta", {}) + return { + "data": data.get("data", []), + "available": meta.get("available", True), + "error": None, + } + def get_popular_tools(self, limit: int = 10) -> List[ToolInfo]: """ Get most popular tools. diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..b72fd28 --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,205 @@ +"""Unit tests for semantic search embeddings module.""" + +import json +import sqlite3 + +import pytest + +from cmdforge.registry.embeddings import ( + build_embed_text, + cosine_similarity, + pack_embedding, + store_embedding, + unpack_embedding, +) + + +# --------------------------------------------------------------------------- +# pack / unpack +# --------------------------------------------------------------------------- + +def test_pack_unpack_roundtrip(): + """Pack a vector, unpack it, verify identical.""" + original = [0.1, 0.2, 0.3, -0.5, 1.0, 0.0] + blob = pack_embedding(original) + result = unpack_embedding(blob) + assert result is not None + assert len(result) == len(original) + for a, b in zip(original, result): + assert abs(a - b) < 1e-6 + + +def test_unpack_empty_blob(): + """Empty blob returns None.""" + assert unpack_embedding(b"") is None + + +def test_unpack_malformed_blob(): + """Blob not divisible by 4 returns None.""" + assert unpack_embedding(b"\x00\x01\x02") is None + + +# --------------------------------------------------------------------------- +# cosine_similarity +# --------------------------------------------------------------------------- + +def test_cosine_similarity_identical(): + """Same vector should give similarity of 1.0.""" + vec = [1.0, 2.0, 3.0] + assert abs(cosine_similarity(vec, vec) - 1.0) < 1e-6 + + +def test_cosine_similarity_orthogonal(): + """Orthogonal vectors should give similarity of 0.0.""" + a = [1.0, 0.0, 0.0] + b = [0.0, 1.0, 0.0] + assert abs(cosine_similarity(a, b)) < 1e-6 + + +def test_cosine_similarity_opposite(): + """Opposite vectors should give similarity of -1.0.""" + a = [1.0, 0.0] + b = [-1.0, 0.0] + assert abs(cosine_similarity(a, b) - (-1.0)) < 1e-6 + + +def test_cosine_similarity_dimension_mismatch(): + """Different lengths should return 0.0.""" + a = [1.0, 2.0] + b = [1.0, 2.0, 3.0] + assert cosine_similarity(a, b) == 0.0 + + +def test_cosine_similarity_zero_vector(): + """Zero vector should return 0.0.""" + a = [0.0, 0.0, 0.0] + b = [1.0, 2.0, 3.0] + assert cosine_similarity(a, b) == 0.0 + + +# --------------------------------------------------------------------------- +# build_embed_text +# --------------------------------------------------------------------------- + +def test_build_embed_text(): + """Verify text construction from name/desc/tags.""" + text = build_embed_text("my-tool", "Analyzes CSV files", ["csv", "data"]) + assert text == "my-tool | Analyzes CSV files | csv, data" + + +def test_build_embed_text_no_tags(): + """Handles None tags.""" + text = build_embed_text("my-tool", "A description", None) + assert text == "my-tool | A description" + + +def test_build_embed_text_empty_tags(): + """Handles empty tags list.""" + text = build_embed_text("my-tool", "A description", []) + assert text == "my-tool | A description" + + +def test_build_embed_text_all_empty(): + """All empty/None fields returns None.""" + assert build_embed_text(None, None, None) is None + assert build_embed_text("", "", []) is None + + +def test_build_embed_text_only_name(): + """Name only.""" + text = build_embed_text("my-tool", None, None) + assert text == "my-tool" + + +# --------------------------------------------------------------------------- +# store and retrieve (in-memory SQLite) +# --------------------------------------------------------------------------- + +def _make_db(): + """Create an in-memory SQLite DB with the embeddings table.""" + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA foreign_keys=OFF") # No FK to tools table in test + conn.executescript(""" + CREATE TABLE tools ( + id INTEGER PRIMARY KEY, + owner TEXT, name TEXT, version TEXT, + description TEXT, category TEXT, tags TEXT, + downloads INTEGER DEFAULT 0, + visibility TEXT DEFAULT 'public', + moderation_status TEXT DEFAULT 'approved' + ); + CREATE TABLE tool_embeddings ( + tool_id INTEGER PRIMARY KEY, + embedding BLOB NOT NULL, + dimensions INTEGER NOT NULL, + model TEXT NOT NULL, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ); + """) + return conn + + +def test_store_and_retrieve(): + """Store embedding, verify retrieval.""" + conn = _make_db() + conn.execute( + "INSERT INTO tools (id, owner, name, version) VALUES (1, 'test', 'my-tool', '1.0.0')" + ) + conn.commit() + + vec = [0.1, 0.2, 0.3, 0.4, 0.5] + store_embedding(conn, 1, vec, "nomic-embed-text") + + row = conn.execute("SELECT * FROM tool_embeddings WHERE tool_id = 1").fetchone() + assert row is not None + assert row["dimensions"] == 5 + assert row["model"] == "nomic-embed-text" + + recovered = unpack_embedding(row["embedding"]) + assert recovered is not None + assert len(recovered) == 5 + for a, b in zip(vec, recovered): + assert abs(a - b) < 1e-6 + + +def test_store_upsert(): + """Store should update on conflict.""" + conn = _make_db() + conn.execute( + "INSERT INTO tools (id, owner, name, version) VALUES (1, 'test', 'my-tool', '1.0.0')" + ) + conn.commit() + + store_embedding(conn, 1, [1.0, 2.0], "model-a") + store_embedding(conn, 1, [3.0, 4.0, 5.0], "model-b") + + row = conn.execute("SELECT * FROM tool_embeddings WHERE tool_id = 1").fetchone() + assert row["dimensions"] == 3 + assert row["model"] == "model-b" + + +def test_model_mismatch_filtered(): + """Embeddings with different model should be skipped in search-like filtering.""" + conn = _make_db() + conn.execute( + "INSERT INTO tools (id, owner, name, version, description, tags, visibility, moderation_status) " + "VALUES (1, 'test', 'my-tool', '1.0.0', 'test tool', '[]', 'public', 'approved')" + ) + conn.commit() + + store_embedding(conn, 1, [1.0, 0.0, 0.0], "model-a") + + # Query with model-b should skip this embedding + rows = conn.execute( + "SELECT * FROM tool_embeddings WHERE model = ?", + ["model-b"], + ).fetchall() + assert len(rows) == 0 + + # Query with model-a should find it + rows = conn.execute( + "SELECT * FROM tool_embeddings WHERE model = ?", + ["model-a"], + ).fetchall() + assert len(rows) == 1