400 lines
15 KiB
Python
400 lines
15 KiB
Python
"""
|
|
RAG Engine with advanced retrieval strategies.
|
|
Supports vector search, hybrid search, reranking, and agentic modes.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Dict, List, Optional
|
|
|
|
from llama_index.core import Settings, StorageContext, VectorStoreIndex, get_response_synthesizer, PromptTemplate
|
|
from llama_index.core.query_engine import SubQuestionQueryEngine
|
|
from llama_index.core.retrievers import VectorIndexRetriever
|
|
from llama_index.core.tools import QueryEngineTool
|
|
from llama_index.embeddings.ollama import OllamaEmbedding
|
|
from llama_index.llms.ollama import Ollama
|
|
from llama_index.core.postprocessor import SentenceTransformerRerank, SimilarityPostprocessor
|
|
from llama_index.vector_stores.postgres import PGVectorStore
|
|
from sqlalchemy import make_url
|
|
|
|
# BM25 retriever - needs to be installed separately
|
|
try:
|
|
from llama_index.retrievers.bm25 import BM25Retriever
|
|
HAS_BM25 = True
|
|
except ImportError:
|
|
HAS_BM25 = False
|
|
|
|
from app.core.pg_manager import get_pg_manager
|
|
from app.core.config import Config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RAGEngine:
|
|
"""Advanced RAG engine with multiple retrieval strategies."""
|
|
|
|
def __init__(self, collection_name: str):
|
|
"""
|
|
Initialize RAG engine for a specific knowledge base using PGVectorStore.
|
|
|
|
Args:
|
|
collection_name: Name of the knowledge base
|
|
"""
|
|
self.collection_name = collection_name
|
|
|
|
# Initialize Ollama LLM with optimized settings for 8B model
|
|
self.llm = Ollama(
|
|
model=Config.OLLAMA_MODEL,
|
|
base_url=Config.OLLAMA_HOST,
|
|
request_timeout=120.0, # Timeout for larger model
|
|
context_window=8192, # 8B model supports larger context
|
|
num_ctx=8192, # Match context_window
|
|
temperature=0.1, # Very low for factual RAG responses
|
|
num_predict=1000, # Comprehensive answers
|
|
top_k=40, # Standard sampling space
|
|
top_p=0.9, # Nucleus sampling for better quality
|
|
num_thread=8, # Use more CPU threads
|
|
num_gpu=0, # CPU mode (set to 1 if you have GPU)
|
|
)
|
|
Settings.llm = self.llm
|
|
|
|
# Initialize Ollama embeddings
|
|
self.embed_model = OllamaEmbedding(
|
|
model_name=Config.OLLAMA_EMBED_MODEL,
|
|
base_url=Config.OLLAMA_HOST
|
|
)
|
|
Settings.embed_model = self.embed_model
|
|
|
|
# Verify KB exists and get table name
|
|
self.pg_manager = get_pg_manager()
|
|
if not self.pg_manager.collection_exists(collection_name):
|
|
raise ValueError(f"Collection {collection_name} not found")
|
|
|
|
# Get KB metadata
|
|
conn = self.pg_manager.get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"SELECT id, table_name, embed_dim FROM knowledge_bases WHERE name = %s",
|
|
(collection_name,)
|
|
)
|
|
result = cur.fetchone()
|
|
if not result:
|
|
raise ValueError(f"Collection {collection_name} not found")
|
|
self.kb_id, self.table_name, self.embed_dim = result
|
|
|
|
if not self.table_name:
|
|
raise ValueError(f"Collection {collection_name} has not been migrated to PGVectorStore schema")
|
|
finally:
|
|
self.pg_manager.return_connection(conn)
|
|
|
|
logger.info(f"Initializing RAGEngine for KB: {collection_name} (table: {self.table_name})")
|
|
|
|
# Create PGVectorStore instance
|
|
import os
|
|
pg_user = os.getenv("POSTGRES_USER", os.getenv("USER", "postgres"))
|
|
pg_password = os.getenv("POSTGRES_PASSWORD", "")
|
|
pg_host = os.getenv("POSTGRES_HOST", "localhost")
|
|
pg_port = os.getenv("POSTGRES_PORT", "5432")
|
|
pg_db = os.getenv("POSTGRES_DB", "codeforge")
|
|
|
|
logger.info(f"Connecting to PostgreSQL: {pg_user}@{pg_host}:{pg_port}/{pg_db}")
|
|
|
|
# PGVectorStore automatically adds 'data_' prefix to table names
|
|
# Our schema already has 'data_' prefix, so we need to strip it
|
|
pgvector_table_name = self.table_name.removeprefix('data_')
|
|
logger.info(f"Using PGVectorStore table name: {pgvector_table_name} (actual table: {self.table_name})")
|
|
|
|
# Use individual parameters instead of connection_string
|
|
# This ensures PGVectorStore creates both sync and async connection strings correctly
|
|
self.vector_store = PGVectorStore.from_params(
|
|
host=pg_host,
|
|
port=pg_port,
|
|
database=pg_db,
|
|
user=pg_user,
|
|
password=pg_password,
|
|
table_name=pgvector_table_name, # PGVectorStore will add 'data_' prefix
|
|
embed_dim=self.embed_dim,
|
|
hybrid_search=False, # Start with pure vector search
|
|
hnsw_kwargs=None, # Use IVFFlat index (already created)
|
|
perform_setup=False # Don't recreate tables, use existing schema
|
|
)
|
|
|
|
# Create VectorStoreIndex from the vector store
|
|
self.index = VectorStoreIndex.from_vector_store(
|
|
vector_store=self.vector_store,
|
|
embed_model=self.embed_model
|
|
)
|
|
logger.info(f"Created VectorStoreIndex for {collection_name}")
|
|
|
|
# Create vector retriever
|
|
self.vector_retriever = VectorIndexRetriever(
|
|
index=self.index,
|
|
similarity_top_k=Config.TOP_K_RETRIEVAL
|
|
)
|
|
|
|
# Initialize reranker (DISABLED by default due to performance issues - 40+ seconds per query)
|
|
# The reranker downloads a large model and runs on CPU which is extremely slow
|
|
# Enable only if you have GPU acceleration or can tolerate 40+ second delays
|
|
self.reranker = None
|
|
logger.info("Reranker disabled for performance (enable in code if needed)")
|
|
|
|
# Create custom prompt template that PREVENTS hallucination
|
|
qa_prompt_template = PromptTemplate(
|
|
"You are a documentation assistant. Your ONLY job is to extract and present information from the context below.\n\n"
|
|
"Context from documentation:\n"
|
|
"---------------------\n"
|
|
"{context_str}\n"
|
|
"---------------------\n\n"
|
|
"STRICT RULES - VIOLATION WILL CAUSE SYSTEM FAILURE:\n"
|
|
"1. **ONLY use information explicitly stated in the context above**\n"
|
|
"2. **If the answer is not in the context, respond EXACTLY**: \"I don't have specific documentation about that. The available documentation covers: [list topics from context]\"\n"
|
|
"3. **NEVER invent steps, settings, UI elements, or procedures**\n"
|
|
"4. **ALWAYS cite the source** - Start your answer with \"According to [document name]...\" or \"The documentation states...\"\n"
|
|
"5. **Use exact quotes** when possible, especially for settings, field names, and procedures\n"
|
|
"6. **Use Markdown** - ## headings, **bold**, `code`, lists, but ONLY for information from the context\n"
|
|
"7. **If context is incomplete**, say \"The documentation provides partial information: [what's available]. For complete details, you may need to check [suggest where]\"\n\n"
|
|
"Question: {query_str}\n\n"
|
|
"Answer (extract from context only, cite sources):"
|
|
)
|
|
|
|
# Create similarity filter to enforce minimum relevance threshold
|
|
similarity_filter = SimilarityPostprocessor(similarity_cutoff=Config.SIMILARITY_THRESHOLD)
|
|
|
|
# Create query engine with custom prompt and strict filtering
|
|
# Using "simple_summarize" mode for faster responses (single LLM call instead of multiple)
|
|
node_postprocessors = [similarity_filter]
|
|
if self.reranker:
|
|
node_postprocessors.append(self.reranker)
|
|
|
|
self.query_engine = self.index.as_query_engine(
|
|
similarity_top_k=Config.TOP_K_RETRIEVAL,
|
|
node_postprocessors=node_postprocessors,
|
|
response_mode="simple_summarize",
|
|
text_qa_template=qa_prompt_template
|
|
)
|
|
logger.info(f"Created query engine with strict filtering (similarity >= {Config.SIMILARITY_THRESHOLD}, top_k={Config.TOP_K_RETRIEVAL})")
|
|
|
|
# Initialize BM25 retriever (for hybrid search)
|
|
self.bm25_retriever = None
|
|
if HAS_BM25:
|
|
try:
|
|
all_nodes = list(self.index.docstore.docs.values())
|
|
if all_nodes:
|
|
self.bm25_retriever = BM25Retriever.from_defaults(
|
|
nodes=all_nodes,
|
|
similarity_top_k=Config.TOP_K_RETRIEVAL
|
|
)
|
|
logger.info("Initialized BM25 retriever for hybrid search")
|
|
else:
|
|
logger.warning("No documents found for BM25 retriever")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to initialize BM25 retriever: {e}")
|
|
else:
|
|
logger.warning("BM25 retriever not available")
|
|
|
|
# Initialize agentic query engine
|
|
self.sub_question_engine = None
|
|
try:
|
|
query_tool = QueryEngineTool.from_defaults(
|
|
query_engine=self.query_engine,
|
|
name="knowledge_base",
|
|
description=f"Searches the {collection_name} knowledge base for relevant information"
|
|
)
|
|
self.sub_question_engine = SubQuestionQueryEngine.from_defaults(
|
|
query_engine_tools=[query_tool]
|
|
)
|
|
logger.info("Initialized agentic query engine")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to initialize agentic engine: {e}")
|
|
|
|
def _reciprocal_rank_fusion(
|
|
self,
|
|
results_list: List[List],
|
|
k: int = 60
|
|
) -> List:
|
|
"""
|
|
Fuse multiple retrieval results using reciprocal rank fusion.
|
|
|
|
Args:
|
|
results_list: List of retrieval result lists
|
|
k: RRF parameter (default 60)
|
|
|
|
Returns:
|
|
Fused and sorted list of nodes
|
|
"""
|
|
# Aggregate scores by node ID
|
|
node_scores = {}
|
|
|
|
for results in results_list:
|
|
for rank, node in enumerate(results, start=1):
|
|
node_id = node.node_id
|
|
score = 1.0 / (k + rank)
|
|
|
|
if node_id in node_scores:
|
|
node_scores[node_id]["score"] += score
|
|
else:
|
|
node_scores[node_id] = {
|
|
"node": node,
|
|
"score": score
|
|
}
|
|
|
|
# Sort by score
|
|
sorted_nodes = sorted(
|
|
node_scores.values(),
|
|
key=lambda x: x["score"],
|
|
reverse=True
|
|
)
|
|
|
|
return [item["node"] for item in sorted_nodes[:Config.TOP_K_RETRIEVAL]]
|
|
|
|
def _query_hybrid(self, question: str) -> Dict[str, any]:
|
|
"""
|
|
Execute hybrid search (vector + BM25).
|
|
|
|
Args:
|
|
question: Query string
|
|
|
|
Returns:
|
|
dict: Answer and sources
|
|
"""
|
|
if not self.bm25_retriever:
|
|
logger.warning("BM25 retriever not available, falling back to vector search")
|
|
return self._query_base(question)
|
|
|
|
# Retrieve from both methods
|
|
vector_results = self.vector_retriever.retrieve(question)
|
|
bm25_results = self.bm25_retriever.retrieve(question)
|
|
|
|
# Fuse results
|
|
fused_nodes = self._reciprocal_rank_fusion([vector_results, bm25_results])
|
|
|
|
# Synthesize answer
|
|
synthesizer = get_response_synthesizer()
|
|
response = synthesizer.synthesize(question, nodes=fused_nodes)
|
|
|
|
return self._format_response(response, fused_nodes)
|
|
|
|
def _apply_reranking(self, nodes: List, query: str) -> List:
|
|
"""
|
|
Apply reranking to retrieved nodes.
|
|
|
|
Args:
|
|
nodes: Retrieved nodes
|
|
query: Query string
|
|
|
|
Returns:
|
|
Reranked nodes
|
|
"""
|
|
if not self.reranker:
|
|
logger.warning("Reranker not available")
|
|
return nodes
|
|
|
|
reranked = self.reranker.postprocess_nodes(nodes, query_str=query)
|
|
return reranked
|
|
|
|
def _query_agentic(self, question: str) -> Dict[str, any]:
|
|
"""
|
|
Execute agentic RAG with sub-question decomposition.
|
|
|
|
Args:
|
|
question: Query string
|
|
|
|
Returns:
|
|
dict: Answer, sub-questions, and sources
|
|
"""
|
|
if not self.sub_question_engine:
|
|
logger.warning("Agentic engine not available, falling back to base query")
|
|
return self._query_base(question)
|
|
|
|
response = self.sub_question_engine.query(question)
|
|
|
|
# Extract sub-questions if available
|
|
sub_questions = []
|
|
if hasattr(response, "metadata") and response.metadata:
|
|
sub_qa = response.metadata.get("sub_qa", [])
|
|
sub_questions = [
|
|
{"question": item[0], "answer": item[1]}
|
|
for item in sub_qa
|
|
]
|
|
|
|
result = self._format_response(response, response.source_nodes)
|
|
result["sub_questions"] = sub_questions
|
|
|
|
return result
|
|
|
|
def _query_base(self, question: str) -> Dict[str, any]:
|
|
"""
|
|
Execute base vector search query using PGVectorStore and llama_index query engine.
|
|
|
|
Args:
|
|
question: Query string
|
|
|
|
Returns:
|
|
dict: Answer and sources
|
|
"""
|
|
logger.info(f"Executing base query: {question}")
|
|
|
|
# Use llama_index query engine (with reranking if available)
|
|
response = self.query_engine.query(question)
|
|
|
|
logger.info(f"Generated answer: {str(response)[:200]}")
|
|
|
|
# Format response
|
|
return self._format_response(response, response.source_nodes)
|
|
|
|
def _format_response(self, response, source_nodes) -> Dict[str, any]:
|
|
"""Format query response for return."""
|
|
sources = []
|
|
for node in source_nodes:
|
|
# Convert score to Python float for JSON serialization
|
|
score = node.score if hasattr(node, "score") else None
|
|
if score is not None:
|
|
score = float(score)
|
|
|
|
sources.append({
|
|
"text": node.text,
|
|
"metadata": node.metadata if hasattr(node, "metadata") else {},
|
|
"score": score
|
|
})
|
|
|
|
return {
|
|
"answer": str(response),
|
|
"sources": sources
|
|
}
|
|
|
|
def query(
|
|
self,
|
|
question: str,
|
|
use_hybrid: bool = False,
|
|
use_rerank: bool = False,
|
|
use_agentic: bool = False
|
|
) -> Dict[str, any]:
|
|
"""
|
|
Execute query with specified strategies.
|
|
|
|
Args:
|
|
question: Query string
|
|
use_hybrid: Enable hybrid search (not implemented yet)
|
|
use_rerank: Enable reranking (always enabled if reranker is available)
|
|
use_agentic: Enable agentic RAG
|
|
|
|
Returns:
|
|
dict: Answer and sources
|
|
"""
|
|
try:
|
|
# Route to appropriate query method
|
|
# Note: Reranking is always enabled if reranker is available (via node_postprocessors)
|
|
if use_agentic:
|
|
return self._query_agentic(question)
|
|
elif use_hybrid:
|
|
return self._query_hybrid(question)
|
|
else:
|
|
return self._query_base(question)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Query failed: {e}")
|
|
return {
|
|
"answer": f"Error executing query: {str(e)}",
|
|
"sources": []
|
|
}
|
|
|