Update services/masterllm.py

services/masterllm.py

@@ -1,287 +1,287 @@
-# # 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,
+# # 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,
     }