feat: LLM retry logic + refactor chunked extraction

src/llm/chunked_extraction.py

@@ -3,7 +3,6 @@
 import sys
 import gc
 import json
-import joblib
 import xgboost as xgb
 from pathlib import Path
 from collections import Counter
@@ -13,21 +12,7 @@
 sys.path.insert(0, str(project_root))
 
 from src.llm.llm_client import extract_metrics_from_text
-
-
-def load_classifier(model_dir="src/model/models"):
-    """Load XGBoost classifier and TF-IDF vectorizer."""
-    model_path = Path(model_dir) / "pdf_classifier.json"
-    vectorizer_path = Path(model_dir) / "tfidf_vectorizer.pkl"
-
-    if not model_path.exists():
-        raise FileNotFoundError(f"Model not found: {model_path}")
-
-    model = xgb.Booster()
-    model.load_model(str(model_path))
-    vectorizer = joblib.load(vectorizer_path)
-
-    return model, vectorizer
+from src.model.pdf_classifier import load_classifier
 
 
 def chunk_text(text, chunk_size=3000, overlap=300):
@@ -100,15 +85,15 @@ def extract_with_chunking(
     text,
     model_dir="src/model/models",
     llm_model="qwen2.5:7b",  # Changed from biomistral
-    num_ctx=2048,
+    num_ctx=8192,
     top_n=3,
     chunk_size=3000,
     overlap=300,
 ):
     """Main extraction with chunking pipeline."""
 
     print(f"  [CHUNK] Loading classifier...", file=sys.stderr)
-    model, vectorizer = load_classifier(model_dir)
+    model, vectorizer, _encoder = load_classifier(model_dir)
 
     chunks = chunk_text(text, chunk_size, overlap)
     print(f"  [CHUNK] Split into {len(chunks)} chunks", file=sys.stderr)
@@ -139,7 +124,7 @@ def extract_with_chunking(
 
         try:
             metrics = extract_metrics_from_text(
-                text=chunk[:2500],
+                text=chunk,
                 model=llm_model,
                 num_ctx=num_ctx,
             )

src/llm/llm_client.py

@@ -16,6 +16,8 @@
 import logging
 import re
 import sys
+import time
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
 from pathlib import Path
 
 from ollama import chat
@@ -29,6 +31,32 @@
 
 log = logging.getLogger(__name__)
 
+_OLLAMA_TIMEOUT = 120  # seconds
+_MAX_RETRIES = 3
+_BACKOFF_BASE = 2  # seconds
+
+
+def _call_ollama_with_retry(model, messages, format, options):
+    """Call Ollama with timeout and exponential backoff on transient failures."""
+    transient_errors = (ConnectionError, OSError, TimeoutError, FuturesTimeoutError)
+
+    for attempt in range(_MAX_RETRIES):
+        try:
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(chat, messages=messages, model=model, format=format, options=options)
+                return future.result(timeout=_OLLAMA_TIMEOUT)
+        except FuturesTimeoutError:
+            log.warning("Ollama call timed out (attempt %d/%d)", attempt + 1, _MAX_RETRIES)
+        except transient_errors as e:
+            log.warning("Transient Ollama error (attempt %d/%d): %s", attempt + 1, _MAX_RETRIES, e)
+
+        if attempt < _MAX_RETRIES - 1:
+            wait = _BACKOFF_BASE**attempt
+            log.info("Retrying in %ds...", wait)
+            time.sleep(wait)
+
+    raise RuntimeError(f"Ollama call failed after {_MAX_RETRIES} attempts")
+
 
 def extract_metrics_from_text(
     text: str,
@@ -125,9 +153,9 @@ def extract_metrics_from_text(
 TEXT
 {text}
 """
-    response = chat(
-        messages=[{"role": "user", "content": prompt}],
+    response = _call_ollama_with_retry(
         model=model,
+        messages=[{"role": "user", "content": prompt}],
         format=PredatorDietMetrics.model_json_schema(),
         options={"num_ctx": num_ctx},
     )
@@ -184,9 +212,9 @@ def extract_metrics_from_text(
             retry_prompt += _hints.get(field, "")
         retry_prompt += f"\nTEXT\n{text}"
 
-        retry_response = chat(
-            messages=[{"role": "user", "content": retry_prompt}],
+        retry_response = _call_ollama_with_retry(
             model=model,
+            messages=[{"role": "user", "content": retry_prompt}],
             format=PredatorDietMetrics.model_json_schema(),
             options={"num_ctx": num_ctx},
         )