Upload folder using huggingface_hub

rag_system/cache.py

@@ -101,6 +101,10 @@ def _exact_key(query: str, collection_key: str, params_key: str) -> str:
     return EXACT_PREFIX + hashlib.sha256(payload.encode()).hexdigest()[:32]
 
 
+def _tag_hash(value: str) -> str:
+    return hashlib.sha1(value.encode()).hexdigest()[:16]
+
+
 def _vector_bytes(vec: list[float]) -> bytes:
     return np.array(vec, dtype=np.float32).tobytes()
 
@@ -146,9 +150,10 @@ def get_semantic(query_vec: list[float], collection_key: str, params_key: str) -
 
     vec = _vector_bytes(query_vec)
     k = 4
+    params_tag = _tag_hash(params_key)
     query = (
         f"@{COLLECTION_FIELD}:{{{collection_key}}} "
-        f"@{PARAMS_FIELD}:{{{params_key}}}"
+        f"@{PARAMS_FIELD}:{{{params_tag}}} "
         f"=>[KNN {k} @{VECTOR_FIELD} $vec AS score]"
     )
 
@@ -223,7 +228,7 @@ def set_semantic(
         mapping={
             "query": query,
             COLLECTION_FIELD: collection_key,
-            PARAMS_FIELD: params_key,
+            PARAMS_FIELD: _tag_hash(params_key),
             VECTOR_FIELD: _vector_bytes(query_vec),
             PAYLOAD_FIELD: payload,
         },

rag_system/models.py

@@ -47,6 +47,10 @@ class QueryRequest(BaseModel):
     retrieval_mode: RetrievalMode = RetrievalMode.HYBRID
     embedding_mode: Optional[EmbeddingMode] = None
     top_k: Optional[int] = None
+    top_k_retrieval: Optional[int] = None
+    mmr_lambda: Optional[float] = None
+    bm25_weight: Optional[float] = None
+    vector_weight: Optional[float] = None
     doc_collections: Optional[list[str]] = None  # per-doc sub-collections; None = legacy single-collection mode
     history: list[ChatMessage] = Field(default_factory=list)
     stream: bool = False
@@ -56,6 +60,24 @@ class QueryRequest(BaseModel):
     def sanitize_query(cls,v):
         return v.strip()
 
+    @field_validator("top_k", "top_k_retrieval")
+    @classmethod
+    def validate_top_k(cls, v):
+        if v is None:
+            return v
+        if v < 1:
+            raise ValueError("top_k values must be >= 1")
+        return v
+
+    @field_validator("mmr_lambda", "bm25_weight", "vector_weight")
+    @classmethod
+    def validate_weights(cls, v):
+        if v is None:
+            return v
+        if v < 0 or v > 1:
+            raise ValueError("weights must be between 0 and 1")
+        return v
+
 class SourceDocument(BaseModel):
     doc_id: str
     content: str

rag_system/query_engine.py

@@ -71,9 +71,26 @@ def _cache_collection_key(collections: list[str]) -> str:
     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}"
+def _fmt_param(value: Optional[float]) -> str:
+    if value is None:
+        return "-"
+    return f"{value:.3f}"
+
+
+def _cache_params_key_v2(
+    mode: str,
+    top_k: Optional[int],
+    top_k_retrieval: Optional[int],
+    mmr_lambda: Optional[float],
+    bm25_weight: Optional[float],
+    vector_weight: Optional[float],
+) -> str:
+    k_final = top_k if top_k is not None else settings.top_k_rerank
+    k_retrieve = top_k_retrieval if top_k_retrieval is not None else settings.top_k_retrieval
+    return (
+        f"{mode}:{k_final}:{k_retrieve}:"
+        f"{_fmt_param(mmr_lambda)}:{_fmt_param(bm25_weight)}:{_fmt_param(vector_weight)}"
+    )
 
 # Query rewriting
 async def rewrite_query(query: str) -> str:
@@ -195,7 +212,14 @@ async def query(
     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_params_key = _cache_params_key_v2(
+        mode_val,
+        request.top_k,
+        request.top_k_retrieval,
+        request.mmr_lambda,
+        request.bm25_weight,
+        request.vector_weight,
+    )
     cache_query_vec = None
 
     # 1. Input guardrail
@@ -250,6 +274,10 @@ async def query(
             collections=scoped,
             mode=request.retrieval_mode.value if hasattr(request.retrieval_mode, "value") else str(request.retrieval_mode),
             k_per_collection=k_per,
+            top_k_retrieval=request.top_k_retrieval,
+            mmr_lambda=request.mmr_lambda,
+            bm25_weight=request.bm25_weight,
+            vector_weight=request.vector_weight,
             use_reranker=True,
             expand_context=True,
         )
@@ -260,6 +288,10 @@ async def query(
             collection=collections[0],
             mode=request.retrieval_mode,
             top_k=request.top_k,
+            top_k_retrieval=request.top_k_retrieval,
+            mmr_lambda=request.mmr_lambda,
+            bm25_weight=request.bm25_weight,
+            vector_weight=request.vector_weight,
             use_reranker=True,
             expand_context=True,
         )
@@ -368,7 +400,14 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
     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)
+    cache_params_key = _cache_params_key_v2(
+        mode_val,
+        request.top_k,
+        request.top_k_retrieval,
+        request.mmr_lambda,
+        request.bm25_weight,
+        request.vector_weight,
+    )
 
     yield emit("pipeline_start", "in_progress", {
         "query": request.query,
@@ -451,7 +490,8 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
         # --- Retrieval ---
         yield emit("retrieval_start", "in_progress", {
             "mode": mode_val,
-            "top_k": request.top_k or settings.top_k_retrieval,
+            "top_k": request.top_k or settings.top_k_rerank,
+            "top_k_retrieval": request.top_k_retrieval or settings.top_k_retrieval,
             "collections": len(scoped),
         })
 
@@ -462,6 +502,10 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
                 collections=scoped,
                 mode=mode_val,
                 k_per_collection=k_per,
+                top_k_retrieval=request.top_k_retrieval,
+                mmr_lambda=request.mmr_lambda,
+                bm25_weight=request.bm25_weight,
+                vector_weight=request.vector_weight,
                 use_reranker=True,
                 expand_context=True,
             )
@@ -471,6 +515,10 @@ async def pipeline_stream_query(request: QueryRequest) -> AsyncIterator[str]:
                 collection=scoped[0],
                 mode=request.retrieval_mode,
                 top_k=request.top_k,
+                top_k_retrieval=request.top_k_retrieval,
+                mmr_lambda=request.mmr_lambda,
+                bm25_weight=request.bm25_weight,
+                vector_weight=request.vector_weight,
                 use_reranker=True,
                 expand_context=True,
             )

rag_system/retriever.py

@@ -49,13 +49,30 @@ def bm25_retrieve(query: str, collection: str,k: int) -> list[tuple[Document,flo
 def _rrf_score(rank: int, k: int = 60) -> float:
     return 1.0 / (k + rank + 1) #here 1 is added to handle rank 1 which here comes as 0
 
+
+def _resolve_hybrid_weights(
+    bm25_weight: Optional[float],
+    vector_weight: Optional[float],
+) -> tuple[float, float]:
+    bw = settings.bm25_weight if bm25_weight is None else bm25_weight
+    vw = settings.vector_weight if vector_weight is None else vector_weight
+    total = bw + vw
+    if total <= 0:
+        bw = settings.bm25_weight
+        vw = settings.vector_weight
+        total = bw + vw
+    return bw / total, vw / total
+
 #Hybrid Retrieval
 def hybrid_retrieve(
     query: str,
     collection: str,
-    k: int
+    k: int,
+    bm25_weight: Optional[float] = None,
+    vector_weight: Optional[float] = None,
 ) -> list[tuple[Document,float]]:
     """Reciprocal Rank fusion of BM25 and Vector results"""
+    bw, vw = _resolve_hybrid_weights(bm25_weight, vector_weight)
     pool_size = k*3 #casting a wide net before fusing
     vec_results = similarity_search_with_scores(query,collection,k=pool_size)
     bm25_results = bm25_retrieve(query,collection,k=pool_size)
@@ -65,12 +82,12 @@ def hybrid_retrieve(
 
     for rank, (doc, _) in enumerate(vec_results):
         did = doc.metadata.get("doc_id",id(doc))
-        rrf_scores[did] = rrf_scores.get(did,0) + settings.vector_weight * _rrf_score(rank) #check the existing score first and then add the fresh score
+        rrf_scores[did] = rrf_scores.get(did,0) + vw * _rrf_score(rank) #check the existing score first and then add the fresh score
         doc_map[did] = doc
     
     for rank, (doc, _) in enumerate(bm25_results):
         did = doc.metadata.get("doc_id",id(doc))
-        rrf_scores[did] = rrf_scores.get(did,0) + settings.bm25_weight * _rrf_score(rank) 
+        rrf_scores[did] = rrf_scores.get(did,0) + bw * _rrf_score(rank) 
         doc_map[did] = doc
     
     sorted_ids = sorted(rrf_scores, key=lambda x: rrf_scores[x], reverse=True)[:k]
@@ -81,14 +98,15 @@ async def mmr_retrieve(
     query: str,
     collection: str,
     k: int,
-    lambda_mult: float = None,
+    lambda_mult: Optional[float] = None,
+    fetch_k: Optional[int] = None,
 ) -> list[tuple[Document, float]]:
     """
     Maximal Marginal Relevance: balance relevance vs diversity.
     """
     # 1. Setup parameters
     lam = lambda_mult or settings.mmr_lambda
-    fetch_k = settings.top_k_retrieval
+    fetch_k = fetch_k or settings.top_k_retrieval
     store = get_store(collection)
 
     # 2. Get original scores to map them back later
@@ -195,10 +213,14 @@ async def retrieve(
     collection: str = "default",
     mode: str = "hybrid",
     top_k: Optional[int] = None,
+    top_k_retrieval: Optional[int] = None,
+    mmr_lambda: Optional[float] = None,
+    bm25_weight: Optional[float] = None,
+    vector_weight: Optional[float] = None,
     use_reranker: bool = True,
     expand_context: bool = True,
 ) -> list[tuple[Document,float]]:
-    k_retrieve = settings.top_k_retrieval
+    k_retrieve = top_k_retrieval or settings.top_k_retrieval
     k_final = top_k or settings.top_k_rerank
 
     if mode == "vector":
@@ -206,10 +228,16 @@ async def retrieve(
     elif mode == "bm25":
         results = bm25_retrieve(query,collection,k=k_retrieve)
     elif mode == "mmr":
-        results = await mmr_retrieve(query,collection,k=k_final)
+        results = await mmr_retrieve(query,collection,k=k_final,lambda_mult=mmr_lambda,fetch_k=k_retrieve)
         return results # MMR alrady handles diversity , skip reranker
     else: #go with hybrid
-        results = hybrid_retrieve(query,collection,k=k_retrieve)
+        results = hybrid_retrieve(
+            query,
+            collection,
+            k=k_retrieve,
+            bm25_weight=bm25_weight,
+            vector_weight=vector_weight,
+        )
 
     #Rerank
     if use_reranker:
@@ -259,6 +287,10 @@ async def multi_collection_retrieve(
     collections: list[str],
     mode: str = "hybrid",
     k_per_collection: int = 5,
+    top_k_retrieval: Optional[int] = None,
+    mmr_lambda: Optional[float] = None,
+    bm25_weight: Optional[float] = None,
+    vector_weight: Optional[float] = None,
     use_reranker: bool = True,
     expand_context: bool = True,
 ) -> list[tuple[Document, float]]:
@@ -275,6 +307,10 @@ async def multi_collection_retrieve(
                 collection=coll,
                 mode=mode,
                 top_k=k_per_collection,
+                top_k_retrieval=top_k_retrieval,
+                mmr_lambda=mmr_lambda,
+                bm25_weight=bm25_weight,
+                vector_weight=vector_weight,
                 use_reranker=False,   # defer reranking until after merge
                 expand_context=expand_context,
             )