""" Answer generation module using Ollama for local LLM inference. This module provides answer generation with citation support for RAG systems, optimized for technical documentation Q&A on Apple Silicon. """ import json import logging from dataclasses import dataclass from typing import List, Dict, Any, Optional, Generator, Tuple import ollama from datetime import datetime import re from pathlib import Path import sys # Import calibration framework try: from src.confidence_calibration import ConfidenceCalibrator except ImportError: # Fallback - disable calibration for deployment ConfidenceCalibrator = None logger = logging.getLogger(__name__) @dataclass class Citation: """Represents a citation to a source document chunk.""" chunk_id: str page_number: int source_file: str relevance_score: float text_snippet: str @dataclass class GeneratedAnswer: """Represents a generated answer with citations.""" answer: str citations: List[Citation] confidence_score: float generation_time: float model_used: str context_used: List[Dict[str, Any]] class AnswerGenerator: """ Generates answers using local LLMs via Ollama with citation support. Optimized for technical documentation Q&A with: - Streaming response support - Citation extraction and formatting - Confidence scoring - Fallback model support """ def __init__( self, primary_model: str = "llama3.2:3b", fallback_model: str = "mistral:latest", temperature: float = 0.3, max_tokens: int = 1024, stream: bool = True, enable_calibration: bool = True ): """ Initialize the answer generator. Args: primary_model: Primary Ollama model to use fallback_model: Fallback model for complex queries temperature: Generation temperature (0.0-1.0) max_tokens: Maximum tokens to generate stream: Whether to stream responses enable_calibration: Whether to enable confidence calibration """ self.primary_model = primary_model self.fallback_model = fallback_model self.temperature = temperature self.max_tokens = max_tokens self.stream = stream self.client = ollama.Client() # Initialize confidence calibration self.enable_calibration = enable_calibration self.calibrator = None if enable_calibration and ConfidenceCalibrator is not None: try: self.calibrator = ConfidenceCalibrator() logger.info("Confidence calibration enabled") except Exception as e: logger.warning(f"Failed to initialize calibration: {e}") self.enable_calibration = False elif enable_calibration and ConfidenceCalibrator is None: logger.warning("Calibration requested but ConfidenceCalibrator not available - disabling") self.enable_calibration = False # Verify models are available self._verify_models() def _verify_models(self) -> None: """Verify that required models are available.""" try: model_list = self.client.list() available_models = [] # Handle Ollama's ListResponse object if hasattr(model_list, 'models'): for model in model_list.models: if hasattr(model, 'model'): available_models.append(model.model) elif isinstance(model, dict) and 'model' in model: available_models.append(model['model']) if self.primary_model not in available_models: logger.warning(f"Primary model {self.primary_model} not found. Available models: {available_models}") raise ValueError(f"Model {self.primary_model} not available. Please run: ollama pull {self.primary_model}") if self.fallback_model not in available_models: logger.warning(f"Fallback model {self.fallback_model} not found in: {available_models}") except Exception as e: logger.error(f"Error verifying models: {e}") raise def _create_system_prompt(self) -> str: """Create system prompt for technical documentation Q&A.""" return """You are a technical documentation assistant that provides clear, accurate answers based on the provided context. CORE PRINCIPLES: 1. ANSWER DIRECTLY: If context contains the answer, provide it clearly and confidently 2. BE CONCISE: Keep responses focused and avoid unnecessary uncertainty language 3. CITE ACCURATELY: Use [chunk_X] citations for every fact from context RESPONSE GUIDELINES: - If context has sufficient information → Answer directly and confidently - If context has partial information → Answer what's available, note what's missing briefly - If context is irrelevant → Brief refusal: "This information isn't available in the provided documents" CITATION FORMAT: - Use [chunk_1], [chunk_2] etc. for all facts from context - Example: "According to [chunk_1], RISC-V is an open-source architecture." WHAT TO AVOID: - Do NOT add details not in context - Do NOT second-guess yourself if context is clear - Do NOT use phrases like "does not contain sufficient information" when context clearly answers the question - Do NOT be overly cautious when context is adequate Be direct, confident, and accurate. If the context answers the question, provide that answer clearly.""" def _format_context(self, chunks: List[Dict[str, Any]]) -> str: """ Format retrieved chunks into context for the LLM. Args: chunks: List of retrieved chunks with metadata Returns: Formatted context string """ context_parts = [] for i, chunk in enumerate(chunks): chunk_text = chunk.get('content', chunk.get('text', '')) page_num = chunk.get('metadata', {}).get('page_number', 'unknown') source = chunk.get('metadata', {}).get('source', 'unknown') context_parts.append( f"[chunk_{i+1}] (Page {page_num} from {source}):\n{chunk_text}\n" ) return "\n---\n".join(context_parts) def _extract_citations(self, answer: str, chunks: List[Dict[str, Any]]) -> Tuple[str, List[Citation]]: """ Extract citations from the generated answer and integrate them naturally. Args: answer: Generated answer with [chunk_X] citations chunks: Original chunks used for context Returns: Tuple of (natural_answer, citations) """ citations = [] citation_pattern = r'\[chunk_(\d+)\]' cited_chunks = set() # Find [chunk_X] citations and collect cited chunks matches = re.finditer(citation_pattern, answer) for match in matches: chunk_idx = int(match.group(1)) - 1 # Convert to 0-based index if 0 <= chunk_idx < len(chunks): cited_chunks.add(chunk_idx) # Create Citation objects for each cited chunk chunk_to_source = {} for idx in cited_chunks: chunk = chunks[idx] citation = Citation( chunk_id=chunk.get('id', f'chunk_{idx}'), page_number=chunk.get('metadata', {}).get('page_number', 0), source_file=chunk.get('metadata', {}).get('source', 'unknown'), relevance_score=chunk.get('score', 0.0), text_snippet=chunk.get('content', chunk.get('text', ''))[:200] + '...' ) citations.append(citation) # Map chunk reference to natural source name source_name = chunk.get('metadata', {}).get('source', 'unknown') if source_name != 'unknown': # Use just the filename without extension for natural reference natural_name = Path(source_name).stem.replace('-', ' ').replace('_', ' ') chunk_to_source[f'[chunk_{idx+1}]'] = f"the {natural_name} documentation" else: chunk_to_source[f'[chunk_{idx+1}]'] = "the documentation" # Replace [chunk_X] with natural references instead of removing them natural_answer = answer for chunk_ref, natural_ref in chunk_to_source.items(): natural_answer = natural_answer.replace(chunk_ref, natural_ref) # Clean up any remaining unreferenced citations (fallback) natural_answer = re.sub(r'\[chunk_\d+\]', 'the documentation', natural_answer) # Clean up multiple spaces and formatting natural_answer = re.sub(r'\s+', ' ', natural_answer).strip() return natural_answer, citations def _calculate_confidence(self, answer: str, citations: List[Citation], chunks: List[Dict[str, Any]]) -> float: """ Calculate confidence score for the generated answer with improved calibration. Args: answer: Generated answer citations: Extracted citations chunks: Retrieved chunks Returns: Confidence score (0.0-1.0) """ # Check if no chunks were provided first if not chunks: return 0.05 # No context = very low confidence # Assess context quality to determine base confidence scores = [chunk.get('score', 0) for chunk in chunks] max_relevance = max(scores) if scores else 0 avg_relevance = sum(scores) / len(scores) if scores else 0 # Dynamic base confidence based on context quality if max_relevance >= 0.8: confidence = 0.6 # High-quality context starts high elif max_relevance >= 0.6: confidence = 0.4 # Good context starts moderately elif max_relevance >= 0.4: confidence = 0.2 # Fair context starts low else: confidence = 0.05 # Poor context starts very low # Strong uncertainty and explicit refusal indicators strong_uncertainty_phrases = [ "does not contain sufficient information", "context does not provide", "insufficient information", "cannot determine", "refuse to answer", "cannot answer", "does not contain relevant", "no relevant context", "missing from the provided context" ] # Weak uncertainty phrases that might be in nuanced but correct answers weak_uncertainty_phrases = [ "unclear", "conflicting", "not specified", "questionable", "not contained", "no mention", "no relevant", "missing", "not explicitly" ] # Check for strong uncertainty - these should drastically reduce confidence if any(phrase in answer.lower() for phrase in strong_uncertainty_phrases): return min(0.1, confidence * 0.2) # Max 10% for explicit refusal/uncertainty # Check for weak uncertainty - reduce but don't destroy confidence for good context weak_uncertainty_count = sum(1 for phrase in weak_uncertainty_phrases if phrase in answer.lower()) if weak_uncertainty_count > 0: if max_relevance >= 0.7 and citations: # Good context with citations - reduce less severely confidence *= (0.8 ** weak_uncertainty_count) # Moderate penalty else: # Poor context - reduce more severely confidence *= (0.5 ** weak_uncertainty_count) # Strong penalty # If all chunks have very low relevance scores, cap confidence low if max_relevance < 0.4: return min(0.08, confidence) # Max 8% for low relevance context # Factor 1: Citation quality and coverage if citations and chunks: citation_ratio = len(citations) / min(len(chunks), 3) # Strong boost for high-relevance citations relevant_chunks = [c for c in chunks if c.get('score', 0) > 0.6] if relevant_chunks: # Significant boost for citing relevant chunks confidence += 0.25 * citation_ratio # Extra boost if citing majority of relevant chunks if len(citations) >= len(relevant_chunks) * 0.5: confidence += 0.15 else: # Small boost for citations to lower-relevance chunks confidence += 0.1 * citation_ratio else: # No citations = reduce confidence unless it's a simple factual statement if max_relevance >= 0.8 and len(answer.split()) < 20: confidence *= 0.8 # Gentle penalty for uncited but simple answers else: confidence *= 0.6 # Stronger penalty for complex uncited answers # Factor 2: Relevance score reinforcement if citations: avg_citation_relevance = sum(c.relevance_score for c in citations) / len(citations) if avg_citation_relevance > 0.8: confidence += 0.2 # Strong boost for highly relevant citations elif avg_citation_relevance > 0.6: confidence += 0.1 # Moderate boost elif avg_citation_relevance < 0.4: confidence *= 0.6 # Penalty for low-relevance citations # Factor 3: Context utilization quality if chunks: avg_chunk_length = sum(len(chunk.get('content', chunk.get('text', ''))) for chunk in chunks) / len(chunks) # Boost for substantial, high-quality context if avg_chunk_length > 200 and max_relevance > 0.8: confidence += 0.1 elif avg_chunk_length < 50: # Very short chunks confidence *= 0.8 # Factor 4: Answer characteristics answer_words = len(answer.split()) if answer_words < 10: confidence *= 0.9 # Slight penalty for very short answers elif answer_words > 50 and citations: confidence += 0.05 # Small boost for detailed cited answers # Factor 5: High-quality scenario bonus if (max_relevance >= 0.8 and citations and len(citations) > 0 and not any(phrase in answer.lower() for phrase in strong_uncertainty_phrases)): # This is a high-quality response scenario confidence += 0.15 raw_confidence = min(confidence, 0.95) # Cap at 95% to maintain some uncertainty # Apply temperature scaling calibration if available if self.enable_calibration and self.calibrator and self.calibrator.is_fitted: try: calibrated_confidence = self.calibrator.calibrate_confidence(raw_confidence) logger.debug(f"Confidence calibrated: {raw_confidence:.3f} -> {calibrated_confidence:.3f}") return calibrated_confidence except Exception as e: logger.warning(f"Calibration failed, using raw confidence: {e}") return raw_confidence def fit_calibration(self, validation_data: List[Dict[str, Any]]) -> float: """ Fit temperature scaling calibration using validation data. Args: validation_data: List of dicts with 'confidence' and 'correctness' keys Returns: Optimal temperature parameter """ if not self.enable_calibration or not self.calibrator: logger.warning("Calibration not enabled or not available") return 1.0 try: confidences = [item['confidence'] for item in validation_data] correctness = [item['correctness'] for item in validation_data] optimal_temp = self.calibrator.fit_temperature_scaling(confidences, correctness) logger.info(f"Calibration fitted with temperature: {optimal_temp:.3f}") return optimal_temp except Exception as e: logger.error(f"Failed to fit calibration: {e}") return 1.0 def save_calibration(self, filepath: str) -> bool: """Save fitted calibration to file.""" if not self.calibrator or not self.calibrator.is_fitted: logger.warning("No fitted calibration to save") return False try: calibration_data = { 'temperature': self.calibrator.temperature, 'is_fitted': self.calibrator.is_fitted, 'model_info': { 'primary_model': self.primary_model, 'fallback_model': self.fallback_model } } with open(filepath, 'w') as f: json.dump(calibration_data, f, indent=2) logger.info(f"Calibration saved to {filepath}") return True except Exception as e: logger.error(f"Failed to save calibration: {e}") return False def load_calibration(self, filepath: str) -> bool: """Load fitted calibration from file.""" if not self.enable_calibration or not self.calibrator: logger.warning("Calibration not enabled") return False try: with open(filepath, 'r') as f: calibration_data = json.load(f) self.calibrator.temperature = calibration_data['temperature'] self.calibrator.is_fitted = calibration_data['is_fitted'] logger.info(f"Calibration loaded from {filepath} (temp: {self.calibrator.temperature:.3f})") return True except Exception as e: logger.error(f"Failed to load calibration: {e}") return False def generate( self, query: str, chunks: List[Dict[str, Any]], use_fallback: bool = False ) -> GeneratedAnswer: """ Generate an answer based on the query and retrieved chunks. Args: query: User's question chunks: Retrieved document chunks use_fallback: Whether to use fallback model Returns: GeneratedAnswer object with answer, citations, and metadata """ start_time = datetime.now() model = self.fallback_model if use_fallback else self.primary_model # Check for no-context or very poor context situation if not chunks or all(len(chunk.get('content', chunk.get('text', ''))) < 20 for chunk in chunks): # Handle no-context situation with brief, professional refusal user_prompt = f"""Context: [NO RELEVANT CONTEXT FOUND] Question: {query} INSTRUCTION: Respond with exactly this brief message: "This information isn't available in the provided documents." DO NOT elaborate, explain, or add any other information.""" else: # Format context from chunks context = self._format_context(chunks) # Create concise prompt for faster generation user_prompt = f"""Context: {context} Question: {query} Instructions: Answer using only the context above. Cite with [chunk_1], [chunk_2] etc. Answer:""" try: # Generate response response = self.client.chat( model=model, messages=[ {"role": "system", "content": self._create_system_prompt()}, {"role": "user", "content": user_prompt} ], options={ "temperature": self.temperature, "num_predict": min(self.max_tokens, 300), # Reduce max tokens for speed "top_k": 40, # Optimize sampling for speed "top_p": 0.9, "repeat_penalty": 1.1 }, stream=False # Get complete response for processing ) # Extract answer answer_with_citations = response['message']['content'] # Extract and clean citations clean_answer, citations = self._extract_citations(answer_with_citations, chunks) # Calculate confidence confidence = self._calculate_confidence(clean_answer, citations, chunks) # Calculate generation time generation_time = (datetime.now() - start_time).total_seconds() return GeneratedAnswer( answer=clean_answer, citations=citations, confidence_score=confidence, generation_time=generation_time, model_used=model, context_used=chunks ) except Exception as e: logger.error(f"Error generating answer: {e}") # Return a fallback response return GeneratedAnswer( answer="I apologize, but I encountered an error while generating the answer. Please try again.", citations=[], confidence_score=0.0, generation_time=0.0, model_used=model, context_used=chunks ) def generate_stream( self, query: str, chunks: List[Dict[str, Any]], use_fallback: bool = False ) -> Generator[str, None, GeneratedAnswer]: """ Generate an answer with streaming support. Args: query: User's question chunks: Retrieved document chunks use_fallback: Whether to use fallback model Yields: Partial answer strings Returns: Final GeneratedAnswer object """ start_time = datetime.now() model = self.fallback_model if use_fallback else self.primary_model # Check for no-context or very poor context situation if not chunks or all(len(chunk.get('content', chunk.get('text', ''))) < 20 for chunk in chunks): # Handle no-context situation with brief, professional refusal user_prompt = f"""Context: [NO RELEVANT CONTEXT FOUND] Question: {query} INSTRUCTION: Respond with exactly this brief message: "This information isn't available in the provided documents." DO NOT elaborate, explain, or add any other information.""" else: # Format context from chunks context = self._format_context(chunks) # Create concise prompt for faster generation user_prompt = f"""Context: {context} Question: {query} Instructions: Answer using only the context above. Cite with [chunk_1], [chunk_2] etc. Answer:""" try: # Generate streaming response stream = self.client.chat( model=model, messages=[ {"role": "system", "content": self._create_system_prompt()}, {"role": "user", "content": user_prompt} ], options={ "temperature": self.temperature, "num_predict": min(self.max_tokens, 300), # Reduce max tokens for speed "top_k": 40, # Optimize sampling for speed "top_p": 0.9, "repeat_penalty": 1.1 }, stream=True ) # Collect full answer while streaming full_answer = "" for chunk in stream: if 'message' in chunk and 'content' in chunk['message']: partial = chunk['message']['content'] full_answer += partial yield partial # Process complete answer clean_answer, citations = self._extract_citations(full_answer, chunks) confidence = self._calculate_confidence(clean_answer, citations, chunks) generation_time = (datetime.now() - start_time).total_seconds() return GeneratedAnswer( answer=clean_answer, citations=citations, confidence_score=confidence, generation_time=generation_time, model_used=model, context_used=chunks ) except Exception as e: logger.error(f"Error in streaming generation: {e}") yield "I apologize, but I encountered an error while generating the answer." return GeneratedAnswer( answer="Error occurred during generation.", citations=[], confidence_score=0.0, generation_time=0.0, model_used=model, context_used=chunks ) def format_answer_with_citations(self, generated_answer: GeneratedAnswer) -> str: """ Format the generated answer with citations for display. Args: generated_answer: GeneratedAnswer object Returns: Formatted string with answer and citations """ formatted = f"{generated_answer.answer}\n\n" if generated_answer.citations: formatted += "**Sources:**\n" for i, citation in enumerate(generated_answer.citations, 1): formatted += f"{i}. {citation.source_file} (Page {citation.page_number})\n" formatted += f"\n*Confidence: {generated_answer.confidence_score:.1%} | " formatted += f"Model: {generated_answer.model_used} | " formatted += f"Time: {generated_answer.generation_time:.2f}s*" return formatted if __name__ == "__main__": # Example usage generator = AnswerGenerator() # Example chunks (would come from retrieval system) example_chunks = [ { "id": "chunk_1", "content": "RISC-V is an open-source instruction set architecture (ISA) based on reduced instruction set computer (RISC) principles.", "metadata": {"page_number": 1, "source": "riscv-spec.pdf"}, "score": 0.95 }, { "id": "chunk_2", "content": "The RISC-V ISA is designed to support a wide range of implementations including 32-bit, 64-bit, and 128-bit variants.", "metadata": {"page_number": 2, "source": "riscv-spec.pdf"}, "score": 0.89 } ] # Generate answer result = generator.generate( query="What is RISC-V?", chunks=example_chunks ) # Display formatted result print(generator.format_answer_with_citations(result))