mirror of
https://github.com/mendableai/firecrawl.git
synced 2024-11-15 19:22:19 +08:00
Create claude_stock_analyzer.py
This commit is contained in:
Eric Ciarla
parent
c96b36d045
commit
983f344fa8
180
examples/claude_stock_analyzer/claude_stock_analyzer.py
Normal file
180
examples/claude_stock_analyzer/claude_stock_analyzer.py
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
import os
|
||||
from firecrawl import FirecrawlApp
|
||||
import json
|
||||
from dotenv import load_dotenv
|
||||
import anthropic
|
||||
from e2b_code_interpreter import Sandbox
|
||||
import base64
|
||||
|
||||
# ANSI color codes
|
||||
class Colors:
|
||||
CYAN = '\033[96m'
|
||||
YELLOW = '\033[93m'
|
||||
GREEN = '\033[92m'
|
||||
RED = '\033[91m'
|
||||
MAGENTA = '\033[95m'
|
||||
BLUE = '\033[94m'
|
||||
RESET = '\033[0m'
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Retrieve API keys from environment variables
|
||||
firecrawl_api_key = os.getenv("FIRECRAWL_API_KEY")
|
||||
anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
|
||||
e2b_api_key = os.getenv("E2B_API_KEY")
|
||||
|
||||
# Initialize the FirecrawlApp and Anthropic client
|
||||
app = FirecrawlApp(api_key=firecrawl_api_key)
|
||||
client = anthropic.Anthropic(api_key=anthropic_api_key)
|
||||
sandbox = Sandbox(api_key=e2b_api_key)
|
||||
|
||||
# Find the relevant stock pages via map
|
||||
def find_relevant_page_via_map(stock_search_term, url, app):
|
||||
try:
|
||||
print(f"{Colors.CYAN}Searching for stock: {stock_search_term}{Colors.RESET}")
|
||||
print(f"{Colors.CYAN}Initiating search on the website: {url}{Colors.RESET}")
|
||||
|
||||
map_search_parameter = stock_search_term
|
||||
|
||||
print(f"{Colors.GREEN}Search parameter: {map_search_parameter}{Colors.RESET}")
|
||||
|
||||
print(f"{Colors.YELLOW}Mapping website using the identified search parameter...{Colors.RESET}")
|
||||
map_website = app.map_url(url, params={"search": map_search_parameter})
|
||||
print(f"{Colors.GREEN}Website mapping completed successfully.{Colors.RESET}")
|
||||
print(f"{Colors.GREEN}Located {len(map_website['links'])} relevant links.{Colors.RESET}")
|
||||
return map_website['links']
|
||||
except Exception as e:
|
||||
print(f"{Colors.RED}Error encountered during relevant page identification: {str(e)}{Colors.RESET}")
|
||||
return None
|
||||
|
||||
# Function to plot the scores using e2b
|
||||
def plot_scores(stock_names, stock_scores):
|
||||
print(f"{Colors.YELLOW}Plotting scores...{Colors.RESET}")
|
||||
code_to_run = f"""
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
stock_names = {stock_names}
|
||||
stock_scores = {stock_scores}
|
||||
|
||||
plt.figure(figsize=(10, 5))
|
||||
plt.bar(stock_names, stock_scores, color='blue')
|
||||
plt.xlabel('Stock Names')
|
||||
plt.ylabel('Scores')
|
||||
plt.title('Stock Investment Scores')
|
||||
plt.xticks(rotation=45)
|
||||
plt.tight_layout()
|
||||
plt.savefig('chart.png')
|
||||
plt.show()
|
||||
"""
|
||||
# Run the code inside the sandbox
|
||||
execution = sandbox.run_code(code_to_run)
|
||||
|
||||
# Check if there are any results
|
||||
if execution.results and execution.results[0].png:
|
||||
first_result = execution.results[0]
|
||||
|
||||
# Get the directory where the current python file is located
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
# Save the png to a file in the examples directory. The png is in base64 format.
|
||||
with open(os.path.join(current_dir, 'chart.png'), 'wb') as f:
|
||||
f.write(base64.b64decode(first_result.png))
|
||||
print('Chart saved as examples/chart.png')
|
||||
else:
|
||||
print(f"{Colors.RED}No results returned from the sandbox execution.{Colors.RESET}")
|
||||
|
||||
# Analyze the top stocks and provide investment recommendation
|
||||
def analyze_top_stocks(map_website, app, client):
|
||||
try:
|
||||
# Get top 5 links from the map result
|
||||
top_links = map_website[:10]
|
||||
print(f"{Colors.CYAN}Proceeding to analyze top {len(top_links)} links: {top_links}{Colors.RESET}")
|
||||
|
||||
# Scrape the pages in batch
|
||||
batch_scrape_result = app.batch_scrape_urls(top_links, {'formats': ['markdown']})
|
||||
print(f"{Colors.GREEN}Batch page scraping completed successfully.{Colors.RESET}")
|
||||
|
||||
# Prepare content for LLM
|
||||
stock_contents = []
|
||||
for scrape_result in batch_scrape_result['data']:
|
||||
stock_contents.append({
|
||||
'content': scrape_result['markdown']
|
||||
})
|
||||
|
||||
# Pass all the content to the LLM to analyze and decide which stock to invest in
|
||||
analyze_prompt = f"""
|
||||
Based on the following information about different stocks from their Robinhood pages, analyze and determine which stock is the best investment opportunity. DO NOT include any other text, just the JSON.
|
||||
|
||||
Return the result in the following JSON format. Only return the JSON, nothing else. Do not include backticks or any other formatting, just the JSON.
|
||||
{{
|
||||
"scores": [
|
||||
{{
|
||||
"stock_name": "<stock_name>",
|
||||
"score": <score-out-of-100>
|
||||
}},
|
||||
...
|
||||
]
|
||||
}}
|
||||
|
||||
Stock Information:
|
||||
"""
|
||||
|
||||
for stock in stock_contents:
|
||||
analyze_prompt += f"Content:\n{stock['content']}\n"
|
||||
|
||||
print(f"{Colors.YELLOW}Analyzing stock information with LLM...{Colors.RESET}")
|
||||
analyze_prompt += f"\n\nStart JSON:\n"
|
||||
completion = client.messages.create(
|
||||
model="claude-3-5-sonnet-20240620",
|
||||
max_tokens=1000,
|
||||
temperature=0,
|
||||
system="You are a financial analyst. Only return the JSON, nothing else.",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": analyze_prompt
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
result = completion.content[0].text
|
||||
print(f"{Colors.GREEN}Analysis completed. Here is the recommendation:{Colors.RESET}")
|
||||
print(f"{Colors.MAGENTA}{result}{Colors.RESET}")
|
||||
|
||||
# Plot the scores using e2b
|
||||
try:
|
||||
result_json = json.loads(result)
|
||||
scores = result_json['scores']
|
||||
stock_names = [score['stock_name'] for score in scores]
|
||||
stock_scores = [score['score'] for score in scores]
|
||||
|
||||
plot_scores(stock_names, stock_scores)
|
||||
except json.JSONDecodeError as json_err:
|
||||
print(f"{Colors.RED}Error decoding JSON response: {str(json_err)}{Colors.RESET}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{Colors.RED}Error encountered during stock analysis: {str(e)}{Colors.RESET}")
|
||||
|
||||
# Main function to execute the process
|
||||
def main():
|
||||
# Get user input
|
||||
stock_search_term = input(f"{Colors.BLUE}Enter the stock you're interested in: {Colors.RESET}")
|
||||
if not stock_search_term.strip():
|
||||
print(f"{Colors.RED}No stock entered. Exiting.{Colors.RESET}")
|
||||
return
|
||||
|
||||
url = "https://robinhood.com/stocks"
|
||||
|
||||
print(f"{Colors.YELLOW}Initiating stock analysis process...{Colors.RESET}")
|
||||
# Find the relevant pages
|
||||
map_website = find_relevant_page_via_map(stock_search_term, url, app)
|
||||
|
||||
if map_website:
|
||||
print(f"{Colors.GREEN}Relevant stock pages identified. Proceeding with detailed analysis...{Colors.RESET}")
|
||||
# Analyze top stocks
|
||||
analyze_top_stocks(map_website, app, client)
|
||||
else:
|
||||
print(f"{Colors.RED}No relevant stock pages identified. Consider refining the search term or trying a different stock.{Colors.RESET}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user