Add semantic search (RAG) for registry tool discovery
Users can now find tools by describing what they need in natural language. Uses Ollama embeddings (nomic-embed-text) on AI-Server for vector similarity search. Available via CLI (registry describe), GUI (AI search row), and API. New files: - registry/embeddings.py: core embedding logic (Ollama API, cosine similarity, pack/unpack vectors, backfill) - tests/test_embeddings.py: 16 unit tests Modified: - registry/db.py: tool_embeddings table (schema + migration) - registry/settings.py: embeddings.* settings (ollama_url, model, enabled, min_score) - registry/app.py: semantic-search endpoint, publish hook, admin backfill/status - registry_client.py: semantic_search() method with error surfacing - gui/pages/registry_page.py: AI search row with SemanticSearchWorker - cli/__init__.py + registry_commands.py: registry describe subcommand Backfill required after deploy: POST /api/v1/admin/embeddings/backfill Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
parent
f24450c19c
commit
5159450feb
|
|
@ -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.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)
|
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
|
# registry my-tools
|
||||||
p_reg_mytools = registry_sub.add_parser("my-tools", help="List your published tools")
|
p_reg_mytools = registry_sub.add_parser("my-tools", help="List your published tools")
|
||||||
p_reg_mytools.set_defaults(func=cmd_registry)
|
p_reg_mytools.set_defaults(func=cmd_registry)
|
||||||
|
|
|
||||||
|
|
@ -37,6 +37,8 @@ def cmd_registry(args):
|
||||||
return _cmd_registry_status(args)
|
return _cmd_registry_status(args)
|
||||||
elif args.registry_cmd == "browse":
|
elif args.registry_cmd == "browse":
|
||||||
return _cmd_registry_browse(args)
|
return _cmd_registry_browse(args)
|
||||||
|
elif args.registry_cmd == "describe":
|
||||||
|
return _cmd_registry_describe(args)
|
||||||
elif args.registry_cmd == "config":
|
elif args.registry_cmd == "config":
|
||||||
return _cmd_registry_config(args)
|
return _cmd_registry_config(args)
|
||||||
else:
|
else:
|
||||||
|
|
@ -53,6 +55,7 @@ def cmd_registry(args):
|
||||||
print(" my-tools List your published tools")
|
print(" my-tools List your published tools")
|
||||||
print(" status <tool> Check moderation status of a tool")
|
print(" status <tool> Check moderation status of a tool")
|
||||||
print(" browse Browse tools (GUI)")
|
print(" browse Browse tools (GUI)")
|
||||||
|
print(" describe <query> Find tools by describing what you need (AI)")
|
||||||
print(" config [action] Manage registry settings (admin)")
|
print(" config [action] Manage registry settings (admin)")
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
@ -894,6 +897,67 @@ def _cmd_registry_status(args):
|
||||||
return 0
|
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 <query>")
|
||||||
|
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):
|
def _cmd_registry_browse(args):
|
||||||
"""Browse tools (GUI)."""
|
"""Browse tools (GUI)."""
|
||||||
from ..gui import run_gui
|
from ..gui import run_gui
|
||||||
|
|
|
||||||
|
|
@ -108,6 +108,24 @@ class InstallWorker(QThread):
|
||||||
self.error.emit(str(e))
|
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):
|
class RegistryPage(QWidget):
|
||||||
"""Registry browser page."""
|
"""Registry browser page."""
|
||||||
|
|
||||||
|
|
@ -115,6 +133,7 @@ class RegistryPage(QWidget):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.main_window = main_window
|
self.main_window = main_window
|
||||||
self._search_worker = None
|
self._search_worker = None
|
||||||
|
self._semantic_worker = None
|
||||||
self._install_worker = None
|
self._install_worker = None
|
||||||
self._versions_worker = None
|
self._versions_worker = None
|
||||||
self._selected_tool = None
|
self._selected_tool = None
|
||||||
|
|
@ -198,6 +217,36 @@ class RegistryPage(QWidget):
|
||||||
|
|
||||||
layout.addWidget(filters_box)
|
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
|
# Active tags display
|
||||||
self.tags_widget = QWidget()
|
self.tags_widget = QWidget()
|
||||||
self.tags_layout = QHBoxLayout(self.tags_widget)
|
self.tags_layout = QHBoxLayout(self.tags_widget)
|
||||||
|
|
@ -638,6 +687,98 @@ class RegistryPage(QWidget):
|
||||||
self.btn_browse.setEnabled(True)
|
self.btn_browse.setEnabled(True)
|
||||||
self.status_label.setText(f"Error: {error}")
|
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):
|
def _prev_page(self):
|
||||||
"""Go to previous page."""
|
"""Go to previous page."""
|
||||||
if self._current_page > 1:
|
if self._current_page > 1:
|
||||||
|
|
|
||||||
|
|
@ -695,6 +695,55 @@ def create_app() -> Flask:
|
||||||
|
|
||||||
return jsonify({"data": data, "meta": paginate(page, per_page, total)})
|
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"])
|
@app.route("/api/v1/tools/search", methods=["GET"])
|
||||||
def search_tools() -> Response:
|
def search_tools() -> Response:
|
||||||
query_text = request.args.get("q", "").strip()
|
query_text = request.args.get("q", "").strip()
|
||||||
|
|
@ -2641,6 +2690,34 @@ def create_app() -> Flask:
|
||||||
)
|
)
|
||||||
g.db.commit()
|
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({
|
response = jsonify({
|
||||||
"data": {
|
"data": {
|
||||||
"owner": owner,
|
"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"])
|
@app.route("/api/v1/admin/reports", methods=["GET"])
|
||||||
@require_moderator
|
@require_moderator
|
||||||
def admin_list_reports() -> Response:
|
def admin_list_reports() -> Response:
|
||||||
|
|
|
||||||
|
|
@ -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_hash ON password_reset_tokens(token_hash);
|
||||||
CREATE INDEX IF NOT EXISTS idx_reset_tokens_publisher ON password_reset_tokens(publisher_id);
|
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:
|
except sqlite3.OperationalError:
|
||||||
pass
|
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
|
# Ensure indexes exist
|
||||||
try:
|
try:
|
||||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_tools_owner ON tools(owner)")
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_tools_owner ON tools(owner)")
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
@ -158,6 +158,36 @@ DEFAULT_SETTINGS: List[Setting] = [
|
||||||
category="moderation",
|
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)
|
# Rate limits (can override defaults)
|
||||||
Setting(
|
Setting(
|
||||||
key="rate_limit.publish.limit",
|
key="rate_limit.publish.limit",
|
||||||
|
|
|
||||||
|
|
@ -1022,6 +1022,47 @@ class RegistryClient:
|
||||||
|
|
||||||
return response.json().get("data", {})
|
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]:
|
def get_popular_tools(self, limit: int = 10) -> List[ToolInfo]:
|
||||||
"""
|
"""
|
||||||
Get most popular tools.
|
Get most popular tools.
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
Loading…
Reference in New Issue