Upload folder using huggingface_hub

rag_system/cache.py

@@ -1,133 +1,274 @@
 """
-Two-layer cache:
-1. Exact-match hash cache (Redis/in-memory fallback)
-2. Semantic near-duplicate cache using cosine similarity on query embeddings
-
-Semantic caching prevents re-querying the LLM for paraphrased versions of the 
-same question - a major cost & latency win in production
+Try Docs cache (Redis-only):
+1. Exact-match cache in Redis
+2. Semantic cache using Redis Vector Search (RediSearch)
 """
+from __future__ import annotations
 
 import hashlib
-import json 
+import json
 import logging
-import time
+import uuid
 from typing import Optional
 
-from google_crc32c import value
 import numpy as np
 
 from .config import get_settings
-from .embeddings import cosine_similarity
 
 logger = logging.getLogger(__name__)
 settings = get_settings()
 
-# In memory fallback (used when Redis is unavailable)
-
-class InMemoryCache:
-    def __init__(self,ttl: int = 3600, max_size: int = 1000):
-        self._store: dict[str,tuple[str, float]] = {} # Key -> (value, expiry)
-        self.ttl = ttl
-        self.max_size = max_size
-    
-    def get(self,key: str) -> Optional[str]:
-        entry = self._store.get(key)
-        if entry is None:
-            return None
-        value, expiry = entry
-        if time.time() > expiry:
-            del self._store[key]
-            return None
-        return value
-    
-    def set(self,key: str, value: str) -> None:
-        if len(self._store) >= self.max_size:
-            oldest = next(iter(self._store))
-            del self._store[oldest]
-        self._store[key] = (value, time.time() + self.ttl)
-    
-    def ping(self) -> bool:
-        return True
-    
+CACHE_EMBEDDING_MODE = "openai-small"
+EXACT_PREFIX = "rag:try:exact:"
+SEMANTIC_INDEX = "rag:try:semantic"
+SEMANTIC_PREFIX = "rag:try:sem:"
+VECTOR_FIELD = "vector"
+PAYLOAD_FIELD = "payload"
+COLLECTION_FIELD = "collection_key"
+PARAMS_FIELD = "params_key"
+
+
 def _build_redis_client():
     try:
         import redis
-        client = redis.from_url(settings.redis_url, decode_responses=True)
+
+        client = redis.from_url(settings.redis_url, decode_responses=False)
         client.ping()
         logger.info("Redis cache connected")
         return client
-    except Exception as e:
-        logger.warning(f"Redis unavalaible ({e}) - using in-memory cache.")
-        return InMemoryCache(ttl=settings.cache_ttl_seconds)
-
-_cache_client = _build_redis_client()
-
-# Exact match cache
-def _cache_key(query: str, collection: str, mode: str, embedding_mode: str) -> str:
-    payload = f"{query}::{collection}::{mode}::{embedding_mode}"
-    return "rag:exact:" + hashlib.sha256(payload.encode()).hexdigest()[:32]
-
-def get_exact(query: str, collection: str, mode: str, embedding_mode: str) -> Optional[dict]:
-    key = _cache_key(query, collection, mode, embedding_mode)
-    raw = _cache_client.get(key)
-    if raw:
-        logger.debug(f"Exact cache hit: {key[:16]}...")
-        return json.loads(raw)
-    return None
+    except Exception as exc:
+        logger.warning("Redis unavailable (%s); caching disabled", exc)
+        return None
+
+
+_client = _build_redis_client()
+_semantic_ready = False
+_semantic_failed = False
+
+
+def _ensure_semantic_index() -> bool:
+    global _semantic_ready, _semantic_failed
+    if _semantic_ready:
+        return True
+    if _semantic_failed or _client is None:
+        return False
+
+    try:
+        _client.execute_command("FT.INFO", SEMANTIC_INDEX)
+        _semantic_ready = True
+        return True
+    except Exception:
+        pass
+
+    try:
+        dim = str(int(settings.embedding_dimensions_openai))
+        _client.execute_command(
+            "FT.CREATE",
+            SEMANTIC_INDEX,
+            "ON",
+            "HASH",
+            "PREFIX",
+            1,
+            SEMANTIC_PREFIX,
+            "SCHEMA",
+            "query",
+            "TEXT",
+            COLLECTION_FIELD,
+            "TAG",
+            PARAMS_FIELD,
+            "TAG",
+            VECTOR_FIELD,
+            "VECTOR",
+            "HNSW",
+            6,
+            "TYPE",
+            "FLOAT32",
+            "DIM",
+            dim,
+            "DISTANCE_METRIC",
+            "COSINE",
+        )
+        _semantic_ready = True
+        return True
+    except Exception as exc:
+        logger.warning("Failed to create Redis vector index: %s", exc)
+        _semantic_failed = True
+        return False
+
+
+def _exact_key(query: str, collection_key: str, params_key: str) -> str:
+    payload = f"{query}::{collection_key}::{params_key}"
+    return EXACT_PREFIX + hashlib.sha256(payload.encode()).hexdigest()[:32]
+
+
+def _vector_bytes(vec: list[float]) -> bytes:
+    return np.array(vec, dtype=np.float32).tobytes()
+
+
+def _decode(value: bytes | str | None) -> Optional[str]:
+    if value is None:
+        return None
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    return value
+
+
+# Exact cache
+
+def get_exact(query: str, collection_key: str, params_key: str) -> Optional[dict]:
+    if _client is None:
+        return None
+    key = _exact_key(query, collection_key, params_key)
+    raw = _client.get(key)
+    if not raw:
+        return None
+    payload = _decode(raw)
+    if payload is None:
+        return None
+    return json.loads(payload)
 
-def set_exact(query: str, collection: str, mode: str, embedding_mode: str, value: str) -> None:
-    key = _cache_key(query, collection, mode, embedding_mode)
+
+def set_exact(query: str, collection_key: str, params_key: str, value: dict) -> None:
+    if _client is None:
+        return
+    key = _exact_key(query, collection_key, params_key)
     serialized = json.dumps(value)
-    if hasattr(_cache_client,"setex"):
-        _cache_client.setex(key,settings.cache_ttl_seconds,serialized)
-    else:
-        _cache_client.set(key,serialized)
-
-# Semantic Cache
-# stores (embedding, serialized_response) pairs keyed by short hash
-_semantic_index: dict[str, list[tuple[list[float],str,dict]]] = {} # mode -> (vec,key,response)
-
-def get_semantic(query_vec: list[float], embedding_mode: str) -> Optional[dict]:
-    """Return the cache response if cosine similarity > threshold"""
-    pool = _semantic_index.get(embedding_mode, [])
-    best_score = 0.0
-    best_response = None
-    for vec, _,response in pool:
-        score = cosine_similarity(query_vec,vec)
-        if score > best_score:
-            best_score = score
-            best_response = response
-    if best_score >= settings.semantic_cache_threshold:
-        logger.info(f"Semantic Cache hit (score={best_score:.3f})")
-        return best_response
+    _client.set(key, serialized.encode("utf-8"), ex=settings.cache_ttl_seconds)
+
+
+# Semantic cache
+
+def get_semantic(query_vec: list[float], collection_key: str, params_key: str) -> Optional[dict]:
+    if _client is None:
+        return None
+    if not _ensure_semantic_index():
+        return None
+
+    vec = _vector_bytes(query_vec)
+    k = 4
+    query = (
+        f"@{COLLECTION_FIELD}:{{{collection_key}}} "
+        f"@{PARAMS_FIELD}:{{{params_key}}}"
+        f"=>[KNN {k} @{VECTOR_FIELD} $vec AS score]"
+    )
+
+    try:
+        res = _client.execute_command(
+            "FT.SEARCH",
+            SEMANTIC_INDEX,
+            query,
+            "PARAMS",
+            2,
+            "vec",
+            vec,
+            "RETURN",
+            2,
+            PAYLOAD_FIELD,
+            "score",
+            "DIALECT",
+            2,
+        )
+    except Exception as exc:
+        logger.warning("Redis semantic search failed: %s", exc)
+        return None
+
+    if not res or res[0] == 0:
+        return None
+
+    best_similarity = 0.0
+    best_payload = None
+
+    for i in range(1, len(res), 2):
+        fields = res[i + 1]
+        payload = None
+        distance = None
+        for j in range(0, len(fields), 2):
+            name = _decode(fields[j]) or ""
+            value = fields[j + 1]
+            if name == PAYLOAD_FIELD:
+                payload = _decode(value)
+            elif name == "score":
+                distance = float(_decode(value) or 0)
+        if payload is None or distance is None:
+            continue
+        similarity = 1.0 - distance
+        if similarity > best_similarity:
+            best_similarity = similarity
+            best_payload = payload
+
+    if best_payload and best_similarity >= settings.semantic_cache_threshold:
+        logger.info("Semantic cache hit (score=%.3f)", best_similarity)
+        return json.loads(best_payload)
+
     return None
 
-def set_semantic(query_vec: list[float], query: str, response: dict, embedding_mode: str) -> None:
-    h = hashlib.md5(query.encode()).hexdigest()[:8]
-    pool = _semantic_index.setdefault(embedding_mode, [])
-    pool.append((query_vec, h, response))
-    if len(pool) > 5000: # cap memory
-        pool.pop(0)
+
+def set_semantic(
+    query_vec: list[float],
+    query: str,
+    collection_key: str,
+    params_key: str,
+    response: dict,
+) -> None:
+    if _client is None:
+        return
+    if not _ensure_semantic_index():
+        return
+
+    key = f"{SEMANTIC_PREFIX}{uuid.uuid4().hex}"
+    payload = json.dumps(response)
+
+    _client.hset(
+        key,
+        mapping={
+            "query": query,
+            COLLECTION_FIELD: collection_key,
+            PARAMS_FIELD: params_key,
+            VECTOR_FIELD: _vector_bytes(query_vec),
+            PAYLOAD_FIELD: payload,
+        },
+    )
+    _client.expire(key, settings.cache_ttl_seconds)
+
 
 def cache_connected() -> bool:
+    if _client is None:
+        return False
     try:
-        return bool(_cache_client.ping())
+        return bool(_client.ping())
     except Exception:
         return False
 
+
+def _info_to_dict(raw: list) -> dict:
+    info = {}
+    if not raw:
+        return info
+    for i in range(0, len(raw), 2):
+        key = _decode(raw[i]) or ""
+        info[key] = raw[i + 1]
+    return info
+
+
 def get_cache_stats() -> dict:
-    stats = {}
-    if isinstance(_cache_client, InMemoryCache):
-        stats["system"] = "in-memory (python dictionary)"
-        stats["exact_matches_cached"] = len(_cache_client._store)
-    else:
-        stats["system"] = "redis"
-        try:
-            stats["exact_matches_cached"] = _cache_client.dbsize()
-        except:
-            stats["exact_matches_cached"] = "unknown"
-            
-    stats["semantic_matches_cached"] = sum(len(v) for v in _semantic_index.values())
+    if _client is None:
+        return {"system": "disabled", "exact_matches_cached": 0, "semantic_matches_cached": 0}
+
+    stats = {"system": "redis"}
+    try:
+        stats["exact_matches_cached"] = _client.dbsize()
+    except Exception:
+        stats["exact_matches_cached"] = "unknown"
+
+    try:
+        info = _client.execute_command("FT.INFO", SEMANTIC_INDEX)
+        info_map = _info_to_dict(info)
+        num_docs = info_map.get("num_docs", 0)
+        stats["semantic_matches_cached"] = int(num_docs) if num_docs is not None else 0
+    except Exception:
+        stats["semantic_matches_cached"] = "unknown"
+
     return stats
 
-print("[cache] Module ready")
+
+print("[cache] Module ready")

rag_system/query_engine.py

@@ -7,6 +7,7 @@ Core RAG query pipeline:
 5. Generate answer (sync or streaming)
 6. Return answer + sources
 """
+import hashlib
 import logging
 import re
 import time
@@ -23,7 +24,13 @@ from .retriever import retrieve, detect_query_scope, multi_collection_retrieve
 from .vector_store import resolve_embedding_mode_for_collections
 from .memory import resolve_standalone_question,trim_history_to_budget, build_lc_messages
 from .guardrails import check_query, check_context, redact_pii
-from .cache import get_exact,set_exact,get_semantic,set_semantic
+from .cache import (
+    CACHE_EMBEDDING_MODE,
+    get_exact,
+    set_exact,
+    get_semantic,
+    set_semantic,
+)
 from .embeddings import embed_query
 
 logger = logging.getLogger(__name__)
@@ -53,6 +60,21 @@ def _should_preserve_exact_reference(query: str) -> bool:
     """
     return bool(_SECTION_REF_RE.search(query) and _SECTION_HINT_RE.search(query))
 
+
+def _is_try_docs_scope(collections: list[str]) -> bool:
+    prefix = settings.try_docs_prefix
+    return bool(collections) and all(c.startswith(prefix) for c in collections)
+
+
+def _cache_collection_key(collections: list[str]) -> str:
+    raw = "|".join(sorted(collections))
+    return hashlib.sha1(raw.encode()).hexdigest()[:16]
+
+
+def _cache_params_key(mode: str, top_k: Optional[int]) -> str:
+    k = top_k if top_k is not None else settings.top_k_rerank
+    return f"{mode}:{k}"
+
 # Query rewriting
 async def rewrite_query(query: str) -> str:
     """
@@ -170,6 +192,11 @@ async def query(
 
     collections = request.doc_collections or [request.collection_name]
     embedding_mode = resolve_embedding_mode_for_collections(collections, request.embedding_mode)
+    mode_val = request.retrieval_mode.value if hasattr(request.retrieval_mode, "value") else str(request.retrieval_mode)
+    cache_allowed = settings.cache_enabled and _is_try_docs_scope(collections)
+    cache_collection_key = _cache_collection_key(collections)
+    cache_params_key = _cache_params_key(mode_val, request.top_k)
+    cache_query_vec = None
 
     # 1. Input guardrail
     guard = check_query(request.query)
@@ -182,8 +209,8 @@ async def query(
         )
     
     # 2. Exact cache check
-    if settings.cache_enabled:
-        cached = get_exact(request.query, request.collection_name, request.retrieval_mode, embedding_mode)
+    if cache_allowed:
+        cached = get_exact(request.query, cache_collection_key, cache_params_key)
         if cached:
             logger.info(f"Exact cache hit for query: '{request.query}'")
             cached["cached"] = True
@@ -191,9 +218,9 @@ async def query(
             return QueryResponse(**cached)
     
     # 3. Embed query for semantic cache + later retrieval
-    query_vec = await embed_query(request.query, embedding_mode)
-    if settings.cache_enabled:
-        semantic_hit = get_semantic(query_vec, embedding_mode)
+    if cache_allowed:
+        cache_query_vec = await embed_query(request.query, CACHE_EMBEDDING_MODE)
+        semantic_hit = get_semantic(cache_query_vec, cache_collection_key, cache_params_key)
         if semantic_hit:
             logger.info(f"Semantic cache hit for query: '{request.query}'")
             semantic_hit["cached"] = True
@@ -294,10 +321,12 @@ async def query(
     )
 
     # 9. Cache the result
-    if settings.cache_enabled:
+    if cache_allowed:
         result_dict = result.model_dump()
-        set_exact(request.query, request.collection_name, request.retrieval_mode, embedding_mode, result_dict)
-        set_semantic(query_vec, request.query, result_dict, embedding_mode)
+        set_exact(request.query, cache_collection_key, cache_params_key, result_dict)
+        if cache_query_vec is None:
+            cache_query_vec = await embed_query(request.query, CACHE_EMBEDDING_MODE)
+        set_semantic(cache_query_vec, request.query, cache_collection_key, cache_params_key, result_dict)
 
     return result
 
@@ -337,6 +366,9 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
     mode_val = request.retrieval_mode.value if hasattr(request.retrieval_mode, "value") else str(request.retrieval_mode)
     collections = request.doc_collections or [request.collection_name]
     embedding_mode = resolve_embedding_mode_for_collections(collections, request.embedding_mode)
+    cache_allowed = settings.cache_enabled and _is_try_docs_scope(collections)
+    cache_collection_key = _cache_collection_key(collections)
+    cache_params_key = _cache_params_key(mode_val, request.top_k)
 
     yield emit("pipeline_start", "in_progress", {
         "query": request.query,
@@ -360,9 +392,9 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
         yield emit("guardrail_check", "passed", {})
 
         # --- Cache check ---
-        query_vec = None
-        if settings.cache_enabled:
-            cached = get_exact(request.query, request.collection_name, request.retrieval_mode, embedding_mode)
+        cache_query_vec = None
+        if cache_allowed:
+            cached = get_exact(request.query, cache_collection_key, cache_params_key)
             if cached:
                 cached["cached"] = True
                 cached["latency_ms"] = round((time.monotonic() - start) * 1000, 2)
@@ -371,8 +403,8 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
                 yield "data: [DONE]\n\n"
                 return
 
-            query_vec = await embed_query(request.query, embedding_mode)
-            semantic_hit = get_semantic(query_vec, embedding_mode)
+            cache_query_vec = await embed_query(request.query, CACHE_EMBEDDING_MODE)
+            semantic_hit = get_semantic(cache_query_vec, cache_collection_key, cache_params_key)
             if semantic_hit:
                 semantic_hit["cached"] = True
                 semantic_hit["latency_ms"] = round((time.monotonic() - start) * 1000, 2)
@@ -508,7 +540,7 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
         sources_data = [s.model_dump() for s in sources]
 
         # Cache result — failure must not crash the stream
-        if settings.cache_enabled:
+        if cache_allowed:
             try:
                 result_dict = {
                     "answer": full_answer,
@@ -519,10 +551,10 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
                     "latency_ms": latency_ms,
                     "eval_scores": None,
                 }
-                if query_vec is None:
-                    query_vec = await embed_query(request.query, embedding_mode)
-                set_exact(request.query, request.collection_name, request.retrieval_mode, embedding_mode, result_dict)
-                set_semantic(query_vec, request.query, result_dict, embedding_mode)
+                if cache_query_vec is None:
+                    cache_query_vec = await embed_query(request.query, CACHE_EMBEDDING_MODE)
+                set_exact(request.query, cache_collection_key, cache_params_key, result_dict)
+                set_semantic(cache_query_vec, request.query, cache_collection_key, cache_params_key, result_dict)
             except Exception:
                 logger.warning("Cache write failed (non-fatal)", exc_info=True)
 

rag_system/vector_store.py

@@ -3,6 +3,7 @@ import json
 import logging
 import time
 import os
+import warnings
 from pathlib import Path
 from typing import Optional
 
@@ -225,7 +226,13 @@ def similarity_search_with_scores(
     if store is None:
         raise ValueError(f"Collection '{collection}' not loaded. Ingest documents first.")
     _last_used[collection] = time.time()
-    return store.similarity_search_with_relevance_scores(query, k=k)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"Relevance scores must be between 0 and 1, got.*",
+            category=UserWarning,
+        )
+        return store.similarity_search_with_relevance_scores(query, k=k)
 
 def get_store(collection: str = "default") -> Optional[FAISS]:
     return _stores.get(collection)