brain-template/brain-engine/search.py

#!/usr/bin/env python3
"""
brain-engine/search.py — Recherche sémantique BE-2d
Embed une query → cosine similarity sur brain.db → top-K chunks

Usage :
  python3 brain-engine/search.py "décisions archi SuperOAuth"
  python3 brain-engine/search.py "cold start" --top 10
  python3 brain-engine/search.py "agents helloWorld" --mode file
  python3 brain-engine/search.py "sessions metabolism" --mode json

Modes :
  human  (défaut) → tableau lisible : score | filepath | extrait
  file            → filepaths dédupliqués, triés par score (pour Claude : charger ces fichiers)
  json            → JSON brut : [{score, filepath, title, chunk_text}]

Headless : zéro dépendance display/Wayland.
OLLAMA_URL : variable d'env (défaut localhost:11434).
"""

import os
import sys
import json
import struct
import argparse
import sqlite3
import urllib.request
import urllib.error
from pathlib import Path

BRAIN_ROOT  = Path(__file__).parent.parent
DB_PATH     = BRAIN_ROOT / 'brain.db'
OLLAMA_URL  = os.getenv('OLLAMA_URL', 'http://localhost:11434')
EMBED_MODEL = os.getenv('EMBED_MODEL', 'nomic-embed-text')

# Guardrail — cohérent avec embed.py
_BLOCKED_MODELS = ['mistral', 'qwen', 'llama', 'gemma', 'phi', 'deepseek']
if any(b in EMBED_MODEL.lower() for b in _BLOCKED_MODELS):
    sys.exit(f"❌ EMBED_MODEL='{EMBED_MODEL}' interdit — utiliser nomic-embed-text ou mxbai-embed-large")


# ── Maths ─────────────────────────────────────────────────────────────────────

def cosine_sim(a: list[float], b: list[float]) -> float:
    dot    = sum(x * y for x, y in zip(a, b))
    norm_a = sum(x * x for x in a) ** 0.5
    norm_b = sum(x * x for x in b) ** 0.5
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


def blob_to_vector(blob: bytes) -> list[float]:
    n = len(blob) // 4
    return list(struct.unpack(f'{n}f', blob))


# ── Ollama ─────────────────────────────────────────────────────────────────────

def embed_query(text: str) -> list[float] | None:
    url     = f"{OLLAMA_URL}/api/embeddings"
    payload = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
    req     = urllib.request.Request(url, data=payload,
                                     headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read())
            return data.get('embedding')
    except (urllib.error.URLError, TimeoutError) as e:
        print(f"❌ Ollama indisponible ({OLLAMA_URL}) : {e}", file=sys.stderr)
        return None


# ── SQLite ─────────────────────────────────────────────────────────────────────

def load_vectors(conn: sqlite3.Connection,
                 allowed_scopes: list[str] | None = None,
                 include_historical: bool = False) -> list[dict]:
    """Charge les chunks indexés depuis brain.db, filtrés par scope si fourni.
    Shadow indexing (ADR-037) : scope='historical' exclu par défaut."""
    historical_filter = "" if include_historical else "AND scope != 'historical'"
    if allowed_scopes:
        placeholders = ','.join('?' * len(allowed_scopes))
        rows = conn.execute(f"""
            SELECT chunk_id, filepath, title, chunk_text, vector
            FROM embeddings
            WHERE indexed = 1 AND vector IS NOT NULL
              AND scope IN ({placeholders})
              {historical_filter}
        """, allowed_scopes).fetchall()
    else:
        rows = conn.execute(f"""
            SELECT chunk_id, filepath, title, chunk_text, vector
            FROM embeddings
            WHERE indexed = 1 AND vector IS NOT NULL
              {historical_filter}
        """).fetchall()
    result = []
    for row in rows:
        result.append({
            'chunk_id':   row['chunk_id'],
            'filepath':   row['filepath'],
            'title':      row['title'] or '',
            'chunk_text': row['chunk_text'],
            'vector':     blob_to_vector(row['vector']),
        })
    return result


# ── Search ─────────────────────────────────────────────────────────────────────

def search(query: str, top_k: int = 5, min_score: float = 0.0,
           allowed_scopes: list[str] | None = None) -> list[dict]:
    """Retourne les top-K chunks les plus proches de la query."""
    # 1. Embed la query
    q_vec = embed_query(query)
    if q_vec is None:
        return []

    # 2. Charger les vecteurs (filtrés par scope si fourni)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    chunks = load_vectors(conn, allowed_scopes=allowed_scopes)
    conn.close()

    if not chunks:
        print("⚠️  Index vide — lancer embed.py d'abord", file=sys.stderr)
        return []

    # 3. Cosine similarity
    scored = []
    for chunk in chunks:
        score = cosine_sim(q_vec, chunk['vector'])
        if score >= min_score:
            scored.append({**chunk, 'score': score})

    # 4. Trier, dédupliquer par chunk_id (déjà unique), retourner top-K
    scored.sort(key=lambda x: x['score'], reverse=True)
    top_results = scored[:top_k]

    # 5. Tracking V1 (ADR-037) — hit_count + last_queried_at sur les chunks retournés
    if top_results:
        try:
            track_conn = sqlite3.connect(DB_PATH)
            chunk_ids = [r['chunk_id'] for r in top_results if r.get('chunk_id')]
            if chunk_ids:
                placeholders = ','.join('?' * len(chunk_ids))
                track_conn.execute(f"""
                    UPDATE embeddings
                    SET hit_count = COALESCE(hit_count, 0) + 1,
                        last_queried_at = datetime('now')
                    WHERE chunk_id IN ({placeholders})
                """, chunk_ids)
                track_conn.commit()
            track_conn.close()
        except Exception:
            pass  # tracking is best-effort — never breaks search

    return top_results


# ── Output ─────────────────────────────────────────────────────────────────────

def print_human(results: list[dict], query: str):
    if not results:
        print(f"Aucun résultat pour : {query!r}")
        return
    print(f"\nRecherche : {query!r}  ({len(results)} résultat(s))\n")
    print(f"{'Score':>6}  {'Fichier':<50}  Extrait")
    print("─" * 100)
    for r in results:
        score   = f"{r['score']:.3f}"
        fp      = r['filepath']
        if len(fp) > 50:
            fp = '…' + fp[-49:]
        title   = r['title']
        excerpt = r['chunk_text'].replace('\n', ' ')[:80]
        if title:
            excerpt = f"[{title}] {excerpt}"
        print(f"{score:>6}  {fp:<50}  {excerpt}")
    print()


def print_files(results: list[dict]):
    """Filepaths dédupliqués, ordre par meilleur score."""
    seen = []
    for r in results:
        if r['filepath'] not in seen:
            seen.append(r['filepath'])
    for fp in seen:
        print(fp)


def print_json(results: list[dict]):
    out = [{
        'score':      round(r['score'], 4),
        'filepath':   r['filepath'],
        'title':      r['title'],
        'chunk_text': r['chunk_text'],
    } for r in results]
    print(json.dumps(out, ensure_ascii=False, indent=2))


# ── CLI ────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description='brain-engine search — BE-2d')
    parser.add_argument('query',                          help='Requête en langage naturel')
    parser.add_argument('--top',    type=int, default=5,  help='Nombre de résultats (défaut: 5)')
    parser.add_argument('--mode',   choices=['human', 'file', 'json'], default='human',
                        help='Format de sortie (défaut: human)')
    parser.add_argument('--min-score', type=float, default=0.0,
                        help='Score minimum cosine (0.0–1.0, défaut: 0.0)')
    args = parser.parse_args()

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

    if args.mode == 'file':
        print_files(results)
    elif args.mode == 'json':
        print_json(results)
    else:
        print_human(results, args.query)


if __name__ == '__main__':
    main()