main.py

from typing import Annotated
from typing_extensions import TypedDict
import duckdb
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_aws import ChatBedrock


@tool
def query_liquor_sales(query: str):
    """
    This tool executes a SQL query on the `liquor_sales` table and returns the results.

    Wrap every column name in double quotes. For example: SELECT "Store Number" from liquor_sales.

    The table contains the following columns:
    |      column_name      | column_type | null | description |
    |-----------------------|-------------|------|-------------|
    | Invoice/Item Number   | VARCHAR     | YES  | Concatenated invoice and line number associated with the liquor order. This provides a unique identifier for the individual liquor products included in the store order    |
    | Date                  | DATE        | YES  | Date of order      |
    | Store Number          | BIGINT      | YES  | Unique number assigned to the store who ordered the liquor.    |
    | Store Name            | VARCHAR     | YES  | Name of store who ordered the liquor.    |
    | Address               | VARCHAR     | YES  | Address of store who ordered the liquor.    |
    | City                  | VARCHAR     | YES  | City where the store who ordered the liquor is located    |
    | Zip Code              | VARCHAR     | YES  | Zip code where the store who ordered the liquor is located    |
    | Store Location        | VARCHAR     | YES  | Location of store who ordered the liquor. The Address, City, State and Zip Code are geocoded to provide geographic coordinates. Accuracy of geocoding is dependent on how well the address is interpreted and the completeness of the reference data used. Left NULL where unable to provide point location.    |
    | County Number         | VARCHAR     | YES  | Iowa county number for the county where store who ordered the liquor is located    |
    | County                | VARCHAR     | YES  | County where the store who ordered the liquor is located    |
    | Category              | BIGINT      | YES  | Category code associated with the liquor ordered    |
    | Category Name         | VARCHAR     | YES  | Category of the liquor ordered.    |
    | Vendor Number         | VARCHAR     | YES  | The vendor number of the company for the brand of liquor ordered    |
    | Vendor Name           | VARCHAR     | YES  | The vendor name of the company for the brand of liquor ordered    |
    | Item Number           | VARCHAR     | YES  | Item number for the individual liquor product ordered.    |
    | Item Description      | VARCHAR     | YES  | Text    Description of the individual liquor product ordered.im_desc    |
    | Pack                  | BIGINT      | YES  | The number of bottles in a case for the liquor ordered    |
    | Bottle Volume (ml)    | BIGINT      | YES  | Volume of each liquor bottle ordered in milliliters.    |
    | State Bottle Cost     | DOUBLE      | YES  | The amount that Alcoholic Beverages Division paid for each bottle of liquor ordered    |
    | State Bottle Retail   | DOUBLE      | YES  | The amount the store paid for each bottle of liquor ordered    |
    | Bottles Sold          | BIGINT      | YES  | The number of bottles of liquor ordered by the store    |
    | Sale (Dollars)        | DOUBLE      | YES  | Total cost of liquor order (number of bottles multiplied by the state bottle retail)    |
    | Volume Sold (Liters)  | DOUBLE      | YES  | Total volume of liquor ordered in liters. (i.e. (Bottle Volume (ml) x Bottles Sold)/1,000)    |
    | Volume Sold (Gallons) | DOUBLE      | YES  | Total volume of liquor ordered in gallons. (i.e. (Bottle Volume (ml) x Bottles Sold)/3785.411784)    |
    """
    # print("-- Executing: ", query)
    liquor_sales = duckdb.sql("SELECT * FROM '../liquor_sales.parquet';")
    return duckdb.sql(query).fetchall()


tools = [query_liquor_sales]


class State(TypedDict):
    messages: Annotated[list, add_messages]

llm = ChatBedrock(
    model_id="anthropic.claude-3-sonnet-20240229-v1:0",
    model_kwargs=dict(temperature=0),
).bind_tools(tools)


def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}


graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", ToolNode(tools=tools))
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_conditional_edges("chatbot", tools_condition)
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()


def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [("user", user_input)]}):
        for value in event.values():
            content = value["messages"][-1].content
            print("Assistant: ", content)


while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("Goodbye!")
        break

    stream_graph_updates(user_input)