Update app.py

app.py

@@ -11,6 +11,8 @@ import io
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Keep your existing CONVERSATION_PROMPT and EXTRACTION_PROMPT here
+
 CONVERSATION_PROMPT = """You are LOSS DOG, a professional profile builder. Your goal is to have natural conversations
 with users to gather information about their professional background across 9 categories:
 
@@ -135,26 +137,47 @@ Return the data in this exact structure:
 }
 
 IMPORTANT: Return ONLY the JSON. Do not add any explanation text."""
-import json
-import logging
-from datetime import datetime
-from typing import Dict, List, Optional, Any
-import gradio as gr
-from openai import AsyncOpenAI
-import PyPDF2
-import io
-
-# ... (previous imports and prompts remain the same)
-
 class ProfileBuilder:
     def __init__(self):
         self.conversation_history = []
         self.client = None
+        self.pdf_text = None
     
     def _initialize_client(self, api_key: str) -> None:
-        if not api_key.startswith("sk-"):
-            raise ValueError("Invalid API key format")
-        self.client = AsyncOpenAI(api_key=api_key)
+        """Initialize OpenAI client if not already initialized"""
+        if not self.client:
+            if not api_key.startswith("sk-"):
+                raise ValueError("Invalid API key format")
+            self.client = AsyncOpenAI(api_key=api_key)
+
+    async def process_message(self, message: str, api_key: str) -> Dict[str, Any]:
+        """Process a chat message"""
+        try:
+            # Initialize client if needed
+            self._initialize_client(api_key)
+
+            # Add user message to history
+            self.conversation_history.append({"role": "user", "content": message})
+            
+            # Get AI response
+            completion = await self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": CONVERSATION_PROMPT},
+                    *self.conversation_history
+                ],
+                temperature=0.7
+            )
+            
+            # Extract and store response
+            ai_message = completion.choices[0].message.content
+            self.conversation_history.append({"role": "assistant", "content": ai_message})
+            
+            return {"response": ai_message}
+
+        except Exception as e:
+            logger.error(f"Error in message processing: {str(e)}")
+            return {"error": str(e)}
 
     async def extract_from_pdf(self, pdf_content: bytes) -> str:
         """Extract text from PDF file"""
@@ -164,22 +187,22 @@ class ProfileBuilder:
             text = ""
             for page in pdf_reader.pages:
                 text += page.extract_text()
+            self.pdf_text = text  # Store the PDF text
             return text
         except Exception as e:
-            logger.error(f"Error extracting PDF: {str(e)}")
+            logger.error(f"PDF extraction error: {str(e)}")
             raise
 
     async def process_pdf(self, pdf_path: str, api_key: str) -> Dict[str, Any]:
-        """Process PDF resume and extract information"""
+        """Process PDF resume"""
         try:
-            if not self.client:
-                self._initialize_client(api_key)
+            self._initialize_client(api_key)
 
+            # Read and extract PDF content
             with open(pdf_path, 'rb') as file:
-                pdf_content = file.read()
-                resume_text = await self.extract_from_pdf(pdf_content)
+                resume_text = await self.extract_from_pdf(file.read())
 
-            # Use the extraction prompt directly on PDF content
+            # Process with AI
             completion = await self.client.chat.completions.create(
                 model="gpt-4o-mini",
                 messages=[
@@ -189,10 +212,12 @@ class ProfileBuilder:
                 temperature=0.3
             )
 
+            # Parse response
             response_text = completion.choices[0].message.content.strip()
             profile_data = json.loads(response_text)
 
-            return {
+            # Create profile object
+            profile = {
                 "profile_data": profile_data,
                 "metadata": {
                     "generated_at": datetime.now().isoformat(),
@@ -200,11 +225,65 @@ class ProfileBuilder:
                 }
             }
 
+            return profile
+
         except Exception as e:
-            logger.error(f"Error processing PDF: {str(e)}")
+            logger.error(f"PDF processing error: {str(e)}")
             return {"error": str(e)}
 
-    # ... (rest of the ProfileBuilder class remains the same)
+    async def generate_profile(self) -> tuple[Dict[str, Any], Optional[str]]:
+        """Generate profile from conversation or PDF"""
+        try:
+            if not self.client:
+                raise ValueError("OpenAI client not initialized")
+
+            # Determine source and prepare content
+            if self.conversation_history:
+                content = "\n".join(
+                    f"{msg['role']}: {msg['content']}"
+                    for msg in self.conversation_history
+                )
+                source = "conversation"
+            elif self.pdf_text:
+                content = self.pdf_text
+                source = "pdf"
+            else:
+                raise ValueError("No content available for profile generation")
+
+            # Get AI extraction
+            completion = await self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": EXTRACTION_PROMPT},
+                    {"role": "user", "content": f"Extract profile information from this {source}:\n\n{content}"}
+                ],
+                temperature=0.3
+            )
+
+            # Parse response
+            response_text = completion.choices[0].message.content.strip()
+            profile_data = json.loads(response_text)
+
+            # Create profile
+            profile = {
+                "profile_data": profile_data,
+                "metadata": {
+                    "generated_at": datetime.now().isoformat(),
+                    "source": source
+                }
+            }
+
+            # Save to file
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"profile_{timestamp}.json"
+            with open(filename, 'w', encoding='utf-8') as f:
+                json.dump(profile, f, indent=2)
+
+            return profile, filename
+
+        except Exception as e:
+            logger.error(f"Profile generation error: {str(e)}")
+            return {"error": str(e)}, None
 
 def create_gradio_interface():
     builder = ProfileBuilder()
@@ -212,15 +291,18 @@ def create_gradio_interface():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("# 🐕 LOSS DOG - Professional Profile Builder")
         
+        # Common API key input
         api_key = gr.Textbox(
             label="OpenAI API Key",
             type="password",
             placeholder="Enter your OpenAI API key"
         )
 
+        # Tab interface
         with gr.Tabs() as tabs:
+            # Resume Upload Tab
             with gr.Tab("Upload Resume"):
-                upload_text = gr.Markdown("""
+                gr.Markdown("""
                 # Upload Your Resume
                 Upload your existing resume in PDF format and let LOSS DOG extract your professional profile.
                 """)
@@ -230,12 +312,16 @@ def create_gradio_interface():
                 )
                 process_pdf_btn = gr.Button("Process Resume")
 
+            # Chat Tab
             with gr.Tab("Chat with AI"):
-                chat_text = gr.Markdown("""
+                gr.Markdown("""
                 # Chat with LOSS DOG
                 Start a conversation with LOSS DOG to build your professional profile from scratch.
                 """)
-                chatbot = gr.Chatbot(label="Conversation")
+                chatbot = gr.Chatbot(
+                    label="Conversation",
+                    height=400
+                )
                 with gr.Row():
                     msg = gr.Textbox(
                         label="Message",
@@ -243,14 +329,28 @@ def create_gradio_interface():
                     )
                     send = gr.Button("Send")
 
+        # Common output section
         with gr.Column():
-            generate_btn = gr.Button("Generate Profile")
+            generate_btn = gr.Button("Generate Profile", variant="primary")
             profile_output = gr.JSON(label="Generated Profile")
             download_btn = gr.File(label="Download Profile")
 
+        # Event handlers
+        async def on_message(message: str, history: List[List[str]], key: str):
+            if not message.strip():
+                return history, None, None
+
+            result = await builder.process_message(message, key)
+            
+            if "error" in result:
+                return history, {"error": result["error"]}, None
+
+            new_history = history + [[message, result["response"]]]
+            return new_history, None, None
+
         async def on_pdf_upload(pdf, key):
             if not pdf:
-                return {"error": "No PDF file uploaded"}
+                return {"error": "No PDF file uploaded"}, None
             
             try:
                 result = await builder.process_pdf(pdf.name, key)
@@ -267,37 +367,24 @@ def create_gradio_interface():
             except Exception as e:
                 return {"error": str(e)}, None
 
-        # Event handlers remain the same for chat functionality
-        async def on_message(message: str, history: List[List[str]], key: str):
-            if not message.strip():
-                return history, None
-
-            result = await builder.process_message(message, key)
-            
-            if "error" in result:
-                return history, {"error": result["error"]}
-
-            history = history + [[message, result["response"]]]
-            return history, None
-
         async def on_generate():
             profile, filename = await builder.generate_profile()
             if "error" in profile:
-                return profile, None
+                return {"error": profile["error"]}, None
             return profile["profile_data"], filename
 
         # Bind events
         msg.submit(
             on_message,
             inputs=[msg, chatbot, api_key],
-            outputs=[chatbot, profile_output]
-        ).then(lambda: "", None, msg)
+            outputs=[chatbot, profile_output, msg]
+        )
         
         send.click(
             on_message,
             inputs=[msg, chatbot, api_key],
-            outputs=[chatbot, profile_output]
-        ).then(lambda: "", None, msg)
+            outputs=[chatbot, profile_output, msg]
+        )
         
         process_pdf_btn.click(
             on_pdf_upload,
@@ -315,4 +402,8 @@ def create_gradio_interface():
 if __name__ == "__main__":
     demo = create_gradio_interface()
     demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True
+    )