Cohere Tools on LangChain (Integration Guide)

1
from langchain.agents import AgentExecutor
2
from langchain_cohere.react_multi_hop.agent import (
3
    create_cohere_react_agent,
4
)
5
from langchain_core.prompts import ChatPromptTemplate
6
from langchain_cohere import ChatCohere
7
8
from langchain_community.tools.tavily_search import (
9
    TavilySearchResults,
10
)
11
from pydantic import BaseModel, Field
12
13
os.environ["TAVILY_API_KEY"] = "TAVILY_API_KEY"
14
15
internet_search = TavilySearchResults()
16
internet_search.name = "internet_search"
17
internet_search.description = "Returns a list of relevant document snippets for a textual query retrieved from the internet."
18
19
20
class TavilySearchInput(BaseModel):
21
    query: str = Field(
22
        description="Query to search the internet with"
23
    )
24
25
26
internet_search.args_schema = TavilySearchInput
27
28
# Define the Cohere LLM
29
llm = ChatCohere(
30
    cohere_api_key="COHERE_API_KEY",
31
    model="command-a-03-2025",
32
    temperature=0,
33
)
34
# Preamble
35
preamble = """
36
You are an expert who answers the user's question with the most relevant datasource. You are equipped with an internet search tool and a special vectorstore of information about how to write good essays.
37
"""
38
39
# Prompt template
40
prompt = ChatPromptTemplate.from_template("{input}")
41
42
# Create the ReAct agent
43
agent = create_cohere_react_agent(
44
    llm=llm,
45
    tools=[internet_search],
46
    prompt=prompt,
47
)
48
49
agent_executor = AgentExecutor(
50
    agent=agent, tools=[internet_search], verbose=True
51
)
52
53
response = agent_executor.invoke(
54
    {
55
        "input": "Who is the mayor of the capital of Ontario",
56
        "preamble": preamble,
57
    }
58
)
59
60
print(response["output"])

1
from langchain_cohere import ChatCohere
2
from langchain_core.messages import HumanMessage
3
from pydantic import BaseModel, Field
4
5
6
# Data model
7
class web_search(BaseModel):
8
    """
9
    The internet. Use web_search for questions that are related to anything else than agents, prompt engineering, and adversarial attacks.
10
    """
11
12
    query: str = Field(
13
        description="The query to use when searching the internet."
14
    )
15
16
17
class vectorstore(BaseModel):
18
    """
19
    A vectorstore containing documents related to agents, prompt engineering, and adversarial attacks. Use the vectorstore for questions on these topics.
20
    """
21
22
    query: str = Field(
23
        description="The query to use when searching the vectorstore."
24
    )
25
26
27
# Preamble
28
preamble = """You are an expert at routing a user question to a vectorstore or web search.
29
The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.
30
Use the vectorstore for questions on these topics. Otherwise, use web-search."""
31
32
# LLM with tool use and preamble
33
# Define the Cohere LLM
34
llm = ChatCohere(
35
    cohere_api_key="COHERE_API_KEY", model="command-a-03-2025"
36
)
37
38
llm_with_tools = llm.bind_tools(
39
    tools=[web_search, vectorstore], preamble=preamble
40
)
41
42
messages = [
43
    HumanMessage("Who will the Bears draft first in the NFL draft?")
44
]
45
response = llm_with_tools.invoke(messages, force_single_step=True)
46
print(response.response_metadata["tool_calls"])
47
48
messages = [HumanMessage("What are the types of agent memory?")]
49
response = llm_with_tools.invoke(messages, force_single_step=True)
50
print(response.response_metadata["tool_calls"])
51
52
messages = [HumanMessage("Hi, How are you?")]
53
response = llm_with_tools.invoke(messages, force_single_step=True)
54
print("tool_calls" in response.response_metadata)

1
from langchain_cohere import ChatCohere, create_sql_agent
2
from langchain_community.utilities import SQLDatabase
3
import urllib.request
4
import pandas as pd
5
import sqlite3
6
7
# Download the Chinook SQLite database
8
url = "https://github.com/lerocha/chinook-database/raw/master/ChinookDatabase/DataSources/Chinook_Sqlite.sqlite"
9
urllib.request.urlretrieve(url, "Chinook.db")
10
print("Chinook database downloaded successfully.")
11
12
db = SQLDatabase.from_uri("sqlite:///Chinook.db")
13
print(db.dialect)
14
print(db.get_usable_table_names())
15
db.run("SELECT * FROM Artist LIMIT 10;")
16
17
# Define the Cohere LLM
18
llm = ChatCohere(
19
    cohere_api_key="COHERE_API_KEY",
20
    model="command-a-03-2025",
21
    temperature=0,
22
)
23
24
agent_executor = create_sql_agent(llm, db=db, verbose=True)
25
26
resp = agent_executor.invoke(
27
    "Show me the first 5 rows of the Album table."
28
)
29
print(resp)

1
from langchain_cohere import ChatCohere, create_csv_agent
2
3
# Define the Cohere LLM
4
llm = ChatCohere(
5
    cohere_api_key="COHERE_API_KEY",
6
    model="command-a-03-2025",
7
    temperature=0,
8
)
9
10
agent_executor = create_csv_agent(
11
    llm,
12
    "titanic.csv",  # https://github.com/langchain-ai/langchain/blob/master/templates/csv-agent/titanic.csv
13
)
14
15
resp = agent_executor.invoke(
16
    {"input": "How many people were on the titanic?"}
17
)
18
print(resp.get("output"))

1
from langchain_core.tools import tool
2
from langchain_cohere import ChatCohere
3
4
5
@tool
6
def add(a: int, b: int) -> int:
7
    """Adds a and b."""
8
    return a + b
9
10
11
@tool
12
def multiply(a: int, b: int) -> int:
13
    """Multiplies a and b."""
14
    return a * b
15
16
17
tools = [add, multiply]
18
19
# Define the Cohere LLM
20
llm = ChatCohere(
21
    cohere_api_key="COHERE_API_KEY",
22
    model="command-a-03-2025",
23
    temperature=0,
24
)
25
26
llm_with_tools = llm.bind_tools(tools)
27
28
query = "What is 3 * 12? Also, what is 11 + 49?"
29
30
for chunk in llm_with_tools.stream(query):
31
    if chunk.tool_call_chunks:
32
        print(chunk.tool_call_chunks)

1
from typing import Annotated
2
from typing_extensions import TypedDict
3
from langgraph.graph import StateGraph, START, END
4
from langgraph.graph.message import add_messages
5
from langchain_cohere import ChatCohere
6
7
8
# Create a state graph
9
class State(TypedDict):
10
    messages: Annotated[list, add_messages]
11
12
13
graph_builder = StateGraph(State)
14
15
# Define the Cohere LLM
16
llm = ChatCohere(
17
    cohere_api_key="COHERE_API_KEY", model="command-a-03-2025"
18
)
19
20
21
# Add nodes
22
def chatbot(state: State):
23
    return {"messages": [llm.invoke(state["messages"])]}
24
25
26
graph_builder.add_node("chatbot", chatbot)
27
graph_builder.add_edge(START, "chatbot")
28
graph_builder.add_edge("chatbot", END)
29
30
# Compile the graph
31
graph = graph_builder.compile()
32
33
# Run the chatbot
34
while True:
35
    user_input = input("User: ")
36
    print("User: " + user_input)
37
    if user_input.lower() in ["quit", "exit", "q"]:
38
        print("Goodbye!")
39
        break
40
    for event in graph.stream({"messages": ("user", user_input)}):
41
        for value in event.values():
42
            print("Assistant:", value["messages"][-1].content)

1
from langchain_community.tools.tavily_search import (
2
    TavilySearchResults,
3
)
4
from langchain_cohere import ChatCohere
5
from langgraph.graph import StateGraph, START
6
from langgraph.graph.message import add_messages
7
from langchain_core.messages import ToolMessage
8
from langchain_core.messages import BaseMessage
9
from typing import Annotated, Literal
10
from typing_extensions import TypedDict
11
import json
12
13
# Create a tool
14
tool = TavilySearchResults(max_results=2)
15
tools = [tool]
16
17
18
# Create a state graph
19
class State(TypedDict):
20
    messages: Annotated[list, add_messages]
21
22
23
graph_builder = StateGraph(State)
24
25
# Define the LLM
26
llm = ChatCohere(
27
    cohere_api_key="COHERE_API_KEY", model="command-a-03-2025"
28
)
29
30
# Bind the tools to the LLM
31
llm_with_tools = llm.bind_tools(tools)
32
33
34
# Add nodes
35
def chatbot(state: State):
36
    return {"messages": [llm_with_tools.invoke(state["messages"])]}
37
38
39
graph_builder.add_node("chatbot", chatbot)
40
41
42
class BasicToolNode:
43
    """A node that runs the tools requested in the last AIMessage."""
44
45
    def __init__(self, tools: list) -> None:
46
        self.tools_by_name = {tool.name: tool for tool in tools}
47
48
    def __call__(self, inputs: dict):
49
        if messages := inputs.get("messages", []):
50
            message = messages[-1]
51
        else:
52
            raise ValueError("No message found in input")
53
        outputs = []
54
        for tool_call in message.tool_calls:
55
            tool_result = self.tools_by_name[
56
                tool_call["name"]
57
            ].invoke(tool_call["args"])
58
            outputs.append(
59
                ToolMessage(
60
                    content=json.dumps(tool_result),
61
                    name=tool_call["name"],
62
                    tool_call_id=tool_call["id"],
63
                )
64
            )
65
        return {"messages": outputs}
66
67
68
tool_node = BasicToolNode(tools=[tool])
69
graph_builder.add_node("tools", tool_node)
70
71
72
def route_tools(
73
    state: State,
74
) -> Literal["tools", "__end__"]:
75
    """
76
    Use in the conditional_edge to route to the ToolNode if the last message
77
    has tool calls. Otherwise, route to the end.
78
    """
79
    if isinstance(state, list):
80
        ai_message = state[-1]
81
    elif messages := state.get("messages", []):
82
        ai_message = messages[-1]
83
    else:
84
        raise ValueError(
85
            f"No messages found in input state to tool_edge: {state}"
86
        )
87
    if (
88
        hasattr(ai_message, "tool_calls")
89
        and len(ai_message.tool_calls) > 0
90
    ):
91
        return "tools"
92
    return "__end__"
93
94
95
graph_builder.add_conditional_edges(
96
    "chatbot",
97
    route_tools,
98
    {"tools": "tools", "__end__": "__end__"},
99
)
100
graph_builder.add_edge("tools", "chatbot")
101
graph_builder.add_edge(START, "chatbot")
102
103
# Compile the graph
104
graph = graph_builder.compile()
105
106
# Run the chatbot
107
while True:
108
    user_input = input("User: ")
109
    if user_input.lower() in ["quit", "exit", "q"]:
110
        print("Goodbye!")
111
        break
112
    for event in graph.stream({"messages": [("user", user_input)]}):
113
        for value in event.values():
114
            if isinstance(value["messages"][-1], BaseMessage):
115
                print("Assistant:", value["messages"][-1].content)