"""
Retrieval pipeline with confidence scoring and filtering.
"""
from typing import List, Dict, Any, Optional, Tuple
import logging
import re
from app.config import settings
from app.rag.embeddings import get_embedding_service
from app.rag.vectorstore import get_vector_store
from app.rag.intent import detect_intents, check_direct_match, get_intent_keywords
from app.models.schemas import RetrievalResult
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class RetrievalService:
"""
Handles document retrieval with confidence scoring.
Implements threshold-based filtering for quality control.
"""
def __init__(
self,
top_k: int = settings.TOP_K,
similarity_threshold: float = settings.SIMILARITY_THRESHOLD
):
"""
Initialize the retrieval service.
Args:
top_k: Number of results to retrieve
similarity_threshold: Minimum similarity score to consider relevant
"""
self.top_k = top_k
self.similarity_threshold = similarity_threshold
self.embedding_service = get_embedding_service()
self.vector_store = get_vector_store()
def retrieve(
self,
query: str,
tenant_id: str, # CRITICAL: Multi-tenant isolation
kb_id: str,
user_id: str,
top_k: Optional[int] = None
) -> Tuple[List[RetrievalResult], float, bool]:
"""
Retrieve relevant documents for a query.
Args:
query: User's question
tenant_id: Tenant ID for multi-tenant isolation (CRITICAL)
kb_id: Knowledge base ID to search
user_id: User ID for filtering
top_k: Optional override for number of results
Returns:
Tuple of (results, average_confidence, has_relevant_results)
"""
k = top_k or self.top_k
# Generate query embedding
logger.info(f"Generating embedding for query: {query[:50]}...")
query_embedding = self.embedding_service.embed_query(query)
# Search vector store with filters - MUST include tenant_id for isolation
filter_dict = {
"tenant_id": tenant_id, # CRITICAL: Multi-tenant isolation
"kb_id": kb_id,
"user_id": user_id
}
logger.info(f"Searching vector store with filters: {filter_dict}")
raw_results = self.vector_store.search(
query_embedding=query_embedding,
top_k=k,
filter_dict=filter_dict
)
if not raw_results:
logger.warning(f"No results found for query in kb_id={kb_id}")
return [], 0.0, False
# Convert to RetrievalResult objects
results = []
for r in raw_results:
results.append(RetrievalResult(
chunk_id=r['id'],
content=r['content'],
metadata=r['metadata'],
similarity_score=r['similarity_score']
))
# HEAVY CONFIDENCE MODE: Use maximum similarity score from top results
# This ensures confidence reflects the best match found, not dragged down by weaker results
if results:
# Get top 3 results and use the maximum similarity score
# This gives maximum confidence if there's at least one strong match
top_results = results[:3]
max_score = max(r.similarity_score for r in top_results)
# If max score is good (>=0.4), use it directly
# Otherwise, use weighted average of top 3 to avoid over-inflating weak matches
if max_score >= 0.4:
avg_confidence = max_score
else:
# For weaker matches, use weighted average of top 3
scores = [r.similarity_score for r in top_results]
weights = [1.0, 0.7, 0.5][:len(scores)] # Aggressive weighting
weighted_sum = sum(s * w for s, w in zip(scores, weights))
total_weight = sum(weights[:len(scores)])
avg_confidence = weighted_sum / total_weight if total_weight > 0 else max_score
else:
avg_confidence = 0.0
# Filter results above threshold
filtered_results = [
r for r in results
if r.similarity_score >= self.similarity_threshold
]
# If no results pass threshold but we have results, use top results anyway
# This prevents over-filtering when threshold is too strict
if not filtered_results and results:
logger.warning(f"No results passed threshold {self.similarity_threshold}, using top {min(3, len(results))} results anyway")
filtered_results = results[:min(3, len(results))]
# Recalculate confidence with the fallback results
if filtered_results:
scores = [r.similarity_score for r in filtered_results]
avg_confidence = sum(scores) / len(scores) if scores else 0.0
# DIRECT MATCH GATE: Check if at least one chunk directly matches query intent
# For integration/API questions, this gate is stricter
has_direct_match = False
if filtered_results:
chunk_texts = [r.content for r in filtered_results]
intents = detect_intents(query)
intent_keywords = get_intent_keywords(intents)
# For integration/API questions, require direct match
if "integration" in intents or "api" in query.lower():
has_direct_match = check_direct_match(query, chunk_texts, intent_keywords)
logger.info(f"Direct match check (strict for integration): {has_direct_match} (intents: {intents})")
else:
# For other questions, be more lenient - just check if important words match
query_words = set(re.findall(r'\b\w+\b', query.lower()))
stop_words = {"the", "a", "an", "is", "are", "was", "were", "be", "been",
"to", "of", "and", "or", "but", "in", "on", "at", "for",
"with", "how", "what", "when", "where", "why", "do", "does"}
important_words = query_words - stop_words
# Check if at least one important word appears in chunks
for chunk in chunk_texts:
chunk_lower = chunk.lower()
matches = sum(1 for word in important_words if word in chunk_lower)
if matches >= 1 and len(important_words) > 0: # At least one important word
has_direct_match = True
break
logger.info(f"Direct match check (lenient): {has_direct_match} (intents: {intents})")
# Only consider relevant if we have filtered results AND (direct match OR high confidence)
# High confidence (>0.40) can bypass direct match requirement for non-integration questions
has_relevant = len(filtered_results) > 0 and (has_direct_match or avg_confidence > 0.40)
logger.info(
f"Retrieved {len(results)} results, "
f"{len(filtered_results)} above threshold ({self.similarity_threshold}), "
f"avg confidence: {avg_confidence:.3f}, "
f"direct match: {has_direct_match}"
)
return filtered_results, avg_confidence, has_relevant
def get_context_for_llm(
self,
results: List[RetrievalResult],
max_tokens: int = settings.MAX_CONTEXT_TOKENS
) -> Tuple[str, List[Dict[str, Any]]]:
"""
Format retrieved results into context for the LLM.
Args:
results: List of retrieval results
max_tokens: Maximum tokens for context
Returns:
Tuple of (formatted_context, citation_info)
"""
if not results:
return "", []
context_parts = []
citations = []
current_tokens = 0
# Estimate tokens (rough approximation: 1 token ≈ 4 chars)
for i, result in enumerate(results):
chunk_text = result.content
estimated_tokens = len(chunk_text) // 4
if current_tokens + estimated_tokens > max_tokens:
logger.info(f"Truncating context at {i} chunks due to token limit")
break
# Format chunk with source info
source_info = f"[Source {i+1}: {result.metadata.get('file_name', 'Unknown')}]"
if result.metadata.get('page_number'):
source_info += f" (Page {result.metadata['page_number']})"
context_parts.append(f"{source_info}\n{chunk_text}")
# Build citation info
citations.append({
"index": i + 1,
"file_name": result.metadata.get('file_name', 'Unknown'),
"chunk_id": result.chunk_id,
"page_number": result.metadata.get('page_number'),
"similarity_score": result.similarity_score,
"excerpt": chunk_text[:200] + "..." if len(chunk_text) > 200 else chunk_text
})
current_tokens += estimated_tokens
formatted_context = "\n\n---\n\n".join(context_parts)
return formatted_context, citations
# Global retrieval service instance
_retrieval_service: Optional[RetrievalService] = None
def get_retrieval_service() -> RetrievalService:
"""Get the global retrieval service instance."""
global _retrieval_service
if _retrieval_service is None:
_retrieval_service = RetrievalService()
return _retrieval_service
