sivan22/Ituria
16
1import os2import logging3from typing import List, Dict, Any, Optional, Tuple, Callable, Union4from dotenv import load_dotenv5from llm_providers import LLMProvider6from langchain_core.messages import HumanMessage7from tantivy_search_agent import TantivySearchAgent8 9load_dotenv()10 11class SearchAgent:12 def __init__(self, tantivy_agent: TantivySearchAgent, provider_name: str = "Gemini", api_keys: Dict[str, str] = None):13 """Initialize the search agent with Tantivy agent and LLM client"""14 self.tantivy_agent = tantivy_agent15 self.logger = logging.getLogger(__name__)16 17 # Initialize LLM provider with API keys18 self.llm_provider = LLMProvider(api_keys)19 self.llm = None20 self.set_provider(provider_name)21 22 self.min_confidence_threshold = 0.723 24 def set_provider(self, provider_name: str) -> None:25 self.llm = self.llm_provider.get_provider(provider_name)26 if not self.llm:27 raise ValueError(f"Provider {provider_name} not available")28 self.current_provider = provider_name29 30 def get_available_providers(self) -> list[str]:31 return self.llm_provider.get_available_providers()32 33 def get_query(self, query: str, failed_queries: List[Dict[str, str]] = []) -> str:34 """Generate a Tantivy query using Claude, considering previously failed queries"""35 try:36 if not self.llm:37 raise ValueError("LLM provider not initialized")38 39 prompt = (40 "Create a query for this search request with the following restrictions:\n"+41 self.tantivy_agent.get_query_instructions()+ 42 "\n\nAdditional instructions: \n"43 "1. return only the search query without any other text\n"44 "2. Use only Hebrew terms for the search query\n"45 "3. the corpus to search in is an ancient Hebrew corpus - Tora and Talmud. so Try to use ancient Hebrew terms and or Talmudic expressions."46 "4. prevent modern words that are not common in talmudic texts \n" 47 f"the search request: {query}"48 )49 50 if failed_queries:51 prompt += (52 f"\n\nPrevious failed queries:\n"+53 "------------------------\n"+ 54 '\n'.join(f"Query: {q['query']}, Reason: {q['reason']}" for q in failed_queries)+55 "\n\n"56 "Please generate an alternative query that:\n"57 "1. Uses different Hebrew synonyms or related terms\n"58 "2. Tries broader or more general terms\n"59 "3. Adjusts proximity values or uses wildcards\n"60 "4. Prevents using modern words that are not common in ancient hebrew and talmud texts\n"61 )62 63 response = self.llm.invoke([HumanMessage(content=prompt)])64 tantivy_query = response.content.strip() 65 self.logger.info(f"Generated Tantivy query: {tantivy_query}")66 return tantivy_query67 68 except Exception as e:69 self.logger.error(f"Error generating query: {e}")70 # Fallback to basic quoted search71 return f'"{query}"'72 73 def _evaluate_results(self, results: List[Dict[str, Any]], query: str) -> Dict[str, Any]:74 """Evaluate search results using Claude with confidence scoring"""75 if not self.llm:76 raise ValueError("LLM provider not initialized")77 78 # Prepare context from results79 context = "\n".join(f"Result {i}. Source: {r.get('reference',[])}\n Text: {r.get('text', [])}"80 for i, r in enumerate(results)81 )82 83 try:84 message = self.llm.invoke([HumanMessage(content=f"""Evaluate the search results for answering this question:85 Question: {query}86 87 Search Results:88 {context}89 90 Provide evaluation in this format (3 lines):91 Confidence score (0.0 to 1.0) indicating how well the results can answer the question. this line should include only the number return, don't include '[line 1]'92 ACCEPT if score >= {self.min_confidence_threshold}, REFINE if score < {self.min_confidence_threshold}. return only the word ACCEPT or REFINE.93 Detailed explanation of what information is present or missing, don't include '[line 3]'. it should be only in Hebrew94 """)])95 lines = message.content.strip().replace('\n\n', '\n').split('\n')96 confidence = float(lines[0])97 decision = lines[1].upper()98 explanation = lines[2]99 100 is_good = decision == 'ACCEPT'101 102 self.logger.info(f"Evaluation: Confidence={confidence}, Decision={decision}")103 self.logger.info(f"Explanation: {explanation}")104 105 return {106 "confidence": confidence,107 "is_sufficient": is_good,108 "explanation": explanation,109 110 }111 112 except Exception as e:113 self.logger.error(f"Error evaluating results: {e}")114 # Fallback to simple evaluation115 return {116 "confidence": 0.0,117 "is_sufficient": False,118 "explanation": "", 119 }120 121 def _generate_answer(self, query: str, results: List[Dict[str, Any]]) -> str:122 """Generate answer using Claude with improved context utilization"""123 if not self.llm:124 raise ValueError("LLM provider not initialized")125 126 if not results:127 return "לא נמצאו תוצאות"128 129 # Prepare context from results130 context = "\n".join(f"Result {i+1}. Source: {r.get('reference',[])}\n Text: {r.get('text', [])}"131 for i, r in enumerate(results)132 )133 134 try:135 message = self.llm.invoke([HumanMessage(content=f"""Based on these search results, answer this question:136 Question: {query}137 138 Search Results:139 {context}140 141 Requirements for your answer:142 1. Use only information from the search results143 2. Be comprehensive but concise144 3. Structure the answer clearly145 4. If any aspect of the question cannot be fully answered, acknowledge this146 5. cite sources for each fact or information you use147 6. The answer should be only in Hebrew148 """)])149 return message.content.strip()150 151 except Exception as e:152 self.logger.error(f"Error generating answer: {e}")153 return f"I encountered an error generating the answer: {str(e)}"154 155 def search_and_answer(self, query: str, num_results: int = 10, max_iterations: int = 3, 156 on_step: Optional[Callable[[Dict[str, Any]], None]] = None) -> Dict[str, Any]:157 """Execute multi-step search process using Tantivy with streaming updates"""158 steps = []159 all_results = []160 161 # Step 1: Generate Tantivy query162 initial_query = self.get_query(query)163 step = {164 'action': 'יצירת שאילתת חיפוש',165 'description': 'נוצרה שאילתת חיפוש עבור מנוע החיפוש',166 'results': [{'type': 'query', 'content': initial_query}]167 }168 steps.append(step)169 if on_step:170 on_step(step)171 172 # Step 2: Initial search with Tantivy query173 results = self.tantivy_agent.search(initial_query, num_results)174 175 step = {176 'action': 'חיפוש במאגר',177 'description': f'חיפוש במאגר עבור שאילתת חיפוש: {initial_query}',178 'results': [{'type': 'document', 'content': {179 'title': r['title'],180 'reference': r['reference'],181 'topics': r['topics'],182 'highlights': r['highlights'],183 'score': r['score']184 }} for r in results]185 }186 steps.append(step)187 if on_step:188 on_step(step)189 190 failed_queries = []191 192 if results.__len__() == 0:193 failed_queries.append({'query': initial_query, 'reason': 'no results'})194 is_sufficient = False195 else: 196 all_results.extend(results)197 198 # Step 3: Evaluate results199 evaluation = self._evaluate_results(results, query)200 confidence = evaluation['confidence']201 is_sufficient = evaluation['is_sufficient']202 explanation = evaluation['explanation']203 204 step = {205 'action': 'דירוג תוצאות',206 'description': 'דירוג תוצאות חיפוש',207 'results': [{208 'type': 'evaluation',209 'content': {210 'status': 'accepted' if is_sufficient else 'insufficient',211 'confidence': confidence,212 'explanation': explanation,213 }214 }]215 }216 steps.append(step)217 if on_step:218 on_step(step)219 220 if not is_sufficient:221 failed_queries.append({'query': initial_query, 'reason': explanation})222 223 # Step 4: Additional searches if needed224 attempt = 2225 while not is_sufficient and attempt < max_iterations:226 # Generate new query227 new_query = self.get_query(query, failed_queries)228 229 step = {230 'action': f'יצירת שאילתה מחדש (ניסיון {attempt})',231 'description': 'נוצרה שאילתת חיפוש נוספת עבור מנוע החיפוש',232 'results': [233 {'type': 'new_query', 'content': new_query}234 ]235 }236 steps.append(step)237 if on_step:238 on_step(step)239 240 # Search with new query241 results = self.tantivy_agent.search(new_query, num_results)242 243 step = {244 'action': f'חיפוש נוסף (ניסיון {attempt}) ',245 'description': f'מחפש במאגר עבור שאילתת חיפוש: {new_query}',246 'results': [{'type': 'document', 'content': {247 'title': r['title'],248 'reference': r['reference'],249 'topics': r['topics'],250 'highlights': r['highlights'],251 'score': r['score']252 }} for r in results]253 }254 steps.append(step)255 if on_step:256 on_step(step)257 258 if results.__len__() == 0:259 failed_queries.append({'query': new_query, 'reason': 'no results'})260 261 else:262 all_results.extend(results)263 264 # Re-evaluate with current results265 evaluation = self._evaluate_results(results, query)266 confidence = evaluation['confidence']267 is_sufficient = evaluation['is_sufficient']268 explanation = evaluation['explanation']269 270 step = {271 'action': f'דירוג תוצאות (ניסיון {attempt})',272 'description': 'דירוג תוצאות חיפוש לניסיון זה',273 'explanation': explanation,274 'results': [{275 'type': 'evaluation',276 'content': {277 'status': 'accepted' if is_sufficient else 'insufficient',278 'confidence': confidence,279 'explanation': explanation,280 }281 }]282 }283 steps.append(step)284 if on_step:285 on_step(step)286 287 if not is_sufficient:288 failed_queries.append({'query': new_query, 'reason': explanation})289 290 attempt += 1291 292 # Step 5: Generate final answer293 answer = self._generate_answer(query, all_results)294 295 final_result = {296 'steps': steps,297 'answer': answer,298 'sources': [{299 'title': r['title'],300 'reference': r['reference'],301 'topics': r['topics'],302 'path': r['file_path'],303 'highlights': r['highlights'],304 'text': r['text'],305 'score': r['score']306 } for r in all_results]307 }308 309 # Send final result through callback310 if on_step:311 on_step({312 'action': 'סיום',313 'description': 'החיפוש הושלם',314 'final_result': final_result315 })316 317 return final_result318