Initial commit — kb-app RAG server

2026-05-10 23:07:43 +00:00
commit 89331c1fa5
13 changed files with 1055 additions and 0 deletions
@@ -0,0 +1,83 @@
+"""Reranking via Qwen3-Reranker over Ollama."""
+import requests
+import os
+import re
+from .config import CONFIG
+
+RERANK_MODEL = os.getenv("KB_RERANK_MODEL", "dengcao/Qwen3-Reranker-4B:Q5_K_M")
+
+
+def score_relevance(query: str, document: str) -> float:
+    """Get a relevance score in [0, 1] for (query, document) pair."""
+    if len(document) > 4000:
+        document = document[:4000] + "..."
+    
+    # Pre-fill <think></think> as empty so the model jumps straight to answer
+    prompt = (
+        f"<|im_start|>system\n"
+        f"You evaluate whether a document is relevant to a query about quantitative trading. "
+        f"Be strict: the document must specifically address the query. "
+        f"Output only \"yes\" or \"no\".<|im_end|>\n"
+        f"<|im_start|>user\n"
+        f"<Query>: {query}\n"
+        f"<Document>: {document}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+        f"<think>\n\n</think>\n\n"
+    )
+    
+    try:
+        resp = requests.post(
+            f"{CONFIG.ollama_url}/api/generate",
+            json={
+                "model": RERANK_MODEL,
+                "prompt": prompt,
+                "stream": False,
+                "options": {
+                    "num_predict": 5,
+                    "temperature": 0.0,
+                    "top_p": 1.0,
+                },
+                "raw": True,
+            },
+            timeout=60,
+        )
+        resp.raise_for_status()
+        text = resp.json().get("response", "").strip().lower()
+        text = re.sub(r"</?think>", "", text).strip()
+        resp.raise_for_status()
+        text = resp.json().get("response", "").strip()
+        text = re.sub(r"</?think>", "", text).strip()
+        
+        # DEBUG: print every response so we can see what's happening
+        print(f"DEBUG rerank response: {text[:80]!r}", flush=True)
+        
+        text = text.lower()
+        
+        if text.startswith("yes"):
+            return 1.0
+        elif text.startswith("no"):
+            return 0.0
+        else:
+            print(f"Rerank ambiguous response: {text[:100]!r}", flush=True)
+            return 0.5
+    except Exception as e:
+        print(f"Rerank error: {e}", flush=True)
+        return 0.5
+
+
+def rerank(query: str, candidates: list, top_k: int = 10) -> list:
+    if not candidates:
+        return []
+    
+    scored = []
+    for cand in candidates:
+        chunk_meta = cand.get("chunk_metadata") or {}
+        doc_text = chunk_meta.get("raw_chunk") or cand.get("content", "")
+        score = score_relevance(query, doc_text)
+        cand_with_score = dict(cand)
+        cand_with_score["rerank_score"] = score
+        cand_with_score["original_similarity"] = cand.get("similarity")
+        scored.append(cand_with_score)
+    
+    scored.sort(key=lambda x: x["rerank_score"], reverse=True)
+    return scored[:top_k]