xxjmpxx__
xxjmpxx__

Reputation: 11

How can I implement conditional edges in Langgraph for agent decision?

Im trying to implement a toy AI agent which is going to decide between it has to implement a script in a specific programming language (where it is going to call some differents tools as analyst, developer, etc) or whether it has to just reply general information about a language query, for example, when was Python released. Im trying to implement it with Langgraph in the following way:

load_dotenv(override=True)

# initialize pinecone client and connect to pinecone index
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Define state
class GraphState(TypedDict):
    messages: Annotated[list, add_messages]

# Create the graph
builder = StateGraph(GraphState)

# Create the LLM
llm = ChatOpenAI(
    api_key=OPENAI_API_KEY,
    model="gpt-3.5-turbo",
    temperature=0.0,
    max_tokens=250,
    max_retries=2
)

def create_node(state, system_prompt):
    human_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
    ai_messages = [msg for msg in state["messages"] if isinstance(msg, AIMessage)]
    system_message = [SystemMessage(content=system_prompt)]
    messages = system_message + human_messages + ai_messages
    message = llm.invoke(messages)
    return {"messages": [message]}

# Read the CSV file with programming language data
def read_programming_language_info():
    language_info = {}
    try:
        with open('programming_languages.csv', newline='', encoding='utf-8') as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                language_info[row['Language'].lower()] = {
                    'release_year': row['Release Year'],
                    'creator': row['Creator'],
                    'paradigm': row['Paradigm']
                }
    except FileNotFoundError:
        print("CSV file not found.")
    return language_info

# A decision-making node to classify if the query is about language info or development task
def query_decision_func(state):
    user_query = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)][-1].content

    system_message = """
    You are an AI assistant that decides whether the query is about software development tasks or a specific question about a programming language.
    Based on the context of the user's query, determine whether it is about:
    - A specific programming language (e.g., "When was Java released?")
    - A software development task (e.g., writing code, designing architecture, etc.)
    
    Provide your decision on how the query should be processed.
    If it's related to a programming language, return "language_info".
    If it's related to software development tasks, return "development_tasks".
    """
    
    # Get the model's response to the query and system message
    response = llm.invoke([SystemMessage(content=system_message)] + [HumanMessage(content=user_query)])
    
    # Analyze the response from the model
    model_decision = response["messages"][0].lower()
    
    # Based on the model's decision, route to the appropriate node
    if "language_info" in model_decision:
        return "language_info"
    elif "development_tasks" in model_decision:
        return "analyst"
    else:
        # Default case: route to development tasks if unsure
        return "analyst"

# Node to handle programming language queries
def language_info_node(state):
    user_query = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)][-1].content
    language_info = read_programming_language_info()

    # Find the language in the user query (without explicitly checking for keywords)
    for language in language_info:
        if language in user_query.lower():
            return {"messages": [f"{language.capitalize()} was released in {language_info[language]['release_year']} and was created by {language_info[language]['creator']}. It is a {language_info[language]['paradigm']} language."]}

    return {"messages": ["I couldn't find information about that programming language."]}

# Nodes for development tasks
analyst = lambda state: create_node(state, "You are a software requirements analyst...")
architect = lambda state: create_node(state, "You are a Software Architect...")
developer = lambda state: create_node(state, "You are an Full Stack Developer...")
reviewer = lambda state: create_node(state, "You are an experienced developer...")
tester = lambda state: create_node(state, "You are a test automation expert...")
diagram_designer = lambda state: create_node(state, "You are a Software Designer...")
summary_writer = lambda state: create_node(state, "You are an expert in creating technical documentation...")

# Add nodes to the graph
builder.add_node("analyst", analyst)
builder.add_node("architect", architect)
builder.add_node("developer", developer)
builder.add_node("reviewer", reviewer)
builder.add_node("tester", tester)
builder.add_node("diagram_designer", diagram_designer)
builder.add_node("summary_writer", summary_writer)
builder.add_node("language_info", language_info_node)

# Set entry point and edges
builder.add_conditional_edges(
    START,
    query_decision_func
)
builder.add_edge("analyst", "architect")
builder.add_edge("architect", "developer")
builder.add_edge("developer", "reviewer")
builder.add_edge("reviewer", "tester")
builder.add_edge("tester", "diagram_designer")
builder.add_edge("diagram_designer", "summary_writer")
builder.add_edge("summary_writer", END)
builder.add_edge("language_info", END)

# Compile and run the builder
graph = builder.compile()

# Create a main loop
def main_loop():
    ...

When I execute it, Im getting a

line 77, in query_decision_func
    model_decision = response["messages"][0].lower()
                     ~~~~~~~~^^^^^^^^^^^^
TypeError: 'AIMessage' object is not subscriptable
During task with name '__start__' and id '5a416178-5200-267a-5235-b816bb41cbfe'.

I dont know how could I fix it. Thanks in advance!

Upvotes: 0

Views: 119

Answers (0)

Related Questions