Akshayram1's picture
Update app.py
a9fba1e verified
import streamlit as st
import yfinance as yf
from crewai import Agent, Task, Crew, Process, LLM
from crewai_tools import CodeInterpreterTool
from pydantic import BaseModel, Field
from dotenv import load_dotenv
import os
# Load environment variables
load_dotenv()
# Streamlit UI
st.title("AI-Powered Stock Analysis")
# Get Sambanova API Key from User
sambanova_key = st.text_input("Enter your Sambanova API Key", type="password")
query = st.text_input("Enter a stock query (e.g., 'Plot YTD stock gain of Tesla')")
# Define Query Output Model
class QueryAnalysisOutput(BaseModel):
symbol: str = Field(..., description="Stock ticker symbol")
timeframe: str = Field(..., description="Time period (e.g., '1d', '1mo', '1y')")
action: str = Field(..., description="Action to be performed (e.g., 'fetch', 'plot')")
# Define LLM Model
if sambanova_key:
llm = LLM(model="sambanova/DeepSeek-R1-Distill-Llama-70B", temperature=0.7, api_key=sambanova_key)
# Define CrewAI Agents with backstory
query_parser_agent = Agent(
role="Stock Data Analyst",
goal="Extract stock details from user query.",
backstory="An expert in analyzing stock data and trends.",
llm=llm,
verbose=True,
memory=True,
)
query_parsing_task = Task(
description="Analyze the user query and extract stock details.",
output_pydantic=QueryAnalysisOutput,
agent=query_parser_agent,
)
code_writer_agent = Agent(
role="Senior Python Developer",
goal="Write Python code to visualize stock data.",
backstory="A seasoned developer with expertise in financial data visualization.",
llm=llm,
verbose=True,
)
code_writer_task = Task(
description="Generate Python code for stock visualization.",
agent=code_writer_agent,
)
code_interpreter_tool = CodeInterpreterTool()
code_execution_agent = Agent(
role="Code Execution Expert",
goal="Execute the generated code to visualize stock data.",
backstory="Specializes in running and debugging Python scripts for data analysis.",
tools=[code_interpreter_tool],
allow_code_execution=True,
llm=llm,
verbose=True,
)
code_execution_task = Task(
description="Execute the generated Python script.",
agent=code_execution_agent,
)
# Create Crew
crew = Crew(
agents=[query_parser_agent, code_writer_agent, code_execution_agent],
tasks=[query_parsing_task, code_writer_task, code_execution_task],
process=Process.sequential,
)
if st.button("Analyze Stock"):
result = crew.kickoff(inputs={"query": query})
st.write("Results:", result)
else:
st.warning("Please enter your Sambanova API Key to proceed.")