A basic tool calling agent

Use Nvidia NIMs and Langchain to create an agent that uses tools to find you a coffee shop

Today I set out to build a simple AI agent based on this excellent example.

While the example in that post uses LangGraph to compose a series of agents, I ended up taking a slightly different approach.

I wrote a single “agent” that is invoked multiple times. The agent is responsible for deciding whether its prompt includes enough information to answer the user’s question. If it does, the agent responds. If not, the agent invokes a LangChain “tool”.

In this toy example, the prompt asks the agent to recommend a coffee shop near an address. The tools available are two “mock” functions: - one that returns coffee shops - one that returns coffee shop reviews

While the example uses mock functions, I think it would be possible to use the Google Maps and Search APIs to bring this to life.

The sample code is available in this gist and at the bottom of this post.

A follow up is to more closely follow the pattern in the original post, breaking the steps out into separate agents with fewer responsibilities and storing the intermediates in graph state instead of in a single mega-prompt.

Here is a sample callstack for the address “1 E 161st St, Bronx, NY 10451”

Current Iteration: 1 

## Tools Called: ['search_coffee_shops']

Current Iteration: 2 

## Tools Called: ['coffee_shop_review']

Current Iteration: 3 

## Tools Called: ['coffee_shop_review']

RECOMMENDING...... Based on the reviews, I recommend 'RoastBean' 
because it has a positive review praising its friendly barista, 
rich and smooth espresso, and modern atmosphere.

The fictional reviews are fairly fun as well (also generated by a helper LLM):

{'CremaRoast': "I'm extremely disappointed with my recent visit to CremaRoast,
located near 1 E 161st St in the Bronx. As a coffee aficionado, I was excited to
try out this shop, but unfortunately, it fell short of my expectations in nearly
every aspect. The atmosphere was lacking, with a drab and uninviting decor that
made me feel like I was in a bland, corporate office rather than a cozy coffee
shop. The tables were cramped and uncomfortable, and the background music was
overpowering and annoying. But the real letdown was the coffee itself. I ordered
a cappuccino, which was over-extracted and tasted more like burnt coffee than
the rich, velvety drink I was craving. The milk was also not steamed properly,
resulting in a lukewarm and unappetizing texture. To make matters worse, the
barista seemed completely uninterested in my experience, barely acknowledging my
presence and not even bothering to ask how my drink was. The prices were also
pretty steep, especially considering the subpar quality of the coffee. Overall,
I would not recommend CremaRoast to anyone looking for a good cup of coffee in
the Bronx. With so many other excellent coffee shops in the area, it's just not
worth the visit. Rating: 1/5 stars."}

from langchain_nvidia_ai_endpoints import ChatNVIDIA
from typing_extensions import TypedDict
from typing import List, Mapping
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.tools import tool
import random
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers.openai_tools import PydanticToolsParser


API_KEY = "your-api-key"

llm = ChatNVIDIA(api_key=API_KEY, model="meta/llama-3.3-70b-instruct")


@tool
def search_coffee_shops(address: str, n: int) -> List[Mapping[str, float]]:
    """Provides a list of length n where each element of the list is the name and distance of a coffee shop near the input address"""
    if n > 3:
        raise ValueError("Agent asked for too many coffee shops in search result")
    shops = []
    for i in range(1, n):
        shops.append(get_shop(address, i))
    return shops


def get_shop(address: str, result_index: int) -> Mapping[str, float]:
    """Mock utility that looks up coffee shops by address"""
    # TODO implement this tool using the google maps API
    shop_name_eles: List[str] = random.choices(
        ["Coffee", "Roast", "Bean", "Vanilla", "Mocha", "Grande", "Crema", "Taste"], k=2
    )
    shop_name: str = "".join(shop_name_eles)
    return {shop_name: random.gammavariate(1, 2)}


class FictionalReviewWriter(BaseModel):
    """Writes fictional reviews"""

    review: str = Field(description="Review for a fictional coffee shop")


review_writer = llm.with_structured_output(FictionalReviewWriter)


@tool
def coffee_shop_review(shop_name: str, address: str) -> Mapping[str, str]:
    """Given the name of a coffee shop and a nearby address, returns a review of the shop"""
    # TODO implement this tool using the google search API
    sentiment = random.choice(["positive", "negative", "neutral"])
    prompt = f""" Write a fictional review of a coffee shop named {shop_name} near the address {address} with overall {sentiment} sentiment"""
    review = review_writer.invoke(prompt)
    return {shop_name: review.review}


class LLMCoffeeState(TypedDict):
    """
    Represents the state between calls to the LLM as it solves the user's question, potentially including tool invocations
    """

    shops: List[Mapping[str, float]]
    reviews: List[Mapping[str, str]]


system_prompt = """
    You are a bot designed to find coffee shops near a supplied address and then recommend one of those shops based on shop reviews. Include a reason along with the name of the recommended shop.
    In order to accomplish this task you have access to two tools. 
    The first tool search_coffee_shops will give a list of shops and distances. Call this tool if no shops are listed under SHOPS. Never call this tool with n > 3.
    The second tool coffee_shop_review will supply a review. Call this tool with the name of each shop until you have one review for each candidate coffee shop.
    Once you have a list of shops, and a review of each shop, reply with a recommendation and reason. Do not make further tool calls.
    
    SHOPS
    {shops}
    SHOP REVIEWS
    {reviews} 
"""

core_agent = llm.bind_tools([search_coffee_shops, coffee_shop_review])


def main(address: str):
    """Invoke the agent to find a coffee shop near the supplied address"""

    state = LLMCoffeeState(shops=[], reviews=[])

    calls = 0

    while calls <= 10:
        # update prompt
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system_prompt),
                ("human", "I want a coffee shop recommendation for {address}"),
            ]
        )

        # Uncomment to see current prompt
        # prompt_value = prompt.invoke({
        #     "address": address,
        #     "shops": str(state["shops"]),
        #     "reviews": str(state["reviews"])
        # })
        # print(prompt_value)

        agent = prompt | core_agent

        result = agent.invoke(
            {
                "address": address,
                "shops": str(state["shops"]),
                "reviews": str(state["reviews"]),
            }
        )

        if result.content != "":
            print(f"RECOMMENDING...... {result.content}")
            return

        if len(result.tool_calls):
            tools_called = []
            for tool_call in result.tool_calls:
                tool_name = tool_call["name"].lower()
                tools_called.append(tool_name)
                selected_tool = {
                    "coffee_shop_review": coffee_shop_review,
                    "search_coffee_shops": search_coffee_shops,
                }[tool_name]
                tool_msg = selected_tool.invoke(tool_call)
                if tool_name == "search_coffee_shops":
                    state["shops"] = eval(tool_msg.content)

                if tool_name == "coffee_shop_review":
                    state["reviews"].append(eval(tool_msg.content))

        calls += 1
        print(f"""
Current Iteration: {calls}
___________________________
Tools Called: {tools_called}
---------------------------
            """)


main("1 E 161st St, Bronx, NY 10451")