Update code to use environment variables and improve error handling

app.py

@@ -3,6 +3,9 @@ import chainlit as cl
 from agent_graph import agent_node
 from dotenv import load_dotenv  # Import dotenv
 from typing import List, Dict
+import gradio as gr
+from test_agent import process_query
+import logging
 
 # Load environment variables from .env file
 load_dotenv()
@@ -16,6 +19,10 @@ if not openai_api_key:
 # Store chat history
 chat_histories: Dict[str, List[Dict[str, str]]] = {}
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 @cl.on_chat_start
 async def start_chat():
     # Initialize empty chat history for this session
@@ -75,3 +82,30 @@ async def end_chat():
     session_id = cl.user_session.get("id")
     if session_id in chat_histories:
         del chat_histories[session_id]
+
+def respond(message, history):
+    """Process the user's message and return the response."""
+    try:
+        logger.info(f"Received message: {message}")
+        response = process_query(message)
+        logger.info(f"Sending response: {response}")
+        return response
+    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()

requirements.txt

@@ -1,7 +1,5 @@
-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
+langchain==0.0.267
+openai==0.28.0
+python-dotenv==1.0.0
+yfinance==0.2.18
+gradio==3.50.2

test_agent.py

@@ -3,16 +3,26 @@ 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
-llm = ChatOpenAI(
-    temperature=0,
-    model_name="gpt-3.5-turbo",
-    openai_api_key=os.getenv("OPENAI_API_KEY")
-)
+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
 
 # Define tools
 tools = [
@@ -29,13 +39,33 @@ tools = [
 ]
 
 # Initialize the agent
-agent = initialize_agent(
-    tools,
-    llm,
-    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
-    verbose=True
-)
+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
 
-# Run a test
-response = agent.run("How much is AAPL?")
-print(response)
+# 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)