diff --git a/core/Agents/neoagent.py b/core/Agents/neoagent.py index 9de9874..7d72e11 100644 --- a/core/Agents/neoagent.py +++ b/core/Agents/neoagent.py @@ -6,7 +6,8 @@ from langchain_core.tools import tool from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import MessagesState, StateGraph, START, END -from langgraph.prebuilt import ToolNode +from langchain_core.messages import BaseMessage, AIMessageChunk, HumanMessage, AIMessage, ToolMessage +from langgraph.prebuilt import ToolNode, tools_condition from models import Model #Models for chatGPT @@ -24,28 +25,111 @@ ReAct is a simple multi-step agent architecture. Smaller graphs are often better understood by the LLMs. """ -# Defining the model TODO: Make this configurable with Llama, Grok, Gemini, Claude -model = ChatOpenAI( - model = Model.gpt_4o, - temperature=0, - max_tokens=16384, # Max tokens for mini. For gpt4o it's 128k -) # Using ChatGPT hardcoded (TODO: Make this dynamic) -# Defining the checkpoint memory saver. -memory = MemorySaver() +class Neoagent: + def __init__(self): + print(""" +------------------------------ +Instantiated NeoAgent.... +------------------------------ + """) + system_prompt = "You are Jarvis, an AI assistant here to help the human accomplish tasks. Respond in a conversational, natural style that sounds good when spoken aloud. Keep responses short and to the point, using clear, engaging language. When explaining your thought process, be concise and only describe essential steps to maintain a conversational flow." + # Defining the model TODO: Make this configurable with Llama, Grok, Gemini, Claude + model = ChatOpenAI( + model = Model.gpt_4o, + temperature=0, + max_tokens=16384, # Max tokens for mini. For gpt4o it's 128k + ) # Using ChatGPT hardcoded (TODO: Make this dynamic) + # Defining the checkpoint memory saver. + memory = MemorySaver() -# Defining the tavily web-search tool -tavily = TavilySearchResults(max_results=2) + # Defining the tavily web-search tool + tavily = TavilySearchResults(max_results=2) -tools = [add, tavily] -tool_node = ToolNode(tools) + # Adding tools and creating the tool node. + tools = [add, tavily] + tool_node = ToolNode(tools) + llm_with_tools = model.bind_tools(tools) -bound_model = model.bind_tools(tools) + class State(TypedDict): + messages: Annotated[list, add_messages] -class State(TypedDict): - messages: Annotated[list, add_messages] + graph_builder = StateGraph(State) -graph_builder = StateGraph(State) + #Executive node that thinks about the problem or query at hand. + def executive_node(state: State): + if not state["messages"]: + state["messages"] = [("system", system_prompt)] + return {"messages": [llm_with_tools.invoke(state["messages"])]} + + graph_builder.add_node("executive_node", executive_node) + graph_builder.add_node("tools", tool_node) # The prebuilt tool node added as "tools" + graph_builder.add_conditional_edges( + "executive_node", + tools_condition, + ) + + # add conditionals, entry point and compile the graph. Exit is defined in the tools node if required. + graph_builder.add_edge("tools", "executive_node") + graph_builder.set_entry_point("executive_node") + self.graph = graph_builder.compile(checkpointer=memory) + + # Draws the graph visually + with open("neoagent.png", 'wb') as f: + f.write(self.graph.get_graph().draw_mermaid_png()) + + # Streams graph updates using websockets. + def stream_graph_updates(self, user_input: str): + config = {"configurable": {"thread_id": "1"}} # TODO: Remove. This is just a placeholder + for event in self.graph.stream({"messages": [("user", user_input)]}, config): + for value in event.values(): + print("Assistant:", value["messages"][-1].content) + + async def run(self, user_prompt: str, socketio): + """ + Run the agent with a user prompt and emit the response and total tokens via socket + """ + + # TODO: Make the chats saved and restored, using this ID as the guiding values. + # Sets the thread_id for the conversation + config = {"configurable": {"thread_id": "1"}} + + try: + input = {"messages": [("human", user_prompt)]} + socketio.emit("start_message", " ") + config = {"configurable": {"thread_id": "1"}} # Thread here is hardcoded for now. + async for event in self.graph.astream_events(input, config, version='v2'): # The config uses the memory checkpoint to save chat state. Only in-memory, not persistent yet. + event_type = event.get('event') + # Focuses only on the 'on_chain_stream'-events. + # There may be better events to base the response on + if event_type == 'on_chain_end' and event['name'] == 'LangGraph': + ai_message = event['data']['output']['messages'][-1] + + if isinstance(ai_message, AIMessage): + print(ai_message) + if 'tool_calls' in ai_message.additional_kwargs: + try: + tool_call = ai_message.additional_kwargs['tool_calls'][0]['function'] + #tool_call_id = ai_message.additional_kwargs['call_tool'][0]['tool_call_id'] + socketio.emit("tool_call", tool_call) + continue + except Exception as e: + return e + + socketio.emit("chunk", ai_message.content) + socketio.emit("tokens", ai_message.usage_metadata['total_tokens']) + continue + + if event_type == 'on_chain_stream' and event['name'] == 'tools': + tool_response = event['data']['chunk']['messages'][-1] + if isinstance(tool_response, ToolMessage): + socketio.emit("tool_response", tool_response.content) + continue + + return "success" + except Exception as e: + print(e) + return e """ # Updating the state requires creating a new state (following state immutability for history and checkpoints) diff --git a/core/graphAgent.py b/core/graphAgent.py index 81be601..0c74ca4 100644 --- a/core/graphAgent.py +++ b/core/graphAgent.py @@ -83,28 +83,12 @@ def run_stream_only(self, user_prompt: str): for chunk in self.llm.stream(user_prompt): yield chunk.content - #for running the agent comment out for testing in terminal async def run(self, user_prompt: str, socketio): """ Run the agent with a user prompt and emit the response and total tokens via socket """ try: - if 'system:' in user_prompt: - user_prompt = user_prompt.replace('system:', 'suggested prompt for my ai:') - - # TODO: Link to the current chatID - # Graph.trim_history to remove excess tokens above the limit. - chat_history = [("human", "How many planets are there in the solar system?"), - ("ai", "There are eight planets in our solar system. They are Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune.")] - - input = {"messages": [ - ("system", """ - You are Jarvis, an AI assistant here to help the human accomplish tasks. - Respond in a conversational, natural style that sounds good when spoken aloud. - Keep responses short and to the point, using clear, engaging language. - When explaining your thought process, be concise and only describe essential steps to maintain a conversational flow. - """) - ] + chat_history + [("human", user_prompt)]} + input = {"messages": [("human", user_prompt)]} socketio.emit("start_message", " ") config = {"configurable": {"thread_id": "1"}} # Thread here is hardcoded for now. async for event in self.graph.astream_events(input, config, version='v2'): # The config uses the memory checkpoint to save chat state. Only in-memory, not persistent yet. diff --git a/core/main.py b/core/main.py index cc81baa..0bb6ae0 100644 --- a/core/main.py +++ b/core/main.py @@ -1,5 +1,6 @@ from flask import Flask, request, url_for, jsonify from graphAgent import Graph +from Agents.neoagent import Neoagent from models import Model from summarize_chat import summarize_chat from rag import embed_and_store @@ -29,7 +30,10 @@ socketio = SocketIO(app, cors_allowed_origins="*") # Enable CORS for WebSocket # Agent instantiation +# Graph() contains all complex tools +# Neoagent() is a simple ReAct agent that only has websearch and the add tool. For testing purposes. jarvis = Graph() # API key is configured in agent.py +#jarvis = Neoagent() # Initialize active_chats with the correct format active_chats = defaultdict(lambda: {"chat_history": []}) diff --git a/core/neoagent.png b/core/neoagent.png new file mode 100644 index 0000000..115bf25 Binary files /dev/null and b/core/neoagent.png differ