Update llm_client.py

llm_client.py

@@ -1,198 +1,184 @@
-#!/usr/bin/env python3
-"""
-Local LLM client abstraction (NO OpenAI/Claude).
-
-Providers:
-- ollama        : calls a local Ollama server (your Windows dev)
-- transformers  : runs a local HF model in-process (best for Hugging Face Spaces CPU)
-
-Env:
-  LOCAL_LLM_PROVIDER=ollama|transformers
-
-Ollama:
-  OLLAMA_HOST=http://localhost:11434
-  OLLAMA_MODEL=llama3.2:1b
-
-Transformers:
-  HF_LLM_MODEL=TinyLlama/TinyLlama-1.1B-Chat-v1.0   (recommended CPU default)
-  HF_MAX_NEW_TOKENS=450
-"""
-
-from __future__ import annotations
-
-import json
-import os
-import re
-from typing import Any, Dict, Optional
-
-import requests
-
-
-class LocalLLMClient:
-    def __init__(
-        self,
-        provider: Optional[str] = None,
-        model: Optional[str] = None,
-        host: Optional[str] = None,
-        timeout_sec: int = 120,
-    ):
-        self.provider = (provider or os.getenv("LOCAL_LLM_PROVIDER", "ollama")).lower().strip()
-        self.timeout_sec = int(os.getenv("LLM_TIMEOUT_SEC", str(timeout_sec)))
-
-        # Ollama settings
-        self.host = (host or os.getenv("OLLAMA_HOST", "http://localhost:11434")).strip()
-        self.model = (model or os.getenv("OLLAMA_MODEL", "llama3.2:1b")).strip()
-
-        # Transformers settings (HF Spaces)
-        self.hf_model_id = (os.getenv("HF_LLM_MODEL", "TinyLlama/TinyLlama-1.1B-Chat-v1.0")).strip()
-        self.hf_max_new_tokens = int(os.getenv("HF_MAX_NEW_TOKENS", "450"))
-
-        self._hf_pipe = None  # lazy init
-
-        if self.provider not in {"ollama", "transformers"}:
-            raise ValueError(
-                f"Unsupported LOCAL_LLM_PROVIDER='{self.provider}'. "
-                "Supported: ollama, transformers."
-            )
-
-    # --------------------------- Public API ---------------------------
-
-    def generate(self, prompt: str, temperature: float = 0.2, max_tokens: int = 900) -> str:
-        prompt = (prompt or "").strip()
-        if not prompt:
-            return ""
-
-        if self.provider == "ollama":
-            return self._generate_ollama(prompt, temperature=temperature, max_tokens=max_tokens)
-
-        # transformers
-        return self._generate_transformers(prompt, temperature=temperature, max_tokens=max_tokens)
-
-    # --------------------------- Ollama ---------------------------
-
-    def _generate_ollama(self, prompt: str, temperature: float, max_tokens: int) -> str:
-        url = self.host.rstrip("/") + "/api/generate"
-        payload: Dict[str, Any] = {
-            "model": self.model,
-            "prompt": prompt,
-            "stream": False,
-            "options": {
-                "temperature": float(temperature),
-                "num_predict": int(max_tokens),
-            },
-        }
-
-        try:
-            r = requests.post(url, json=payload, timeout=self.timeout_sec)
-        except requests.RequestException as e:
-            raise RuntimeError(
-                "Failed to connect to local Ollama.\n"
-                f"Tried: {url}\n"
-                "Fix:\n"
-                "  - Ensure Ollama is running\n"
-                "  - Confirm endpoint: iwr http://localhost:11434/api/tags -UseBasicParsing\n"
-                f"Error: {repr(e)}"
-            ) from e
-
-        if r.status_code != 200:
-            body = (r.text or "").strip()
-            msg = body
-            try:
-                j = r.json()
-                if isinstance(j, dict):
-                    msg = j.get("error") or j.get("message") or body
-            except Exception:
-                pass
-            raise RuntimeError(
-                "Ollama returned an error.\n"
-                f"URL: {url}\n"
-                f"HTTP: {r.status_code}\n"
-                f"Model: {self.model}\n"
-                f"Details: {msg}"
-            )
-
-        data = r.json()
-        return (data.get("response") or "").strip()
-
-    # --------------------------- Transformers (HF Spaces) ---------------------------
-
-    def _lazy_init_hf(self):
-        if self._hf_pipe is not None:
-            return
-
-        # Lazy import to keep local installs lighter
-        from transformers import pipeline
-
-        # CPU inference; use bfloat16 only if supported (some spaces may not)
-        # Keep it simple and robust.
-        self._hf_pipe = pipeline(
-            "text-generation",
-            model=self.hf_model_id,
-            device=-1,  # CPU
-        )
-
-    def _generate_transformers(self, prompt: str, temperature: float, max_tokens: int) -> str:
-        self._lazy_init_hf()
-
-        # Cap generation for HF CPU
-        max_new = min(int(max_tokens), int(self.hf_max_new_tokens))
-
-        # Many instruct/chat models work better with a simple instruction wrapper.
-        wrapped = (
-            "You are a helpful assistant.\n\n"
-            f"{prompt}\n\n"
-            "Answer:"
-        )
-
-        out = self._hf_pipe(
-            wrapped,
-            max_new_tokens=max_new,
-            do_sample=True,
-            temperature=float(max(0.05, temperature)),
-            top_p=0.9,
-            repetition_penalty=1.1,
-        )
-
-        if not out:
-            return ""
-
-        # pipeline returns list[{"generated_text": "..."}]
-        text = out[0].get("generated_text", "")
-        text = (text or "").strip()
-
-        # Remove the prompt prefix if the model echoed it
-        if text.startswith(wrapped):
-            text = text[len(wrapped):].strip()
-
-        return text
-
-    # --------------------------- JSON helpers ---------------------------
-
-    @staticmethod
-    def _strip_code_fences(text: str) -> str:
-        t = text.strip()
-        t = re.sub(r"^```(?:json)?\s*", "", t, flags=re.IGNORECASE)
-        t = re.sub(r"\s*```$", "", t)
-        return t.strip()
-
-    def safe_json_loads(self, text: str) -> Dict[str, Any]:
-        if not text:
-            return {}
-
-        t = self._strip_code_fences(text)
-
-        try:
-            out = json.loads(t)
-            return out if isinstance(out, dict) else {}
-        except Exception:
-            pass
-
-        m = re.search(r"\{.*\}", t, flags=re.DOTALL)
-        if m:
-            try:
-                out = json.loads(m.group(0))
-                return out if isinstance(out, dict) else {}
-            except Exception:
-                return {}
-
-        return {}
+#!/usr/bin/env python3
+"""
+Local LLM client abstraction (NO OpenAI/Claude).
+
+Providers:
+- ollama        : calls a local Ollama server (for Windows/local dev)
+- transformers  : runs a local HF model in-process (required for Hugging Face Spaces)
+
+Env:
+  LOCAL_LLM_PROVIDER=ollama|transformers
+
+Ollama:
+  OLLAMA_HOST=http://localhost:11434
+  OLLAMA_MODEL=llama3.2:1b
+
+Transformers (HF Spaces):
+  HF_LLM_MODEL=TinyLlama/TinyLlama-1.1B-Chat-v1.0
+  HF_MAX_NEW_TOKENS=450
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from typing import Any, Dict, Optional
+
+import requests
+
+
+class LocalLLMClient:
+    def __init__(
+        self,
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+        host: Optional[str] = None,
+        timeout_sec: int = 120,
+    ):
+        self.provider = (provider or os.getenv("LOCAL_LLM_PROVIDER", "ollama")).lower().strip()
+        self.timeout_sec = int(os.getenv("LLM_TIMEOUT_SEC", str(timeout_sec)))
+
+        # Ollama settings (local)
+        self.host = (host or os.getenv("OLLAMA_HOST", "http://localhost:11434")).strip()
+        self.model = (model or os.getenv("OLLAMA_MODEL", "llama3.2:1b")).strip()
+
+        # Transformers settings (HF Spaces)
+        self.hf_model_id = (os.getenv("HF_LLM_MODEL", "TinyLlama/TinyLlama-1.1B-Chat-v1.0")).strip()
+        self.hf_max_new_tokens = int(os.getenv("HF_MAX_NEW_TOKENS", "450"))
+
+        self._hf_pipe = None  # lazy init
+
+        if self.provider not in {"ollama", "transformers"}:
+            raise ValueError(
+                f"Unsupported LOCAL_LLM_PROVIDER='{self.provider}'. Supported: ollama, transformers."
+            )
+
+    def generate(self, prompt: str, temperature: float = 0.2, max_tokens: int = 900) -> str:
+        prompt = (prompt or "").strip()
+        if not prompt:
+            return ""
+
+        if self.provider == "ollama":
+            return self._generate_ollama(prompt, temperature=temperature, max_tokens=max_tokens)
+
+        return self._generate_transformers(prompt, temperature=temperature, max_tokens=max_tokens)
+
+    # --------------------------- Ollama ---------------------------
+
+    def _generate_ollama(self, prompt: str, temperature: float, max_tokens: int) -> str:
+        url = self.host.rstrip("/") + "/api/generate"
+        payload: Dict[str, Any] = {
+            "model": self.model,
+            "prompt": prompt,
+            "stream": False,
+            "options": {
+                "temperature": float(temperature),
+                "num_predict": int(max_tokens),
+            },
+        }
+
+        try:
+            r = requests.post(url, json=payload, timeout=self.timeout_sec)
+        except requests.RequestException as e:
+            raise RuntimeError(
+                "Failed to connect to local Ollama.\n"
+                f"Tried: {url}\n"
+                "Fix:\n"
+                "  - Ensure Ollama is running\n"
+                "  - Confirm endpoint: iwr http://localhost:11434/api/tags -UseBasicParsing\n"
+                f"Error: {repr(e)}"
+            ) from e
+
+        if r.status_code != 200:
+            body = (r.text or "").strip()
+            msg = body
+            try:
+                j = r.json()
+                if isinstance(j, dict):
+                    msg = j.get("error") or j.get("message") or body
+            except Exception:
+                pass
+            raise RuntimeError(
+                "Ollama returned an error.\n"
+                f"URL: {url}\n"
+                f"HTTP: {r.status_code}\n"
+                f"Model: {self.model}\n"
+                f"Details: {msg}"
+            )
+
+        data = r.json()
+        return (data.get("response") or "").strip()
+
+    # --------------------------- Transformers (HF Spaces) ---------------------------
+
+    def _lazy_init_hf(self):
+        if self._hf_pipe is not None:
+            return
+
+        from transformers import pipeline
+
+        self._hf_pipe = pipeline(
+            "text-generation",
+            model=self.hf_model_id,
+            device=-1,  # CPU
+        )
+
+    def _generate_transformers(self, prompt: str, temperature: float, max_tokens: int) -> str:
+        self._lazy_init_hf()
+
+        max_new = min(int(max_tokens), int(self.hf_max_new_tokens))
+
+        wrapped = (
+            "You are a helpful assistant.\n\n"
+            f"{prompt}\n\n"
+            "Answer:"
+        )
+
+        out = self._hf_pipe(
+            wrapped,
+            max_new_tokens=max_new,
+            do_sample=True,
+            temperature=float(max(0.05, temperature)),
+            top_p=0.9,
+            repetition_penalty=1.1,
+        )
+
+        if not out:
+            return ""
+
+        text = (out[0].get("generated_text", "") or "").strip()
+        if text.startswith(wrapped):
+            text = text[len(wrapped):].strip()
+        return text
+
+    # --------------------------- JSON helpers ---------------------------
+
+    @staticmethod
+    def _strip_code_fences(text: str) -> str:
+        t = text.strip()
+        t = re.sub(r"^```(?:json)?\s*", "", t, flags=re.IGNORECASE)
+        t = re.sub(r"\s*```$", "", t)
+        return t.strip()
+
+    def safe_json_loads(self, text: str) -> Dict[str, Any]:
+        if not text:
+            return {}
+
+        t = self._strip_code_fences(text)
+
+        try:
+            out = json.loads(t)
+            return out if isinstance(out, dict) else {}
+        except Exception:
+            pass
+
+        m = re.search(r"\{.*\}", t, flags=re.DOTALL)
+        if m:
+            try:
+                out = json.loads(m.group(0))
+                return out if isinstance(out, dict) else {}
+            except Exception:
+                return {}
+
+        return {}