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)}")