Update app.py

app.py

@@ -9,278 +9,221 @@ import openai
 import gradio as gr
 from epub2txt import epub2txt
 
-class AppLogic:
+# --- Configuration and Constants ---
+MODEL_NAME = os.getenv("POE_MODEL", "GPT-5-mini")
+PROMPT = os.getenv("prompt", "Summarize the following text:")
+KEYS_FILE = "user_keys.json"
+
+# --- Helper Functions for Security and User Management ---
+# By moving logic out of a class, we avoid complex state issues.
+
+def get_user_id(username):
+    """Generates a unique, stable ID from a username."""
+    return hashlib.sha256(username.encode()).hexdigest()
+
+def generate_fernet_key(user_id):
+    """Generates a deterministic encryption key from a user ID."""
+    salt = user_id.encode()[:16].ljust(16, b'\0') # Salt must be 16 bytes
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=100000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(user_id.encode()))
+    return Fernet(key)
+
+def load_user_keys():
+    """Loads the entire user key database."""
+    if os.path.exists(KEYS_FILE):
+        with open(KEYS_FILE, 'r') as f:
+            return json.load(f)
+    return {}
+
+def save_user_keys(keys_data):
+    """Saves the entire user key database."""
+    with open(KEYS_FILE, 'w') as f:
+        json.dump(keys_data, f)
+
+def get_decrypted_api_key(user_id):
+    """Retrieves and decrypts a single user's API key."""
+    keys_data = load_user_keys()
+    encrypted_key_b64 = keys_data.get(user_id)
+    if encrypted_key_b64:
+        try:
+            cipher = generate_fernet_key(user_id)
+            encrypted_key = base64.urlsafe_b64decode(encrypted_key_b64.encode())
+            return cipher.decrypt(encrypted_key).decode()
+        except Exception as e:
+            print(f"Decryption failed for user {user_id}: {e}")
+    return None
+
+def set_encrypted_api_key(user_id, api_key):
+    """Encrypts and saves a single user's API key."""
+    keys_data = load_user_keys()
+    cipher = generate_fernet_key(user_id)
+    encrypted_key = cipher.encrypt(api_key.encode())
+    keys_data[user_id] = base64.urlsafe_b64encode(encrypted_key).decode()
+    save_user_keys(keys_data)
+
+def delete_api_key(user_id):
+    """Deletes a user's API key."""
+    keys_data = load_user_keys()
+    if user_id in keys_data:
+        del keys_data[user_id]
+        save_user_keys(keys_data)
+
+# --- Gradio Event Handlers ---
+
+def initialize_client(api_key):
+    """Initializes and returns an OpenAI client or None on failure."""
+    if not api_key:
+        return None
+    try:
+        client = openai.OpenAI(api_key=api_key, base_url="https://api.poe.com/v1")
+        # Test connection
+        client.models.list()
+        return client
+    except Exception as e:
+        print(f"Failed to initialize client: {e}")
+        return None
+
+def update_ui_on_load(request: gr.Request):
     """
-    This class now ONLY contains the application logic (handling keys, processing files).
-    The UI definition has been moved into the `build_ui` function.
+    This function runs after the user logs in and the page loads.
+    It configures the UI based on whether the user has a saved API key.
     """
-    def __init__(self, request: gr.Request):
-        # Configuration
-        self.model_name = os.getenv("POE_MODEL", "GPT-5-mini")
-        self.prompt = os.getenv("prompt", "Summarize the following text:")
-        self.client = None
-        self.api_key = None
-        self.keys_file = "user_keys.json"
+    user_id = get_user_id(request.username)
+    api_key = get_decrypted_api_key(user_id)
+    client = initialize_client(api_key)
+
+    if client:
+        # Key is valid and loaded
+        welcome_msg = f"Welcome back, **{request.username}**! Your saved API key is loaded."
+        return {
+            welcome_md: gr.update(value=welcome_msg),
+            api_key_section: gr.update(visible=False),
+            inp: gr.update(visible=True),
+            clear_key_btn: gr.update(visible=True)
+        }
+    else:
+        # No key or invalid key found
+        welcome_msg = f"Welcome, **{request.username}**! Please provide your Poe API key to begin."
+        return {
+            welcome_md: gr.update(value=welcome_msg),
+            api_key_section: gr.update(visible=True),
+            inp: gr.update(visible=False),
+            clear_key_btn: gr.update(visible=False)
+        }
+
+def set_api_key(api_key, remember_key, request: gr.Request):
+    """Handles the 'Set API Key' button click."""
+    user_id = get_user_id(request.username)
+    client = initialize_client(api_key)
+
+    if not client:
+        return {out: gr.update(value="❌ Invalid or incorrect API key. Please check it and try again.")}
+
+    if remember_key:
+        set_encrypted_api_key(user_id, api_key)
+        msg = "✅ API key validated and saved! You can now upload an ePub."
+    else:
+        delete_api_key(user_id) # Ensure no old key is stored if user unchecks
+        msg = "✅ API key validated for this session. You can now upload an ePub."
+
+    return {
+        out: gr.update(value=msg),
+        inp: gr.update(visible=True),
+        api_key_section: gr.update(visible=False),
+        clear_key_btn: gr.update(visible=remember_key)
+    }
+
+def clear_saved_key(request: gr.Request):
+    """Handles the 'Clear Saved Key' button click."""
+    user_id = get_user_id(request.username)
+    delete_api_key(user_id)
+    return {
+        out: gr.update(value="✅ Saved API key cleared. Please enter an API key to continue."),
+        api_key_section: gr.update(visible=True),
+        inp: gr.update(visible=False),
+        clear_key_btn: gr.update(visible=False),
+        api_key_input: gr.update(value="")
+    }
+
+def process_epub(file, request: gr.Request):
+    """Processes the uploaded ePub file. This is a generator function."""
+    user_id = get_user_id(request.username)
+    api_key = get_decrypted_api_key(user_id)
+    
+    # Re-initialize client to ensure API key is available for this long-running task
+    client = initialize_client(api_key)
+    if not client:
+        yield "⚠️ Your API key is missing or invalid. Please clear and set your key again."
+        return
+    
+    # ... (The rest of your ePub processing logic is identical) ...
+    try:
+        ch_list = epub2txt(file.name, outputlist=True)
+        title = epub2txt.title
+        yield f"# {title}\n\nProcessing ePub file..."
+        sm_list = []
+        for text in ch_list[2:]:
+            if text.strip():
+                response = client.chat.completions.create(
+                    model=MODEL_NAME,
+                    messages=[{"role": "user", "content": PROMPT + "\n\n" + text[:4000]}], # Simple chunking
+                ).choices[0].message.content
+                sm_list.append(response)
         
-        # User-specific properties derived from the request
-        self.username = request.username
-        self.user_id = self._get_user_id(self.username)
-
-        # Initialize by loading the saved key, if it exists
-        self.load_and_initialize_saved_key()
-
-    def _get_user_id(self, username):
-        """Generate a unique user ID from the username."""
-        if not username:
-            return None
-        return hashlib.sha256(username.encode()).hexdigest()
+        yield f"# {title}\n\n" + "\n\n---\n\n".join(sm_list)
 
-    def _generate_key_from_user(self, user_id):
-        """Generate encryption key from user ID."""
-        salt = user_id.encode()[:32].ljust(32, b'0')
-        kdf = PBKDF2HMAC(
-            algorithm=hashes.SHA256(),
-            length=32,
-            salt=salt,
-            iterations=100000,
-        )
-        key = base64.urlsafe_b64encode(kdf.derive(user_id.encode()))
-        return Fernet(key)
-
-    def _load_user_keys(self):
-        """Load encrypted user keys from file."""
-        if os.path.exists(self.keys_file):
-            with open(self.keys_file, 'r') as f:
-                return json.load(f)
-        return {}
-
-    def _save_user_keys(self, keys_data):
-        """Save encrypted user keys to file."""
-        with open(self.keys_file, 'w') as f:
-            json.dump(keys_data, f)
-
-    def _get_saved_key(self):
-        """Retrieve and decrypt user's API key."""
-        keys_data = self._load_user_keys()
-        if self.user_id in keys_data:
-            try:
-                cipher_suite = self._generate_key_from_user(self.user_id)
-                encrypted_key = base64.urlsafe_b64decode(keys_data[self.user_id].encode())
-                return cipher_suite.decrypt(encrypted_key).decode()
-            except Exception as e:
-                print(f"Error retrieving saved key for {self.username}: {e}")
-        return None
-
-    def _save_encrypted_key(self, api_key):
-        """Encrypt and save user's API key."""
-        keys_data = self._load_user_keys()
-        cipher_suite = self._generate_key_from_user(self.user_id)
-        encrypted_key = cipher_suite.encrypt(api_key.encode())
-        keys_data[self.user_id] = base64.urlsafe_b64encode(encrypted_key).decode()
-        self._save_user_keys(keys_data)
-        return True
-
-    def _delete_saved_key(self):
-        """Delete user's saved API key."""
-        keys_data = self._load_user_keys()
-        if self.user_id in keys_data:
-            del keys_data[self.user_id]
-            self._save_user_keys(keys_data)
-        return True
-    
-    def _initialize_client(self):
-        """Initialize the Poe API client if an API key is present."""
-        if not self.api_key:
-            return False
-        try:
-            self.client = openai.OpenAI(
-                api_key=self.api_key,
-                base_url="https://api.poe.com/v1",
-            )
-            # A quick check to see if the client is functional without making a call
-            return hasattr(self.client, 'chat')
-        except Exception as e:
-            print(f"Error initializing Poe client for {self.username}: {e}")
-            return False
+    except Exception as e:
+        yield f"An error occurred: {e}"
 
-    def load_and_initialize_saved_key(self):
-        """Checks for a saved key and initializes the client if found."""
-        saved_key = self._get_saved_key()
-        if saved_key:
-            self.api_key = saved_key
-            self._initialize_client()
 
-    def set_api_key(self, api_key, remember_key):
-        """Set and validate the API key."""
-        if not api_key or not api_key.strip():
-            return (
-                gr.update(value="⚠️ Please enter a valid API key."),
-                gr.update(visible=False),
-                gr.update(visible=True),
-                gr.update(visible=False),
-            )
-        
-        self.api_key = api_key.strip()
-        
-        if self._initialize_client():
-            try:
-                # Test the API key with a simple request
-                self.client.chat.completions.create(
-                    model=self.model_name, messages=[{"role": "user", "content": "Hello"}], max_tokens=10
-                )
-                
-                if remember_key:
-                    self._save_encrypted_key(self.api_key)
-                    msg = "✅ API key validated and saved! You can now upload an ePub file."
-                else:
-                    self._delete_saved_key() # Ensure no old key is lingering if user unchecks
-                    msg = "✅ API key validated for this session! You can now upload an ePub file."
+# --- Build the Gradio Interface ---
 
-                return (
-                    gr.update(value=msg),
-                    gr.update(visible=True),  # Show file input
-                    gr.update(visible=False), # Hide API key section
-                    gr.update(visible=remember_key), # Show clear button only if saved
-                )
-            except Exception as e:
-                self.client = None # Reset client on failure
-                error_msg = str(e)
-                if "401" in error_msg:
-                    return gr.update(value="❌ Invalid API key. Please check it and try again."), gr.update(), gr.update(), gr.update()
-                return gr.update(value=f"❌ API connection error: {error_msg}"), gr.update(), gr.update(), gr.update()
-        else:
-            return gr.update(value="❌ Failed to initialize API client."), gr.update(), gr.update(), gr.update()
+with gr.Blocks(title="ePub Summarizer") as demo:
+    welcome_md = gr.Markdown("Welcome! Please log in to continue.")
 
-    def clear_saved_key(self):
-        """Clear the user's saved API key."""
-        self._delete_saved_key()
-        self.api_key = None
-        self.client = None
-        return (
-            gr.update(value="✅ Saved API key cleared. Please enter your API key to continue."),
-            gr.update(visible=True),   # Show API key section
-            gr.update(visible=False),  # Hide clear button
-            gr.update(value=""),       # Clear API key input field
-            gr.update(visible=False),  # Hide file input
-        )
-    
-    # ... (get_model_response and process methods are unchanged) ...
-    def get_model_response(self, text: str) -> str:
-        """Get response from Poe API for the given text."""
-        if not self.client:
-            return "Error: API client not initialized. Please set your API key."
-        try:
-            chat = self.client.chat.completions.create(
-                model=self.model_name,
-                messages=[{"role": "user", "content": text}],
-            )
-            return chat.choices[0].message.content
-        except Exception as e:
-            return f"Error calling Poe API: {str(e)}"
+    # API Key Section (initially hidden, made visible by update_ui_on_load)
+    with gr.Column(visible=False) as api_key_section:
+        api_key_input = gr.Textbox(label="Poe API Key", type="password")
+        remember_key = gr.Checkbox(label="Remember my API key", value=True)
+        api_key_btn = gr.Button("Set API Key", variant="primary")
 
-    def process(self, file):
-        """Processes the uploaded ePub file."""
-        if not self.client:
-            yield gr.update(value='⚠️ Please set your Poe API key first.')
-            return
-        if file is None:
-            yield gr.update(value='Please upload an ePub file.')
-            return
+    # Main App Section
+    out = gr.Markdown()
+    inp = gr.File(file_types=['.epub'], visible=False, label="Upload ePub File")
+    clear_key_btn = gr.Button("Clear Saved Key", variant="secondary", visible=False)
 
-        try:
-            ch_list = epub2txt(file.name, outputlist=True)
-            chapter_titles = epub2txt.content_titles
-            title = epub2txt.title
-            
-            yield gr.update(value=f"# {title}\n\nProcessing ePub file...")
-            sm_list = []
-            
-            # Skip first 2 chapters (usually metadata)
-            for idx, text in enumerate(ch_list[2:], 1):
-                if not text.strip(): continue
-                yield gr.update(value=f"# {title}\n\nProcessing chapter {idx}...")
-                
-                docs = []
-                chunk_size = 2000
-                for i in range(0, len(text), chunk_size):
-                    chunk = text[i:i + 2048]
-                    if len(chunk.strip()) > 0:
-                        response = self.get_model_response(self.prompt + "\n\n" + chunk)
-                        docs.append(response)
-                
-                if docs:
-                    hist = docs[0]
-                    for doc in docs[1:]:
-                        combined_text = f"{self.prompt}\n\nCombine these summaries:\n\n{hist}\n\n{doc}"
-                        hist = self.get_model_response(combined_text)
-                    sm_list.append(hist)
-                
-                final_summaries = "\n\n".join([f"## {ct}\n\n{sm}" for ct, sm in zip(chapter_titles[2:idx+1], sm_list)])
-                yield gr.update(value=f"# {title}\n\n{final_summaries}")
-            
-            if sm_list:
-                complete_summary = f"# {title}\n\n" + "\n\n".join([f"## {ct}\n\n{sm}" for ct, sm in zip(chapter_titles[2:len(sm_list)+2], sm_list)])
-                yield gr.update(value=complete_summary)
-            else:
-                yield gr.update(value=f"# {title}\n\nNo content found to summarize.")
-                
-        except Exception as e:
-            yield gr.update(value=f"Error processing file: {str(e)}")
+    # --- Event Wiring ---
 
-# --- UI Building Function ---
-def build_ui(request: gr.Request):
-    """
-    This function is the main entry point for the UI.
-    It's only called AFTER the user has logged in.
-    """
-    # Instantiate the logic class for this specific user session
-    logic = AppLogic(request)
-    
-    # Determine the initial state based on whether a key was loaded
-    has_saved_key = logic.client is not None
-    
-    welcome_message = (f"# ePub Summarization Tool\n\nWelcome back, **{logic.username}**! "
-                       "✅ Your saved API key is loaded and ready.") if has_saved_key else \
-                      (f"# ePub Summarization Tool\n\nWelcome, **{logic.username}**! "
-                       "Please set up your Poe API key to begin.")
-                       
-    with gr.Blocks(title="ePub Summarizer") as demo:
-        gr.Markdown(welcome_message)
-        
-        # API Key input section
-        with gr.Column(visible=not has_saved_key) as api_key_section:
-            api_key_input = gr.Textbox(label="Poe API Key", type="password")
-            remember_key = gr.Checkbox(label="Remember my API key", value=True)
-            api_key_btn = gr.Button("Set API Key", variant="primary")
-        
-        # Main content area
-        out = gr.Markdown()
-        inp = gr.File(file_types=['.epub'], visible=has_saved_key, label="Upload ePub File")
-        clear_key_btn = gr.Button("Clear Saved Key", variant="secondary", visible=has_saved_key)
+    # 1. After login, the 'load' event fires, calling update_ui_on_load
+    demo.load(
+        fn=update_ui_on_load,
+        outputs=[welcome_md, api_key_section, inp, clear_key_btn]
+    )
 
-        # --- Event Handlers ---
-        api_key_btn.click(
-            logic.set_api_key, 
-            inputs=[api_key_input, remember_key], 
-            outputs=[out, inp, api_key_section, clear_key_btn]
-        )
-        clear_key_btn.click(
-            logic.clear_saved_key,
-            outputs=[out, api_key_section, clear_key_btn, api_key_input, inp]
-        )
-        inp.change(logic.process, inputs=[inp], outputs=[out])
+    # 2. User interactions trigger other handlers
+    api_key_btn.click(
+        fn=set_api_key,
+        inputs=[api_key_input, remember_key],
+        outputs=[out, inp, api_key_section, clear_key_btn]
+    )
+    clear_key_btn.click(
+        fn=clear_saved_key,
+        outputs=[out, api_key_section, inp, clear_key_btn, api_key_input]
+    )
+    inp.change(fn=process_epub, inputs=[inp], outputs=[out])
 
-    return demo
+# --- Main Execution ---
 
-# --- Main Application Execution ---
 if __name__ == "__main__":
-    # This dummy function is for local testing. On Hugging Face, it's ignored.
+    # This dummy function is used for local testing.
+    # On Hugging Face Spaces, their login system is used automatically.
     def auth(username, password):
         return True
 
-    # Use gr.mount_gradio_app for FastAPI integration if needed, or just launch directly
-    with gr.Blocks() as demo:
-        # The auth() method gates the UI. build_ui is only called after successful login.
-        demo.auth(build_ui, auth)
-
-    demo.queue().launch(show_error=True)
+    # The 'auth' parameter is passed to launch(), not as a separate method call.
+    demo.queue().launch(auth=auth, show_error=True)