Final_Assignment_Template / main_agent.py
AGiorni's picture
successfull submission
825032d
# imports
from typing import TypedDict, Annotated, Optional
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, SystemMessage
from langchain_openai import AzureChatOpenAI
from langgraph.graph import START, StateGraph
from tools import duckduck_tool, wiki_RAG_tool, image_analyser_tool, audio_transcriber_tool, python_script_opener
from langgraph.prebuilt import ToolNode, tools_condition
import prompts_lib as my_prompts
import os
from dotenv import load_dotenv
# load environment variables
load_dotenv() # take environment variables
# define state
class State(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
file_path: Optional[str]
# create llm interface
llm = AzureChatOpenAI(
deployment_name = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME"),
openai_api_key = os.environ.get("AZURE_OPENAI_API_KEY"),
azure_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT"),
openai_api_version = os.environ.get("OPENAI_API_VERSION"),
temperature=0
)
# bild tools
tools = [duckduck_tool, wiki_RAG_tool, image_analyser_tool, audio_transcriber_tool, python_script_opener]
chat_w_tools = llm.bind_tools(tools)
# load system prompt
system_prompt = my_prompts.system_prompt2
system_message = SystemMessage(content=system_prompt)
# define nodes
def assistant(state: State):
file_path = state.get("file_path", None)
if file_path:
state["messages"].append(SystemMessage(content=f"File path provided: {file_path}"))
return {
"messages": [chat_w_tools.invoke([system_message] + state["messages"])],
}
# define graph
builder = StateGraph(State)
# add nodes
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
# define edges
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition,
{"tools": "tools", "__end__": "__end__"})
builder.add_edge("tools", "assistant")
# compile gtaph
# agent = builder.compile()
class Agent:
def __init__(self):
self.builder = builder.compile()