Aniruddha7/QueryLens-Text2SQL_DocVQA-V2
0
1"""2Agentic Workflow Main Module3---------------------------4This file implements the core agentic workflow for text-to-SQL, including:5- LLM orchestration (Ollama, MCP, fallback)6- Memory management and safety checks7- SQL cleaning, validation, and artifact stripping8- Agentic orchestration (plan/micro/LLM phases)9- Robust fallback and error handling10 11All agent logic is preserved. Comments clarify key logic and utility functions.12"""13 14import os15import asyncio16import re17import gc # Explicit garbage collection import18import signal # For timeouts19import threading # For thread management20import weakref # For weak references to avoid memory leaks21import time # For timing operations22from typing import Dict, Any, List, Optional23import json24import uuid25from datetime import datetime26from pathlib import Path27import psutil # Import for memory monitoring28import requests29 30# Database-related imports31from sqlalchemy import create_engine, text, Table, MetaData, inspect32from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession33from sqlalchemy.orm import sessionmaker34from sqlalchemy.exc import NoSuchTableError35 36# LLM imports37from llama_index.core import Settings38from llama_index.llms.ollama import Ollama39from llama_index.core.agent import ReActAgent40from llama_index.core.tools import FunctionTool41from llama_index.core.llms import ChatMessage, MessageRole, LLMMetadata, CompletionResponse42try:43 from llama_index.core.llms.custom import CustomLLM44except ImportError:45 from llama_index.core.llms import CustomLLM46from llama_index.core.storage.chat_store import SimpleChatStore47try:48 from huggingface_hub import InferenceClient49except ImportError:50 InferenceClient = None51 52try:53 from .invariants import enforce_invariants, classify_question54except (ImportError, ValueError):55 try:56 from Agent.invariants import enforce_invariants, classify_question57 except ImportError:58 def enforce_invariants(*args, **kwargs): return args[0] if args else None59 def classify_question(*args, **kwargs): return "unknown"60 61from dotenv import load_dotenv62import os as _os63_dotenv_candidates = [64 _os.path.join(_os.path.dirname(__file__), '..', '.env'), # project root .env65 _os.path.join(_os.path.dirname(__file__), '.env'), # Agent/.env (most specific)66]67for _dotenv_path in _dotenv_candidates:68 if _os.path.isfile(_dotenv_path):69 print(f"[DEBUG] Loading .env from: {_dotenv_path}")70 load_dotenv(dotenv_path=_dotenv_path, override=False) # Environment variables (Docker) now win over .env files71del _os, _dotenv_candidates, _dotenv_path72print(f"[DEBUG] Before CWD load_dotenv, DB_CONNECTION_URL={os.environ.get('DB_CONNECTION_URL')}")73load_dotenv() # also load CWD .env for any remaining vars74print(f"[DEBUG] After CWD load_dotenv, DB_CONNECTION_URL={os.environ.get('DB_CONNECTION_URL')}")75 76# Simple initialization77DB_CONNECTION_URL = os.environ.get("DB_CONNECTION_URL", None)78print(f"[DEBUG] Final DB_CONNECTION_URL={DB_CONNECTION_URL}")79if not DB_CONNECTION_URL:80 raise ValueError("DB_CONNECTION_URL environment variable is not set")81 82# Import our memory management utilities (robust import with fallbacks)83import importlib84import sys85 86memory_manager = None87safe_check_memory = None88TimeoutManager = None89_last_import_error = None90 91# Try several import strategies: package-relative, package-absolute, and top-level92try_names = []93if __package__:94 # e.g. when imported as Agent.agentic_workflow, __package__ == 'Agent'95 try_names.append(f"{__package__}.memory_manager")96# Also try the absolute package path97try_names.append("Agent.memory_manager")98# Finally try top-level fallback99try_names.append("memory_manager")100 101for modname in try_names:102 try:103 mod = importlib.import_module(modname)104 # Extract expected symbols if present105 memory_manager = getattr(mod, 'memory_manager', None)106 safe_check_memory = getattr(mod, 'safe_check_memory', None)107 TimeoutManager = getattr(mod, 'TimeoutManager', None)108 if memory_manager is not None and safe_check_memory is not None and TimeoutManager is not None:109 break110 except Exception as e:111 _last_import_error = e112 113if memory_manager is None or safe_check_memory is None or TimeoutManager is None:114 # Provide a clear diagnostic so startup logs show why import failed115 print("Failed to import Agent.memory_manager via tried paths:", try_names)116 print("cwd=", __import__('os').getcwd())117 print("sys.path[0]=", sys.path[0] if len(sys.path) > 0 else None)118 if _last_import_error:119 # Re-raise the last error to preserve traceback120 raise _last_import_error 121 else:122 raise ImportError("Could not locate memory_manager module; ensure Agent/memory_manager.py exists and package imports are correct")123 124# Ollama configuration125OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")126 127 128# --- Employee Database Schema for Prompt Injection ---129EMPLOYEE_DB_SCHEMA = (130 "Database schema:\n"131 "- employee(emp_no, birth_date, first_name, last_name, gender, hire_date)\n"132 "- department(dept_no, dept_name)\n"133 "- dept_emp(emp_no, dept_no, from_date, to_date)\n"134 "- dept_manager(emp_no, dept_no, from_date, to_date)\n"135 "- salary(emp_no, salary, from_date, to_date) # column name is 'salary' in this dataset\n"136 "- title(emp_no, title, from_date, to_date)\n"137 "- views: current_dept_emp (current department per employee), dept_emp_latest_date\n"138 "\n"139 "CRITICAL Rules:\n"140 "- The 'department' table does NOT have a 'to_date' column. Never filter department by to_date.\n"141 "- To get CURRENT records (managers, salaries, department assignments), ALWAYS filter by `to_date = '9999-01-01'` in the 'salary', 'dept_emp', 'dept_manager', or 'title' tables.\n"142 "- To get MANAGER NAMES, you MUST join `dept_manager` (dm) with `employee` (e) on `emp_no`.\n"143 "- To get DEPARTMENT NAMES for employees, join `dept_emp` (de) with `department` (d) on `dept_no`.\n"144 "- The 'employee' table does NOT have a 'dept_no' column. To find an employee's department name, you MUST join 'dept_emp' (de) with 'department' (d) on 'dept_no'.\n"145 "- Salaries are in table `salary` (column `salary`). Always use `s.to_date = '9999-01-01'` for current pay.\n"146)147 148# Allow a dedicated finetuned model name (merged LoRA) via FINETUNED_MODEL_NAME env.149# Resolution order:150# 1. FINETUNED_MODEL_NAME (e.g., gemma1b-text2sql)151# 2. OLLAMA_MODEL_NAME (defaults to base gemma3:1b)152# 3. Hard fallback 'gemma3:1b'153FINETUNED_MODEL_NAME = os.environ.get("FINETUNED_MODEL_NAME")154BASE_MODEL_FALLBACK = "qwen2.5-coder:1.5b" #"gemma3:1b-it-qat"155OLLAMA_MODEL_NAME = FINETUNED_MODEL_NAME or os.environ.get("OLLAMA_MODEL_NAME", BASE_MODEL_FALLBACK)156FALLBACK_MODEL_NAME = BASE_MODEL_FALLBACK # Keep fallback stable157if FINETUNED_MODEL_NAME:158 print(f"🦙 Using finetuned Ollama model: {OLLAMA_MODEL_NAME} (base fallback: {FALLBACK_MODEL_NAME})")159else:160 print(f"🦙 Using base Ollama model: {OLLAMA_MODEL_NAME} (fallback: {FALLBACK_MODEL_NAME})")161 162# OCR Q&A model configuration163# PRIMARY: Granite Vision VLM (GRANITE_MODEL_PATH env) — used when image file is available164# FALLBACK: lightweight Ollama text model — used for old docs without a saved image165GRANITE_MODEL_PATH = (166 os.environ.get("GRANITE_MODEL_PATH") or167 os.environ.get("GRANITE_MODEL") or168 "qwen2-vl:2b" # local Ollama fallback (dev only)169)170if os.environ.get("USE_HF_CLOUD", "0") in ("1", "true", "True"):171 print(f"[GRANITE_VQA] Target Model: {os.environ.get('GROQ_VISION_MODEL', 'meta-llama/llama-4-scout-17b-16e-instruct')} (via Groq Vision)")172else:173 print(f"[GRANITE_VQA] Target Model: {GRANITE_MODEL_PATH} (via Ollama)")174# Text-only fallback: Ollama model used when no image file exists for the doc.175# NOT the same as GRANITE_MODEL_PATH — this runs via Ollama REST API, not HuggingFace.176OLLAMA_OCR_QA_MODEL = os.environ.get("OLLAMA_OCR_QA_MODEL", "gemma3:1b-it-qat")177print(f"[OCR_QA] Text fallback model (Ollama): {OLLAMA_OCR_QA_MODEL}")178OCR_QA_TIMEOUT = float(os.environ.get("OCR_QA_TIMEOUT", "45"))179print(f"[OCR_QA] Text fallback timeout={OCR_QA_TIMEOUT}s")180 181def check_memory_safety():182 """Check if there's enough memory to safely process a request"""183 # Use our improved memory checking system184 return safe_check_memory()185 186# Global variables for database connection187available_tables = []188inspector = None189sql_database = None190async_engine = None191async_session = None192 193# Initialize chat store directory and files194chat_store_dir = Path('chat_store')195chat_store_dir.mkdir(parents=True, exist_ok=True)196 197# Define chat store file paths198public_store_path = chat_store_dir/'chat_store_public.json'199private_store_path = chat_store_dir/'chat_store_private.json'200# Initialize empty chat store structure with conversation key201initial_store = {"conversation": []}202 203# Create JSON files if they don't exist204if not public_store_path.exists(): 205 with open(public_store_path, 'w') as f:206 json.dump(initial_store, f, indent=2)207 208if not private_store_path.exists(): 209 with open(private_store_path, 'w') as f:210 json.dump(initial_store, f, indent=2)211 212# Initialize chat stores213print("Debug: Initializing chat stores")214chat_store = SimpleChatStore()215chat_store_public = chat_store.from_persist_path(str(public_store_path))216chat_store_private = chat_store.from_persist_path(str(private_store_path))217print("Debug: Chat stores initialized")218 219# Utility: robustly extract clean SQL from LLM/artifact wrappers220def extract_clean_sql(raw: str) -> str:221 """Remove <<<SQL_START>>>, <<<SQL_END>>>, code fences, and similar wrappers from SQL text."""222 import re223 if not raw or not isinstance(raw, str):224 return raw225 s = raw.strip()226 # Remove <<<SQL_START>>> ... <<<SQL_END>>> blocks227 s = re.sub(r"<<<SQL_START>>>[\s\S]*?<<<SQL_END>>>", lambda m: m.group(0).replace('<<<SQL_START>>>','').replace('<<<SQL_END>>>','').strip(), s, flags=re.IGNORECASE)228 # Remove <<<SQL_REVISED_START>>> ... <<<SQL_REVISED_END>>> blocks229 s = re.sub(r"<<<SQL_REVISED_START>>>[\s\S]*?<<<SQL_REVISED_END>>>", lambda m: m.group(0).replace('<<<SQL_REVISED_START>>>','').replace('<<<SQL_REVISED_END>>>','').strip(), s, flags=re.IGNORECASE)230 # Remove code fences (```sql ... ``` or ```)231 s = re.sub(r"```(?:sql)?([\s\S]*?)```", lambda m: m.group(1).strip(), s, flags=re.IGNORECASE)232 # Remove any remaining <<<...>>> wrappers233 s = re.sub(r"<<<[A-Z_]+>>>", "", s)234 235 # Extract only the SQL portion before any explanations.236 # Handles both plain "Explanation:" and markdown "### Explanation:" headers.237 _explanation_pattern = r';\s*\n[\s\S]*?(?:#{1,3}\s*[Ee]xplanation|[Ee]xplanation)\s*:'238 if re.search(_explanation_pattern, s):239 s = re.split(_explanation_pattern, s)[0].rstrip() + ';'240 elif re.search(r';\s*\n\s*[Ee]xplanation:', s):241 s = re.split(r';\s*\n\s*[Ee]xplanation:', s)[0] + ';'242 243 244 # Remove leading/trailing whitespace and repeated newlines245 s = s.strip()246 # Remove any leading/trailing code fence lines247 s = re.sub(r"^```[a-zA-Z]*\s*|\s*```$", "", s, flags=re.MULTILINE)248 249 # Remove any trailing explanation text that might still be there250 lines = s.split('\n')251 sql_lines = []252 for line in lines:253 line = line.strip()254 # Stop at explanation markers — including markdown ### headers255 if re.match(r'^#{1,3}\s*(explanation|note|this query)', line, re.IGNORECASE):256 break257 if re.match(r'^(explanation|note|this query):', line, re.IGNORECASE):258 break259 if line.lower().startswith('explanation:'):260 break261 if re.match(r'^\d+\.\s+', line): # Numbered explanation steps262 break263 sql_lines.append(line)264 265 s = '\n'.join(sql_lines).strip()266 267 # NEW: Safety check for multiple commands. 268 # If the LLM outputted multiple queries (separated by semicolons), 269 # we only want the LAST valid one (which is usually the most refined).270 if s.count(';') > 1:271 # Split by semicolon, keep only those that look like queries272 statements = [stmt.strip() for stmt in s.split(';') if 'select' in stmt.lower()]273 if statements:274 # Use the last one as it's typically the final refinement275 s = statements[-1] + ';'276 print(f"[CLEAN_SQL] Multiple queries detected. Selected last valid statement.")277 278 return s.strip()279 280# -----------------------------281# OCR Q&A dedicated helpers282# -----------------------------283def _get_ocr_text_by_doc_id(doc_id: str) -> str:284 try:285 p = Path('chat_store') / 'docs' / f"{doc_id}.txt"286 if not p.exists():287 return ""288 return p.read_text(encoding='utf-8')289 except Exception:290 return ""291 292def _get_image_path_by_doc_id(doc_id: str) -> str:293 """Return the saved image file path for doc_id, or empty string if not found."""294 try:295 p = Path('chat_store') / 'docs' / f"{doc_id}.img_path"296 if not p.exists():297 return ""298 return p.read_text(encoding='utf-8').strip()299 except Exception:300 return ""301 302def _granite_vision_qa(image_path: str, question: str) -> str:303 """Answer a question about an image using IBM Granite Vision 3.1-2b (direct VLM Q&A).304 305 Loads the model from local HF cache with aggressive quantization to stay within306 the RAM budget (4-bit → 8-bit → fp16 fallback chain). The model is explicitly307 deleted after generation to free memory for the SQL model to reload.308 """309 import gc310 try:311 import torch312 from transformers import AutoProcessor, AutoModelForVision2Seq313 from PIL import Image as PILImage314 except ImportError as e:315 return f"[ERROR] Granite Vision dependencies not installed: {e}"316 317 model_path = GRANITE_MODEL_PATH318 cache_dir = os.environ.get("HF_HOME") or r"D:\hf_cache"319 offload_dir = os.path.join(cache_dir, "offload")320 os.makedirs(offload_dir, exist_ok=True)321 322 device = "cuda" if torch.cuda.is_available() else "cpu"323 print(f"[GRANITE_VQA] Loading {model_path} on {device} from {cache_dir}")324 325 processor = vlm_model = None326 try:327 processor = AutoProcessor.from_pretrained(model_path, local_files_only=True, cache_dir=cache_dir)328 329 # Try 4-bit → 8-bit → plain fp16 to fit within RAM budget330 load_kwargs = dict(local_files_only=True, cache_dir=cache_dir, low_cpu_mem_usage=True)331 loaded = False332 for attempt, extra in enumerate([333 {"load_in_4bit": True},334 {"load_in_8bit": True},335 {"device_map": "auto", "offload_folder": offload_dir},336 ]):337 try:338 vlm_model = AutoModelForVision2Seq.from_pretrained(model_path, **load_kwargs, **extra)339 print(f"[GRANITE_VQA] Loaded (attempt {attempt+1}: {list(extra.keys())[0]})")340 loaded = True341 break342 except Exception as le:343 print(f"[GRANITE_VQA] Load attempt {attempt+1} failed: {le}")344 if not loaded:345 return "[ERROR] Granite Vision could not be loaded — all quantisation attempts failed."346 347 img = PILImage.open(image_path).convert("RGB")348 conversation = [349 {350 "role": "user",351 "content": [352 {"type": "image", "url": image_path},353 {"type": "text", "text": (354 f"{question}\n\n"355 "Answer concisely using only information visible in the image. "356 "If you see a table, read the correct row and column carefully."357 )},358 ],359 }360 ]361 inputs = processor.apply_chat_template(362 conversation,363 add_generation_prompt=True,364 tokenize=True,365 return_dict=True,366 return_tensors="pt",367 )368 # Move inputs to model device369 try:370 model_device = next(vlm_model.parameters()).device371 inputs = {k: v.to(model_device) if hasattr(v, 'to') else v for k, v in inputs.items()}372 except Exception:373 pass374 375 with torch.no_grad():376 output_ids = vlm_model.generate(**inputs, max_new_tokens=256, do_sample=False)377 answer = processor.decode(output_ids[0], skip_special_tokens=True)378 379 # Strip echoed prompt (Granite echoes conversation before the answer)380 if "ASSISTANT" in answer.upper():381 answer = answer.split("ASSISTANT")[-1].strip(" :").strip()382 elif question.lower()[:20] in answer.lower():383 idx = answer.lower().rfind(question.lower()[:20])384 answer = answer[idx + len(question):].strip(" :").strip() if idx != -1 else answer385 386 print(f"[GRANITE_VQA] Answer (chars={len(answer)}): {answer[:200]}")387 return answer.strip() or "[ERROR] Granite Vision returned an empty response."388 389 except Exception as e:390 print(f"[GRANITE_VQA] Generation failed: {e}")391 return f"[ERROR] Granite Vision Q&A failed: {e}"392 finally:393 # Aggressively free VLM memory so SQL model can reload394 try:395 del vlm_model, processor396 except Exception:397 pass398 gc.collect()399 try:400 import torch401 if torch.cuda.is_available():402 torch.cuda.empty_cache()403 except Exception:404 pass405 print("[GRANITE_VQA] Model unloaded; RAM released.")406 407 408def _structure_invoice_text(raw: str) -> str:409 """Convert flat Tesseract OCR output into a structured KEY: VALUE format.410 411 Tesseract merges field labels and values on the same line and squashes412 table rows together, causing a 1B LLM to confuse GST numbers with phone413 numbers, amounts with dates, etc. This function:414 - Recognises common invoice field patterns via regex415 - Emits explicit 'FIELD_NAME: value' lines for the LLM416 - Preserves the rest of the text verbatim417 So the model receives unambiguous context like:418 GSTIN (Seller): 24HDE7487RE5RT4419 Customer Phone: 9372346666420 Customer GSTIN: 07AOLCC1206D126421 """422 import re423 lines = raw.replace('\r\n', '\n').replace('\r', '\n').split('\n')424 structured = []425 seen_keys = set()426 427 def emit(key: str, val: str):428 tag = f"{key}: {val.strip()}"429 if tag not in seen_keys:430 seen_keys.add(tag)431 structured.append(tag)432 433 # Regex patterns for common Indian invoice fields434 patterns = [435 # GSTIN – 15-char alphanumeric starting with 2 digits436 (r'GSTIN\s*(?:No\.?|:)?\s*([0-9]{2}[A-Z0-9]{13})', 'GSTIN'),437 # Phone numbers (10-digit Indian mobile or landline with STD)438 (r'(?:PHONE|TEL|MOB(?:ILE)?)\s*[:\.]?\s*([0-9\-\+\s]{8,15})', 'Phone'),439 # Invoice number440 (r'Invoice\s*No\.?\s*[:\.]?\s*([A-Z0-9/\-]+)', 'Invoice Number'),441 # Challan number442 (r'Challan\s*No\.?\s*[:\.]?\s*([0-9]+)', 'Challan Number'),443 # Invoice / Challan dates444 (r'Invoice\s*Date\s*[:\.]?\s*(\d{1,2}[-/]\w{3,9}[-/]\d{2,4})', 'Invoice Date'),445 (r'Challan\s*Date\s*[:\.]?\s*(\d{1,2}[-/]\w{3,9}[-/]\d{2,4})', 'Challan Date'),446 (r'Due\s*Date\s*[:\.]?\s*(\d{1,2}[-/]\w{3,9}[-/]\d{2,4})', 'Due Date'),447 (r'Delivery\s*Date\s*[:\.]?\s*(\d{1,2}[-/]\w{3,9}[-/]\d{2,4})', 'Delivery Date'),448 # E-Way number449 (r'E-?Way\s*No\.?\s*[:\.]?\s*([A-Z0-9]+)', 'E-Way Number'),450 # LR / PO numbers451 (r'L\.?R\.?\s*No\.?\s*[:\.]?\s*([0-9]+)', 'LR Number'),452 (r'P\.?O\.?\s*No\.?\s*[:\.]?\s*([0-9]+)', 'PO Number'),453 # PAN454 (r'PAN\s*[:\.]?\s*([A-Z]{5}[0-9]{4}[A-Z])', 'PAN'),455 # Bank details456 (r'Bank\s*(?:Account|A/?C)\s*(?:No\.?|Number)?\s*[:\.]?\s*([0-9]+)', 'Bank Account Number'),457 (r'IFSC\s*[:\.]?\s*([A-Z]{4}0[A-Z0-9]{6})', 'Bank IFSC'),458 # Totals459 (r'Total\s+(?:Amount)?\s*After\s*Tax\s*[:\.]?\s*[\u20B9Rs\.]*\s*([\d,\.]+)', 'Total Amount After Tax'),460 (r'Taxable\s+(?:Amount|Value)\s*[:\.]?\s*([\d,\.]+)', 'Taxable Amount'),461 (r'Add\s*[:\.]?\s*IGST\s*[:\.]?\s*([\d,\.]+)', 'IGST Amount'),462 ]463 464 # Scan all lines for known field patterns465 full_text = '\n'.join(lines)466 # Track which GSTINs we've seen so we can label seller vs customer467 gstin_count = 0468 for line in lines:469 for pat, label in patterns:470 m = re.search(pat, line, re.IGNORECASE)471 if m:472 val = m.group(1).strip()473 if label == 'GSTIN':474 gstin_count += 1475 lbl = 'Seller GSTIN' if gstin_count == 1 else 'Customer GSTIN'476 emit(lbl, val)477 elif label == 'Phone':478 emit('Customer Phone', val)479 else:480 emit(label, val)481 482 # Extract company name (typically first non-empty ALL-CAPS-ish line)483 for line in lines:484 stripped = line.strip()485 if len(stripped) > 4 and stripped.isupper() and not re.search(r'\d{5,}', stripped):486 emit('Company Name', stripped)487 break488 489 # Extract customer name after 'M/S' or 'Bill to'490 for i, line in enumerate(lines):491 if re.search(r'^(M/?S|Bill\s*to)\b', line.strip(), re.IGNORECASE):492 # Customer name is likely the next non-empty line or on same line after M/S493 rest = re.sub(r'^(M/?S|Bill\s*to)\s*', '', line.strip(), flags=re.IGNORECASE).strip()494 if rest:495 emit('Customer Name', rest)496 elif i + 1 < len(lines) and lines[i+1].strip():497 emit('Customer Name', lines[i+1].strip())498 break499 500 # --- Product line items ---501 # Scan for rows like: "2 | Stanley Hammer 8295 1.00 PCS 568.00 568.00 9.00 51.12 619.12"502 # OCR often squashes all rows onto one long line; use the item-number+pipe as separator.503 item_pat = re.compile(504 r'\b(\d{1,3})\s*[|Ill1]\s*' # item number + pipe (I/l/1 are common OCR artefacts for |)505 r'([A-Za-z][A-Za-z\s\-]{2,40}?)\s+' # product name506 r'(\d{4,8})\s+' # HSN / SAC code507 r'(\d+(?:\.\d+)?)\s+' # quantity508 r'(?:PCS|NOS|KG|LT|MT|SET|BOX|PC|NO|EA)\.?\s+' # unit (consumed but not captured)509 r'([\d,]+(?:\.\d+)?)' # rate (first numeric after unit)510 r'((?:\s+[\d,]+(?:\.\d+)?){1,4})', # remaining numbers (taxable value, GST, total etc.)511 re.IGNORECASE512 )513 for m in item_pat.finditer(full_text):514 item_no = m.group(1)515 name = m.group(2).strip()516 hsn = m.group(3)517 qty = m.group(4)518 rate = m.group(5)519 rest_nums = re.findall(r'[\d,]+(?:\.\d+)?', m.group(6) or '')520 # rest_nums order depends on invoice: [taxable_value, gst%, gst_amount, total]521 taxable = rest_nums[0] if len(rest_nums) > 0 else rate522 total = rest_nums[-1] if len(rest_nums) > 1 else taxable523 label = f"Product {item_no} - {name}"524 emit(label, f"HSN={hsn}, Qty={qty}, Rate={rate}, Taxable Value={taxable}, Total={total}")525 526 # Build the final structured block527 header = "=== STRUCTURED DOCUMENT FIELDS ==="528 body = "=== RAW OCR TEXT ===\n" + raw.strip()529 if structured:530 return header + '\n' + '\n'.join(structured) + '\n\n' + body531 return body532 533 534def _ollama_unload_model(model: str) -> bool:535 """Ask Ollama to immediately unload a model from RAM (keep_alive=0).536 537 Uses /api/generate with a minimal prompt and keep_alive=0 to signal Ollama538 to release the model. Works for models loaded by any client (LlamaIndex etc).539 Returns True if the request succeeded, False otherwise (non-fatal).540 """541 base = OLLAMA_BASE_URL.rstrip('/')542 try:543 # A single-space prompt ensures Ollama actually processes the request544 # (empty string may be skipped); keep_alive=0 causes immediate unload.545 resp = requests.post(546 f"{base}/api/generate",547 json={"model": model, "prompt": " ", "keep_alive": 0, "stream": False},548 timeout=90.0,549 )550 print(f"[OCR_QA] Unloaded model '{model}' from Ollama RAM (status={resp.status_code})")551 return resp.status_code < 300552 except Exception as e:553 print(f"[OCR_QA][WARN] Failed to unload '{model}': {e}")554 return False555 556 557def _ollama_generate_text(model: str, prompt: str, timeout: float) -> str:558 """Minimal REST call to Ollama generate API to avoid llama_index routing.559 560 Uses OLLAMA_BASE_URL env and returns the concatenated response text.561 num_ctx is chosen adaptively based on available RAM so we don't cause562 excessive paging on memory-constrained machines.563 """564 565 url = f"{OLLAMA_BASE_URL.rstrip('/')}/api/generate"566 567 # Pick num_ctx based on current available memory568 # ≥800 MB free -> 2048 (full invoice fits comfortably)569 # 400–800 MB -> 1024 (still fits most invoices)570 # <400 MB -> 512 (minimal; avoid paging)571 try:572 stats = memory_manager.get_memory_stats()573 avail_mb = stats.get("available_mb", 999.0)574 except Exception:575 avail_mb = 999.0576 577 if avail_mb >= 800:578 num_ctx = 1024579 num_predict = 200580 elif avail_mb >= 400:581 num_ctx = 512582 num_predict = 150583 else:584 num_ctx = 256585 num_predict = 100586 587 payload = {588 "model": model,589 "prompt": prompt,590 "stream": False,591 "keep_alive": 0, # unload immediately after response — prevents gemma3 squatting in RAM592 "options": {593 "temperature": 0.2,594 "num_ctx": num_ctx,595 "num_predict": num_predict,596 "num_gpu": -1, # keep on CPU same as Qwen; prevents GPU/CPU swap latency597 "num_thread": 2, # 2 threads is sufficient for 1B model on CPU598 }599 }600 try:601 # Debug request meta (do not print full prompt to avoid huge logs here)602 try:603 print(f"[DEBUG][OLLAMA_REQUEST] url={url} model={model} prompt_chars={len(prompt)} timeout={timeout} num_ctx={num_ctx} avail_mb={avail_mb:.0f}")604 except Exception:605 pass606 r = requests.post(url, json=payload, timeout=timeout)607 r.raise_for_status()608 data = r.json()609 # Log only response length, not full body (can be very large)610 resp_text = data.get('response') or data.get('text') or str(data)611 print(f"[DEBUG][OLLAMA_RESPONSE] model={model} response_chars={len(resp_text)}")612 # Non-stream returns {response: "..."}613 txt = resp_text614 return txt615 except Exception as e:616 print(f"[ERROR][OLLAMA_API] Ollama generate failed: {e}")617 return f"Error: Ollama generate failed: {e}"618 619def ocr_agent_qa(question: str, doc_id: str, *, model: str | None = None, timeout: float | None = None) -> str:620 """Answer a question about a previously uploaded document.621 622 PRIMARY: If the original image is available (.img_path), routes to Granite Vision623 via MCP (granite_vision.qa tool) for direct visual Q&A.624 FALLBACK: If only OCR text is available (.txt), builds a grounded prompt and625 calls OCR_TEXT_FALLBACK_MODEL (gemma3) via Ollama.626 """627 text = _get_ocr_text_by_doc_id(doc_id)628 629 # PRIMARY PATH: Granite Vision direct visual Q&A630 image_path = _get_image_path_by_doc_id(doc_id)631 632 # ── DIAGNOSTIC (temporary) ──────────────────────────────────────────633 _img_path_file = Path('chat_store') / 'docs' / f"{doc_id}.img_path"634 print(f"[OCR_QA_DIAG] doc_id = {doc_id}")635 print(f"[OCR_QA_DIAG] .img_path file = {_img_path_file.resolve()} | exists={_img_path_file.exists()}")636 print(f"[OCR_QA_DIAG] image_path = {image_path!r}")637 print(f"[OCR_QA_DIAG] isfile(image) = {os.path.isfile(image_path) if image_path else 'N/A (empty)'}")638 # ────────────────────────────────────────────────────────────────────639 640 if image_path and os.path.isfile(image_path):641 print(f"[OCR_QA] Using Granite Vision for doc_id={doc_id}, image={image_path}")642 # Only perform RAM swapping/unloading if we are NOT in cloud mode643 if not (os.environ.get("USE_HF_CLOUD", "0") in ("1", "true", "True")):644 print(f"[OCR_QA] Swapping: unloading SQL model '{OLLAMA_MODEL_NAME}' to free RAM for Granite Vision")645 _ollama_unload_model(OLLAMA_MODEL_NAME)646 time.sleep(1.0) # let OS reclaim pages647 else:648 print(f"[OCR_QA] 🌩️ CLOUD MODE: Skipping local RAM swap.")649 # Route through MCP so Granite Vision runs as a proper registered agent tool650 try:651 from .mcp_client import mcp_call_tool652 except ImportError:653 from Agent.mcp_client import mcp_call_tool654 try:655 ctx = {656 "tool_call": {657 "name": "granite_vision.qa",658 "args": {659 "image_path": image_path,660 "question": question,661 "model_path": GRANITE_MODEL_PATH,662 },663 }664 }665 resp = mcp_call_tool(prompt="run granite vision qa", timeout=300.0, context=ctx)666 tool_result = resp.get("tool_result") if isinstance(resp, dict) else None667 if isinstance(tool_result, dict):668 return tool_result.get("answer", "[ERROR] No answer returned from granite_vision.qa")669 return str(tool_result) if tool_result else "[ERROR] granite_vision.qa returned no result"670 except Exception as e:671 print(f"[OCR_QA] MCP granite_vision.qa failed: {e}")672 return f"[ERROR] Vision Q&A timed out or failed. Please ensure Docker has enough RAM (6GB+) and try again. ({e})"673 674 # No image available — this doc was uploaded before the Granite Vision upgrade.675 # Returning a clear message instead of a wrong answer from the old gemma3 pipeline.676 print(f"[OCR_QA] No image found for doc_id={doc_id} — document predates Granite Vision upgrade")677 return (678 "⚠️ This document was uploaded before the Granite Vision upgrade. "679 "Please click 'Upload Image' to re-upload the image — "680 "Granite Vision will then answer your question directly from the image."681 )682 683# ---------------------------------------------------------------------------684# Conversation Memory Index (improves recall vs naive substring search)685# ---------------------------------------------------------------------------686question_memory_index = {} # normalized_question -> last answer687 688def normalize_question(q: str) -> str:689 """Create a stable normalization of a user question for indexing.690 Steps: lowercase, remove punctuation except underscores, collapse whitespace.691 """692 import re693 q = q.lower()694 # remove any document snippet brackets if present (they make later repeats unmatchable)695 q = re.sub(r"\[document .*?excerpt end\]", "", q, flags=re.IGNORECASE|re.DOTALL)696 # strip bracketed metadata sections697 q = re.sub(r"\[[^\]]+\]", "", q)698 # punctuation -> space (keep underscores)699 q = re.sub(r"[^a-z0-9_]+", " ", q)700 q = re.sub(r"\s+", " ", q).strip()701 return q702 703def rebuild_memory_index():704 """Rebuild in-memory index from persisted chat history."""705 question_memory_index.clear()706 msgs = chat_store_private.get_messages(key="conversation")707 last_user = None708 for m in msgs:709 if m.role == MessageRole.USER:710 last_user = m.content711 elif m.role == MessageRole.ASSISTANT and last_user:712 key = normalize_question(last_user)713 question_memory_index[key] = m.content714 last_user = None715 print(f"[MEM-INDEX] Rebuilt memory index with {len(question_memory_index)} entries")716 717# Build initial index718try:719 rebuild_memory_index()720except Exception as e:721 print(f"[MEM-INDEX][WARN] Failed to build initial index: {e}")722 723# ---------------------------------------------------------------------------724# Feature flags / global controls725# ---------------------------------------------------------------------------726# Disable ALL template-based fallback logic to enforce pure LLM agentic path727DISABLE_TEMPLATE_FALLBACK = True728# Feature toggles729ENABLE_INVARIANTS = True # can be turned off for debugging730USE_SIMPLE_LLM_CALL = True # bypass TimeoutManager for main generation731# Configurable generation timeouts (seconds); can override via env vars732FIRST_ATTEMPT_TIMEOUT = float(os.environ.get("FIRST_ATTEMPT_TIMEOUT", 180)) # Increased to 180s for initial attempt733SECOND_ATTEMPT_TIMEOUT = float(os.environ.get("SECOND_ATTEMPT_TIMEOUT", 150)) # Increased to 150s for retry734THIRD_ATTEMPT_TIMEOUT = float(os.environ.get("THIRD_ATTEMPT_TIMEOUT", 120)) # Increased to 120s for final attempt735DEFAULT_LLM_REQUEST_TIMEOUT = float(os.environ.get("LLM_REQUEST_TIMEOUT", 150))736 737# Micro-verification timeouts738MICRO_VERIFY_MCP_TIMEOUT = float(os.environ.get("MICRO_VERIFY_MCP_TIMEOUT", 30))739MICRO_VERIFY_LLM_TIMEOUT = float(os.environ.get("MICRO_VERIFY_LLM_TIMEOUT", 60))740 741# Debug: Print timeout values to verify they're correct742print(f"[TIMEOUT_DEBUG] FIRST_ATTEMPT_TIMEOUT={FIRST_ATTEMPT_TIMEOUT}s")743print(f"[TIMEOUT_DEBUG] SECOND_ATTEMPT_TIMEOUT={SECOND_ATTEMPT_TIMEOUT}s")744print(f"[TIMEOUT_DEBUG] THIRD_ATTEMPT_TIMEOUT={THIRD_ATTEMPT_TIMEOUT}s")745STABLE_REUSE_LLM = True # don't re-init on pure timeout, only on MemoryError746# Absolute safety thresholds to avoid system instability / BSOD risk747ABS_MIN_FREE_MB = float(os.environ.get("LLM_ABSOLUTE_MIN_FREE_MB", "200")) # Lowered for GPU mode748SAFE_POST_QUERY_TARGET_MB = float(os.environ.get("LLM_SAFE_POST_QUERY_TARGET_MB", "300")) # Lowered for GPU mode749# Placeholder tokens that should never appear in final executable SQL750PLACEHOLDER_TOKENS = {751 "your title", "some title", "sample title", "your table", "table_name",752 "column_name", "some_column", "placeholder", "value_here"753}754PLACEHOLDER_REGEXES = [755 r"your_[a-z0-9_]*", # generic your_ placeholder tokens756]757 758# (Duplicate EMPLOYEE_DB_SCHEMA removed to prevent overwriting)759 760# Global readiness flag so we can re-attempt handshake before first real generation761LLM_HANDSHAKE_COMPLETE = False762LAST_GENERATED_SQL = None # updated each time a final SQL statement is produced763# Cached real department names from DB — populated during initialize_database764# Fallback: known employee-DB departments (used if DB query fails)765_cached_dept_names: set = {766 'Customer Service', 'Development', 'Finance', 'Human Resources',767 'Marketing', 'Production', 'Quality Management', 'Research', 'Sales'768}769 770# Function to initialize LLM with better memory management771def initialize_llm(progressive_reduction=0):772 """773 Initialize the LLM model with Ollama with proper memory management.774 775 Args:776 progressive_reduction: Level of resource reduction (0=normal, 1=reduced, 2=minimal)777 778 Returns:779 An initialized Ollama LLM instance780 """781 # Make sure we have a clean slate782 gc.collect()783 784 # Get current memory stats785 global LLM_HANDSHAKE_COMPLETE786 stats = memory_manager.get_memory_stats()787 available_mb = stats["available_mb"]788 789 # CLOUD MODE: If USE_HF_CLOUD is enabled, skip local GPU/memory checks790 if os.environ.get("USE_HF_CLOUD", "0") in ("1", "true", "True"):791 hf_token = os.environ.get("HF_API_TOKEN")792 hf_model = os.environ.get("HF_LLM_MODEL", "Qwen/Qwen2.5-Coder-7B-Instruct")793 794 if not hf_token:795 print("[WARN] HF_API_TOKEN not found! Falling back to local Ollama even though USE_HF_CLOUD=1")796 elif InferenceClient is None:797 print("[WARN] huggingface_hub not installed! Falling back to local Ollama.")798 else:799 print(f"[INFO] CLOUD MODE: Initializing Official HF InferenceClient ({hf_model})")800 # Create a simple client801 client = InferenceClient(model=hf_model, token=hf_token)802 803 class HFClientWrapper(CustomLLM):804 # Use a dummy model name that LlamaIndex accepts805 model_name: str = hf_model806 807 @property808 def metadata(self) -> LLMMetadata:809 return LLMMetadata(810 context_window=4096,811 num_output=1024,812 model_name=self.model_name,813 )814 815 def complete(self, prompt: str, **kwargs) -> CompletionResponse:816 # Talk directly to HF serverless pool using the Chat Completion API817 # This satisfies providers that require the "conversational" task818 hf_client = InferenceClient(model=self.model_name, token=hf_token)819 resp = hf_client.chat_completion(820 messages=[{"role": "user", "content": prompt}],821 max_tokens=1024822 )823 return CompletionResponse(text=resp.choices[0].message.content)824 825 def stream_complete(self, prompt: str, **kwargs):826 # Not needed for our agent, but required by base class827 raise NotImplementedError()828 829 def chat(self, messages, **kwargs):830 # Convert ChatMessage list to string prompt for stability831 prompt = ""832 for m in messages:833 role = getattr(m, 'role', 'user')834 content = getattr(m, 'content', str(m))835 prompt += f"{role}: {content}\n"836 return self.complete(prompt)837 838 llm = HFClientWrapper()839 # Register it in Settings so downstream tools can use it840 Settings.llm = llm841 global LLM_HANDSHAKE_COMPLETE842 LLM_HANDSHAKE_COMPLETE = True843 return llm844 845 # Hard safety abort first846 if available_mb < ABS_MIN_FREE_MB:847 raise MemoryError(f"Refusing to initialize LLM: free {available_mb:.1f}MB < ABS_MIN_FREE_MB {ABS_MIN_FREE_MB}MB (system safety)")848 849 # Automatically determine if we need to increase resource reduction based on memory850 if available_mb < memory_manager.critical_mb * 1.5:851 # Force at least level 2 (minimal) if memory is near critical852 progressive_reduction = max(progressive_reduction, 2)853 print(f"Memory critically low ({available_mb:.2f}MB): Forcing minimal resource mode")854 elif available_mb < memory_manager.threshold_mb:855 # Force at least level 1 (reduced) if memory is below threshold856 progressive_reduction = max(progressive_reduction, 1)857 print(f"Memory low ({available_mb:.2f}MB): Using reduced resource mode")858 859 # Determine resource usage level based on available memory and progressive reduction860 # Allow environment overrides for mode thresholds861 try:862 env_ultra = float(os.environ.get("LLM_ULTRA_MB", "800"))863 env_min = float(os.environ.get("LLM_MIN_MB", "600"))864 except Exception:865 env_ultra, env_min = 400.0, 200.0866 ultra_low_memory_mode = available_mb < env_ultra or progressive_reduction >= 1867 minimal_mode = available_mb < env_min or progressive_reduction >= 2868 869 # Adjust parameters based on available memory870 # Allow explicit overrides so user can test larger context to reduce repeated token delays871 ctx_override = os.environ.get("LLM_CTX" )872 thread_override = os.environ.get("LLM_THREADS")873 try:874 # Set absolute floor at 512 to prevent "forgetting" the schema hints875 ctx_size = int(ctx_override) if ctx_override else (512 if minimal_mode else 768 if ultra_low_memory_mode else 1024)876 except Exception:877 ctx_size = 512878 879 try:880 num_thread = int(thread_override) if thread_override else 1881 except Exception:882 num_thread = 1883 timeout = 12.0 if minimal_mode else 18.0 if ultra_low_memory_mode else 25.0884 885 # Log mode based on memory constraints886 mode_desc = "minimal" if minimal_mode else "ultra-low" if ultra_low_memory_mode else "low"887 print(f"Initializing LLM with ctx_size={ctx_size}, num_thread={num_thread}, timeout={timeout:.1f}s in {mode_desc} memory mode")888 889 # Use smaller model if in minimal mode and memory is critically low890 # Permit explicit model override via env for experimentation (e.g. smaller model to prevent timeouts)891 model_name = os.environ.get("LLM_MODEL_OVERRIDE") or (892 FALLBACK_MODEL_NAME if minimal_mode and available_mb < memory_manager.critical_mb * 1.2 else OLLAMA_MODEL_NAME893 )894 895 # Create the LLM instance with optimized settings but do NOT register yet.896 # We perform a lightweight handshake first and then register only if memory allows.897 llm = None898 try:899 llm = Ollama(900 model=model_name,901 base_url=OLLAMA_BASE_URL,902 temperature=0.05, # Restored from 7.7903 request_timeout=DEFAULT_LLM_REQUEST_TIMEOUT,904 additional_kwargs={905 # Optimized parameters for memory efficiency906 "num_ctx": ctx_size,907 "num_batch": 1, # Restored from 7.7908 "num_gpu": -1, # Keep GPU but with stable batching909 "f16_kv": True, # Half-precision for key/value cache910 "mirostat": 0, # Disable mirostat sampling911 "num_thread": num_thread,912 "seed": 42 # Consistent seed913 }914 )915 916 # Handshake strategy: if strict mode enabled, we block until success (within max timeout budget);917 # otherwise we warn and continue, with a later retry on first real query.918 strict_mode = os.environ.get("LLM_HANDSHAKE_STRICT", "0") not in ("0", "false", "False")919 base_timeout = float(os.environ.get("LLM_HANDSHAKE_TIMEOUT", "30")) # single-attempt timeout920 max_total = float(os.environ.get("LLM_HANDSHAKE_MAX_TOTAL", "120")) # total budget across attempts921 handshake_prompts = ["SELECT 1;", "-- warmup\nSELECT 1;"]922 attempt = 0923 start_all = time.time()924 handshake_ok = False925 while not handshake_ok and (time.time() - start_all) < max_total and attempt < len(handshake_prompts):926 hp = handshake_prompts[attempt]927 attempt += 1928 remaining_budget = max_total - (time.time() - start_all)929 this_timeout = min(base_timeout, remaining_budget)930 print(f"[HS] Handshake attempt {attempt} timeout={this_timeout:.1f}s remaining_budget={remaining_budget:.1f}s")931 try:932 # Use direct call bounded by TimeoutManager to avoid internal indefinite waits.933 resp = TimeoutManager.run_with_timeout(lambda: llm.complete(hp), timeout=this_timeout)934 # Some wrappers return object with .text, others raw string935 _txt = getattr(resp, 'text', None)936 if resp is not None and (_txt is None or 'error' not in str(_txt).lower()):937 elapsed_all = time.time() - start_all938 print(f"[HS] LLM handshake succeeded in {elapsed_all:.2f}s (model={model_name})")939 handshake_ok = True940 break941 except TimeoutError:942 print(f"[HS][WARN] Handshake attempt {attempt} timed out after {this_timeout:.1f}s")943 except Exception as e_hs:944 print(f"[HS][WARN] Handshake attempt {attempt} error: {e_hs}")945 # If model just started loading, give it a short pause before next attempt946 time.sleep(1.5)947 if not handshake_ok:948 msg = "[HS][WARN] Handshake not successful within budget; will retry on first real query." if not strict_mode else "[HS][ERROR] Strict handshake mode enabled and handshake failed." 949 print(msg)950 if strict_mode:951 # In strict mode, raise to stop startup (user explicitly wants guaranteed readiness)952 raise TimeoutError("Strict LLM handshake failed")953 else:954 # Mark readiness955 LLM_HANDSHAKE_COMPLETE = True956 957 # Check memory after initialization/handshake to ensure we didn't consume too much958 after_stats = memory_manager.get_memory_stats()959 memory_used = available_mb - after_stats.get("available_mb", 0)960 if memory_used > 100: # If we used more than 100MB just for initialization961 print(f"⚠️ Warning: LLM initialization consumed {memory_used:.2f}MB")962 # Post-init safety check963 if after_stats.get('available_mb', 0) < ABS_MIN_FREE_MB:964 print(f"[SAFETY] Post-init free memory {after_stats.get('available_mb',0):.1f}MB < ABS_MIN_FREE_MB {ABS_MIN_FREE_MB}MB; tearing down instance")965 try:966 if hasattr(llm, 'close'): llm.close()967 except Exception:968 pass969 raise MemoryError("Post-init memory below absolute floor; aborted")970 971 # Attempt to register the instance with the memory manager now that it's healthy.972 try:973 memory_manager.register_llm_instance(llm)974 except MemoryError as me:975 # Registration refused due to low memory; leave the llm unregistered but return it976 print(f"MemoryManager refused to register LLM: {me}")977 except Exception as reg_err:978 # Non-fatal: log and continue returning the llm (unregistered)979 print(f"Warning: unexpected error registering LLM: {reg_err}")980 981 return llm982 983 except Exception:984 # Ensure we free any partial resources on failure985 try:986 if llm is not None:987 try:988 memory_manager.unregister_llm_instance(llm)989 except Exception:990 pass991 finally:992 # Best-effort GC and working set trim993 gc.collect()994 try:995 memory_manager.force_collect_garbage()996 except Exception:997 pass998 raise999 1000# Create a function to run LLM safely with proper timeouts and cleanup1001def run_llm_with_timeout(llm, prompt, timeout=30.0, retries=1):1002 """1003 Run an LLM with proper timeout handling and resource cleanup.1004 1005 Args:1006 llm: The LLM instance to use1007 prompt: The prompt to send to the LLM1008 timeout: Timeout in seconds1009 retries: Number of retries on timeout or memory error1010 1011 Returns:1012 The LLM response1013 1014 Raises:1015 TimeoutError: If all attempts time out1016 MemoryError: If memory is insufficient after retries1017 """1018 last_error = None1019 1020 for attempt in range(retries + 1):1021 print(f"[LLM][ATTEMPT] Attempt {attempt+1}/{retries+1} for prompt (len={len(prompt)})")1022 # Check memory before attempt and log status1023 stats = memory_manager.get_memory_stats()1024 if stats['available_mb'] < ABS_MIN_FREE_MB:1025 print(f"[LLM][MEMORY] Insufficient memory before attempt: {stats['available_mb']:.1f}MB < ABS_MIN_FREE_MB {ABS_MIN_FREE_MB}MB")1026 raise MemoryError(f"Available memory {stats['available_mb']:.1f}MB < ABS_MIN_FREE_MB {ABS_MIN_FREE_MB}MB; aborting generation")1027 print(f"Memory check before attempt {attempt+1}: {stats['available_mb']:.2f}MB available (ABS_MIN_FREE_MB={ABS_MIN_FREE_MB}MB)")1028 1029 if stats['available_mb'] < memory_manager.threshold_mb:1030 # Try with more aggressive memory reduction1031 reduction_level = attempt + 11032 print(f"[LLM][MEMORY] Memory low, creating LLM with reduction level {reduction_level}")1033 1034 # Release the current LLM instance if it exists1035 if 'llm' in locals() and llm is not None:1036 try:1037 memory_manager.unregister_llm_instance(llm)1038 llm = None1039 except Exception as e:1040 print(f"Error unregistering LLM: {str(e)}")1041 1042 # Force garbage collection1043 gc.collect()1044 memory_manager.force_collect_garbage()1045 1046 # Check if memory is critically low1047 stats = memory_manager.get_memory_stats()1048 if stats['available_mb'] < memory_manager.critical_mb:1049 print(f"🚨 CRITICAL MEMORY CONDITION: Only {stats['available_mb']:.2f}MB available!")1050 raise MemoryError(f"System memory critically low: {stats['available_mb']:.2f}MB available")1051 1052 # Create a new LLM with more aggressive resource reduction1053 llm = initialize_llm(progressive_reduction=reduction_level)1054 1055 # Register with memory manager1056 memory_manager.register_llm_instance(llm)1057 1058 try:1059 # Run the LLM with timeout using ThreadManager for safety1060 print(f"[LLM][RUN] Running LLM (attempt {attempt+1}/{retries+1}) with {timeout:.1f}s timeout")1061 1062 # Start timing1063 start_time = time.time()1064 1065 # First, try talking to an external MCP server if available. This centralizes1066 # model & tool orchestration. If MCP is unreachable, fall back to local LLM.1067 # Prefer local LLM first (preserves pre-MCP behaviour). If local LLM fails or times out,1068 # attempt to use MCP as a fallback so an unavailable MCP won't break everything.1069 with memory_manager.suspend_emergency("llm_generation"):1070 result = None # Initialize result to avoid UnboundLocalError1071 try:1072 print("Attempting local LLM completion first...")1073 # Wrap the LLM call to ensure we pass prompt correctly and bound it with our TimeoutManager1074 result = TimeoutManager.run_with_timeout(lambda: llm.complete(prompt), timeout=timeout)1075 except TimeoutError as local_to:1076 print(f"Local LLM completion timed out after {timeout}s: {local_to}; attempting MCP fallback...")1077 try:1078 from .mcp_client import mcp_complete1079 mcp_text = mcp_complete(prompt, timeout=min( max(5.0, timeout*0.5 ), timeout ))1080 class _Resp:1081 def __init__(self, text):1082 self.text = text1083 result = _Resp(mcp_text)1084 print(f"MCP fallback succeeded after local LLM timeout")1085 except Exception as mcp_err:1086 print(f"MCP completion failed or unavailable: {mcp_err}")1087 # If MCP is clearly unreachable (connection refused), don't treat as critical1088 if "connection refused" in str(mcp_err).lower() or "max retries exceeded" in str(mcp_err).lower():1089 print("[MCP] Server appears to be down - continuing without MCP fallback")1090 # Re-raise the original timeout error for retry logic1091 raise local_to1092 except Exception as local_err:1093 print(f"Local LLM completion failed: {local_err}; attempting MCP fallback...")1094 try:1095 from .mcp_client import mcp_complete1096 mcp_text = mcp_complete(prompt, timeout=min( max(5.0, timeout*0.5 ), timeout ))1097 class _Resp:1098 def __init__(self, text):1099 self.text = text1100 result = _Resp(mcp_text)1101 print(f"MCP fallback succeeded after local LLM failure")1102 except Exception as mcp_err:1103 print(f"MCP completion failed or unavailable: {mcp_err}")1104 # If MCP is clearly unreachable (connection refused), don't treat as critical1105 if "connection refused" in str(mcp_err).lower() or "max retries exceeded" in str(mcp_err).lower():1106 print("[MCP] Server appears to be down - continuing without MCP fallback")1107 # For timeout errors, try a longer timeout retry instead of immediate failure1108 if isinstance(local_err, TimeoutError) and attempt < retries:1109 print(f"[RETRY] Will retry with longer timeout due to initial timeout")1110 # Don't re-raise immediately, let the retry logic handle it1111 raise local_err1112 else:1113 # Re-raise original local error to be handled by outer retry logic1114 raise local_err1115 1116 # Log completion time1117 elapsed = time.time() - start_time1118 print(f"LLM completed in {elapsed:.2f}s")1119 1120 # Check memory after completion1121 stats = memory_manager.get_memory_stats()1122 print(f"Memory after LLM: {stats['available_mb']:.2f}MB available")1123 # If after execution we are below safety target, proactively free (soft protection)1124 if stats['available_mb'] < SAFE_POST_QUERY_TARGET_MB:1125 try:1126 print(f"[SAFETY] Free memory {stats['available_mb']:.1f}MB < SAFE_POST_QUERY_TARGET_MB {SAFE_POST_QUERY_TARGET_MB}MB; releasing LLM to protect system")1127 memory_manager.unregister_llm_instance(llm)1128 except Exception:1129 pass1130 memory_manager.force_collect_garbage()1131 1132 # Force cleanup to prevent memory leaks1133 memory_manager.force_collect_garbage()1134 1135 return result1136 1137 except (TimeoutError, MemoryError) as e:1138 last_error = e1139 print(f"[LLM][FAIL] Attempt {attempt+1} failed: {type(e).__name__}: {str(e)}")1140 if isinstance(e, TimeoutError):1141 print(f"[LLM][TIMEOUT] Timeout occurred for prompt (len={len(prompt)}). Consider simplifying the query or increasing system resources.")1142 1143 # Unregister LLM from memory manager1144 try:1145 memory_manager.unregister_llm_instance(llm)1146 except Exception as unreg_error:1147 print(f"Error unregistering LLM: {str(unreg_error)}")1148 1149 # Aggressive cleanup after error1150 for _ in range(3): # Multiple cleanup attempts1151 gc.collect()1152 time.sleep(0.1) # Brief pause to allow OS to reclaim memory1153 1154 memory_manager.force_collect_garbage()1155 1156 if attempt < retries:1157 # INCREASE timeout for subsequent attempts to handle initialization overhead1158 timeout = min(180.0, timeout * 1.5) # Increase by 50% each retry, cap at 3 minutes1159 print(f"[LLM][RETRY] Retrying with increased timeout={timeout:.1f}s (attempt {attempt+2})")1160 # Escalate LLM reduction after 2 timeouts1161 if isinstance(e, TimeoutError) and attempt >= 1:1162 print(f"[LLM][ESCALATE] Escalating to more aggressive LLM reduction after repeated timeouts.")1163 try:1164 llm = initialize_llm(progressive_reduction=attempt+2)1165 memory_manager.register_llm_instance(llm)1166 except Exception as esc_err:1167 print(f"[LLM][ESCALATE][FAIL] Could not escalate LLM: {esc_err}")1168 else:1169 # Last attempt failed1170 print("[LLM][FAIL] All LLM attempts failed for this prompt.")1171 print(f"[LLM][USER] The LLM failed to respond after multiple attempts. Please try a simpler question, check your system resources, or increase the timeout settings if possible.")1172 # Final cleanup before raising error1173 memory_manager.force_collect_garbage()1174 raise last_error1175 1176 except Exception as e:1177 # For other exceptions, cleanup and don't retry1178 print(f"LLM error: {type(e).__name__}: {str(e)}")1179 1180 # Unregister LLM from memory manager1181 try:1182 memory_manager.unregister_llm_instance(llm)1183 except Exception as unreg_error:1184 print(f"Error unregistering LLM: {str(unreg_error)}")1185 1186 # Cleanup1187 memory_manager.force_collect_garbage()1188 1189 raise1190 1191# Initialize LLM — guard against double-import (Agent.agentic_workflow vs agentic_workflow)1192# Both module paths share the same Ollama process, so re-running initialize_llm() causes1193# redundant handshakes and wastes ~440MB of KV-cache allocation each time.1194_LLM_INIT_SENTINEL = "Agent.__agentic_workflow_llm_initialized__"1195import sys as _sys1196if _sys.modules.get(_LLM_INIT_SENTINEL) is None:1197 _sys.modules[_LLM_INIT_SENTINEL] = True # mark as done before call (prevents re-entry)1198 llm = initialize_llm()1199 Settings.llm = llm1200else: