Update gemini_utils.py

gemini_utils.py

@@ -1,10 +1,11 @@
 import os
 import json
 import re
-from typing import Optional, Tuple
+from typing import Optional
 
 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")
@@ -12,156 +13,63 @@ 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 to your domain. These help clamp to the main heading if present.
+# Customize anchors for your domain if you have reliable headings
 _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_doc(
+def enhance_with_nim(
     extracted_text: str,
     user_prompt: str,
-    nim_model: Optional[str] = None,
+    model: Optional[str] = None,
     timeout_s: int = 60,
 ) -> str:
     """
-    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
+    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.
     """
     if not NVIDIA_API_KEY:
+        # No key available -> return original text
         return extracted_text
 
-    model_name = nim_model or NIM_MODEL_DEFAULT
+    model_name = 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 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]]]."
+        "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."
     )
 
     user = f"""User instructions:
@@ -171,22 +79,54 @@ Original document:
 {extracted_text}
 """
 
-    # 1) Strict function-calling attempt
     try:
-        out = _nim_call_with_tools(client, model_name, system, user, timeout_s)
-        if out:
-            return _finalize(out, extracted_text)
-    except Exception:
-        # Some NIM deployments may not support tool_choice strictly; fall back.
-        pass
+        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
 
-    # 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
+        # On any error, return original text
+        return extracted_text
+
 
-    # 3) Last resort
-    return extracted_text
+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)