Update app.py

app.py

@@ -100,62 +100,153 @@ def detect_chart_type(client, image_data):
         st.error(f"Error in chart type detection: {str(e)}")
         return "Other"
 
-def analyze_chart_with_claude(client, image_data, prompt, chart_type=None):
-    """Analyze chart using Claude Vision"""
+def continue_analysis_with_claude(client, question, previous_analysis, image_data=None):
+    """Continue the analysis based on a follow-up question"""
     try:
-        encoded_image = base64.b64encode(image_data).decode('utf-8')
+        content = [
+            {
+                "type": "text",
+                "text": f"""Previous analysis: {previous_analysis}
+
+User's follow-up question: {question}
+
+Please provide a detailed answer to the follow-up question, maintaining the context of the previous analysis."""
+            }
+        ]
         
-        # If chart type wasn't provided, detect it first
-        if not chart_type:
-            chart_type = detect_chart_type(client, image_data)
-            st.info(f"Detected chart type: {chart_type}")
+        # Add image to the content if available
+        if image_data:
+            encoded_image = base64.b64encode(image_data).decode('utf-8')
+            content.append({
+                "type": "image",
+                "source": {
+                    "type": "base64",
+                    "media_type": "image/jpeg",
+                    "data": encoded_image
+                }
+            })
         
         message = client.messages.create(
             model="claude-3-opus-20240229",
             max_tokens=1000,
             messages=[{
                 "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": prompt.format(chart_type=chart_type)
-                    },
-                    {
-                        "type": "image",
-                        "source": {
-                            "type": "base64",
-                            "media_type": "image/jpeg",
-                            "data": encoded_image
-                        }
-                    }
-                ]
+                "content": content
             }]
         )
         
-        return message.content[0].text, chart_type
+        return message.content[0].text
     except Exception as e:
-        st.error(f"Error in Claude analysis: {str(e)}")
-        return None, None
+        st.error(f"Error in follow-up analysis: {str(e)}")
+        return None
+
+def get_trading_education(client, concept):
+    """Get educational content about trading concepts"""
+    try:
+        message = client.messages.create(
+            model="claude-3-opus-20240229",
+            max_tokens=1000,
+            messages=[{
+                "role": "user",
+                "content": f"""Please explain the trading concept '{concept}' in a clear, educational way. Structure your response as follows:
+
+1. Basic Definition
+2. How it Works
+3. Key Characteristics
+4. When to Look for It
+5. Trading Implications
+6. Common Mistakes to Avoid
+7. Real-World Example
 
-def save_chat_history(chat_history, filename=None):
-    """Saves chat history to a JSON file"""
+If relevant, describe what a typical chart pattern for this concept looks like.
+Include any important formulas or calculations if applicable.
+
+Please make this explanation suitable for beginners while also including enough depth for intermediate traders."""
+            }]
+        )
+        
+        return message.content[0].text
+    except Exception as e:
+        st.error(f"Error in getting educational content: {str(e)}")
+        return None
+
+def extract_stock_info(analysis_text):
+    """Extract stock name and other metadata from analysis text"""
+    # This is a simple implementation - can be made more sophisticated
+    stock_name = "Unknown"
+    try:
+        # Look for common stock name patterns
+        if "analyzing" in analysis_text.lower():
+            words = analysis_text.split()
+            for i, word in enumerate(words):
+                if word.lower() == "analyzing":
+                    stock_name = words[i + 1].strip("(),.:")
+    except:
+        pass
+    return stock_name
+
+def save_chat_history(chat_history, image_data=None, filename=None):
+    """Saves chat history and associated image to JSON and image files"""
     if not os.path.exists("chat_histories"):
         os.makedirs("chat_histories")
+    if not os.path.exists("chat_images"):
+        os.makedirs("chat_images")
         
-    if filename is None:
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = f"chat_history_{timestamp}.json"
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     
-    filepath = os.path.join("chat_histories", filename)
+    # Get the stock name from the latest analysis
+    stock_name = "Unknown"
+    if chat_history:
+        latest_analysis = chat_history[-1]['analysis']
+        stock_name = extract_stock_info(latest_analysis)
+    
+    # Create filename with metadata
+    if filename:
+        base_filename = filename
+    else:
+        base_filename = f"{stock_name}_{timestamp}"
+    
+    # Save image if provided
+    image_filename = None
+    if image_data:
+        image_filename = f"{base_filename}.jpg"
+        image_path = os.path.join("chat_images", image_filename)
+        with open(image_path, "wb") as f:
+            f.write(image_data)
+    
+    # Add metadata to chat history
+    chat_data = {
+        'metadata': {
+            'stock_name': stock_name,
+            'date_created': timestamp,
+            'image_file': image_filename
+        },
+        'conversations': chat_history
+    }
+    
+    # Save chat history
+    json_filename = f"{base_filename}.json"
+    filepath = os.path.join("chat_histories", json_filename)
     with open(filepath, "w") as f:
-        json.dump(chat_history, f)
-    return filename
+        json.dump(chat_data, f)
+    
+    return json_filename
 
 def load_chat_history(filename):
-    """Loads chat history from a JSON file"""
+    """Loads chat history and associated image"""
     filepath = os.path.join("chat_histories", filename)
     with open(filepath, "r") as f:
-        return json.load(f)
+        chat_data = json.load(f)
+        
+    # Load associated image if it exists
+    image_data = None
+    if chat_data.get('metadata', {}).get('image_file'):
+        image_path = os.path.join("chat_images", chat_data['metadata']['image_file'])
+        if os.path.exists(image_path):
+            with open(image_path, "rb") as f:
+                image_data = f.read()
+                
+    return chat_data, image_data
 
 def main():
     st.set_page_config(
@@ -170,10 +261,167 @@ def main():
         st.error("Failed to initialize Anthropic client. Please check your API key configuration.")
         return
     
-    # Initialize session state for chat history
+    # Initialize session state
     if 'chat_history' not in st.session_state:
         st.session_state.chat_history = []
+    if 'current_image' not in st.session_state:
+        st.session_state.current_image = None
+    if 'current_analysis' not in st.session_state:
+        st.session_state.current_analysis = None
+    
+    # Tab selection
+    tab1, tab2 = st.tabs(["Chart Analysis", "Learn Trading"])
     
+    with tab1:
+        # Initialize variables
+        uploaded_file = None
+        screenshot_taken = False
+
+        # Sidebar
+        with st.sidebar:
+            st.title("🚀 Chart Analysis AI")
+            upload_option = st.radio(
+                "Choose input method:",
+                ("Upload Image", "Take Screenshot", "Ask Question")
+            )
+            
+            # File uploader
+            if upload_option == "Upload Image":
+                uploaded_file = st.file_uploader("Upload your chart", type=["png", "jpg", "jpeg"])
+                if uploaded_file:
+                    st.session_state.current_image = uploaded_file.getvalue()
+            elif upload_option == "Take Screenshot":
+                if st.button("Take Screenshot", key="screenshot"):
+                    st.info("Feature coming soon! For now, please use the Upload Image option.")
+                    screenshot_taken = False
+            
+            # Analysis Options
+            st.subheader("Analysis Options")
+            patterns = st.multiselect(
+                "Patterns to Look For",
+                ["Double Top/Bottom", "Head and Shoulders", "Triangle", 
+                 "Flag", "Wedge", "Channel", "Support/Resistance"]
+            )
+            
+            indicators = st.multiselect(
+                "Technical Indicators",
+                ["Moving Averages", "RSI", "MACD", "Bollinger Bands", 
+                 "Volume", "Stochastic", "ADX"]
+            )
+
+        # Main content area
+        st.title("📈 Stock Chart Analysis Assistant")
+        
+        # Create two columns for layout
+        col1, col2 = st.columns([2, 1])
+        
+        with col1:
+            if upload_option == "Ask Question":
+                user_question = st.text_input("What would you like to know about your chart?")
+            
+            # Display uploaded image
+            if uploaded_file is not None:
+                st.image(uploaded_file, caption="Uploaded Chart", use_container_width=True)
+            
+            # Continue chat section
+            if st.session_state.current_analysis:
+                st.subheader("Continue Analysis")
+                follow_up = st.text_input("Ask a follow-up question about this chart:")
+                if st.button("Send Follow-up"):
+                    if follow_up:
+                        with st.spinner("Analyzing..."):
+                            follow_up_response = continue_analysis_with_claude(
+                                client,
+                                follow_up,
+                                st.session_state.current_analysis,
+                                st.session_state.current_image
+                            )
+                            if follow_up_response:
+                                st.write(follow_up_response)
+                                # Add to chat history
+                                st.session_state.chat_history.append({
+                                    'timestamp': datetime.now().isoformat(),
+                                    'question': follow_up,
+                                    'analysis': follow_up_response
+                                })
+            
+            if st.button("Analyze"):
+                if upload_option == "Ask Question" and user_question:
+                    st.info("Question-based analysis feature coming soon!")
+                elif uploaded_file is None and not screenshot_taken:
+                    st.warning("Please upload an image or take a screenshot first.")
+                else:
+                    with st.spinner("Analyzing chart..."):
+                        # Generate prompt
+                        prompt = create_prompt_template(patterns, indicators)
+                        
+                        if uploaded_file:
+                            # Process image and get analysis
+                            analysis_result, chart_type = analyze_chart_with_claude(
+                                client,
+                                uploaded_file.getvalue(),
+                                prompt
+                            )
+                            
+                            if analysis_result:
+                                # Store current analysis
+                                st.session_state.current_analysis = analysis_result
+                                
+                                # Add to chat history
+                                st.session_state.chat_history.append({
+                                    'timestamp': datetime.now().isoformat(),
+                                    'chart_type': chart_type,
+                                    'analysis': analysis_result
+                                })
+                                
+                                # Display analysis
+                                st.subheader("Analysis Results")
+                                st.write(analysis_result)
+                                
+                                # Risk warning
+                                st.warning(
+                                    "⚠️ This analysis is AI-generated and for informational purposes only. "
+                                    "Do not make trading decisions solely based on this information."
+                                )
+        
+        with col2:
+            st.subheader("Chat History")
+            
+            # Display chat history
+            for chat in st.session_state.chat_history:
+                timestamp = datetime.fromisoformat(chat['timestamp']).strftime("%Y-%m-%d %H:%M")
+                with st.expander(f"Analysis from {timestamp}"):
+                    st.write(chat['analysis'])
+                    if 'question' in chat:
+                        st.write(f"Follow-up: {chat['question']}")
+            
+            # Save chat options
+            save_name = st.text_input("Save chat as (optional):")
+            if st.button("Save Chat"):
+                if st.session_state.chat_history:
+                    filename = save_chat_history(
+                        st.session_state.chat_history,
+                        st.session_state.current_image,
+                        f"{save_name}.json" if save_name else None
+                    )
+                    st.success(f"Chat saved as {filename}")
+                else:
+                    st.warning("No chat history to save.")
+    
+    with tab2:
+        st.title("📚 Learn Trading")
+        
+        # Search or select trading concept
+        concept = st.text_input("Enter a trading concept you'd like to learn about (e.g., 'evening star pattern', 'RSI', 'MACD'):")
+        if st.button("Learn"):
+            if concept:
+                with st.spinner("Getting educational content..."):
+                    education_content = get_trading_education(client, concept)
+                    if education_content:
+                        st.markdown(education_content)
+
+if __name__ == "__main__":
+    main()
     # Initialize variables
     uploaded_file = None
     screenshot_taken = False