| |
|
|
| |
| import sys |
| try: |
| __import__('pysqlite3') |
| sys.modules['sqlite3'] = sys.modules.pop('pysqlite3') |
| except ImportError: |
| pass |
| import os |
| os.environ["ANONYMIZED_TELEMETRY"] = "False" |
| os.environ["CHROMA_TELEMETRY"] = "false" |
| |
|
|
| import re |
| import json |
| import time |
| import hashlib |
| from datetime import datetime, timezone |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from functools import wraps |
|
|
| import requests |
| import dirtyjson |
| import chromadb |
| from chromadb.config import Settings |
| from sentence_transformers import CrossEncoder |
| from langchain_mistralai.chat_models import ChatMistralAI |
| from langchain_core.messages import HumanMessage |
| from markdown_it import MarkdownIt |
| from rank_bm25 import BM25Okapi |
| import numpy as np |
| import markdown |
| import trafilatura |
| from bs4 import BeautifulSoup |
| import spacy |
|
|
| from config import ( |
| MISTRAL_MODEL, MAX_RETRIES, DEBUG, MAX_CONCURRENT_CRAWLERS, |
| TOP_N_URLS_TO_PROCESS, CRAWL_CACHE_DIR, SPACY_MODEL, |
| VECTOR_DB_PATH, CROSS_ENCODER_MODEL, RAG_CANDIDATE_POOL_SIZE, |
| RAG_FINAL_EVIDENCE_COUNT, SUPABASE_CONFIGURED, |
| OPENROUTER_API_KEY, PERPLEXITY_SONAR_MODEL, QWEN_EMBEDDING_MODEL, LLAMA_VERIFIER_MODEL, |
| MISTRAL_API_KEY, SEARCH_API_KEY, CSE_ID, SONAR_ANALYTICS_DIR |
| ) |
| from schemas import ( |
| DEFAULT_BLANK_FIELDS, CHOICE_OPTIONS, INFERABLE_FIELDS, BLACKLISTED_DOMAINS |
| ) |
| |
| from supabase_client import get_knowledge_cache, set_knowledge_cache, download_from_storage, upload_to_storage, fetch_races_db_fields |
|
|
| class AgentInitializationError(Exception): pass |
|
|
| def clean_markdown_with_trafilatura(md_text: str) -> str: |
| try: |
| md_text = re.sub(r'!\[.*?\]\(.*?\)', ' ', md_text) |
| html = markdown.markdown(md_text) |
| soup = BeautifulSoup(html, "html.parser") |
| for tag in soup(["img", "nav", "footer", "header", "aside", "script", "style"]): |
| tag.decompose() |
| clean_html = str(soup) |
| extracted = trafilatura.extract(clean_html, include_links=False, include_images=False, include_comments=False, favor_precision=True) |
| if not extracted: return md_text |
| extracted = re.sub(r'\n{3,}', '\n\n', extracted) |
| extracted = re.sub(r'[ \t]+', ' ', extracted) |
| return extracted |
| except Exception as e: |
| print(f" - [WARNING] Trafilatura text cleaning failed: {e}") |
| return md_text |
|
|
| def retry(retries=MAX_RETRIES, delay=5): |
| def decorator(f): |
| @wraps(f) |
| def wrapper(*args, **kwargs): |
| for i in range(retries): |
| try: |
| return f(*args, **kwargs) |
| except Exception as e: |
| if i < retries - 1: |
| time.sleep(delay * (2 ** i)) |
| else: |
| print(f" [ERROR] Function '{f.__name__}' failed after {retries} retries."); |
| return None |
| return wrapper |
| return decorator |
|
|
| class Field: |
| def __init__(self, value=None, confidence=0.0, sources=None, inferred_by=""): |
| self.value = value |
| self.confidence = confidence |
| self.sources = sources or[] |
| self.inferred_by = inferred_by |
| self.last_updated = datetime.now(timezone.utc).isoformat() |
|
|
| def to_dict(self): |
| return { |
| "value": self.value, "confidence": self.confidence, |
| "sources": self.sources, "inferred_by": self.inferred_by, "last_updated": self.last_updated |
| } |
|
|
| @classmethod |
| def from_dict(cls, data): |
| return cls(value=data.get('value'), confidence=data.get('confidence', 0.0), sources=data.get('sources',[]), inferred_by=data.get('inferred_by', '')) |
|
|
| class OpenRouterEmbedder: |
| def __init__(self, api_key, model_name): |
| self.api_key = api_key |
| self.model_name = model_name |
| self.url = "https://openrouter.ai/api/v1/embeddings" |
| |
| def encode(self, texts): |
| if not self.api_key: |
| raise Exception("OPENROUTER_API_KEY is missing for embeddings.") |
| if isinstance(texts, str): texts = [texts] |
| |
| embeddings =[] |
| headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"} |
| |
| for i in range(0, len(texts), 50): |
| batch = texts[i:i+50] |
| data = {"model": self.model_name, "input": batch} |
| resp = requests.post(self.url, headers=headers, json=data, timeout=60) |
| if resp.status_code == 200: |
| data_list = resp.json().get("data",[]) |
| data_list = sorted(data_list, key=lambda x: x["index"]) |
| embeddings.extend([item["embedding"] for item in data_list]) |
| else: |
| raise Exception(f"Embedding API Error: {resp.status_code} - {resp.text}") |
| return np.array(embeddings) |
|
|
|
|
| class MistralAnalystAgent: |
| def __init__(self, mistral_key: str, search_key: str, cse_id: str, schema: list, enable_fallback: bool = False): |
| if not all([mistral_key, search_key, cse_id]): |
| raise AgentInitializationError("One or more API keys (Mistral, Search) are missing.") |
|
|
| self.enable_fallback = enable_fallback |
| self.llm_client = ChatMistralAI(api_key=mistral_key, model=MISTRAL_MODEL, temperature=0.0) |
| self.search_api_key = search_key |
| self.cse_id = cse_id |
| self.schema = schema |
| |
| |
| self.field_instructions = self._generate_field_instructions() |
| |
| self.invalid_years =[str(y) for y in range(2015, 2025)] |
|
|
| print("[INFO] Initializing ML models and VectorDB...") |
| try: |
| self.nlp = spacy.load(SPACY_MODEL) |
| except (OSError, IOError) as e: |
| raise AgentInitializationError(f"[FATAL] Failed to load spaCy model '{SPACY_MODEL}'. Error: {e}") |
|
|
| self.chroma_client = chromadb.Client(Settings(persist_directory=VECTOR_DB_PATH, anonymized_telemetry=False, is_persistent=True, allow_reset=True)) |
| self.embedding_model = OpenRouterEmbedder(OPENROUTER_API_KEY, QWEN_EMBEDDING_MODEL) |
| self.cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL) |
| self.md_parser = MarkdownIt() |
| self.chroma_collection = None |
| self.bm25_index = None |
| self.mission_corpus =[] |
| self.corpus_map = {} |
| self.mission_inference_cache = {} |
| print(" -[SUCCESS] Models initialized.") |
|
|
| def shutdown(self): |
| try: self.chroma_client.reset() |
| except: pass |
|
|
| def get_legacy_caching_key(self, event_name: str) -> str: |
| base_name = re.sub(r'sprint|standard|olympic|full iron|half iron|70\.3', '', event_name, flags=re.IGNORECASE) |
| return re.sub(r'[^a-z0-9]+', '-', base_name.lower()).strip('-') |
|
|
| def get_caching_key(self, event_name: str, url: str = "") -> str: |
| if not event_name: event_name = "unknown_event" |
| base_name = re.sub(r'sprint|standard|olympic|full iron|half iron|70\.3', '', event_name, flags=re.IGNORECASE) |
| slug = re.sub(r'[^a-z0-9]+', '-', base_name.lower()).strip('-') |
| name_hash = hashlib.md5(slug.encode('utf-8')).hexdigest()[:8] |
| url_hash = "" |
| if url: url_hash = "-" + hashlib.md5(url.encode('utf-8')).hexdigest()[:8] |
| safe_slug = slug[:40].strip('-') |
| if not safe_slug: safe_slug = "event" |
| return f"{safe_slug}-{name_hash}{url_hash}" |
|
|
| def _generate_field_instructions(self) -> dict: |
| instructions = {} |
| db_fields = fetch_races_db_fields() |
| |
| |
| if not db_fields and not self.enable_fallback: |
| raise AgentInitializationError("Failed to fetch schema fields from Database2 and Fallback is Disabled. Check Database2 settings.") |
|
|
| meta_map = {row['field']: row for row in db_fields} if db_fields else {} |
|
|
| for key in self.schema: |
| if key in DEFAULT_BLANK_FIELDS: continue |
| |
| meta = meta_map.get(key) |
| if meta: |
| inst = f"Extract '{key}'" |
| if meta.get('display_name'): inst += f" (also referred to as '{meta['display_name']}')." |
| else: inst += "." |
| |
| if meta.get('question_text'): inst += f" Context constraint: {meta['question_text']}." |
| |
| opts = meta.get('data_options') |
| if opts: |
| if isinstance(opts, str): |
| if opts.strip().startswith('[') and opts.strip().endswith(']'): |
| try: |
| opt_list = json.loads(opts) |
| opt_str = ', '.join([str(o) for o in opt_list]) |
| except: opt_str = opts.replace('\n', ', ') |
| else: opt_str = opts.replace('\n', ', ') |
| else: opt_str = str(opts) |
| inst += f" MUST be exactly one of: [{opt_str}]." |
| elif key in CHOICE_OPTIONS: |
| inst += f" MUST be exactly one of: {', '.join(CHOICE_OPTIONS[key])}." |
| |
| dt = meta.get('data_type') |
| df = meta.get('data_format') |
| if dt or df: inst += f" Required format/type: {dt or ''} {df or ''}." |
| |
| instructions[key] = inst.strip() |
| else: |
| |
| if not self.enable_fallback: |
| raise AgentInitializationError(f"Field '{key}' not found in Database2 schema and Fallback is Disabled.") |
|
|
| if key in CHOICE_OPTIONS: instructions[key] = f"Extract '{key}'. MUST be one of: {', '.join(CHOICE_OPTIONS[key])}." |
| else: instructions[key] = f"Extract '{key}'." |
| |
| return instructions |
|
|
| def _normalize_key(self, key_str: str) -> str: |
| return re.sub(r'[^a-z0-9]', '', str(key_str).lower()) |
|
|
| def _clean_citations(self, text: str) -> str: |
| if not isinstance(text, str): return text |
| return re.sub(r'\[\d+\]', '', text).strip() |
|
|
| def _clean_dirtyjson(self, obj): |
| if isinstance(obj, dict): |
| return {self._clean_citations(k): self._clean_dirtyjson(v) for k, v in obj.items()} |
| elif isinstance(obj, list): |
| return[self._clean_dirtyjson(v) for v in obj] |
| elif isinstance(obj, str): |
| return self._clean_citations(obj) |
| else: |
| return obj |
|
|
| def _parse_json_response(self, response_text: str): |
| if not response_text: return None |
| clean_text = self._clean_citations(response_text) |
| match = re.search(r'\{.*\}|\[.*\]', clean_text, re.DOTALL) |
| if not match: return None |
| raw = match.group(0) |
| try: |
| return self._clean_dirtyjson(json.loads(raw)) |
| except json.JSONDecodeError: |
| try: |
| parsed = dirtyjson.loads(raw) |
| return self._clean_dirtyjson(parsed) |
| except Exception: |
| return None |
|
|
| @retry() |
| def _call_llm(self, prompt: str) -> str: |
| return self.llm_client.invoke([HumanMessage(content=prompt)]).content |
|
|
| @retry() |
| def _call_openrouter_gap_filler(self, prompt: str) -> str: |
| if not OPENROUTER_API_KEY: return None |
| headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"} |
| data = { |
| "model": PERPLEXITY_SONAR_MODEL, |
| "messages":[{"role": "user", "content": prompt}], |
| "temperature": 0.1, |
| "max_tokens": 2000 |
| } |
| try: |
| resp = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=data, timeout=60) |
| if resp.status_code == 200: return resp.json()['choices'][0]['message']['content'] |
| else: |
| print(f" - [ERROR] Perplexity API Error: {resp.status_code} - {resp.text}") |
| return None |
| except Exception as e: |
| print(f" -[ERROR] Perplexity connection failed: {e}") |
| return None |
|
|
| @retry() |
| def _call_openrouter_llama(self, prompt: str) -> str: |
| if not OPENROUTER_API_KEY: return None |
| headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json"} |
| data = { |
| "model": LLAMA_VERIFIER_MODEL, |
| "messages":[{"role": "user", "content": prompt}], |
| "temperature": 0.0, |
| "max_tokens": 3000 |
| } |
| try: |
| resp = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=data, timeout=60) |
| if resp.status_code == 200: return resp.json()['choices'][0]['message']['content'] |
| else: |
| print(f" -[ERROR] Llama API Error: {resp.status_code} - {resp.text}") |
| return None |
| except Exception as e: |
| print(f" - [ERROR] Llama connection failed: {e}") |
| return None |
|
|
| def _step_1a_initial_search(self, race_info: dict) -> list: |
| event_name = race_info.get("Festival") |
| print(f"\n[STEP 1A] Performing initial search for '{event_name}'") |
| query = f'{event_name} 2025 OR 2026' |
| try: |
| url = "https://www.googleapis.com/customsearch/v1" |
| params = {"key": self.search_api_key, "cx": self.cse_id, "q": query, "num": 10} |
| response = requests.get(url, params=params) |
| response.raise_for_status() |
| |
| raw_results =[{"title": i.get("title"), "link": i.get("link"), "snippet": i.get("snippet")} for i in response.json().get("items", [])] |
| clean_results =[r for r in raw_results if not any(d in r['link'].lower() for d in BLACKLISTED_DOMAINS)] |
| return clean_results |
| except requests.HTTPError as e: return[] |
|
|
| def _step_1b_validate_and_select_urls(self, event_name: str, search_results: list) -> list: |
| print("[STEP 1B] Validating search results with LLM...") |
| if not search_results: return[] |
|
|
| prompt = f"Identify the most relevant websites for '{event_name}'. Select the single best 'primary_url' (official page) and up to three 'secondary_urls'.\n\nSearch Results:\n{json.dumps(search_results, indent=2)}\n\nResponse MUST be JSON: {{'primary_url': '...', 'secondary_urls': [...]}}" |
| response_text = self._call_llm(prompt) |
| parsed = self._parse_json_response(response_text) |
| if parsed and isinstance(parsed, dict): |
| primary = parsed.get("primary_url") |
| secondaries = parsed.get("secondary_urls",[]) |
| final_urls = list(dict.fromkeys([u for u in ([primary] + secondaries if primary else secondaries) if u])) |
| return final_urls[:TOP_N_URLS_TO_PROCESS] |
| |
| if response_text: |
| salvaged_urls = re.findall(r'https?://[^\s"\'\)\],]+', response_text) |
| if salvaged_urls: |
| return list(dict.fromkeys(salvaged_urls))[:TOP_N_URLS_TO_PROCESS] |
| |
| return [r['link'] for r in search_results[:TOP_N_URLS_TO_PROCESS] if r.get('link')] |
|
|
| @retry(retries=2, delay=10) |
| def _get_content_from_url(self, url: str, is_web_research: bool) -> str | None: |
| if not url.strip(): return None |
| url_hash = hashlib.md5((url + str(is_web_research)).encode()).hexdigest() |
| storage_file_path = f"{url_hash}.md" |
| |
| if SUPABASE_CONFIGURED: |
| content = download_from_storage(storage_file_path) |
| if content: return content |
| |
| local_cache_path = os.path.join(CRAWL_CACHE_DIR, f"{url_hash}.md") |
| if os.path.exists(local_cache_path): |
| with open(local_cache_path, 'r', encoding='utf-8') as f: return f.read() |
| |
| print(f" - No cache found. Crawling: {url}") |
| try: |
| api_url = f"https://r.jina.ai/{url}" |
| response = requests.get(api_url, timeout=60) |
| if response.status_code == 200 and response.text: |
| content = response.text |
| if is_web_research: |
| print(" - Applying Trafilatura semantic extraction...") |
| content = clean_markdown_with_trafilatura(content) |
| with open(local_cache_path, 'w', encoding='utf-8') as f: f.write(content) |
| if SUPABASE_CONFIGURED: upload_to_storage(storage_file_path, content) |
| return content |
| return None |
| except requests.RequestException: return None |
|
|
| def _chunk_and_index_text(self, text: str, url: str, event_id_str: str): |
| if not text: return |
| from langchain_text_splitters import RecursiveCharacterTextSplitter |
| chunks = RecursiveCharacterTextSplitter(chunk_size=768, chunk_overlap=100).split_text(text) |
| if not chunks: return |
| |
| chunk_ids =[f"{event_id_str}_{hashlib.md5(chunk.encode()).hexdigest()}" for chunk in chunks] |
| unique_chunk_ids = list(set(chunk_ids)) |
| existing_chunks = self.chroma_collection.get(ids=unique_chunk_ids) |
| existing_ids = set(existing_chunks['ids']) |
| |
| new_chunks_to_add =[] |
| start_idx = len(self.mission_corpus) |
|
|
| for i, chunk_id in enumerate(chunk_ids): |
| self.mission_corpus.append(chunks[i]) |
| self.corpus_map[start_idx + i] = {'id': chunk_id, 'snippet': chunks[i]} |
| if chunk_id not in existing_ids: |
| new_chunks_to_add.append({'id': chunk_id, 'chunk': chunks[i]}) |
| existing_ids.add(chunk_id) |
| |
| if new_chunks_to_add: |
| print(f" - Indexing {len(new_chunks_to_add)} new passages from {url}") |
| new_ids = [item['id'] for item in new_chunks_to_add] |
| new_documents =[item['chunk'] for item in new_chunks_to_add] |
| try: |
| new_embeddings = self.embedding_model.encode(new_documents).tolist() |
| new_metadatas =[{"source_url": url, "event_id": event_id_str} for _ in new_ids] |
| if new_ids: |
| self.chroma_collection.add(ids=new_ids, embeddings=new_embeddings, documents=new_documents, metadatas=new_metadatas) |
| except Exception as e: |
| print(f" -[ERROR] Failed to embed documents: {e}") |
|
|
| def _build_bm25_index(self): |
| if self.mission_corpus: |
| tokenized_corpus =[doc.lower().split() for doc in self.mission_corpus] |
| self.bm25_index = BM25Okapi(tokenized_corpus) |
|
|
| def _retrieve_and_fuse_evidence(self, query: str, top_k: int) -> list[dict]: |
| collection_count = self.chroma_collection.count() |
| if collection_count == 0: return[] |
| |
| vector_hits = {} |
| try: |
| query_embedding = self.embedding_model.encode([query]).tolist() |
| n_results = min(top_k * 2, collection_count) |
| chroma_results = self.chroma_collection.query(query_embeddings=query_embedding, n_results=n_results) |
| if chroma_results['ids']: |
| for rank, (doc_id, doc) in enumerate(zip(chroma_results['ids'][0], chroma_results['documents'][0])): |
| vector_hits[doc_id] = {'snippet': doc, 'rank': rank} |
| except Exception: pass |
|
|
| keyword_hits = {} |
| if self.bm25_index: |
| tokenized_query = query.lower().split() |
| doc_scores = self.bm25_index.get_scores(tokenized_query) |
| top_indices = sorted(range(len(doc_scores)), key=lambda i: doc_scores[i], reverse=True)[:top_k*2] |
| for rank, idx in enumerate(top_indices): |
| if idx in self.corpus_map: |
| item = self.corpus_map[idx] |
| keyword_hits[item['id']] = {'snippet': item['snippet'], 'rank': rank} |
|
|
| fused_scores = {} |
| k = 60 |
| for doc_id, hit in vector_hits.items(): |
| if doc_id not in fused_scores: fused_scores[doc_id] = 0 |
| fused_scores[doc_id] += 1 / (k + hit['rank'] + 1) |
| for doc_id, hit in keyword_hits.items(): |
| if doc_id not in fused_scores: fused_scores[doc_id] = 0 |
| fused_scores[doc_id] += 1 / (k + hit['rank'] + 1) |
| |
| sorted_ids = sorted(fused_scores.keys(), key=lambda x: fused_scores[x], reverse=True)[:top_k] |
| final_results =[] |
| for doc_id in sorted_ids: |
| if doc_id in vector_hits: final_results.append({'id': doc_id, 'snippet': vector_hits[doc_id]['snippet']}) |
| elif doc_id in keyword_hits: final_results.append({'id': doc_id, 'snippet': keyword_hits[doc_id]['snippet']}) |
| return final_results |
|
|
| def _rerank_evidence_with_cross_encoder(self, query: str, evidence: list[dict]) -> list[dict]: |
| if not evidence: return [] |
| pairs =[(query, item['snippet']) for item in evidence] |
| scores = self.cross_encoder.predict(pairs) |
| for i, item in enumerate(evidence): |
| item['rerank_score'] = float(scores[i]) |
| return sorted(evidence, key=lambda x: x['rerank_score'], reverse=True) |
|
|
| def _perform_holistic_extraction(self, knowledge_base: dict, event_name: str, variant_name: str): |
| print(f" - Pass 1: Holistic RAG extraction for '{variant_name}'...") |
| all_fields_query = f"Extract all available information for the event '{event_name}', specifically for the '{variant_name}' category." |
| candidate_evidence = self._retrieve_and_fuse_evidence(all_fields_query, top_k=RAG_CANDIDATE_POOL_SIZE) |
| reranked_evidence = self._rerank_evidence_with_cross_encoder(all_fields_query, candidate_evidence) |
| final_evidence = reranked_evidence[:RAG_FINAL_EVIDENCE_COUNT + 5] |
|
|
| if not final_evidence: return knowledge_base |
|
|
| evidence_prompt = "\n".join([f"Evidence Snippet:\n---\n{e['snippet']}\n---" for e in final_evidence]) |
| fields_to_extract =[f"- {key}: {desc}" for key, desc in self.field_instructions.items()] |
| fields_prompt = "\n".join(fields_to_extract) |
|
|
| prompt = f"""You are a strict data extraction assistant. Based on the evidence below, extract all specified fields for '{variant_name}' of '{event_name}'. |
| CRITICAL RULES: |
| - Extract ONLY verifiable information directly from the provided source. |
| - DO NOT infer, assume, estimate, or create missing values. |
| - If a value is not explicitly present, omit the key or set it to null. |
| - **IMPORTANT**: For prices, distances, and cutoffs, ensure they match the EXACT variant '{variant_name}'. If the evidence shows data for a DIFFERENT distance, return null. |
| - Return ONLY flat strings. NO nested JSON/dictionaries inside values. |
| |
| ## Evidence |
| {evidence_prompt} |
| |
| ## Task |
| Extract these fields into a JSON object: |
| {fields_prompt} |
| |
| ## Response |
| JSON object only.""" |
|
|
| response_text = self._call_llm(prompt) |
| extracted_data = self._parse_json_response(response_text) |
| |
| if extracted_data and isinstance(extracted_data, dict): |
| norm_schema_keys = {self._normalize_key(k): k for k in self.schema} |
| filled = 0 |
| confidence = 0.85 |
| for key, value in extracted_data.items(): |
| norm_key = norm_schema_keys.get(self._normalize_key(key)) |
| if norm_key and value and str(value).lower() not in["null", "none", "unknown", ""]: |
| if not knowledge_base[variant_name].get(norm_key, Field()).value: |
| val_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value) |
| knowledge_base[variant_name][norm_key] = Field(value=val_str, confidence=confidence, inferred_by="rag_holistic", sources=final_evidence) |
| filled += 1 |
| print(f" - [SUCCESS] Holistic extracted {filled} fields.") |
| return knowledge_base |
|
|
| def _perform_targeted_recovery(self, knowledge_base: dict, event_name: str, variant_name: str): |
| print(f" - Pass 2: Targeted Recovery for missing fields...") |
| missing_fields =[k for k in self.schema if k not in DEFAULT_BLANK_FIELDS and not knowledge_base[variant_name].get(k, Field()).value] |
| |
| if not missing_fields: return knowledge_base |
| |
| for field in missing_fields: |
| query = f"What is the {field} for {variant_name} in {event_name}?" |
| evidence = self._retrieve_and_fuse_evidence(query, top_k=10) |
| reranked = self._rerank_evidence_with_cross_encoder(query, evidence)[:3] |
| |
| if not reranked: continue |
| evidence_text = "\n".join([e['snippet'] for e in reranked]) |
| |
| instruction = self.field_instructions.get(field, f"Extract '{field}'.") |
|
|
| prompt = f"""Based on the evidence, find the exact value for '{field}' for the race '{variant_name}'. |
| {instruction} |
| |
| CRITICAL RULES: |
| - Extract ONLY verifiable information directly from the provided source. |
| - Ensure the value applies EXACTLY to the '{variant_name}'. If the evidence is for a different distance/variant, return null. |
| - Return flat strings only. No dicts. |
| |
| Evidence: {evidence_text} |
| Return a JSON object: {{"answer": "value"}} or {{"answer": null}} if not found.""" |
| |
| response = self._call_llm(prompt) |
| parsed = self._parse_json_response(response) |
| if parsed and isinstance(parsed, dict): |
| ans = parsed.get('answer') |
| if ans and str(ans).lower() not in['null', 'none', 'unknown', '']: |
| val_str = json.dumps(ans) if isinstance(ans, (dict, list)) else str(ans) |
| knowledge_base[variant_name][field] = Field(value=val_str, confidence=0.9, inferred_by="rag_targeted", sources=reranked) |
| return knowledge_base |
|
|
| def _perform_batched_gap_fill(self, knowledge_base: dict, event_name: str): |
| if not OPENROUTER_API_KEY: return knowledge_base |
|
|
| print(f" - Pass 3: Batched Gap Filling ({PERPLEXITY_SONAR_MODEL})...") |
| |
| missing_data_map = {} |
| for variant, fields in knowledge_base.items(): |
| missing_fields =[k for k in self.schema if k not in DEFAULT_BLANK_FIELDS and not fields.get(k, Field()).value] |
| if missing_fields: missing_data_map[variant] = missing_fields |
|
|
| if not missing_data_map: return knowledge_base |
|
|
| records_str = "" |
| all_missing_fields_set = set() |
| |
| for i, (variant, m_fields) in enumerate(missing_data_map.items()): |
| field_list_str = "\n - ".join(m_fields) |
| records_str += f"\n{i+1}. Record ID: \"{variant}\"\n Missing fields:\n - {field_list_str}\n" |
| for f in m_fields: all_missing_fields_set.add(f) |
|
|
| fields_to_extract =[f"- {key}: {self.field_instructions.get(key, 'Extract exact value')}" for key in all_missing_fields_set] |
| field_specific_instructions = "\n".join(fields_to_extract) |
|
|
| prompt = f"""Search the live web to fill the following missing fields for the athletic event '{event_name}'. |
| |
| Records: |
| {records_str} |
| |
| Search Methodology: |
| - Use your ability to search the web to find official, accurate data for this specific event. |
| - Do not repeat or re-fetch fields not listed. |
| |
| CRITICAL Extraction Rules: |
| - DIFFERENT VARIANTS HAVE DIFFERENT PRICES AND CUTOFFS. You MUST extract the specific registrationCost and Cutoff time for EACH individual variant. Do not blindly copy one variant's fee to another. |
| - Extract ONLY verifiable information found online. |
| - DO NOT infer, assume, estimate, or create missing values. |
| - If specific information is unavailable on the web, return: NULL |
| - Return ONLY flat strings. NO nested objects! |
| |
| Field Definitions (Strictly follow these definitions): |
| {field_specific_instructions} |
| |
| Output Format: |
| Return strict JSON matching exactly this schema: |
| {{ |
| "<record_id_exactly_as_written_above>": {{ |
| "<field_name_exactly_as_written>": "<exact_value_or_NULL>" |
| }} |
| }}""" |
| response_text = self._call_openrouter_gap_filler(prompt) |
| |
| if response_text: |
| os.makedirs(SONAR_ANALYTICS_DIR, exist_ok=True) |
| clean_name = re.sub(r'[^a-zA-Z0-9]+', '_', event_name[:20]).strip('_') |
| analytics_file = os.path.join(SONAR_ANALYTICS_DIR, f"sonar_gap_{clean_name}_{int(time.time())}.json") |
| with open(analytics_file, 'w', encoding='utf-8') as f: |
| f.write(response_text) |
|
|
| filled_data = self._parse_json_response(response_text) |
| |
| if filled_data and isinstance(filled_data, dict): |
| filled_count = 0 |
| norm_kb_keys = {self._normalize_key(k): k for k in knowledge_base.keys()} |
| norm_schema_keys = {self._normalize_key(k): k for k in self.schema} |
|
|
| for variant, fields in filled_data.items(): |
| norm_var = self._normalize_key(variant) |
| actual_var = norm_kb_keys.get(norm_var) |
| |
| if actual_var and isinstance(fields, dict): |
| for key, value in fields.items(): |
| norm_key = norm_schema_keys.get(self._normalize_key(key)) |
| |
| if norm_key and value and str(value).lower() not in["null", "none", "unknown", ""]: |
| val_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value) |
| knowledge_base[actual_var][norm_key] = Field(value=val_str, confidence=0.95, inferred_by="sonar_gap_fill") |
| filled_count += 1 |
| print(f" - [SUCCESS] Perplexity Sonar filled {filled_count} missing gaps.") |
| return knowledge_base |
|
|
| def _verify_hallucinations(self, knowledge_base: dict): |
| if not OPENROUTER_API_KEY: return knowledge_base |
| print(f" - Pass 4: Hallucination Check ({LLAMA_VERIFIER_MODEL})...") |
| |
| kb_json = {} |
| for variant, fields in knowledge_base.items(): |
| kb_json[variant] = {} |
| for k, f in fields.items(): |
| if f.value and k not in DEFAULT_BLANK_FIELDS: |
| kb_json[variant][k] = f.value |
| |
| if not kb_json: return knowledge_base |
|
|
| prompt = f"""You are a strict Data Integrity Verifier. |
| Review the following extracted event data for hallucinations, logical impossibilities, or invalid formats. |
| |
| Data: |
| {json.dumps(kb_json, indent=2)} |
| |
| Rules for Rejection (Set to NULL if violated): |
| - Dates must be logically possible. |
| - Prices must be realistic. |
| - Distances must match reality (e.g., ultras can be 50k, 75k, 100k, 100M). Do not reject standard ultra distances. |
| - Text fields must not contain conversational filler (e.g., "The price is $50"). |
| |
| Return ONLY a JSON object with the exact same structure. Replace any hallucinated or invalid values with NULL. Do not add any extra text or markdown outside the JSON. |
| """ |
| response = self._call_openrouter_llama(prompt) |
| if response: |
| verified_data = self._parse_json_response(response) |
| |
| if verified_data and isinstance(verified_data, dict): |
| rejected_count = 0 |
| |
| norm_kb_keys = {self._normalize_key(k): k for k in knowledge_base.keys()} |
| norm_schema_keys = {self._normalize_key(k): k for k in self.schema} |
| |
| for variant, fields in verified_data.items(): |
| actual_var = norm_kb_keys.get(self._normalize_key(variant)) |
| if actual_var and isinstance(fields, dict): |
| for key, val in fields.items(): |
| actual_key = norm_schema_keys.get(self._normalize_key(key)) |
| if actual_key and (val is None or str(val).lower() == "null"): |
| knowledge_base[actual_var][actual_key] = Field() |
| rejected_count += 1 |
| if rejected_count > 0: |
| print(f" - [WARNING] Llama rejected {rejected_count} fields as hallucinations.") |
| else: |
| print(f" -[SUCCESS] Llama verified data integrity.") |
| return knowledge_base |
|
|
| def _discover_and_validate_variants_from_content(self, content: str, knowledge_base: dict): |
| print(" - Discovering all race variants from content...") |
| prompt = f"""You are a strict data extractor. Identify the specific, official race categories or distances explicitly available for registration in the text provided. |
| CRITICAL RULES: |
| 1. Extract ONLY variants explicitly mentioned in the text. |
| 2. DO NOT invent, assume, or infer distances (e.g., do not add "Half Marathon" if only "50K" is listed). |
| 3. DO NOT list general distances, only those meant for this specific event. |
| Return ONLY a JSON list of strings. If no variants are found, return an empty list[]. |
| Text: {content[:15000]}...""" |
| |
| response_text = self._call_llm(prompt) |
| variants = self._parse_json_response(response_text) |
| |
| if isinstance(variants, list): |
| cleaned_variants =[v for v in variants if isinstance(v, str) and len(v) < 100] |
| if cleaned_variants: |
| print(f" - [SUCCESS] Found variants: {', '.join(cleaned_variants)}") |
| for variant_name in cleaned_variants: |
| if variant_name not in knowledge_base: |
| knowledge_base[variant_name] = {field: Field() for field in self.schema} |
| return |
| print(" - [WARNING] Variant discovery failed. Using default.") |
|
|
| def _inject_pre_filled_data(self, knowledge_base: dict, pre_filled_data: dict, pre_filled_confidence: float = 0.3): |
| if not pre_filled_data: return |
| print(f" - Injecting pre-filled data (confidence={pre_filled_confidence})...") |
| for variant_name in knowledge_base.keys(): |
| for key, value in pre_filled_data.items(): |
| if key in self.schema and value: |
| if knowledge_base[variant_name].get(key, Field()).confidence < pre_filled_confidence: |
| val_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value) |
| knowledge_base[variant_name][key] = Field(value=val_str, confidence=pre_filled_confidence, inferred_by="pre_processed_data") |
|
|
| def _run_inferential_filling(self, knowledge_base: dict): |
| print("\n[INFERENCE] Running final analysis to infer missing data...") |
| for variant, data in knowledge_base.items(): |
| city = data.get('city', Field()).value |
| country = data.get('country', Field()).value |
| if city and not country: |
| resp = self._call_llm(f"What country is {city} in? Return ONLY the country name.") |
| if resp: |
| knowledge_base[variant]['country'] = Field(value=resp.strip(), confidence=0.8, inferred_by="inference") |
| return knowledge_base |
|
|
| def determine_event_type(self, url: str) -> str | None: |
| print(f" - Determining event type for URL: {url}") |
| content = self._get_content_from_url(url, is_web_research=False) |
| if not content: return None |
| valid_types = ", ".join(CHOICE_OPTIONS["type"]) |
| prompt = f"Analyze the text and determine the athletic event type. Answer MUST be one of: {valid_types}.\nText: {content[:2000]}" |
| response = self._call_llm(prompt) |
| if not response: return None |
| for et in CHOICE_OPTIONS["type"]: |
| if et.lower() in response.lower(): |
| print(f" - [SUCCESS] Determined type: {et}") |
| return et |
| return None |
|
|
| def _check_supabase_for_cache(self, event_key: str, legacy_key: str = None) -> dict | None: |
| if not SUPABASE_CONFIGURED: return None |
| data = get_knowledge_cache(event_key) |
| if data: |
| print(f"[SUCCESS] Found cache for '{event_key}' in Supabase.") |
| return data |
| if legacy_key and legacy_key != event_key: |
| print(f" - Primary key missing. Checking legacy key: '{legacy_key}'...") |
| data = get_knowledge_cache(legacy_key) |
| if data: |
| print(f"[SUCCESS] Found legacy cache for '{legacy_key}' in Supabase.") |
| return data |
| return None |
|
|
| def run(self, race_info: dict) -> dict | None: |
| """Called by Web Research Agent.""" |
| event_name = race_info.get("Festival") |
| event_key = self.get_caching_key(event_name) |
| legacy_key = self.get_legacy_caching_key(event_name) |
| |
| cached_data = self._check_supabase_for_cache(event_key, legacy_key) |
| if cached_data: |
| return {v: {f: Field.from_dict(d) for f, d in fs.items()} for v, fs in cached_data.items()} |
| |
| search_results = self._step_1a_initial_search(race_info) |
| if not search_results: return None |
| validated_urls = self._step_1b_validate_and_select_urls(event_name, search_results) |
| if not validated_urls: return None |
| |
| knowledge_base = self._crawl_and_extract(validated_urls, race_info, is_web_research=True) |
| |
| if knowledge_base and SUPABASE_CONFIGURED: |
| serializable_kb = {v: {f: field.to_dict() for f, field in fs.items()} for v, fs in knowledge_base.items()} |
| set_knowledge_cache(event_key, serializable_kb) |
| return knowledge_base |
|
|
| def run_direct(self, race_info: dict, direct_urls: list, pre_filled_data: dict = None, pre_filled_confidence: float = 0.3) -> dict | None: |
| """Called by Prescraped Data Processor.""" |
| event_name = race_info.get("Festival") |
| url_part = hashlib.md5(direct_urls[0].encode()).hexdigest()[:8] |
| event_key = self.get_caching_key(f"{event_name}-{url_part}") |
| legacy_key = self.get_legacy_caching_key(event_name) |
| |
| cached_data = self._check_supabase_for_cache(event_key, legacy_key) |
| if cached_data: |
| return {v: {f: Field.from_dict(d) for f, d in fs.items()} for v, fs in cached_data.items()} |
|
|
| validated_urls =[url for url in direct_urls if self._is_valid_url(url)] |
| if not validated_urls: return None |
| |
| knowledge_base = self._crawl_and_extract(validated_urls, race_info, pre_filled_data, pre_filled_confidence, is_web_research=False) |
| |
| if knowledge_base and SUPABASE_CONFIGURED: |
| serializable_kb = {v: {f: field.to_dict() for f, field in fs.items()} for v, fs in knowledge_base.items()} |
| set_knowledge_cache(event_key, serializable_kb) |
| return knowledge_base |
|
|
| def _is_valid_url(self, url: str) -> bool: |
| if not url or not url.startswith(('http://', 'https://')): return False |
| if any(year in url for year in self.invalid_years): return False |
| return True |
|
|
| def _crawl_and_extract(self, urls: list, race_info: dict, pre_filled_data: dict = None, pre_filled_confidence: float = 0.3, is_web_research: bool = False) -> dict: |
| event_name = race_info.get("Festival") |
| event_id_str = self.get_caching_key(event_name) |
| self.chroma_collection = self.chroma_client.get_or_create_collection(name=f"coll_{event_id_str}") |
| print(f"\n[STEP 2] Starting RAG processing for '{event_name}' (Collection: coll_{event_id_str})") |
| |
| self.mission_corpus =[] |
| self.corpus_map = {} |
| self.bm25_index = None |
|
|
| knowledge_base = {} |
| all_content =[] |
| with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_CRAWLERS) as executor: |
| futures = {executor.submit(self._get_content_from_url, url, is_web_research): url for url in urls} |
| for future in as_completed(futures): |
| if content := future.result(): |
| all_content.append(content) |
| self._chunk_and_index_text(content, futures[future], event_id_str) |
|
|
| if not all_content: |
| print(" -[WARNING] No content crawled.") |
| return {} |
|
|
| self._build_bm25_index() |
|
|
| combined_content = "\n\n".join(all_content) |
| self._discover_and_validate_variants_from_content(combined_content, knowledge_base) |
|
|
| if not knowledge_base: |
| knowledge_base[event_name] = {field: Field() for field in self.schema} |
|
|
| if pre_filled_data: |
| self._inject_pre_filled_data(knowledge_base, pre_filled_data, pre_filled_confidence) |
|
|
| for variant in list(knowledge_base.keys()): |
| self._perform_holistic_extraction(knowledge_base, event_name, variant) |
| self._perform_targeted_recovery(knowledge_base, event_name, variant) |
|
|
| self._perform_batched_gap_fill(knowledge_base, event_name) |
| self._verify_hallucinations(knowledge_base) |
| self._run_inferential_filling(knowledge_base) |
| |
| print("\n[SUCCESS] Extraction complete.") |
| return knowledge_base |
