""" LLM reasoning layer for answering questions with citations. Ensures all responses are grounded in retrieved context. """ import os import re from typing import List, Dict, Optional from src.models import QueryRequest, AgentResponse, Citation, ContextUnit from src.retrieval import RetrievalResult, ContextBuilder from src.groq_integration import GroqReasoningEngine class CitationExtractor: """Extract cell references from LLM responses.""" @staticmethod def extract_citations(response_text: str, retrieved_units: List[ContextUnit]) -> List[Citation]: """ Extract cell citations from LLM response. Looks for patterns like "Cell X", "cell_X", "(X)", etc. """ citations = [] cell_ids = {u.cell.cell_id for u in retrieved_units} # Pattern 1: "Cell X" or "cell X" pattern1 = r'[Cc]ell\s+([a-zA-Z_0-9]+)' matches1 = re.findall(pattern1, response_text) # Pattern 2: "(cell_X)" or similar pattern2 = r'\(([a-zA-Z_0-9]+)\)' matches2 = re.findall(pattern2, response_text) # Combine and deduplicate potential_cells = set(matches1 + matches2) # Validate against retrieved cells for cell_id in potential_cells: if cell_id in cell_ids: # Find the unit unit = next((u for u in retrieved_units if u.cell.cell_id == cell_id), None) if unit: citation = Citation( cell_id=cell_id, cell_type=unit.cell.cell_type, content_snippet=CitationExtractor._get_snippet(unit), intent=unit.intent if unit.intent != "[Pending intent inference]" else None ) citations.append(citation) return citations @staticmethod def _get_snippet(unit: ContextUnit, max_length: int = 100) -> str: """Get content snippet from unit.""" return unit.cell.source[:max_length] class HallucinationDetector: """Detect potential hallucinations in responses.""" @staticmethod def check_for_unsupported_claims(response: str, context: str) -> bool: """Check if response makes claims not supported by context.""" # Simplified check - in real implementation, use more sophisticated methods response_lower = response.lower() context_lower = context.lower() # Check for common hallucination indicators hallucination_indicators = [ "according to", "experts say", "research shows", "it's known that", "generally", "typically" ] for indicator in hallucination_indicators: if indicator in response_lower and indicator not in context_lower: return True return False class ReasoningEngine: """LLM reasoning engine for answering questions.""" def __init__(self): self.groq_client = self._init_groq() self.openai_client = self._init_openai() def _init_groq(self): """Initialize Groq client (preferred for speed and cost).""" try: return GroqReasoningEngine() except Exception: return None def _init_openai(self): """Initialize OpenAI client (fallback).""" try: api_key = os.getenv("OPENAI_API_KEY") if api_key and not api_key.startswith("sk-placeholder"): from openai import OpenAI return OpenAI(api_key=api_key) except Exception: pass return None def reason(self, query: str, retrieval_result: RetrievalResult, conversation_history: Optional[List[Dict]] = None) -> AgentResponse: """ Reason about a question given retrieved context. Returns response with citations. """ # Build context for LLM context = ContextBuilder.build_context_for_llm( retrieval_result.units, query ) # Try to use Groq first (fast and free) if self.groq_client: try: groq_result = self.groq_client.reason_with_context(query, context, conversation_history=conversation_history) answer = groq_result["answer"] except Exception as e: print(f"Groq query failed: {e}. Using OpenAI fallback.") answer = self._query_openai(query, context) if self.openai_client else self._generate_answer_fallback(query, retrieval_result) elif self.openai_client: try: answer = self._query_openai(query, context) except Exception as e: print(f"OpenAI query failed: {e}. Using fallback.") answer = self._generate_answer_fallback(query, retrieval_result) else: # Use fallback reasoning answer = self._generate_answer_fallback(query, retrieval_result) # Extract citations citations = CitationExtractor.extract_citations(answer, retrieval_result.units) # Check for hallucination risk has_hallucination_risk = HallucinationDetector.check_for_unsupported_claims( answer, context ) # Calculate confidence confidence = self._calculate_confidence( len(citations), len(retrieval_result.units), ) return AgentResponse( answer=answer, citations=citations, confidence=confidence, has_hallucination_risk=has_hallucination_risk, retrieved_units=retrieval_result.units ) def _query_openai(self, query: str, context: str) -> str: """Query OpenAI API.""" if not self.openai_client: raise Exception("OpenAI client not available") prompt = f""" Based on the following notebook context, answer the question. Cite specific cells when referencing information. Context: {context} Question: {query} Answer:""" response = self.openai_client.chat.completions.create( model="gpt-4", messages=[{"role": "user", "content": prompt}], max_tokens=500 ) return response.choices[0].message.content def _generate_answer_fallback(self, query: str, retrieval_result: RetrievalResult) -> str: """Generate answer using simple fallback logic.""" query_lower = query.lower() # Handle specific question types # Match: "what is this notebook about", "whats the notebook about", "what's this about", etc. if any(phrase in query_lower for phrase in [ "what is this notebook about", "what does this notebook", "whats the notebook", "what's this", "what about this notebook", "what is this about", "describe this notebook", "tell me about this" ]): return self._summarize_notebook(retrieval_result) if "why" in query_lower: return self._explain_decision(retrieval_result, query) # Default: find relevant code snippets and summarize all units if retrieval_result.units: return self._summarize_notebook(retrieval_result) return "I couldn't find specific information about that in the notebook context." def _summarize_notebook(self, retrieval_result: RetrievalResult) -> str: """Generate a comprehensive summary of what the notebook is about.""" data_sources = [] data_operations = [] models = [] metrics = [] visualizations = [] code_cells = 0 markdown_cells = 0 for unit in retrieval_result.units: intent = unit.intent.lower() if unit.intent else "" source = unit.cell.source.lower() if unit.cell.cell_type == "code": code_cells += 1 elif unit.cell.cell_type == "markdown": markdown_cells += 1 # Data loading/sources if "load data" in intent or "read" in source or "dataset" in source: if "iris" in source: data_sources.append("Iris dataset") elif "csv" in source or "pd.read_csv" in source: data_sources.append("CSV data files") elif "excel" in source or "xlsx" in source: data_sources.append("Excel spreadsheets") else: data_sources.append("external datasets") # Data operations if "preprocess" in intent or "clean" in source or "drop" in source: data_operations.append("data cleaning and preprocessing") if "filter" in source or "select" in source: data_operations.append("data filtering") if "merge" in source or "join" in source: data_operations.append("data merging") # Models if "model" in intent or "fit" in source or "train" in source: if "randomforest" in source: models.append("Random Forest classifier") elif "regression" in source: models.append("regression model") elif "neural" in source or "nn" in source: models.append("neural network") else: models.append("machine learning model") # Evaluation metrics if "accuracy" in source or "precision" in source or "recall" in source or "f1" in source: metrics.append("classification metrics") if "rmse" in source or "mse" in source: metrics.append("regression metrics") if "auc" in source or "roc" in source: metrics.append("ROC/AUC analysis") # Visualizations if "visualize" in intent or "plot" in source or "matplotlib" in source or "seaborn" in source: if "scatter" in source: visualizations.append("scatter plots") elif "hist" in source: visualizations.append("histograms") elif "bar" in source: visualizations.append("bar charts") else: visualizations.append("data visualizations") # Build comprehensive summary summary = [] # Main purpose if data_sources and models: summary.append(f"This is a machine learning notebook that analyzes {', '.join(set(data_sources))}") elif data_sources: summary.append(f"This notebook analyzes {', '.join(set(data_sources))}") elif models: summary.append("This notebook demonstrates machine learning model development and evaluation") else: summary.append("This is a data analysis notebook") # Data operations if data_operations: summary.append(f"It includes {', '.join(set(data_operations))}") # Models and evaluation if models or metrics: model_desc = f"Uses {', '.join(set(models))}" if models else "Includes model training" if metrics: model_desc += f" with {', '.join(set(metrics))}" summary.append(model_desc) # Visualizations if visualizations: summary.append(f"Includes {', '.join(set(visualizations))} for data exploration and results visualization") # Notebook structure total_cells = code_cells + markdown_cells if total_cells > 0: summary.append(f"\n**Notebook Structure:** {code_cells} code cells, {markdown_cells} documentation cells") return ". ".join(summary) + "." def _explain_decision(self, retrieval_result: RetrievalResult, query: str) -> str: """Explain why certain decisions were made.""" query_lower = query.lower() # Look for common decisions if "remove" in query_lower or "drop" in query_lower: for unit in retrieval_result.units: if "drop" in unit.cell.source.lower() or "remove" in unit.cell.source.lower(): return f"Data was removed/cleaned as shown in: {unit.cell.source[:150]}" return "The notebook shows standard data preprocessing and modeling steps." def _calculate_confidence(self, num_citations: int, num_units: int) -> float: """Calculate confidence score.""" if num_units == 0: return 0.0 # If we have units but no explicit citations, give baseline confidence (0.7) if num_citations == 0 and num_units > 0: return 0.7 # With citations, confidence increases base_confidence = min(num_citations / max(num_units, 1), 1.0) return max(base_confidence, 0.7) class ContextualAnsweringSystem: """End-to-end system for context-aware question answering.""" def __init__(self, retrieval_engine, use_llm: bool = True): self.retrieval_engine = retrieval_engine self.reasoning_engine = ReasoningEngine() self.use_llm = use_llm def answer_question(self, query: str, top_k: int = 5, conversation_history: Optional[List[Dict]] = None) -> AgentResponse: """ Answer a question about the notebook context. Args: query: User's natural language question top_k: Number of cells to retrieve for context conversation_history: Previous conversation for context Returns: AgentResponse with answer, citations, and context """ # Step 1: Retrieve relevant context retrieval_result = self.retrieval_engine.retrieve(query, top_k=top_k) # Step 2: Reason and generate answer with conversation context response = self.reasoning_engine.reason(query, retrieval_result, conversation_history) return response