Update app.py

app.py

@@ -1,3 +1,5 @@
+from openai import OpenAI
+from os import getenv
 from flask import Flask, request, jsonify, render_template
 import fitz  # PyMuPDF for PDF text extraction
 import faiss  # FAISS for vector search
@@ -11,18 +13,19 @@ app = Flask(__name__, template_folder=os.getcwd())
 
 # Default settings
 class ChatConfig:
-    MODEL = "meta-llama/Llama-3.1-8B-Instruct"  # Change back to Gemma
+    MODEL = "google/gemma-7b-it:free"  # Use OpenRouter's Gemma model
     DEFAULT_SYSTEM_MSG = "You are an AI assistant answering only based on the uploaded PDF."
     DEFAULT_MAX_TOKENS = 512
     DEFAULT_TEMP = 0.3
     DEFAULT_TOP_P = 0.95
 
 # Get the token from environment variable
-HF_TOKEN = os.getenv('HUGGINGFACE_TOKEN')
-client = InferenceClient(
-    ChatConfig.MODEL,
-    token=HF_TOKEN
+OPENROUTER_API_KEY = getenv('OPENROUTER_API_KEY')
+client = OpenAI(
+    base_url="https://openrouter.ai/api/v1",
+    api_key=OPENROUTER_API_KEY,
 )
+
 embed_model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder="/tmp")
 vector_dim = 384  # Embedding size
 index = faiss.IndexFlatL2(vector_dim)  # FAISS index
@@ -92,21 +95,16 @@ def generate_response(
             messages.append({"role": "assistant", "content": bot_msg})
 
     try:
-        response = ""
-        for chunk in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            token = chunk.choices[0].delta.content or ""
-            response += token
-            yield response
+        # Use OpenRouter to get the response
+        completion = client.chat.completions.create(
+            model="google/gemma-7b-it:free",
+            messages=messages
+        )
+        return completion.choices[0].message.content
     except Exception as e:
         print(f"Error generating response: {str(e)}")
-        yield "I apologize, but I encountered an error while generating the response. Please try again."
-        
+        return "I apologize, but I encountered an error while generating the response. Please try again."
+
 @app.route('/')
 def index():
     """Serve the HTML page for the user interface"""
@@ -148,7 +146,7 @@ def ask_question():
     message = request.json.get('message')
     history = request.json.get('history', [])
     response = generate_response(message, history)
-    return jsonify({"response": "".join(response)})  # Join all streamed responses
+    return jsonify({"response": response})
 
 if __name__ == '__main__':
     app.run(debug=True)