User Question → LLM → Answer (Limited to training data)

User Question → Retrieve Relevant Documents → LLM + Context → Answer (Grounded in knowledge base)

┌─────────────────────────────────────────────────────────────┐

│                    KNOWLEDGE BASE                            │

│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐         │

│  │  Documents  │  │   Vectors   │  │  Metadata   │         │

│  └─────────────┘  └─────────────┘  └─────────────┘         │

└─────────────────────────────────────────────────────────────┘

                        │

                        │ Ingestion Pipeline

                        ▼

┌─────────────────────────────────────────────────────────────┐

│                  PROCESSING LAYER                            │

│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────┐  │

│  │  Chunk   │→│ Embed    │→│  Index   │→│  Store   │  │

│  └──────────┘  └──────────┘  └──────────┘  └──────────┘  │

└─────────────────────────────────────────────────────────────┘

                        │

                        │ Query

                        ▼

┌─────────────────────────────────────────────────────────────┐

│                  RETRIEVAL LAYER                            │

│  ┌────────────┐  ┌────────────┐  ┌────────────┐           │

│  │  Query     │→│  Semantic  │→│  Hybrid    │           │

│  │ Embedding  │  │  Search    │  │  Search    │           │

│  └────────────┘  └────────────┘  └────────────┘           │

│                      │                                      │

│                      ▼                                      │

│              ┌──────────────┐                               │

│              │  Rerank &    │                               │

│              │  Filter      │                               │

│              └──────────────┘                               │

└─────────────────────────────────────────────────────────────┘

                        │

                        │ Context

                        ▼

┌─────────────────────────────────────────────────────────────┐

│                  GENERATION LAYER                            │

│  ┌────────────┐  ┌────────────┐  ┌────────────┐           │

│  │  Prompt    │→│    LLM     │→│  Response  │           │

│  │  Building  │  │  Inference │  │  Synthesis │           │

│  └────────────┘  └────────────┘  └────────────┘           │

└─────────────────────────────────────────────────────────────┘

                        │

                        ▼

                    User Response

# ingestion/pipeline.py

from typing import List, Dict

from pathlib import Path

import hashlib



class DocumentIngestionPipeline:

    """Process and ingest documents into knowledge base"""



    def __init__(self, config: Dict):

        self.chunker = DocumentChunker(config['chunking'])

        self.embedder = EmbeddingGenerator(config['embeddings'])

        self.vector_store = VectorStore(config['vector_db'])



    async def ingest_document(self, document: Dict) -> List[str]:

        """

        Ingest a document into the knowledge base



        Returns: List of chunk IDs

        """

        # 1. Extract text and metadata

        text = document['content']

        metadata = {

            'title': document['title'],

            'source': document['source'],

            'author': document.get('author'),

            'created_at': document.get('created_at'),

            'doc_type': document.get('type', 'unknown'),

            'language': document.get('language', 'en')

        }



        # 2. Split into chunks

        chunks = self.chunker.chunk(text)



        # 3. Generate embeddings

        chunk_texts = [chunk['text'] for chunk in chunks]

        embeddings = await self.embedder.generate_batch(chunk_texts)



        # 4. Prepare records for storage

        records = []

        for chunk, embedding in zip(chunks, embeddings):

            record = {

                'id': self._generate_chunk_id(document['id'], chunk['index']),

                'document_id': document['id'],

                'text': chunk['text'],

                'embedding': embedding,

                'metadata': {

                    **metadata,

                    'chunk_index': chunk['index'],

                    'chunk_size': len(chunk['text']),

                    'start_char': chunk['start'],

                    'end_char': chunk['end']

                }

            }

            records.append(record)



        # 5. Store in vector database

        chunk_ids = await self.vector_store.insert(records)



        return chunk_ids



    def _generate_chunk_id(self, doc_id: str, chunk_index: int) -> str:

        """Generate unique chunk ID"""

        hash_input = f"{doc_id}_{chunk_index}"

        return hashlib.sha256(hash_input.encode()).hexdigest()[:32]

# retrieval/engine.py

from typing import List, Dict, Optional

import numpy as np



class RetrievalEngine:

    """Retrieve relevant documents for queries"""



    def __init__(self, vector_store, embedder, config: Dict):

        self.vector_store = vector_store

        self.embedder = embedder

        self.config = config

        self.reranker = Reranker(config.get('reranking'))



    async def retrieve(

        self,

        query: str,

        top_k: int = 10,

        filters: Optional[Dict] = None

    ) -> List[Dict]:

        """

        Retrieve relevant chunks for a query



        Args:

            query: User query

            top_k: Number of results to return

            filters: Metadata filters (e.g., {category: 'technology'})



        Returns:

            List of retrieved chunks with scores

        """

        # 1. Generate query embedding

        query_embedding = await self.embedder.generate(query)



        # 2. Semantic search

        results = await self.vector_store.similarity_search(

            query_embedding,

            top_k=top_k * 2,  # Retrieve more for reranking

            filters=filters

        )



        # 3. Rerank if configured

        if self.reranker and len(results) > top_k:

            results = await self.reranker.rerank(query, results, top_k)



        return results[:top_k]



    async def retrieve_with_hybrid_search(

        self,

        query: str,

        top_k: int = 10,

        alpha: float = 0.5,

        filters: Optional[Dict] = None

    ) -> List[Dict]:

        """

        Hybrid retrieval combining semantic and keyword search



        Args:

            query: User query

            top_k: Number of results

            alpha: Weight for semantic search (0-1)

            filters: Metadata filters



        Returns:

            Reranked combined results

        """

        # 1. Semantic search

        semantic_results = await self.vector_store.similarity_search(

            await self.embedder.generate(query),

            top_k=top_k * 2,

            filters=filters

        )



        # 2. Keyword search

        keyword_results = await self.vector_store.keyword_search(

            query,

            top_k=top_k * 2,

            filters=filters

        )



        # 3. Combine and rerank

        combined = self._combine_results(

            semantic_results,

            keyword_results,

            alpha

        )



        # 4. Rerank combined results

        if self.reranker:

            combined = await self.reranker.rerank(query, combined, top_k)



        return combined[:top_k]



    def _combine_results(

        self,

        semantic_results: List[Dict],

        keyword_results: List[Dict],

        alpha: float

    ) -> List[Dict]:

        """Combine semantic and keyword search results"""

        # Score normalization

        sem_scores = np.array([r['score'] for r in semantic_results])

        key_scores = np.array([r['score'] for r in keyword_results])



        sem_normalized = (sem_scores - sem_scores.min()) / (sem_scores.max() - sem_scores.min())

        key_normalized = (key_scores - key_scores.min()) / (key_scores.max() - key_scores.min())



        # Combine scores

        for i, result in enumerate(semantic_results):

            result['combined_score'] = alpha * sem_normalized[i]



        for i, result in enumerate(keyword_results):

            result['combined_score'] += (1 - alpha) * key_normalized[i]



        # Merge and sort by combined score

        seen = set()

        combined = []

        for result in semantic_results + keyword_results:

            if result['id'] not in seen:

                seen.add(result['id'])

                combined.append(result)



        combined.sort(key=lambda x: x['combined_score'], reverse=True)

        return combined

# generation/generator.py

from typing import List, Dict

import openai



class ResponseGenerator:

    """Generate responses using retrieved context"""



    def __init__(self, config: Dict):

        self.client = openai.AsyncClient(api_key=config['api_key'])

        self.model = config['model']

        self.temperature = config.get('temperature', 0.3)

        self.max_tokens = config.get('max_tokens', 1000)



    async def generate_response(

        self,

        query: str,

        context: List[Dict],

        conversation_history: Optional[List[Dict]] = None

    ) -> Dict:

        """

        Generate response using retrieved context



        Args:

            query: User query

            context: Retrieved chunks

            conversation_history: Previous messages (for chat)



        Returns:

            Generated response with citations

        """

        # 1. Build prompt with context

        prompt = self._build_prompt(query, context)



        # 2. Generate response

        messages = self._build_messages(prompt, conversation_history)



        response = await self.client.chat.completions.create(

            model=self.model,

            messages=messages,

            temperature=self.temperature,

            max_tokens=self.max_tokens

        )



        # 3. Extract response and citations

        answer = response.choices[0].message.content

        citations = self._extract_citations(response, context)



        return {

            'answer': answer,

            'citations': citations,

            'sources': self._get_unique_sources(context),

            'model': self.model,

            'tokens_used': response.usage.total_tokens

        }



    def _build_prompt(self, query: str, context: List[Dict]) -> str:

        """Build prompt with context"""

        context_str = "\n\n".join([

            f"[Source {i+1}]\n{chunk['text']}"

            for i, chunk in enumerate(context)

        ])



        prompt = f"""You are a helpful assistant that answers questions based on the provided context.



Context:

{context_str}



Question: {query}



Instructions:

1. Answer the question using only the provided context

2. If the context doesn't contain enough information, say so

3. Cite sources using [Source X] notation

4. Be concise and accurate

5. If asked for sources, provide them



Answer:"""



        return prompt



    def _build_messages(

        self,

        prompt: str,

        history: Optional[List[Dict]] = None

    ) -> List[Dict]:

        """Build message list for API"""

        messages = []



        if history:

            messages.extend(history)



        messages.append({

            "role": "user",

            "content": prompt

        })



        return messages



    def _extract_citations(

        self,

        response: openai.ChatCompletion,

        context: List[Dict]

    ) -> List[Dict]:

        """Extract citations from response"""

        answer = response.choices[0].message.content



        # Find source references like [Source 1], [Source 2], etc.

        import re

        citations = re.findall(r'\[Source (\d+)\]', answer)



        # Map to actual source chunks

        unique_citations = []

        for citation in set(citations):

            idx = int(citation) - 1  # Convert to 0-based index

            if idx < len(context):

                unique_citations.append({

                    'index': int(citation),

                    'chunk_id': context[idx]['id'],

                    'document_id': context[idx]['metadata']['document_id'],

                    'title': context[idx]['metadata']['title'],

                    'source': context[idx]['metadata']['source']

                })



        return unique_citations



    def _get_unique_sources(self, context: List[Dict]) -> List[Dict]:

        """Get unique sources from context"""

        seen = set()

        sources = []



        for chunk in context:

            doc_id = chunk['metadata']['document_id']

            if doc_id not in seen:

                seen.add(doc_id)

                sources.append({

                    'document_id': doc_id,

                    'title': chunk['metadata']['title'],

                    'source': chunk['metadata']['source'],

                    'author': chunk['metadata'].get('author'),

                    'created_at': chunk['metadata'].get('created_at')

                })



        return sources

-- Single query for vectors + metadata

SELECT

  d.title,

  d.content,

  d.metadata->>'category' as category,

  1 - (d.embedding <=> query_embedding) as similarity

FROM documents d

JOIN document_tags dt ON d.id = dt.document_id

WHERE d.status = 'published'

  AND dt.tag_id = ANY(SELECT id FROM tags WHERE name IN ('AI', 'ML'))

  AND d.created_at > NOW() - INTERVAL '1 year'

ORDER BY d.embedding <=> query_embedding

LIMIT 20;

-- With proper indexing

CREATE INDEX idx_documents_embedding_hnsw ON documents

  USING hnsw(embedding vector_cosine_ops)

  WITH (m = 16, ef_construction = 64);



-- Query performance: 15-30ms for 1M vectors

embeddings = OpenAIEmbeddings(

    model="text-embedding-3-small",

    dimensions=1536  # Can truncate to 512 for faster search

)

# For legal/medical/technical content

from sentence_transformers import SentenceTransformer



model = SentenceTransformer('BAAI/bge-large-en-v1.5')

embeddings = model.encode(texts)

# Reduce from 1536 to 512 dimensions (faster search, lower storage)

import numpy as np

from sklearn.decomposition import PCA



def reduce_dimensions(embeddings: np.ndarray, target_dim: int = 512) -> np.ndarray:

    """Reduce embedding dimensions using PCA"""

    pca = PCA(n_components=target_dim)

    return pca.fit_transform(embeddings)



# Usage

full_embeddings = np.array([...])  # (N, 1536)

reduced_embeddings = reduce_dimensions(full_embeddings, 512)

# Combine semantic and keyword embeddings

class HybridEmbedding:

    def __init__(self):

        self.semantic_model = OpenAIEmbeddings(model="text-embedding-3-small")

        self.bm25 = BM25Encoder()



    def embed_documents(self, texts: List[str]) -> Dict[str, np.ndarray]:

        """Generate both semantic and keyword embeddings"""

        semantic = self.semantic_model.embed_documents(texts)

        keyword = self.bm25.encode_documents(texts)



        return {

            'semantic': np.array(semantic),

            'keyword': np.array(keyword)

        }

# Expand queries with related terms for better retrieval

async def expand_query(query: str, llm) -> List[str]:

    """Generate query variations"""

    prompt = f"""Generate 3-5 alternative queries for: "{query}"



    Consider:

    - Synonyms

    - Related concepts

    - Different phrasings

    - Broader/narrower terms



    Return one query per line."""



    response = await llm.generate(prompt)

    variations = [line.strip() for line in response.split('\n') if line.strip()]



    return [query] + variations



# Usage

query_variations = await expand_query("How to implement RAG?", llm)

# Returns: [

#   "How to implement RAG?",

#   "Building retrieval-augmented generation systems",

#   "RAG implementation guide",

#   "Creating RAG applications",

#   "RAG system architecture"

# ]

# chunking/fixed_size.py

from typing import List, Dict



class FixedSizeChunker:

    """Split text into fixed-size chunks"""



    def __init__(self, chunk_size: int = 1000, overlap: int = 200):

        self.chunk_size = chunk_size

        self.overlap = overlap



    def chunk(self, text: str) -> List[Dict]:

        """Split text into chunks"""

        chunks = []

        start = 0

        chunk_index = 0



        while start < len(text):

            end = start + self.chunk_size

            chunk_text = text[start:end]



            chunks.append({

                'text': chunk_text,

                'index': chunk_index,

                'start': start,

                'end': end,

                'size': len(chunk_text)

            })



            chunk_index += 1

            start = end - self.overlap



        return chunks



# Pros: Simple, predictable

# Cons: May split sentences, loses context

# chunking/sentence.py

import re

from typing import List, Dict



class SentenceChunker:

    """Split text into sentence-based chunks"""



    def __init__(self, sentences_per_chunk: int = 5, overlap: int = 1):

        self.sentences_per_chunk = sentences_per_chunk

        self.overlap = overlap



    def chunk(self, text: str) -> List[Dict]:

        """Split text into sentence-based chunks"""

        # Split into sentences

        sentences = re.split(r'(?<=[.!?])\s+', text)



        chunks = []

        chunk_index = 0

        i = 0



        while i < len(sentences):

            # Get sentences for this chunk

            end = min(i + self.sentences_per_chunk, len(sentences))

            chunk_sentences = sentences[i:end]

            chunk_text = ' '.join(chunk_sentences)



            start_char = text.find(chunk_sentences[0])

            end_char = start_char + len(chunk_text)



            chunks.append({

                'text': chunk_text,

                'index': chunk_index,

                'start': start_char,

                'end': end_char,

                'size': len(chunk_text),

                'sentence_count': len(chunk_sentences)

            })



            chunk_index += 1

            i += self.sentences_per_chunk - self.overlap



        return chunks



# Pros: Preserves sentence boundaries, better context

# Cons: Variable chunk sizes, may be too short/long

# chunking/semantic.py

from typing import List, Dict

import numpy as np



class SemanticChunker:

    """Split text into semantically coherent chunks"""



    def __init__(self, embedder, max_chunk_size: int = 1500, threshold: float = 0.7):

        self.embedder = embedder

        self.max_chunk_size = max_chunk_size

        self.threshold = threshold



    async def chunk(self, text: str) -> List[Dict]:

        """Split text into semantic chunks"""

        # 1. Split into sentences

        sentences = self._split_sentences(text)



        # 2. Generate embeddings for each sentence

        sentence_embeddings = await self.embedder.embed_documents(sentences)



        # 3. Calculate similarities between consecutive sentences

        similarities = [

            self._cosine_similarity(sentence_embeddings[i], sentence_embeddings[i+1])

            for i in range(len(sentence_embeddings) - 1)

        ]



        # 4. Identify chunk boundaries (where similarity drops below threshold)

        boundaries = [0]

        for i, sim in enumerate(similarities):

            if sim < self.threshold:

                boundaries.append(i + 1)

        boundaries.append(len(sentences))



        # 5. Create chunks

        chunks = []

        chunk_index = 0



        for i in range(len(boundaries) - 1):

            start_idx = boundaries[i]

            end_idx = boundaries[i+1]



            # Combine sentences in this segment

            chunk_sentences = sentences[start_idx:end_idx]

            chunk_text = ' '.join(chunk_sentences)



            # Further split if chunk is too long

            if len(chunk_text) > self.max_chunk_size:

                sub_chunks = self._split_long_chunk(chunk_text, self.max_chunk_size)

                for sub_chunk in sub_chunks:

                    chunks.append({

                        'text': sub_chunk,

                        'index': chunk_index,

                        'type': 'semantic'

                    })

                    chunk_index += 1

            else:

                chunks.append({

                    'text': chunk_text,

                    'index': chunk_index,

                    'sentence_count': len(chunk_sentences),

                    'type': 'semantic'

                })

                chunk_index += 1



        return chunks



    def _split_sentences(self, text: str) -> List[str]:

        """Split text into sentences"""

        import re

        return [s.strip() for s in re.split(r'(?<=[.!?])\s+', text) if s.strip()]



    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:

        """Calculate cosine similarity"""

        return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))



    def _split_long_chunk(self, text: str, max_size: int) -> List[str]:

        """Split long chunk into smaller pieces"""

        # Fallback to fixed-size splitting

        chunks = []

        start = 0

        while start < len(text):

            end = start + max_size

            chunks.append(text[start:end])

            start = end - 200  # Add overlap

        return chunks



# Pros: Semantically coherent, better retrieval

# Cons: Slower (requires embeddings), more complex

# chunking/hierarchical.py

class HierarchicalChunker:

    """Create multi-level chunk hierarchy for different use cases"""



    def __init__(self, embedder):

        self.embedder = embedder



    async def chunk(self, text: str, document_id: str) -> Dict[str, List[Dict]]:

        """Create hierarchical chunks"""

        # Level 1: Document-level (for broad queries)

        doc_chunk = {

            'id': f"{document_id}_doc",

            'level': 'document',

            'text': text[:2000],  # Summary/first part

            'metadata': {'type': 'document_summary'}

        }



        # Level 2: Section-level (for medium queries)

        section_chunks = self._chunk_by_sections(text)



        # Level 3: Paragraph-level (for specific queries)

        paragraph_chunks = self._chunk_by_paragraphs(text)



        # Level 4: Sentence-level (for precise queries)

        sentence_chunks = self._chunk_by_sentences(text)



        return {

            'document': [doc_chunk],

            'sections': section_chunks,

            'paragraphs': paragraph_chunks,

            'sentences': sentence_chunks

        }



    def _chunk_by_sections(self, text: str) -> List[Dict]:

        """Split by markdown/document sections"""

        import re

        sections = re.split(r'\n#{1,3}\s+', text)

        return [{'text': s, 'level': 'section'} for s in sections if s.strip()]



    def _chunk_by_paragraphs(self, text: str) -> List[Dict]:

        """Split by paragraphs"""

        paragraphs = text.split('\n\n')

        return [{'text': p, 'level': 'paragraph'} for p in paragraphs if p.strip()]



    def _chunk_by_sentences(self, text: str) -> List[Dict]:

        """Split by sentences"""

        import re

        sentences = re.split(r'(?<=[.!?])\s+', text)

        return [{'text': s, 'level': 'sentence'} for s in sentences if s.strip()]



# Usage: Store all levels, retrieve based on query type

chunker = SemanticChunker(

    embedder=OpenAIEmbeddings(model="text-embedding-3-small"),

    max_chunk_size=1000,

    threshold=0.75

)

chunker = HierarchicalChunker(embedder)

# Allows retrieval at appropriate granularity

# retrieval/query_rewrite.py

from typing import List



class QueryRewriter:

    """Rewrite queries for better retrieval"""



    def __init__(self, llm):

        self.llm = llm



    async def rewrite(self, query: str, context: str = "") -> str:

        """Rewrite query to improve retrieval"""

        prompt = f"""Rewrite the following query to improve information retrieval.



Original query: {query}



Context: {context}



Guidelines:

1. Make the query more specific

2. Add relevant domain terms

3. Fix grammatical issues

4. Expand abbreviations

5. Keep the intent the same



Rewritten query:"""



        response = await self.llm.generate(prompt)

        return response.strip()



    async def decompose(self, query: str) -> List[str]:

        """Decompose complex query into sub-queries"""

        prompt = f"""Break down the following query into 2-4 simpler sub-queries.



Query: {query}



Return one sub-query per line."""



        response = await self.llm.generate(prompt)

        sub_queries = [line.strip() for line in response.split('\n') if line.strip()]



        return sub_queries

# retrieval/reranker.py

from typing import List, Dict

import openai



class Reranker:

    """Rerank retrieved results for better relevance"""



    def __init__(self, model: str = "gpt-4o-mini"):

        self.client = openai.AsyncClient()

        self.model = model



    async def rerank(

        self,

        query: str,

        results: List[Dict],

        top_k: int = 10

    ) -> List[Dict]:

        """Rerank results based on query relevance"""

        if len(results) <= top_k:

            return results



        # Prepare reranking prompt

        result_texts = [

            f"[{i+1}] {r['text'][:500]}"

            for i, r in enumerate(results[:20])  # Rerank top 20

        ]



        prompt = f"""Rank the following passages by their relevance to the query.



Query: {query}



Passages:

{chr(10).join(result_texts)}



Instructions:

1. Rank passages from most relevant (1) to least relevant (20)

2. Return only the rankings as a comma-separated list

3. Consider: direct answers, completeness, specificity



Rankings:"""



        response = await self.client.chat.completions.create(

            model=self.model,

            messages=[{"role": "user", "content": prompt}],

            temperature=0.1

        )



        # Parse rankings

        rankings = response.choices[0].message.content

        ranked_indices = [int(x.strip()) - 1 for x in rankings.split(',')]



        # Reorder results

        reranked = [results[i] for i in ranked_indices if i < len(results)]



        return reranked[:top_k]

# retrieval/hybrid.py

class HybridSearcher:

    """Combine semantic and keyword search"""



    def __init__(self, vector_store, keyword_index):

        self.vector_store = vector_store

        self.keyword_index = keyword_index  # BM25 or similar



    async def search(

        self,

        query: str,

        top_k: int = 10,

        alpha: float = 0.5

    ) -> List[Dict]:

        """

        Hybrid search combining semantic and keyword



        Args:

            query: Search query

            top_k: Number of results

            alpha: Semantic search weight (0-1)



        Returns:

            Combined and reranked results

        """

        # Semantic search

        semantic_results = await self.vector_store.search(query, top_k * 2)



        # Keyword search

        keyword_results = await self.keyword_index.search(query, top_k * 2)



        # Combine scores

        combined = self._combine_results(

            semantic_results,

            keyword_results,

            alpha

        )



        # Remove duplicates and sort

        seen = set()

        unique_results = []

        for result in combined:

            if result['id'] not in seen:

                seen.add(result['id'])

                unique_results.append(result)



        return unique_results[:top_k]



    def _combine_results(

        self,

        semantic: List[Dict],

        keyword: List[Dict],

        alpha: float

    ) -> List[Dict]:

        """Combine and score results"""

        # Normalize scores

        sem_scores = [r['score'] for r in semantic]

        key_scores = [r['score'] for r in keyword]



        sem_max, sem_min = max(sem_scores), min(sem_scores)

        key_max, key_min = max(key_scores), min(key_scores)



        # Normalize to 0-1

        for r in semantic:

            r['normalized_score'] = (r['score'] - sem_min) / (sem_max - sem_min) if sem_max > sem_min else 0



        for r in keyword:

            r['normalized_score'] = (r['score'] - key_min) / (key_max - key_min) if key_max > key_min else 0



        # Combine

        combined = {}



        for r in semantic:

            combined[r['id']] = {

                **r,

                'combined_score': alpha * r['normalized_score']

            }



        for r in keyword:

            if r['id'] in combined:

                combined[r['id']]['combined_score'] += (1 - alpha) * r['normalized_score']

            else:

                combined[r['id']] = {

                    **r,

                    'combined_score': (1 - alpha) * r['normalized_score']

                }



        # Sort by combined score

        results = list(combined.values())

        results.sort(key=lambda x: x['combined_score'], reverse=True)



        return results

# retrieval/filtering.py

from typing import Dict, List, Any



class MetadataFilter:

    """Apply metadata filters to search results"""



    @staticmethod

    def apply_filters(

        results: List[Dict],

        filters: Dict[str, Any]

    ) -> List[Dict]:

        """Filter results based on metadata"""

        filtered = results



        for key, value in filters.items():

            if isinstance(value, list):

                # Filter: value in list

                filtered = [

                    r for r in filtered

                    if r['metadata'].get(key) in value

                ]

            elif isinstance(value, dict):

                # Range filter

                if '$gte' in value:

                    filtered = [

                        r for r in filtered

                        if r['metadata'].get(key, 0) >= value['$gte']

                    ]

                if '$lte' in value:

                    filtered = [

                        r for r in filtered

                        if r['metadata'].get(key, float('inf')) <= value['$lte']

                    ]

            else:

                # Exact match

                filtered = [

                    r for r in filtered

                    if r['metadata'].get(key) == value

                ]



        return filtered



# Usage

filtered_results = MetadataFilter.apply_filters(results, {

    'category': ['technology', 'ai'],

    'created_at': {'$gte': '2025-01-01'},

    'status': 'published'

})

def build_basic_rag_prompt(query: str, context: List[Dict]) -> str:

    """Build basic RAG prompt"""

    context_str = "\n\n---\n\n".join([

        f"Document: {chunk['metadata']['title']}\n{chunk['text']}"

        for chunk in context

    ])



    return f"""You are a helpful assistant. Answer the following question using the provided context.



Context:

{context_str}



Question: {query}



Instructions:

1. Base your answer only on the provided context

2. If the context doesn't contain the answer, say "I don't have enough information to answer this"

3. Cite sources using [Document X] notation

4. Be accurate and concise



Answer:"""

def build_advanced_rag_prompt(

    query: str,

    context: List[Dict],

    conversation_history: List[Dict] = None

) -> str:

    """Build advanced RAG prompt with conversation history"""



    context_by_source = {}

    for chunk in context:

        source = chunk['metadata']['source']

        if source not in context_by_source:

            context_by_source[source] = []

        context_by_source[source].append(chunk)



    context_str = "\n\n".join([

        f"## {source}\n" + "\n\n".join([c['text'] for c in chunks])

        for source, chunks in context_by_source.items()

    ])



    history_str = ""

    if conversation_history:

        history_str = "\n\n".join([

            f"{msg['role']}: {msg['content']}"

            for msg in conversation_history[-5:]  # Last 5 messages

        ])



    return f"""You are an expert knowledge assistant. Help answer questions by synthesizing information from multiple sources.



### Conversation History

{history_str}



### Available Sources

{context_str}



### Current Question

{query}



### Instructions

1. Synthesize information from multiple sources when relevant

2. Acknowledge when sources disagree or conflict

3. Prioritize recent and authoritative sources

4. Use [Source: Document Title] citations

5. If information is missing, explicitly state what's unknown

6. Provide a clear, well-structured answer



### Answer Format

- Start with a direct answer

- Follow with supporting details

- Include source citations

- End with limitations (if any)



Answer:"""

def build_cot_rag_prompt(query: str, context: List[Dict]) -> str:

    """Build chain-of-thought RAG prompt"""

    return f"""Answer the following question using the provided context. Show your reasoning.



Context:

{' '.join([c['text'] for c in context[:3]])}



Question: {query}



Think step by step:

1. What is the question asking?

2. What relevant information is in the context?

3. What can I conclude from this information?

4. What information is missing?



Answer:"""

# generation/post_process.py

from typing import Dict, List

import re



class ResponsePostProcessor:

    """Post-process generated responses"""



    @staticmethod

    def add_citations(response: str, context: List[Dict]) -> Dict:

        """Add citation links to response"""

        # Find [Source X] references

        citations = re.findall(r'\[Source (\d+)\]', response)



        # Create citation mapping

        citation_map = {}

        for cit in set(citations):

            idx = int(cit) - 1

            if idx < len(context):

                citation_map[cit] = {

                    'title': context[idx]['metadata']['title'],

                    'source': context[idx]['metadata']['source'],

                    'url': context[idx]['metadata'].get('url', '#')

                }



        return {

            'response': response,

            'citations': citation_map

        }



    @staticmethod

    def extract_key_points(response: str) -> List[str]:

        """Extract key points from response"""

        prompt = f"""Extract the key points from the following response.



Response:

{response}



Return one key point per line."""



        # Use LLM to extract

        # Implementation depends on your LLM setup

        return []



    @staticmethod

    def format_response(

        response: str,

        citations: List[Dict],

        sources: List[Dict]

    ) -> Dict:

        """Format final response for API"""

        return {

            'answer': response,

            'citations': [

                {

                    'index': i + 1,

                    'title': c['title'],

                    'source': c['source'],

                    'url': c.get('url')

                }

                for i, c in enumerate(citations)

            ],

            'sources': sources,

            'answer_length': len(response),

            'citation_count': len(citations)

        }

# evaluation/retrieval.py

from typing import List, Dict



class RetrievalEvaluator:

    """Evaluate retrieval quality"""



    @staticmethod

    def precision_at_k(retrieved: List[Dict], relevant: List[str], k: int) -> float:

        """Calculate Precision@K"""

        retrieved_ids = [r['id'] for r in retrieved[:k]]

        relevant_retrieved = set(retrieved_ids) & set(relevant)

        return len(relevant_retrieved) / k



    @staticmethod

    def recall_at_k(retrieved: List[Dict], relevant: List[str], k: int) -> float:

        """Calculate Recall@K"""

        retrieved_ids = [r['id'] for r in retrieved[:k]]

        relevant_retrieved = set(retrieved_ids) & set(relevant)

        return len(relevant_retrieved) / len(relevant)



    @staticmethod

    def mrr(retrieved: List[Dict], relevant: List[str]) -> float:

        """Calculate Mean Reciprocal Rank"""

        retrieved_ids = [r['id'] for r in retrieved]

        for i, doc_id in enumerate(retrieved_ids, 1):

            if doc_id in relevant:

                return 1 / i

        return 0.0



    @staticmethod

    def ndcg(retrieved: List[Dict], relevant: List[str], k: int) -> float:

        """Calculate Normalized DCG"""

        retrieved_ids = [r['id'] for r in retrieved[:k]]

        dcg = 0.0

        for i, doc_id in enumerate(retrieved_ids, 1):

            if doc_id in relevant:

                dcg += 1 / np.log2(i + 1)



        # Ideal DCG

        idcg = sum(1 / np.log2(i + 1) for i in range(1, min(len(relevant), k) + 1))



        return dcg / idcg if idcg > 0 else 0.0

# evaluation/generation.py

import openai



class GenerationEvaluator:

    """Evaluate generation quality using LLM-as-a-judge"""



    def __init__(self, model: str = "gpt-4o"):

        self.client = openai.Client()

        self.model = model



    def evaluate_relevance(self, query: str, response: str, context: List[Dict]) -> Dict:

        """Evaluate if response is relevant to query"""

        prompt = f"""Rate the relevance of the following response to the query.



Query: {query}



Response: {response}



Available Context:

{' '.join([c['text'][:200] for c in context[:3]])}



Rate on:

1. Relevance (0-100): Does it answer the question?

2. Accuracy (0-100): Is it factually correct based on context?

3. Completeness (0-100): Does it provide sufficient detail?

4. Citation Quality (0-100): Are citations appropriate?



Return as JSON:

{{"relevance": X, "accuracy": Y, "completeness": Z, "citation_quality": W}}"""



        response = self.client.chat.completions.create(

            model=self.model,

            messages=[{"role": "user", "content": prompt}],

            response_format={"type": "json_object"}

        )



        import json

        return json.loads(response.choices[0].message.content)



    def evaluate_hallucination(self, response: str, context: List[Dict]) -> Dict:

        """Check for hallucinations"""

        context_text = " ".join([c['text'] for c in context])



        prompt = f"""Analyze the following response for hallucinations (information not supported by context).



Response:

{response}



Context:

{context_text}



Identify:

1. Factual claims not in context

2. Invented sources or citations

3. Contradictions to context

4. Speculative statements presented as fact



Return as JSON with hallucinations list and severity (low/medium/high)."""



        response = self.client.chat.completions.create(

            model=self.model,

            messages=[{"role": "user", "content": prompt}],

            response_format={"type": "json_object"}

        )



        import json

        return json.loads(response.choices[0].message.content)