Update main.py

main.py

@@ -13,6 +13,7 @@ from pydantic import BaseModel
 import os
 from dotenv import load_dotenv
 import warnings
+import uvicorn
 
 # Logging configuration
 logging.basicConfig(level=logging.DEBUG)
@@ -27,41 +28,56 @@ warnings.filterwarnings("ignore", category=DeprecationWarning)
 
 # Load environment variables
 load_dotenv()
+TOGETHER_AI_API = os.getenv("TOGETHER_AI")
+HF_HOME = os.getenv("HF_HOME", "./cache")
 
-# Set HF_HOME for Hugging Face cache directory
-HF_HOME = os.getenv("HF_HOME", "/tmp/cache")
+# Set cache directory for Hugging Face
 os.environ["HF_HOME"] = HF_HOME
 
-# Ensure the directory exists
-os.makedirs(HF_HOME, exist_ok=True)
+# Ensure HF_HOME exists and is writable
+if not os.path.exists(HF_HOME):
+    os.makedirs(HF_HOME, exist_ok=True)
 
-# Validate Together API environment variable
-TOGETHER_AI_API = os.getenv("TOGETHER_AI")
+# Validate environment variables
 if not TOGETHER_AI_API:
     raise ValueError("Environment variable TOGETHER_AI_API is missing. Please set it in your .env file.")
 
-# Initialize embeddings and vectorstore
-embeddings = HuggingFaceEmbeddings(
-    model_name="nomic-ai/nomic-embed-text-v1",
-    model_kwargs={"trust_remote_code": True, "revision": "289f532e14dbbbd5a04753fa58739e9ba766f3c7"},
-)
+# Initialize embeddings
+try:
+    embeddings = HuggingFaceEmbeddings(
+        model_name="nomic-ai/nomic-embed-text-v1",
+        model_kwargs={"trust_remote_code": True},
+    )
+except Exception as e:
+    logger.error(f"Error loading embeddings: {e}")
+    raise RuntimeError("Error initializing HuggingFaceEmbeddings.")
 
-# Ensure FAISS vectorstore is loaded properly
+# Ensure FAISS vectorstore is loaded or created
 try:
     db = FAISS.load_local("ipc_vector_db", embeddings, allow_dangerous_deserialization=True)
     db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 2, "max_length": 512})
 except Exception as e:
     logger.error(f"Error loading FAISS vectorstore: {e}")
-    raise RuntimeError("FAISS vectorstore could not be loaded. Ensure the vector database exists.")
+    # If not found, create a new vectorstore
+    try:
+        loader = DirectoryLoader('./data')
+        text_splitter = RecursiveCharacterTextSplitter()
+        documents = text_splitter.split_documents(loader.load())
+        db = FAISS.from_documents(documents, embeddings)
+        db.save_local("ipc_vector_db")
+        db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 2, "max_length": 512})
+    except Exception as inner_e:
+        logger.error(f"Error creating FAISS vectorstore: {inner_e}")
+        raise RuntimeError("FAISS vectorstore could not be created or loaded.")
 
 # Define the prompt template
 prompt_template = """<s>[INST]As a legal chatbot specializing in the Indian Penal Code, provide a concise and accurate answer based on the given context. Avoid unnecessary details or unrelated content. Only respond if the answer can be derived from the provided context; otherwise, say "The information is not available in the provided context." 
-CONTEXT: {context}
-CHAT HISTORY: {chat_history}
-QUESTION: {question}
-ANSWER:
-</s>[INST]
-"""
+    CONTEXT: {context}
+    CHAT HISTORY: {chat_history}
+    QUESTION: {question}
+    ANSWER:
+    </s>[INST]
+    """
 prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question", "chat_history"])
 
 # Initialize the Together API
@@ -111,3 +127,9 @@ async def chat(request: ChatRequest):
     except Exception as e:
         logger.error(f"Error during chat invocation: {e}")
         raise HTTPException(status_code=500, detail="Internal server error")
+
+# Start Uvicorn server if run directly
+if __name__ == "__main__":
+    ENV = os.getenv("ENV", "prod")
+    PORT = int(os.environ.get("PORT", 10000))  # Use the default port 10000 or the environment port
+    uvicorn.run("main:app", host="0.0.0.0", port=PORT, reload=(ENV == "dev"))