Update main.py

main.py

@@ -1,19 +1,19 @@
 import os
+os.environ["HF_HOME"] = "/tmp/huggingface"  # Prevent permission error in HF Spaces
+
 import fitz  # PyMuPDF
 import uuid
 from fastapi import FastAPI, UploadFile, File, Form, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import HTMLResponse, JSONResponse
-from pydantic import BaseModel
-from typing import List
 from dotenv import load_dotenv
+from typing import List
 
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_core.documents import Document
-
 from anthropic import Anthropic
 
 # ---- Load API Keys ----
@@ -25,18 +25,16 @@ CLAUDE_MODEL = "claude-3-haiku-20240307"
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-# Create static directory if it doesn't exist
+# Mount static directory (if needed for frontend)
 os.makedirs(os.path.join(os.path.dirname(__file__), "static"), exist_ok=True)
-
-# Mount static files directory
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # ---- In-Memory Stores ----
-db_store = {}
-chat_store = {}
-general_chat_sessions = {}
+db_store = {}                # session_id → Chroma vector DB
+chat_store = {}              # session_id → chat messages
+general_chat_sessions = {}   # session_id → general (no PDF) flag
 
-# ---- Utils ----
+# ---- Utility Functions ----
 
 def extract_text_from_pdf(file) -> str:
     """Extracts text from the first page of a PDF."""
@@ -47,9 +45,7 @@ def build_vector_db(text: str, collection_name: str) -> Chroma:
     """Chunks, embeds, and stores text in ChromaDB."""
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     docs = splitter.create_documents([text])
-
-    # Using a standard model that should be available publicly
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     vectordb = Chroma.from_documents(docs, embeddings, collection_name=collection_name)
     return vectordb
 
@@ -62,11 +58,8 @@ def create_session(is_pdf: bool = True) -> str:
     """Creates a new unique session ID."""
     sid = str(uuid.uuid4())
     chat_store[sid] = []
-    
-    # Track if this is a general chat session (without PDF)
     if not is_pdf:
         general_chat_sessions[sid] = True
-    
     return sid
 
 def append_chat(session_id: str, role: str, msg: str):
@@ -80,62 +73,55 @@ def delete_session(session_id: str):
     db_store.pop(session_id, None)
     general_chat_sessions.pop(session_id, None)
 
-# ---- API Routes ----
+# ---- API Endpoints ----
 
 @app.get("/", response_class=HTMLResponse)
 async def get_home():
-    with open(os.path.join(os.path.dirname(__file__), "static", "index.html")) as f:
-        return f.read()
+    try:
+        with open(os.path.join(os.path.dirname(__file__), "static", "index.html")) as f:
+            return f.read()
+    except FileNotFoundError:
+        return HTMLResponse(content="<h1>RAG Chatbot API</h1><p>Upload a PDF or start a chat.</p>")
 
 @app.post("/start-chat/")
 async def start_general_chat():
     """Starts a general chat session without PDF."""
     session_id = create_session(is_pdf=False)
     return {"session_id": session_id, "message": "General chat session started."}
-        
+
 @app.post("/upload/")
 async def upload_pdf(file: UploadFile = File(...), current_session_id: str = Form(None)):
     """Handles PDF upload and indexing with chat continuity."""
-    # Extract text from PDF
     text = extract_text_from_pdf(file)
-    
-    # Handle session continuity
+
     if current_session_id and current_session_id in chat_store:
-        # Continue with existing session
         session_id = current_session_id
-        # Remove from general chat sessions if it was one
-        if session_id in general_chat_sessions:
-            general_chat_sessions.pop(session_id)
+        general_chat_sessions.pop(session_id, None)  # upgrade to PDF mode
     else:
-        # Create a new session
         session_id = create_session()
-    
-    # Create and store the vector database
+
     vectordb = build_vector_db(text, collection_name=session_id)
     db_store[session_id] = vectordb
-    
+
     return {"session_id": session_id, "message": "PDF indexed."}
 
 @app.post("/chat/")
 async def chat(session_id: str = Form(...), prompt: str = Form(...)):
-    """Handles user chat prompt, fetches relevant info, calls Claude."""
-    # Check if this is a general chat or PDF chat
     is_general_chat = session_id in general_chat_sessions
     is_pdf_chat = session_id in db_store
-    
+
     if not is_general_chat and not is_pdf_chat:
         return {"error": "Invalid session ID"}
-    
+
     append_chat(session_id, "user", prompt)
-    
-    # Ensure we have an API key and initialize with proper parameters
+
     if not ANTHROPIC_API_KEY:
         return JSONResponse(status_code=500, content={"error": "Missing ANTHROPIC_API_KEY environment variable"})
-    
+
     client = Anthropic(api_key=ANTHROPIC_API_KEY.strip())
-    
+
     if is_general_chat:
-        # General chat without PDF context
+        # No context, just send prompt
         response = client.messages.create(
             model=CLAUDE_MODEL,
             max_tokens=512,
@@ -143,7 +129,6 @@ async def chat(session_id: str = Form(...), prompt: str = Form(...)):
             messages=[{"role": "user", "content": prompt}]
         )
     else:
-        # PDF-based chat with context
         context = retrieve_context(db_store[session_id], prompt)
         response = client.messages.create(
             model=CLAUDE_MODEL,
@@ -162,3 +147,4 @@ async def end_chat(session_id: str = Form(...)):
     """Ends session and deletes associated data."""
     delete_session(session_id)
     return {"message": "Session cleared."}
+