search.py

#!/usr/bin/env python3
"""Semantic search over SurrealDB memory entries using Qwen3-Embedding-4B."""

import argparse
import gc
import json
import sys
import time
import requests
import torch
from sentence_transformers import SentenceTransformer

SURREAL_URL = "http://localhost:8000"
SURREAL_NS = "kit"
SURREAL_DB = "memory"
SURREAL_USER = "root"
SURREAL_PASS = "root"

MODEL_NAME = "Qwen/Qwen3-Embedding-4B"
EMBEDDING_DIM = 1024


def surreal_query(sql: str) -> list:
    headers = {
        "Accept": "application/json",
        "surreal-ns": SURREAL_NS,
        "surreal-db": SURREAL_DB,
    }
    resp = requests.post(
        f"{SURREAL_URL}/sql",
        headers=headers,
        data=sql,
        auth=(SURREAL_USER, SURREAL_PASS),
    )
    resp.raise_for_status()
    return resp.json()


def load_model() -> SentenceTransformer:
    print(f"Loading {MODEL_NAME}...")
    t0 = time.time()
    model = SentenceTransformer(MODEL_NAME, trust_remote_code=True, model_kwargs={"torch_dtype": torch.float16})
    model.max_seq_length = 8192
    model.truncate_dim = EMBEDDING_DIM
    print(f"Model loaded in {time.time() - t0:.1f}s")
    return model


def unload_model(model):
    del model
    gc.collect()
    torch.cuda.empty_cache()


def search(query: str, top_k: int = 5, model: SentenceTransformer | None = None):
    own_model = model is None
    if own_model:
        model = load_model()

    try:
        # Embed query with instruction prefix for retrieval
        query_emb = model.encode(
            [query],
            prompt_name="query",
            normalize_embeddings=True,
        )[0].tolist()

        vec_str = json.dumps(query_emb)
        sql = f"""
            SELECT id, title, content, tags, kind, created_at,
                   vector::similarity::cosine(embedding, {vec_str}) AS score
            FROM memory
            WHERE embedding IS NOT NONE
            ORDER BY score DESC
            LIMIT {top_k};
        """
        results = surreal_query(sql)
        return results[0].get("result", [])
    finally:
        if own_model:
            unload_model(model)


def main():
    parser = argparse.ArgumentParser(description="Semantic search over memory entries")
    parser.add_argument("query", help="Search query string")
    parser.add_argument("-k", "--top-k", type=int, default=5, help="Number of results")
    args = parser.parse_args()

    results = search(args.query, top_k=args.top_k)

    if not results:
        print("No results found.")
        return

    for i, r in enumerate(results, 1):
        score = r.get("score", 0)
        print(f"\n{'='*60}")
        print(f"#{i} [{score:.4f}] {r.get('title', 'Untitled')}")
        print(f"   Kind: {r.get('kind')} | Tags: {', '.join(r.get('tags', []))}")
        print(f"   Created: {r.get('created_at', 'unknown')}")
        print(f"   {r.get('content', '')[:200]}")


if __name__ == "__main__":
    main()