Update main.py

main.py

@@ -8,6 +8,7 @@ from requests import JSONDecodeError
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_community.llms import HuggingFaceEndpoint
+from langchain_openai import ChatOpenAI
 
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
@@ -60,17 +61,27 @@ def response_generator(query: str) -> str:
     logger.info("Using HF model %s", model)
 
     # prepare HF text-generation LLM
-    hf = HuggingFaceEndpoint(
-        # endpoint_url=f"https://api-inference.huggingface.co/models/{model}",
-        endpoint_url=f"https://router.huggingface.co/hf-inference/models/{model}",
-        task="text-generation",
-        huggingfacehub_api_token=hf_api_key,
-        model_kwargs={
-            "temperature": temperature,
-            "max_new_tokens": max_tokens,
-            "return_full_text": False,
-        },
-    )
+    # hf = HuggingFaceEndpoint(
+    #     # endpoint_url=f"https://api-inference.huggingface.co/models/{model}",
+    #     endpoint_url=f"https://router.huggingface.co/hf-inference/models/{model}",
+    #     task="text-generation",
+    #     huggingfacehub_api_token=hf_api_key,
+    #     model_kwargs={
+    #         "temperature": temperature,
+    #         "max_new_tokens": max_tokens,
+    #         "return_full_text": False,
+    #     },
+    # )
+
+    hf = ChatOpenAI(
+            base_url=f"https://router.huggingface.co/hf-inference/models/{model}/v1",
+            api_key=hf_api_key,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            timeout=30,  # Add timeout
+            max_retries=3,  # Built-in retry logic
+        )
 
     # conversational RAG chain
     qa = ConversationalRetrievalChain.from_llm(