Update app.py

app.py

@@ -4,16 +4,45 @@ import time
 import pandas as pd
 import sqlite3
 import logging
+import requests  # for HTTP calls to Gemini
 
-from langchain.document_loaders import OnlinePDFLoader  # for loading the PDF text
+from langchain.document_loaders import OnlinePDFLoader  # for loading PDF text
 from langchain.embeddings import HuggingFaceEmbeddings  # open source embedding model
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.vectorstores import Chroma  # updated import for vectorization
+from langchain_community.vectorstores import Chroma  # vectorization from langchain_community
 from langchain.chains import RetrievalQA  # for QA chain
-from langchain_community.chat_models import ChatOpenAI  # updated import for ChatOpenAI
 from langchain_core.prompts import PromptTemplate  # prompt template import
 
-# Setup basic logging
+# ------------------------------
+# Gemini API Wrapper
+# ------------------------------
+class ChatGemini:
+    def __init__(self, api_key, temperature=0, model_name="gemini-2.0-flash"):
+        self.api_key = api_key
+        self.temperature = temperature
+        self.model_name = model_name
+
+    def generate(self, prompt):
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model_name}:generateContent?key={self.api_key}"
+        payload = {
+            "contents": [{
+                "parts": [{"text": prompt}]
+            }]
+        }
+        headers = {"Content-Type": "application/json"}
+        response = requests.post(url, json=payload, headers=headers)
+        if response.status_code != 200:
+            raise Exception(f"Gemini API error: {response.status_code} - {response.text}")
+        data = response.json()
+        candidate = data.get("candidates", [{}])[0]
+        return candidate.get("output", {}).get("text", "No output from Gemini API")
+    
+    def __call__(self, prompt, **kwargs):
+        return self.generate(prompt)
+
+# ------------------------------
+# Setup Logging
+# ------------------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 log_messages = ""  # global log collector
@@ -23,12 +52,12 @@ def update_log(message):
     log_messages += message + "\n"
     logger.info(message)
 
-def load_pdf_and_generate_embeddings(pdf_doc, open_ai_key, relevant_pages):
+# ------------------------------
+# PDF Embedding & QA Chain (No OCR)
+# ------------------------------
+def load_pdf_and_generate_embeddings(pdf_doc, gemini_api_key, relevant_pages):
     try:
-        if open_ai_key is not None:
-            os.environ['OPENAI_API_KEY'] = open_ai_key
-        
-        # Use the file path directly as OCR is removed; text is extracted via the document loader.
+        # Use the PDF file's path to extract text.
         pdf_path = pdf_doc.name
         loader = OnlinePDFLoader(pdf_path)
         pages = loader.load_and_split()
@@ -59,18 +88,21 @@ def load_pdf_and_generate_embeddings(pdf_doc, open_ai_key, relevant_pages):
     
         global pdf_qa
         pdf_qa = RetrievalQA.from_chain_type(
-            llm=ChatOpenAI(temperature=0, model_name="gpt-4"),
+            llm=ChatGemini(api_key=gemini_api_key, temperature=0, model_name="gemini-2.0-flash"),
             chain_type="stuff",
             retriever=vectordb.as_retriever(search_kwargs={"k": 5}),
             chain_type_kwargs=chain_type_kwargs,
             return_source_documents=False
         )
-        update_log("PDF embeddings generated and QA chain initialized.")
+        update_log("PDF embeddings generated and QA chain initialized using Gemini.")
         return "Ready"
     except Exception as e:
         update_log(f"Error in load_pdf_and_generate_embeddings: {str(e)}")
         return f"Error: {str(e)}"
 
+# ------------------------------
+# SQLite Question Set Functions
+# ------------------------------
 def create_db_connection():
     DB_FILE = "./questionset.db"
     connection = sqlite3.connect(DB_FILE, check_same_thread=False)
@@ -226,6 +258,9 @@ def answer_query(query):
 def get_log():
     return log_messages
 
+# ------------------------------
+# Gradio Interface
+# ------------------------------
 css = """
 #col-container {max-width: 700px; margin: auto;}
 """
@@ -243,7 +278,7 @@ with gr.Blocks(css=css, theme=gr.themes.Monochrome()) as demo:
     
     with gr.Tab("Chatbot"):
         with gr.Column():
-            open_ai_key = gr.Textbox(label="Your GPT-4 API Key", type="password")
+            gemini_api_key = gr.Textbox(label="Your Gemini API Key", type="password")
             pdf_doc = gr.File(label="Load a PDF", file_types=['.pdf'], type='filepath')
             relevant_pages = gr.Textbox(label="Optional: Comma separated page numbers")
         
@@ -272,15 +307,13 @@ with gr.Blocks(css=css, theme=gr.themes.Monochrome()) as demo:
             
             log_window = gr.Textbox(label="Log Window", interactive=False, lines=10)
     
-    with gr.Tab("OCR Converter"):
-        # This tab is now repurposed (or can be removed)
+    with gr.Tab("Text Extractor"):
         with gr.Column():
-            image_pdf = gr.File(label="Load PDF for Conversion", file_types=['.pdf'], type='filepath')
+            image_pdf = gr.File(label="Load PDF for Text Extraction", file_types=['.pdf'], type='filepath')
         with gr.Row():
             extracted_text = gr.Textbox(label="Extracted Text", lines=10)
             extract_btn = gr.Button("Extract Text")
     
-        # For demonstration, extract text using OnlinePDFLoader
         def extract_text(pdf_file):
             try:
                 loader = OnlinePDFLoader(pdf_file.name)
@@ -306,7 +339,7 @@ with gr.Blocks(css=css, theme=gr.themes.Monochrome()) as demo:
     refresh_log_btn = gr.Button("Refresh Log")
     refresh_log_btn.click(get_log, outputs=log_window)
     
-    load_pdf_btn.click(load_pdf_and_generate_embeddings, inputs=[pdf_doc, open_ai_key, relevant_pages], outputs=status)
+    load_pdf_btn.click(load_pdf_and_generate_embeddings, inputs=[pdf_doc, gemini_api_key, relevant_pages], outputs=status)
     summarize_pdf_btn.click(summarize_contents, outputs=summary)
     submit_query_btn.click(answer_query, inputs=input_query, outputs=output_answer)