Updateds

src/agentic_multiwriter/models/llm_client.py

@@ -1,156 +1,145 @@
-# src/agentic_multiwriter/models/llm_client.py
-
-from __future__ import annotations
-
-import logging
 import os
+import logging
 from dataclasses import dataclass
+from typing import Optional
 
 from huggingface_hub import InferenceClient
 from langchain_ollama import ChatOllama
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import SystemMessage, HumanMessage, BaseMessage
+from langchain_core.messages import SystemMessage, HumanMessage
 
 logger = logging.getLogger(__name__)
 
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Settings
-# ──────────────────────────────────────────────────────────────────────────────
-
-
 @dataclass
 class LLMSettings:
-    """Configuration for the LLM client."""
-
-    llm_provider: str = os.getenv("AMW_LLM_PROVIDER", "ollama").strip() or "ollama"
-    llm_model: str = os.getenv("AMW_LLM_MODEL", "llama3").strip() or "llama3"
+    provider: str = os.getenv("AMW_LLM_PROVIDER", "ollama")  # 'ollama', 'openai', 'hf_endpoint'
+    llm_model: str = os.getenv("AMW_LLM_MODEL", "llama3")
     temperature: float = float(os.getenv("AMW_TEMPERATURE", "0.4"))
-    max_tokens: int = int(os.getenv("AMW_MAX_TOKENS", "1024"))
+    max_tokens: int = int(os.getenv("AMW_MAX_TOKENS", "768"))
 
+    # HF token is optional; if not set, HF will still work for some public models
+    hf_api_token: Optional[str] = os.getenv("HUGGINGFACEHUB_API_TOKEN", None)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Client wrapper
-# ──────────────────────────────────────────────────────────────────────────────
+    # OpenAI key is optional unless provider='openai'
+    openai_api_key: Optional[str] = os.getenv("OPENAI_API_KEY", None)
 
 
 class LLMClient:
     """
-    Thin wrapper that normalizes access to different LLM providers:
+    Thin wrapper over different backends:
 
-    - openai        → ChatOpenAI (OpenAI API)
-    - ollama        → ChatOllama (local server)
-    - hf_endpoint   → huggingface_hub.InferenceClient (Hugging Face endpoints / Spaces)
+    - provider='ollama'      -> local Ollama (ChatOllama)
+    - provider='openai'      -> OpenAI ChatCompletion models
+    - provider='hf_endpoint' -> Hugging Face Inference API (text_generation)
     """
 
-    def __init__(self, settings: LLMSettings | None = None) -> None:
+    def __init__(self, settings: Optional[LLMSettings] = None) -> None:
         self.settings = settings or LLMSettings()
-        self._client = self._build_client()
-
-    # --------------------------------------------------------------------- #
-    # Client factory
-    # --------------------------------------------------------------------- #
-
-    def _build_client(self):
-        provider = self.settings.llm_provider.lower()
-        model = self.settings.llm_model
+        self.provider = self.settings.provider.lower()
+        self.model = self.settings.llm_model
+        self.temperature = self.settings.temperature
 
         logger.info(
             "LLMClient initialized with provider='%s', model='%s', temperature=%.2f",
-            provider,
-            model,
-            self.settings.temperature,
+            self.provider,
+            self.model,
+            self.temperature,
         )
 
-        if provider == "openai":
-            # Uses OPENAI_API_KEY from env
-            return ChatOpenAI(
-                model=model,
-                temperature=self.settings.temperature,
-                max_tokens=self.settings.max_tokens,
+        if self.provider == "ollama":
+            self._client = ChatOllama(
+                model=self.model,
+                temperature=self.temperature,
             )
 
-        if provider == "ollama":
-            # Expects a local Ollama server (not available on HF Spaces)
-            return ChatOllama(
-                model=model,
-                temperature=self.settings.temperature,
+        elif self.provider == "openai":
+            if not self.settings.openai_api_key:
+                logger.warning(
+                    "OPENAI_API_KEY not set but provider='openai'. "
+                    "Requests will fail until the key is configured."
+                )
+            self._client = ChatOpenAI(
+                model=self.model,
+                temperature=self.temperature,
+                api_key=self.settings.openai_api_key,
             )
 
-        if provider == "hf_endpoint":
-            # Hugging Face Inference API / Spaces
-            # HUGGINGFACEHUB_API_TOKEN is recommended, but on Spaces the
-            # internal token usually works without specifying it.
-            hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-            if hf_token:
-                client = InferenceClient(token=hf_token)
+        elif self.provider in {"hf_endpoint", "huggingface", "hf"}:
+            # Bind the client directly to the model so we use the model's
+            # Inference API endpoint (not the generic router).
+            if self.settings.hf_api_token:
                 logger.info("Using explicit HUGGINGFACEHUB_API_TOKEN for hf_endpoint.")
             else:
-                client = InferenceClient()
-                logger.info(
-                    "No HUGGINGFACEHUB_API_TOKEN set; relying on default HF Space token."
+                logger.warning(
+                    "HUGGINGFACEHUB_API_TOKEN not set. For reliable HF Inference, "
+                    "set it as a secret in your Space or local environment."
                 )
-            return client
-
-        raise ValueError(f"Unsupported LLM provider: {provider}")
 
-    # --------------------------------------------------------------------- #
-    # Public API
-    # --------------------------------------------------------------------- #
+            self._client = InferenceClient(
+                model=self.model,
+                token=self.settings.hf_api_token,
+            )
+        else:
+            raise ValueError(f"Unknown LLM provider: {self.provider}")
 
+    # ---------------------------------------------------------------------
+    # Unified generate() API
+    # ---------------------------------------------------------------------
     def generate(self, system_prompt: str, user_prompt: str) -> str:
         """
-        Generate text using the configured provider, given a system and user prompt.
-        Returns only the text content (stripped).
+        Generates a single string response from the configured backend.
         """
-        provider = self.settings.llm_provider.lower()
-
-        if provider == "openai":
-            messages: list[BaseMessage] = [
-                SystemMessage(content=system_prompt),
-                HumanMessage(content=user_prompt),
-            ]
-            response = self._client.invoke(messages)
-            text = response.content
 
-        elif provider == "ollama":
-            messages: list[BaseMessage] = [
+        if self.provider == "ollama":
+            messages = [
                 SystemMessage(content=system_prompt),
                 HumanMessage(content=user_prompt),
             ]
             response = self._client.invoke(messages)
-            text = response.content
+            return response.content  # type: ignore[return-value]
 
-        elif provider == "hf_endpoint":
-            # IMPORTANT: Zephyr and many HF models in Spaces are exposed as
-            # conversational / chat models. We therefore use `chat_completion`
-            # instead of `text_generation`, which fixes the "task text-generation
-            # not supported, use conversational" error you saw.
+        if self.provider == "openai":
             messages = [
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": user_prompt},
             ]
-
-            response = self._client.chat_completion(
-                model=self.settings.llm_model,
-                messages=messages,
-                max_tokens=self.settings.max_tokens,
-                temperature=self.settings.temperature,
-                top_p=0.95,
+            response = self._client.invoke(messages)
+            # langchain-openai returns AIMessage
+            return response.content  # type: ignore[return-value]
+
+        if self.provider in {"hf_endpoint", "huggingface", "hf"}:
+            # For HF Inference we use plain text-generation.
+            # We concatenate system + user into a single prompt.
+            prompt = (
+                system_prompt.strip()
+                + "\n\nUser:\n"
+                + user_prompt.strip()
+                + "\n\nAssistant:"
             )
 
-            # huggingface_hub ChatCompletionOutput is a small dataclass-like
-            # object. Its first choice contains the assistant message.
-            choice = response.choices[0]
-            message = choice.message
-
-            # Handle both dict-style and attribute-style message objects
-            if isinstance(message, dict):
-                text = message.get("content", "")
-            else:
-                text = getattr(message, "content", "")
-
-        else:
-            raise ValueError(f"Unsupported LLM provider: {provider}")
-
-        return (text or "").strip()
+            try:
+                text = self._client.text_generation(
+                    prompt,
+                    max_new_tokens=self.settings.max_tokens,
+                    temperature=self.temperature,
+                    do_sample=self.temperature > 0,
+                    repetition_penalty=1.05,
+                    return_full_text=False,  # only new tokens
+                )
+                # text_generation returns a plain string when return_full_text=False
+                return text.strip()
+            except Exception as e:
+                logger.exception(
+                    "Error while calling Hugging Face Inference API for model '%s': %s",
+                    self.model,
+                    e,
+                )
+                raise RuntimeError(
+                    f"Hugging Face Inference error for model '{self.model}'. "
+                    f"Check that the model supports text-generation and that "
+                    f"your token has Inference permissions."
+                ) from e
+
+        # Should never reach here
+        raise RuntimeError(f"Unsupported provider: {self.provider}")