rag.py

"""
LoomChat — LightRAG wired to a self-hosted Ollama Space.

No external API keys. No token limits. Completely free.

LLM:        Ollama Space  →  qwen2.5:3b   (~12 tok/s on HF CPU)
Embeddings: Ollama Space  →  nomic-embed-text  (768-dim, fast on CPU)
Graph:      Neo4j Aura Free
Vectors:    Qdrant Cloud Free
KV/Cache:   HF Persistent Disk  /data  (50 GB, survives restarts)
"""

import json
import os
import re
from pathlib import Path

import httpx
import numpy as np
from lightrag import LightRAG, QueryParam
from lightrag.utils import EmbeddingFunc

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

# Ollama Space URL — set this to your deployed ollama-space HF URL
# e.g. https://your-username-loomchat-ollama.hf.space
OLLAMA_BASE_URL = os.environ["OLLAMA_SPACE_URL"].rstrip("/")
OLLAMA_API_BASE = f"{OLLAMA_BASE_URL}/v1"

# LLM endpoint — defaults to Ollama Space but can be any OpenAI-compatible URL
# e.g. set LLM_BASE_URL=https://idnameraj-copaw-agent.hf.space to use CoPaw
LLM_BASE_URL = os.environ.get("LLM_BASE_URL", OLLAMA_API_BASE).rstrip("/")
# Model name is optional — leave blank to omit from request (e.g. CoPaw doesn't need it)
LLM_MODEL   = os.environ.get("OLLAMA_LLM_MODEL", "")

EMBED_MODEL = os.environ.get("OLLAMA_EMBED_MODEL",  "nomic-embed-text")
EMBED_DIM   = int(os.environ.get("OLLAMA_EMBED_DIM", "768"))

NEO4J_URI      = os.environ["NEO4J_URI"]
NEO4J_USERNAME = os.environ["NEO4J_USERNAME"]
NEO4J_PASSWORD = os.environ["NEO4J_PASSWORD"]

QDRANT_URL        = os.environ["QDRANT_URL"]
QDRANT_API_KEY    = os.environ["QDRANT_API_KEY"]
QDRANT_COLLECTION = os.environ.get("QDRANT_COLLECTION", "loomchat")

LIGHTRAG_WORKING_DIR = os.environ.get("LIGHTRAG_WORKING_DIR", "/data/lightrag")

# Propagate Neo4j env vars — LightRAG's Neo4JStorage reads these directly
os.environ["NEO4J_URI"]      = NEO4J_URI
os.environ["NEO4J_USERNAME"] = NEO4J_USERNAME
os.environ["NEO4J_PASSWORD"] = NEO4J_PASSWORD


# ---------------------------------------------------------------------------
# LLM — raw httpx call (avoids openai_complete double-model kwarg conflict).
# LLM_BASE_URL can point to Ollama Space or CoPaw Space.
# LLM_MODEL is optional — omitted from request if empty (CoPaw doesn't need it).
# ---------------------------------------------------------------------------

async def llm_model_func(
    prompt: str,
    system_prompt: str | None = None,
    history_messages: list | None = None,
    **kwargs,
) -> str:
    # Pop LightRAG internal kwargs we don't forward to the HTTP backend
    kwargs.pop("hashing_kv", None)
    kwargs.pop("cache_type", None)
    kwargs.pop("chunk_id", None)
    kwargs.pop("_priority", None)
    keyword_extraction = kwargs.pop("keyword_extraction", False)

    messages: list[dict] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    if history_messages:
        messages.extend(history_messages)
    messages.append({"role": "user", "content": prompt})

    body: dict = {"messages": messages, "stream": False}
    if LLM_MODEL:
        body["model"] = LLM_MODEL
    body["max_tokens"]  = kwargs.get("max_tokens", 1024)
    body["temperature"] = kwargs.get("temperature", 0.1)

    async with httpx.AsyncClient(timeout=180.0) as client:
        resp = await client.post(
            f"{LLM_BASE_URL}/v1/chat/completions",
            json=body,
        )
        resp.raise_for_status()
        content = resp.json()["choices"][0]["message"]["content"]

    if not keyword_extraction:
        return content

    # LightRAG expects {"high_level_keywords": [...], "low_level_keywords": [...]}
    # Strip markdown code fences the model may add
    cleaned = re.sub(r"```(?:json)?|```", "", content).strip()
    try:
        parsed = json.loads(cleaned)
        if isinstance(parsed, list):
            # Model returned a list — flatten into expected dict shape
            flat = [k for item in parsed for k in (item if isinstance(item, list) else [str(item)])]
            return json.dumps({"high_level_keywords": flat, "low_level_keywords": []})
        return json.dumps(parsed)   # already a dict, re-serialise cleanly
    except json.JSONDecodeError:
        # Fallback: treat whole response as a single keyword
        return json.dumps({"high_level_keywords": [content.strip()], "low_level_keywords": []})


# ---------------------------------------------------------------------------
# Embeddings — LightRAG's built-in openai_embed, pre-bound to Ollama Space.
# ---------------------------------------------------------------------------

async def embed_model_func(texts: list[str]) -> list[list[float]]:
    async with httpx.AsyncClient(timeout=60.0) as client:
        resp = await client.post(
            f"{OLLAMA_API_BASE}/embeddings",
            json={"model": EMBED_MODEL, "input": texts},
            headers={"Authorization": "Bearer ollama"},
        )
        resp.raise_for_status()
        data = resp.json()["data"]
        vectors = [item["embedding"] for item in sorted(data, key=lambda x: x["index"])]
        return np.array(vectors, dtype=np.float32)


# ---------------------------------------------------------------------------
# LightRAG singleton
# ---------------------------------------------------------------------------

_rag_instance: LightRAG | None = None
QUERY_MODES  = {"naive", "local", "global", "hybrid", "mix"}
DEFAULT_MODE = "mix"


def get_rag() -> LightRAG:
    if _rag_instance is None:
        raise RuntimeError("RAG not initialised — await init_rag() at startup")
    return _rag_instance


async def init_rag() -> LightRAG:
    global _rag_instance

    Path(LIGHTRAG_WORKING_DIR).mkdir(parents=True, exist_ok=True)

    _rag_instance = LightRAG(
        working_dir=LIGHTRAG_WORKING_DIR,

        # LLM — qwen2.5:3b via Ollama
        llm_model_func=llm_model_func,
        llm_model_max_async=2,   # CPU is single-threaded; keep concurrency low

        # Embeddings — nomic-embed-text via Ollama (768-dim)
        embedding_func=EmbeddingFunc(
            embedding_dim=EMBED_DIM,
            max_token_size=8192,
            func=embed_model_func,
        ),
        embedding_batch_num=8,
        embedding_func_max_async=2,

        # Graph → Neo4j Aura Free
        graph_storage="Neo4JStorage",

        # Vectors → Qdrant Cloud Free
        vector_storage="QdrantVectorDBStorage",
        vector_db_storage_cls_kwargs={
            "collection_name": QDRANT_COLLECTION,
            "url": QDRANT_URL,
            "api_key": QDRANT_API_KEY,
        },

        # KV + chunk cache → /data persistent disk (JsonKVStorage default)
    )

    await _rag_instance.initialize_storages()
    return _rag_instance


async def query(
    question: str,
    mode: str = DEFAULT_MODE,
    top_k: int = 10,
    only_need_context: bool = False,
) -> str:
    if mode not in QUERY_MODES:
        raise ValueError(f"mode must be one of {QUERY_MODES}")
    return await get_rag().aquery(
        question,
        param=QueryParam(
            mode=mode,
            top_k=top_k,
            only_need_context=only_need_context,
        ),
    )