import asyncio import logging import os import time from abc import ABC, abstractmethod from datetime import datetime from json import load from pathlib import Path from tempfile import NamedTemporaryFile from textwrap import dedent from typing import Any, Dict, List, Optional, Protocol, Union import mlflow import mlflow.config import numpy as np import pandas as pd from dotenv import load_dotenv from langchain_core.output_parsers import JsonOutputParser, StrOutputParser from langchain_core.prompts import PromptTemplate from langchain_openai import ChatOpenAI from pydantic import BaseModel, ConfigDict, Field from tqdm import tqdm from tqdm.asyncio import tqdm_asyncio from src.config import Config, load_spreadsheet from src.config.logging import setup_logging from src.config.spreadsheets import ( GoogleSpreadsheetManager, GoogleSpreadsheetManagerMLFlow, ) from src.rag import ( ChromaSlideStore, HyperbolicScorer, MinScorer, PresentationRetriever, ScorerTypes, ) from src.rag.storage import LLMPresentationRetriever, RetrieverTypes logger = logging.getLogger(__name__) class MetricResult(BaseModel): """Container for metric calculation results""" name: str score: float explanation: Optional[str] = None model_config = ConfigDict(arbitrary_types_allowed=True) class BaseMetric(BaseModel): """Base class for evaluation metrics""" model_config = ConfigDict(arbitrary_types_allowed=True) @property def name(self) -> str: """Get metric name""" return self.__class__.__name__.lower() @abstractmethod async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: """Calculate metric value asynchronously""" pass def calculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: """Synchronous wrapper for calculate""" return asyncio.run(self.acalculate(run_output, ground_truth)) class PresentationMatch(BaseMetric): """Check if top-1 retrieved presentation matches ground truth""" async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: best_pres_info = run_output["contexts"][0] best_pres_name = best_pres_info["pres_name"] score = float(best_pres_name == ground_truth["pres_name"]) return MetricResult( name=self.name, score=score, explanation=f"Retrieved: {best_pres_name}, Expected: {ground_truth['pres_name']}", ) class PresentationFound(BaseMetric): """Check if ground truth presentation is in top-k""" async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: found_pres_names = [c["pres_name"] for c in run_output["contexts"]] score = float(ground_truth["pres_name"] in found_pres_names) return MetricResult( name=self.name, score=score, explanation=f"Found in positions: {[i for i, p in enumerate(found_pres_names) if p == ground_truth['pres_name']]}", ) class PresentationIdx(BaseMetric): async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: found_pres_names = [c["pres_name"] for c in run_output["contexts"]] score = float("nan") for i, pres in enumerate(found_pres_names): if pres == ground_truth["pres_name"]: score = float(i + 1) return MetricResult( name=self.name, score=score, explanation=( f"Presentation was found at position {score}" if score != float("nan") else "Presentation was not found" ), ) class PageMatch(BaseMetric): """Check if best page matches ground truth""" async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: score = 0.0 explanation = "" for pres_info in run_output["contexts"]: best_page_found = pres_info["pages"][0] if pres_info["pres_name"] == ground_truth["pres_name"]: reference_pages = ground_truth["pages"] if not reference_pages: score = 1.0 explanation = "No specific page required" elif best_page_found in reference_pages: score = 1.0 explanation = f"Found correct page {best_page_found}" else: explanation = f"Page mismatch: found {best_page_found}, expected {reference_pages}" return MetricResult(name=self.name, score=score, explanation=explanation) class PageFound(BaseMetric): """Check if any of ground truth pages are found in retrieved results The page is considered found if it appears in ANY position in the correct presentation. This is less strict than PageMatch which checks best matching page. """ async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: """Calculate metric value""" score = 0.0 explanation = "" # Get all pages from each presentation's results for pres_info in run_output["contexts"]: # Only check pages from the correct presentation if pres_info["pres_name"] == ground_truth["pres_name"]: found_pages = pres_info["pages"] reference_pages = ground_truth["pages"] # Handle case when no specific page required if not reference_pages: score = 1.0 explanation = "No specific page required" break # Check if any reference page is found matching_pages = set(found_pages) & set(reference_pages) if matching_pages: score = 1.0 explanation = f"Found pages {matching_pages} in positions {[found_pages.index(p)+1 for p in matching_pages]}" break else: explanation = f"No matching pages found. Retrieved: {found_pages}, Expected: {reference_pages}" return MetricResult(name=self.name, score=score, explanation=explanation) class PresentationCount(BaseMetric): """Count number of retrieved presentations""" async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: """Count presentations in retrieved results""" n_pres = len(run_output["contexts"]) return MetricResult( name=self.name, score=float(n_pres), explanation=f"Retrieved {n_pres} presentations", ) class BestChunkMatch(BaseMetric): async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: """Count presentations in retrieved results""" best_pres = run_output["contexts"][0] best_chunk = best_pres["best_chunk"] true_content_type = ground_truth["content_type"] found_content_type = best_chunk["chunk_type"] score = 0 if true_content_type in found_content_type: # text_content and visual_content score = 1 if true_content_type == "general" and found_content_type in [ "general_description", "conclusions_and_insights", "layout_and_composition", ]: score = 1 return MetricResult( name=self.name, score=float(score), explanation=f"Found content type '{found_content_type}' matches ground truth '{true_content_type}'", ) class LLMRelevance(BaseMetric): """LLM-based relevance scoring""" class RelevanceOutput(BaseModel): explanation: str = Field( description="Detailed explanation of why the content is/isn't relevant and how it relates to the query" ) relevance_score: int = Field(description="Relevance score from 0-10") model_config = ConfigDict(arbitrary_types_allowed=True) llm: ChatOpenAI = Field(description="LLM for relevance scoring") n_contexts: int = Field(default=-1, description="Number of contexts to evaluate") rate_limit_timeout: float = Field( default=-1.0, description="Rate limit timeout in seconds" ) def model_post_init(self, __context: Any): # fmt: off prompt_template = PromptTemplate.from_template(dedent( """\ You are an expert relevance assessor for a presentation retrieval system. Your task is to evaluate whether the retrieved slide descriptions contain relevant information that answers the user's query. Analyze all provided slide descriptions as a collective unit and provide a detailed explanation along with a relevance score. Each slide description contains these equally weighted sections: - Text Content: The actual text present on the slide - Visual Content: Description of images, charts, or other visual elements - Topic Overview: Main themes and subjects covered - Insights and Conclusions: Key takeaways and conclusions - Layout and Composition: Structural organization of the slide Scoring Guidelines: - 9-10: Perfect match - Content directly and comprehensively answers the query (e.g., query asks about sales trends, and slides show exact sales data and analysis) - 7-8: Strong relevance - Content clearly relates to the query but may miss minor details (e.g., query asks about complete workflow, slides show most but not all steps) - 5-6: Moderate relevance - Content addresses the query partially or indirectly (e.g., query asks about specific feature, slides discuss it briefly among other topics) - 3-4: Weak relevance - Content touches the topic but doesn't provide substantial answer (e.g., query asks about implementation details, slides only mention the concept) - 1-2: Minimal relevance - Only slight connection to the query (e.g., query asks about specific metric, slides only mention related general category) - 0: No relevance - Content has no connection to the query Evaluation Rules: 1. Award points if ANY section (text, visual, etc.) contains relevant information 2. In your explanation, cite specific sections and content that justify your score 3. Treat all sections equally - a match in visual content is as valuable as a match in text content 4. Consider all slides collectively - relevant information might be spread across multiple slides 5. Partial matches are valuable if they provide any useful information related to the query # Slide Descriptions {context_str} --- END OF SLIDE DESCRIPTIONS --- Question: {query_str} Output formatting: {format_instructions} """)) # fmt: on self._parser = JsonOutputParser(pydantic_object=self.RelevanceOutput) self._chain = prompt_template | self.llm.with_structured_output( self.RelevanceOutput ) async def acalculate(self, run_output: Dict, ground_truth: Dict) -> MetricResult: """Evaluate relevance of retrieved content""" if self.rate_limit_timeout > 0: time.sleep(1.05) # Rate limiting question = ground_truth["question"] pres = run_output["contexts"][0] contexts_used = ( pres["contexts"] if self.n_contexts <= 0 else pres["contexts"][: self.n_contexts] ) pres_context = "\n\n---\n\n".join(contexts_used) llm_out = await self._chain.ainvoke( dict( query_str=question, context_str=pres_context, format_instructions=self._parser.get_format_instructions(), ) ) llm_out_dict = llm_out.model_dump() return MetricResult( name=self.name, score=float(llm_out_dict["relevance_score"]), explanation=llm_out_dict["explanation"], ) class MetricsRegistry: """Factory for creating metric instances""" _metrics = { "presentationmatch": PresentationMatch, "presentationfound": PresentationFound, "presentationidx": PresentationIdx, "pagematch": PageMatch, "pagefound": PageFound, "presentationcount": PresentationCount, "bestchunkmatch": BestChunkMatch, "llmrelevance": LLMRelevance, } @classmethod def create(cls, metric_name: str, **kwargs) -> BaseMetric: """Create metric instance by name""" metric_cls = cls._metrics.get(metric_name.lower()) if metric_cls is None: raise ValueError(f"Unknown metric: {metric_name}") return metric_cls(**kwargs) class MetricPresets: """Available metric combinations for evaluation""" BASIC = [ "presentationmatch", "presentationfound", "presentationidx", "pagematch", "pagefound", "presentationcount", "bestchunkmatch", ] LLM = ["llmrelevance"] ALL = BASIC + LLM @classmethod def get_preset(cls, name: str) -> List[str]: """Get metric names from preset""" try: return getattr(cls, name.upper()) except AttributeError: raise ValueError(f"Unknown preset name: {name}") @classmethod def parse_specs(cls, specs: List[str]) -> List[str]: """Parse metric specifications Args: specs: List of metric specifications. Each item can be: - Preset name: "basic", "llm", "full" - Metric name: "presentationmatch", "llmrelevance", etc Returns: List of metric names with duplicates removed """ metrics = [] for spec in specs: # Check if spec is a preset name if hasattr(cls, spec.upper()): metrics.extend(cls.get_preset(spec)) else: metrics.append(spec.lower()) # Remove duplicates while preserving order seen = set() return [m for m in metrics if not (m in seen or seen.add(m))] # type: ignore class MlflowConfig(BaseModel): """Configuration for RAG evaluation""" experiment_name: str = "RAG_test" tracking_uri: str = f"sqlite:////{Config().navigator.eval_runs / 'mlruns.db'}" artifacts_uri: str = f"file:////{Config().navigator.eval_artifacts}" scorers: List[ScorerTypes] retriever: Union[PresentationRetriever, LLMPresentationRetriever] metrics: List[str] = ["presentationmatch", "pagematch"] n_judge_contexts: int = 10 write_to_google: bool = False metric_args: Dict[str, Any] = {} model_config = ConfigDict(arbitrary_types_allowed=True) def get_retriever_with_scorer(self, scorer: ScorerTypes) -> PresentationRetriever: self.retriever.set_scorer(scorer) return self.retriever def model_post_init(self, __context: Any) -> None: """Process metric specifications after initialization""" self.metrics = MetricPresets.parse_specs(self.metrics) logger.info(f"Using metrics: {self.metrics}") return super().model_post_init(__context) def get_log_params(self) -> Dict[str, Any]: """Get parameters for MLflow logging""" return { "experiment_name": self.experiment_name, "n_judge_contexts": self.n_judge_contexts, "metrics": ",".join(self.metrics), "metric_args": self.metric_args, } class RAGEvaluatorMlflow: """MLFlow-based evaluator for RAG pipeline""" def __init__( self, config: MlflowConfig, llm: Optional[ChatOpenAI] = None, max_concurrent: int = 5, ): load_dotenv() # Setup logging self._logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}") # Setup Evaluation self.config = config self.llm = llm or Config().model_config.load_vsegpt(model="openai/gpt-4o-mini") self._max_concurrent = max_concurrent # Setup GoogleSheets eval_spreadsheet_id = os.getenv("EVAL_SPREADSHEET_ID") if eval_spreadsheet_id is not None: self.gsheets = GoogleSpreadsheetManagerMLFlow(eval_spreadsheet_id) else: raise FileNotFoundError("no eval_spreadsheet_id in .env") # Setup MLFlow mlflow.set_tracking_uri(config.tracking_uri) mlflow.config.enable_async_logging(True) self._logger.info( f"MLflow tracking URI: {config.tracking_uri}, artifacts: {config.artifacts_uri}" ) # Initialize metrics self.metrics: List[BaseMetric] = [] for metric_name in config.metrics: kwargs = {} if "llm" in metric_name and llm: kwargs = dict(llm=self.llm, n_contexts=config.n_judge_contexts) self.metrics.append(MetricsRegistry.create(metric_name, **kwargs)) self._logger.info(f"Initialized metric: {metric_name}") @staticmethod def load_questions_from_sheet(*args, **kwargs) -> pd.DataFrame: """Load evaluation questions from spreadsheet""" df = load_spreadsheet(*args, **kwargs) df.fillna(dict(page=""), inplace=True) return df async def evaluate_single( self, output: Dict[str, Any], question: str, ground_truth: Dict ) -> Dict[str, MetricResult]: """Evaluate single search result against ground truth. Args: output: Dictionary with retrieval results including: - contexts: List of presentation results with metadata question: Original search query ground_truth: Dictionary with: - pres_name: Expected presentation name - pages: List of expected page numbers - question: Original question Returns: Dictionary mapping metric names to MetricResult objects """ # Log evaluation start self._logger.info(f"Evaluating question: '{question}'") results = {} # Calculate each metric for metric in self.metrics: try: result = await metric.acalculate(output, ground_truth) results[metric.name] = result # Log metric result log_msg = f"Metric {metric.name}: {result.score}" if result.explanation: log_msg += f" ({result.explanation[:200]})" self._logger.info(log_msg) except Exception as e: self._logger.error( f"Failed to calculate metric {metric.name}: {str(e)}" ) # Create failure result results[metric.name] = MetricResult( name=metric.name, score=0.0, explanation=f"Calculation failed: {str(e)}", ) return results async def process_question( self, retriever: RetrieverTypes, row: pd.Series, metric_values: Dict[str, List[float]], results_log: List[Dict], question_idx: int, total_questions: int, semaphore: asyncio.Semaphore, ) -> None: """Process single question with semaphore-controlled concurrency""" async with semaphore: self._logger.info( f"Processing question {question_idx+1}/{total_questions}: " f"{row['question'][:50]}..." ) ground_truth = { "question": row["question"], "pres_name": row["pres_name"], "pages": [int(x) for x in row["page"].split(",") if x], "content_type": row["content"], } try: # Retrieve asynchronously output = await retriever.aretrieve(query=row["question"]) # Evaluate results results = await self.evaluate_single( output=output, question=row["question"], ground_truth=ground_truth, ) # Update aggregated results result_row = { "question": row["question"], "expected_content_type": row["content"], "expected_presentation": row["pres_name"], "expected_pages": row["page"], "retrieved_presentations": [ p["pres_name"] for p in output["contexts"] ], "retrieved_pages": [ ",".join(map(str, p["pages"])) for p in output["contexts"] ], "best_chunk_type": output["contexts"][0]["best_chunk"]["chunk_type"], # fmt: skip } for metric_name, metric_result in results.items(): result_row[f"metric_{metric_name}_score"] = metric_result.score if metric_result.explanation: result_row[f"metric_{metric_name}_explanation"] = ( metric_result.explanation ) metric_values[metric_name].append(metric_result.score) results_log.append(result_row) except Exception as e: self._logger.error( f"Failed to process question {question_idx+1}: {str(e)}" ) async def process_questions_batch( self, retriever: RetrieverTypes, questions_df: pd.DataFrame, metric_values: Dict[str, List[float]], results_log: List[Dict], ) -> None: """Process questions with controlled concurrency""" # Create semaphore within the async context semaphore = asyncio.Semaphore(self._max_concurrent) # Create tasks for all questions tasks = [ self.process_question( retriever=retriever, row=row, metric_values=metric_values, results_log=results_log, question_idx=idx, total_questions=len(questions_df), semaphore=semaphore, ) for idx, (_, row) in enumerate(questions_df.iterrows()) ] # Wait for all tasks to complete await tqdm_asyncio.gather( *tasks, desc=f"Processing questions for '{retriever.scorer.id[:15]}' (max {self._max_concurrent} concurrent)", total=len(questions_df), dynamic_ncols=True, # Adjust width automatically ) def run_evaluation(self, questions_df: pd.DataFrame) -> None: """Run evaluation with async LLM queries and controlled concurrency""" timestamp = datetime.now().replace(microsecond=0).isoformat() self._logger.info(f"Starting evaluation with {len(questions_df)} questions") # MLflow setup experiment = mlflow.get_experiment_by_name(self.config.experiment_name) if experiment is not None: experiment_id = experiment.experiment_id self._logger.info( f"Using existing experiment: {self.config.experiment_name}" ) else: experiment_id = mlflow.create_experiment( self.config.experiment_name, artifact_location=self.config.artifacts_uri, ) experiment = mlflow.get_experiment_by_name(self.config.experiment_name) self._logger.info(f"Created new experiment: {self.config.experiment_name}") mlflow.set_experiment(experiment_id=experiment_id) for scorer in self.config.scorers: self._logger.info(f"Evaluating with scorer: {scorer.id}") # Initialize retriever retriever = self.config.get_retriever_with_scorer(scorer) # Get preprocessor id preprocessor_id = ( retriever.storage.query_preprocessor.id if retriever.storage.query_preprocessor else "None" ) with mlflow.start_run( run_name=f"scorer_{scorer.id}__retriever_{retriever.id}__preprocessor_{preprocessor_id}" ): # Log preprocessor mlflow.log_params({"preprocessing": preprocessor_id}) self._logger.info(f"Using preprocessor: {preprocessor_id}") # Log config parameters mlflow.log_params( {f"config_{k}": v for k, v in self.config.get_log_params().items()} ) self._logger.debug("Logged config parameters") # Log scorer parameters mlflow.log_params( {f"scorer_{k}": v for k, v in scorer.model_dump().items()} ) self._logger.debug("Logged scorer parameters") # Initialize retriever and log its parameters mlflow.log_params( {f"retriever_{k}": v for k, v in retriever.get_log_params().items()} ) self._logger.debug("Logged retriever parameters") # Initialize aggregation containers results_log = [] metric_values = {m.name: [] for m in self.metrics} # Process questions with async handling asyncio.run( self.process_questions_batch( retriever, questions_df, metric_values, results_log ) ) # Calculate n_errors n_errors = ( len(questions_df) - len(results_log) if results_log else len(questions_df) ) if n_errors > 1: logger.error(f"{n_errors} while processing {retriever.id}") # Process results results_df = pd.DataFrame(results_log) results_df["experiment_name"] = ( experiment.name if experiment is not None else "no_name" ) results_df["experiment_id"] = experiment_id results_df["scorer"] = scorer.id results_df["retriever"] = retriever.id results_df["timestamp"] = timestamp # Save results with NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f: results_df.to_csv(f.name, index=False) fpath = str("detailed_results") mlflow.log_artifact(f.name, fpath) self._logger.info(f"Saved detailed results to {fpath}") # Write to google sheets if enabled if self.config.write_to_google: self.gsheets.write_evaluation_results( results_df=results_df, metric_values=metric_values, experiment_name=self.config.experiment_name, ) # Log metrics for name, values in metric_values.items(): if values: mean_value = np.nanmean(values) mlflow.log_metric(f"mean_{name}", mean_value) mlflow.log_metric(f"n_questions", len(questions_df)) mlflow.log_metric(f"error_rate", n_errors / len(questions_df)) self._logger.info(f"Mean {name}: {mean_value:.3f}")