Stock_Market_AI / agents.py
roger33303's picture
Upload 4 files
56015a3 verified
raw
history blame
5.9 kB
from langgraph.graph import StateGraph, MessagesState, START, END
from langmem.short_term import SummarizationNode
from langchain_core.messages.utils import count_tokens_approximately
from langgraph.prebuilt.chat_agent_executor import AgentState
from langgraph.checkpoint.memory import InMemorySaver
from typing import Any
import os
from langchain.chat_models import init_chat_model
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor
import custom_tools as custom_tools
from dotenv import load_dotenv
load_dotenv()
supervisor_prompt = os.getenv("SUPERVISOR_PROMPT")
def build_graph():
a_debug = False
model = init_chat_model("llama3-70b-8192", model_provider="groq")
summarization_node = SummarizationNode(token_counter=count_tokens_approximately,
model=model,
max_tokens=2048,
max_summary_tokens=1024,
output_messages_key="llm_input_messages",)
class State(AgentState):
context: dict[str, Any]
checkpointer = InMemorySaver()
agent0= create_react_agent(model=model,
tools= [custom_tools.stock_data_tool],
prompt=("Give the company name to the stock tool produce the stock fundamentals. Summarize all the information in one paragraph in 200 words."),
name="Stock Analyst",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
agent1= create_react_agent(model=model,
tools= [custom_tools.web_search],
prompt='''You are a geek that can find anything from the internet. Do the research and provide the summery in 250 words and in only one paragraph.''',
name="Web Surfer",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
agent2= create_react_agent(model=model,
tools= [custom_tools.reddit_search_tool],
prompt='''You are a social media analyst. your job is to search social media and find out the sentiments of public on any topic. Do the research and provide the summery in 250 words and in only one paragraph.''',
name="Social Media Analyst",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
agent3= create_react_agent(model=model,
tools= [custom_tools.tech_news_tool],
prompt='''You are a Technology news analyst. You find the latest news related to a topic and summerize it to findout the overall picture. Do the research and provide the summery in 250 words and in only one paragraph.''',
name="Tech Journalist",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
agent4= create_react_agent(model=model,
tools= [custom_tools.politics_news_tool],
prompt='''You are a Politics news analyst. You find the latest news related to a topic and summerize it to findout the overall picture. Do the research and provide the summery in 250 words and in only one paragraph.''',
name="Political Journalist",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
agent5= create_react_agent(model=model,
tools= [custom_tools.business_news_tool],
prompt='''You are a Business news analyst. You find the latest news related to a topic and summerize it to findout the overall picture. Do the research and provide the summery in 250 words and in only one paragraph.''',
name="Business Journalist",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
agent6= create_react_agent(model=model,
tools= [custom_tools.world_news_tool],
prompt='''You are a Geopolitical news analyst. You find the latest news related to a topic and summerize it to findout the overall picture. Do the research and provide the summery in 250 words and in only one paragraph.''',
name="Geopolitical Journalist",
pre_model_hook= summarization_node,
state_schema=State,
checkpointer=checkpointer,
debug=a_debug)
supervisor_prompt = '''
'''
graph = create_supervisor(agents=[agent0,agent1,agent2,agent3,agent4,agent5,agent6],
model=model,
prompt= supervisor_prompt)
return graph.compile()