services/masterllm.py

# # services/masterllm.py
# import json
# import requests
# import os
# import re

# # Required: set MISTRAL_API_KEY in the environment
# MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
# if not MISTRAL_API_KEY:
#     raise RuntimeError("Missing MISTRAL_API_KEY environment variable.")

# MISTRAL_ENDPOINT = os.getenv("MISTRAL_ENDPOINT", "https://api.mistral.ai/v1/chat/completions")
# MISTRAL_MODEL = os.getenv("MISTRAL_MODEL", "mistral-small")

# # Steps we support
# ALLOWED_STEPS = {"text", "table", "describe", "summarize", "ner", "classify", "translate"}

# def build_prompt(instruction: str) -> str:
#     return f"""You are a document‑processing assistant.
# Return exactly one JSON object and nothing else — no markdown, no code fences, no explanation, no extra keys.
# Use only the steps the user asks for in the instruction. Do not add any steps not mentioned.
# Valid steps (dash‑separated): {', '.join(sorted(ALLOWED_STEPS))}
# Output schema:
# {{
#   "pipeline": "<dash‑separated‑steps>",
#   "tools": {{ /* object or null */ }},
#   "start_page": <int>,
#   "end_page": <int>,
#   "target_lang": <string or null>
# }}
# Instruction:
# \"\"\"{instruction.strip()}\"\"\"
# """

# def extract_json_block(text: str) -> dict:
#     # Grab everything between the first { and last }
#     start = text.find("{")
#     end = text.rfind("}")
#     if start == -1 or end == -1:
#         return {"error": "no JSON braces found", "raw": text}
#     snippet = text[start:end + 1]
#     try:
#         return json.loads(snippet)
#     except json.JSONDecodeError as e:
#         # attempt to fix common "tools": {null} → "tools": {}
#         cleaned = re.sub(r'"tools"\s*:\s*\{null\}', '"tools": {}', snippet)
#         try:
#             return json.loads(cleaned)
#         except json.JSONDecodeError:
#             return {"error": f"json decode error: {e}", "raw": snippet}

# def validate_pipeline(cfg: dict) -> dict:
#     pipe = cfg.get("pipeline")
#     if isinstance(pipe, list):
#         pipe = "-".join(pipe)
#         cfg["pipeline"] = pipe
#     if not isinstance(pipe, str):
#         return {"error": "pipeline must be a string"}

#     steps = pipe.split("-")
#     bad = [s for s in steps if s not in ALLOWED_STEPS]
#     if bad:
#         return {"error": f"invalid steps: {bad}"}

#     # translate requires target_lang
#     if "translate" in steps and not cfg.get("target_lang"):
#         return {"error": "target_lang required for translate"}
#     return {"ok": True}

# def _sanitize_config(cfg: dict) -> dict:
#     # Defaults and types
#     try:
#         sp = int(cfg.get("start_page", 1))
#     except Exception:
#         sp = 1
#     try:
#         ep = int(cfg.get("end_page", sp))
#     except Exception:
#         ep = sp
#     if sp < 1:
#         sp = 1
#     if ep < sp:
#         ep = sp
#     cfg["start_page"] = sp
#     cfg["end_page"] = ep

#     # Ensure tools is an object
#     if cfg.get("tools") is None:
#         cfg["tools"] = {}

#     # Normalize pipeline separators (commas, spaces → dashes)
#     raw_pipe = cfg.get("pipeline", "")
#     steps = [s.strip() for s in re.split(r"[,\s\-]+", raw_pipe) if s.strip()]
#     # Deduplicate while preserving order
#     dedup = []
#     for s in steps:
#         if s in ALLOWED_STEPS and s not in dedup:
#             dedup.append(s)
#     cfg["pipeline"] = "-".join(dedup)

#     # Normalize target_lang
#     if "target_lang" in cfg and cfg["target_lang"] is not None:
#         t = str(cfg["target_lang"]).strip()
#         cfg["target_lang"] = t if t else None

#     return cfg

# def generate_pipeline(instruction: str) -> dict:
#     prompt = build_prompt(instruction)
#     res = requests.post(
#         MISTRAL_ENDPOINT,
#         headers={
#             "Authorization": f"Bearer {MISTRAL_API_KEY}",
#             "Content-Type": "application/json",
#         },
#         json={
#             "model": MISTRAL_MODEL,
#             "messages": [{"role": "user", "content": prompt}],
#             "temperature": 0.0,
#             "max_tokens": 256,
#         },
#         timeout=60,
#     )
#     res.raise_for_status()
#     content = res.json()["choices"][0]["message"]["content"]

#     parsed = extract_json_block(content)
#     if "error" in parsed:
#         raise RuntimeError(f"PARSE_ERROR: {parsed['error']}\nRAW_OUTPUT:\n{parsed.get('raw', content)}")

#     # Sanitize and normalize
#     parsed = _sanitize_config(parsed)

#     check = validate_pipeline(parsed)
#     if "error" in check:
#         raise RuntimeError(f"PARSE_ERROR: {check['error']}\nRAW_OUTPUT:\n{content}")

#     return parsed


# services/masterllm.py
import json
import os
import re
from typing import Dict, Any, List

import requests

# Google Gemini API configuration
# Free tier: 15 RPM, 1M TPM, 1500 RPD for gemini-1.5-flash
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-2.0-flash")
GEMINI_ENDPOINT = f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_MODEL}:generateContent"

_TOOL_TO_TOKEN = {
    "extract_text": "text",
    "extract_tables": "table",
    "describe_images": "describe",
    "summarize_text": "summarize",
    "classify_text": "classify",
    "extract_entities": "ner",
    "translate_text": "translate",
    "signature_verification": "signature",
    "stamp_detection": "stamp",
}

_ALLOWED_TOOLS = list(_TOOL_TO_TOKEN.keys())


def _invoke_gemini(prompt: str) -> str:
    """
    Invoke Google Gemini API for pipeline planning.
    Free tier: 15 RPM, 1M TPM, 1500 RPD for gemini-1.5-flash
    """
    if not GEMINI_API_KEY:
        raise RuntimeError("Missing GEMINI_API_KEY or GOOGLE_API_KEY environment variable")
    
    headers = {
        "Content-Type": "application/json",
    }
    
    payload = {
        "contents": [{
            "parts": [{"text": prompt}]
        }],
        "generationConfig": {
            "temperature": 0.0,
            "maxOutputTokens": 512,
        }
    }
    
    response = requests.post(
        f"{GEMINI_ENDPOINT}?key={GEMINI_API_KEY}",
        headers=headers,
        json=payload,
        timeout=60,
    )
    
    if response.status_code != 200:
        raise RuntimeError(f"Gemini API error: {response.status_code} - {response.text}")
    
    result = response.json()
    
    # Extract text from Gemini response
    try:
        return result["candidates"][0]["content"]["parts"][0]["text"]
    except (KeyError, IndexError) as e:
        raise RuntimeError(f"Failed to parse Gemini response: {e}\nResponse: {result}")


def generate_pipeline(user_instruction: str) -> Dict[str, Any]:
    """
    Produce a proposed plan as a compact pipeline string + config.
    Output example:
    {
      "pipeline": "text-table-summarize",
      "start_page": 1,
      "end_page": 3,
      "target_lang": null,
      "tools": ["extract_text", "extract_tables", "summarize_text"],
      "reason": "..."
    }
    """
    system_prompt = f"""You design a tool execution plan for MasterLLM. 
Return STRICT JSON with keys:
- pipeline: string of hyphen-joined steps using tokens: text, table, describe, summarize, classify, ner, translate, signature, stamp
- tools: array of tool names from: {", ".join(_ALLOWED_TOOLS)}
- start_page: integer (default 1)
- end_page: integer (default start_page)
- target_lang: string or null
- reason: short rationale
Extract any page range or language from the user's request.

User instruction: {user_instruction}

Return only the JSON object, no markdown or explanation."""

    raw = _invoke_gemini(system_prompt)

    # best-effort JSON extraction
    try:
        data = json.loads(raw)
    except Exception:
        match = re.search(r"\{.*\}", raw, re.S)
        data = json.loads(match.group(0)) if match else {}

    # Fallbacks / validation
    tools: List[str] = data.get("tools") or []
    # Map tools -> pipeline tokens
    tokens = [_TOOL_TO_TOKEN[t] for t in tools if t in _TOOL_TO_TOKEN]
    if not tokens:
        # heuristic fallback
        text_lower = user_instruction.lower()
        if "table" in text_lower:
            tokens.append("table")
        if any(w in text_lower for w in ["text", "extract", "read", "content"]):
            tokens.insert(0, "text")
        if any(w in text_lower for w in ["summarize", "summary"]):
            tokens.append("summarize")
        if any(w in text_lower for w in ["translate", "spanish", "french", "german"]):
            tokens.append("translate")
        if any(w in text_lower for w in ["classify", "category", "categories"]):
            tokens.append("classify")
        if any(w in text_lower for w in ["ner", "entity", "entities"]):
            tokens.append("ner")
        if any(w in text_lower for w in ["image", "figure", "diagram", "photo"]):
            tokens.append("describe")
    pipeline = "-".join(tokens) if tokens else "text"

    start_page = int(data.get("start_page") or 1)
    end_page = int(data.get("end_page") or start_page)
    target_lang = data.get("target_lang") if data.get("target_lang") not in ["", "none", None] else None

    # if tools empty but tokens present, infer tools from tokens
    if not tools and tokens:
        inv = {v: k for k, v in _TOOL_TO_TOKEN.items()}
        tools = [inv[t] for t in tokens if t in inv]

    return {
        "pipeline": pipeline,
        "start_page": start_page,
        "end_page": end_page,
        "target_lang": target_lang,
        "tools": tools,
        "reason": data.get("reason") or "Auto-generated plan.",
        "raw_instruction": user_instruction,
    }