Update app.py

app.py

@@ -4,6 +4,30 @@ from langchain_community.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import RetrievalQA
 from langchain_community.llms import HuggingFaceHub
+try:
+    from langchain_huggingface import HuggingFaceEndpoint
+    HUGGINGFACE_ENDPOINT_AVAILABLE = True
+except ImportError:
+    HUGGINGFACE_ENDPOINT_AVAILABLE = False
+    print("langchain-huggingface not available, using fallback")
+
+from langchain.schema import BaseLanguageModel
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from typing import Any, List, Optional
+
+# Simple mock LLM for testing
+class MockLLM(BaseLanguageModel):
+    def _llm_type(self) -> str:
+        return "mock"
+    
+    def _call(
+        self, 
+        prompt: str, 
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        return f"Based on the document context, here's a response to your query: {prompt[:100]}... [This is a mock response for testing. Please configure a proper LLM.]"
 from langchain_community.embeddings import HuggingFaceEmbeddings
 
 # You can use this section to suppress warnings generated by your code:
@@ -24,15 +48,60 @@ def get_llm():
     Initializes and returns a Hugging Face Hub LLM model.
     Using a conversational model suitable for legal advice.
     """
-    repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-    llm = HuggingFaceHub(
-        repo_id=repo_id,
-        # The 'task' parameter is required to specify the model's function.
-        # This fixes the 'Value error, Got invalid task None' error.
-        task="text-generation",
-        model_kwargs={"temperature": 0.1, "max_length": 512}
-    )
-    return llm
+    # Check if API token is properly set
+    api_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+    if not api_token or api_token == "hf_YOUR_HUGGINGFACE_TOKEN":
+        print("Warning: No valid API token found. Using mock LLM for testing.")
+        return MockLLM()
+    
+    if HUGGINGFACE_ENDPOINT_AVAILABLE:
+        try:
+            # Try the newer HuggingFaceEndpoint first
+            llm = HuggingFaceEndpoint(
+                repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
+                max_length=512,
+                temperature=0.1,
+                huggingfacehub_api_token=api_token
+            )
+            print("Successfully initialized HuggingFaceEndpoint")
+            return llm
+        except Exception as e:
+            print(f"HuggingFaceEndpoint failed: {e}")
+    
+    try:
+        # Fallback to traditional HuggingFaceHub
+        llm = HuggingFaceHub(
+            repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
+            task="text-generation",
+            model_kwargs={
+                "temperature": 0.1, 
+                "max_length": 512,
+                "do_sample": True,
+                "top_p": 0.9
+            },
+            huggingfacehub_api_token=api_token
+        )
+        print("Successfully initialized HuggingFaceHub")
+        return llm
+    except Exception as e:
+        print(f"HuggingFaceHub with Mixtral failed: {e}")
+        
+        # Try with a smaller, more reliable model
+        try:
+            llm = HuggingFaceHub(
+                repo_id="google/flan-t5-base",
+                task="text2text-generation",
+                model_kwargs={
+                    "temperature": 0.1,
+                    "max_length": 512
+                },
+                huggingfacehub_api_token=api_token
+            )
+            print("Successfully initialized HuggingFaceHub with FLAN-T5")
+            return llm
+        except Exception as e2:
+            print(f"All HuggingFace models failed. Using mock LLM. Errors: {e}, {e2}")
+            return MockLLM()
 
 ## Document loader
 def document_loader(file_path):