wehub-resource-sync
103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
import re
|
|
import json
|
|
import os
|
|
import yfinance as yf
|
|
from pydantic import BaseModel, Field
|
|
from crewai import Agent, Task, Crew, Process, LLM
|
|
from crewai_tools import CodeInterpreterTool, FileReadTool
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
class QueryAnalysisOutput(BaseModel):
|
|
"""Structured output for the query analysis task."""
|
|
symbols: list[str] = Field(..., description="List of stock ticker symbols (e.g., ['TSLA', 'AAPL']).")
|
|
timeframe: str = Field(..., description="Time period (e.g., '1d', '1mo', '1y').")
|
|
action: str = Field(..., description="Action to be performed (e.g., 'fetch', 'plot').")
|
|
|
|
llm = LLM(
|
|
model="ollama/deepseek-r1:7b",
|
|
base_url="http://localhost:11434",
|
|
# temperature=0.7
|
|
)
|
|
|
|
# llm = LLM(
|
|
# model="openai/gpt-4o",
|
|
# # temperature=0.7
|
|
# )
|
|
|
|
# 1) Query parser agent
|
|
query_parser_agent = Agent(
|
|
role="Stock Data Analyst",
|
|
goal="Extract stock details and fetch required data from this user query: {query}.",
|
|
backstory="You are a financial analyst specializing in stock market data retrieval.",
|
|
llm=llm,
|
|
verbose=True,
|
|
memory=True,
|
|
)
|
|
|
|
query_parsing_task = Task(
|
|
description="Analyze the user query and extract stock details.",
|
|
expected_output="A dictionary with keys: 'symbol', 'timeframe', 'action'.",
|
|
output_pydantic=QueryAnalysisOutput,
|
|
agent=query_parser_agent,
|
|
)
|
|
|
|
|
|
# 2) Code writer agent
|
|
code_writer_agent = Agent(
|
|
role="Senior Python Developer",
|
|
goal="Write Python code to visualize stock data.",
|
|
backstory="""You are a Senior Python developer specializing in stock market data visualization.
|
|
You are also a Pandas, Matplotlib and yfinance library expert.
|
|
You are skilled at writing production-ready Python code""",
|
|
llm=llm,
|
|
verbose=True,
|
|
)
|
|
|
|
code_writer_task = Task(
|
|
description="""Write Python code to visualize stock data based on the inputs from the stock analyst
|
|
where you would find stock symbol, timeframe and action.""",
|
|
expected_output="A clean and executable Python script file (.py) for stock visualization.",
|
|
agent=code_writer_agent,
|
|
)
|
|
|
|
|
|
# 3) Code interpreter agent (uses code interpreter tool from crewai)
|
|
code_interpreter_tool = CodeInterpreterTool()
|
|
|
|
code_execution_agent = Agent(
|
|
role="Senior Code Execution Expert",
|
|
goal="Review and execute the generated Python code by code writer agent to visualize stock data and fix any errors encountered. It can delegate tasks to code writer agent if needed.",
|
|
backstory="You are a code execution expert. You are skilled at executing Python code.",
|
|
# tools=[code_interpreter_tool],
|
|
allow_code_execution=True, # This automatically adds the CodeInterpreterTool
|
|
allow_delegation=True,
|
|
llm=llm,
|
|
verbose=True,
|
|
)
|
|
|
|
code_execution_task = Task(
|
|
description="""Review and execute the generated Python code by code writer agent to visualize stock data and fix any errors encountered.""",
|
|
expected_output="A clean, working and executable Python script file (.py) for stock visualization.",
|
|
agent=code_execution_agent,
|
|
)
|
|
|
|
# Create the 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
|
|
)
|
|
|
|
# Function to be wrapped inside MCP tool
|
|
def run_financial_analysis(query):
|
|
result = crew.kickoff(inputs={"query": query})
|
|
return result.raw
|
|
|
|
if __name__ == "__main__":
|
|
# Run the crew with a query
|
|
# query = input("Enter the stock to analyze: ")
|
|
result = crew.kickoff(inputs={"query": "Plot YTD stock gain of Tesla"})
|
|
print(result.raw) |