Implement dense-first retrieval, citation enforcement, and Redis caching

api/adapters/pg_adapter.py

@@ -26,6 +26,7 @@ def get_pool() -> asyncpg.Pool:
     "hybrid_search_trusted_v5",
     "hybrid_search_with_graph_v5",
     "hybrid_search_trusted_with_graph_v5",
+    "dense_search_v5",
     "queue_document",
     "dequeue_document",
     "complete_document",

api/config.py

@@ -37,6 +37,15 @@ class Settings(BaseSettings):
     stream_fallback_budget_seconds: int = 8
     trusted_proxies: str = ""
 
+    # Retrieval strategy: dense-first with hybrid fallback.
+    dense_hit_min_results: int = 5
+    dense_hit_min_similarity: float = 0.5
+    rerank_skip_similarity: float = 0.8
+
+    # Query-result cache (Redis) and per-key daily token quotas.
+    query_cache_ttl_seconds: int = 3600
+    quota_enabled: bool = True
+
     redis_url: str = "redis://localhost:6379"
     database_url: str = "postgresql://depthapi:depthapi@localhost:5432/depthapi"
     cache_ttl: int = 86400

api/services/cache.py

@@ -0,0 +1,146 @@
+"""Query-result cache and per-key token quotas backed by Redis.
+
+Both features fail open: if Redis is unreachable, queries run uncached and
+quotas are not enforced (with a loud warning). Availability first; enforcement
+is best-effort until Redis becomes a hard deployment dependency.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from fastapi import HTTPException
+
+from api.config import get_settings
+
+log = logging.getLogger(__name__)
+
+CACHE_VERSION = 1
+
+_client = None
+
+
+def get_client():
+    """Lazy Redis singleton; None when Redis is unreachable. Monkeypatchable."""
+    global _client
+    if _client is not None:
+        return _client
+    try:
+        import redis
+
+        client = redis.Redis.from_url(
+            get_settings().redis_url,
+            socket_connect_timeout=2,
+            socket_timeout=2,
+            decode_responses=True,
+        )
+        client.ping()
+        _client = client
+        return _client
+    except Exception as exc:
+        log.warning("Redis unavailable, running without cache/quotas: %s", exc)
+        return None
+
+
+def reset_client() -> None:
+    """Drop the cached client (tests)."""
+    global _client
+    _client = None
+
+
+def cache_key(tenant_id: str, query: str, collection_id: str | None, depth: int,
+              temperature: float, rerank: bool, use_trusted: bool,
+              graph_hops: int | None, llm_model: str) -> str:
+    """Generate a deterministic Redis cache key scoped by tenant/API key ID."""
+    material = "|".join([
+        f"v{CACHE_VERSION}", query.strip(),
+        collection_id or "", str(depth), f"{temperature:.2f}",
+        str(rerank), str(use_trusted), str(graph_hops), llm_model,
+    ])
+    digest = hashlib.sha256(material.encode("utf-8"), usedforsecurity=False).hexdigest()
+    return f"depthapi:q:{tenant_id}:{digest}"
+
+
+def get_cached(key: str) -> dict[str, Any] | None:
+    client = get_client()
+    if client is None:
+        return None
+    try:
+        raw = client.get(key)
+        if not raw:
+            return None
+        payload = json.loads(raw)
+        if not isinstance(payload, dict) or "answer" not in payload or "contexts" not in payload:
+            return None
+        return payload
+    except Exception as exc:
+        log.warning("Cache read failed, treating as miss: %s", exc)
+        return None
+
+
+def put_cached(key: str, payload: dict[str, Any]) -> None:
+    client = get_client()
+    if client is None:
+        return
+    try:
+        client.set(key, json.dumps(payload, default=str), ex=get_settings().query_cache_ttl_seconds)
+    except Exception as exc:
+        log.warning("Cache write failed: %s", exc)
+
+
+def quota_limit(is_pro: bool) -> int:
+    settings = get_settings()
+    return settings.pro_daily_token_quota if is_pro else settings.daily_token_quota_per_user
+
+
+def count_tokens(text: str, model: str) -> int:
+    try:
+        import tiktoken
+
+        try:
+            enc = tiktoken.encoding_for_model(model)
+        except Exception:
+            enc = tiktoken.get_encoding("cl100k_base")
+        return len(enc.encode(text or ""))
+    except Exception as exc:
+        log.warning("Token counting failed, assuming 0: %s", exc)
+        return 0
+
+
+def _quota_redis_key(tenant_id: str) -> str:
+    day = datetime.now(timezone.utc).strftime("%Y%m%d")
+    return f"depthapi:quota:{tenant_id}:{day}"
+
+
+def check_quota(tenant_id: str, is_pro: bool, estimated_tokens: int) -> None:
+    """Raise 429 when the key would exceed its daily token budget. Fail-open."""
+    limit = quota_limit(is_pro)
+    if limit <= 0:
+        return
+    client = get_client()
+    if client is None:
+        return
+    try:
+        used = int(client.get(_quota_redis_key(tenant_id)) or 0)
+    except Exception as exc:
+        log.warning("Quota read failed, allowing request: %s", exc)
+        return
+    if used + estimated_tokens > limit:
+        raise HTTPException(429, "Daily token quota exceeded")
+
+
+def consume_quota(tenant_id: str, tokens: int) -> None:
+    if tokens <= 0:
+        return
+    client = get_client()
+    if client is None:
+        return
+    try:
+        key = _quota_redis_key(tenant_id)
+        client.incrby(key, tokens)
+        client.expire(key, 172800)
+    except Exception as exc:
+        log.warning("Quota accounting failed: %s", exc)

api/services/inference/inference.py

@@ -1,35 +1,100 @@
 """Mode-free response generation for retrieved contexts."""
+import re
 from collections.abc import AsyncIterator
 from typing import Any
 
 from api.config import get_settings
 
+CITATION_PATTERN = re.compile(r"\[\d+\]")
+
+_ABSTENTION_MARKERS = (
+    "could not find sufficient",
+    "no matching knowledge",
+    "insufficient",
+)
+
 
 def _fallback_response(contexts: list[dict[str, Any]]) -> str:
     if not contexts:
         return "No matching knowledge was found."
     excerpts = [str(context.get("content", "")).strip() for context in contexts]
     return "\n\n".join(excerpt for excerpt in excerpts if excerpt) or "No matching knowledge was found."
 
-async def generate_response(query: str, contexts: list[dict[str, Any]], temperature: float = 0.7) -> str:
+
+def has_citation_markers(answer: str) -> bool:
+    """True when the answer cites at least one numbered source like [1]."""
+    return CITATION_PATTERN.search(answer or "") is not None
+
+
+def looks_like_abstention(answer: str) -> bool:
+    """True for honest no-match answers, which carry no citations by design."""
+    lowered = (answer or "").lower()
+    return any(marker in lowered for marker in _ABSTENTION_MARKERS)
+
+
+def _numbered_sources(contexts: list[dict[str, Any]]) -> str:
+    chunks = []
+    for idx, item in enumerate(contexts, 1):
+        chunks.append(f"[{idx}] {str(item.get('content', ''))[:6000]}")
+    return "\n\n".join(chunks)
+
+
+def _citation_system_prompt() -> str:
+    return (
+        "Answer using only the supplied knowledge, which is numbered [1], [2], ... . "
+        "Cite every factual claim inline with its source number, e.g. [1]. "
+        "If the knowledge is insufficient, say so plainly without citations."
+    )
+
+
+async def _complete_once(
+    client: Any, model: str, temperature: float, query: str, source_text: str, nudge: str = ""
+) -> str | None:
+    user_content = f"Question: {query}\n\nKnowledge:\n{source_text}"
+    if nudge:
+        user_content += f"\n\n{nudge}"
+    response = await client.chat.completions.create(
+        model=model,
+        temperature=temperature,
+        messages=[
+            {"role": "system", "content": _citation_system_prompt()},
+            {"role": "user", "content": user_content},
+        ],
+    )
+    answer = response.choices[0].message.content if response.choices else None
+    return answer.strip() if answer else None
+
+
+async def generate_response(
+    query: str, contexts: list[dict[str, Any]], temperature: float = 0.7, enforce_citations: bool = True
+) -> str:
     settings = get_settings()
     api_key = settings.openai_api_key.get_secret_value()
     if not api_key or not contexts:
         return _fallback_response(contexts)
     try:
         from openai import AsyncOpenAI
         client = AsyncOpenAI(api_key=api_key, timeout=settings.llm_timeout_seconds)
-        source_text = "\n\n".join(str(item.get("content", ""))[:6000] for item in contexts)
-        response = await client.chat.completions.create(
-            model=settings.llm_model,
-            temperature=temperature,
-            messages=[
-                {"role": "system", "content": "Answer using only the supplied knowledge. If it is insufficient, say so."},
-                {"role": "user", "content": f"Question: {query}\n\nKnowledge:\n{source_text}"},
-            ],
-        )
-        answer = response.choices[0].message.content if response.choices else None
-        return answer.strip() if answer else _fallback_response(contexts)
+        source_text = _numbered_sources(contexts)
+        answer = await _complete_once(client, settings.llm_model, temperature, query, source_text)
+        if not answer:
+            return _fallback_response(contexts)
+        if (
+            enforce_citations
+            and not has_citation_markers(answer)
+            and not looks_like_abstention(answer)
+        ):
+            retried = await _complete_once(
+                client,
+                settings.llm_model,
+                temperature,
+                query,
+                source_text,
+                nudge="Your previous answer contained no citations like [1]. Answer again, citing every factual claim with its source number.",
+            )
+            if retried and (has_citation_markers(retried) or looks_like_abstention(retried)):
+                return retried
+        return answer
     except Exception:
         return _fallback_response(contexts)
 
@@ -47,13 +112,13 @@ async def generate_stream_response(query: str, contexts: list[dict[str, Any]], t
     try:
         from openai import AsyncOpenAI
         client = AsyncOpenAI(api_key=api_key, timeout=settings.llm_timeout_seconds)
-        source_text = "\n\n".join(str(item.get("content", ""))[:6000] for item in contexts)
+        source_text = _numbered_sources(contexts)
         stream = await client.chat.completions.create(
             model=settings.llm_model,
             temperature=temperature,
             stream=True,
             messages=[
-                {"role": "system", "content": "Answer using only the supplied knowledge. If it is insufficient, say so."},
+                {"role": "system", "content": _citation_system_prompt()},
                 {"role": "user", "content": f"Question: {query}\n\nKnowledge:\n{source_text}"},
             ],
         )

crates/depth_engine/src/retrieval/crag.rs

@@ -54,6 +54,11 @@ pub fn evaluate_confidence(scores: &[f64], is_reranked: bool) -> ConfidenceEvalu
 }
 
 /// Helper to evaluate confidence directly from a vector of context dictionaries.
+///
+/// Score scales differ by retrieval source: dense cosine similarity
+/// (~0.3-1.0, tagged with `"match_source": "dense"`) versus small RRF-fused
+/// scores for hybrid/graph retrieval. Each scale uses its own thresholds so
+/// a strong dense hit is not misread as a weak fused one (or vice versa).
 pub fn evaluate_contexts_confidence(contexts: &[serde_json::Value]) -> ConfidenceEvaluation {
     if contexts.is_empty() {
         return ConfidenceEvaluation {
@@ -72,10 +77,31 @@ pub fn evaluate_contexts_confidence(contexts: &[serde_json::Value]) -> Confidenc
             .collect();
         evaluate_confidence(&scores, true)
     } else {
+        let is_dense = contexts.iter().any(|c| {
+            c.get("match_source").and_then(|v| v.as_str()) == Some("dense")
+        });
         let scores: Vec<f64> = contexts
             .iter()
             .filter_map(|c| c.get("score").and_then(|v| v.as_f64()))
             .collect();
+        if is_dense {
+            if scores.is_empty() {
+                return evaluate_confidence(&[], false);
+            }
+            let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
+            let tier = if max_score < 0.55 {
+                "low"
+            } else if max_score < 0.7 {
+                "medium"
+            } else {
+                "high"
+            };
+            return ConfidenceEvaluation {
+                confidence: tier.to_string(),
+                is_insufficient: false,
+                max_score: Some(max_score),
+            };
+        }
         evaluate_confidence(&scores, false)
     }
 }
@@ -117,4 +143,26 @@ mod tests {
         let low = evaluate_confidence(&[0.005, 0.002], false);
         assert_eq!(low.confidence, "low");
     }
+
+    fn dense_context(score: f64) -> serde_json::Value {
+        serde_json::json!({"content": "x", "score": score, "match_source": "dense"})
+    }
+
+    #[test]
+    fn test_crag_dense_similarity_scale() {
+        let high = evaluate_contexts_confidence(&[dense_context(0.9), dense_context(0.8)]);
+        assert_eq!(high.confidence, "high");
+
+        let medium = evaluate_contexts_confidence(&[dense_context(0.6)]);
+        assert_eq!(medium.confidence, "medium");
+
+        let low = evaluate_contexts_confidence(&[dense_context(0.4)]);
+        assert_eq!(low.confidence, "low");
+    }
+
+    #[test]
+    fn test_crag_fused_scale_unchanged() {
+        let fused = serde_json::json!({"content": "x", "score": 0.03});
+        assert_eq!(evaluate_contexts_confidence(&[fused]).confidence, "high");
+    }
 }

db/migrations/004_dense_search.sql

@@ -0,0 +1,23 @@
+-- Migration 004: dense-only retrieval for dense-first strategy.
+-- hybrid_search_v5 fuses dense+lexical in one RRF step, which dilutes a
+-- strong dense signal with a weak lexical one. The application now tries
+-- a dense-only lookup first and falls back to the hybrid functions only
+-- when dense coverage is weak. Score is cosine similarity (1 - distance),
+-- so higher is better, matching the lexical ts_rank direction.
+CREATE OR REPLACE FUNCTION dense_search_v5(
+    query_embedding vector(768),
+    collection_filter uuid DEFAULT NULL,
+    api_key_filter uuid DEFAULT NULL
+) RETURNS TABLE(content text, document_id uuid, source_url text, score real) LANGUAGE sql STABLE AS $$
+  SELECT c.content, c.document_id, d.source_url,
+         (1.0 - (c.embedding <=> query_embedding))::real AS score
+  FROM knowledge_chunks c
+  JOIN knowledge_documents d ON d.id = c.document_id
+  JOIN knowledge_collections k ON k.id = d.collection_id
+  WHERE api_key_filter IS NOT NULL
+    AND k.api_key_id = api_key_filter
+    AND (collection_filter IS NULL OR d.collection_id = collection_filter)
+    AND c.embedding IS NOT NULL
+  ORDER BY c.embedding <=> query_embedding ASC
+  LIMIT 10
+$$;