Update code to use environment variables and improve error handling

Dockerfile

@@ -10,7 +10,7 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 # Create a non-root user
-RUN useradd -m -u 1000 app_user
+RUN useradd -m -u 1000 chainlit_user
 
 # Copy requirements first for better caching
 COPY requirements.txt /code/requirements.txt
@@ -21,16 +21,17 @@ COPY . /code/
 
 # Create necessary directories and set permissions
 RUN mkdir -p /code/.files && \
-    chown -R app_user:app_user /code && \
+    chown -R chainlit_user:chainlit_user /code && \
     chmod -R 755 /code
 
 # Set environment variables
 ENV HOST=0.0.0.0
 ENV PORT=7860
 ENV PYTHONPATH=/code
+ENV OPENAI_API_KEY=${OPENAI_API_KEY}
 
 # Switch to non-root user
-USER app_user
+USER chainlit_user
 
 # Command to run the application
-CMD python app.py 
+CMD chainlit run app.py --host $HOST --port $PORT 

agent_graph.py

@@ -18,7 +18,7 @@ if not openai_api_key:
 
 # Initialize the LLM
 llm = ChatOpenAI(
-    model="gpt-4",
+    model="gpt-3.5-turbo",
     temperature=0,
     openai_api_key=openai_api_key
 )

app.py

@@ -1,39 +1,77 @@
-import gradio as gr
-from test_agent import process_query
-import logging
 import os
-from dotenv import load_dotenv
+import chainlit as cl
+from agent_graph import agent_node
+from dotenv import load_dotenv  # Import dotenv
+from typing import List, Dict
 
-# Load environment variables
+# Load environment variables from .env file
 load_dotenv()
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Ensure your OpenAI API key is set up in environment variables
+openai_api_key = os.getenv("OPENAI_API_KEY")
 
-def respond(message, history):
-    """Process the user's message and return the response."""
+if not openai_api_key:
+    raise ValueError("OpenAI API key is missing in the .env file")
+
+# Store chat history
+chat_histories: Dict[str, List[Dict[str, str]]] = {}
+
+@cl.on_chat_start
+async def start_chat():
+    # Initialize empty chat history for this session
+    chat_histories[cl.user_session.get("id")] = []
+    
+    welcome_message = """👋 Welcome to the Stock Price Calculator!
+
+I can help you with:
+• Getting real-time stock prices
+• Calculating how many shares you can buy
+
+Try these examples:
+• Type 'AAPL' to get Apple's stock price
+• Ask 'How many MSFT shares can I buy with $10000?'
+
+What would you like to know?"""
+    
+    await cl.Message(content=welcome_message).send()
+
+@cl.on_message
+async def handle_message(message: cl.Message):
     try:
-        logger.info(f"Received message: {message}")
-        response = process_query(message)
-        logger.info(f"Sending response: {response}")
-        return response
+        # Get chat history for this session
+        session_id = cl.user_session.get("id")
+        history = chat_histories.get(session_id, [])
+        
+        # Add current message to history
+        history.append({"role": "user", "content": message.content})
+        
+        # Create state dictionary with history
+        state = {
+            "input": message.content,
+            "chat_history": history
+        }
+        
+        # Process the message with agent_node
+        response = agent_node(state)
+        
+        # Send the response back to the user
+        if isinstance(response, dict) and "output" in response:
+            # Add response to history
+            history.append({"role": "assistant", "content": response["output"]})
+            await cl.Message(content=response["output"]).send()
+        else:
+            await cl.Message(content="Received an invalid response format from the agent.").send()
+            
+        # Update history in storage
+        chat_histories[session_id] = history
+            
     except Exception as e:
-        error_msg = f"Error processing message: {str(e)}"
-        logger.error(error_msg)
-        return error_msg
-
-# Create the Gradio interface
-demo = gr.ChatInterface(
-    fn=respond,
-    title="AI Stock Agent",
-    description="Ask me about stock prices and buying power calculations!",
-    examples=[
-        "What is the current price of AAPL?",
-        "How many shares of MSFT can I buy with $5000?",
-    ],
-)
-
-# Launch the interface
-if __name__ == "__main__":
-    demo.launch()
+        print(f"Error occurred: {e}")
+        await cl.Message(content="Sorry, something went wrong while processing your request.").send()
+
+@cl.on_chat_end
+async def end_chat():
+    # Clean up chat history when session ends
+    session_id = cl.user_session.get("id")
+    if session_id in chat_histories:
+        del chat_histories[session_id]

chainlit.md

@@ -0,0 +1,55 @@
+# Stock Price & Share Calculator 📈
+
+Welcome! This app helps you:
+1. Check real-time stock prices
+2. Calculate how many shares you can buy with a specific amount
+
+## Quick Start Guide 🚀
+
+### 1️⃣ Get Stock Prices
+Just type a ticker symbol or company name:
+```
+AAPL
+```
+or
+```
+What's the price of Apple?
+```
+
+### 2️⃣ Calculate Shares
+Ask how many shares you can afford:
+```
+How many AAPL shares can I buy with $10000?
+```
+
+## Popular Stocks You Can Try 💎
+
+### Tech Companies
+- `AAPL` - Apple
+- `MSFT` - Microsoft
+- `GOOGL` - Google
+- `AMZN` - Amazon
+- `META` - Meta/Facebook
+- `TSLA` - Tesla
+- `NVDA` - NVIDIA
+
+### Financial
+- `JPM` - JPMorgan
+- `BAC` - Bank of America
+- `GS` - Goldman Sachs
+
+### Retail
+- `WMT` - Walmart
+- `COST` - Costco
+- `SBUX` - Starbucks
+
+## Tips 💡
+- You can use either the ticker (e.g., `AAPL`) or company name (e.g., "Apple")
+- Include a dollar amount to calculate shares (e.g., "AAPL $5000")
+- All prices are real-time from the market
+
+Try it now! Type a ticker symbol or ask about any company listed above.
+
+## Welcome screen
+
+To modify the welcome screen, edit the `chainlit.md` file at the root of your project. If you do not want a welcome screen, just leave this file empty.

requirements.txt

@@ -1,6 +1,7 @@
-langchain==0.0.267
-openai==0.28.0
-python-dotenv==1.0.0
-yfinance==0.2.18
-gradio==3.50.2
-requests==2.31.0
+chainlit
+langchain-openai>=0.0.3
+langchain-core>=0.1.4
+openai>=1.7.0
+python-dotenv>=1.0.0
+yfinance>=0.2.36
+typing-extensions>=4.9.0

test_agent.py

@@ -3,26 +3,16 @@ from langchain.agents import initialize_agent, Tool, AgentType
 from tools import get_price, buying_power_tool
 from dotenv import load_dotenv
 import os
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
 # Load environment variables from .env file
 load_dotenv()
 
 # Initialize the language model with API key from environment variables
-try:
-    llm = ChatOpenAI(
-        temperature=0,
-        model_name="gpt-3.5-turbo",
-        openai_api_key=os.getenv("OPENAI_API_KEY")
-    )
-    logger.info("Successfully initialized ChatOpenAI")
-except Exception as e:
-    logger.error(f"Failed to initialize ChatOpenAI: {str(e)}")
-    raise
+llm = ChatOpenAI(
+    temperature=0,
+    model_name="gpt-3.5-turbo",
+    openai_api_key=os.getenv("OPENAI_API_KEY")
+)
 
 # Define tools
 tools = [
@@ -39,33 +29,13 @@ tools = [
 ]
 
 # Initialize the agent
-try:
-    agent = initialize_agent(
-        tools,
-        llm,
-        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
-        verbose=True
-    )
-    logger.info("Successfully initialized agent")
-except Exception as e:
-    logger.error(f"Failed to initialize agent: {str(e)}")
-    raise
-
-def process_query(query):
-    """Process a user query and return the response."""
-    try:
-        logger.info(f"Processing query: {query}")
-        response = agent.run(query)
-        logger.info(f"Response: {response}")
-        return response
-    except Exception as e:
-        error_msg = f"Error processing query: {str(e)}"
-        logger.error(error_msg)
-        return error_msg
+agent = initialize_agent(
+    tools,
+    llm,
+    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+    verbose=True
+)
 
-# Run a test if this is the main module
-if __name__ == "__main__":
-    test_query = "How much is AAPL?"
-    logger.info(f"Running test with query: {test_query}")
-    response = process_query(test_query)
-    print(response)
+# Run a test
+response = agent.run("How much is AAPL?")
+print(response)