search_server.py

"""Knowledge Engine Search UI — personal search engine over the knowledge base."""

import sqlite3
import os
import json
import re
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS

app = Flask(__name__, static_folder="search_ui")
CORS(app)

DB_PATH = os.path.join(os.path.dirname(__file__), "knowledge.db")


def get_conn():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    return conn


def highlight_snippet(content, query, max_len=300):
    """Extract a relevant snippet from content with query terms highlighted."""
    terms = query.lower().split()
    lines = content.split("\n")

    # Find best matching line
    best_line = ""
    best_score = -1
    for line in lines:
        line_lower = line.lower()
        score = sum(1 for t in terms if t in line_lower)
        if score > best_score and len(line.strip()) > 20:
            best_score = score
            best_line = line.strip()

    if not best_line and lines:
        # Fall back to first non-empty line
        for line in lines:
            if len(line.strip()) > 20:
                best_line = line.strip()
                break

    snippet = best_line[:max_len]
    if len(best_line) > max_len:
        snippet += "..."

    # Bold the matching terms
    for term in terms:
        snippet = re.sub(
            f"({re.escape(term)})",
            r"<mark>\1</mark>",
            snippet,
            flags=re.IGNORECASE,
        )

    return snippet


@app.route("/")
def index():
    return send_from_directory("search_ui", "index.html")


@app.route("/images/<path:filename>")
def serve_image(filename):
    return send_from_directory(os.path.join(os.path.dirname(__file__), "images"), filename)


@app.route("/api/search")
def search():
    q = request.args.get("q", "").strip()
    limit = min(int(request.args.get("limit", 20)), 50)

    if not q:
        return jsonify({"results": [], "query": "", "count": 0})

    conn = get_conn()

    # FTS5 search with BM25 ranking + relevance boost
    try:
        rows = conn.execute(
            """
            SELECT d.id, d.title, d.content, d.url, d.author,
                   d.keywords, d.relevance_score, d.created_at, d.updated_at,
                   s.name as source_name, s.type as source_type,
                   bm25(documents_fts, 5.0, 1.0, 3.0) as bm25_score
            FROM documents_fts fts
            JOIN documents d ON d.id = fts.rowid
            LEFT JOIN sources s ON d.source_id = s.id
            WHERE documents_fts MATCH ?
            ORDER BY (bm25(documents_fts, 5.0, 1.0, 3.0) * (0.5 + d.relevance_score))
            LIMIT ?
        """,
            (q, limit),
        ).fetchall()
    except Exception:
        # If FTS match fails (bad syntax), try simple LIKE
        rows = conn.execute(
            """
            SELECT d.id, d.title, d.content, d.url, d.author,
                   d.keywords, d.relevance_score, d.created_at, d.updated_at,
                   s.name as source_name, s.type as source_type,
                   0 as bm25_score
            FROM documents d
            LEFT JOIN sources s ON d.source_id = s.id
            WHERE d.title LIKE ? OR d.content LIKE ?
            ORDER BY d.relevance_score DESC
            LIMIT ?
        """,
            (f"%{q}%", f"%{q}%", limit),
        ).fetchall()

    # Log the query
    conn.execute(
        "INSERT INTO query_log (query, result_count) VALUES (?, ?)", (q, len(rows))
    )
    conn.commit()

    results = []
    for r in rows:
        # Extract image paths — local images from research + any web URLs
        images = re.findall(r"/images/[^\s)]+", r["content"] or "")
        images += re.findall(
            r"https?://[^\s)]+\.(?:png|jpg|jpeg|gif|webp)", r["content"] or ""
        )
        # Extract source URLs from content
        urls = re.findall(r"https?://[^\s)>\"]+", r["content"] or "")[:5]

        results.append(
            {
                "id": r["id"],
                "title": r["title"] or "Untitled",
                "snippet": highlight_snippet(r["content"] or "", q),
                "url": r["url"],
                "author": r["author"],
                "source": r["source_name"],
                "source_type": r["source_type"],
                "relevance": round(r["relevance_score"] or 0, 2),
                "keywords": json.loads(r["keywords"]) if r["keywords"] else [],
                "created": r["created_at"],
                "updated": r["updated_at"],
                "images": images[:3],
                "source_urls": urls,
                "content_length": len(r["content"] or ""),
            }
        )

    # Get related domains
    domains = []
    if rows:
        doc_ids = [r["id"] for r in rows[:5]]
        placeholders = ",".join("?" * len(doc_ids))
        domain_rows = conn.execute(
            f"""
            SELECT DISTINCT kd.name, kd.description
            FROM document_domains dd
            JOIN knowledge_domains kd ON dd.domain_id = kd.id
            WHERE dd.document_id IN ({placeholders})
        """,
            doc_ids,
        ).fetchall()
        domains = [{"name": d["name"], "description": d["description"]} for d in domain_rows]

    conn.close()

    return jsonify(
        {
            "results": results,
            "query": q,
            "count": len(results),
            "domains": domains,
        }
    )


@app.route("/api/document/<int:doc_id>")
def get_document(doc_id):
    conn = get_conn()
    row = conn.execute(
        """
        SELECT d.*, s.name as source_name, s.type as source_type
        FROM documents d
        LEFT JOIN sources s ON d.source_id = s.id
        WHERE d.id = ?
    """,
        (doc_id,),
    ).fetchone()

    if not row:
        conn.close()
        return jsonify({"error": "Not found"}), 404

    # Get connections
    connections = conn.execute(
        """
        SELECT kg.relationship, kg.weight, d.id, d.title
        FROM knowledge_graph kg
        JOIN documents d ON (d.id = kg.to_doc_id OR d.id = kg.from_doc_id)
        WHERE (kg.from_doc_id = ? OR kg.to_doc_id = ?) AND d.id != ?
    """,
        (doc_id, doc_id, doc_id),
    ).fetchall()

    conn.close()

    return jsonify(
        {
            "id": row["id"],
            "title": row["title"],
            "content": row["content"],
            "url": row["url"],
            "author": row["author"],
            "relevance": round(row["relevance_score"] or 0, 2),
            "keywords": json.loads(row["keywords"]) if row["keywords"] else [],
            "created": row["created_at"],
            "updated": row["updated_at"],
            "source": row["source_name"],
            "connections": [
                {
                    "relationship": c["relationship"],
                    "weight": c["weight"],
                    "doc_id": c["id"],
                    "title": c["title"],
                }
                for c in connections
            ],
        }
    )


def format_db_size():
    """Return human-friendly file size of the knowledge DB."""
    try:
        size_bytes = os.path.getsize(DB_PATH)
    except OSError:
        return "unknown"
    if size_bytes >= 1_073_741_824:
        return f"{size_bytes / 1_073_741_824:.2f} GB"
    elif size_bytes >= 1_048_576:
        return f"{size_bytes / 1_048_576:.1f} MB"
    elif size_bytes >= 1024:
        return f"{size_bytes / 1024:.1f} KB"
    return f"{size_bytes} B"


@app.route("/api/stats")
def stats():
    conn = get_conn()
    s = {
        "db_size": format_db_size(),
        "documents": conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0],
        "sources": conn.execute("SELECT COUNT(*) FROM sources").fetchone()[0],
        "edges": conn.execute("SELECT COUNT(*) FROM knowledge_graph").fetchone()[0],
        "domains": conn.execute("SELECT COUNT(*) FROM knowledge_domains").fetchone()[0],
        "queries": conn.execute("SELECT COUNT(*) FROM query_log").fetchone()[0],
        "avg_relevance": round(
            conn.execute(
                "SELECT COALESCE(AVG(relevance_score), 0) FROM documents"
            ).fetchone()[0],
            2,
        ),
        "recent_docs": [
            {"id": r[0], "title": r[1], "created": r[2]}
            for r in conn.execute(
                "SELECT id, title, created_at FROM documents ORDER BY created_at DESC LIMIT 10"
            ).fetchall()
        ],
        "top_domains": [
            {"name": r[0], "count": r[1], "target": r[2]}
            for r in conn.execute(
                """SELECT kd.name, COUNT(dd.document_id), kd.target_doc_count
                   FROM knowledge_domains kd
                   LEFT JOIN document_domains dd ON kd.id = dd.domain_id
                   GROUP BY kd.id ORDER BY COUNT(dd.document_id) DESC LIMIT 15"""
            ).fetchall()
        ],
    }
    conn.close()
    return jsonify(s)


@app.route("/api/feedback", methods=["POST"])
def feedback():
    data = request.json
    doc_id = data.get("doc_id")
    query = data.get("query", "")
    relevant = data.get("relevant", True)

    if not doc_id:
        return jsonify({"error": "doc_id required"}), 400

    conn = get_conn()

    # Log feedback
    conn.execute(
        "INSERT INTO feedback_log (document_id, query, relevant) VALUES (?, ?, ?)",
        (doc_id, query, 1 if relevant else 0),
    )

    # Update relevance with EMA
    doc = conn.execute(
        "SELECT relevance_score, feedback_count, positive_feedback FROM documents WHERE id = ?",
        (doc_id,),
    ).fetchone()

    if doc:
        count = doc["feedback_count"] + 1
        positive = doc["positive_feedback"] + (1 if relevant else 0)
        alpha = 0.3
        new_score = alpha * (1.0 if relevant else 0.0) + (1 - alpha) * doc[
            "relevance_score"
        ]
        conn.execute(
            """UPDATE documents SET relevance_score = ?, feedback_count = ?,
               positive_feedback = ?, updated_at = datetime('now') WHERE id = ?""",
            (new_score, count, positive, doc_id),
        )

    conn.commit()
    conn.close()
    return jsonify({"ok": True, "new_score": round(new_score, 2) if doc else None})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8585, debug=False)