src/retrieval.py

import logging
from typing import List, Dict, Any, Tuple, Optional
from pathlib import Path
from flashrank import Ranker, RerankRequest

from config import (
    TOP_K_INITIAL,
    TOP_K_RERANKED,
    RERANKER_MODEL_NAME,
    ENABLE_QUERY_REWRITING,
    DUAL_QUERY_MERGE_STRATEGY
)
from vector_store import get_vector_store
from query_rewriter import get_query_rewriter

FLASHRANK_CACHE = str(Path(__file__).resolve().parent / "flashrank_cache")

logger = logging.getLogger(__name__)


class RetrievalEngine:
    
    def __init__(self, vector_store=None):
        self.vector_store = vector_store or get_vector_store()
        self._reranker = None
    
    @property
    def reranker(self):
        if self._reranker is None:
            self._reranker = Ranker(
                model_name=RERANKER_MODEL_NAME,
                cache_dir=FLASHRANK_CACHE
            )
        return self._reranker
    
    def _merge_and_deduplicate(
        self,
        results_a: List[Dict[str, Any]],
        results_b: List[Dict[str, Any]],
        strategy: str = "score"
    ) -> List[Dict[str, Any]]:
        seen_ids = set()
        merged = []
        
        if strategy == "interleave":
            max_len = max(len(results_a), len(results_b))
            for i in range(max_len):
                if i < len(results_a):
                    chunk_id = (results_a[i].get("source", ""), results_a[i].get("chunk_index", 0))
                    if chunk_id not in seen_ids:
                        seen_ids.add(chunk_id)
                        merged.append(results_a[i])
                if i < len(results_b):
                    chunk_id = (results_b[i].get("source", ""), results_b[i].get("chunk_index", 0))
                    if chunk_id not in seen_ids:
                        seen_ids.add(chunk_id)
                        merged.append(results_b[i])
        else:
            combined = results_a + results_b
            combined.sort(key=lambda x: x.get("score", 0), reverse=True)
            for result in combined:
                chunk_id = (result.get("source", ""), result.get("chunk_index", 0))
                if chunk_id not in seen_ids:
                    seen_ids.add(chunk_id)
                    merged.append(result)
        
        return merged
    
    def _single_query_search(
        self,
        query: str,
        use_reranking: bool,
        top_k_initial: int,
        top_k_final: int
    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
        initial_results = self.vector_store.search(query=query, top_k=top_k_initial)
        
        stage_debug = {
            "stage1_count": len(initial_results),
            "stage1_top_score": initial_results[0]["score"] if initial_results else 0
        }
        
        if not use_reranking or len(initial_results) == 0:
            return initial_results[:top_k_final], stage_debug
        
        passages = [
            {"id": i, "text": r["text"], "meta": r}
            for i, r in enumerate(initial_results)
        ]
        
        rerank_request = RerankRequest(query=query, passages=passages)
        reranked = self.reranker.rerank(rerank_request)
        
        reranked_results = []
        for item in reranked:
            original = item["meta"]
            reranked_results.append({
                "id": original["id"],
                "score": item["score"],
                "original_score": original["score"],
                "text": item["text"],
                "source": original["source"],
                "chunk_index": original["chunk_index"],
                "page_number": original.get("page_number", -1),
                "metadata": original["metadata"]
            })
        
        reranked_results.sort(key=lambda x: x["score"], reverse=True)
        stage_debug["stage2_top_score"] = reranked_results[0]["score"] if reranked_results else 0
        for rank, item in enumerate(reranked_results[:top_k_final], 1):
            logger.debug("reranker rank=%d score=%.4f source=%s chunk=%d", rank, item["score"], item.get("source", "?"), item.get("chunk_index", -1))
        
        return reranked_results[:top_k_final], stage_debug
    
    def retrieve(
        self,
        query: str,
        chat_history: Optional[List[Dict[str, str]]] = None,
        use_reranking: bool = True,
        top_k_initial: int = TOP_K_INITIAL,
        top_k_final: int = TOP_K_RERANKED,
        skip_rewrite: bool = False,
    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
        chat_history = chat_history or []
        
        search_query = query
        rewrite_debug = None
        use_dual_query = False
        
        if not skip_rewrite and ENABLE_QUERY_REWRITING and chat_history:
            rewriter = get_query_rewriter()
            search_query, rewrite_debug = rewriter.rewrite(query, chat_history)
            use_dual_query = ENABLE_QUERY_REWRITING and search_query != query
        
        debug_info = {
            "query": query,
            "search_query": search_query,
            "query_rewritten": search_query != query,
            "rewrite_info": rewrite_debug,
            "use_reranking": use_reranking,
            "use_dual_query": use_dual_query,
            "top_k_initial": top_k_initial,
            "top_k_final": top_k_final
        }
        
        if use_dual_query:
            original_results, original_debug = self._single_query_search(
                query, use_reranking, top_k_initial, top_k_final * 2
            )
            rewritten_results, rewritten_debug = self._single_query_search(
                search_query, use_reranking, top_k_initial, top_k_final * 2
            )
            
            merged_results = self._merge_and_deduplicate(
                rewritten_results, original_results, DUAL_QUERY_MERGE_STRATEGY
            )
            final_results = merged_results[:top_k_final]
            
            debug_info["dual_query_stats"] = {
                "original_count": len(original_results),
                "rewritten_count": len(rewritten_results),
                "merged_count": len(merged_results),
                "merge_strategy": DUAL_QUERY_MERGE_STRATEGY
            }
            debug_info["stage1_count"] = original_debug.get("stage1_count", 0)
            debug_info["stage1_top_score"] = max(
                original_debug.get("stage1_top_score", 0),
                rewritten_debug.get("stage1_top_score", 0)
            )
            debug_info["final_method"] = "dual_query_reranked" if use_reranking else "dual_query_vector"
            if use_reranking:
                debug_info["stage2_top_score"] = final_results[0]["score"] if final_results else 0
            
            return final_results, debug_info
        
        final_results, stage_debug = self._single_query_search(
            search_query, use_reranking, top_k_initial, top_k_final
        )
        
        debug_info.update(stage_debug)
        debug_info["final_method"] = "reranked" if use_reranking else "vector_only"
        
        return final_results, debug_info
    
    def build_context(self, results: List[Dict[str, Any]]) -> str:
        if not results:
            return "Bağlam bulunamadı."
        
        parts = []
        for i, r in enumerate(results, 1):
            source = r.get("source", "Bilinmeyen")
            page = r.get("page_number", -1)
            text = r.get("text", "")
            
            if page > 0:
                parts.append(f"[Kaynak {i}: {source}, Sayfa {page}]\n{text}")
            else:
                parts.append(f"[Kaynak {i}: {source}]\n{text}")
        
        return "\n\n---\n\n".join(parts)
    
    def format_sources(self, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        sources = []
        for i, r in enumerate(results, 1):
            sources.append({
                "index": i,
                "source": r.get("source", "Bilinmeyen"),
                "page_number": r.get("page_number", -1),
                "text": r.get("text", ""),
                "score": r.get("score", 0),
                "original_score": r.get("original_score", r.get("score", 0))
            })
        return sources


_retrieval_engine_instance = None


def get_retrieval_engine() -> RetrievalEngine:
    global _retrieval_engine_instance
    if _retrieval_engine_instance is None:
        _retrieval_engine_instance = RetrievalEngine()
    return _retrieval_engine_instance


def reset_retrieval_engine():
    global _retrieval_engine_instance
    _retrieval_engine_instance = None


if __name__ == "__main__":
    from vector_store import VectorStore
    from embeddings import get_embedder
    
    store = VectorStore(use_memory=True)
    store.add_documents(
        ["Atlas ERP muhasebe modülü.", "Finans raporlama özellikleri."],
        [{"source": "test.pdf", "chunk_index": 0}, {"source": "test.pdf", "chunk_index": 1}]
    )
    
    engine = RetrievalEngine(vector_store=store)
    results, debug = engine.retrieve("muhasebe", use_reranking=True, top_k_final=2)
    
    print(f"Debug: {debug}")
    print(f"Context:\n{engine.build_context(results)}")