Updated LLMs Client

src/agentic_multiwriter/models/llm_client.py

@@ -5,8 +5,8 @@ import os
 from dataclasses import dataclass
 from typing import Literal
 
+from huggingface_hub import InferenceClient
 from langchain_community.chat_models import ChatOllama
-from langchain_community.llms import HuggingFaceEndpoint
 from langchain_openai import ChatOpenAI
 
 from agentic_multiwriter.tools import get_logger
@@ -27,7 +27,7 @@ class LLMClient:
     Small wrapper around different LLM backends:
 
     - provider="ollama"      -> local ChatOllama
-    - provider="hf_endpoint" -> Hugging Face Inference endpoint
+    - provider="hf_endpoint" -> Hugging Face Inference API (InferenceClient)
     - provider="openai"      -> OpenAI Chat model
     """
 
@@ -45,7 +45,6 @@ class LLMClient:
             )
 
         elif provider == "hf_endpoint":
-            # Uses HuggingFaceHosted Inference API (text-generation)
             token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
             if not token:
                 raise RuntimeError(
@@ -53,13 +52,10 @@ class LLMClient:
                 )
 
             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=model,
-                huggingfacehub_api_token=token,
-                temperature=temperature,
-                max_new_tokens=800,
+            # Direct low-level client; no LangChain wrapper here.
+            self._client = InferenceClient(
+                model=model,
+                token=token,
             )
 
         elif provider == "openai":
@@ -100,18 +96,24 @@ class LLMClient:
                 ("user", user_prompt),
             ]
             response = self._client.invoke(messages)
-            # Both ChatOllama and ChatOpenAI return an object with `.content`
             return getattr(response, "content", str(response))
 
         elif self._mode == "hf_endpoint":
-            # HuggingFaceEndpoint expects a single text prompt
+            # Use HF InferenceClient directly (text generation)
             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
+
+            text = self._client.text_generation(
+                prompt,
+                max_new_tokens=800,
+                temperature=self.settings.temperature,
+                top_p=0.95,
+                do_sample=True,
+            )
+            # InferenceClient.text_generation returns a plain string
             return text.strip() if isinstance(text, str) else str(text)
 
         else: