from openai import OpenAI
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any
from enum import Enum
import asyncio
import asyncpg
import json
from datetime import datetime
class RetrievalStrategy(Enum):
VECTOR = "vector" # Best for semantic/meaning-based queries
KEYWORD = "keyword" # Best for exact terms, names, IDs
HYBRID = "hybrid" # Best overall -- combines both strengths
@dataclass
class Document:
id: str
content: str
metadata: Dict[str, Any]
score: float = 0.0
@dataclass
class RAGResponse:
answer: str
sources: List[Dict[str, Any]]
confidence: float
retrieval_time_ms: float
generation_time_ms: float
total_tokens: int
model: str
@dataclass
class RAGConfig:
"""Configuration for RAG pipeline"""
# Retrieval -- cast a wide net first, then narrow down
retrieval_strategy: RetrievalStrategy = RetrievalStrategy.HYBRID
top_k: int = 10 # How many docs to retrieve initially
similarity_threshold: float = 0.7 # Below this = probably irrelevant noise
# Re-ranking -- expensive but dramatically improves precision
enable_reranking: bool = True
rerank_top_k: int = 5 # Keep only top 5 after re-ranking (what the LLM sees)
# Generation
model: str = "gpt-4o"
temperature: float = 0 # Zero for factual accuracy; increase for creative tasks
max_tokens: int = 1000
# Query processing -- helps when user queries are vague or use different terms
enable_query_expansion: bool = True
enable_hyde: bool = False # Hypothetical Document Embeddings -- powerful but adds latency
class ProductionRAG:
"""Enterprise-grade RAG system"""
def __init__(self, database_url: str, config: RAGConfig = None):
self.database_url = database_url
self.config = config or RAGConfig()
self.openai = OpenAI()
self.pool: asyncpg.Pool = None
async def initialize(self):
"""Initialize database pool"""
self.pool = await asyncpg.create_pool(
self.database_url,
min_size=5,
max_size=20
)
async def close(self):
if self.pool:
await self.pool.close()
# ============= EMBEDDING =============
def _get_embedding(self, text: str) -> List[float]:
response = self.openai.embeddings.create(
model="text-embedding-3-small",
input=text
)
return response.data[0].embedding
def _get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
response = self.openai.embeddings.create(
model="text-embedding-3-small",
input=texts
)
return [e.embedding for e in response.data]
# ============= QUERY PROCESSING =============
async def _process_query(self, query: str) -> List[str]:
"""Transform query for better retrieval"""
queries = [query]
if self.config.enable_query_expansion:
expanded = await self._expand_query(query)
queries.extend(expanded)
if self.config.enable_hyde:
hyde_doc = await self._generate_hyde(query)
queries.append(hyde_doc)
return queries
async def _expand_query(self, query: str) -> List[str]:
"""Generate query variations.
Why: Users say "vacation policy" but docs say "PTO guidelines."
Expansion bridges this vocabulary gap by generating synonyms and
alternative phrasings so retrieval catches more relevant docs.
"""
response = self.openai.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": """Generate 3 alternative phrasings of this search query.
Return as JSON: {"queries": ["...", "...", "..."]}"""
},
{"role": "user", "content": query}
],
response_format={"type": "json_object"},
temperature=0.7
)
result = json.loads(response.choices[0].message.content)
return result.get("queries", [])
async def _generate_hyde(self, query: str) -> str:
"""Hypothetical Document Embeddings -- generate an ideal answer first.
Why: A question's embedding lives in "question space," but your docs
live in "answer space." HyDE generates a fake answer, embeds that instead,
and searches for real docs that look like that answer. Clever trick that
significantly improves recall for open-ended questions.
"""
response = self.openai.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": """You are a helpful assistant. Write a short, factual paragraph
that would be the perfect answer to the user's question."""
},
{"role": "user", "content": query}
],
max_tokens=200
)
return response.choices[0].message.content
# ============= RETRIEVAL =============
async def _retrieve_vector(
self,
query_embedding: List[float],
top_k: int
) -> List[Document]:
"""Vector similarity search"""
async with self.pool.acquire() as conn:
rows = await conn.fetch("""
SELECT
id::text,
content,
metadata,
1 - (embedding <=> $1::vector) as score
FROM documents
WHERE 1 - (embedding <=> $1::vector) > $2
ORDER BY embedding <=> $1::vector
LIMIT $3
""", str(query_embedding), self.config.similarity_threshold, top_k)
return [
Document(
id=row['id'],
content=row['content'],
metadata=json.loads(row['metadata']),
score=row['score']
)
for row in rows
]
async def _retrieve_keyword(
self,
query: str,
top_k: int
) -> List[Document]:
"""Full-text keyword search"""
async with self.pool.acquire() as conn:
rows = await conn.fetch("""
SELECT
id::text,
content,
metadata,
ts_rank(to_tsvector('english', content),
plainto_tsquery('english', $1)) as score
FROM documents
WHERE to_tsvector('english', content) @@ plainto_tsquery('english', $1)
ORDER BY score DESC
LIMIT $2
""", query, top_k)
return [
Document(
id=row['id'],
content=row['content'],
metadata=json.loads(row['metadata']),
score=row['score']
)
for row in rows
]
async def _retrieve_hybrid(
self,
query: str,
query_embedding: List[float],
top_k: int
) -> List[Document]:
"""Hybrid search with Reciprocal Rank Fusion.
Why hybrid? Vector search understands meaning ("PTO" matches "vacation")
but misses exact terms. Keyword search nails exact matches ("error code
E-4012") but misses synonyms. Combining both with RRF consistently
outperforms either approach alone.
"""
# Fetch 2x candidates from each method -- RRF will fuse and trim
vector_results = await self._retrieve_vector(query_embedding, top_k * 2)
keyword_results = await self._retrieve_keyword(query, top_k * 2)
# Build document map
doc_map: Dict[str, Document] = {}
vector_ranks: Dict[str, int] = {}
keyword_ranks: Dict[str, int] = {}
for i, doc in enumerate(vector_results):
doc_map[doc.id] = doc
vector_ranks[doc.id] = i + 1
for i, doc in enumerate(keyword_results):
if doc.id not in doc_map:
doc_map[doc.id] = doc
keyword_ranks[doc.id] = i + 1
# Reciprocal Rank Fusion: score = sum(1/(k+rank)) across all rankings.
# k=60 is the standard constant (from the original paper) -- it dampens
# the effect of very high ranks so one method can't dominate.
k = 60
rrf_scores: Dict[str, float] = {}
for doc_id in doc_map:
v_rank = vector_ranks.get(doc_id, len(vector_results) + 1)
k_rank = keyword_ranks.get(doc_id, len(keyword_results) + 1)
rrf_scores[doc_id] = 1/(k + v_rank) + 1/(k + k_rank)
# Sort by RRF score
sorted_ids = sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True)
return [doc_map[doc_id] for doc_id in sorted_ids[:top_k]]
async def _retrieve(
self,
queries: List[str]
) -> List[Document]:
"""Main retrieval method"""
import time
start = time.time()
all_docs: Dict[str, Document] = {}
for query in queries:
query_embedding = self._get_embedding(query)
if self.config.retrieval_strategy == RetrievalStrategy.VECTOR:
docs = await self._retrieve_vector(query_embedding, self.config.top_k)
elif self.config.retrieval_strategy == RetrievalStrategy.KEYWORD:
docs = await self._retrieve_keyword(query, self.config.top_k)
else: # HYBRID
docs = await self._retrieve_hybrid(query, query_embedding, self.config.top_k)
for doc in docs:
if doc.id not in all_docs or doc.score > all_docs[doc.id].score:
all_docs[doc.id] = doc
# Sort by score and limit
sorted_docs = sorted(all_docs.values(), key=lambda x: x.score, reverse=True)
return sorted_docs[:self.config.top_k]
# ============= RE-RANKING =============
async def _rerank(
self,
query: str,
documents: List[Document]
) -> List[Document]:
"""Re-rank documents using LLM as judge"""
if not self.config.enable_reranking or len(documents) <= self.config.rerank_top_k:
return documents[:self.config.rerank_top_k]
# Use LLM to score relevance
doc_texts = "\n\n".join([
f"[Doc {i+1}] {doc.content[:500]}"
for i, doc in enumerate(documents)
])
response = self.openai.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": """Rate the relevance of each document to the query.
Return JSON: {"rankings": [{"doc": 1, "score": 0.9}, ...]}
Score from 0 (irrelevant) to 1 (highly relevant)."""
},
{
"role": "user",
"content": f"Query: {query}\n\nDocuments:\n{doc_texts}"
}
],
response_format={"type": "json_object"}
)
rankings = json.loads(response.choices[0].message.content)
# Apply new scores
for ranking in rankings.get("rankings", []):
idx = ranking["doc"] - 1
if 0 <= idx < len(documents):
documents[idx].score = ranking["score"]
# Re-sort and return top_k
return sorted(documents, key=lambda x: x.score, reverse=True)[:self.config.rerank_top_k]
# ============= GENERATION =============
def _build_context(self, documents: List[Document]) -> tuple[str, List[Dict]]:
"""Build context string and sources list"""
context_parts = []
sources = []
for i, doc in enumerate(documents, 1):
context_parts.append(f"[Source {i}]\n{doc.content}")
sources.append({
"id": i,
"doc_id": doc.id,
"title": doc.metadata.get("title", "Untitled"),
"source": doc.metadata.get("source", "Unknown"),
"chunk_index": doc.metadata.get("chunk_index"),
"score": round(doc.score, 3)
})
return "\n\n---\n\n".join(context_parts), sources
async def _generate(
self,
query: str,
context: str,
sources: List[Dict]
) -> RAGResponse:
"""Generate answer using LLM"""
import time
start = time.time()
system_prompt = """You are a helpful assistant that answers questions based on provided sources.
RULES:
1. Only use information from the provided sources
2. Always cite sources using [Source N] format
3. If sources don't contain the answer, say "I don't have information about that in my sources"
4. Be concise but thorough
5. If multiple sources agree, cite all of them"""
response = self.openai.chat.completions.create(
model=self.config.model,
messages=[
{"role": "system", "content": system_prompt},
{
"role": "user",
"content": f"""Sources:
{context}
Question: {query}
Please answer based on the sources above:"""
}
],
temperature=self.config.temperature,
max_tokens=self.config.max_tokens
)
gen_time = (time.time() - start) * 1000
return RAGResponse(
answer=response.choices[0].message.content,
sources=sources,
confidence=self._estimate_confidence(response.choices[0].message.content, sources),
retrieval_time_ms=0, # Set by caller
generation_time_ms=gen_time,
total_tokens=response.usage.total_tokens,
model=self.config.model
)
def _estimate_confidence(self, answer: str, sources: List[Dict]) -> float:
"""Estimate answer confidence based on source citations"""
# Count how many sources are cited
cited = sum(1 for s in sources if f"[Source {s['id']}]" in answer)
citation_ratio = cited / len(sources) if sources else 0
# Check for uncertainty markers
uncertainty_phrases = [
"i don't have information",
"not mentioned",
"cannot find",
"unclear",
"may not be"
]
has_uncertainty = any(p in answer.lower() for p in uncertainty_phrases)
if has_uncertainty:
return 0.3
elif citation_ratio >= 0.5:
return 0.9
elif citation_ratio > 0:
return 0.7
else:
return 0.5
# ============= MAIN QUERY METHOD =============
async def query(self, question: str) -> RAGResponse:
"""Main RAG query method"""
import time
# 1. Process query
queries = await self._process_query(question)
# 2. Retrieve
retrieval_start = time.time()
documents = await self._retrieve(queries)
retrieval_time = (time.time() - retrieval_start) * 1000
if not documents:
return RAGResponse(
answer="I couldn't find any relevant information to answer your question.",
sources=[],
confidence=0.0,
retrieval_time_ms=retrieval_time,
generation_time_ms=0,
total_tokens=0,
model=self.config.model
)
# 3. Re-rank
if self.config.enable_reranking:
documents = await self._rerank(question, documents)
# 4. Build context
context, sources = self._build_context(documents)
# 5. Generate
response = await self._generate(question, context, sources)
response.retrieval_time_ms = retrieval_time
return response
# Usage
async def main():
rag = ProductionRAG(
database_url="postgresql://user:pass@localhost/docs",
config=RAGConfig(
retrieval_strategy=RetrievalStrategy.HYBRID,
enable_reranking=True,
enable_query_expansion=True
)
)
await rag.initialize()
response = await rag.query("What is our company's vacation policy?")
print(f"Answer: {response.answer}")
print(f"\nSources:")
for source in response.sources:
print(f" [{source['id']}] {source['title']} (score: {source['score']})")
print(f"\nConfidence: {response.confidence}")
print(f"Retrieval: {response.retrieval_time_ms:.1f}ms")
print(f"Generation: {response.generation_time_ms:.1f}ms")
await rag.close()
asyncio.run(main())