Update app.py

app.py

@@ -229,11 +229,13 @@ from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from llama_index.core.llms import LLM
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.llms.huggingface import HuggingFaceInferenceAPI
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.templating import Jinja2Templates
+from huggingface_hub import InferenceClient
 import datetime
 from simple_salesforce import Salesforce, SalesforceLogin
 
@@ -241,6 +243,41 @@ from simple_salesforce import Salesforce, SalesforceLogin
 class MessageRequest(BaseModel):
     message: str
 
+# Custom LLM class for Hugging Face Inference API
+class HuggingFaceInferenceLLM(LLM):
+    def __init__(self, model_name: str, token: str):
+        super().__init__()
+        self.client = InferenceClient(model=model_name, token=token)
+        self.model_name = model_name
+
+    def chat(self, messages: list[ChatMessage], **kwargs) -> str:
+        prompt = ""
+        for msg in messages:
+            role = "user" if msg.role == MessageRole.USER else "assistant"
+            prompt += f"{role}: {msg.content}\n"
+        try:
+            response = self.client.text_generation(
+                prompt,
+                max_new_tokens=512,
+                temperature=0.1,
+                do_sample=True,
+                stop_sequences=["\n"]
+            )
+            return response
+        except Exception as e:
+            return f"Error in API call: {str(e)}"
+
+    async def achat(self, messages: list[ChatMessage], **kwargs) -> str:
+        return self.chat(messages, **kwargs)
+
+    @property
+    def metadata(self):
+        return {
+            "model_name": self.model_name,
+            "context_window": 3000,
+            "max_new_tokens": 512
+        }
+
 # Validate environment variables
 required_env_vars = ["HF_TOKEN", "username", "password", "security_token", "domain"]
 for var in required_env_vars:
@@ -288,11 +325,9 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="static")
 
 # LlamaIndex configuration
-Settings.llm = HuggingFaceInferenceAPI(
+Settings.llm = HuggingFaceInferenceLLM(
     model_name="meta-llama/Meta-Llama-3-8B-Instruct",
-    token=os.getenv("HF_TOKEN"),
-    max_new_tokens=512,
-    temperature=0.1
+    token=os.getenv("HF_TOKEN")
 )
 
 Settings.embed_model = HuggingFaceEmbedding(
@@ -404,7 +439,6 @@ async def receive_form_data(request: Request):
             'Phone': form_data.get('phone', '').strip(),
             'Email': form_data.get('email', ''),
         }
-        
         result = sf.Lead.create(data)
         return JSONResponse({"id": result['id']})
     except Exception as e: