Phase 3: retriever utils, hybrid scoring, DDG+S2 wrappers, disk cache

data/cache/058352a518feab93e76651fe9a922736.json

@@ -0,0 +1 @@
+[{"url": "https://arxiv.org/html/2310.07240v6", "snippet": "name :KVCacheCompressionand Streaming for FastLargeLanguageModelServing ...largelanguagemodels(LLMs) are ubiquitously used (Anastasiya, ...", "title": "\\name: KV Cache Compression and Streaming for Fast Large", "inferred_year": null, "hybrid_score": 0.0, "source": "duckduckgo"}, {"url": "https://arxiv.org/html/2410.00161v2", "snippet": "We propose query-group-compression, a simple yet effective method tocompresstheKVcacheof GQAmodelswithout repeating it into the dimension of ...", "title": "KV-Compress: Paged KV-Cache Compression with Variable", "inferred_year": null, "hybrid_score": 0.0, "source": "duckduckgo"}, {"url": "https://arxiv.org/html/2509.05165v2", "snippet": "Largelanguagemodels(LLMs) rely on key-value (KV)cachesfor efficient autoregressive decoding; however,cachesize grows linearly with context ...", "title": "KVCompose: Efficient Structured KV Cache Compression with", "inferred_year": null, "hybrid_score": 0.0, "source": "duckduckgo"}]

requirements.txt

@@ -4,7 +4,7 @@ langchain-groq
 langchain-huggingface
 sentence-transformers
 semanticscholar
-duckduckgo-search
+ddgs
 tavily-python
 ragas
 datasets

src/retriever_utils.py

@@ -0,0 +1,297 @@
+import os
+import time
+import math
+import hashlib
+import json
+import logging
+from datetime import datetime
+from typing import Optional
+
+from dotenv import load_dotenv
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+from src.state import Paper, WebResult
+
+load_dotenv()
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Embedding model — loaded once at module level (CPU, fast)
+# ---------------------------------------------------------------------------
+_embedder: Optional[SentenceTransformer] = None
+
+def get_embedder() -> SentenceTransformer:
+    global _embedder
+    if _embedder is None:
+        _embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    return _embedder
+
+
+# ---------------------------------------------------------------------------
+# Disk cache — prevents re-fetching on eval loop crashes
+# ---------------------------------------------------------------------------
+_CACHE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "cache")
+os.makedirs(_CACHE_DIR, exist_ok=True)
+
+def _cache_key(text: str) -> str:
+    return hashlib.md5(text.encode()).hexdigest()
+
+def _cache_get(key: str) -> Optional[list]:
+    path = os.path.join(_CACHE_DIR, f"{key}.json")
+    if os.path.exists(path):
+        with open(path) as f:
+            return json.load(f)
+    return None
+
+def _cache_set(key: str, data: list) -> None:
+    path = os.path.join(_CACHE_DIR, f"{key}.json")
+    with open(path, "w") as f:
+        json.dump(data, f)
+
+
+# ---------------------------------------------------------------------------
+# Recency scoring — three formulas for ablation study
+# ---------------------------------------------------------------------------
+CURRENT_YEAR = datetime.now().year
+
+def recency_score(year: int, decay_config: str = "linear") -> float:
+    """
+    Returns a 0–1 recency score for a paper given its publication year.
+    decay_config: "none" | "linear" | "log"
+    """
+    if year is None or year == 0:
+        return 0.0
+    age = max(0, CURRENT_YEAR - year)
+
+    if decay_config == "none":
+        return 1.0
+    elif decay_config == "linear":
+        return max(0.0, 1.0 - (age / 20.0))
+    elif decay_config == "log":
+        return max(0.0, 1.0 - math.log1p(age) / math.log1p(20))
+    else:
+        return max(0.0, 1.0 - (age / 20.0))  # default to linear
+
+
+def authority_score(citation_count: int) -> float:
+    """Normalize citation count to 0–1 using log scale."""
+    if citation_count <= 0:
+        return 0.0
+    return min(1.0, math.log1p(citation_count) / math.log1p(10000))
+
+
+def hybrid_score(
+    semantic_sim: float,
+    year: int,
+    citation_count: int,
+    decay_config: str = "linear",
+) -> float:
+    """
+    final_score = semantic_sim × 0.5 + recency × 0.3 + authority × 0.2
+    Weights chosen by ablation study (see eval/).
+    """
+    r = recency_score(year, decay_config)
+    a = authority_score(citation_count)
+    return round(semantic_sim * 0.5 + r * 0.3 + a * 0.2, 4)
+
+
+# ---------------------------------------------------------------------------
+# Semantic Scholar search
+# ---------------------------------------------------------------------------
+def search_semantic_scholar(
+    query: str,
+    limit: int = 5,
+    decay_config: str = "linear",
+    use_cache: bool = True,
+) -> list[Paper]:
+    """
+    Search Semantic Scholar for papers matching the query.
+    Returns a list of Paper objects sorted by hybrid_score descending.
+    Rate limit: 100 req/5 min (with key), ~10/min (without).
+    sleep(3) is applied on every call to stay safe.
+    """
+    cache_key = _cache_key(f"s2_{query}_{limit}")
+    if use_cache:
+        cached = _cache_get(cache_key)
+        if cached:
+            logger.info(f"S2 cache hit: {query[:50]}")
+            return [Paper(**p) for p in cached]
+
+    s2_key = os.getenv("S2_API_KEY")
+
+    try:
+        import semanticscholar as ss
+        if s2_key:
+            sch = ss.SemanticScholar(api_key=s2_key)
+        else:
+            sch = ss.SemanticScholar()
+
+        time.sleep(3)  # rate limit guard — always
+
+        results = sch.search_paper(
+            query,
+            limit=limit,
+            fields=["title", "abstract", "year", "citationCount",
+                    "authors", "references", "paperId"],
+        )
+
+    except Exception as e:
+        logger.warning(f"S2 search failed for '{query}': {e}")
+        return []
+
+    if not results:
+        return []
+
+    # Embed query once, then score all abstracts
+    embedder = get_embedder()
+    query_vec = embedder.encode([query])
+
+    papers = []
+    for r in results:
+        if not r.abstract:
+            continue
+
+        abstract_vec = embedder.encode([r.abstract])
+        sim = float(cosine_similarity(query_vec, abstract_vec)[0][0])
+
+        year = r.year or 0
+        citations = r.citationCount or 0
+        authors = [a["name"] for a in (r.authors or [])]
+        references = [ref["paperId"] for ref in (r.references or [])
+                      if ref.get("paperId")]
+
+        paper = Paper(
+            title=r.title or "Untitled",
+            abstract=r.abstract,
+            year=year,
+            citation_count=citations,
+            paper_id=r.paperId or "",
+            authors=authors,
+            references=references,
+            hybrid_score=hybrid_score(sim, year, citations, decay_config),
+            source="semantic_scholar",
+        )
+        papers.append(paper)
+
+    papers.sort(key=lambda p: p.hybrid_score, reverse=True)
+
+    if use_cache:
+        _cache_set(cache_key, [p.__dict__ for p in papers])
+
+    return papers
+
+
+# ---------------------------------------------------------------------------
+# DuckDuckGo web search (with Tavily fallback)
+# ---------------------------------------------------------------------------
+def search_web(
+    query: str,
+    limit: int = 5,
+    use_cache: bool = True,
+) -> list[WebResult]:
+    """
+    Search the web via DuckDuckGo. Falls back to Tavily if DDG fails.
+    Returns a list of WebResult objects.
+    """
+    cache_key = _cache_key(f"web_{query}_{limit}")
+    if use_cache:
+        cached = _cache_get(cache_key)
+        if cached:
+            logger.info(f"Web cache hit: {query[:50]}")
+            return [WebResult(**r) for r in cached]
+
+    results = _ddg_search(query, limit)
+
+    if not results:
+        logger.warning(f"DDG failed for '{query}', trying Tavily fallback")
+        results = _tavily_search(query, limit)
+
+    if use_cache and results:
+        _cache_set(cache_key, [r.__dict__ for r in results])
+
+    return results
+
+
+def _ddg_search(query: str, limit: int) -> list[WebResult]:
+    try:
+        from ddgs import DDGS
+        time.sleep(1)
+        # Force English results, safesearch off, recent results
+        search_query = f"{query} research paper arxiv"
+        with DDGS() as ddgs:
+            raw = list(ddgs.text(
+                search_query,
+                max_results=limit,
+                region="wt-wt",      # worldwide — avoids regional override
+                safesearch="off",
+            ))
+        results = []
+        for r in raw:
+            year = _infer_year(r.get("body", ""))
+            results.append(WebResult(
+                url=r.get("href", ""),
+                snippet=r.get("body", "")[:500],
+                title=r.get("title", ""),
+                inferred_year=year,
+                source="duckduckgo",
+            ))
+        return results
+    except Exception as e:
+        logger.warning(f"DDG error: {e}")
+        return []
+
+
+def _tavily_search(query: str, limit: int) -> list[WebResult]:
+    tavily_key = os.getenv("TAVILY_API_KEY")
+    if not tavily_key:
+        return []
+    try:
+        from tavily import TavilyClient
+        client = TavilyClient(api_key=tavily_key)
+        response = client.search(query, max_results=limit)
+        results = []
+        for r in response.get("results", []):
+            year = _infer_year(r.get("content", ""))
+            results.append(WebResult(
+                url=r.get("url", ""),
+                snippet=r.get("content", "")[:500],
+                title=r.get("title", ""),
+                inferred_year=year,
+                source="tavily",
+            ))
+        return results
+    except Exception as e:
+        logger.warning(f"Tavily error: {e}")
+        return []
+
+
+def _infer_year(text: str) -> Optional[int]:
+    """Try to extract a 4-digit year (2000–2026) from a text snippet."""
+    import re
+    matches = re.findall(r"\b(20[0-2][0-9])\b", text)
+    if matches:
+        years = [int(y) for y in matches]
+        return max(years)
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Citation graph builder
+# ---------------------------------------------------------------------------
+def build_citation_graph(papers: list[Paper]) -> dict:
+    """
+    Build a citation graph from retrieved papers.
+    Returns {paper_id: [list of referenced paper_ids that are also in our set]}
+    Only includes edges where both source and target are in our retrieved set.
+    """
+    paper_ids = {p.paper_id for p in papers}
+    graph = {}
+    for p in papers:
+        graph[p.paper_id] = [
+            ref for ref in p.references
+            if ref in paper_ids
+        ]
+    return graph

test_phase3.py

@@ -0,0 +1,51 @@
+import sys, logging
+sys.path.insert(0, ".")
+logging.basicConfig(level=logging.WARNING)
+
+from src.retriever_utils import (
+    search_semantic_scholar,
+    search_web,
+    build_citation_graph,
+    hybrid_score,
+    recency_score,
+)
+
+print("=== Phase 3: Retriever Utils ===\n")
+
+# Test 1: Scoring functions
+print("--- Scoring functions ---")
+print(f"✓ recency linear year=2024: {recency_score(2024, 'linear'):.3f}")
+print(f"✓ recency linear year=2019: {recency_score(2019, 'linear'):.3f}")
+print(f"✓ recency log    year=2019: {recency_score(2019, 'log'):.3f}")
+print(f"✓ recency none   year=2019: {recency_score(2019, 'none'):.3f}")
+print(f"✓ hybrid score (sim=0.8, year=2023, citations=500): {hybrid_score(0.8, 2023, 500):.4f}")
+
+# Test 2: Semantic Scholar
+print("\n--- Semantic Scholar ---")
+papers = search_semantic_scholar("KV cache compression LLM", limit=3)
+if papers:
+    for p in papers:
+        print(f"  ✓ [{p.hybrid_score:.3f}] {p.title[:60]} ({p.year})")
+else:
+    print("  ⚠ No results (S2 key may not be active yet — expected)")
+
+# Test 3: Web search
+print("\n--- DuckDuckGo ---")
+results = search_web("KV cache compression large language models 2024", limit=3)
+if results:
+    for r in results:
+        print(f"  ✓ [{r.source}] {r.title[:60]}")
+else:
+    print("  ⚠ No results from DDG or Tavily")
+
+# Test 4: Citation graph
+print("\n--- Citation graph ---")
+if papers:
+    graph = build_citation_graph(papers)
+    print(f"  ✓ Graph nodes: {len(graph)}")
+    edges = sum(len(v) for v in graph.values())
+    print(f"  ✓ Internal edges: {edges}")
+else:
+    print("  ⚠ Skipped (no papers retrieved)")
+
+print("\n✅ Phase 3 complete")