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