cAIc/rag.py

"""
JarvisChat - RAG pipeline: Qdrant vector search + system prompt assembly.
"""
import asyncio
import logging
from datetime import datetime, timezone, timedelta

import httpx

from db import get_db, get_setting, list_skills_with_state, format_active_skills_prompt
from memory import search_memories
from config import (
    MAX_SKILL_PROMPT_CHARS,
    RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER,
    RAG_EVICTION_BATCH, RAG_PINNED_SOURCES, RAG_GRACE_HOURS,
    RAG_ACCESS_WEIGHT, RAG_AGE_WEIGHT,
)

log = logging.getLogger("jarvischat")

QDRANT_URL = "http://192.168.50.108:6333"
EMBED_URL = "http://192.168.50.210:11434"
EMBED_MODEL = "mxbai-embed-large"
RAG_COLLECTION = "jarvis_rag"
RAG_SCORE_THRESHOLD = 0.25

eviction_lock = asyncio.Lock()
EVICTION_LOG: list[dict] = []


def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list:
    words = text.split()
    target_words = int(chunk_size / 1.3)
    overlap_words = int(overlap / 1.3)
    if not words:
        return []
    chunks = []
    start = 0
    while start < len(words):
        end = min(start + target_words, len(words))
        chunks.append(" ".join(words[start:end]))
        if end == len(words):
            break
        start += target_words - overlap_words
    return chunks


async def _update_retrieval_count(point_id: str, current_count: int = 0):
    """Fire-and-forget increment of retrieval_count."""
    try:
        async with httpx.AsyncClient() as client:
            payload = {
                "retrieval_count": current_count + 1,
                "last_accessed": datetime.now(timezone.utc).isoformat(),
            }
            resp = await client.put(
                f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/payload",
                json={"points": [point_id], "payload": payload},
                timeout=5.0,
            )
            if resp.status_code not in (200, 201):
                log.warning(f"Failed to increment retrieval count for {point_id}: {resp.status_code}")
    except Exception as e:
        log.warning(f"Error incrementing retrieval count for {point_id}: {e}")


async def query_rag(query: str, limit: int = 3) -> list:
    try:
        async with httpx.AsyncClient() as client:
            embed_resp = await client.post(
                f"{EMBED_URL}/api/embeddings",
                json={"model": EMBED_MODEL, "prompt": query},
                timeout=10.0,
            )
            if embed_resp.status_code != 200:
                return []
            vector = embed_resp.json()["embedding"]
            search_resp = await client.post(
                f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/search",
                json={"vector": vector, "limit": limit, "with_payload": True},
                timeout=10.0,
            )
            if search_resp.status_code != 200:
                return []
            results = search_resp.json().get("result", [])
            for r in results:
                pid = r.get("id")
                if pid:
                    current = r.get("payload", {}).get("retrieval_count", 0) or 0
                    asyncio.ensure_future(_update_retrieval_count(pid, current))
            return results
    except Exception as e:
        log.warning(f"RAG query error: {e}")
        return []


async def get_collection_count() -> int:
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(
                f"{QDRANT_URL}/collections/{RAG_COLLECTION}",
                timeout=10.0,
            )
            if resp.status_code == 200:
                return resp.json().get("result", {}).get("vectors_count", 0)
    except Exception as e:
        log.warning(f"get_collection_count error: {e}")
    return 0


async def get_collection_stats() -> dict:
    count = await get_collection_count()
    high_water_pct = int(RAG_EVICTION_HIGH_WATER * 100)
    low_water_pct = int(RAG_EVICTION_LOW_WATER * 100)
    percent_full = round((count / RAG_MAX_VECTORS) * 100, 1) if RAG_MAX_VECTORS > 0 else 0
    return {
        "vector_count": count,
        "max_vectors": RAG_MAX_VECTORS,
        "high_water_mark": int(RAG_MAX_VECTORS * RAG_EVICTION_HIGH_WATER),
        "low_water_mark": int(RAG_MAX_VECTORS * RAG_EVICTION_LOW_WATER),
        "high_water_pct": high_water_pct,
        "low_water_pct": low_water_pct,
        "percent_full": percent_full,
        "pinned_sources": list(RAG_PINNED_SOURCES),
    }


async def evict_batch(batch_size: int) -> int:
    """Scroll non-pinned, out-of-grace-period vectors, compute scores, delete lowest-scoring."""
    filter_conditions = {
        "must_not": [
            {"match": {"key": "source", "value": src}}
            for src in RAG_PINNED_SOURCES
        ]
    }
    try:
        async with httpx.AsyncClient() as client:
            scroll_resp = await client.post(
                f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll",
                json={
                    "filter": filter_conditions,
                    "limit": min(batch_size * 10, 10000),
                    "with_payload": True,
                    "with_vector": False,
                },
                timeout=30.0,
            )
            if scroll_resp.status_code != 200:
                log.warning(f"Eviction scroll failed: {scroll_resp.status_code}")
                return 0

            points = scroll_resp.json().get("result", {}).get("points", [])
            if not points:
                return 0

            now = datetime.now(timezone.utc)
            scored = []
            for p in points:
                payload = p.get("payload", {})
                date_str = payload.get("ingest_date") or payload.get("upload_date", "")
                if date_str:
                    age_hours = (now - datetime.fromisoformat(date_str)).total_seconds() / 3600
                else:
                    age_hours = 999999

                if age_hours < RAG_GRACE_HOURS:
                    continue

                retrieval_count = payload.get("retrieval_count", 0) or 0
                score = retrieval_count * RAG_ACCESS_WEIGHT + age_hours * RAG_AGE_WEIGHT
                last_accessed = payload.get("last_accessed", date_str)
                scored.append((score, last_accessed, p["id"]))

            if not scored:
                log.warning("No evictable vectors found (all pinned or newborn)")
                return 0

            scored.sort(key=lambda x: (x[0], x[1]))
            to_delete = [p[2] for p in scored[:batch_size]]
            if not to_delete:
                return 0

            delete_resp = await client.post(
                f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/delete",
                json={"points": to_delete},
                timeout=30.0,
            )
            if delete_resp.status_code not in (200, 201):
                log.warning(f"Eviction delete failed: {delete_resp.status_code}")
                return 0

            return len(to_delete)
    except Exception as e:
        log.warning(f"evict_batch error: {e}")
        return 0


async def maybe_evict() -> int:
    if RAG_MAX_VECTORS <= 0:
        return 0
    effective_batch = max(RAG_EVICTION_BATCH, 1)

    async with eviction_lock:
        count = await get_collection_count()
        threshold_high = int(RAG_MAX_VECTORS * RAG_EVICTION_HIGH_WATER)
        threshold_low = int(RAG_MAX_VECTORS * RAG_EVICTION_LOW_WATER)

        if count < threshold_high:
            return 0

        total_evicted = 0
        while count >= threshold_low:
            if total_evicted > 0 and count < threshold_low:
                break
            deleted = await evict_batch(effective_batch)
            if deleted == 0:
                break
            total_evicted += deleted
            count -= deleted
            if count < threshold_high and total_evicted > 0:
                break
            if count < threshold_low:
                break

        if total_evicted > 0:
            entry = {
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "count": total_evicted,
                "remaining": count,
            }
            EVICTION_LOG.append(entry)
            if len(EVICTION_LOG) > 1000:
                EVICTION_LOG.pop(0)
            log.info(f"Evicted {total_evicted} vectors ({count} remaining)")

        return total_evicted


async def get_rag_operational_stats() -> dict:
    stats = await get_collection_stats()
    now = datetime.now(timezone.utc)
    cutoff_1m = now - timedelta(minutes=1)
    cutoff_5m = now - timedelta(minutes=5)
    cutoff_30m = now - timedelta(minutes=30)

    eviction_1m = sum(
        e["count"] for e in EVICTION_LOG
        if datetime.fromisoformat(e["timestamp"]) > cutoff_1m
    )
    eviction_5m = sum(
        e["count"] for e in EVICTION_LOG
        if datetime.fromisoformat(e["timestamp"]) > cutoff_5m
    )
    eviction_30m = sum(
        e["count"] for e in EVICTION_LOG
        if datetime.fromisoformat(e["timestamp"]) > cutoff_30m
    )

    stats.update({
        "grace_hours": RAG_GRACE_HOURS,
        "eviction_counts_last_1m": eviction_1m,
        "eviction_counts_last_5m": eviction_5m,
        "eviction_counts_last_30m": eviction_30m,
    })
    return stats


async def build_system_prompt(db, extra_prompt: str = "", user_message: str = "") -> str:
    parts = []
    settings = {row["key"]: row["value"] for row in db.execute("SELECT key, value FROM settings").fetchall()}

    if settings.get("profile_enabled", "true") == "true":
        profile = db.execute("SELECT content FROM profile WHERE id = 1").fetchone()
        if profile and profile["content"].strip():
            parts.append(profile["content"].strip())

    if settings.get("memory_enabled", "true") == "true" and user_message:
        memories = search_memories(user_message, limit=5)
        if memories:
            memory_lines = [f"- {m['fact']}" for m in memories]
            parts.append("## Relevant Context from Memory\n" + "\n".join(memory_lines))
            log.debug(f"Injected {len(memories)} memories into context")

    if user_message:
        try:
            rag_results = await query_rag(user_message)
            if rag_results:
                rag_lines = [r["payload"]["text"] for r in rag_results if r["score"] > RAG_SCORE_THRESHOLD]
                if rag_lines:
                    parts.append("## Retrieved Context\n" + "\n\n---\n\n".join(rag_lines))
                    log.info(f"RAG injected {len(rag_lines)} chunks into context")
        except Exception as e:
            log.warning(f"RAG injection error: {e}")

    if settings.get("skills_enabled", "true") == "true":
        active_skills = [s for s in list_skills_with_state(db) if s["enabled"]]
        if active_skills:
            parts.append(format_active_skills_prompt(active_skills))

    if extra_prompt and extra_prompt.strip():
        parts.append(extra_prompt.strip())

    return "\n\n---\n\n".join(parts) if parts else ""