Update app.py

app.py

@@ -1,44 +1,31 @@
-from flask import Flask, request, jsonify, render_template, session
+import gradio as gr
 import uuid
 import logging
 from datetime import datetime
 import os
-import threading
-import time
 
 from src.graphs.finalAgentGraph import sparrowAgent
 from langchain_core.messages import HumanMessage, AIMessage
 
-
-app = Flask(__name__)
-app.secret_key = os.environ.get('FLASK_SECRET_KEY', 'your-secret-key-here')
-
-
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-conversations = {}
-conversations_lock = threading.Lock()
-
-
 def ensure_langchain_message(message):
     """Ensure a message is a proper LangChain message object"""
     if isinstance(message, (HumanMessage, AIMessage)):
         return message
     elif isinstance(message, dict):
-        
         content = message.get('content', str(message))
-        message_type = message.get('type', 'ai')  
+        message_type = message.get('type', 'ai')
         if message_type == 'human':
             return HumanMessage(content=content)
         else:
             return AIMessage(content=content)
     elif isinstance(message, str):
-        
         return AIMessage(content=message)
     else:
-        
         return AIMessage(content=str(message))
 
 
@@ -51,235 +38,239 @@ def clean_messages_list(messages):
     return cleaned_messages
 
 
-@app.route('/')
-def index():
-    """Serve the main chat interface"""
-    return render_template('index.html')
+def initialize_conversation():
+    """Initialize a new conversation state"""
+    return {
+        'thread_id': str(uuid.uuid4()),
+        'messages': [],
+        'notes': [],
+        'query_brief': '',
+        'final_message': '',
+        'created_at': datetime.now(),
+        'last_updated': datetime.now()
+    }
 
 
-@app.route('/chat', methods=['POST'])
-def chat():
-    """Handle chat messages"""
+def process_message(user_message, history, conversation_state):
+    """
+    Process user message and return response
+    
+    Args:
+        user_message: The user's input message
+        history: Gradio chat history (list of [user_msg, bot_msg] pairs)
+        conversation_state: Dictionary containing conversation context
+    
+    Returns:
+        Tuple of (empty string, updated history, updated conversation state, status message)
+    """
     try:
-        data = request.get_json()
-        user_message = data.get('message', '').strip()
-
-        if not user_message:
-            return jsonify({'success': False, 'error': 'Empty message'})
-
+        if not user_message or not user_message.strip():
+            return "", history, conversation_state, "Please enter a message"
         
-        thread_id = session.get('thread_id')
-        if not thread_id:
-            thread_id = str(uuid.uuid4())
-            session['thread_id'] = thread_id
-
+        # Initialize conversation state if None
+        if conversation_state is None:
+            conversation_state = initialize_conversation()
         
-        with conversations_lock:
-            if thread_id not in conversations:
-                conversations[thread_id] = {
-                    'messages': [],
-                    'notes': [],
-                    'query_brief': '',
-                    'final_message': '',
-                    'created_at': datetime.now(),
-                    'last_updated': datetime.now()
-                }
-            conversation = conversations[thread_id]
-
+        thread_id = conversation_state['thread_id']
         
+        # Add user message to conversation
         human_message = HumanMessage(content=user_message)
-        conversation['messages'].append(human_message)
-        conversation['last_updated'] = datetime.now()
-
+        conversation_state['messages'].append(human_message)
+        conversation_state['last_updated'] = datetime.now()
         
-        cleaned_messages = clean_messages_list(conversation['messages'])
-
+        # Clean messages
+        cleaned_messages = clean_messages_list(conversation_state['messages'])
         
+        # Prepare input for sparrow agent
         sparrow_input = {
             'messages': cleaned_messages,
-            'notes': conversation.get('notes', []),
-            'query_brief': conversation.get('query_brief', ''),
-            'final_message': conversation.get('final_message', '')
+            'notes': conversation_state.get('notes', []),
+            'query_brief': conversation_state.get('query_brief', ''),
+            'final_message': conversation_state.get('final_message', '')
         }
-
+        
         logger.info(f"[{thread_id}] Processing message: {user_message[:100]}")
         logger.info(f"[{thread_id}] Input messages count: {len(cleaned_messages)}")
-
         
+        # Invoke the sparrow agent
         result = sparrowAgent.invoke(sparrow_input)
-
         
+        # Extract response message
         response_message = ""
         ai_message = None
-
+        
         if result.get('final_message'):
             response_message = result['final_message']
             ai_message = AIMessage(content=response_message)
         else:
-            
             result_messages = clean_messages_list(result.get('messages', []))
             
-            
+            # Find last user message index
             last_user_index = -1
             for i, msg in enumerate(result_messages):
                 if isinstance(msg, HumanMessage):
                     last_user_index = i
             
-            
+            # Get first AI message after last user message
             for i in range(last_user_index + 1, len(result_messages)):
                 msg = result_messages[i]
                 if isinstance(msg, AIMessage) and msg.content and msg.content.strip():
                     response_message = msg.content
                     ai_message = msg
                     break
-
+        
         if not response_message:
             response_message = "I'm processing your request. Could you provide more details?"
             ai_message = AIMessage(content=response_message)
-
         
-        status_info = ""
+        # Update conversation state
+        if result.get('messages'):
+            conversation_state['messages'] = clean_messages_list(result['messages'])
+        else:
+            conversation_state['messages'].append(ai_message)
+        
+        # Remove consecutive duplicates
+        cleaned_conversation_messages = []
+        prev_content = None
+        prev_type = None
+        
+        for msg in conversation_state['messages']:
+            current_content = msg.content if hasattr(msg, 'content') else str(msg)
+            current_type = type(msg).__name__
+            
+            if current_content != prev_content or current_type != prev_type:
+                cleaned_conversation_messages.append(msg)
+                prev_content = current_content
+                prev_type = current_type
+        
+        conversation_state['messages'] = cleaned_conversation_messages
+        conversation_state['notes'] = result.get('notes', conversation_state.get('notes', []))
+        conversation_state['query_brief'] = result.get('query_brief', conversation_state.get('query_brief', ''))
+        conversation_state['final_message'] = result.get('final_message', conversation_state.get('final_message', ''))
+        conversation_state['last_updated'] = datetime.now()
+        
+        # Update Gradio chat history
+        history.append([user_message, response_message])
+        
+        # Create status message
+        status_info = f"Thread: {thread_id[:8]}... | Messages: {len(conversation_state['messages'])}"
         if result.get('execution_jobs'):
-            status_info = f"Executed: {', '.join(result['execution_jobs'])}"
+            status_info += f" | Executed: {', '.join(result['execution_jobs'])}"
         elif result.get('notes') and isinstance(result['notes'], list) and result['notes']:
-            status_info = str(result['notes'][-1])
-
+            status_info += f" | Note: {str(result['notes'][-1])[:50]}"
         
-        with conversations_lock:
-            
-            if result.get('messages'):
-                conversation['messages'] = clean_messages_list(result['messages'])
-            else:
-                
-                conversation['messages'].append(ai_message)
-            
-            
-            cleaned_conversation_messages = []
-            prev_content = None
-            prev_type = None
-            
-            for msg in conversation['messages']:
-                current_content = msg.content if hasattr(msg, 'content') else str(msg)
-                current_type = type(msg).__name__
-                
-                
-                if current_content != prev_content or current_type != prev_type:
-                    cleaned_conversation_messages.append(msg)
-                    prev_content = current_content
-                    prev_type = current_type
-            
-            conversation['messages'] = cleaned_conversation_messages
-            
-            
-            conversation['notes'] = result.get('notes', conversation.get('notes', []))
-            conversation['query_brief'] = result.get('query_brief', conversation.get('query_brief', ''))
-            conversation['final_message'] = result.get('final_message', conversation.get('final_message', ''))
-            conversation['last_updated'] = datetime.now()
-            conversations[thread_id] = conversation
-
         logger.info(f"[{thread_id}] Response generated: {response_message[:100]}")
-        logger.info(f"[{thread_id}] Final messages count: {len(conversation['messages'])}")
-
+        logger.info(f"[{thread_id}] Final messages count: {len(conversation_state['messages'])}")
         
-        message_types = [f"{type(msg).__name__}: {msg.content[:50] if hasattr(msg, 'content') else str(msg)[:50]}" 
-                        for msg in conversation['messages'][-3:]]  # Show last 3 messages
-        logger.info(f"[{thread_id}] Recent message types: {message_types}")
+        return "", history, conversation_state, status_info
+    
+    except Exception as e:
+        logger.error(f"Error processing message: {str(e)}", exc_info=True)
+        error_msg = f"An error occurred: {str(e)}"
+        history.append([user_message, error_msg])
+        return "", history, conversation_state, f"Error: {str(e)}"
 
-        return jsonify({
-            'success': True,
-            'response': response_message,
-            'status': status_info,
-            'thread_id': thread_id
-        })
 
-    except Exception as e:
-        logger.error(f"Error in chat endpoint: {str(e)}", exc_info=True)
-        return jsonify({'success': False, 'error': f"An error occurred: {str(e)}"})
+def clear_conversation():
+    """Clear conversation and start fresh"""
+    new_state = initialize_conversation()
+    logger.info(f"[{new_state['thread_id']}] New conversation started")
+    return [], new_state, f"New conversation started (ID: {new_state['thread_id'][:8]}...)"
 
 
-@app.route('/new_conversation', methods=['POST'])
-def new_conversation():
-    """Start a new conversation thread"""
-    old_thread_id = session.get('thread_id')
-    session.pop('thread_id', None)
+def get_conversation_info(conversation_state):
+    """Get current conversation information"""
+    if conversation_state is None:
+        return "No active conversation"
     
-    if old_thread_id:
-        logger.info(f"[{old_thread_id}] Starting new conversation")
+    info_lines = [
+        f"**Thread ID:** {conversation_state['thread_id']}",
+        f"**Messages:** {len(conversation_state.get('messages', []))}",
+        f"**Notes:** {len(conversation_state.get('notes', []))}",
+        f"**Has Query Brief:** {bool(conversation_state.get('query_brief'))}",
+        f"**Has Final Message:** {bool(conversation_state.get('final_message'))}",
+        f"**Created:** {conversation_state.get('created_at', 'N/A')}",
+        f"**Last Updated:** {conversation_state.get('last_updated', 'N/A')}"
+    ]
     
-    return jsonify({'success': True, 'message': 'New conversation started'})
+    return "\n\n".join(info_lines)
 
 
-@app.route('/conversation_info', methods=['GET'])
-def conversation_info():
-    """Get current conversation information (for debugging)"""
-    thread_id = session.get('thread_id')
-    if not thread_id:
-        return jsonify({'error': 'No active conversation'})
+# Create Gradio interface
+with gr.Blocks(title="Sparrow Agent Chat", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🦜 Sparrow Agent Chat")
+    gr.Markdown("Interact with the Sparrow AI Agent. Ask questions and get intelligent responses!")
     
-    with conversations_lock:
-        conversation = conversations.get(thread_id, {})
-        if not conversation:
-            return jsonify({'error': 'Conversation not found'})
-        
+    # State to store conversation context
+    conversation_state = gr.State(initialize_conversation())
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(
+                label="Conversation",
+                height=500,
+                show_copy_button=True
+            )
+            
+            with gr.Row():
+                msg = gr.Textbox(
+                    label="Your Message",
+                    placeholder="Type your message here...",
+                    lines=2,
+                    scale=4
+                )
+                submit_btn = gr.Button("Send", variant="primary", scale=1)
+            
+            with gr.Row():
+                clear_btn = gr.Button("New Conversation", variant="secondary")
+                status_box = gr.Textbox(
+                    label="Status",
+                    interactive=False,
+                    lines=1
+                )
         
-        info = {
-            'thread_id': thread_id,
-            'message_count': len(conversation.get('messages', [])),
-            'message_types': [type(msg).__name__ for msg in conversation.get('messages', [])],
-            'notes_count': len(conversation.get('notes', [])),
-            'has_query_brief': bool(conversation.get('query_brief')),
-            'has_final_message': bool(conversation.get('final_message')),
-            'created_at': conversation.get('created_at', '').isoformat() if conversation.get('created_at') else '',
-            'last_updated': conversation.get('last_updated', '').isoformat() if conversation.get('last_updated') else ''
-        }
+        with gr.Column(scale=1):
+            gr.Markdown("### Debug Info")
+            info_btn = gr.Button("Show Conversation Info")
+            info_display = gr.Markdown("Click button to show info")
     
-    return jsonify(info)
-
-
-@app.route('/health')
-def health():
-    """Health check endpoint"""
-    with conversations_lock:
-        active_conversations = len(conversations)
-    return jsonify({
-        'status': 'healthy',
-        'timestamp': datetime.now().isoformat(),
-        'active_conversations': active_conversations
-    })
-
-
-@app.errorhandler(404)
-def not_found(error):
-    return jsonify({'error': 'Endpoint not found'}), 404
-
-
-@app.errorhandler(500)
-def internal_error(error):
-    logger.error(f"Internal server error: {error}")
-    return jsonify({'error': 'Internal server error'}), 500
-
-
-def cleanup_conversations():
-    """Remove old conversations older than 24 hours"""
-    while True:
-        time.sleep(3600)  # every hour
-        cutoff = datetime.now().timestamp() - 24 * 3600
-        removed_count = 0
-        with conversations_lock:
-            threads_to_remove = [tid for tid, conv in conversations.items()
-                                 if conv['last_updated'].timestamp() < cutoff]
-            for tid in threads_to_remove:
-                del conversations[tid]
-                removed_count += 1
-        if removed_count > 0:
-            logger.info(f"Cleaned up {removed_count} old conversations")
-
-
-if __name__ == '__main__':
-    cleanup_thread = threading.Thread(target=cleanup_conversations, daemon=True)
-    cleanup_thread.start()
-
-    port = int(os.environ.get('PORT', 5000))
-    debug = os.environ.get('FLASK_DEBUG', 'False').lower() == 'true'
-    logger.info(f"Starting Sparrow Agent Flask app on port {port}")
-    app.run(host='0.0.0.0', port=port, debug=debug)
+    # Event handlers
+    submit_btn.click(
+        fn=process_message,
+        inputs=[msg, chatbot, conversation_state],
+        outputs=[msg, chatbot, conversation_state, status_box]
+    )
+    
+    msg.submit(
+        fn=process_message,
+        inputs=[msg, chatbot, conversation_state],
+        outputs=[msg, chatbot, conversation_state, status_box]
+    )
+    
+    clear_btn.click(
+        fn=clear_conversation,
+        inputs=[],
+        outputs=[chatbot, conversation_state, status_box]
+    )
+    
+    info_btn.click(
+        fn=get_conversation_info,
+        inputs=[conversation_state],
+        outputs=[info_display]
+    )
+    
+    # Initialize status on load
+    demo.load(
+        fn=lambda state: f"Ready | Thread: {state['thread_id'][:8]}...",
+        inputs=[conversation_state],
+        outputs=[status_box]
+    )
+
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.environ.get('PORT', 7860)),
+        share=False
+    )