update LLM api

app.py

@@ -2,6 +2,7 @@
 
 import os
 import json
+import gzip
 import requests
 import gradio as gr
 from openai import OpenAI
@@ -18,9 +19,15 @@ def _require_env(name: str) -> str:
     return value
 
 AUDIOGEN_API_URL           = _require_env("AUDIOGEN_API_URL")
-LLM_BASE_URL               = _require_env("LLM_BASE_URL")
+LLM_BASE_URL               = os.environ.get("LLM_BASE_URL", "")
+CLAW_API_URL               = os.environ.get("CLAW_API_URL", "")
 PROMPT_REFINER_MAX_RETRIES = 3
 
+# Prompt Refiner calling mode:
+#   "claw"   — POST plain text to CLAW_API_URL, no auth required (default)
+#   "openai" — OpenAI-compatible chat completions via LLM_BASE_URL + API_KEY
+PROMPT_REFINER_MODE = os.environ.get("PROMPT_REFINER_MODE", "claw")
+
 # Special token order and mapping
 SPECIAL_TOKEN_ORDER = ["caption", "speech", "sfx", "music", "env", "asr"]
 SPECIAL_TOKEN_MAP = {
@@ -82,31 +89,112 @@ def call_audiogen(structured_prompt):
         return None, f"Error: {str(e)}"
 
 
-def call_prompt_refiner(user_input, max_retries=PROMPT_REFINER_MAX_RETRIES):
-    """Call the Prompt Refiner via LLM API.
+def _parse_and_validate(raw_content: str, attempt: int):
+    """Parse JSON and validate required 'caption' field. Returns (dict|None, error_str)."""
+    try:
+        parsed = json.loads(raw_content)
+    except json.JSONDecodeError as e:
+        return None, f"Invalid JSON on attempt {attempt}: {e}"
+
+    normalized = {
+        k.lower(): v
+        for k, v in parsed.items()
+        if v is not None and str(v).strip()
+    }
+    if not normalized.get("caption"):
+        return None, f"Missing required 'Caption' field on attempt {attempt}."
+    return normalized, None
+
 
-    Returns a dict with lowercase keys: caption, speech, sfx, music, env, asr.
-    Retries up to `max_retries` times on JSON / validation errors.
-    Raises EnvironmentError if required env vars are missing.
-    Raises RuntimeError on unrecoverable API or repeated failures.
+def _decode_claw_response_json(response: requests.Response) -> dict:
+    """Decode CLAW response robustly, including mis-labeled gzip responses."""
+    raw_bytes = response.raw.read(decode_content=False)
+
+    candidates = []
+    # 1) Treat as plain utf-8 text first (some responses are plain text but mislabeled)
+    candidates.append(raw_bytes.decode("utf-8", errors="replace"))
+    # 2) Try gzip decode as fallback when content-encoding is incorrect/mixed
+    try:
+        candidates.append(gzip.decompress(raw_bytes).decode("utf-8", errors="replace"))
+    except Exception:
+        pass
+
+    last_err = None
+    for text in candidates:
+        try:
+            return json.loads(text)
+        except Exception as e:
+            last_err = e
+    raise ValueError(f"Unable to decode CLAW JSON response: {last_err}")
+
+
+def _call_prompt_refiner_claw(user_input: str, max_retries: int) -> dict:
+    """Call Prompt Refiner via CLAW endpoint (no auth required).
+    Sends the full prompt template with user input substituted as plain text.
     """
+    # Substitute user input into the prompt template
+    full_prompt = _PROMPT_REFINER_SYSTEM.replace("{{user_input}}", user_input).strip()
+
+    last_error = None
+    for attempt in range(1, max_retries + 1):
+        try:
+            response = requests.post(
+                CLAW_API_URL,
+                headers={"Content-Type": "text/plain"},
+                data=full_prompt.encode("utf-8"),
+                timeout=60,
+                stream=True,
+            )
+            response.raise_for_status()
+
+            # Response has same structure as OpenAI: choices[0].message.content
+            resp_json = _decode_claw_response_json(response)
+            raw_content = resp_json["choices"][0]["message"]["content"]
+
+            result, err = _parse_and_validate(raw_content, attempt)
+            if err:
+                last_error = err
+                continue
+            return result
+
+        except requests.exceptions.HTTPError as e:
+            code = e.response.status_code
+            raise RuntimeError(f"CLAW API HTTP error {code}: {e.response.reason}") from e
+        except requests.exceptions.ConnectionError as e:
+            raise RuntimeError(f"CLAW API connection error: {e}") from e
+        except requests.exceptions.Timeout:
+            last_error = f"CLAW API timed out on attempt {attempt}."
+        except Exception as e:
+            last_error = f"CLAW API error on attempt {attempt}: {e}"
+
+    raise RuntimeError(
+        f"Prompt Refiner (claw) failed after {max_retries} attempt(s). "
+        f"Last error: {last_error}"
+    )
+
+
+def _call_prompt_refiner_openai(user_input: str, max_retries: int) -> dict:
+    """Call Prompt Refiner via OpenAI-compatible chat completions endpoint."""
     api_key    = os.environ.get("API_KEY")
     model_name = os.environ.get("MODEL_NAME")
 
     if not api_key:
         raise EnvironmentError(
             "API_KEY environment variable is not set. "
-            "Please set it before using Auto Mode."
+            "Please set it before using Auto Mode (openai mode)."
         )
     if not model_name:
         raise EnvironmentError(
             "MODEL_NAME environment variable is not set. "
-            "Please set it before using Auto Mode."
+            "Please set it before using Auto Mode (openai mode)."
+        )
+    if not LLM_BASE_URL:
+        raise EnvironmentError(
+            "LLM_BASE_URL environment variable is not set. "
+            "Please set it before using Auto Mode (openai mode)."
         )
 
     client = OpenAI(api_key=api_key, base_url=LLM_BASE_URL)
-
-    # Remove the {{user_input}} placeholder line; user text goes as user message
     system_content = _PROMPT_REFINER_SYSTEM.replace("{{user_input}}", "").strip()
 
     last_error = None
@@ -123,43 +211,45 @@ def call_prompt_refiner(user_input, max_retries=PROMPT_REFINER_MAX_RETRIES):
             )
             raw_content = completion.choices[0].message.content
 
-            # ── JSON decode validation ────────────────────────────────────────
-            try:
-                parsed = json.loads(raw_content)
-            except json.JSONDecodeError as e:
-                last_error = f"Invalid JSON on attempt {attempt}: {e}"
-                continue  # retry
-
-            # ── Field validation ──────────────────────────────────────────────
-            # Normalize all keys to lowercase; drop null values
-            normalized = {
-                k.lower(): v
-                for k, v in parsed.items()
-                if v is not None and str(v).strip()
-            }
-            if not normalized.get("caption"):
-                last_error = f"Missing required 'Caption' field on attempt {attempt}."
-                continue  # retry
-
-            return normalized
+            result, err = _parse_and_validate(raw_content, attempt)
+            if err:
+                last_error = err
+                continue
+            return result
 
         except EnvironmentError:
-            raise  # propagate immediately
+            raise
         except Exception as e:
             err_str = str(e).lower()
-            # Don't retry on auth / key errors
             if any(kw in err_str for kw in ("authentication", "api_key", "unauthorized", "403", "401")):
                 raise RuntimeError(f"Prompt Refiner auth error: {e}") from e
             last_error = f"API error on attempt {attempt}: {e}"
-            if attempt == max_retries:
-                break
 
     raise RuntimeError(
-        f"Prompt Refiner failed after {max_retries} attempt(s). "
+        f"Prompt Refiner (openai) failed after {max_retries} attempt(s). "
         f"Last error: {last_error}"
     )
 
 
+def call_prompt_refiner(user_input, max_retries=PROMPT_REFINER_MAX_RETRIES):
+    """Dispatch to the configured Prompt Refiner backend.
+
+    Mode is controlled by the PROMPT_REFINER_MODE environment variable:
+      'claw'   — CLAW plain-text endpoint, no auth required (default)
+      'openai' — OpenAI-compatible chat completions endpoint
+    """
+    mode = PROMPT_REFINER_MODE.lower()
+    if mode == "openai":
+        return _call_prompt_refiner_openai(user_input, max_retries)
+    elif mode == "claw":
+        return _call_prompt_refiner_claw(user_input, max_retries)
+    else:
+        raise ValueError(
+            f"Unknown PROMPT_REFINER_MODE '{mode}'. "
+            "Valid values: 'claw' (default), 'openai'."
+        )
+
+
 def build_caption_from_refined(refined: dict) -> str:
     """Build the full structured prompt string from a refined dict.
     This is a convenience wrapper around build_structured_prompt."""