Update LLM Client

src/agentic_multiwriter/models/llm_client.py

@@ -1,77 +1,118 @@
+# src/agentic_multiwriter/models/llm_client.py
 from __future__ import annotations
 
-from typing import Optional
+import os
+from dataclasses import dataclass
+from typing import Literal
 
-from langchain_ollama import ChatOllama
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_community.chat_models import ChatOllama
 from langchain_community.llms import HuggingFaceEndpoint
+from langchain_openai import ChatOpenAI
 
-from ..config import settings
-from ..tools import get_logger
+from agentic_multiwriter.tools import get_logger
 
 logger = get_logger()
 
 
+@dataclass
+class LLMSettings:
+    """Configuration for the LLM backend."""
+    llm_provider: str = os.getenv("AMW_LLM_PROVIDER", "ollama").lower()
+    llm_model: str = os.getenv("AMW_LLM_MODEL", "llama3")
+    temperature: float = float(os.getenv("AMW_TEMPERATURE", "0.4"))
+
+
 class LLMClient:
     """
-    Wrapper to abstract over multiple LLM providers:
-    - ollama: local Ollama server
-    - openai: OpenAI Chat API
-    - hf_endpoint: Hugging Face Inference Endpoint / hosted model
+    Small wrapper around different LLM backends:
+
+    - provider="ollama"      -> local ChatOllama
+    - provider="hf_endpoint" -> Hugging Face Inference endpoint
+    - provider="openai"      -> OpenAI Chat model
     """
 
-    def __init__(
-        self,
-        provider: Optional[str] = None,
-        model: Optional[str] = None,
-        temperature: Optional[float] = None,
-    ) -> None:
-        self.provider = (provider or settings.llm_provider).lower()
-        self.model = model or settings.llm_model
-        self.temperature = temperature if temperature is not None else settings.temperature
-
-        settings.validate()
-        self._init_client()
-        logger.info(
-            "LLMClient initialized with provider='%s', model='%s', temperature=%.2f",
-            self.provider,
-            self.model,
-            self.temperature,
-        )
+    def __init__(self, settings: LLMSettings | None = None) -> None:
+        self.settings = settings or LLMSettings()
+        provider = self.settings.llm_provider
+        model = self.settings.llm_model
+        temperature = self.settings.temperature
 
-    def _init_client(self) -> None:
-        if self.provider == "ollama":
+        if provider == "ollama":
+            self._mode: Literal["ollama", "hf_endpoint", "openai"] = "ollama"
             self._client = ChatOllama(
-                model=self.model, temperature=self.temperature)
-        elif self.provider == "openai":
-            self._client = ChatOpenAI(
-                model=self.model, temperature=self.temperature)
-        elif self.provider == "hf_endpoint":
-            # Uses Hugging Face hosted model via Inference API
-            # Model id should be a HF repo id, e.g. "HuggingFaceH4/zephyr-7b-beta"
+                model=model,
+                temperature=temperature,
+            )
+
+        elif provider == "hf_endpoint":
+            # Uses HuggingFaceHosted Inference API (text-generation)
+            token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+            if not token:
+                raise RuntimeError(
+                    "HUGGINGFACEHUB_API_TOKEN is required when AMW_LLM_PROVIDER=hf_endpoint"
+                )
+
+            self._mode = "hf_endpoint"
+            # IMPORTANT: do NOT pass our own client object here; let
+            # HuggingFaceEndpoint build the correct internal client.
             self._client = HuggingFaceEndpoint(
-                repo_id=self.model,
-                temperature=self.temperature,
+                repo_id=model,
+                huggingfacehub_api_token=token,
+                temperature=temperature,
+                max_new_tokens=800,
             )
+
+        elif provider == "openai":
+            api_key = os.getenv("OPENAI_API_KEY")
+            if not api_key:
+                raise RuntimeError(
+                    "OPENAI_API_KEY is required when AMW_LLM_PROVIDER=openai"
+                )
+
+            self._mode = "openai"
+            self._client = ChatOpenAI(
+                model=model,
+                temperature=temperature,
+                api_key=api_key,
+            )
+
         else:
-            raise ValueError(f"Unsupported provider: {self.provider}")
+            raise ValueError(f"Unknown AMW_LLM_PROVIDER='{provider}'")
+
+        logger.info(
+            "LLMClient initialized with provider='%s', model='%s', temperature=%.2f",
+            provider,
+            model,
+            temperature,
+        )
 
-    def generate(self, system_prompt: str, user_prompt: str) -> str:
-        if self.provider in {"ollama", "openai"}:
+    def generate(self, *, system_prompt: str, user_prompt: str) -> str:
+        """
+        Unified generate() interface for all providers.
+        """
+        system_prompt = system_prompt.strip()
+        user_prompt = user_prompt.strip()
+
+        if self._mode in ("ollama", "openai"):
+            # Chat-style models (ChatOllama / ChatOpenAI)
             messages = [
-                SystemMessage(content=system_prompt),
-                HumanMessage(content=user_prompt),
+                ("system", system_prompt),
+                ("user", user_prompt),
             ]
             response = self._client.invoke(messages)
-            return response.content or ""
+            # Both ChatOllama and ChatOpenAI return an object with `.content`
+            return getattr(response, "content", str(response))
 
-        if self.provider == "hf_endpoint":
-            # HuggingFaceEndpoint is a text-completion-style model,
-            # so we concatenate system + user into a single prompt.
-            prompt = f"{system_prompt.strip()}\n\n{user_prompt.strip()}"
-            response = self._client.invoke(prompt)
-            return response or ""
+        elif self._mode == "hf_endpoint":
+            # HuggingFaceEndpoint expects a single text prompt
+            prompt = (
+                f"{system_prompt}\n\n"
+                f"User:\n{user_prompt}\n\n"
+                f"Assistant:"
+            )
+            text = self._client.invoke(prompt)
+            # HuggingFaceEndpoint typically returns raw text
+            return text.strip() if isinstance(text, str) else str(text)
 
-        raise ValueError(
-            f"Unsupported provider at generate(): {self.provider}")
+        else:
+            raise RuntimeError("Unsupported LLM provider mode")