Update app.py

app.py

@@ -1,58 +1,223 @@
 import os
-import subprocess
+import openai
 import gradio as gr
 from epub2txt import epub2txt
-from websocket import create_connection
-
-if not os.path.exists(os.getenv("checkpoint_path")):
-    os.system("git clone --recurse-submodules https://github.com/ztxz16/fastllm.git")
-    os.system("cd fastllm; mkdir build; cd build; cmake ..; make -j; cd tools; python setup.py install --user --prefix=")
-    os.system("wget https://huggingface.co/huangyuyang/chatglm2-6b-int4.flm/resolve/main/chatglm2-6b-int4.flm")
-    subprocess.Popen(["uvicorn", "api:app"])
 
 class GUI:
     def __init__(self, *args, **kwargs):
-        with gr.Blocks() as demo:
+        # 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 = os.getenv("POE_API_KEY")
+        
+        with gr.Blocks(title="ePub Summarizer") as demo:
             with gr.Row():
                 gr.Markdown(scale=2).attach_load_event(self.hello, None)
                 gr.LoginButton()
                 gr.LogoutButton()
+            
+            # API Key input section
+            with gr.Row(visible=not self.api_key):
+                with gr.Column():
+                    gr.Markdown("""
+                    ### Poe API Key Required
+                    
+                    To use this tool, you need a Poe API key:
+                    
+                    1. Visit [https://poe.com/api_key](https://poe.com/api_key)
+                    2. If you don't have an account, create one first
+                    3. Generate a new API key or copy your existing one
+                    4. Paste it in the field below
+                    
+                    **Note**: Your API key is only stored temporarily during this session for security.
+                    """)
+                    api_key_input = gr.Textbox(
+                        label="Poe API Key",
+                        placeholder="Enter your Poe API key here...",
+                        type="password"
+                    )
+                    api_key_btn = gr.Button("Set API Key", variant="primary")
+            
             out = gr.Markdown()
-            inp = gr.File(file_types=['.epub'])
+            inp = gr.File(file_types=['.epub'], visible=bool(self.api_key))
+            
+            # Set up event handlers
+            if not self.api_key:
+                api_key_btn.click(
+                    self.set_api_key, 
+                    inputs=[api_key_input], 
+                    outputs=[out, inp]
+                )
+            else:
+                self._initialize_client()
+            
             inp.change(self.process, inp, out)
-            self.ws = None
-            self.out = []
+            
         demo.queue(concurrency_count=2).launch()
 
-    def process(self, file, profile: gr.OAuthProfile | None):
+    def _initialize_client(self):
+        """Initialize the Poe API client"""
+        try:
+            self.client = openai.OpenAI(
+                api_key=self.api_key,
+                base_url="https://api.poe.com/v1",
+            )
+            return True
+        except Exception as e:
+            print(f"Error initializing Poe client: {e}")
+            return False
+
+    def set_api_key(self, api_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)
+        
+        self.api_key = api_key.strip()
+        
+        if self._initialize_client():
+            # Test the API key with a simple request
+            try:
+                test_chat = self.client.chat.completions.create(
+                    model=self.model_name,
+                    messages=[{"role": "user", "content": "Hello"}],
+                    max_tokens=10
+                )
+                return (
+                    gr.update(value="✅ API key validated successfully! You can now upload an ePub file."), 
+                    gr.update(visible=True)
+                )
+            except Exception as e:
+                error_msg = str(e)
+                if "401" in error_msg or "unauthorized" in error_msg.lower():
+                    return (
+                        gr.update(value="❌ Invalid API key. Please check your key and try again."), 
+                        gr.update(visible=False)
+                    )
+                elif "quota" in error_msg.lower() or "limit" in error_msg.lower():
+                    return (
+                        gr.update(value="⚠️ API key valid but quota exceeded. Please check your Poe account."), 
+                        gr.update(visible=False)
+                    )
+                else:
+                    return (
+                        gr.update(value=f"❌ API connection error: {error_msg}"), 
+                        gr.update(visible=False)
+                    )
+        else:
+            return (
+                gr.update(value="❌ Failed to initialize API client."), 
+                gr.update(visible=False)
+            )
+
+    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"
+            
+        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:
+            print(f"Error calling Poe API: {e}")
+            return f"Error processing text: {str(e)}"
+
+    def process(self, file, profile: gr.OAuthProfile | None = None):
         if profile is None:
             return gr.update(value='Login to access the tool.')
 
-        h = ''
-        chapter_titles = epub2txt.content_titles
-        title = epub2txt.title
-        if self.ws is None:
-            self.ws = create_connection(f"ws://localhost:8000/ws")
-            self.ws.send(file.name)
-        res = ''
-        while 'output: ' not in res:
-            res = self.ws.recv()
-            if 'chsum: ' in res:
-                self.out.append(res.replace("chsum: ", ""))
-            elif 'draft_sum: ' in res:
-                h = res[11:]
-            elif 'output: ' in res:
-                self.ws.close()
-                self.ws = None
-                self.out = []
-                yield gr.update(value=res)
-            yield gr.update(
-                value=f"# {title}\n\n" +h+ "\n\n".join(
-                    [f"## {ct}\n\n{c}" for ct, c in zip(chapter_titles, self.out)]))
-
-    def hello(self, profile: gr.OAuthProfile | None):
+        if not self.client:
+            return gr.update(value='⚠️ Please set your Poe API key first.')
+
+        if file is None:
+            return gr.update(value='Please upload an ePub file.')
+
+        try:
+            # Extract content from ePub
+            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 = []
+            
+            # Process each chapter (skip first 2 as they're 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 = []
+                # Split chapter into chunks for processing
+                chunk_size = 2000
+                for i in range(0, len(text), chunk_size):
+                    chunk = text[i:i+2048]  # Slight overlap for context
+                    if len(chunk.strip()) > 0:
+                        response = self.get_model_response(self.prompt + "\n\n" + chunk)
+                        docs.append(response)
+                        
+                        # Update UI with current progress
+                        current_summaries = "\n\n".join([
+                            f"## {ct}\n\n{sm}" 
+                            for ct, sm in zip(chapter_titles[2:idx+1], sm_list + [f"Processing chunk {len(docs)}..."])
+                        ])
+                        yield gr.update(value=f"# {title}\n\n{current_summaries}")
+                
+                # Combine chunk summaries into chapter summary
+                if docs:
+                    if len(docs) == 1:
+                        hist = docs[0]
+                    else:
+                        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)
+                            
+                            # Update UI with draft summary
+                            current_summaries = "\n\n".join([
+                                f"## {ct}\n\n{sm}" 
+                                for ct, sm in zip(chapter_titles[2:idx+1], sm_list + [f"Draft: {hist}"])
+                            ])
+                            yield gr.update(value=f"# {title}\n\n{current_summaries}")
+                    
+                    sm_list.append(hist)
+                
+                # Update final output for this chapter
+                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}")
+            
+            # Final complete summary
+            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)}")
+
+    def hello(self, profile: gr.OAuthProfile | None = None):
         if profile is None:
-            return '# ePub summarization tool\n\nLogin to access the tool.'
-        return f"# ePub summarization tool\n\nWelcome {profile.name}!!"
+            return '# ePub Summarization Tool\n\nLogin to access the tool.'
+        
+        if not self.api_key:
+            return f"# ePub Summarization Tool\n\nWelcome {profile.name}!\n\nPlease set your Poe API key below to get started."
+        
+        return f"# ePub Summarization Tool\n\nWelcome {profile.name}! Ready to summarize ePub files."
 
-GUI()
+# Run the application
+if __name__ == "__main__":
+    GUI()