Support for LangGraph

I’ve been trying to work through one of the LangGraph examples on using humans in the loop (https://langchain-ai.github.io/langgraph/how-tos/human_in_the_loop/time-travel/#build-the-agent) because I am trying to control the LLM response after a few tool calls.

I’ve managed to solve a few issues regarding the difference between SambaNova’s wrapper and the expected OpenAI wrapper, by providing a simple prompt (e.g. “You are a helpful agent”) and adding an AIMessage after the tool call completes inside the should_continue function. This means I get the LLM to reproduce the example.

For example, replace:

with:

model = ChatSambaNovaCloud(
    base_url="https://api.sambanova.ai/v1/",  
    api_key=api_key,
    streaming=False,
    temperature=0.01,
    model="Meta-Llama-3.1-70B-Instruct",
)

Change, call_model(state) to:

# Define the function that calls the model
def call_model(state):
    messages = state["messages"]
    if isinstance(messages[-1], ToolMessage):
        # We have to add an AI prompt here for SambaNova. Otherwise, it won't know what to do
        response_message = AIMessage(content="Please respond conversationally to the user")
        messages.append(response_message)

    response = model.invoke(messages)
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}

And then use:

function_calling_system_prompt = """
Your answer should be in the same language as the initial query.
Either call a tool or respond to the user.
You are an helpful assistant.
"""
history = function_calling_chat_template.format_prompt().to_messages()
input_message = HumanMessage(content="Can you play Taylor Swift's most popular song?")
history.append(input_message)

for event in app.stream({"messages": history}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

to generate the responses.

I can make one of the inputs fail by replacing the return:

@tool#(args_schema=GetAppleMusicSongSchema)
def play_song_on_apple(song: str):
    """Play a song on Apple Music"""
    # Call the apple music API ...
    #return f"Successfully played {song} on Apple Music!"
    return f"Sorry, you don't have access to {song} on Apple Music"

The system properly responds by asking the user if they would prefer to try the other service, creating the Messages:

================================= Tool Message =================================
Name: play_song_on_apple

Sorry, you don't have access to Shake It Off by Taylor Swift on Apple Music

================================== Ai Message ==================================

I'm sorry, you don't have access to Shake It Off by Taylor Swift on Apple Music. Would you like to try another song or use a different music service?

This is great, but unfortunately, when I provide a human message such as “Yes, please”, I get a response of [200] with no payload from the SambaNova endpoint, which generates the RuntimeError:

RuntimeError: Sambanova /complete call failed couldn't get JSON response

So, my problem is that the system seems to be only responding to the command “Yes, please” for which it cannot generate a meaningful response. I’m not sure it’s properly using the checkpointer, but state does have the correct message history. I’ve tried to manually grab that history and feed it back as the message array, but it doesn’t seem to make any difference.

Am I doing something obviously wrong here or do I have to manually control the history, such as in the ai-starter-kit RAG examples, by using ConversationSummaryMemory ?

Thanks

1 Like

@afdreher Welcome to the community and thank you for y your development efforts. I am going to reach out to our starter kit team , as they are better at LangChain and LangGraph than I am , to see if they can give you some guidance.

-Coby

@afdreher are you using ChatSambaNovaCloud from Langchain community or from AI-Starter-Kit util?

Thanks Coby!

I’m trying to follow some of the LangChain / LangGraph examples to get up to speed, but there are just little differences between SambaNova and the OpenAI systems that seem to keep coming up.

I am using ChatSambaNovaCloud from the ai-starter-kit. I’m running a Jupyter notebook alongside and importing from utils.model_wrappers.langchain_chat_models.

@afdreher do you have the latest utils.model_wrappers.langchain_chat_models utils version? if not, can you please update and check if the error persists

we did a patch last week solving one issue with the assistant, and tools messages handling and is likely to be related to your issue

Hi Jorge,

So, I just deleted my local directory and pulled a new copy from the main branch of the repository, but unfortunately, the issue is still there.

My notebook is just a minor modification to the LangGraph example, but would it be helpful to send you the sheet somehow?

As far as I can tell, the issue is that the agent is not utilizing the entire message history because I get the same kind of response I get when the example replays the tool_message without me injecting an AIMessage dictating what the LLM should do.

@afdreher We have opened a Jira ticket with our engineering department. We apologize for this occurrence and will let you know the status as we find out more.

-Coby

I actually think I may have found the issue by manually tracing the messages. I’ve inserted the tools into the function_calling example from the ai-starter-kit. It doesn’t behave quite like the LangGraph example expects, since it asks for both the song name and the platform, but it does perform with the history.

I’ll try later to see if it’s a prompt issue.

Thanks for the help!

1 Like

We also have isolated an error on our side and are debugging.

1 Like

@afdreher I guess the issue you’re facing is due to the way you’re handling the human input and the model’s response. When the model fails to generate a response, it’s not properly updating the state, which is causing the issue.
After the user provides a human message, you need to update the state with the new message. You can do this by creating a new AIMessage with the user’s response and appending it to the state’s message history.
also ensure each interaction explicitly provides sufficient context. Update the call_model function to handle fallback scenarios.
I’ve also tried with two different models response
Meta-Llama-3.1-8B-Instruct
Meta-Llama-3.1-70B-Instruct

Below I’ve provided with little changes in the code

from langchain_community.chat_models import ChatSambaNovaCloud
from langchain.schema import HumanMessage, AIMessage, BaseMessage
from langchain.tools import tool
from pydantic import BaseModel
import os


api_key = os.getenv("SAMBANOVA_API_KEY")
if not api_key:
    raise ValueError("SAMBANOVA_API_KEY environment variable not set.")

# Initialize the SambaNova model
model = ChatSambaNovaCloud(
    base_url="https://api.sambanova.ai/v1/",
    api_key=api_key,
    streaming=False,
    temperature=0.01,
    model="Meta-Llama-3.1-8B-Instruct",
)

# Define the schema for the tool arguments
class PlaySongInput(BaseModel):
    song: str

# Define a tool to simulate playing songs on Apple Music
@tool(args_schema=PlaySongInput)
def play_song_on_apple(song: str):
    """Play a song on Apple Music."""
    available_songs = ["Shake It Off", "Blank Space"]
    if song.lower() not in map(str.lower, available_songs):
        return f"Sorry, {song} is not available on Apple Music. Would you like to try Spotify instead?"
    return f"Successfully played {song} on Apple Music!"

def format_prompt():
    """Format the initial system prompt."""
    return """
Your answer should be in the same language as the initial query.
Either call a tool or respond to the user.
You are a helpful assistant.
"""

def call_model(state):
    """Invoke the SambaNova model with the current state."""
    messages = state["messages"]

    # Handle fallback if the last message isn't recognized
    if isinstance(messages[-1], BaseMessage):
        response_message = AIMessage(content="Please respond conversationally to the user.")
        messages.append(response_message)

    # Invoke the model
    try:
        response = model.invoke(messages)
    except Exception as e:
        raise RuntimeError(f"Error during model invocation: {e}")

    # Append the model's response to the message history
    state["messages"].append(response)
    return state

def handle_human_input(state, content):
    """Append a HumanMessage to the state's message history."""
    state["messages"].append(HumanMessage(content=content))
    return state

def handle_tool_response(state, tool_response):
    """Append a tool's response as an AIMessage to the state's message history."""
    state["messages"].append(AIMessage(content=tool_response))
    return state

def main():
    # Initialize state
    state = {"messages": []}

   
    state["messages"].append(AIMessage(content=format_prompt()))

   
    input_message = "Can you play Taylor Swift's most popular song?"
    state = handle_human_input(state, input_message)

   
    state = call_model(state)

    # Simulate tool usage based on the response
    song_request = {"song": "Shake It Off"}  # Example: assume model suggests "Shake It Off"
    tool_response = play_song_on_apple.invoke(song_request)  # Use invoke() and pass a dictionary
    state = handle_tool_response(state, tool_response)

    # Final call to the model with updated state
    state = call_model(state)

    
    for message in state["messages"]:
        print(f"{message.__class__.__name__}: {message.content}")

if __name__ == "__main__":
    main()

I hope this will help to resolve your queries. If anything is required more, I’m happy to connect.

2 Likes

Hi! Thanks for the response. Unfortunately, that’s not quite the problem that’s occurring here. I initially thought it might be related to LangGraph and the state or a hidden prompt issue, but it’s not. I’ve tried to replay the messages manually with different system prompts, but I get the same result.

I’ve finally been able to distill down the issue more, and the problem is that the system isn’t utilizing the full context of the message array. I’ve created an annotated script from the Jupyter notebook I’ve been playing with.

As you can see, the tool calls will behave just fine. The problem occurs when the human responds to the last AIMessage with something that isn’t a command but rather a response to the agent’s suggestion.

import getpass
import os
import sys

from dotenv import load_dotenv
from pydantic import BaseModel, Field

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langchain.globals import set_verbose
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, StateGraph
from langgraph.graph import MessagesState, START
from langgraph.prebuilt import ToolNode


set_verbose(True)

# Get absolute paths for ai-starter-kit. Your location may vary
current_dir = os.getcwd()
parent_dir = os.path.abspath(os.path.join(current_dir, '..'))
ai_kit_dir = os.path.abspath(os.path.join(parent_dir, 'ai-starter-kit'))
env_dir = parent_dir
sys.path.append(ai_kit_dir)

from utils.model_wrappers.langchain_chat_models import ChatSambaNovaCloud

os.environ["LANGCHAIN_TRACING_V2"] = "false"

# load env variables from a .env file into Python environment 
if load_dotenv(os.path.join(env_dir, '.env')):
    api_key = os.getenv('SAMBANOVA_API_KEY') 
else:
    os.environ["SAMBANOVA_API_KEY"] = getpass.getpass()
    api_key = os.environ.get("SAMBANOVA_API_KEY")


# Get the SambaNova chat client.
# Here, I've chosen to use the 70B version to have more context
model = ChatSambaNovaCloud(
    base_url="https://api.sambanova.ai/v1/",  
    api_key=api_key,
    streaming=False,
    temperature=0.01,
    model="Meta-Llama-3.1-70B-Instruct",
)


# Argument schema
# NOTE: These aren't technically necessary, but most of the SambaNova examples
# use them, whereas the newer LangGraph examples skip them.

class GetSongSchema(BaseModel):
    """Play a song"""

    song: str = Field(description='song name to play')

class GetSpotifySongSchema(GetSongSchema):
    """Play a song on Spotify"""
    pass

class GetAppleMusicSongSchema(GetSongSchema):
    """Play a song on Apple Music"""
    pass


# Define the available tools.  These are just mocks for a real tool call.
# We want the system to fail with a call to Apple Music, attempt to recover by 
# suggesting Spotify, and then respond to the user's follow-up message.

@tool(args_schema=GetSpotifySongSchema)
def play_song_on_spotify(song: str):
    """Play a song on Spotify"""
    # Call the Spotify API ...
    return f"Successfully played {song} on Spotify!"


@tool(args_schema=GetAppleMusicSongSchema)
def play_song_on_apple(song: str):
    """Play a song on Apple Music"""
    # Call the Apple Music API ...

    # We always want this to fail so that the agent has to respond with
    # context. The actual song choice is irrelevant.
    if False: #
        return f"Successfully played {song} on Apple Music!"
    else:
        return f"Sorry, you don't have access to {song} on Apple Music"

# By placing Apple first, the agent will try the failing method before trying
# Spotify, which always succeeds.
tools = [play_song_on_apple, play_song_on_spotify]
tool_node = ToolNode(tools)

# Bind the tools. Do not enable parallel calls because we want the model to
# behave as if it is calling a real service.
model = model.bind_tools(tools, parallel_tool_calls=False)

# Now we need to define the LangGraph components. This taken from the example
# and modified slightly to provide better instrumentation and wrap the tool
# call with an AIMessage, which is important for SambaNova's client.

# Define the function that determines whether to continue or not
def should_continue(state) -> str:
    messages = state["messages"]
    last_message = messages[-1]
    
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "end"
    # Otherwise if there is, we continue
    else:
        return "continue"
    
    # The values "end" and "continue" are condition names in the graph


# Define the function that calls the model
def call_model(state):
    messages = state["messages"]

    if isinstance(messages[-1], ToolMessage):
        # We have to add an AI prompt here for SambaNova. Otherwise, it won't 
        # know what to do. This should only be appended after a ToolMessage 
        # because otherwise the LLM gets very chatty about irrelevant topics.
        response_message = AIMessage(
            content="Please respond conversationally to the user"
        )
        messages.append(response_message)

    response = call_llm(model, messages)

    # We return a list, because this will get added to the existing list
    return {"messages": [response]}

# I'm putting this here to print the type of messages that we're invoking
# This is the only place in the workflow where the LLM is queried
def call_llm(llm, messages):
    # Uncomment these to show all of the message types being passed to the LLM
    # message_types = [str(type(message)) for message in messages]
    # print(f"[messages] [{', '.join(message_types)}]")
    return llm.invoke(messages) # Send back the response as usual

def create_workflow():
    """Get a new LangGraph workflow"""

    # Define a new graph
    workflow = StateGraph(MessagesState)

    # Define the two nodes we will cycle between
    workflow.add_node("agent", call_model)
    workflow.add_node("action", tool_node)

    # Set the entrypoint as `agent`
    # This means that this node is the first one called
    workflow.add_edge(START, "agent")

    # We now add a conditional edge
    workflow.add_conditional_edges(
        # First, we define the start node. We use `agent`.
        # This means these are the edges taken after the `agent` node is called.
        "agent",
        # Next, we pass in the function that will determine which node is called next.
        should_continue,
        # Finally we pass in a mapping.
        # The keys are strings, and the values are other nodes.
        # END is a special node marking that the graph should finish.
        # What will happen is we will call `should_continue`, and then the output of that
        # will be matched against the keys in this mapping.
        # Based on which one it matches, that node will then be called.
        {
            # If `tools`, then we call the tool node.
            "continue": "action",
            # Otherwise we finish.
            "end": END,
        },
    )

    # We now add a normal edge from `tools` to `agent`.
    # This means that after `tools` is called, `agent` node is called next.
    workflow.add_edge("action", "agent")

    # Set up memory
    memory = MemorySaver()

    # Finally, we compile it!
    # This compiles it into a LangChain Runnable,
    # meaning you can use it as you would any other runnable

    # We add in `interrupt_before=["action"]`
    # This will add a breakpoint before the `action` node is called
    app = workflow.compile(checkpointer=memory)

    return app

# This prompt is good enough. You can certainly use the more elaborate
# prompt from the function calling 
simple_assistant_prompt = """
Your answer should be in the same language as the initial query.
Either call a tool or respond to the user.
You are a helpful assistant.
"""

if __name__ == '__main__':

    app = create_workflow()

    # For the configuration, any thread id will work. Just invent something.
    config = {"configurable": {"thread_id": "13542"}}

    # Instantiate the template
    chat_template = ChatPromptTemplate.from_messages([('system', simple_assistant_prompt)])
    history = chat_template.format_prompt().to_messages()

    initial_human_message = HumanMessage(content="Can you play Taylor Swift's most popular song?")
    history.append(initial_human_message)
    
    # Call the system
    for event in app.stream({"messages": history}, config, stream_mode="values"):
        event["messages"][-1].pretty_print()

    # The result at this point should look something like this:
    #
    # ================================ Human Message =================================
    #
    # Can you play Taylor Swift's most popular song?
    #
    # ================================== Ai Message ==================================
    # Tool Calls:
    #   play_song_on_apple (call_ad7294371a424cfcb3)
    #  Call ID: call_ad7294371a424cfcb3
    #   Args:
    #     song: Shake It Off
    # ================================= Tool Message =================================
    # Name: play_song_on_apple
    #
    # Sorry, you don't have access to Shake It Off on Apple Music
    #
    # ================================== Ai Message ==================================
    #
    # I apologize, but it seems that you don't have access to "Shake It Off" on Apple Music. 
    # Would you like to try playing it on Spotify instead?
    #
    # If you look at the state and the messages by calling
    # `app.get_state(config).values['messages']`, you should find the messages
    # [SystemMessage, HumanMessage, AIMessage, ToolMessage, AIMessage] as shown
    # above.

    acknowledge_message = HumanMessage(content="Yes, please!")
    # Note here that we don't append the human message directory to the history
    # because the history is already stored in `config`.  You can verify that
    # message is correctly appended by setting a breakpoint or printing the
    # message array in call_llm(...)

    try:
        for event in app.stream({"messages": [acknowledge_message]}, config, stream_mode="values"):
            event["messages"][-1].pretty_print()
    except:
        # This will be triggered because the response from the LLM is empty
        # Note, however, that the LLM has all of the information it needs.
        print("An error occurred trying to recover using spotify")
1 Like

For anyone playing along, the issue here seems to be related to the prompt and to me expecting too much from Llama. I’ve stripped everything to just the messages, and the system can reflect back on the chat. It just doesn’t seem to grab the correct context from messages or perform the function call without the explicit directive.

I say I expected to much because it doesn’t correctly answer this query. So… back to the drawing board.

Thanks for the help!

from openai import OpenAI

client = OpenAI(
    base_url="https://api.sambanova.ai/v1", 
    api_key=<YOUR-KEY>
)

completion = client.chat.completions.create(
    model = "Meta-Llama-3.1-70B-Instruct",
    messages = [
         {'role': 'system', 'content': '\nYour answer should be in the same language as the initial query.\n  Think step by step.\nDo not call a tool if the input depends on another tool output you dont have yet.\nDo not try to answer until you get tools output, if you dont have an answer yet you can continue calling tools until you do.\nYour answer should be in the same language as the initial query.\nYou are a helpful assistant.\n'}, 
         {'role': 'user', 'content': "Can you play Taylor Swift's most popular song?"}, 
         {'role': 'assistant', 'content': None, 'tool_calls': [{'function': {'arguments': '{"song": "Shake It Off"}', 'name': 'play_song_on_apple'}, 'id': 'call_2e7d05544fe2479188', 'type': 'function'}]}, 
         {'role': 'tool', 'content': "Sorry, you don't have access to Shake It Off on Apple Music", 'tool_call_id': 'call_2e7d05544fe2479188'}, 
         {'role': 'system', 'content': 'Please respond conversationally to the user'},
         {'role': 'assistant', 'content': 'I apologize, but it seems that you don\'t have access to "Shake It Off" on Apple Music. Would you like to try playing it on Spotify instead?'},
         {'role': 'user', 'content': "What action did you recommend?"}
     ], 
    max_tokens = 1024, 
    temperature = 0.01, 
    tools = [
         {'type': 'function', 'function': {'name': 'play_song_on_apple', 'description': 'Play a song on Apple Music', 'parameters': {'properties': {'song': {'description': 'song name to play on apple music', 'type': 'string'}}, 'required': ['song'], 'type': 'object'}}}, 
         {'type': 'function', 'function': {'name': 'play_song_on_spotify', 'description': 'Play a song on Spotify', 'parameters': {'properties': {'song': {'description': 'song name to play on spotify', 'type': 'string'}}, 'required': ['song'], 'type': 'object'}}}], 
    tool_choice = 'auto', 
    parallel_tool_calls = False,
    stream = False
)
1 Like