sdk/python/foundation-models/cohere/command_tools-langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Command R+ tool/function calling using LangChain\t\n",
    "\n",
    "You can use Cohere Command R or Cohere Command R+ models deployed in Azure AI and Azure ML with `langchain` to create more sophisticated intelligent applications.\n",
    "\n",
    "> Review the [documentation](https://learn.microsoft.com/en-us/azure/ai-studio/how-to/deploy-models-cohere-command) for the Cohere family of models at for AI Studio and for ML Studio for details on how to provision inference endpoints, regional availability, pricing and inference schema reference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "Before we start, there are certain steps we need to take to deploy the models:\n",
    "\n",
    "* Register for a valid Azure account with subscription \n",
    "* Make sure you have access to [Azure AI Studio](https://learn.microsoft.com/en-us/azure/ai-studio/what-is-ai-studio?tabs=home)\n",
    "* Create a project and resource group\n",
    "* Select `Cohere Command R` or `Cohere Command R+` and `Cohere-embed-v3-english` or `Cohere-embed-v3-multilingual`.\n",
    "* Deploy with \"Pay-as-you-go\"\n",
    "\n",
    "Once deployed successfully, you should be assigned for an API endpoint and a security key for inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example\n",
    "\n",
    "Tool use allows developers to connect Cohere's models to external tools like search engines, APIs, functions, databases, etc.\n",
    "\n",
    "Multi-step tool use is an extension of this basic idea, and allows the model to call more than one tool in a sequence of steps, using the results from one tool call in a subsequent step. This process allows the language model to reason, perform dynamic actions, and quickly adapt on the basis of information coming from external sources.\n",
    "\n",
    "The recommended way to achieve [multi-step tool use with Cohere](https://docs.cohere.com/docs/multi-step-tool-use) is by leveraging the [Langchain framework](https://python.langchain.com/docs/integrations/providers/cohere/#react-agent) in Python.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install --quiet langchain langchain_cohere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_cohere.chat_models import ChatCohere\n",
    "from langchain_cohere import CohereEmbeddings\n",
    "from langchain_core.documents.base import Document\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.agents import AgentExecutor\n",
    "from langchain_cohere.react_multi_hop.agent import create_cohere_react_agent\n",
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants\n",
    "We will set the values for the keys and models that we will use in our RAG pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create instances of our Chat and Embed models deployed in Azure AI or Azure ML. Use the `langchain_cohere` package and configure it as follows:\n",
    "\n",
    "- `embed_endpoint` and `command_endpoint`: Use the endpoint URL from your deployment. Include `/v1` at the end of the endpoint URL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed\n",
    "azure_cohere_embed_endpoint = \"https://<endpoint>.<region>.inference.ai.azure.com/v1\"\n",
    "azure_cohere_embed_key = \"<key>\"\n",
    "\n",
    "# Command\n",
    "azure_cohere_command_endpoint = \"https://<endpoint>.<region>.inference.ai.azure.com/v1\"\n",
    "azure_cohere_command_key = \"<key>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the command model to be used in a chat\n",
    "chat_model = ChatCohere(\n",
    "    base_url=azure_cohere_command_endpoint, cohere_api_key=azure_cohere_command_key\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the embedding model to be used in the vector index\n",
    "embed_model = CohereEmbeddings(\n",
    "    base_url=azure_cohere_embed_endpoint, cohere_api_key=azure_cohere_embed_key\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define tools\n",
    "Your agent will be equipped with the following tools. The model can pick between\n",
    "\n",
    "* doing a search on the web\n",
    "* and/or retrieving data from a vector store\n",
    "* and/or directly answering [ this tool comes out of the box! ]\n",
    "\n",
    "Plus the model can self-reflect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector search\n",
    "\n",
    "You can easily equip your agent with a vector store!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_descriptions = [\n",
    "    \"Project Alpha: Implementing a new CRM system to improve customer relationships.\",\n",
    "    \"Project Beta: Developing an AI-based forecasting tool for sales trends.\",\n",
    "    \"Project Gamma: Overhauling the company website for better user engagement.\",\n",
    "    \"Project Delta: Integrating blockchain technology for enhanced supply chain transparency.\",\n",
    "    \"Project Epsilon: Launching a cloud migration strategy to enhance data accessibility and security.\",\n",
    "    \"Project Zeta: Implementing machine learning algorithms to optimize inventory management systems.\",\n",
    "    \"Project Eta: Developing a mobile application to improve field service operations and customer reporting.\",\n",
    "    \"Project Theta: Upgrading cybersecurity measures to protect against emerging digital threats.\",\n",
    "    \"Project Iota: Creating a virtual reality training program to enhance employee skill sets and engagement.\",\n",
    "    \"Project Kappa: Automating financial processes using AI to increase efficiency and reduce errors.\",\n",
    "]\n",
    "\n",
    "project_docs = []\n",
    "for doc in project_descriptions:\n",
    "    project_docs.append(Document(page_content=doc, metadata={\"source\": \"local\"}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ingest documents with Embed endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = FAISS.from_documents(project_docs, embed_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also convert the vectorstore into a Retriever class. This allows us to easily use it in other LangChain methods, which largely work with retrievers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db_retriever = db.as_retriever(search_kwargs={\"k\": 2})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.tools.retriever import create_retriever_tool\n",
    "\n",
    "vectorstore_search = create_retriever_tool(\n",
    "    retriever=db_retriever,\n",
    "    name=\"vectorstore_search\",\n",
    "    description=\"Retrieve relevant info from a vectorstore that contains information about Projects.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Web search\n",
    "You can easily equip your agent with web search!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\n",
    "    \"TAVILY_API_KEY\"\n",
    "] = \"API_KEY\"  # you can create an API key for free on Tavily's website\n",
    "\n",
    "internet_search = TavilySearchResults()\n",
    "internet_search.name = \"internet_search\"\n",
    "internet_search.description = \"Returns a list of relevant document snippets for a textual query retrieved from the internet.\"\n",
    "\n",
    "\n",
    "class TavilySearchInput(BaseModel):\n",
    "    query: str = Field(description=\"Query to search the internet with\")\n",
    "\n",
    "\n",
    "internet_search.args_schema = TavilySearchInput"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create ReAct Agent\n",
    "\n",
    "The model can smartly pick the right tool(s) for the user query, call them in any sequence, analyze the results and self-reflect.\n",
    "Once the model considers it has enough information to answer the user question, it generates the final answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preamble\n",
    "preamble = \"\"\"\n",
    "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 Projects.\n",
    "Use internet search to define strategies.\n",
    "\"\"\"\n",
    "\n",
    "# Prompt template\n",
    "prompt = ChatPromptTemplate.from_template(\"{input}\")\n",
    "\n",
    "# Create the ReAct agent\n",
    "agent = create_cohere_react_agent(\n",
    "    llm=chat_model,\n",
    "    tools=[internet_search, vectorstore_search],\n",
    "    prompt=prompt,\n",
    ")\n",
    "\n",
    "agent_executor = AgentExecutor(\n",
    "    agent=agent, tools=[internet_search, vectorstore_search], verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ask a standalone question to the ReAct agent\n",
    "A question that requires using a predefined tool from Langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent_executor.invoke(\n",
    "    {\n",
    "        \"input\": \"Summarize the AI based projects we are currently working on\",\n",
    "        \"preamble\": preamble,\n",
    "    }\n",
    ")\n",
    "\n",
    "response[\"output\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent_executor.invoke(\n",
    "    {\n",
    "        \"input\": \"How can we automate financial processes on Project Kappa? Then define the strategy to execite that process.\",\n",
    "        \"preamble\": preamble,\n",
    "    }\n",
    ")\n",
    "\n",
    "response[\"output\"]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rag-app",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}