Generator Refactor HF Inference Client API

app/services/generator.py

@@ -1,88 +1,77 @@
-"""LLM generation service using Hugging Face Inference API"""
-import requests
-from typing import Dict, Any, Optional
+"""LLM generation service using Hugging Face Inference Client SDK"""
+import os
+from typing import Optional
+from huggingface_hub import InferenceClient
 from app.config import settings
 from app.utils.logger import setup_logger
 
 logger = setup_logger(__name__)
 
+
 class GeneratorService:
-    """Handles text generation using Hugging Face models"""
-    
+    """Handles text generation using Hugging Face InferenceClient"""
+
     def __init__(self):
-        self.api_url = settings.HF_API_URL
-        self.headers = {"Authorization": f"Bearer {settings.HF_TOKEN}"}
-    
+        # Create a single reusable inference client
+        self.client = InferenceClient(api_key=settings.HF_TOKEN)
+
+        # Use model from settings or fallback
+        self.model = getattr(settings, "HF_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
+
     def generate(
         self,
         prompt: str,
         max_tokens: int = 512,
-        temperature: float = 0.7
+        temperature: float = 0.7,
     ) -> str:
-        """Generate text using the LLM"""
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": max_tokens,
-                "temperature": temperature,
-                "return_full_text": False
-            }
-        }
-        
+        """Generate text using HF chat-completion API"""
+
         try:
-            logger.info("Calling Hugging Face API...")
-            response = requests.post(
-                self.api_url,
-                headers=self.headers,
-                json=payload,
-                timeout=30
+            logger.info(f"Calling HF InferenceClient (model={self.model})...")
+
+            completion = self.client.chat.completions.create(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=max_tokens,
+                temperature=temperature,
             )
-            response.raise_for_status()
-            
-            result = response.json()
-            
-            # Handle different response formats
-            if isinstance(result, list) and len(result) > 0:
-                generated_text = result[0].get('generated_text', '')
-            elif isinstance(result, dict):
-                generated_text = result.get('generated_text', '')
-            else:
-                generated_text = str(result)
-            
+
+            generated_text = completion.choices[0].message.content
             logger.info("Generation successful")
+
             return generated_text.strip()
-        
-        except requests.exceptions.RequestException as e:
-            logger.error(f"API request failed: {str(e)}")
-            # Fallback to simple response
+
+        except Exception as e:
+            logger.error(f"HF Generation failed: {str(e)}")
             return self._fallback_response(prompt)
-    
+
     def _fallback_response(self, prompt: str) -> str:
-        """Fallback response when API fails"""
-        return "I apologize, but I'm unable to generate a response at the moment. Please try again later."
-    
-    def generate_rag_response(
-        self,
-        query: str,
-        context: str
-    ) -> str:
-        """Generate response using RAG pattern"""
+        """Fallback response when LLM API fails"""
+        return (
+            "I apologize, but I'm unable to generate a response at the moment. "
+            "Please try again later."
+        )
+
+    def generate_rag_response(self, query: str, context: str) -> str:
+        """Generate response using RAG-style prompt formatting"""
         prompt = self._build_rag_prompt(query, context)
         return self.generate(prompt)
-    
+
     def _build_rag_prompt(self, query: str, context: str) -> str:
-        """Build RAG prompt template"""
-        prompt = f"""<s>[INST] You are WorkWise, an AI assistant specialized in analyzing Jira project data. Answer the user's question based on the provided context.
+        """Build WorkWise-style RAG prompt"""
+        return f"""
+You are WorkWise, an AI assistant specialized in analyzing Jira project data.
+Answer the user's question based only on the context.
 
 Context:
 {context}
 
 User Question: {query}
 
-Provide a clear, concise answer based on the context. If the context doesn't contain enough information, say so. [/INST]</s>
+Provide a clear, concise answer. 
+If the context doesn't contain enough information, say so.
+""".strip()
 
-Answer:"""
-        return prompt
 
 # Global instance
-generator = GeneratorService()
+generator = GeneratorService()

app/services/generator.py.legacyJSON

@@ -0,0 +1,88 @@
+"""LLM generation service using Hugging Face Inference API"""
+import requests
+from typing import Dict, Any, Optional
+from app.config import settings
+from app.utils.logger import setup_logger
+
+logger = setup_logger(__name__)
+
+class GeneratorService:
+    """Handles text generation using Hugging Face models"""
+    
+    def __init__(self):
+        self.api_url = settings.HF_API_URL
+        self.headers = {"Authorization": f"Bearer {settings.HF_TOKEN}"}
+    
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int = 512,
+        temperature: float = 0.7
+    ) -> str:
+        """Generate text using the LLM"""
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": max_tokens,
+                "temperature": temperature,
+                "return_full_text": False
+            }
+        }
+        
+        try:
+            logger.info("Calling Hugging Face API...")
+            response = requests.post(
+                self.api_url,
+                headers=self.headers,
+                json=payload,
+                timeout=30
+            )
+            response.raise_for_status()
+            
+            result = response.json()
+            
+            # Handle different response formats
+            if isinstance(result, list) and len(result) > 0:
+                generated_text = result[0].get('generated_text', '')
+            elif isinstance(result, dict):
+                generated_text = result.get('generated_text', '')
+            else:
+                generated_text = str(result)
+            
+            logger.info("Generation successful")
+            return generated_text.strip()
+        
+        except requests.exceptions.RequestException as e:
+            logger.error(f"API request failed: {str(e)}")
+            # Fallback to simple response
+            return self._fallback_response(prompt)
+    
+    def _fallback_response(self, prompt: str) -> str:
+        """Fallback response when API fails"""
+        return "I apologize, but I'm unable to generate a response at the moment. Please try again later."
+    
+    def generate_rag_response(
+        self,
+        query: str,
+        context: str
+    ) -> str:
+        """Generate response using RAG pattern"""
+        prompt = self._build_rag_prompt(query, context)
+        return self.generate(prompt)
+    
+    def _build_rag_prompt(self, query: str, context: str) -> str:
+        """Build RAG prompt template"""
+        prompt = f"""<s>[INST] You are WorkWise, an AI assistant specialized in analyzing Jira project data. Answer the user's question based on the provided context.
+
+Context:
+{context}
+
+User Question: {query}
+
+Provide a clear, concise answer based on the context. If the context doesn't contain enough information, say so. [/INST]</s>
+
+Answer:"""
+        return prompt
+
+# Global instance
+generator = GeneratorService()

requirements.txt

@@ -1,8 +1,10 @@
+huggingface-hub>=0.26.0
 fastapi==0.109.0
 uvicorn[standard]==0.27.0
 python-dotenv==1.0.0
 python-multipart==0.0.6    # if you accept file uploads
 
+
 # === Data / utilities ===
 pandas==2.2.0
 numpy==1.26.3