Rag Agent Builder

"""
Agentic RAG Implementation

Agent that makes intelligent decisions about what to retrieve and when.
Supports iterative refinement for complex questions.
"""

from typing import List, Tuple, Optional, Dict
from enum import Enum


class RetrievalDecision(Enum):
    """Decisions the agent can make about retrieval."""

    RETRIEVE = "retrieve"
    REFINE_QUERY = "refine_query"
    GENERATE = "generate"
    SEARCH_EXPANSION = "search_expansion"


class AgenticRAG:
    """RAG system where agent controls retrieval strategy."""

    def __init__(self, llm_client, vector_db, max_iterations: int = 3):
        """
        Initialize agentic RAG.

        Args:
            llm_client: LLM client for reasoning
            vector_db: Vector database for retrieval
            max_iterations: Maximum refinement iterations
        """
        self.llm = llm_client
        self.vector_db = vector_db
        self.max_iterations = max_iterations
        self.conversation_history = []

    def decide_retrieval(self, query: str, context: Optional[str] = None) -> Dict:
        """
        Use LLM to decide on retrieval strategy.

        Args:
            query: User query
            context: Previous context if iterating

        Returns:
            Decision dict with action and parameters
        """
        decision_prompt = f"""Based on this query, decide what retrieval strategy to use:

Query: {query}
{"Previous context: " + context if context else ""}

Decide one of:
1. RETRIEVE - directly retrieve documents
2. REFINE_QUERY - reformulate query first
3. SEARCH_EXPANSION - expand query with related terms
4. GENERATE - answer without retrieval (general knowledge)

Respond with decision and reasoning."""

        response = self.llm.generate(decision_prompt)
        return self._parse_decision(response, query)

    def _parse_decision(self, response: str, original_query: str) -> Dict:
        """Parse LLM decision response."""
        # Placeholder logic
        if "refine" in response.lower():
            return {"action": RetrievalDecision.REFINE_QUERY, "query": original_query}
        elif "expand" in response.lower():
            return {
                "action": RetrievalDecision.SEARCH_EXPANSION,
                "query": original_query,
            }
        elif "generate" in response.lower():
            return {"action": RetrievalDecision.GENERATE, "query": original_query}
        else:
            return {"action": RetrievalDecision.RETRIEVE, "query": original_query}

    def refine_query(self, original_query: str) -> str:
        """
        Refine query for better retrieval.

        Args:
            original_query: Original user query

        Returns:
            Refined query
        """
        refine_prompt = f"""Reformulate this query to be more specific and retrieval-friendly:

Original: {original_query}

Refined query:"""

        refined = self.llm.generate(refine_prompt)
        return refined.strip()

    def expand_query(self, query: str, num_expansions: int = 2) -> List[str]:
        """
        Generate related search queries.

        Args:
            query: Original query
            num_expansions: Number of related queries

        Returns:
            List of expanded queries
        """
        expand_prompt = f"""Generate {num_expansions} related search queries for:

Query: {query}

Related queries (one per line):"""

        response = self.llm.generate(expand_prompt)
        queries = response.strip().split("\n")
        return [q.strip() for q in queries if q.strip()]

    def retrieve_documents(self, query: str, k: int = 5) -> List[Tuple[str, float]]:
        """
        Retrieve documents for query.

        Args:
            query: Query string
            k: Number of results

        Returns:
            List of (document, score) tuples
        """
        # Convert query to embedding and retrieve
        results = self.vector_db.query(query, k=k)
        return results

    def evaluate_retrieved_docs(
        self, documents: List[str], query: str
    ) -> Tuple[bool, float]:
        """
        Evaluate if retrieved documents are sufficient.

        Args:
            documents: Retrieved documents
            query: Original query

        Returns:
            Tuple of (sufficient, confidence_score)
        """
        eval_prompt = f"""Are these documents sufficient to answer the query?

Query: {query}

Documents:
{"".join(f"{i+1}. {doc[:200]}..." for i, doc in enumerate(documents))}

Respond with: YES or NO, and confidence (0-1)"""

        response = self.llm.generate(eval_prompt)

        # Parse response
        sufficient = "yes" in response.lower()
        confidence = 0.8 if sufficient else 0.4  # Placeholder

        return sufficient, confidence

    def execute(self, query: str) -> Dict:
        """
        Execute agentic RAG pipeline.

        Args:
            query: User query

        Returns:
            Dict with answer and metadata
        """
        current_query = query
        iteration = 0
        all_retrieved_docs = []

        while iteration < self.max_iterations:
            # Decide retrieval strategy
            decision = self.decide_retrieval(current_query)

            if decision["action"] == RetrievalDecision.REFINE_QUERY:
                current_query = self.refine_query(current_query)
                iteration += 1
                continue

            elif decision["action"] == RetrievalDecision.SEARCH_EXPANSION:
                expanded_queries = self.expand_query(current_query)
                # Retrieve for all queries
                for expanded_q in expanded_queries:
                    docs = self.retrieve_documents(expanded_q, k=3)
                    all_retrieved_docs.extend(docs)
                iteration += 1
                continue

            elif decision["action"] == RetrievalDecision.RETRIEVE:
                docs = self.retrieve_documents(current_query, k=5)
                all_retrieved_docs.extend(docs)

                # Evaluate retrieved documents
                doc_texts = [doc[0] for doc in docs]
                sufficient, confidence = self.evaluate_retrieved_docs(
                    doc_texts, query
                )

                if sufficient or iteration >= self.max_iterations - 1:
                    break

                # Refine for next iteration
                current_query = self.refine_query(query)
                iteration += 1
                continue

            else:  # GENERATE
                break

        # Generate answer from all retrieved documents
        context = "\n\n".join([doc[0] for doc in all_retrieved_docs])
        answer = self.generate_answer(query, context)

        return {
            "answer": answer,
            "query": query,
            "refined_query": current_query,
            "iterations": iteration,
            "retrieved_docs_count": len(all_retrieved_docs),
            "documents": [doc[0] for doc in all_retrieved_docs],
        }

    def generate_answer(self, query: str, context: str) -> str:
        """
        Generate final answer with context.

        Args:
            query: Original query
            context: Retrieved context

        Returns:
            Generated answer
        """
        answer_prompt = f"""Based on the provided context, answer this question:

Context:
{context}

Question: {query}

Answer:"""

        return self.llm.generate(answer_prompt)

"""
Basic RAG Implementation

Simplest RAG implementation demonstrating document chunking,
embedding generation, vector storage, and retrieval.
"""

from typing import List, Tuple


class BasicRAG:
    """Basic RAG pipeline for document Q&A."""

    def __init__(self, embedding_model="all-MiniLM-L6-v2", vector_db=None):
        """
        Initialize basic RAG system.

        Args:
            embedding_model: Name of embedding model to use
            vector_db: Vector database instance (Chroma, Pinecone, etc.)
        """
        self.embedding_model = embedding_model
        self.vector_db = vector_db
        self.chunks = []

    def chunk_documents(
        self, documents: List[str], chunk_size: int = 1000, overlap: int = 100
    ) -> List[str]:
        """
        Split documents into overlapping chunks.

        Args:
            documents: List of document texts
            chunk_size: Size of each chunk in characters
            overlap: Overlap between consecutive chunks

        Returns:
            List of text chunks
        """
        chunks = []
        for doc in documents:
            # Split each document into chunks
            for i in range(0, len(doc), chunk_size - overlap):
                chunk = doc[i : i + chunk_size]
                if chunk.strip():
                    chunks.append(chunk)

        self.chunks = chunks
        return chunks

    def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """
        Generate embeddings for texts.

        Args:
            texts: List of text segments

        Returns:
            List of embedding vectors
        """
        # Placeholder - implement with actual embedding model
        # In practice, use: from sentence_transformers import SentenceTransformer
        # model = SentenceTransformer(self.embedding_model)
        # embeddings = model.encode(texts)
        return [[0.1] * 384 for _ in texts]  # Dummy embeddings

    def index_documents(self, documents: List[str]) -> None:
        """
        Index documents into vector database.

        Args:
            documents: List of document texts
        """
        chunks = self.chunk_documents(documents)
        embeddings = self.generate_embeddings(chunks)

        # Store in vector database
        for chunk, embedding in zip(chunks, embeddings):
            self.vector_db.add_text(chunk, embedding=embedding)

    def retrieve(self, query: str, k: int = 5) -> List[Tuple[str, float]]:
        """
        Retrieve top-k relevant chunks for query.

        Args:
            query: User query
            k: Number of results to retrieve

        Returns:
            List of (chunk, similarity_score) tuples
        """
        query_embedding = self.generate_embeddings([query])[0]
        results = self.vector_db.query(query_embedding, k=k)
        return results

    def generate_answer(self, query: str, context: str, llm_client) -> str:
        """
        Generate answer using retrieved context.

        Args:
            query: User query
            context: Retrieved context
            llm_client: LLM client instance

        Returns:
            Generated answer
        """
        prompt = f"""You are a helpful assistant. Answer the question based on the provided context.

Context:
{context}

Question: {query}

Answer:"""

        response = llm_client.generate(prompt)
        return response

    def query(self, query: str, llm_client, k: int = 5) -> Tuple[str, List[str]]:
        """
        Complete RAG pipeline: retrieve and generate.

        Args:
            query: User query
            llm_client: LLM client instance
            k: Number of chunks to retrieve

        Returns:
            Tuple of (answer, retrieved_chunks)
        """
        # Retrieve relevant chunks
        retrieved = self.retrieve(query, k=k)
        chunks = [item[0] for item in retrieved]
        context = "\n\n".join(chunks)

        # Generate answer
        answer = self.generate_answer(query, context, llm_client)

        return answer, chunks

"""
Advanced Retrieval Strategies for RAG

Implements hybrid search, reranking, and filtering techniques.
"""

from typing import List, Tuple, Optional
import math


class HybridRetriever:
    """Combines keyword (BM25) and semantic (embedding) search."""

    def __init__(self, bm25_retriever, vector_db, alpha: float = 0.5):
        """
        Initialize hybrid retriever.

        Args:
            bm25_retriever: BM25 keyword search instance
            vector_db: Vector database for semantic search
            alpha: Weight for combining results (0-1)
                  0 = pure keyword, 1 = pure semantic
        """
        self.bm25 = bm25_retriever
        self.vector_db = vector_db
        self.alpha = alpha

    def retrieve(self, query: str, k: int = 5) -> List[Tuple[str, float]]:
        """
        Retrieve using both keyword and semantic search.

        Args:
            query: User query
            k: Number of results

        Returns:
            Ranked list of (document, score) tuples
        """
        # Keyword search results
        keyword_results = self.bm25.search(query, k=k)
        keyword_scores = {doc: score for doc, score in keyword_results}

        # Semantic search results
        semantic_results = self.vector_db.query(query, k=k)
        semantic_scores = {doc: score for doc, score in semantic_results}

        # Combine scores
        combined_scores = {}
        all_docs = set(keyword_scores.keys()) | set(semantic_scores.keys())

        for doc in all_docs:
            keyword_score = keyword_scores.get(doc, 0.0)
            semantic_score = semantic_scores.get(doc, 0.0)
            combined = self.alpha * semantic_score + (1 - self.alpha) * keyword_score
            combined_scores[doc] = combined

        # Sort and return top-k
        sorted_results = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)
        return sorted_results[:k]


class DocumentReranker:
    """Rerank retrieved documents by relevance."""

    def __init__(self, reranker_model="cross-encoder/ms-marco-MiniLM-L-12-v2"):
        """
        Initialize reranker.

        Args:
            reranker_model: Cross-encoder model name
        """
        self.model = reranker_model

    def rerank(
        self, query: str, documents: List[str], top_k: int = 5
    ) -> List[Tuple[str, float]]:
        """
        Rerank documents by relevance to query.

        Args:
            query: User query
            documents: List of candidate documents
            top_k: Number of results to return

        Returns:
            Reranked list of (document, score) tuples
        """
        # Placeholder - use sentence_transformers.CrossEncoder in production
        # from sentence_transformers import CrossEncoder
        # model = CrossEncoder(self.model)
        # scores = model.predict([(query, doc) for doc in documents])

        # Simple placeholder: rank by query-document overlap
        scores = []
        query_words = set(query.lower().split())
        for doc in documents:
            doc_words = set(doc.lower().split())
            overlap = len(query_words & doc_words) / len(query_words)
            scores.append(overlap)

        ranked = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
        return ranked[:top_k]


class RelevanceFilter:
    """Filter retrieved documents by relevance threshold."""

    def __init__(self, threshold: float = 0.7):
        """
        Initialize filter.

        Args:
            threshold: Minimum relevance score (0-1)
        """
        self.threshold = threshold

    def filter_results(
        self, results: List[Tuple[str, float]]
    ) -> List[Tuple[str, float]]:
        """
        Filter results by relevance threshold.

        Args:
            results: List of (document, score) tuples

        Returns:
            Filtered results above threshold
        """
        return [
            (doc, score) for doc, score in results if score >= self.threshold
        ]

    def filter_and_log(
        self, results: List[Tuple[str, float]], query: str = ""
    ) -> Tuple[List[Tuple[str, float]], dict]:
        """
        Filter and return statistics.

        Args:
            results: List of (document, score) tuples
            query: Original query (for logging)

        Returns:
            Tuple of (filtered_results, statistics)
        """
        original_count = len(results)
        filtered = self.filter_results(results)
        filtered_count = len(filtered)

        stats = {
            "original_count": original_count,
            "filtered_count": filtered_count,
            "filtered_out": original_count - filtered_count,
            "threshold": self.threshold,
        }

        return filtered, stats


class ContextWindowManager:
    """Manage retrieved documents to fit within LLM context window."""

    def __init__(self, max_tokens: int = 3000, tokens_per_char: float = 0.25):
        """
        Initialize context manager.

        Args:
            max_tokens: Maximum tokens for context
            tokens_per_char: Approximate tokens per character
        """
        self.max_tokens = max_tokens
        self.tokens_per_char = tokens_per_char

    def estimate_tokens(self, text: str) -> int:
        """Estimate token count for text."""
        return int(len(text) * self.tokens_per_char)

    def fit_documents(self, documents: List[str]) -> Tuple[List[str], dict]:
        """
        Fit retrieved documents into context window.

        Args:
            documents: List of document chunks

        Returns:
            Tuple of (selected_documents, statistics)
        """
        selected = []
        total_tokens = 0

        for doc in documents:
            doc_tokens = self.estimate_tokens(doc)
            if total_tokens + doc_tokens <= self.max_tokens:
                selected.append(doc)
                total_tokens += doc_tokens
            else:
                break

        stats = {
            "total_tokens_available": self.max_tokens,
            "total_tokens_used": total_tokens,
            "documents_selected": len(selected),
            "documents_truncated": len(documents) - len(selected),
        }

        return selected, stats

    def prepare_context(self, documents: List[str]) -> str:
        """
        Prepare context string from documents.

        Args:
            documents: List of document chunks

        Returns:
            Formatted context string
        """
        selected, stats = self.fit_documents(documents)
        context = "\n\n".join(selected)
        return context

"""
Embedding Management for RAG

Handles embedding generation, normalization, and caching.
"""

from typing import List, Dict, Optional, Tuple
import math


class EmbeddingManager:
    """Manages embedding generation and storage."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2", cache_embeddings: bool = True):
        """
        Initialize embedding manager.

        Args:
            model_name: Name of embedding model
            cache_embeddings: Whether to cache generated embeddings
        """
        self.model_name = model_name
        self.cache = {} if cache_embeddings else None

    def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """
        Generate embeddings for texts.

        Args:
            texts: List of text segments

        Returns:
            List of embedding vectors
        """
        embeddings = []
        for text in texts:
            if self.cache is not None and text in self.cache:
                embeddings.append(self.cache[text])
            else:
                # Placeholder - use actual embedding model
                # from sentence_transformers import SentenceTransformer
                # model = SentenceTransformer(self.model_name)
                # embedding = model.encode(text)

                embedding = self._dummy_embedding(text)
                if self.cache is not None:
                    self.cache[text] = embedding
                embeddings.append(embedding)

        return embeddings

    def _dummy_embedding(self, text: str) -> List[float]:
        """Generate placeholder embedding."""
        # Create deterministic dummy embedding based on text
        seed = sum(ord(c) for c in text[:10])
        return [math.sin(seed + i) for i in range(384)]

    def normalize_embeddings(
        self, embeddings: List[List[float]]
    ) -> List[List[float]]:
        """
        Normalize embeddings to unit vectors.

        Args:
            embeddings: List of embedding vectors

        Returns:
            Normalized embeddings
        """
        normalized = []
        for embedding in embeddings:
            norm = math.sqrt(sum(x**2 for x in embedding))
            if norm > 0:
                normalized_emb = [x / norm for x in embedding]
            else:
                normalized_emb = embedding
            normalized.append(normalized_emb)

        return normalized

    def compute_similarity(
        self, embedding1: List[float], embedding2: List[float]
    ) -> float:
        """
        Compute cosine similarity between embeddings.

        Args:
            embedding1: First embedding
            embedding2: Second embedding

        Returns:
            Similarity score (0-1)
        """
        dot_product = sum(a * b for a, b in zip(embedding1, embedding2))
        norm1 = math.sqrt(sum(x**2 for x in embedding1))
        norm2 = math.sqrt(sum(x**2 for x in embedding2))

        if norm1 == 0 or norm2 == 0:
            return 0.0

        return dot_product / (norm1 * norm2)

    def get_cache_stats(self) -> Dict:
        """Get embedding cache statistics."""
        if self.cache is None:
            return {"caching_enabled": False}

        return {
            "caching_enabled": True,
            "cached_embeddings": len(self.cache),
            "cache_size_estimate_mb": len(self.cache) * 384 * 4 / (1024 * 1024),
        }

    def clear_cache(self) -> None:
        """Clear embedding cache."""
        if self.cache is not None:
            self.cache.clear()


class EmbeddingQualityAssessment:
    """Assess quality of embeddings."""

    @staticmethod
    def embedding_variance(embeddings: List[List[float]]) -> float:
        """
        Calculate variance across embeddings.

        Args:
            embeddings: List of embedding vectors

        Returns:
            Variance metric
        """
        if not embeddings:
            return 0.0

        # Calculate mean embedding
        dim = len(embeddings[0])
        mean_embedding = [
            sum(emb[i] for emb in embeddings) / len(embeddings) for i in range(dim)
        ]

        # Calculate variance
        variance = sum(
            sum((emb[i] - mean_embedding[i]) ** 2 for i in range(dim))
            for emb in embeddings
        ) / len(embeddings)

        return variance

    @staticmethod
    def embedding_distribution_quality(embeddings: List[List[float]]) -> Dict:
        """
        Assess quality of embedding distribution.

        Args:
            embeddings: List of embedding vectors

        Returns:
            Quality assessment dict
        """
        if not embeddings:
            return {"quality": "unknown", "metrics": {}}

        # Calculate pairwise similarities
        similarities = []
        for i in range(len(embeddings)):
            for j in range(i + 1, len(embeddings)):
                sim = EmbeddingManager.compute_similarity(
                    embeddings[i], embeddings[j]
                )
                similarities.append(sim)

        if not similarities:
            return {"quality": "unknown", "metrics": {}}

        avg_similarity = sum(similarities) / len(similarities)
        max_similarity = max(similarities)
        min_similarity = min(similarities)

        # Assess quality
        if avg_similarity < 0.3:
            quality = "good"  # Diverse embeddings
        elif avg_similarity < 0.6:
            quality = "moderate"
        else:
            quality = "poor"  # Too similar

        return {
            "quality": quality,
            "metrics": {
                "avg_similarity": avg_similarity,
                "max_similarity": max_similarity,
                "min_similarity": min_similarity,
                "diversity": 1 - avg_similarity,
            },
        }

    @staticmethod
    def detect_dead_dimensions(
        embeddings: List[List[float]], threshold: float = 1e-6
    ) -> List[int]:
        """
        Detect dimensions with low variance (dead dimensions).

        Args:
            embeddings: List of embedding vectors
            threshold: Variance threshold

        Returns:
            List of dead dimension indices
        """
        if not embeddings:
            return []

        dim = len(embeddings[0])
        dead_dims = []

        for d in range(dim):
            values = [emb[d] for emb in embeddings]
            variance = sum((v - sum(values) / len(values)) ** 2 for v in values) / len(
                values
            )
            if variance < threshold:
                dead_dims.append(d)

        return dead_dims

"""
RAG Evaluation Metrics

Evaluate retrieval quality and answer generation quality.
"""

from typing import List, Dict, Tuple, Optional
import math


class RetrievalMetrics:
    """Calculate retrieval quality metrics."""

    @staticmethod
    def precision_at_k(retrieved: List[str], relevant: List[str], k: int) -> float:
        """
        Calculate precision@k.

        Args:
            retrieved: List of retrieved documents
            relevant: List of relevant documents
            k: Cutoff position

        Returns:
            Precision@k score
        """
        if k == 0:
            return 0.0

        top_k_retrieved = set(retrieved[:k])
        relevant_set = set(relevant)
        matches = len(top_k_retrieved & relevant_set)
        return matches / k

    @staticmethod
    def recall_at_k(retrieved: List[str], relevant: List[str], k: int) -> float:
        """
        Calculate recall@k.

        Args:
            retrieved: List of retrieved documents
            relevant: List of relevant documents
            k: Cutoff position

        Returns:
            Recall@k score
        """
        if not relevant:
            return 0.0

        top_k_retrieved = set(retrieved[:k])
        relevant_set = set(relevant)
        matches = len(top_k_retrieved & relevant_set)
        return matches / len(relevant_set)

    @staticmethod
    def mean_reciprocal_rank(retrieved: List[str], relevant: List[str]) -> float:
        """
        Calculate Mean Reciprocal Rank (MRR).

        Args:
            retrieved: List of retrieved documents
            relevant: List of relevant documents

        Returns:
            MRR score
        """
        relevant_set = set(relevant)
        for rank, doc in enumerate(retrieved, 1):
            if doc in relevant_set:
                return 1.0 / rank
        return 0.0

    @staticmethod
    def ndcg(retrieved: List[Tuple[str, float]], relevant: List[str], k: int = 10) -> float:
        """
        Calculate Normalized Discounted Cumulative Gain (NDCG).

        Args:
            retrieved: List of (document, score) tuples
            relevant: List of relevant documents
            k: Cutoff position

        Returns:
            NDCG score
        """
        relevant_set = set(relevant)
        docs = [doc for doc, _ in retrieved[:k]]

        # Calculate DCG
        dcg = 0.0
        for rank, doc in enumerate(docs, 1):
            if doc in relevant_set:
                dcg += 1.0 / math.log2(rank + 1)

        # Calculate IDCG (ideal DCG)
        idcg = 0.0
        for rank in range(1, min(k, len(relevant_set)) + 1):
            idcg += 1.0 / math.log2(rank + 1)

        if idcg == 0:
            return 0.0

        return dcg / idcg

    @staticmethod
    def compute_retrieval_scores(
        retrieved: List[Tuple[str, float]], relevant: List[str]
    ) -> Dict[str, float]:
        """
        Compute all retrieval metrics.

        Args:
            retrieved: List of (document, score) tuples
            relevant: List of relevant documents

        Returns:
            Dict of metric names to scores
        """
        docs = [doc for doc, _ in retrieved]

        return {
            "precision_at_5": RetrievalMetrics.precision_at_k(docs, relevant, 5),
            "precision_at_10": RetrievalMetrics.precision_at_k(docs, relevant, 10),
            "recall_at_5": RetrievalMetrics.recall_at_k(docs, relevant, 5),
            "recall_at_10": RetrievalMetrics.recall_at_k(docs, relevant, 10),
            "mrr": RetrievalMetrics.mean_reciprocal_rank(docs, relevant),
            "ndcg_5": RetrievalMetrics.ndcg(retrieved, relevant, k=5),
            "ndcg_10": RetrievalMetrics.ndcg(retrieved, relevant, k=10),
        }


class AnswerQualityMetrics:
    """Evaluate answer quality."""

    @staticmethod
    def has_source_citations(answer: str, source_documents: List[str]) -> bool:
        """
        Check if answer cites sources.

        Args:
            answer: Generated answer
            source_documents: Retrieved source documents

        Returns:
            Whether answer contains source citations
        """
        # Simple check - looks for document references
        citations = ["source", "document", "according to", "cited from", "[1]"]
        return any(citation in answer.lower() for citation in citations)

    @staticmethod
    def answer_length_quality(answer: str, min_length: int = 50) -> Dict:
        """
        Evaluate answer length quality.

        Args:
            answer: Generated answer
            min_length: Minimum acceptable length

        Returns:
            Length quality assessment
        """
        length = len(answer)
        word_count = len(answer.split())

        if length < min_length:
            quality = "too_short"
        elif length > 5000:
            quality = "too_long"
        else:
            quality = "appropriate"

        return {
            "quality": quality,
            "character_count": length,
            "word_count": word_count,
        }

    @staticmethod
    def answer_coherence_score(answer: str) -> float:
        """
        Estimate answer coherence (0-1).

        Args:
            answer: Generated answer

        Returns:
            Coherence score
        """
        # Simple heuristic - score based on structure
        sentences = answer.split(".")
        valid_sentences = [s.strip() for s in sentences if len(s.strip()) > 10]

        if not valid_sentences:
            return 0.0

        # More sentences = better coherence (up to a point)
        coherence = min(len(valid_sentences) / 5.0, 1.0)
        return coherence

    @staticmethod
    def grounding_in_context(answer: str, context: str) -> float:
        """
        Estimate how well answer is grounded in context.

        Args:
            answer: Generated answer
            context: Retrieved context

        Returns:
            Grounding score (0-1)
        """
        answer_words = set(answer.lower().split())
        context_words = set(context.lower().split())

        if not answer_words:
            return 0.0

        overlap = len(answer_words & context_words)
        grounding = overlap / len(answer_words)

        return min(grounding, 1.0)


class RAGEvaluator:
    """Comprehensive RAG evaluation."""

    def __init__(self):
        """Initialize evaluator."""
        self.evaluation_history = []

    def evaluate_retrieval(
        self, retrieved: List[Tuple[str, float]], relevant: List[str], query: str
    ) -> Dict:
        """
        Evaluate retrieval step.

        Args:
            retrieved: List of (document, score) tuples
            relevant: List of relevant documents
            query: Original query

        Returns:
            Evaluation results
        """
        scores = RetrievalMetrics.compute_retrieval_scores(retrieved, relevant)

        result = {
            "step": "retrieval",
            "query": query,
            "metrics": scores,
            "retrieved_count": len(retrieved),
            "relevant_count": len(relevant),
        }

        self.evaluation_history.append(result)
        return result

    def evaluate_answer(self, answer: str, context: str, query: str) -> Dict:
        """
        Evaluate answer generation step.

        Args:
            answer: Generated answer
            context: Retrieved context
            query: Original query

        Returns:
            Evaluation results
        """
        has_citations = AnswerQualityMetrics.has_source_citations(answer, [context])
        length_quality = AnswerQualityMetrics.answer_length_quality(answer)
        coherence = AnswerQualityMetrics.answer_coherence_score(answer)
        grounding = AnswerQualityMetrics.grounding_in_context(answer, context)

        result = {
            "step": "answer",
            "query": query,
            "has_citations": has_citations,
            "length_quality": length_quality,
            "coherence": coherence,
            "grounding": grounding,
            "overall_quality": (coherence + grounding) / 2,
        }

        self.evaluation_history.append(result)
        return result

    def evaluate_rag_pipeline(
        self,
        query: str,
        retrieved: List[Tuple[str, float]],
        relevant_documents: List[str],
        answer: str,
        context: str,
    ) -> Dict:
        """
        Evaluate complete RAG pipeline.

        Args:
            query: User query
            retrieved: Retrieved documents
            relevant_documents: Relevant documents
            answer: Generated answer
            context: Retrieved context

        Returns:
            Complete evaluation
        """
        retrieval_eval = self.evaluate_retrieval(retrieved, relevant_documents, query)
        answer_eval = self.evaluate_answer(answer, context, query)

        return {
            "query": query,
            "retrieval": retrieval_eval,
            "answer": answer_eval,
            "overall_score": (
                retrieval_eval["metrics"]["ndcg_5"]
                + answer_eval["overall_quality"]
            )
            / 2,
        }

    def get_evaluation_summary(self) -> Dict:
        """Get summary of all evaluations."""
        if not self.evaluation_history:
            return {"evaluations_count": 0}

        retrieval_evals = [e for e in self.evaluation_history if e["step"] == "retrieval"]
        answer_evals = [e for e in self.evaluation_history if e["step"] == "answer"]

        avg_precision = (
            sum(e["metrics"]["precision_at_5"] for e in retrieval_evals)
            / len(retrieval_evals)
            if retrieval_evals
            else 0
        )

        avg_quality = (
            sum(e["overall_quality"] for e in answer_evals) / len(answer_evals)
            if answer_evals
            else 0
        )

        return {
            "total_evaluations": len(self.evaluation_history),
            "retrieval_evaluations": len(retrieval_evals),
            "answer_evaluations": len(answer_evals),
            "avg_retrieval_precision": avg_precision,
            "avg_answer_quality": avg_quality,
        }

"""
Vector Database Manager

Abstracts interactions with different vector database backends.
"""

from typing import List, Tuple, Dict, Optional
from abc import ABC, abstractmethod


class VectorDBBackend(ABC):
    """Abstract base for vector database backends."""

    @abstractmethod
    def add_text(self, text: str, embedding: List[float], metadata: Optional[Dict] = None) -> str:
        """Add text with embedding to database."""
        pass

    @abstractmethod
    def query(self, embedding: List[float], k: int = 5) -> List[Tuple[str, float]]:
        """Query database for similar embeddings."""
        pass

    @abstractmethod
    def delete(self, document_id: str) -> bool:
        """Delete document from database."""
        pass

    @abstractmethod
    def get_stats(self) -> Dict:
        """Get database statistics."""
        pass


class InMemoryVectorDB(VectorDBBackend):
    """Simple in-memory vector database for prototyping."""

    def __init__(self):
        """Initialize in-memory database."""
        self.documents = {}  # id -> (text, embedding, metadata)
        self.next_id = 0

    def add_text(
        self, text: str, embedding: List[float], metadata: Optional[Dict] = None
    ) -> str:
        """
        Add text to database.

        Args:
            text: Text content
            embedding: Embedding vector
            metadata: Optional metadata dict

        Returns:
            Document ID
        """
        doc_id = f"doc_{self.next_id}"
        self.documents[doc_id] = {
            "text": text,
            "embedding": embedding,
            "metadata": metadata or {},
        }
        self.next_id += 1
        return doc_id

    def query(self, embedding: List[float], k: int = 5) -> List[Tuple[str, float]]:
        """
        Query for similar documents.

        Args:
            embedding: Query embedding
            k: Number of results

        Returns:
            List of (text, similarity_score) tuples
        """
        if not self.documents:
            return []

        # Calculate similarities
        similarities = []
        for doc_id, doc_data in self.documents.items():
            similarity = self._cosine_similarity(embedding, doc_data["embedding"])
            similarities.append((doc_data["text"], similarity))

        # Sort and return top-k
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:k]

    def delete(self, document_id: str) -> bool:
        """Delete document."""
        if document_id in self.documents:
            del self.documents[document_id]
            return True
        return False

    def get_stats(self) -> Dict:
        """Get database statistics."""
        return {
            "backend": "in_memory",
            "document_count": len(self.documents),
            "next_id": self.next_id,
        }

    @staticmethod
    def _cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
        """Compute cosine similarity."""
        dot_product = sum(a * b for a, b in zip(vec1, vec2))
        norm1 = sum(x**2 for x in vec1) ** 0.5
        norm2 = sum(x**2 for x in vec2) ** 0.5

        if norm1 == 0 or norm2 == 0:
            return 0.0

        return dot_product / (norm1 * norm2)


class VectorDBManager:
    """Unified interface for vector database operations."""

    def __init__(self, backend: VectorDBBackend):
        """
        Initialize manager with backend.

        Args:
            backend: Vector database backend instance
        """
        self.backend = backend
        self.operation_count = 0

    def add_documents(
        self,
        texts: List[str],
        embeddings: List[List[float]],
        metadatas: Optional[List[Dict]] = None,
    ) -> List[str]:
        """
        Add multiple documents.

        Args:
            texts: List of text contents
            embeddings: List of embedding vectors
            metadatas: Optional list of metadata dicts

        Returns:
            List of document IDs
        """
        doc_ids = []
        for i, (text, embedding) in enumerate(zip(texts, embeddings)):
            metadata = metadatas[i] if metadatas else None
            doc_id = self.backend.add_text(text, embedding, metadata)
            doc_ids.append(doc_id)
            self.operation_count += 1

        return doc_ids

    def search(self, query_embedding: List[float], k: int = 5) -> List[Tuple[str, float]]:
        """
        Search for similar documents.

        Args:
            query_embedding: Query embedding vector
            k: Number of results

        Returns:
            List of (text, score) tuples
        """
        self.operation_count += 1
        return self.backend.query(query_embedding, k=k)

    def remove_document(self, document_id: str) -> bool:
        """Remove document."""
        result = self.backend.delete(document_id)
        if result:
            self.operation_count += 1
        return result

    def get_database_info(self) -> Dict:
        """Get database information."""
        stats = self.backend.get_stats()
        stats["manager_operations"] = self.operation_count
        return stats

    def health_check(self) -> bool:
        """Check if database is healthy."""
        try:
            stats = self.backend.get_stats()
            return stats is not None
        except Exception:
            return False


class VectorDBFactory:
    """Factory for creating vector database instances."""

    @staticmethod
    def create_db(db_type: str = "in_memory", **kwargs) -> VectorDBBackend:
        """
        Create vector database instance.

        Args:
            db_type: Type of database ("in_memory", "chroma", "pinecone", etc.)
            **kwargs: Database-specific arguments

        Returns:
            Vector database instance
        """
        if db_type == "in_memory":
            return InMemoryVectorDB()

        elif db_type == "chroma":
            # from chromadb import Client
            # return ChromaVectorDB(**kwargs)
            raise NotImplementedError("Chroma backend not implemented")

        elif db_type == "pinecone":
            # from pinecone import Index
            # return PineconeVectorDB(**kwargs)
            raise NotImplementedError("Pinecone backend not implemented")

        else:
            raise ValueError(f"Unknown database type: {db_type}")