Simple LLM agent to recommend fake coffee shops via tool calling

README.md

See https://loppsided.blog/posts/2024-12-27-a-basic-tool-calling-agent/

coffee_llm_agent.py

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")