""" Base Agent Abstract base class for all mining intelligence agents. Improvements over v1: - JSON mode via response_mime_type (no more regex JSON extraction) - Retry with exponential backoff (3 retries), respecting Gemini retry_delay - Processes full document via pages, not truncated to 3000 chars - Confidence score required in all agent outputs - Proper QuotaExceededError raised (no more silent empty-dict returns) """ import asyncio import logging import re from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional import google.generativeai as genai from google.generativeai.types import GenerationConfig from app.config import settings logger = logging.getLogger(__name__) # Configure Gemini once at module level genai.configure(api_key=settings.GEMINI_API_KEY) # Generation config that forces JSON output — no more regex parsing _JSON_GENERATION_CONFIG = GenerationConfig( response_mime_type="application/json", temperature=0.1, # Low temperature for consistent structured output top_p=0.95, ) _MAX_RETRIES = 3 _RETRY_BASE_DELAY = 2.0 # seconds — minimum delay between retries _MAX_RETRY_DELAY = 120.0 # seconds — cap for retry_delay parsed from API response class QuotaExceededError(RuntimeError): """Raised when the Gemini API quota / rate-limit is exhausted.""" def _parse_retry_delay(err_str: str) -> Optional[float]: """ Extract the suggested retry_delay (in seconds) from a Gemini 429 error message. The error body contains a line like: retry_delay { seconds: 31 } Returns None if no delay can be parsed. """ match = re.search(r"retry_delay\s*\{\s*seconds:\s*(\d+)", err_str) if match: return min(float(match.group(1)), _MAX_RETRY_DELAY) # Fallback: look for "Please retry in X.Xs" match2 = re.search(r"retry in (\d+\.?\d*)s", err_str) if match2: return min(float(match2.group(1)), _MAX_RETRY_DELAY) return None class BaseAgent(ABC): """ Abstract base class for mining document intelligence agents. Each agent is responsible for a specific analysis task: - Classification - Safety Analysis - Entity Extraction - Summarization """ def __init__( self, model_name: str = None, provider: str = "gemini", fallback_model: str = None, fallback_provider: str = None, ): self.provider = provider self.model_name = model_name or settings.GEMINI_MODEL self.name = self.__class__.__name__ # Fallback config (e.g. Cerebras when Groq is rate-limited) self.fallback_model = fallback_model self.fallback_provider = fallback_provider self._fallback_client = None self._using_fallback = False if self.fallback_provider and self.fallback_model: self._init_fallback_client() self._init_client() def _init_client(self): """Initialize the primary provider client.""" if self.provider == "gemini": self.model = genai.GenerativeModel( model_name=self.model_name, generation_config=_JSON_GENERATION_CONFIG, ) elif self.provider == "groq": from app.services.llm_provider import get_groq_client self.client = get_groq_client() elif self.provider == "mistral": from app.services.llm_provider import get_mistral_client self.client = get_mistral_client() elif self.provider == "cerebras": from app.services.llm_provider import get_cerebras_client self.client = get_cerebras_client() def _init_fallback_client(self): """Initialize the fallback provider client.""" if self.fallback_provider == "cerebras": from app.services.llm_provider import get_cerebras_client self._fallback_client = get_cerebras_client() elif self.fallback_provider == "groq": from app.services.llm_provider import get_groq_client self._fallback_client = get_groq_client() elif self.fallback_provider == "mistral": from app.services.llm_provider import get_mistral_client self._fallback_client = get_mistral_client() logger.info( f"{self.name}: Fallback configured — {self.fallback_provider}/{self.fallback_model}" ) @abstractmethod async def analyze( self, text: str, context: Optional[Dict] = None, ) -> Dict[str, Any]: """ Analyze document text and return structured results. Args: text: Document text content (full text or representative sample) context: Additional context (e.g., document category from classifier) Returns: Dictionary with agent-specific analysis results. All results MUST include a 'confidence' key (0.0–1.0). """ @property @abstractmethod def system_prompt(self) -> str: """System prompt defining the agent's role and capabilities.""" # ── Generation with retry + fallback ─────────────────────────────────────── async def _call_openai_compat(self, client, model: str, prompt: str) -> str: """Call an OpenAI-compatible provider and return text response.""" response = await client.chat.completions.create( model=model, messages=[ {"role": "system", "content": self.system_prompt}, {"role": "user", "content": prompt}, ], response_format={"type": "json_object"}, temperature=0.1, ) return response.choices[0].message.content or "" async def _generate_json(self, prompt: str) -> Dict[str, Any]: """ Generate structured JSON output with retry and automatic provider fallback. Primary provider is tried first. On 429/rate-limit errors, if a fallback is configured (e.g. Cerebras when Groq is rate-limited), the request is retried on the fallback provider before raising QuotaExceededError. """ import json last_error: Optional[Exception] = None for attempt in range(1, _MAX_RETRIES + 1): try: full_prompt = f"{self.system_prompt}\n\n{prompt}" if self.provider == "gemini": response = await asyncio.to_thread( self.model.generate_content, full_prompt, ) text_response = getattr(response, "text", "") elif self.provider in ["groq", "mistral", "cerebras"]: text_response = await self._call_openai_compat( self.client, self.model_name, prompt ) try: return json.loads(text_response) except (json.JSONDecodeError, AttributeError) as parse_err: logger.warning( f"{self.name} attempt {attempt}: JSON parse failed — {parse_err}. " f"Raw response: {text_response[:200]}" ) last_error = parse_err delay = _RETRY_BASE_DELAY * (2 ** (attempt - 1)) await asyncio.sleep(delay) continue except Exception as e: last_error = e err_str = str(e) is_quota = ( "429" in err_str or "rate_limit" in err_str.lower() or "quota" in err_str.lower() or "RESOURCE_EXHAUSTED" in err_str ) if is_quota: # Try fallback provider if available and not already using it if ( self._fallback_client and self.fallback_model and not self._using_fallback ): logger.warning( f"{self.name}: Primary provider rate-limited. " f"Falling back to {self.fallback_provider}/{self.fallback_model}" ) try: text_response = await self._call_openai_compat( self._fallback_client, self.fallback_model, prompt ) result = json.loads(text_response) self._using_fallback = True return result except Exception as fb_err: logger.error(f"{self.name}: Fallback also failed: {fb_err}") # Fall through to raise QuotaExceededError # Parse the suggested wait time from the error body suggested_delay = _parse_retry_delay(err_str) if attempt < _MAX_RETRIES and suggested_delay is not None: logger.warning( f"{self.name}: Quota/rate-limit hit (attempt {attempt}/{_MAX_RETRIES}). " f"Waiting {suggested_delay}s..." ) await asyncio.sleep(suggested_delay) continue logger.error(f"{self.name}: All providers exhausted — {e}") raise QuotaExceededError( f"Rate limit exceeded for {self.name}. " "Please try again later." ) from e # Transient non-quota error — exponential backoff delay = _RETRY_BASE_DELAY * (2 ** (attempt - 1)) logger.warning( f"{self.name} attempt {attempt}/{_MAX_RETRIES} failed: {e}. " f"Retrying in {delay}s..." ) await asyncio.sleep(delay) logger.error(f"{self.name} failed after {_MAX_RETRIES} attempts: {last_error}") return {} # ── Text helpers ─────────────────────────────────────────────────────────── def _prepare_text(self, text: str, max_chars: int = 15000) -> str: """ Prepare text for agent analysis. Instead of hard-truncating to 3000 chars (old behavior), we use up to 15000 chars (≈10 pages) to capture much more document content. Long documents get the first 12000 chars + last 3000 chars to include both the opening context and the conclusion/summary sections. """ if len(text) <= max_chars: return text head = text[:12000] tail = text[-3000:] return head + "\n\n[... middle of document omitted for analysis ...]\n\n" + tail # ── Kept for backward compatibility ─────────────────────────────────────── def _parse_json(self, text: str) -> Dict[str, Any]: """Legacy JSON parser — kept for any subclass that still needs it.""" import json import re text = text.strip() if text.startswith("```"): lines = text.split("\n") text = "\n".join(line for line in lines if not line.startswith("```")) try: return json.loads(text) except json.JSONDecodeError: match = re.search(r"\{.*\}", text, re.DOTALL) if match: try: return json.loads(match.group()) except json.JSONDecodeError: pass return {}