Spaces:
Running
Running
| import re | |
| import pycountry | |
| from docx import Document | |
| import json | |
| import os | |
| import numpy as np | |
| import faiss | |
| from collections import defaultdict | |
| import ast # For literal_eval | |
| import math # For ceiling function | |
| import data_preprocess | |
| import mtdna_classifier | |
| # --- IMPORTANT: UNCOMMENT AND CONFIGURE YOUR REAL API KEY --- | |
| import google.generativeai as genai | |
| #genai.configure(api_key=os.getenv("GOOGLE_API_KEY")) | |
| genai.configure(api_key=os.getenv("GOOGLE_API_KEY_BACKUP")) | |
| import nltk | |
| from nltk.corpus import stopwords | |
| try: | |
| nltk.data.find('corpora/stopwords') | |
| except LookupError: | |
| nltk.download('stopwords') | |
| nltk.download('punkt_tab') | |
| # # --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) --- | |
| # # Prices are per 1,000 tokens | |
| # PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens | |
| # PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens | |
| # PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens | |
| # Gemini 2.5 Flash-Lite pricing per 1,000 tokens | |
| PRICE_PER_1K_INPUT_LLM = 0.00010 # $0.10 per 1M input tokens | |
| PRICE_PER_1K_OUTPUT_LLM = 0.00040 # $0.40 per 1M output tokens | |
| # Embedding-001 pricing per 1,000 input tokens | |
| PRICE_PER_1K_EMBEDDING_INPUT = 0.00015 # $0.15 per 1M input tokens | |
| # --- API Functions (REAL API FUNCTIONS) --- | |
| # def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"): | |
| # """Generates an embedding for the given text using a Google embedding model.""" | |
| # try: | |
| # result = genai.embed_content( | |
| # model="models/text-embedding-004", # Specify the embedding model | |
| # content=text, | |
| # task_type=task_type | |
| # ) | |
| # return np.array(result['embedding']).astype('float32') | |
| # except Exception as e: | |
| # print(f"Error getting embedding: {e}") | |
| # return np.zeros(768, dtype='float32') | |
| def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"): | |
| """Safe Gemini 1.5 embedding call with fallback.""" | |
| import numpy as np | |
| try: | |
| if not text or len(text.strip()) == 0: | |
| raise ValueError("Empty text cannot be embedded.") | |
| result = genai.embed_content( | |
| model="models/text-embedding-004", | |
| content=text, | |
| task_type=task_type | |
| ) | |
| return np.array(result['embedding'], dtype='float32') | |
| except Exception as e: | |
| print(f"❌ Embedding error: {e}") | |
| return np.zeros(768, dtype='float32') | |
| def call_llm_api(prompt, model_name="gemini-2.5-flash-lite"):#'gemini-1.5-flash-latest'): | |
| """Calls a Google Gemini LLM with the given prompt.""" | |
| try: | |
| model = genai.GenerativeModel(model_name) | |
| response = model.generate_content(prompt) | |
| return response.text, model # Return model instance for token counting | |
| except Exception as e: | |
| print(f"Error calling LLM: {e}") | |
| return "Error: Could not get response from LLM API.", None | |
| # --- Core Document Processing Functions (All previously provided and fixed) --- | |
| def read_docx_text(path): | |
| """ | |
| Reads text and extracts potential table-like strings from a .docx document. | |
| Separates plain text from structured [ [ ] ] list-like tables. | |
| Also attempts to extract a document title. | |
| """ | |
| doc = Document(path) | |
| plain_text_paragraphs = [] | |
| table_strings = [] | |
| document_title = "Unknown Document Title" # Default | |
| # Attempt to extract the document title from the first few paragraphs | |
| title_paragraphs = [p.text.strip() for p in doc.paragraphs[:5] if p.text.strip()] | |
| if title_paragraphs: | |
| # A heuristic to find a title: often the first or second non-empty paragraph | |
| # or a very long first paragraph if it's the title | |
| if len(title_paragraphs[0]) > 50 and "Human Genetics" not in title_paragraphs[0]: | |
| document_title = title_paragraphs[0] | |
| elif len(title_paragraphs) > 1 and len(title_paragraphs[1]) > 50 and "Human Genetics" not in title_paragraphs[1]: | |
| document_title = title_paragraphs[1] | |
| elif any("Complete mitochondrial genomes" in p for p in title_paragraphs): | |
| # Fallback to a known title phrase if present | |
| document_title = "Complete mitochondrial genomes of Thai and Lao populations indicate an ancient origin of Austroasiatic groups and demic diffusion in the spread of Tai–Kadai languages" | |
| current_table_lines = [] | |
| in_table_parsing_mode = False | |
| for p in doc.paragraphs: | |
| text = p.text.strip() | |
| if not text: | |
| continue | |
| # Condition to start or continue table parsing | |
| if text.startswith("## Table "): # Start of a new table section | |
| if in_table_parsing_mode and current_table_lines: | |
| table_strings.append("\n".join(current_table_lines)) | |
| current_table_lines = [text] # Include the "## Table X" line | |
| in_table_parsing_mode = True | |
| elif in_table_parsing_mode and (text.startswith("[") or text.startswith('"')): | |
| # Continue collecting lines if we're in table mode and it looks like table data | |
| # Table data often starts with '[' for lists, or '"' for quoted strings within lists. | |
| current_table_lines.append(text) | |
| else: | |
| # If not in table mode, or if a line doesn't look like table data, | |
| # then close the current table (if any) and add the line to plain text. | |
| if in_table_parsing_mode and current_table_lines: | |
| table_strings.append("\n".join(current_table_lines)) | |
| current_table_lines = [] | |
| in_table_parsing_mode = False | |
| plain_text_paragraphs.append(text) | |
| # After the loop, add any remaining table lines | |
| if current_table_lines: | |
| table_strings.append("\n".join(current_table_lines)) | |
| return "\n".join(plain_text_paragraphs), table_strings, document_title | |
| # --- Structured Data Extraction and RAG Functions --- | |
| def parse_literal_python_list(table_str): | |
| list_match = re.search(r'(\[\s*\[\s*(?:.|\n)*?\s*\]\s*\])', table_str) | |
| #print("Debug: list_match object (before if check):", list_match) | |
| if not list_match: | |
| if "table" in table_str.lower(): # then the table doest have the "]]" at the end | |
| table_str += "]]" | |
| list_match = re.search(r'(\[\s*\[\s*(?:.|\n)*?\s*\]\s*\])', table_str) | |
| if list_match: | |
| try: | |
| matched_string = list_match.group(1) | |
| #print("Debug: Matched string for literal_eval:", matched_string) | |
| return ast.literal_eval(matched_string) | |
| except (ValueError, SyntaxError) as e: | |
| print(f"Error evaluating literal: {e}") | |
| return [] | |
| return [] | |
| _individual_code_parser = re.compile(r'([A-Z0-9]+?)(\d+)$', re.IGNORECASE) | |
| def _parse_individual_code_parts(code_str): | |
| match = _individual_code_parser.search(code_str) | |
| if match: | |
| return match.group(1), match.group(2) | |
| return None, None | |
| def parse_sample_id_to_population_code(plain_text_content): | |
| sample_id_map = {} | |
| contiguous_ranges_data = defaultdict(list) | |
| #section_start_marker = "The sample identification of each population is as follows:" | |
| section_start_marker = ["The sample identification of each population is as follows:","## table"] | |
| for s in section_start_marker: | |
| relevant_text_search = re.search( | |
| re.escape(s.lower()) + r"\s*(.*?)(?=\n##|\Z)", | |
| plain_text_content.lower(), | |
| re.DOTALL | |
| ) | |
| if relevant_text_search: | |
| break | |
| if not relevant_text_search: | |
| print("Warning: 'Sample ID Population Code' section start marker not found or block empty.") | |
| return sample_id_map, contiguous_ranges_data | |
| relevant_text_block = relevant_text_search.group(1).strip() | |
| # print(f"\nDEBUG_PARSING: --- Start of relevant_text_block (first 500 chars) ---") | |
| # print(relevant_text_block[:500]) | |
| # print(f"DEBUG_PARSING: --- End of relevant_text_block (last 500 chars) ---") | |
| # print(relevant_text_block[-500:]) | |
| # print(f"DEBUG_PARSING: Relevant text block length: {len(relevant_text_block)}") | |
| mapping_pattern = re.compile( | |
| r'\b([A-Z0-9]+\d+)(?:-([A-Z0-9]+\d+))?\s+([A-Z0-9]+)\b', # Changed the last group | |
| re.IGNORECASE) | |
| range_expansion_count = 0 | |
| direct_id_count = 0 | |
| total_matches_found = 0 | |
| for match in mapping_pattern.finditer(relevant_text_block): | |
| total_matches_found += 1 | |
| id1_full_str, id2_full_str_opt, pop_code = match.groups() | |
| #print(f" DEBUG_PARSING: Matched: '{match.group(0)}'") | |
| pop_code_upper = pop_code.upper() | |
| id1_prefix, id1_num_str = _parse_individual_code_parts(id1_full_str) | |
| if id1_prefix is None: | |
| #print(f" DEBUG_PARSING: Failed to parse ID1: {id1_full_str}. Skipping this mapping.") | |
| continue | |
| if id2_full_str_opt: | |
| id2_prefix_opt, id2_num_str_opt = _parse_individual_code_parts(id2_full_str_opt) | |
| if id2_prefix_opt is None: | |
| #print(f" DEBUG_PARSING: Failed to parse ID2: {id2_full_str_opt}. Treating {id1_full_str} as single ID1.") | |
| sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper | |
| direct_id_count += 1 | |
| continue | |
| #print(f" DEBUG_PARSING: Comparing prefixes: '{id1_prefix.lower()}' vs '{id2_prefix_opt.lower()}'") | |
| if id1_prefix.lower() == id2_prefix_opt.lower(): | |
| #print(f" DEBUG_PARSING: ---> Prefixes MATCH for range expansion! Range: {id1_prefix}{id1_num_str}-{id2_prefix_opt}{id2_num_str_opt}") | |
| try: | |
| start_num = int(id1_num_str) | |
| end_num = int(id2_num_str_opt) | |
| for num in range(start_num, end_num + 1): | |
| sample_id = f"{id1_prefix.upper()}{num}" | |
| sample_id_map[sample_id] = pop_code_upper | |
| range_expansion_count += 1 | |
| contiguous_ranges_data[id1_prefix.upper()].append( | |
| (start_num, end_num, pop_code_upper) | |
| ) | |
| except ValueError: | |
| print(f" DEBUG_PARSING: ValueError in range conversion for {id1_num_str}-{id2_num_str_opt}. Adding endpoints only.") | |
| sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper | |
| sample_id_map[f"{id2_prefix_opt.upper()}{id2_num_str_opt}"] = pop_code_upper | |
| direct_id_count += 2 | |
| else: | |
| #print(f" DEBUG_PARSING: Prefixes MISMATCH for range: '{id1_prefix}' vs '{id2_prefix_opt}'. Adding endpoints only.") | |
| sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper | |
| sample_id_map[f"{id2_prefix_opt.upper()}{id2_num_str_opt}"] = pop_code_upper | |
| direct_id_count += 2 | |
| else: | |
| sample_id_map[f"{id1_prefix.upper()}{id1_num_str}"] = pop_code_upper | |
| direct_id_count += 1 | |
| # print(f"DEBUG_PARSING: Total matches found by regex: {total_matches_found}.") | |
| # print(f"DEBUG_PARSING: Parsed sample IDs: {len(sample_id_map)} total entries.") | |
| # print(f"DEBUG_PARSING: (including {range_expansion_count} from range expansion and {direct_id_count} direct ID/endpoint entries).") | |
| return sample_id_map, contiguous_ranges_data | |
| country_keywords_regional_overrides = { | |
| "north thailand": "Thailand", "central thailand": "Thailand", | |
| "northeast thailand": "Thailand", "east myanmar": "Myanmar", "west thailand": "Thailand", | |
| "central india": "India", "east india": "India", "northeast india": "India", | |
| "south sibera": "Russia", "siberia": "Russia", "yunnan": "China", #"tibet": "China", | |
| "sumatra": "Indonesia", "borneo": "Indonesia", | |
| "northern mindanao": "Philippines", "west malaysia": "Malaysia", | |
| "mongolia": "China", | |
| "beijing": "China", | |
| "north laos": "Laos", "central laos": "Laos", | |
| "east myanmar": "Myanmar", "west myanmar": "Myanmar"} | |
| # Updated get_country_from_text function | |
| def get_country_from_text(text): | |
| text_lower = text.lower() | |
| # 1. Use pycountry for official country names and common aliases | |
| for country in pycountry.countries: | |
| # Check full name match first | |
| if text_lower == country.name.lower(): | |
| return country.name | |
| # Safely check for common_name | |
| if hasattr(country, 'common_name') and text_lower == country.common_name.lower(): | |
| return country.common_name | |
| # Safely check for official_name | |
| if hasattr(country, 'official_name') and text_lower == country.official_name.lower(): | |
| return country.official_name | |
| # Check if country name is part of the text (e.g., 'Thailand' in 'Thailand border') | |
| if country.name.lower() in text_lower: | |
| return country.name | |
| # Safely check if common_name is part of the text | |
| if hasattr(country, 'common_name') and country.common_name.lower() in text_lower: | |
| return country.common_name | |
| # 2. Prioritize specific regional overrides | |
| for keyword, country in country_keywords_regional_overrides.items(): | |
| if keyword in text_lower: | |
| return country | |
| # 3. Check for broader regions that you want to map to "unknown" or a specific country | |
| if "north asia" in text_lower or "southeast asia" in text_lower or "east asia" in text_lower: | |
| return "unknown" | |
| return "unknown" | |
| # Get the list of English stop words from NLTK | |
| non_meaningful_pop_names = set(stopwords.words('english')) | |
| def parse_population_code_to_country(plain_text_content, table_strings): | |
| pop_code_country_map = {} | |
| pop_code_ethnicity_map = {} # NEW: To store ethnicity for structured lookup | |
| pop_code_specific_loc_map = {} # NEW: To store specific location for structured lookup | |
| # Regex for parsing population info in structured lists and general text | |
| # This pattern captures: (Pop Name/Ethnicity) (Pop Code) (Region/Specific Location) (Country) (Linguistic Family) | |
| # The 'Pop Name/Ethnicity' (Group 1) is often the ethnicity | |
| pop_info_pattern = re.compile( | |
| r'([A-Za-z\s]+?)\s+([A-Z]+\d*)\s+' # Pop Name (Group 1), Pop Code (Group 2) - Changed \d+ to \d* for codes like 'SH' | |
| r'([A-Za-z\s\(\)\-,\/]+?)\s+' # Region/Specific Location (Group 3) | |
| r'(North+|South+|West+|East+|Thailand|Laos|Cambodia|Myanmar|Philippines|Indonesia|Malaysia|China|India|Taiwan|Vietnam|Russia|Nepal|Japan|South Korea)\b' # Country (Group 4) | |
| r'(?:.*?([A-Za-z\s\-]+))?\s*' # Optional Linguistic Family (Group 5), made optional with ?, followed by optional space | |
| r'(\d+(?:\s+\d+\.?\d*)*)?', # Match all the numbers (Group 6) - made optional | |
| re.IGNORECASE | |
| ) | |
| for table_str in table_strings: | |
| table_data = parse_literal_python_list(table_str) | |
| if table_data: | |
| is_list_of_lists = bool(table_data) and isinstance(table_data[0], list) | |
| if is_list_of_lists: | |
| for row_idx, row in enumerate(table_data): | |
| row_text = " ".join(map(str, row)) | |
| match = pop_info_pattern.search(row_text) | |
| if match: | |
| pop_name = match.group(1).strip() | |
| pop_code = match.group(2).upper() | |
| specific_loc_text = match.group(3).strip() | |
| country_text = match.group(4).strip() | |
| linguistic_family = match.group(5).strip() if match.group(5) else 'unknown' | |
| final_country = get_country_from_text(country_text) | |
| if final_country == 'unknown': # Try specific loc text for country if direct country is not found | |
| final_country = get_country_from_text(specific_loc_text) | |
| if pop_code: | |
| pop_code_country_map[pop_code] = final_country | |
| # Populate ethnicity map (often Pop Name is ethnicity) | |
| pop_code_ethnicity_map[pop_code] = pop_name | |
| # Populate specific location map | |
| pop_code_specific_loc_map[pop_code] = specific_loc_text # Store as is from text | |
| else: | |
| row_text = " ".join(map(str, table_data)) | |
| match = pop_info_pattern.search(row_text) | |
| if match: | |
| pop_name = match.group(1).strip() | |
| pop_code = match.group(2).upper() | |
| specific_loc_text = match.group(3).strip() | |
| country_text = match.group(4).strip() | |
| linguistic_family = match.group(5).strip() if match.group(5) else 'unknown' | |
| final_country = get_country_from_text(country_text) | |
| if final_country == 'unknown': # Try specific loc text for country if direct country is not found | |
| final_country = get_country_from_text(specific_loc_text) | |
| if pop_code: | |
| pop_code_country_map[pop_code] = final_country | |
| # Populate ethnicity map (often Pop Name is ethnicity) | |
| pop_code_ethnicity_map[pop_code] = pop_name | |
| # Populate specific location map | |
| pop_code_specific_loc_map[pop_code] = specific_loc_text # Store as is from text | |
| # # Special case refinements for ethnicity/location if more specific rules are known from document: | |
| # if pop_name.lower() == "khon mueang": # and specific conditions if needed | |
| # pop_code_ethnicity_map[pop_code] = "Khon Mueang" | |
| # # If Khon Mueang has a specific city/district, add here | |
| # # e.g., if 'Chiang Mai' is directly linked to KM1 in a specific table | |
| # # pop_code_specific_loc_map[pop_code] = "Chiang Mai" | |
| # elif pop_name.lower() == "lawa": | |
| # pop_code_ethnicity_map[pop_code] = "Lawa" | |
| # # Add similar specific rules for other populations (e.g., Mon for MO1, MO2, MO3) | |
| # elif pop_name.lower() == "mon": | |
| # pop_code_ethnicity_map[pop_code] = "Mon" | |
| # # For MO2: "West Thailand (Thailand Myanmar border)" -> no city | |
| # # For MO3: "East Myanmar (Thailand Myanmar border)" -> no city | |
| # # If the doc gives "Bangkok" for MO4, add it here for MO4's actual specific_location. | |
| # # etc. | |
| # Fallback to parsing general plain text content (sentences) | |
| sentences = data_preprocess.extract_sentences(plain_text_content) | |
| for s in sentences: # Still focusing on just this one sentence | |
| # Use re.finditer to get all matches | |
| matches = pop_info_pattern.finditer(s) | |
| pop_name, pop_code, specific_loc_text, country_text = "unknown", "unknown", "unknown", "unknown" | |
| for match in matches: | |
| if match.group(1): | |
| pop_name = match.group(1).strip() | |
| if match.group(2): | |
| pop_code = match.group(2).upper() | |
| if match.group(3): | |
| specific_loc_text = match.group(3).strip() | |
| if match.group(4): | |
| country_text = match.group(4).strip() | |
| # linguistic_family = match.group(5).strip() if match.group(5) else 'unknown' # Already captured by pop_info_pattern | |
| final_country = get_country_from_text(country_text) | |
| if final_country == 'unknown': | |
| final_country = get_country_from_text(specific_loc_text) | |
| if pop_code.lower() not in non_meaningful_pop_names: | |
| if final_country.lower() not in non_meaningful_pop_names: | |
| pop_code_country_map[pop_code] = final_country | |
| if pop_name.lower() not in non_meaningful_pop_names: | |
| pop_code_ethnicity_map[pop_code] = pop_name # Default ethnicity from Pop Name | |
| if specific_loc_text.lower() not in non_meaningful_pop_names: | |
| pop_code_specific_loc_map[pop_code] = specific_loc_text | |
| # Specific rules for ethnicity/location in plain text: | |
| if pop_name.lower() == "khon mueang": | |
| pop_code_ethnicity_map[pop_code] = "Khon Mueang" | |
| elif pop_name.lower() == "lawa": | |
| pop_code_ethnicity_map[pop_code] = "Lawa" | |
| elif pop_name.lower() == "mon": | |
| pop_code_ethnicity_map[pop_code] = "Mon" | |
| elif pop_name.lower() == "seak": # Added specific rule for Seak | |
| pop_code_ethnicity_map[pop_code] = "Seak" | |
| elif pop_name.lower() == "nyaw": # Added specific rule for Nyaw | |
| pop_code_ethnicity_map[pop_code] = "Nyaw" | |
| elif pop_name.lower() == "nyahkur": # Added specific rule for Nyahkur | |
| pop_code_ethnicity_map[pop_code] = "Nyahkur" | |
| elif pop_name.lower() == "suay": # Added specific rule for Suay | |
| pop_code_ethnicity_map[pop_code] = "Suay" | |
| elif pop_name.lower() == "soa": # Added specific rule for Soa | |
| pop_code_ethnicity_map[pop_code] = "Soa" | |
| elif pop_name.lower() == "bru": # Added specific rule for Bru | |
| pop_code_ethnicity_map[pop_code] = "Bru" | |
| elif pop_name.lower() == "khamu": # Added specific rule for Khamu | |
| pop_code_ethnicity_map[pop_code] = "Khamu" | |
| return pop_code_country_map, pop_code_ethnicity_map, pop_code_specific_loc_map | |
| def general_parse_population_code_to_country(plain_text_content, table_strings): | |
| pop_code_country_map = {} | |
| pop_code_ethnicity_map = {} | |
| pop_code_specific_loc_map = {} | |
| sample_id_to_pop_code = {} | |
| for table_str in table_strings: | |
| table_data = parse_literal_python_list(table_str) | |
| if not table_data or not isinstance(table_data[0], list): | |
| continue | |
| header_row = [col.lower() for col in table_data[0]] | |
| header_map = {col: idx for idx, col in enumerate(header_row)} | |
| # MJ17: Direct PopCode → Country | |
| if 'id' in header_map and 'country' in header_map: | |
| for row in table_strings[1:]: | |
| row = parse_literal_python_list(row)[0] | |
| if len(row) < len(header_row): | |
| continue | |
| pop_code = str(row[header_map['id']]).strip() | |
| country = str(row[header_map['country']]).strip() | |
| province = row[header_map['province']].strip() if 'province' in header_map else 'unknown' | |
| pop_group = row[header_map['population group / region']].strip() if 'population group / region' in header_map else 'unknown' | |
| pop_code_country_map[pop_code] = country | |
| pop_code_specific_loc_map[pop_code] = province | |
| pop_code_ethnicity_map[pop_code] = pop_group | |
| # A1YU101 or EBK/KSK: SampleID → PopCode | |
| elif 'sample id' in header_map and 'population code' in header_map: | |
| for row in table_strings[1:]: | |
| row = parse_literal_python_list(row)[0] | |
| if len(row) < 2: | |
| continue | |
| sample_id = row[header_map['sample id']].strip().upper() | |
| pop_code = row[header_map['population code']].strip().upper() | |
| sample_id_to_pop_code[sample_id] = pop_code | |
| # PopCode → Country (A1YU101/EBK mapping) | |
| elif 'population code' in header_map and 'country' in header_map: | |
| for row in table_strings[1:]: | |
| row = parse_literal_python_list(row)[0] | |
| if len(row) < 2: | |
| continue | |
| pop_code = row[header_map['population code']].strip().upper() | |
| country = row[header_map['country']].strip() | |
| pop_code_country_map[pop_code] = country | |
| return pop_code_country_map, pop_code_ethnicity_map, pop_code_specific_loc_map, sample_id_to_pop_code | |
| def chunk_text(text, chunk_size=500, overlap=50): | |
| """Splits text into chunks (by words) with overlap.""" | |
| chunks = [] | |
| words = text.split() | |
| num_words = len(words) | |
| start = 0 | |
| while start < num_words: | |
| end = min(start + chunk_size, num_words) | |
| chunk = " ".join(words[start:end]) | |
| chunks.append(chunk) | |
| if end == num_words: | |
| break | |
| start += chunk_size - overlap # Move start by (chunk_size - overlap) | |
| return chunks | |
| def build_vector_index_and_data(doc_path, index_path="faiss_index.bin", chunks_path="document_chunks.json", structured_path="structured_lookup.json"): | |
| """ | |
| Reads document, builds structured lookup, chunks remaining text, embeds chunks, | |
| and builds/saves a FAISS index. | |
| """ | |
| print("Step 1: Reading document and extracting structured data...") | |
| # plain_text_content, table_strings, document_title = read_docx_text(doc_path) # Get document_title here | |
| # sample_id_map, contiguous_ranges_data = parse_sample_id_to_population_code(plain_text_content) | |
| # pop_code_to_country, pop_code_to_ethnicity, pop_code_to_specific_loc = parse_population_code_to_country(plain_text_content, table_strings) | |
| # master_structured_lookup = {} | |
| # master_structured_lookup['document_title'] = document_title # Store document title | |
| # master_structured_lookup['sample_id_map'] = sample_id_map | |
| # master_structured_lookup['contiguous_ranges'] = dict(contiguous_ranges_data) | |
| # master_structured_lookup['pop_code_to_country'] = pop_code_to_country | |
| # master_structured_lookup['pop_code_to_ethnicity'] = pop_code_to_ethnicity # NEW: Store pop_code to ethnicity map | |
| # master_structured_lookup['pop_code_to_specific_loc'] = pop_code_to_specific_loc # NEW: Store pop_code to specific_loc map | |
| # # Final consolidation: Use sample_id_map to derive full info for queries | |
| # final_structured_entries = {} | |
| # for sample_id, pop_code in master_structured_lookup['sample_id_map'].items(): | |
| # country = master_structured_lookup['pop_code_to_country'].get(pop_code, 'unknown') | |
| # ethnicity = master_structured_lookup['pop_code_to_ethnicity'].get(pop_code, 'unknown') # Retrieve ethnicity | |
| # specific_location = master_structured_lookup['pop_code_to_specific_loc'].get(pop_code, 'unknown') # Retrieve specific location | |
| # final_structured_entries[sample_id] = { | |
| # 'population_code': pop_code, | |
| # 'country': country, | |
| # 'type': 'modern', | |
| # 'ethnicity': ethnicity, # Store ethnicity | |
| # 'specific_location': specific_location # Store specific location | |
| # } | |
| # master_structured_lookup['final_structured_entries'] = final_structured_entries | |
| plain_text_content, table_strings, document_title = read_docx_text(doc_path) | |
| pop_code_to_country, pop_code_to_ethnicity, pop_code_to_specific_loc, sample_id_map = general_parse_population_code_to_country(plain_text_content, table_strings) | |
| final_structured_entries = {} | |
| if sample_id_map: | |
| for sample_id, pop_code in sample_id_map.items(): | |
| country = pop_code_to_country.get(pop_code, 'unknown') | |
| ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown') | |
| specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown') | |
| final_structured_entries[sample_id] = { | |
| 'population_code': pop_code, | |
| 'country': country, | |
| 'type': 'modern', | |
| 'ethnicity': ethnicity, | |
| 'specific_location': specific_loc | |
| } | |
| else: | |
| for pop_code in pop_code_to_country.keys(): | |
| country = pop_code_to_country.get(pop_code, 'unknown') | |
| ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown') | |
| specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown') | |
| final_structured_entries[pop_code] = { | |
| 'population_code': pop_code, | |
| 'country': country, | |
| 'type': 'modern', | |
| 'ethnicity': ethnicity, | |
| 'specific_location': specific_loc | |
| } | |
| if not final_structured_entries: | |
| # traditional way of A1YU101 | |
| sample_id_map, contiguous_ranges_data = parse_sample_id_to_population_code(plain_text_content) | |
| pop_code_to_country, pop_code_to_ethnicity, pop_code_to_specific_loc = parse_population_code_to_country(plain_text_content, table_strings) | |
| if sample_id_map: | |
| for sample_id, pop_code in sample_id_map.items(): | |
| country = pop_code_to_country.get(pop_code, 'unknown') | |
| ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown') | |
| specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown') | |
| final_structured_entries[sample_id] = { | |
| 'population_code': pop_code, | |
| 'country': country, | |
| 'type': 'modern', | |
| 'ethnicity': ethnicity, | |
| 'specific_location': specific_loc | |
| } | |
| else: | |
| for pop_code in pop_code_to_country.keys(): | |
| country = pop_code_to_country.get(pop_code, 'unknown') | |
| ethnicity = pop_code_to_ethnicity.get(pop_code, 'unknown') | |
| specific_loc = pop_code_to_specific_loc.get(pop_code, 'unknown') | |
| final_structured_entries[pop_code] = { | |
| 'population_code': pop_code, | |
| 'country': country, | |
| 'type': 'modern', | |
| 'ethnicity': ethnicity, | |
| 'specific_location': specific_loc | |
| } | |
| master_lookup = { | |
| 'document_title': document_title, | |
| 'pop_code_to_country': pop_code_to_country, | |
| 'pop_code_to_ethnicity': pop_code_to_ethnicity, | |
| 'pop_code_to_specific_loc': pop_code_to_specific_loc, | |
| 'sample_id_map': sample_id_map, | |
| 'final_structured_entries': final_structured_entries | |
| } | |
| print(f"Structured lookup built with {len(final_structured_entries)} entries in 'final_structured_entries'.") | |
| with open(structured_path, 'w') as f: | |
| json.dump(master_lookup, f, indent=4) | |
| print(f"Structured lookup saved to {structured_path}.") | |
| print("Step 2: Chunking document for RAG vector index...") | |
| # replace the chunk here with the all_output from process_inputToken and fallback to this traditional chunk | |
| clean_text, clean_table = "", "" | |
| if plain_text_content: | |
| clean_text = data_preprocess.normalize_for_overlap(plain_text_content) | |
| if table_strings: | |
| clean_table = data_preprocess.normalize_for_overlap(". ".join(table_strings)) | |
| all_clean_chunk = clean_text + clean_table | |
| document_chunks = chunk_text(all_clean_chunk) | |
| print(f"Document chunked into {len(document_chunks)} chunks.") | |
| print("Step 3: Generating embeddings for chunks (this might take time and cost API calls)...") | |
| embedding_model_for_chunks = genai.GenerativeModel('models/text-embedding-004') | |
| chunk_embeddings = [] | |
| for i, chunk in enumerate(document_chunks): | |
| embedding = get_embedding(chunk, task_type="RETRIEVAL_DOCUMENT") | |
| if embedding is not None and embedding.shape[0] > 0: | |
| chunk_embeddings.append(embedding) | |
| else: | |
| print(f"Warning: Failed to get valid embedding for chunk {i}. Skipping.") | |
| chunk_embeddings.append(np.zeros(768, dtype='float32')) | |
| if not chunk_embeddings: | |
| raise ValueError("No valid embeddings generated. Check get_embedding function and API.") | |
| embedding_dimension = chunk_embeddings[0].shape[0] | |
| index = faiss.IndexFlatL2(embedding_dimension) | |
| index.add(np.array(chunk_embeddings)) | |
| faiss.write_index(index, index_path) | |
| with open(chunks_path, "w") as f: | |
| json.dump(document_chunks, f) | |
| print(f"FAISS index built and saved to {index_path}.") | |
| print(f"Document chunks saved to {chunks_path}.") | |
| return master_lookup, index, document_chunks, all_clean_chunk | |
| def load_rag_assets(index_path="faiss_index.bin", chunks_path="document_chunks.json", structured_path="structured_lookup.json"): | |
| """Loads pre-built RAG assets (FAISS index, chunks, structured lookup).""" | |
| print("Loading RAG assets...") | |
| master_structured_lookup = {} | |
| if os.path.exists(structured_path): | |
| with open(structured_path, 'r') as f: | |
| master_structured_lookup = json.load(f) | |
| print("Structured lookup loaded.") | |
| else: | |
| print("Structured lookup file not found. Rebuilding is likely needed.") | |
| index = None | |
| chunks = [] | |
| if os.path.exists(index_path) and os.path.exists(chunks_path): | |
| try: | |
| index = faiss.read_index(index_path) | |
| with open(chunks_path, "r") as f: | |
| chunks = json.load(f) | |
| print("FAISS index and chunks loaded.") | |
| except Exception as e: | |
| print(f"Error loading FAISS index or chunks: {e}. Will rebuild.") | |
| index = None | |
| chunks = [] | |
| else: | |
| print("FAISS index or chunks files not found.") | |
| return master_structured_lookup, index, chunks | |
| # Helper function for query_document_info | |
| def exactInContext(text, keyword): | |
| # try keyword_prfix | |
| # code_pattern = re.compile(r'([A-Z0-9]+?)(\d+)$', re.IGNORECASE) | |
| # # Attempt to parse the keyword into its prefix and numerical part using re.search | |
| # keyword_match = code_pattern.search(keyword) | |
| # keyword_prefix = None | |
| # keyword_num = None | |
| # if keyword_match: | |
| # keyword_prefix = keyword_match.group(1).lower() | |
| # keyword_num = int(keyword_match.group(2)) | |
| text = text.lower() | |
| idx = text.find(keyword.lower()) | |
| if idx == -1: | |
| # if keyword_prefix: | |
| # idx = text.find(keyword_prefix) | |
| # if idx == -1: | |
| # return False | |
| return False | |
| return True | |
| def chooseContextLLM(contexts, kw): | |
| # if kw in context | |
| for con in contexts: | |
| context = contexts[con] | |
| if context: | |
| if exactInContext(context, kw): | |
| return con, context | |
| #if cannot find anything related to kw in context, return all output | |
| if contexts["all_output"]: | |
| return "all_output", contexts["all_output"] | |
| else: | |
| # if all_output not exist | |
| # look of chunk and still not exist return document chunk | |
| if contexts["chunk"]: return "chunk", contexts["chunk"] | |
| elif contexts["document_chunk"]: return "document_chunk", contexts["document_chunk"] | |
| else: return None, None | |
| def clean_llm_output(llm_response_text, output_format_str): | |
| results = [] | |
| lines = llm_response_text.strip().split('\n') | |
| output_country, output_type, output_ethnicity, output_specific_location = [],[],[],[] | |
| for line in lines: | |
| extracted_country, extracted_type, extracted_ethnicity, extracted_specific_location = "unknown", "unknown", "unknown", "unknown" | |
| line = line.strip() | |
| if output_format_str == "ethnicity, specific_location/unknown": # Targeted RAG output | |
| parsed_output = re.search(r'^\s*([^,]+?),\s*(.+?)\s*$', llm_response_text) | |
| if parsed_output: | |
| extracted_ethnicity = parsed_output.group(1).strip() | |
| extracted_specific_location = parsed_output.group(2).strip() | |
| else: | |
| print(" DEBUG: LLM did not follow expected 2-field format for targeted RAG. Defaulting to unknown for ethnicity/specific_location.") | |
| extracted_ethnicity = 'unknown' | |
| extracted_specific_location = 'unknown' | |
| elif output_format_str == "modern/ancient/unknown, ethnicity, specific_location/unknown": | |
| parsed_output = re.search(r'^\s*([^,]+?),\s*([^,]+?),\s*(.+?)\s*$', llm_response_text) | |
| if parsed_output: | |
| extracted_type = parsed_output.group(1).strip() | |
| extracted_ethnicity = parsed_output.group(2).strip() | |
| extracted_specific_location = parsed_output.group(3).strip() | |
| else: | |
| # Fallback: check if only 2 fields | |
| parsed_output_2_fields = re.search(r'^\s*([^,]+?),\s*([^,]+?)\s*$', llm_response_text) | |
| if parsed_output_2_fields: | |
| extracted_type = parsed_output_2_fields.group(1).strip() | |
| extracted_ethnicity = parsed_output_2_fields.group(2).strip() | |
| extracted_specific_location = 'unknown' | |
| else: | |
| # even simpler fallback: 1 field only | |
| parsed_output_1_field = re.search(r'^\s*([^,]+?)\s*$', llm_response_text) | |
| if parsed_output_1_field: | |
| extracted_type = parsed_output_1_field.group(1).strip() | |
| extracted_ethnicity = 'unknown' | |
| extracted_specific_location = 'unknown' | |
| else: | |
| print(" DEBUG: LLM did not follow any expected simplified format. Attempting verbose parsing fallback.") | |
| type_match_fallback = re.search(r'Type:\s*([A-Za-z\s-]+)', llm_response_text) | |
| extracted_type = type_match_fallback.group(1).strip() if type_match_fallback else 'unknown' | |
| extracted_ethnicity = 'unknown' | |
| extracted_specific_location = 'unknown' | |
| else: | |
| parsed_output = re.search(r'^\s*([^,]+?),\s*([^,]+?),\s*([^,]+?),\s*(.+?)\s*$', line) | |
| if parsed_output: | |
| extracted_country = parsed_output.group(1).strip() | |
| extracted_type = parsed_output.group(2).strip() | |
| extracted_ethnicity = parsed_output.group(3).strip() | |
| extracted_specific_location = parsed_output.group(4).strip() | |
| else: | |
| print(f" DEBUG: Line did not follow expected 4-field format: {line}") | |
| parsed_output_2_fields = re.search(r'^\s*([^,]+?),\s*([^,]+?)\s*$', line) | |
| if parsed_output_2_fields: | |
| extracted_country = parsed_output_2_fields.group(1).strip() | |
| extracted_type = parsed_output_2_fields.group(2).strip() | |
| extracted_ethnicity = 'unknown' | |
| extracted_specific_location = 'unknown' | |
| else: | |
| print(f" DEBUG: Fallback to verbose-style parsing: {line}") | |
| country_match_fallback = re.search(r'Country:\s*([A-Za-z\s-]+)', line) | |
| type_match_fallback = re.search(r'Type:\s*([A-Za-z\s-]+)', line) | |
| extracted_country = country_match_fallback.group(1).strip() if country_match_fallback else 'unknown' | |
| extracted_type = type_match_fallback.group(1).strip() if type_match_fallback else 'unknown' | |
| extracted_ethnicity = 'unknown' | |
| extracted_specific_location = 'unknown' | |
| results.append({ | |
| "country": extracted_country, | |
| "type": extracted_type, | |
| "ethnicity": extracted_ethnicity, | |
| "specific_location": extracted_specific_location | |
| #"country_explain":extracted_country_explain, | |
| #"type_explain": extracted_type_explain | |
| }) | |
| # if more than 2 results | |
| if output_format_str == "ethnicity, specific_location/unknown": | |
| for result in results: | |
| if result["ethnicity"] not in output_ethnicity: | |
| output_ethnicity.append(result["ethnicity"]) | |
| if result["specific_location"] not in output_specific_location: | |
| output_specific_location.append(result["specific_location"]) | |
| return " or ".join(output_ethnicity), " or ".join(output_specific_location) | |
| elif output_format_str == "modern/ancient/unknown, ethnicity, specific_location/unknown": | |
| for result in results: | |
| if result["type"] not in output_type: | |
| output_type.append(result["type"]) | |
| if result["ethnicity"] not in output_ethnicity: | |
| output_ethnicity.append(result["ethnicity"]) | |
| if result["specific_location"] not in output_specific_location: | |
| output_specific_location.append(result["specific_location"]) | |
| return " or ".join(output_type)," or ".join(output_ethnicity), " or ".join(output_specific_location) | |
| else: | |
| for result in results: | |
| if result["country"] not in output_country: | |
| output_country.append(result["country"]) | |
| if result["type"] not in output_type: | |
| output_type.append(result["type"]) | |
| if result["ethnicity"] not in output_ethnicity: | |
| output_ethnicity.append(result["ethnicity"]) | |
| if result["specific_location"] not in output_specific_location: | |
| output_specific_location.append(result["specific_location"]) | |
| return " or ".join(output_country)," or ".join(output_type)," or ".join(output_ethnicity), " or ".join(output_specific_location) | |
| # def parse_multi_sample_llm_output(raw_response: str, output_format_str): | |
| # """ | |
| # Parse LLM output with possibly multiple metadata lines + shared explanations. | |
| # """ | |
| # lines = [line.strip() for line in raw_response.strip().splitlines() if line.strip()] | |
| # metadata_list = [] | |
| # explanation_lines = [] | |
| # if output_format_str == "country_name, modern/ancient/unknown": | |
| # parts = [x.strip() for x in lines[0].split(",")] | |
| # if len(parts)==2: | |
| # metadata_list.append({ | |
| # "country": parts[0], | |
| # "sample_type": parts[1]#, | |
| # #"ethnicity": parts[2], | |
| # #"location": parts[3] | |
| # }) | |
| # if 1<len(lines): | |
| # line = lines[1] | |
| # if "\n" in line: line = line.split("\n") | |
| # if ". " in line: line = line.split(". ") | |
| # if isinstance(line,str): line = [line] | |
| # explanation_lines += line | |
| # elif output_format_str == "modern/ancient/unknown": | |
| # metadata_list.append({ | |
| # "country": "unknown", | |
| # "sample_type": lines[0]#, | |
| # #"ethnicity": parts[2], | |
| # #"location": parts[3] | |
| # }) | |
| # explanation_lines.append(lines[1]) | |
| # # Assign explanations (optional) to each sample — same explanation reused | |
| # for md in metadata_list: | |
| # md["country_explanation"] = None | |
| # md["sample_type_explanation"] = None | |
| # if md["country"].lower() != "unknown" and len(explanation_lines) >= 1: | |
| # md["country_explanation"] = explanation_lines[0] | |
| # if md["sample_type"].lower() != "unknown": | |
| # if len(explanation_lines) >= 2: | |
| # md["sample_type_explanation"] = explanation_lines[1] | |
| # elif len(explanation_lines) == 1 and md["country"].lower() == "unknown": | |
| # md["sample_type_explanation"] = explanation_lines[0] | |
| # elif len(explanation_lines) == 1: | |
| # md["sample_type_explanation"] = explanation_lines[0] | |
| # return metadata_list | |
| def parse_multi_sample_llm_output(raw_response: str, output_format_str): | |
| """ | |
| Parse LLM output with possibly multiple metadata lines + shared explanations. | |
| """ | |
| metadata_list = {} | |
| explanation_lines = [] | |
| output_answers = raw_response.split("\n")[0].split(", ") | |
| explanation_lines = [x for x in raw_response.split("\n")[1:] if x.strip()] | |
| print("raw explanation line which split by new line: ", explanation_lines) | |
| if len(explanation_lines) == 1: | |
| if len(explanation_lines[0].split(". ")) > len(explanation_lines): | |
| explanation_lines = [x for x in explanation_lines[0].split(". ") if x.strip()] | |
| print("explain line split by dot: ", explanation_lines) | |
| output_formats = output_format_str.split(", ") | |
| explain = "" | |
| # assign output format to its output answer and explanation | |
| if output_format_str: | |
| outputs = output_format_str.split(", ") | |
| for o in range(len(outputs)): | |
| output = outputs[o] | |
| metadata_list[output] = {"answer":"", | |
| output+"_explanation":""} | |
| # assign output answers | |
| if o < len(output_answers): | |
| # check if output_format unexpectedly in the answer such as: | |
| #country_name: Europe, modern/ancient: modern | |
| try: | |
| if ": " in output_answers[o]: | |
| output_answers[o] = output_answers[o].split(": ")[1] | |
| except: | |
| pass | |
| # Europe, modern | |
| metadata_list[output]["answer"] = output_answers[o] | |
| if "unknown" in metadata_list[output]["answer"].lower(): | |
| metadata_list[output]["answer"] = "unknown" | |
| else: | |
| metadata_list[output]["answer"] = "unknown" | |
| # assign explanations | |
| if metadata_list[output]["answer"] != "unknown": | |
| if explanation_lines: | |
| explain = explanation_lines.pop(0) | |
| metadata_list[output][output+"_explanation"] = explain | |
| else: | |
| metadata_list[output][output+"_explanation"] = "unknown" | |
| return metadata_list | |
| def merge_metadata_outputs(metadata_list): | |
| """ | |
| Merge a list of metadata dicts into one, combining differing values with 'or'. | |
| Assumes all dicts have the same keys. | |
| """ | |
| if not metadata_list: | |
| return {} | |
| merged = {} | |
| keys = metadata_list[0].keys() | |
| for key in keys: | |
| values = [md[key] for md in metadata_list if key in md] | |
| unique_values = list(dict.fromkeys(values)) # preserve order, remove dupes | |
| if "unknown" in unique_values: | |
| unique_values.pop(unique_values.index("unknown")) | |
| if len(unique_values) == 1: | |
| merged[key] = unique_values[0] | |
| else: | |
| merged[key] = " or ".join(unique_values) | |
| return merged | |
| def query_document_info(query_word, alternative_query_word, metadata, master_structured_lookup, faiss_index, document_chunks, llm_api_function, chunk=None, all_output=None, model_ai=None): | |
| """ | |
| Queries the document using a hybrid approach: | |
| 1. Local structured lookup (fast, cheap, accurate for known patterns). | |
| 2. RAG with semantic search and LLM (general, flexible, cost-optimized). | |
| """ | |
| if model_ai: | |
| if model_ai == "gemini-1.5-flash-latest": | |
| genai.configure(api_key=os.getenv("GOOGLE_API_KEY")) | |
| PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens | |
| PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens | |
| PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens | |
| global_llm_model_for_counting_tokens = genai.GenerativeModel("gemini-1.5-flash-latest")#('gemini-1.5-flash-latest') | |
| else: | |
| genai.configure(api_key=os.getenv("GOOGLE_API_KEY_BACKUP")) | |
| # Gemini 2.5 Flash-Lite pricing per 1,000 tokens | |
| PRICE_PER_1K_INPUT_LLM = 0.00010 # $0.10 per 1M input tokens | |
| PRICE_PER_1K_OUTPUT_LLM = 0.00040 # $0.40 per 1M output tokens | |
| # Embedding-001 pricing per 1,000 input tokens | |
| PRICE_PER_1K_EMBEDDING_INPUT = 0.00015 # $0.15 per 1M input tokens | |
| global_llm_model_for_counting_tokens = genai.GenerativeModel("gemini-2.5-flash-lite")#('gemini-1.5-flash-latest') | |
| if metadata: | |
| extracted_country, extracted_specific_location, extracted_ethnicity, extracted_type = metadata["country"], metadata["specific_location"], metadata["ethnicity"], metadata["sample_type"] | |
| extracted_col_date, extracted_iso, extracted_title, extracted_features = metadata["collection_date"], metadata["isolate"], metadata["title"], metadata["all_features"] | |
| else: | |
| extracted_country, extracted_specific_location, extracted_ethnicity, extracted_type = "unknown", "unknown", "unknown", "unknown" | |
| extracted_col_date, extracted_iso, extracted_title = "unknown", "unknown", "unknown" | |
| # --- NEW: Pre-process alternative_query_word to remove '.X' suffix if present --- | |
| if alternative_query_word: | |
| alternative_query_word_cleaned = alternative_query_word.split('.')[0] | |
| else: | |
| alternative_query_word_cleaned = alternative_query_word | |
| country_explanation, sample_type_explanation = None, None | |
| # Use the consolidated final_structured_entries for direct lookup | |
| final_structured_entries = master_structured_lookup.get('final_structured_entries', {}) | |
| document_title = master_structured_lookup.get('document_title', 'Unknown Document Title') # Retrieve document title | |
| # Default values for all extracted fields. These will be updated. | |
| method_used = 'unknown' # Will be updated based on the method that yields a result | |
| population_code_from_sl = 'unknown' # To pass to RAG prompt if available | |
| total_query_cost = 0 | |
| # Attempt 1: Try primary query_word (e.g., isolate name) with structured lookup | |
| try: | |
| print("try attempt 1 in model query") | |
| structured_info = final_structured_entries.get(query_word.upper()) | |
| if structured_info: | |
| if extracted_country == 'unknown': | |
| extracted_country = structured_info['country'] | |
| if extracted_type == 'unknown': | |
| extracted_type = structured_info['type'] | |
| # if extracted_ethnicity == 'unknown': | |
| # extracted_ethnicity = structured_info.get('ethnicity', 'unknown') # Get ethnicity from structured lookup | |
| # if extracted_specific_location == 'unknown': | |
| # extracted_specific_location = structured_info.get('specific_location', 'unknown') # Get specific_location from structured lookup | |
| population_code_from_sl = structured_info['population_code'] | |
| method_used = "structured_lookup_direct" | |
| print(f"'{query_word}' found in structured lookup (direct match).") | |
| except: | |
| print("pass attempt 1 in model query") | |
| pass | |
| # Attempt 2: Try primary query_word with heuristic range lookup if direct fails (only if not already resolved) | |
| try: | |
| print("try attempt 2 in model query") | |
| if method_used == 'unknown': | |
| query_prefix, query_num_str = _parse_individual_code_parts(query_word) | |
| if query_prefix is not None and query_num_str is not None: | |
| try: query_num = int(query_num_str) | |
| except ValueError: query_num = None | |
| if query_num is not None: | |
| query_prefix_upper = query_prefix.upper() | |
| contiguous_ranges = master_structured_lookup.get('contiguous_ranges', defaultdict(list)) | |
| pop_code_to_country = master_structured_lookup.get('pop_code_to_country', {}) | |
| pop_code_to_ethnicity = master_structured_lookup.get('pop_code_to_ethnicity', {}) | |
| pop_code_to_specific_loc = master_structured_lookup.get('pop_code_to_specific_loc', {}) | |
| if query_prefix_upper in contiguous_ranges: | |
| for start_num, end_num, pop_code_for_range in contiguous_ranges[query_prefix_upper]: | |
| if start_num <= query_num <= end_num: | |
| country_from_heuristic = pop_code_to_country.get(pop_code_for_range, 'unknown') | |
| if country_from_heuristic != 'unknown': | |
| if extracted_country == 'unknown': | |
| extracted_country = country_from_heuristic | |
| if extracted_type == 'unknown': | |
| extracted_type = 'modern' | |
| # if extracted_ethnicity == 'unknown': | |
| # extracted_ethnicity = pop_code_to_ethnicity.get(pop_code_for_range, 'unknown') | |
| # if extracted_specific_location == 'unknown': | |
| # extracted_specific_location = pop_code_to_specific_loc.get(pop_code_for_range, 'unknown') | |
| population_code_from_sl = pop_code_for_range | |
| method_used = "structured_lookup_heuristic_range_match" | |
| print(f"'{query_word}' not direct. Heuristic: Falls within range {query_prefix_upper}{start_num}-{query_prefix_upper}{end_num}.") | |
| break | |
| else: | |
| print(f"'{query_word}' heuristic match found, but country unknown. Will fall to RAG below.") | |
| except: | |
| print("pass attempt 2 in model query") | |
| pass | |
| # Attempt 3: If primary query_word failed all structured lookups, try alternative_query_word (cleaned) | |
| try: | |
| print("try attempt 3 in model query") | |
| if method_used == 'unknown' and alternative_query_word_cleaned and alternative_query_word_cleaned != query_word: | |
| print(f"'{query_word}' not found in structured (or heuristic). Trying alternative '{alternative_query_word_cleaned}'.") | |
| # Try direct lookup for alternative word | |
| structured_info_alt = final_structured_entries.get(alternative_query_word_cleaned.upper()) | |
| if structured_info_alt: | |
| if extracted_country == 'unknown': | |
| extracted_country = structured_info_alt['country'] | |
| if extracted_type == 'unknown': | |
| extracted_type = structured_info_alt['type'] | |
| # if extracted_ethnicity == 'unknown': | |
| # extracted_ethnicity = structured_info_alt.get('ethnicity', 'unknown') | |
| # if extracted_specific_location == 'unknown': | |
| # extracted_specific_location = structured_info_alt.get('specific_location', 'unknown') | |
| population_code_from_sl = structured_info_alt['population_code'] | |
| method_used = "structured_lookup_alt_direct" | |
| print(f"Alternative '{alternative_query_word_cleaned}' found in structured lookup (direct match).") | |
| else: | |
| # Try heuristic lookup for alternative word | |
| alt_prefix, alt_num_str = _parse_individual_code_parts(alternative_query_word_cleaned) | |
| if alt_prefix is not None and alt_num_str is not None: | |
| try: alt_num = int(alt_num_str) | |
| except ValueError: alt_num = None | |
| if alt_num is not None: | |
| alt_prefix_upper = alt_prefix.upper() | |
| contiguous_ranges = master_structured_lookup.get('contiguous_ranges', defaultdict(list)) | |
| pop_code_to_country = master_structured_lookup.get('pop_code_to_country', {}) | |
| pop_code_to_ethnicity = master_structured_lookup.get('pop_code_to_ethnicity', {}) | |
| pop_code_to_specific_loc = master_structured_lookup.get('pop_code_to_specific_loc', {}) | |
| if alt_prefix_upper in contiguous_ranges: | |
| for start_num, end_num, pop_code_for_range in contiguous_ranges[alt_prefix_upper]: | |
| if start_num <= alt_num <= end_num: | |
| country_from_heuristic_alt = pop_code_to_country.get(pop_code_for_range, 'unknown') | |
| if country_from_heuristic_alt != 'unknown': | |
| if extracted_country == 'unknown': | |
| extracted_country = country_from_heuristic_alt | |
| if extracted_type == 'unknown': | |
| extracted_type = 'modern' | |
| # if extracted_ethnicity == 'unknown': | |
| # extracted_ethnicity = pop_code_to_ethnicity.get(pop_code_for_range, 'unknown') | |
| # if extracted_specific_location == 'unknown': | |
| # extracted_specific_location = pop_code_to_specific_loc.get(pop_code_for_range, 'unknown') | |
| population_code_from_sl = pop_code_for_range | |
| method_used = "structured_lookup_alt_heuristic_range_match" | |
| break | |
| else: | |
| print(f"Alternative '{alternative_query_word_cleaned}' heuristic match found, but country unknown. Will fall to RAG below.") | |
| except: | |
| print("pass attempt 3 in model query") | |
| pass | |
| # use the context_for_llm to detect present_ancient before using llm model | |
| # retrieved_chunks_text = [] | |
| # if document_chunks: | |
| # for idx in range(len(document_chunks)): | |
| # retrieved_chunks_text.append(document_chunks[idx]) | |
| # context_for_llm = "" | |
| # all_context = "\n".join(retrieved_chunks_text) # | |
| # listOfcontexts = {"chunk": chunk, | |
| # "all_output": all_output, | |
| # "document_chunk": all_context} | |
| # label, context_for_llm = chooseContextLLM(listOfcontexts, query_word) | |
| # if not context_for_llm: | |
| # label, context_for_llm = chooseContextLLM(listOfcontexts, alternative_query_word_cleaned) | |
| # if not context_for_llm: | |
| # context_for_llm = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + extracted_features | |
| # if context_for_llm: | |
| # extracted_type, explain = mtdna_classifier.detect_ancient_flag(context_for_llm) | |
| # extracted_type = extracted_type.lower() | |
| # sample_type_explanation = explain | |
| # 5. Execute RAG if needed (either full RAG or targeted RAG for missing fields) | |
| # Determine if a RAG call is necessary | |
| # run_rag = (extracted_country == 'unknown' or extracted_type == 'unknown')# or \ | |
| # #extracted_ethnicity == 'unknown' or extracted_specific_location == 'unknown') | |
| run_rag = True | |
| if run_rag: | |
| print("try run rag") | |
| # Determine the phrase for LLM query | |
| rag_query_phrase = f"'{query_word}'" | |
| if alternative_query_word_cleaned and alternative_query_word_cleaned != query_word: | |
| rag_query_phrase += f" or its alternative word '{alternative_query_word_cleaned}'" | |
| # Construct a more specific semantic query phrase for embedding if structured info is available | |
| semantic_query_for_embedding = rag_query_phrase # Default | |
| # if extracted_country != 'unknown': # If country is known from structured lookup (for targeted RAG) | |
| # if population_code_from_sl != 'unknown': | |
| # semantic_query_for_embedding = f"ethnicity and specific location for {query_word} population {population_code_from_sl} in {extracted_country}" | |
| # else: # If pop_code not found in structured, still use country hint | |
| # semantic_query_for_embedding = f"ethnicity and specific location for {query_word} in {extracted_country}" | |
| # print(f" DEBUG: Semantic query for embedding: '{semantic_query_for_embedding}'") | |
| # Determine fields to ask LLM for and output format based on what's known/needed | |
| prompt_instruction_prefix = "" | |
| output_format_str = "" | |
| # Determine if it's a full RAG or targeted RAG scenario based on what's already extracted | |
| is_full_rag_scenario = True#(extracted_country == 'unknown') | |
| if is_full_rag_scenario: # Full RAG scenario | |
| output_format_str = "country_name, modern/ancient/unknown"#, ethnicity, specific_location/unknown" | |
| method_used = "rag_llm" | |
| print(f"Proceeding to FULL RAG for {rag_query_phrase}.") | |
| # else: # Targeted RAG scenario (country/type already known, need ethnicity/specific_location) | |
| # if extracted_type == "unknown": | |
| # prompt_instruction_prefix = ( | |
| # f"I already know the country is {extracted_country}. " | |
| # f"{f'The population code is {population_code_from_sl}. ' if population_code_from_sl != 'unknown' else ''}" | |
| # ) | |
| # #output_format_str = "modern/ancient/unknown, ethnicity, specific_location/unknown" | |
| # output_format_str = "modern/ancient/unknown" | |
| # # else: | |
| # # prompt_instruction_prefix = ( | |
| # # f"I already know the country is {extracted_country} and the sample type is {extracted_type}. " | |
| # # f"{f'The population code is {population_code_from_sl}. ' if population_code_from_sl != 'unknown' else ''}" | |
| # # ) | |
| # # output_format_str = "ethnicity, specific_location/unknown" | |
| # method_used = "hybrid_sl_rag" | |
| # print(f"Proceeding to TARGETED RAG for {rag_query_phrase}.") | |
| # Calculate embedding cost for the primary query word | |
| current_embedding_cost = 0 | |
| try: | |
| query_embedding_vector = get_embedding(semantic_query_for_embedding, task_type="RETRIEVAL_QUERY") | |
| query_embedding_tokens = global_llm_model_for_counting_tokens.count_tokens(semantic_query_for_embedding).total_tokens | |
| current_embedding_cost += (query_embedding_tokens / 1000) * PRICE_PER_1K_EMBEDDING_INPUT | |
| print(f" DEBUG: Query embedding tokens (for '{semantic_query_for_embedding}'): {query_embedding_tokens}, cost: ${current_embedding_cost:.6f}") | |
| if alternative_query_word_cleaned and alternative_query_word_cleaned != query_word: | |
| alt_embedding_vector = get_embedding(alternative_query_word_cleaned, task_type="RETRIEVAL_QUERY") | |
| alt_embedding_tokens = global_llm_model_for_counting_tokens.count_tokens(alternative_query_word_cleaned).total_tokens | |
| current_embedding_cost += (alt_embedding_tokens / 1000) * PRICE_PER_1K_EMBEDDING_INPUT | |
| print(f" DEBUG: Alternative query ('{alternative_query_word_cleaned}') embedding tokens: {alt_embedding_tokens}, cost: ${current_embedding_cost:.6f}") | |
| except Exception as e: | |
| print(f"Error getting query embedding for RAG: {e}") | |
| return extracted_country, extracted_type, "embedding_failed", extracted_ethnicity, extracted_specific_location, total_query_cost | |
| if query_embedding_vector is None or query_embedding_vector.shape[0] == 0: | |
| return extracted_country, extracted_type, "embedding_failed", extracted_ethnicity, extracted_specific_location, total_query_cost | |
| D, I = faiss_index.search(np.array([query_embedding_vector]), 4) | |
| retrieved_chunks_text = [] | |
| for idx in I[0]: | |
| if 0 <= idx < len(document_chunks): | |
| retrieved_chunks_text.append(document_chunks[idx]) | |
| context_for_llm = "" | |
| all_context = "\n".join(retrieved_chunks_text) # | |
| listOfcontexts = {"chunk": chunk, | |
| "all_output": all_output, | |
| "document_chunk": all_context} | |
| label, context_for_llm = chooseContextLLM(listOfcontexts, query_word) | |
| if not context_for_llm: | |
| label, context_for_llm = chooseContextLLM(listOfcontexts, alternative_query_word_cleaned) | |
| if not context_for_llm: | |
| context_for_llm = "Collection_date: " + col_date +". Isolate: " + iso + ". Title: " + title + ". Features: " + extracted_features | |
| #print("context for llm: ", label) | |
| # prompt_for_llm = ( | |
| # f"{prompt_instruction_prefix}" | |
| # f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} or the mitochondrial DNA sample in general if these specific identifiers are not explicitly found. " | |
| # f"Identify its primary associated country/geographic location. " | |
| # f"Also, determine if the genetic sample or individual mentioned is from a 'modern' (present-day living individual) " | |
| # f"or 'ancient' (e.g., prehistoric remains, archaeological sample) source. " | |
| # f"If the text does not mention whether the sample is ancient or modern, assume the sample is modern unless otherwise explicitly described as ancient or archaeological. " | |
| # f"Additionally, extract its ethnicity and a more specific location (city/district level) within the predicted country. " | |
| # f"If any information is not explicitly present in the provided text snippets, state 'unknown' for that specific piece of information. " | |
| # f"Provide only the country, sample type, ethnicity, and specific location, do not add extra explanations.\n\n" | |
| # f"Text Snippets:\n{context_for_llm}\n\n" | |
| # f"Output Format: {output_format_str}" | |
| # ) | |
| if len(context_for_llm) > 1000*1000: | |
| context_for_llm = context_for_llm[:900000] | |
| # fix the prompt better: | |
| # firstly clarify more by saying which type of organism, prioritize homo sapiens | |
| features = metadata["all_features"] | |
| organism = "general" | |
| if features != "unknown": | |
| if "organism" in features: | |
| try: | |
| organism = features.split("organism: ")[1].split("\n")[0] | |
| except: | |
| organism = features.replace("\n","; ") | |
| explain_list = "country or sample type (modern/ancient)" #or ethnicity or specific location (province/city)" | |
| # prompt_for_llm = ( | |
| # f"{prompt_instruction_prefix}" | |
| # f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} or the mitochondrial DNA sample in general if these specific identifiers are not explicitly found. " | |
| # f"Identify its primary associated country/geographic location. " | |
| # f"Also, determine if the genetic sample or individual mentioned is from a 'modern' (present-day living individual) " | |
| # f"or 'ancient' (e.g., prehistoric remains, archaeological sample) source. " | |
| # f"If the text does not mention whether the sample is ancient or modern, assume the sample is modern unless otherwise explicitly described as ancient or archaeological. " | |
| # f"Provide only {output_format_str}. " | |
| # f"If any information is not explicitly present in the provided text snippets, state 'unknown' for that specific piece of information. " | |
| # f"If the country or sample type (modern/ancient) is not 'unknown', write 1 sentence after the output explaining how you inferred it from the text (one sentence for each)." | |
| # f"\n\nText Snippets:\n{context_for_llm}\n\n" | |
| # f"Output Format: {output_format_str}" | |
| # ) | |
| # prompt_for_llm = ( | |
| # f"{prompt_instruction_prefix}" | |
| # f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} or the mitochondrial DNA sample in {organism} if these specific identifiers are not explicitly found. " | |
| # f"Identify its primary associated country/geographic location. " | |
| # f"Also, determine if the genetic sample or individual mentioned is from a 'modern' (present-day living individual) " | |
| # f"or 'ancient' (e.g., prehistoric remains, archaeological sample) source. " | |
| # f"If the text does not mention whether the sample is ancient or modern, assume the sample is modern unless otherwise explicitly described as ancient or archaeological. " | |
| # f"Provide only {output_format_str}. " | |
| # f"If any information is not explicitly present in the provided text snippets, state 'unknown' for that specific piece of information. " | |
| # f"If the {explain_list} is not 'unknown', write 1 sentence after the output explaining how you inferred it from the text (one sentence for each)." | |
| # f"\n\nText Snippets:\n{context_for_llm}\n\n" | |
| # f"Output Format: {output_format_str}" | |
| # ) | |
| # prompt_for_llm = ( | |
| # f"{prompt_instruction_prefix}" | |
| # f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} " | |
| # f"or the mitochondrial DNA sample in {organism} if these identifiers are not explicitly found. " | |
| # f"Identify its **primary associated geographic location**, preferring the most specific available: " | |
| # f"first try to determine the exact country; if no country is explicitly mentioned, then provide " | |
| # f"the next most specific region, continent, island, or other clear geographic area mentioned. " | |
| # f"If no geographic clues at all are present, state 'unknown' for location. " | |
| # f"Also, determine if the genetic sample is from a 'modern' (present-day living individual) " | |
| # f"or 'ancient' (prehistoric/archaeological) source. " | |
| # f"If the text does not specify ancient or archaeological context, assume 'modern'. " | |
| # f"Provide only {output_format_str}. " | |
| # f"If any information is not explicitly present, use the fallback rules above before defaulting to 'unknown'. " | |
| # f"For each non-'unknown' field in {explain_list}, write one sentence explaining how it was inferred from the text (one sentence for each)." | |
| # f"\n\nText Snippets:\n{context_for_llm}\n\n" | |
| # f"Output Format: {output_format_str}" | |
| # ) | |
| prompt_for_llm = ( | |
| f"{prompt_instruction_prefix}" | |
| f"Given the following text snippets, analyze the entity/concept {rag_query_phrase} " | |
| f"or the mitochondrial DNA sample in {organism} if these identifiers are not explicitly found. " | |
| f"Identify its **primary associated geographic location**, preferring the most specific available: " | |
| f"first try to determine the exact country; if no country is explicitly mentioned, then provide " | |
| f"the next most specific region, continent, island, or other clear geographic area mentioned. " | |
| f"If no geographic clues at all are present, state 'unknown' for location. " | |
| f"Also, determine if the genetic sample is from a 'modern' (present-day living individual) " | |
| f"or 'ancient' (prehistoric/archaeological) source. " | |
| f"If the text does not specify ancient or archaeological context, assume 'modern'. " | |
| f"Provide only {output_format_str}. " | |
| f"If any information is not explicitly present, use the fallback rules above before defaulting to 'unknown'. " | |
| f"For each non-'unknown' field in {explain_list}, write one sentence explaining how it was inferred from the text " | |
| f"(one sentence for each). " | |
| f"Format your answer so that:\n" | |
| f"1. The **first line** contains only the {output_format_str} answer.\n" | |
| f"2. The **second line onward** contains the explanations.\n" | |
| f"\nText Snippets:\n{context_for_llm}\n\n" | |
| f"Output Format Example:\nChina, modern, Daur, Heilongjiang province.\n" | |
| f"The text explicitly states \"chinese Daur ethnic group in Heilongjiang province\", indicating the country, " | |
| f"the ethnicity, and the specific province. The study is published in a journal, implying research on living " | |
| f"individuals, hence modern." | |
| ) | |
| if model_ai: | |
| print("back up to ", model_ai) | |
| llm_response_text, model_instance = call_llm_api(prompt_for_llm, model=model_ai) | |
| else: | |
| print("still 2.5 flash gemini") | |
| llm_response_text, model_instance = call_llm_api(prompt_for_llm) | |
| print("\n--- DEBUG INFO FOR RAG ---") | |
| print("Retrieved Context Sent to LLM (first 500 chars):") | |
| print(context_for_llm[:500] + "..." if len(context_for_llm) > 500 else context_for_llm) | |
| print("\nRaw LLM Response:") | |
| print(llm_response_text) | |
| print("--- END DEBUG INFO ---") | |
| llm_cost = 0 | |
| if model_instance: | |
| try: | |
| input_llm_tokens = global_llm_model_for_counting_tokens.count_tokens(prompt_for_llm).total_tokens | |
| output_llm_tokens = global_llm_model_for_counting_tokens.count_tokens(llm_response_text).total_tokens | |
| print(f" DEBUG: LLM Input tokens: {input_llm_tokens}") | |
| print(f" DEBUG: LLM Output tokens: {output_llm_tokens}") | |
| llm_cost = (input_llm_tokens / 1000) * PRICE_PER_1K_INPUT_LLM + \ | |
| (output_llm_tokens / 1000) * PRICE_PER_1K_OUTPUT_LLM | |
| print(f" DEBUG: Estimated LLM cost: ${llm_cost:.6f}") | |
| except Exception as e: | |
| print(f" DEBUG: Error counting LLM tokens: {e}") | |
| llm_cost = 0 | |
| total_query_cost += current_embedding_cost + llm_cost | |
| print(f" DEBUG: Total estimated cost for this RAG query: ${total_query_cost:.6f}") | |
| # Parse the LLM's response based on the Output Format actually used | |
| # if output_format_str == "ethnicity, specific_location/unknown": # Targeted RAG output | |
| # extracted_ethnicity,extracted_specific_location = clean_llm_output(llm_response_text, output_format_str) | |
| # elif output_format_str == "modern/ancient/unknown, ethnicity, specific_location/unknown": | |
| # extracted_type, extracted_ethnicity,extracted_specific_location=clean_llm_output(llm_response_text, output_format_str) | |
| # else: # Full RAG output (country, type, ethnicity, specific_location) | |
| # extracted_country,extracted_type, extracted_ethnicity,extracted_specific_location=clean_llm_output(llm_response_text, output_format_str) | |
| metadata_list = parse_multi_sample_llm_output(llm_response_text, output_format_str) | |
| # merge_metadata = merge_metadata_outputs(metadata_list) | |
| # if output_format_str == "country_name, modern/ancient/unknown": | |
| # extracted_country, extracted_type = merge_metadata["country"], merge_metadata["sample_type"] | |
| # country_explanation,sample_type_explanation = merge_metadata["country_explanation"], merge_metadata["sample_type_explanation"] | |
| # elif output_format_str == "modern/ancient/unknown": | |
| # extracted_type = merge_metadata["sample_type"] | |
| # sample_type_explanation = merge_metadata["sample_type_explanation"] | |
| # for the output_format that is not default | |
| if output_format_str == "country_name, modern/ancient/unknown": | |
| outputs = output_format_str.split(", ") | |
| extracted_country, extracted_type = metadata_list[outputs[0]]["answer"], metadata_list[outputs[1]]["answer"] | |
| country_explanation,sample_type_explanation = metadata_list[outputs[0]][outputs[0]+"_explanation"], metadata_list[outputs[1]][outputs[1]+"_explanation"] | |
| # extracted_ethnicity, extracted_specific_location = metadata_list[outputs[2]]["answer"], metadata_list[outputs[3]]["answer"] | |
| # ethnicity_explanation, specific_loc_explanation = metadata_list[outputs[2]][outputs[2]+"_explanation"], metadata_list[outputs[3]][outputs[3]+"_explanation"] | |
| # 6. Optional: Second LLM call for specific_location from general knowledge if still unknown | |
| # if extracted_specific_location == 'unknown': | |
| # # Check if we have enough info to ask general knowledge LLM | |
| # if extracted_country != 'unknown' and extracted_ethnicity != 'unknown': | |
| # print(f" DEBUG: Specific location still unknown. Querying general knowledge LLM from '{extracted_ethnicity}' and '{extracted_country}'.") | |
| # general_knowledge_prompt = ( | |
| # f"Based on general knowledge, what is a highly specific location (city or district) " | |
| # f"associated with the ethnicity '{extracted_ethnicity}' in '{extracted_country}'? " | |
| # f"Consider the context of scientific studies on human genetics, if known. " | |
| # f"If no common specific location is known, state 'unknown'. " | |
| # f"Provide only the city or district name, or 'unknown'." | |
| # ) | |
| # general_llm_response, general_llm_model_instance = call_llm_api(general_knowledge_prompt, model_name='gemini-1.5-flash-latest') | |
| # if general_llm_response and general_llm_response.lower().strip() != 'unknown': | |
| # extracted_specific_location = general_llm_response.strip() + " (predicted from general knowledge)" | |
| # # Add cost of this second LLM call | |
| # if general_llm_model_instance: | |
| # try: | |
| # gk_input_tokens = general_llm_model_instance.count_tokens(general_knowledge_prompt).total_tokens | |
| # gk_output_tokens = general_llm_model_instance.count_tokens(general_llm_response).total_tokens | |
| # gk_cost = (gk_input_tokens / 1000) * PRICE_PER_1K_INPUT_LLM + \ | |
| # (gk_output_tokens / 1000) * PRICE_PER_1K_OUTPUT_LLM | |
| # print(f" DEBUG: General Knowledge LLM cost to predict specific location alone: ${gk_cost:.6f}") | |
| # total_query_cost += gk_cost # Accumulate cost | |
| # except Exception as e: | |
| # print(f" DEBUG: Error counting GK LLM tokens: {e}") | |
| # else: | |
| # print(" DEBUG: General knowledge LLM returned unknown or empty for specific location.") | |
| # # 6. Optional: Second LLM call for ethnicity from general knowledge if still unknown | |
| # if extracted_ethnicity == 'unknown': | |
| # # Check if we have enough info to ask general knowledge LLM | |
| # if extracted_country != 'unknown' and extracted_specific_location != 'unknown': | |
| # print(f" DEBUG: Ethnicity still unknown. Querying general knowledge LLM from '{extracted_specific_location}' and '{extracted_country}'.") | |
| # general_knowledge_prompt = ( | |
| # f"Based on general knowledge, what is a highly ethnicity (population) " | |
| # f"associated with the specific location '{extracted_specific_location}' in '{extracted_country}'? " | |
| # f"Consider the context of scientific studies on human genetics, if known. " | |
| # f"If no common ethnicity is known, state 'unknown'. " | |
| # f"Provide only the ethnicity or popluation name, or 'unknown'." | |
| # ) | |
| # general_llm_response, general_llm_model_instance = call_llm_api(general_knowledge_prompt, model_name='gemini-1.5-flash-latest') | |
| # if general_llm_response and general_llm_response.lower().strip() != 'unknown': | |
| # extracted_ethnicity = general_llm_response.strip() + " (predicted from general knowledge)" | |
| # # Add cost of this second LLM call | |
| # if general_llm_model_instance: | |
| # try: | |
| # gk_input_tokens = general_llm_model_instance.count_tokens(general_knowledge_prompt).total_tokens | |
| # gk_output_tokens = general_llm_model_instance.count_tokens(general_llm_response).total_tokens | |
| # gk_cost = (gk_input_tokens / 1000) * PRICE_PER_1K_INPUT_LLM + \ | |
| # (gk_output_tokens / 1000) * PRICE_PER_1K_OUTPUT_LLM | |
| # print(f" DEBUG: General Knowledge LLM cost to predict ethnicity alone: ${gk_cost:.6f}") | |
| # total_query_cost += gk_cost # Accumulate cost | |
| # except Exception as e: | |
| # print(f" DEBUG: Error counting GK LLM tokens: {e}") | |
| # else: | |
| # print(" DEBUG: General knowledge LLM returned unknown or empty for ethnicity.") | |
| #return extracted_country, extracted_type, method_used, extracted_ethnicity, extracted_specific_location, total_query_cost | |
| return extracted_country, extracted_type, method_used, country_explanation, sample_type_explanation, total_query_cost | |