Update code to use environment variables and improve error handling

Dockerfile

@@ -9,9 +9,6 @@ RUN apt-get update && apt-get install -y \
     software-properties-common \
     && rm -rf /var/lib/apt/lists/*
 
-# Create a non-root user
-RUN useradd -m -u 1000 chainlit_user
-
 # Copy requirements first for better caching
 COPY requirements.txt /code/requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
@@ -19,25 +16,10 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Copy the rest of the application
 COPY . /code/
 
-# Create necessary directories and set permissions
-RUN mkdir -p /code/.files && \
-    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}
-ENV PYTHONUNBUFFERED=1
-
-# Create log directory and set permissions
-RUN mkdir -p /code/logs && \
-    chown -R chainlit_user:chainlit_user /code/logs && \
-    chmod -R 755 /code/logs
-
-# Switch to non-root user
-USER chainlit_user
 
 # Command to run the application
 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-3.5-turbo",
+    model="gpt-4",
     temperature=0,
     openai_api_key=openai_api_key
 )

app.py

@@ -3,11 +3,6 @@ import chainlit as cl
 from agent_graph import agent_node
 from dotenv import load_dotenv  # Import dotenv
 from typing import List, Dict
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
 # Load environment variables from .env file
 load_dotenv()
@@ -16,7 +11,6 @@ load_dotenv()
 openai_api_key = os.getenv("OPENAI_API_KEY")
 
 if not openai_api_key:
-    logger.error("OpenAI API key is missing in the .env file")
     raise ValueError("OpenAI API key is missing in the .env file")
 
 # Store chat history
@@ -24,12 +18,10 @@ chat_histories: Dict[str, List[Dict[str, str]]] = {}
 
 @cl.on_chat_start
 async def start_chat():
-    try:
-        # Initialize empty chat history for this session
-        session_id = cl.user_session.get("id")
-        chat_histories[session_id] = []
-        
-        welcome_message = """👋 Welcome to the Stock Price Calculator!
+    # 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
@@ -40,29 +32,15 @@ Try these examples:
 • Ask 'How many MSFT shares can I buy with $10000?'
 
 What would you like to know?"""
-        
-        # Send welcome message
-        await cl.Message(
-            content=welcome_message,
-            author="Assistant",
-            type="text"
-        ).send()
-        
-        logger.info("Chat session started successfully")
-    except Exception as e:
-        logger.error(f"Error in start_chat: {str(e)}")
-        raise
+    
+    await cl.Message(content=welcome_message).send()
 
 @cl.on_message
 async def handle_message(message: cl.Message):
     try:
-        logger.info(f"Received message: {message.content}")
-        
         # Get chat history for this session
         session_id = cl.user_session.get("id")
-        if session_id not in chat_histories:
-            chat_histories[session_id] = []
-        history = chat_histories[session_id]
+        history = chat_histories.get(session_id, [])
         
         # Add current message to history
         history.append({"role": "user", "content": message.content})
@@ -75,78 +53,25 @@ async def handle_message(message: cl.Message):
         
         # Process the message with agent_node
         response = agent_node(state)
-        logger.info(f"Agent response: {response}")
         
         # Send the response back to the user
         if isinstance(response, dict) and "output" in response:
-            # Format the response for better display
-            formatted_response = response["output"]
-            
             # Add response to history
-            history.append({"role": "assistant", "content": formatted_response})
-            
-            # Check if the response contains ticker suggestions
-            if "Did you mean one of these?" in formatted_response:
-                # Split the response into main message and suggestions
-                main_msg, suggestions = formatted_response.split("Did you mean one of these?")
-                
-                # Send the main message
-                await cl.Message(
-                    content=main_msg.strip(),
-                    author="Assistant",
-                    type="text"
-                ).send()
-                
-                # Create a list of suggestions
-                suggestion_list = suggestions.strip().split('\n')
-                elements = []
-                for suggestion in suggestion_list:
-                    if suggestion.strip():
-                        elements.append(cl.Text(name="Suggestion", content=suggestion.strip()))
-                
-                # Send the suggestions
-                if elements:
-                    await cl.Message(
-                        content="",
-                        elements=elements,
-                        author="Assistant",
-                        type="text"
-                    ).send()
-            else:
-                # Send regular response
-                await cl.Message(
-                    content=formatted_response,
-                    author="Assistant",
-                    type="text"
-                ).send()
+            history.append({"role": "assistant", "content": response["output"]})
+            await cl.Message(content=response["output"]).send()
         else:
-            error_msg = "Received an invalid response format from the agent."
-            logger.error(f"{error_msg} Response: {response}")
-            await cl.Message(
-                content=error_msg,
-                author="Assistant",
-                type="text"
-            ).send()
+            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 occurred: {str(e)}"
-        logger.error(error_msg)
-        await cl.Message(
-            content="Sorry, something went wrong while processing your request. Please try again.",
-            author="Assistant",
-            type="text"
-        ).send()
+        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():
-    try:
-        # 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]
-        logger.info("Chat session ended successfully")
-    except Exception as e:
-        logger.error(f"Error in end_chat: {str(e)}")
+    # 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]

requirements.txt

@@ -1,9 +1,7 @@
-chainlit>=0.6.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
-pandas>=2.0.0
-numpy>=1.24.0
-requests>=2.31.0
+typing-extensions>=4.9.0

test_agent.py

@@ -1,18 +1,10 @@
 from langchain.chat_models import ChatOpenAI
 from langchain.agents import initialize_agent, Tool, AgentType
 from tools import get_price, buying_power_tool
-from dotenv import load_dotenv
-import os
-
-# Load environment variables from .env file
-load_dotenv()
 
-# Initialize the language model with API key from environment variables
-llm = ChatOpenAI(
-    temperature=0,
-    model_name="gpt-3.5-turbo",
-    openai_api_key=os.getenv("OPENAI_API_KEY")
-)
+# Make sure OpenAI key is set in your environment
+import os
+os.environ["OPENAI_API_KEY"] = "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"
 
 # Define tools
 tools = [
@@ -23,18 +15,14 @@ tools = [
     ),
     Tool(
         name="BuyingPowerCalculator",
-        func=buying_power_tool,
-        description="Calculate how many shares you can buy. Input format: 'TICKER,AMOUNT' (e.g., 'AAPL,5000')"
+        func=lambda q: buying_power_tool(q.split(',')[0], float(q.split(',')[1])),
+        description="Calculate how many shares you can buy. Input format: 'AAPL,5000'"
     )
 ]
 
-# Initialize the agent
-agent = initialize_agent(
-    tools,
-    llm,
-    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
-    verbose=True
-)
+# Create agent
+llm = ChatOpenAI(model="gpt-4", temperature=0)
+agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
 
 # Run a test
 response = agent.run("How much is AAPL?")

tools.py

@@ -10,13 +10,7 @@ def get_price(ticker):
         stock = yf.Ticker(ticker)
         
         # Get the stock info
-        try:
-            info = stock.info
-        except Exception as e:
-            if "404" in str(e):
-                return f"'{ticker}' is not a valid stock ticker. Please check the symbol and try again."
-            return f"Error fetching data for {ticker}: {str(e)}"
-            
+        info = stock.info
         if 'regularMarketPrice' not in info or info['regularMarketPrice'] is None:
             return f"Could not get the current price for {ticker}. Please verify the ticker symbol."
             
@@ -47,13 +41,8 @@ def buying_power_tool(query):
             
         # Get stock price
         stock = yf.Ticker(ticker)
-        try:
-            info = stock.info
-        except Exception as e:
-            if "404" in str(e):
-                return f"'{ticker}' is not a valid stock ticker. Please check the symbol and try again."
-            return f"Error fetching data for {ticker}: {str(e)}"
-            
+        info = stock.info
+        
         if 'regularMarketPrice' not in info or info['regularMarketPrice'] is None:
             return f"Could not get the current price for {ticker}. Please verify the ticker symbol."