ashen007

{
    'answer': '[{"title": "TELEGRAPH VOYAGE", "revenue": Decimal("1631.25")}, {"title": "PELICAN COMFORTS", "revenue": Decimal("1234.47")}, {"title": "TITANS JERK", "revenue": Decimal("1217.44")}]',
    'messages': [HumanMessage(content='what are the top 3 in-term of revenue in PG rating films?',
                              id='6a0c3175-16b2-49db-b33f-a096ffc65556'), AIMessage(content='', additional_kwargs={
        'tool_calls': [{'id': 'call_8y82', 'function': {
            'arguments': '{"query": "select title, sum(amount * datediff(return_date, rental_date)) as revenue from film join (select i.film_id, r.* from inventory as i join (select rental.inventory_id, rental.rental_id, amount, return_date, rental_date from rental join payment on rental.rental_id = payment.rental_id) as r on i.inventory_id = r.inventory_id) as t on film.film_id = t.film_id where film.rating = \'PG\' group by title order by revenue desc limit 3"}',
            'name': 'db_query_tool'}, 'type': 'function'}]}, response_metadata={

ashen007

"""
Question: which film has the most rentals?
SQL Query:
    select title,
           count(rental_id) as rental_count
    from film
             join (select i.film_id,
                          r.rental_id
                   from inventory as i
                            join rental as r on i.inventory_id = r.inventory_id) as t

ashen007

{
        'answer': '[{"title": "TELEGRAPH VOYAGE", "revenue": Decimal("231.73")}, {"title": "GOODFELLAS SALUTE", "revenue": Decimal("209.69")}, {"title": "TITANS JERK", "revenue": Decimal("201.71")}]',
        'messages': [HumanMessage(content='what are the top 3 in-term of revenue in PG rating films?',
                                  id='724a1d3c-7915-4cd3-85c0-7ba2a6736f12'), AIMessage(content='', additional_kwargs={
            'tool_calls': [{'id': 'call_fvz7', 'function': {
                'arguments': '{"query":"SELECT f.title, SUM(p.amount) as revenue FROM film f JOIN inventory i ON f.film_id = i.film_id JOIN rental r ON i.inventory_id = r.inventory_id JOIN payment p ON r.rental_id = p.rental_id WHERE f.rating=\'PG\' GROUP BY f.title ORDER BY revenue DESC LIMIT 3"}',
                'name': 'db_query_tool'}, 'type': 'function'}]}, response_metadata={
            'token_usage': {'completion_tokens': 135, 'prompt_tokens': 2007, 'total_tokens': 2142,
                            'completion_time': 0.433

ashen007

"""
film table column descriptions:
film_id: Unique identifier for each film
title: Name of the film
description: Brief synopsis of the film's plot
release_year: Year the film was released
language_id: ID of the film's primary language
original_language_id: ID of the film's original language (if different)
rental_duration: Standard rental period in days
rental_rate: Cost to rent the film

ashen007

{
        'answer': '[{"title": "BEHAVIOR RUNAWAY", "rental_rate": Decimal("4.99")}, {"title": "BIRCH ANTITRUST", "rental_rate": Decimal("4.99")}, {"title": "ALI FOREVER", "rental_rate": Decimal("4.99")}]',
        'messages': [HumanMessage(content='what are the top 3 in-term of revenue in PG rating films?',
                                  id='ce9697be-50d0-417f-a168-c1652ae4c35a'), AIMessage(content='', additional_kwargs={
            'tool_calls': [{'id': 'call_p6ts', 'function': {
                'arguments': '{"query":"SELECT title, rental_rate FROM film WHERE rating = \'PG\' ORDER BY rental_rate DESC LIMIT 3;"}',
                'name': 'db_query_tool'}, 'type': 'function'}]}, response_metadata={
            'token_usage': {'completion_tokens': 55, 'prompt_tokens': 1584, 'total_tokens': 1639,
                            'completion_time': 0.180338999, 'prompt_time': 0.131903739,
                            'queue_time': 0.023737848000000006, 'total_time': 0.312242738},

ashen007

from langchain.prompts import ChatPromptTemplate
from utils.db_tools import get_engine
from tools.table_schemas import get_table_schemas

query_generation = """
You are a SQL query expert capable of analyzing questions and generating error-free SQL queries to answer them.

## Main Tasks

1. Thoroughly understand the question before generating SQL queries.

ashen007

from langgraph.graph import StateGraph, START, END
from langgraph.graph.graph import CompiledGraph
from states import State
from tools.database_tools import db_query_tool
from tools.tool_nodes import query_gen_node, should_continue, final_answer_node
from tools.error_handling import create_tool_node_with_fallback

def run() -> CompiledGraph:
    workflow = StateGraph(State)
    workflow.add_node("query_gen", query_gen_node)

ashen007

def should_continue(state: State) -> Literal["final_answer", "query_gen"]:
    messages = state["messages"]
    last_message = messages[-1]

    if last_message.content.startswith("Error:"):
        print('error in query, go into query gen node')
        return "query_gen"

    else:
        print('final answer was gathered, go into final answer node')

ashen007

from langchain_core.messages import HumanMessage
from agent.states import State
from langchain_core.output_parsers import JsonOutputParser

def final_answer_node(state: State):
    last_message = state["messages"][-1]

    try:
        last_message = HumanMessage(content=last_message.content
                                    .replace("'", "\"")

ashen007

def create_tool_node_with_fallback(tools: list) -> RunnableWithFallbacks[Any, dict]:
    """
    Creates a ToolNode with a fallback mechanism to gracefully handle errors encountered
    during tool execution and surface them to the agent for appropriate handling.
    """
    return ToolNode(tools).with_fallbacks(
        [RunnableLambda(handle_tool_error)], exception_key='error'
    )