Create app.py

app.py

@@ -0,0 +1,204 @@
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import os
+# No requests import needed for Ollama connection check if not using Ollama
+
+# Import Hugging Face Transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch # For checking GPU availability
+
+from langchain_community.embeddings import HuggingFaceEmbeddings # Using HF Embeddings now
+from langchain_community.vectorstores import Chroma
+from langchain_core.documents import Document
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+app = Flask(__name__)
+CORS(app)
+
+# --- Model Configuration for Hugging Face Transformers ---
+# CHOOSE A SMALLER MODEL! Gemma 4B is too large for free tier usually.
+# 'google/gemma-2b-it' is a good conversational starting point.
+LLM_MODEL_NAME_HF = "google/gemma-2b-it" 
+EMBEDDING_MODEL_NAME_HF = "sentence-transformers/all-MiniLM-L6-v2" # Standard small embedding model
+
+# Global variables for models
+llm_pipeline = None # Will be a Hugging Face pipeline
+embeddings = None   # Will be a HuggingFaceEmbeddings instance
+
+# --- User-specific Vector Stores Cache ---
+user_vectorstores = {} 
+
+def initialize_models():
+    """
+    Initialize Hugging Face models (LLM pipeline and Embeddings).
+    """
+    global llm_pipeline, embeddings
+    print("Initializing Hugging Face models...")
+    try:
+        # Determine device for LLM: Use GPU if available, otherwise CPU
+        device = 0 if torch.cuda.is_available() else -1
+        print(f"Using device: {'cuda' if device == 0 else 'cpu'}")
+
+        # Initialize LLM Pipeline
+        # This will download the model weights (gemma-2b-it is ~5GB)
+        # It's recommended to do this once at startup.
+        print(f"Loading LLM: {LLM_MODEL_NAME_HF}...")
+        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME_HF)
+        model = AutoModelForCausalLM.from_pretrained(LLM_MODEL_NAME_HF, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32)
+        llm_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=500, # Limit response length
+            device=device,
+            # Add other generation parameters as needed, e.g., do_sample=True, top_p=0.9, temperature=0.7
+        )
+        print("LLM Pipeline initialized successfully!")
+
+        # Initialize Hugging Face Embeddings
+        print(f"Loading Embedding Model: {EMBEDDING_MODEL_NAME_HF}...")
+        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME_HF)
+        print("Embedding Model initialized successfully!")
+
+    except Exception as e:
+        print(f"ERROR: An unexpected error occurred during model initialization: {e}")
+        llm_pipeline = None
+        embeddings = None
+        # Raise the exception to prevent the app from starting if models fail to load
+        raise e 
+
+# --- Helper function to adapt HF pipeline to LangChain's LLM interface ---
+# LangChain's pipeline.py can convert HF pipelines but requires some setup.
+# For simplicity, we'll manually wrap it in the RAG chain
+# We will use it directly in the RAG chain's invoke step.
+
+@app.route('/load_document', methods=['POST'])
+def load_document():
+    # ... (rest of your /load_document function remains largely the same) ...
+    # Ensure 'embeddings' is properly loaded before this.
+    if not embeddings:
+        return jsonify({"error": "Embedding model not initialized. Server might be restarting or failed to load models."}), 500
+
+    data = request.get_json()
+    user_id = data.get("user_id") 
+    text = data.get("text")
+    
+    if not user_id: return jsonify({"error": "User ID (user_id) is required to load a document."}), 400
+    if not text: return jsonify({"error": "No text provided to load."}), 400
+
+    print(f"Loading document for user: {user_id}")
+
+    try:
+        # Create a unique persistence directory for each user's ChromaDB
+        # NOTE: On Hugging Face Spaces, this persist_dir will be within the Space's storage,
+        # which can be ephemeral or reset, depending on space type/resource usage.
+        # For a true persistent solution, you'd need external storage.
+        persist_dir = f"./chroma_db_users/{user_id}/"
+        os.makedirs(persist_dir, exist_ok=True) 
+
+        base_document = Document(page_content=text, metadata={"user_id": user_id, "source": "user_upload"})
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        chunks = text_splitter.split_documents([base_document])
+
+        user_vectorstores[user_id] = Chroma.from_documents(
+            chunks, embedding=embeddings, persist_directory=persist_dir
+        )
+        
+        print(f"Document loaded for user '{user_id}'. Chunks created: {len(chunks)} at {persist_dir}")
+        return jsonify({"message": f"Document loaded successfully for user '{user_id}'.", "chunks_created": len(chunks)})
+    except Exception as e:
+        print(f"Error loading document for user '{user_id}': {e}")
+        return jsonify({"error": f"Error loading document: {e}"}), 500
+
+
+@app.route('/query', methods=['POST'])
+def query():
+    """
+    Query the currently loaded document for a specific user to summarize or answer a question.
+    """
+    if not llm_pipeline or not embeddings:
+        return jsonify({"error": "Models not initialized. Server might be restarting or failed to load models."}), 500
+
+    data = request.get_json()
+    user_id = data.get("user_id") 
+    query_text = data.get("query") 
+
+    if not user_id: return jsonify({"error": "User ID (user_id) is required to query."}), 400
+    if not query_text: return jsonify({"error": "No query text provided."}), 400
+
+    print(f"Query received for user: {user_id}, Query: '{query_text}'")
+
+    current_user_vectorstore = user_vectorstores.get(user_id)
+    if not current_user_vectorstore:
+        user_persist_dir = f"./chroma_db_users/{user_id}/"
+        if os.path.exists(user_persist_dir):
+            try:
+                current_user_vectorstore = Chroma(persist_directory=user_persist_dir, embedding_function=embeddings)
+                user_vectorstores[user_id] = current_user_vectorstore 
+                print(f"Loaded existing vectorstore for user '{user_id}' from disk.")
+            except Exception as e:
+                print(f"Error loading vectorstore from disk for user '{user_id}': {e}")
+                return jsonify({"error": f"Failed to load document for user '{user_id}'. Please try loading it again or check server logs."}), 500
+        else:
+            return jsonify({"error": f"No document loaded for user '{user_id}'. Please load a document first using /load_document."}), 400
+
+    try:
+        retriever = current_user_vectorstore.as_retriever()
+        
+        prompt_template = ChatPromptTemplate.from_template(
+            """Answer the question based ONLY on the following context. If the answer is not available in the provided context, politely state that you cannot find the answer in the provided information.
+
+Context: {context}
+
+Question: {question}
+"""
+        )
+        
+        # --- RAG Chain for Hugging Face Pipeline ---
+        # Get relevant context documents
+        retrieved_docs = retriever.invoke(query_text)
+        context_text = "\n\n".join([doc.page_content for doc in retrieved_docs])
+
+        # Format the prompt using the template and retrieved context
+        formatted_prompt = prompt_template.format(context=context_text, question=query_text)
+
+        # Use the Hugging Face pipeline directly for text generation
+        # Pass the formatted prompt to the pipeline
+        outputs = llm_pipeline(formatted_prompt)
+        
+        # The output from the pipeline needs to be parsed based on its structure
+        # It's usually a list of dictionaries, with 'generated_text' key.
+        # You might need to refine this parsing based on the exact model's output format.
+        generated_text = outputs[0]['generated_text']
+
+        # The model might repeat the prompt or parts of it, extract only the new response.
+        # This is a common challenge with text generation.
+        # A simple way is to find the query in the generated text and take what comes after.
+        response_start_index = generated_text.find(formatted_prompt)
+        if response_start_index != -1:
+            response = generated_text[response_start_index + len(formatted_prompt):].strip()
+        else:
+            response = generated_text.strip() # Fallback if prompt isn't found perfectly
+
+        print(f"Response generated for user '{user_id}'.")
+        return jsonify({"response": response})
+    except Exception as e:
+        print(f"ERROR: An unexpected error occurred during query for user '{user_id}': {e}")
+        import traceback
+        traceback.print_exc()
+        return jsonify({"error": f"Error processing query: {e}"}), 500
+    
+if __name__ == "__main__":
+    # Call initialization function directly (no Flask debug)
+    initialize_models() 
+    print(f"Starting Flask RAG MVP application on http://0.0.0.0:7860 (Hugging Face Spaces default port)")
+    print(f"Using LLM: {LLM_MODEL_NAME_HF}, Embeddings: {EMBEDDING_MODEL_NAME_HF}")
+    print("API endpoints:")
+    print(" - POST /load_document (Requires 'user_id' and 'text')")
+    print(" - POST /query (Requires 'user_id' and 'query')")
+    
+    # Hugging Face Spaces typically runs on port 7860
+    app.run(host="0.0.0.0", port=7860)