Update gemini_utils.py

gemini_utils.py

@@ -1,11 +1,10 @@
 import os
 import json
 import re
-from typing import Optional
+from typing import Optional, Tuple
 
 from openai import OpenAI
 
-
 # Environment configuration
 NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
 NIM_BASE_URL = os.getenv("NIM_BASE_URL", "https://integrate.api.nvidia.com/v1")
@@ -13,63 +12,156 @@ NIM_MODEL_DEFAULT = os.getenv("NIM_MODEL", "meta/llama-3.1-8b-instruct")
 
 # Sanitization utilities
 _PREFACE_RE = re.compile(r"^(okay[, ]|sure[, ]|here(?:'|’)s|summary:?|note:?|context:)\b", re.I)
-# Customize anchors for your domain if you have reliable headings
+# Customize anchors to your domain. These help clamp to the main heading if present.
 _ANCHOR_RE = re.compile(r"\b(meeting\s*minutes|minutes\s*of\s*meeting|invoice|report|summary)\b", re.I)
 _DOC_BLOCK_RE = re.compile(r"\[\[\[DOC\]\]\](.*)\[\[\[\/DOC\]\]\]", re.S)
 
-
 def _sanitize_preface(text: str) -> str:
-    """Remove typical LLM prefaces and trim to a reliable anchor if present."""
     s = (text or "").lstrip()
-
-    # Remove obvious preface lines at the start
     lines = s.splitlines()
     while lines and _PREFACE_RE.match(lines[0].strip()):
         lines.pop(0)
     s = "\n".join(lines).lstrip()
 
-    # If your documents have a reliable heading/anchor, trim to it
     m = _ANCHOR_RE.search(s)
     if m:
         s = s[m.start():]
 
     return s.strip()
 
-
 def _extract_marked_block(text: str) -> Optional[str]:
-    """Keep only [[[DOC]]] ... [[[/DOC]]] if present."""
     m = _DOC_BLOCK_RE.search(text or "")
     if m:
         return m.group(1).strip()
     return None
 
+def _finalize(out: str, fallback: str) -> str:
+    # Hard clamp to markers if present
+    block = _extract_marked_block(out)
+    if block:
+        out = block
+    out = _sanitize_preface(out)
+    return out or fallback
+
+def _nim_call_with_tools(
+    client: OpenAI, model_name: str, system: str, user: str, timeout_s: int
+) -> str:
+    """
+    Strict mode: require a function call so the model must return arguments only.
+    Parse tool_calls and extract enhanced_text from function.arguments.
+    """
+    resp = client.chat.completions.create(
+        model=model_name,
+        messages=[
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ],
+        temperature=0.0,
+        top_p=1.0,
+        max_tokens=8192,
+        tools=[{
+            "type": "function",
+            "function": {
+                "name": "return_enhanced_text",
+                "description": "Return the enhanced document text only.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"enhanced_text": {"type": "string"}},
+                    "required": ["enhanced_text"]
+                }
+            }
+        }],
+        # Require the function; some NIM deployments fully support this,
+        # others may error. We'll catch and fallback if needed.
+        tool_choice={"type": "function", "function": {"name": "return_enhanced_text"}},
+        timeout=timeout_s,
+    )
+
+    msg = resp.choices[0].message
+    # OpenAI-compatible: tool_calls array with function name + arguments JSON
+    tool_calls = getattr(msg, "tool_calls", None)
+    if tool_calls:
+        for tc in tool_calls:
+            fn = tc.function
+            if fn and fn.name == "return_enhanced_text":
+                try:
+                    args = json.loads(fn.arguments or "{}")
+                    val = args.get("enhanced_text")
+                    if isinstance(val, str) and val.strip():
+                        return val.strip()
+                except Exception:
+                    pass
+
+    # If the provider didn’t do a tool call or arguments failed to parse,
+    # try content JSON as a fallback path in this same response.
+    content = (msg.content or "").strip()
+    if content:
+        try:
+            obj = json.loads(content)
+            val = obj.get("enhanced_text")
+            if isinstance(val, str) and val.strip():
+                return val.strip()
+        except json.JSONDecodeError:
+            pass
+
+    # No usable output
+    return (msg.content or "").strip()
+
+def _nim_call_json_only(
+    client: OpenAI, model_name: str, system: str, user: str, timeout_s: int
+) -> str:
+    """
+    JSON-only mode: enforce response_format={"type": "json_object"} and parse enhanced_text.
+    """
+    resp = client.chat.completions.create(
+        model=model_name,
+        messages=[
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ],
+        temperature=0.1,
+        top_p=1.0,
+        max_tokens=8192,
+        response_format={"type": "json_object"},
+        timeout=timeout_s,
+    )
+    content = (resp.choices[0].message.content or "").strip()
+    try:
+        obj = json.loads(content)
+        val = obj.get("enhanced_text")
+        if isinstance(val, str) and val.strip():
+            return val.strip()
+    except json.JSONDecodeError:
+        pass
+    return content
 
-def enhance_with_nim(
+def enhance_doc(
     extracted_text: str,
     user_prompt: str,
-    model: Optional[str] = None,
+    nim_model: Optional[str] = None,
     timeout_s: int = 60,
 ) -> str:
     """
-    Enhance document using NVIDIA NIM (OpenAI-compatible Chat Completions).
-    Enforces JSON-only output: {"enhanced_text": "..."}.
-    Returns only the enhanced text (string). On any failure, returns original text.
+    Public entrypoint (NIM only):
+      1) Try strict function-calling
+      2) Fallback to JSON-only
+      3) Final clamps (markers + preface removal)
+      4) Return enhanced text or original on failure
     """
     if not NVIDIA_API_KEY:
-        # No key available -> return original text
         return extracted_text
 
-    model_name = model or NIM_MODEL_DEFAULT
+    model_name = nim_model or NIM_MODEL_DEFAULT
     client = OpenAI(api_key=NVIDIA_API_KEY, base_url=NIM_BASE_URL)
 
     system = (
         "You are a professional document editor. Edit and improve the provided document according to the user's "
         "instructions while preserving meaning, structure, headings, lists, and tone. "
         "Do not include any preface, summary, or explanation. "
-        "Return only JSON with a single field 'enhanced_text'. "
-        "If you add any extra commentary, it will be ignored.\n"
-        "Optionally, also wrap the final edited document between markers [[[DOC]]] and [[[/DOC]]] "
-        "if you must return any non-JSON content."
+        "Return the result only via a function call named 'return_enhanced_text' with a single string field "
+        "'enhanced_text'. If you cannot call the function, return only JSON with a single field 'enhanced_text'. "
+        "If you produce any extra commentary, it will be discarded.\n"
+        "If you must show raw text, wrap the final edited document strictly between [[[DOC]]] and [[[/DOC]]]."
     )
 
     user = f"""User instructions:
@@ -79,54 +171,22 @@ Original document:
 {extracted_text}
 """
 
+    # 1) Strict function-calling attempt
     try:
-        resp = client.chat.completions.create(
-            model=model_name,
-            messages=[
-                {"role": "system", "content": system},
-                {"role": "user", "content": user},
-            ],
-            temperature=0.1,
-            top_p=1.0,
-            max_tokens=8192,
-            response_format={"type": "json_object"},  # enforce JSON
-            timeout=timeout_s,
-        )
-
-        content = (resp.choices[0].message.content or "").strip()
-
-        # Expect a JSON object like {"enhanced_text": "..."}
-        try:
-            obj = json.loads(content)
-            out = obj.get("enhanced_text")
-            if isinstance(out, str) and out.strip():
-                out = out.strip()
-            else:
-                out = content
-        except json.JSONDecodeError:
-            # If model ignored JSON, use raw (then clamp below)
-            out = content
-
-        # Final clamps
-        block = _extract_marked_block(out)
-        if block:
-            out = block
-        out = _sanitize_preface(out)
-
-        return out or extracted_text
-
+        out = _nim_call_with_tools(client, model_name, system, user, timeout_s)
+        if out:
+            return _finalize(out, extracted_text)
     except Exception:
-        # On any error, return original text
-        return extracted_text
+        # Some NIM deployments may not support tool_choice strictly; fall back.
+        pass
 
+    # 2) JSON-only attempt
+    try:
+        out = _nim_call_json_only(client, model_name, system, user, timeout_s)
+        if out:
+            return _finalize(out, extracted_text)
+    except Exception:
+        pass
 
-def enhance_doc(
-    extracted_text: str,
-    user_prompt: str,
-    nim_model: Optional[str] = None,
-) -> str:
-    """
-    Public entrypoint: enhance via NIM only.
-    Returns the enhanced text or the original text on failure.
-    """
-    return enhance_with_nim(extracted_text, user_prompt, model=nim_model)
+    # 3) Last resort
+    return extracted_text