쿼리문을 작성하는 RAG

현서·2025년 9월 21일
쿼리문을 작성하는 RAG

AI Agent

목록 보기
8/8
post-thumbnail

1. SQLite

SQLite는 서버가 필요 없는 가볍고(self-contained) 독립적인 관계형 데이터베이스 관리 시스템(RDBMS)으로, 애플리케이션 내부에 파일 형태로 데이터베이스를 저장해 별도의 설치나 설정 없이 쉽게 사용할 수 있다.

ANSI-SQL 표준을 대부분 지원하며, 트랜잭션, 인덱스, 뷰, 트리거 같은 기능도 제공하지만 MySQL이나 PostgreSQL 같은 대형 서버형 DBMS에 비해 동시성 처리나 고성능 분산 처리에는 한계가 있다.

작은 규모의 웹·모바일 앱, 프로토타입 개발, 내장형 기기 등에 활용되며, 파일 하나만 복사하면 전체 데이터베이스를 옮길 수 있을 만큼 간편성이 뛰어나다.

import requests

url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"

response = requests.get(url)

if response.status_code == 200:
    with open("Chinook.db", "wb") as file:
        file.write(response.content)
    print("Chinook.db 다운로드")
else:
    print(f"다운로드 실패: {response.status_code}")
!pip install langchain_community
from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db")
print(f"Dialect: {db.dialect}")
print(f"Available tables: {db.get_usable_table_names()}")
print(f'Sample output: {db.run("SELECT * FROM Artist LIMIT 5;")}')
Dialect: sqlite
Available tables: ['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']
Sample output: [(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains')]
from langchain_community.utilities import SQLDatabase
import os

db_path = "./Chinook.db"
print("경로 존재 여부:", os.path.exists(db_path))

db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
print(f"Dialect: {db.dialect}")
print(f"Available tables: {db.get_usable_table_names()}")
print(f'Sample output: {db.run("SELECT * FROM Artist LIMIT 5;")}')
경로 존재 여부: True
Dialect: sqlite
Available tables: ['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']
Sample output: [(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains')]
import getpass
import os

def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

_set_env("OPENAI_API_KEY")

OPENAI_API_KEY 입력

!pip install langchain_openai
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o", temperature=0, max_retries=3)

SQLDatabaseToolkit

  • sql_db_query: SQL 쿼리를 입력하면, 데이터베이스 조회 결과를 출력하는 도구로, 쿼리가 올바르지 않으면 오류 메시지가 반환됨.
  • sql_db_schema: 테이블 목록을 입력하면 해당 테이블의 스키마와 샘플 행이 출력되는 도구.
  • sql_db_list_tables: 데이터베이스 테이블 목록을 반환하는 도구.
  • sql_db_query_checker: 쿼리가 올바른지 확인하는 도구.
from langchain_community.agent_toolkits import SQLDatabaseToolkit

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

tools = toolkit.get_tools()

for tool in tools:
    print(f"{tool.name}: {tool.description}\n")
!pip install langgraph
from langgraph.prebuilt import create_react_agent

system_prompt = f"""
You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct {db.dialect} query to run,
then look at the results of the query and return the answer. Unless the user
specifies a specific number of examples they wish to obtain, always limit your
query to at most 5 results.

You can order the results by a relevant column to return the most interesting
examples in the database. Never query for all the columns from a specific table,
only ask for the relevant columns given the question.

You MUST double check your query before executing it. If you get an error while
executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
database.

To start you should ALWAYS look at the tables in the database to see what you
can query. Do NOT skip this step.

Then you should query the schema of the most relevant tables.
"""

agent = create_react_agent(
    llm,
    tools,
    prompt=system_prompt,
)

agent

question = "2009년에 가장 많은 매출을 올린 영업 사원은 누구인가요?"

for step in agent.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()
================================ Human Message =================================

2009년에 가장 많은 매출을 올린 영업 사원은 누구인가요?
================================== Ai Message ==================================
Tool Calls:
  sql_db_list_tables (call_XSHywV2CRirHFxWq4VjC8NQG)
 Call ID: call_XSHywV2CRirHFxWq4VjC8NQG
  Args:
================================= Tool Message =================================
Name: sql_db_list_tables

Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
================================== Ai Message ==================================
Tool Calls:
  sql_db_schema (call_cjj37wUrt3XhT2U2piXqLplZ)
 Call ID: call_cjj37wUrt3XhT2U2piXqLplZ
  Args:
    table_names: Employee, Invoice
================================= Tool Message =================================
Name: sql_db_schema


CREATE TABLE "Employee" (
	"EmployeeId" INTEGER NOT NULL, 
	"LastName" NVARCHAR(20) NOT NULL, 
	"FirstName" NVARCHAR(20) NOT NULL, 
	"Title" NVARCHAR(30), 
	"ReportsTo" INTEGER, 
	"BirthDate" DATETIME, 
	"HireDate" DATETIME, 
	"Address" NVARCHAR(70), 
	"City" NVARCHAR(40), 
	"State" NVARCHAR(40), 
	"Country" NVARCHAR(40), 
	"PostalCode" NVARCHAR(10), 
	"Phone" NVARCHAR(24), 
	"Fax" NVARCHAR(24), 
	"Email" NVARCHAR(60), 
	PRIMARY KEY ("EmployeeId"), 
	FOREIGN KEY("ReportsTo") REFERENCES "Employee" ("EmployeeId")
)

/*
3 rows from Employee table:
EmployeeId	LastName	FirstName	Title	ReportsTo	BirthDate	HireDate	Address	City	State	Country	PostalCode	Phone	Fax	Email
1	Adams	Andrew	General Manager	None	1962-02-18 00:00:00	2002-08-14 00:00:00	11120 Jasper Ave NW	Edmonton	AB	Canada	T5K 2N1	+1 (780) 428-9482	+1 (780) 428-3457	andrew@chinookcorp.com
2	Edwards	Nancy	Sales Manager	1	1958-12-08 00:00:00	2002-05-01 00:00:00	825 8 Ave SW	Calgary	AB	Canada	T2P 2T3	+1 (403) 262-3443	+1 (403) 262-3322	nancy@chinookcorp.com
3	Peacock	Jane	Sales Support Agent	2	1973-08-29 00:00:00	2002-04-01 00:00:00	1111 6 Ave SW	Calgary	AB	Canada	T2P 5M5	+1 (403) 262-3443	+1 (403) 262-6712	jane@chinookcorp.com
*/


CREATE TABLE "Invoice" (
	"InvoiceId" INTEGER NOT NULL, 
	"CustomerId" INTEGER NOT NULL, 
	"InvoiceDate" DATETIME NOT NULL, 
	"BillingAddress" NVARCHAR(70), 
	"BillingCity" NVARCHAR(40), 
	"BillingState" NVARCHAR(40), 
	"BillingCountry" NVARCHAR(40), 
	"BillingPostalCode" NVARCHAR(10), 
	"Total" NUMERIC(10, 2) NOT NULL, 
	PRIMARY KEY ("InvoiceId"), 
	FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")
)

/*
3 rows from Invoice table:
InvoiceId	CustomerId	InvoiceDate	BillingAddress	BillingCity	BillingState	BillingCountry	BillingPostalCode	Total
1	2	2009-01-01 00:00:00	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	1.98
2	4	2009-01-02 00:00:00	Ullevålsveien 14	Oslo	None	Norway	0171	3.96
3	8	2009-01-03 00:00:00	Grétrystraat 63	Brussels	None	Belgium	1000	5.94
*/
================================== Ai Message ==================================
Tool Calls:
  sql_db_schema (call_Ps0qm8l7Ug7k1aQzH8nzo9iy)
 Call ID: call_Ps0qm8l7Ug7k1aQzH8nzo9iy
  Args:
    table_names: Customer
================================= Tool Message =================================
Name: sql_db_schema


CREATE TABLE "Customer" (
	"CustomerId" INTEGER NOT NULL, 
	"FirstName" NVARCHAR(40) NOT NULL, 
	"LastName" NVARCHAR(20) NOT NULL, 
	"Company" NVARCHAR(80), 
	"Address" NVARCHAR(70), 
	"City" NVARCHAR(40), 
	"State" NVARCHAR(40), 
	"Country" NVARCHAR(40), 
	"PostalCode" NVARCHAR(10), 
	"Phone" NVARCHAR(24), 
	"Fax" NVARCHAR(24), 
	"Email" NVARCHAR(60) NOT NULL, 
	"SupportRepId" INTEGER, 
	PRIMARY KEY ("CustomerId"), 
	FOREIGN KEY("SupportRepId") REFERENCES "Employee" ("EmployeeId")
)

/*
3 rows from Customer table:
CustomerId	FirstName	LastName	Company	Address	City	State	Country	PostalCode	Phone	Fax	Email	SupportRepId
1	Luís	Gonçalves	Embraer - Empresa Brasileira de Aeronáutica S.A.	Av. Brigadeiro Faria Lima, 2170	São José dos Campos	SP	Brazil	12227-000	+55 (12) 3923-5555	+55 (12) 3923-5566	luisg@embraer.com.br	3
2	Leonie	Köhler	None	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	+49 0711 2842222	None	leonekohler@surfeu.de	5
3	François	Tremblay	None	1498 rue Bélanger	Montréal	QC	Canada	H2G 1A7	+1 (514) 721-4711	None	ftremblay@gmail.com	3
*/
================================== Ai Message ==================================
Tool Calls:
  sql_db_query_checker (call_YZ9TtLCgd9IWC8jyFThDBEyn)
 Call ID: call_YZ9TtLCgd9IWC8jyFThDBEyn
  Args:
    query: SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales
FROM Employee e
JOIN Customer c ON e.EmployeeId = c.SupportRepId
JOIN Invoice i ON c.CustomerId = i.CustomerId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY e.EmployeeId
ORDER BY TotalSales DESC
LIMIT 1;
================================= Tool Message =================================
Name: sql_db_query_checker

SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales
FROM Employee e
JOIN Customer c ON e.EmployeeId = c.SupportRepId
JOIN Invoice i ON c.CustomerId = i.CustomerId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY e.EmployeeId
ORDER BY TotalSales DESC
LIMIT 1;

================================== Ai Message ==================================
Tool Calls:
  sql_db_query (call_vicVkeSQyMTYHcGATIAdgfN7)
 Call ID: call_vicVkeSQyMTYHcGATIAdgfN7
  Args:
    query: SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales
FROM Employee e
JOIN Customer c ON e.EmployeeId = c.SupportRepId
JOIN Invoice i ON c.CustomerId = i.CustomerId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY e.EmployeeId
ORDER BY TotalSales DESC
LIMIT 1;
================================= Tool Message =================================
Name: sql_db_query

[('Steve', 'Johnson', 164.33999999999997)]
================================== Ai Message ==================================

2009년에 가장 많은 매출을 올린 영업 사원은 Steve Johnson으로, 총 매출액은 약 164.34입니다.
tools

2. Agent 구축

2-1. ToolNode 생성하기

from langgraph.prebuilt import ToolNode

list_tables_tool = next(tool for tool in tools if tool.name == "sql_db_list_tables")
list_tables_node = ToolNode([list_tables_tool], name="list_tables")

get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
get_schema_node = ToolNode([get_schema_tool], name="get_schema")

run_query_tool = next(tool for tool in tools if tool.name == "sql_db_query")
run_query_node = ToolNode([run_query_tool], name="run_query")

2-2. DB 조회 시작을 결정할 첫 노드 생성하기

from langgraph.graph import END, START, MessagesState, StateGraph

def chatbot(state: MessagesState):
    llm_with_tools = llm.bind_tools([list_tables_tool])
    response = llm_with_tools.invoke(state["messages"]) # tool_calls : DB 조회가 필요 / AIMessage : DB 조회가 필요X
    return {"messages": [response]}

2-3. 스키마 조회를 위한 도구 호출

def call_get_schema(state: MessagesState):
    llm_with_tools = llm.bind_tools([get_schema_tool], tool_choice="any")
    response = llm_with_tools.invoke(state["messages"]) # tool_calls : 필요한 테이블에 대한 스키마 요청

    return {"messages": [response]}

2-4. 사용자 질문 기반 쿼리문 생성하기

from langchain_core.prompts import ChatPromptTemplate

generate_query_system_prompt = f"""
You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct {db.dialect} query to run,
then look at the results of the query and return the answer. Unless the user
specifies a specific number of examples they wish to obtain, always limit your
query to at most 5 results.

You can order the results by a relevant column to return the most interesting
examples in the database. Never query for all the columns from a specific table,
only ask for the relevant columns given the question.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.
DO NOT wrap the response in any backticks or anything else(```). Respond with a SQL statement only!
"""

generate_query_user_prompt = """
User input: {question}
Schema: {schema}

If an error message is given, regenerate the query based on the error message.
History: {history}

SQL query:"""


def generate_query(state: MessagesState):
    print("##### GENERATE QUERY #####")
    print(state["messages"])
    history = ""
    for message in state["messages"][2:]:
        history += message.content + "\n"


    generate_query_msgs = [
        ("system", generate_query_system_prompt),
        ("user", generate_query_user_prompt),
    ]
    generate_propmt = ChatPromptTemplate.from_messages(generate_query_msgs)

    response = llm.invoke(generate_propmt.format_messages(
        question=state["messages"][0].content,
        schema=state["messages"][1].content,
        history=history
    ))
    print("generate_query", response)

    return {"messages": [response]}

2-5. 생성한 쿼리문 실행하기

from langchain_core.messages import AIMessage
import re  # 파일 상단 임포트

def _sanitize_sql(text: str) -> str:
    m = re.search(r"```(?:sql)?\s*(.*?)```", text, flags=re.DOTALL | re.IGNORECASE)
    return (m.group(1) if m else text).strip()

def check_query(state: MessagesState):
    print("##### CHECK QUERY #####")
    raw = state["messages"][-1].content
    query = _sanitize_sql(raw)                 # <-- 코드펜스 제거 포인트
    print("check_query", query)
    response = run_query_tool.invoke({"query": query})
    return {"messages": [AIMessage(content=str(response))]}

2-6. 질문 + 쿼리 실행 결과 기반 답변 생성하기

answer_system_prompt = """
You are a highly intelligent assistant trained to provide concise and accurate answers.
You will be given a context that has been retrieved from a database using a specific SQL query.
Your task is to analyze the context and answer the user’s question based on the information provided in the context.
ANSWER IN KOREAN.
"""

def answer(state: MessagesState):
    print("##### ANSWER #####")
    question = state["messages"][0].content
    context = state["messages"][-1].content
    generated_query = state["messages"][-2].content
    print("context", context)

    answer_msgs = [
        ("system", answer_system_prompt),
        ("user", "User Question: {question} SQL Query: {generated_query} Context: {context}"),
    ]
    answer_propmt = ChatPromptTemplate.from_messages(answer_msgs)

    response = llm.invoke(answer_propmt.format_messages(question=question, generated_query=generated_query, context=context))
    print("response", response)

    return {"messages": [response]}

2-7. Graph Compile

from langgraph.prebuilt import tools_condition

graph_builder = StateGraph(MessagesState)
graph_builder.add_node("chatbot", chatbot)

graph_builder.add_edge(START, "chatbot")

graph_builder.add_node("list_tables", list_tables_node)

graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
    {
        "tools": "list_tables",
        END: END,
    },
)
graph_builder.add_node("call_get_schema", call_get_schema)
graph_builder.add_node("get_schema", get_schema_node)
graph_builder.add_node("generate_query", generate_query)
graph_builder.add_node("check_query", check_query)
graph_builder.add_node("answer", answer)

graph_builder.add_edge("list_tables", "call_get_schema")
graph_builder.add_edge("call_get_schema", "get_schema")
graph_builder.add_edge("get_schema", "generate_query")
graph_builder.add_edge("generate_query", "check_query")


def should_correct(state):
    print(state["messages"][-1].content)
    txt = state["messages"][-1].content
    if "Error:" in txt or "error" in txt.lower(): # DB 조회 결과가 [-1] 에 저장
        return "generate_query"
    else:
        return "answer"

graph_builder.add_conditional_edges(
    "check_query",
    should_correct,
    {
        "generate_query": "generate_query",
        "answer": "answer",
    },
)


graph_builder.add_edge("answer", END)

graph = graph_builder.compile()
graph

question = "2009년에 가장 많은 매출을 올린 영업 사원은 누구인가요?"

for step in graph.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()
================================ Human Message =================================

2009년에 가장 많은 매출을 올린 영업 사원은 누구인가요?
================================== Ai Message ==================================
Tool Calls:
  sql_db_list_tables (call_pliFVULoE2PqBfS8oftPvmHT)
 Call ID: call_pliFVULoE2PqBfS8oftPvmHT
  Args:
================================= Tool Message =================================
Name: sql_db_list_tables

Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
================================== Ai Message ==================================
Tool Calls:
  sql_db_schema (call_Cm9o0m44u9lUFldfUw15Lsj6)
 Call ID: call_Cm9o0m44u9lUFldfUw15Lsj6
  Args:
    table_names: Employee
  sql_db_schema (call_crmN6plZmJp8WdulQenrTWb3)
 Call ID: call_crmN6plZmJp8WdulQenrTWb3
  Args:
    table_names: Invoice
================================= Tool Message =================================
Name: sql_db_schema


CREATE TABLE "Invoice" (
	"InvoiceId" INTEGER NOT NULL, 
	"CustomerId" INTEGER NOT NULL, 
	"InvoiceDate" DATETIME NOT NULL, 
	"BillingAddress" NVARCHAR(70), 
	"BillingCity" NVARCHAR(40), 
	"BillingState" NVARCHAR(40), 
	"BillingCountry" NVARCHAR(40), 
	"BillingPostalCode" NVARCHAR(10), 
	"Total" NUMERIC(10, 2) NOT NULL, 
	PRIMARY KEY ("InvoiceId"), 
	FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")
)

/*
3 rows from Invoice table:
InvoiceId	CustomerId	InvoiceDate	BillingAddress	BillingCity	BillingState	BillingCountry	BillingPostalCode	Total
1	2	2009-01-01 00:00:00	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	1.98
2	4	2009-01-02 00:00:00	Ullevålsveien 14	Oslo	None	Norway	0171	3.96
3	8	2009-01-03 00:00:00	Grétrystraat 63	Brussels	None	Belgium	1000	5.94
*/
##### GENERATE QUERY #####
[HumanMessage(content='2009년에 가장 많은 매출을 올린 영업 사원은 누구인가요?', additional_kwargs={}, response_metadata={}, id='29cf76cf-7b4c-4a3d-a6f4-6a5929686667'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_pliFVULoE2PqBfS8oftPvmHT', 'function': {'arguments': '{}', 'name': 'sql_db_list_tables'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 12, 'prompt_tokens': 86, 'total_tokens': 98, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJvWK4kmVsEiB6OGelEofVqFi5DP', 'service_tier': 'default', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--8e2a3528-c4d5-4156-a961-8cf60dccf053-0', tool_calls=[{'name': 'sql_db_list_tables', 'args': {}, 'id': 'call_pliFVULoE2PqBfS8oftPvmHT', 'type': 'tool_call'}], usage_metadata={'input_tokens': 86, 'output_tokens': 12, 'total_tokens': 98, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), ToolMessage(content='Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track', name='sql_db_list_tables', id='db5a4f5f-fbee-4b1e-92c6-17550a9251dd', tool_call_id='call_pliFVULoE2PqBfS8oftPvmHT'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Cm9o0m44u9lUFldfUw15Lsj6', 'function': {'arguments': '{"table_names": "Employee"}', 'name': 'sql_db_schema'}, 'type': 'function'}, {'id': 'call_crmN6plZmJp8WdulQenrTWb3', 'function': {'arguments': '{"table_names": "Invoice"}', 'name': 'sql_db_schema'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 48, 'prompt_tokens': 183, 'total_tokens': 231, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJvXMiQOwOTkEl1Oo5aj9ATbjDHv', 'service_tier': 'default', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--29f05f43-28e4-4164-88af-689366f96e70-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Employee'}, 'id': 'call_Cm9o0m44u9lUFldfUw15Lsj6', 'type': 'tool_call'}, {'name': 'sql_db_schema', 'args': {'table_names': 'Invoice'}, 'id': 'call_crmN6plZmJp8WdulQenrTWb3', 'type': 'tool_call'}], usage_metadata={'input_tokens': 183, 'output_tokens': 48, 'total_tokens': 231, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), ToolMessage(content='\nCREATE TABLE "Employee" (\n\t"EmployeeId" INTEGER NOT NULL, \n\t"LastName" NVARCHAR(20) NOT NULL, \n\t"FirstName" NVARCHAR(20) NOT NULL, \n\t"Title" NVARCHAR(30), \n\t"ReportsTo" INTEGER, \n\t"BirthDate" DATETIME, \n\t"HireDate" DATETIME, \n\t"Address" NVARCHAR(70), \n\t"City" NVARCHAR(40), \n\t"State" NVARCHAR(40), \n\t"Country" NVARCHAR(40), \n\t"PostalCode" NVARCHAR(10), \n\t"Phone" NVARCHAR(24), \n\t"Fax" NVARCHAR(24), \n\t"Email" NVARCHAR(60), \n\tPRIMARY KEY ("EmployeeId"), \n\tFOREIGN KEY("ReportsTo") REFERENCES "Employee" ("EmployeeId")\n)\n\n/*\n3 rows from Employee table:\nEmployeeId\tLastName\tFirstName\tTitle\tReportsTo\tBirthDate\tHireDate\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\n1\tAdams\tAndrew\tGeneral Manager\tNone\t1962-02-18 00:00:00\t2002-08-14 00:00:00\t11120 Jasper Ave NW\tEdmonton\tAB\tCanada\tT5K 2N1\t+1 (780) 428-9482\t+1 (780) 428-3457\tandrew@chinookcorp.com\n2\tEdwards\tNancy\tSales Manager\t1\t1958-12-08 00:00:00\t2002-05-01 00:00:00\t825 8 Ave SW\tCalgary\tAB\tCanada\tT2P 2T3\t+1 (403) 262-3443\t+1 (403) 262-3322\tnancy@chinookcorp.com\n3\tPeacock\tJane\tSales Support Agent\t2\t1973-08-29 00:00:00\t2002-04-01 00:00:00\t1111 6 Ave SW\tCalgary\tAB\tCanada\tT2P 5M5\t+1 (403) 262-3443\t+1 (403) 262-6712\tjane@chinookcorp.com\n*/', name='sql_db_schema', id='71dcdbab-bc6d-42e2-9ee8-f2fd43d90884', tool_call_id='call_Cm9o0m44u9lUFldfUw15Lsj6'), ToolMessage(content='\nCREATE TABLE "Invoice" (\n\t"InvoiceId" INTEGER NOT NULL, \n\t"CustomerId" INTEGER NOT NULL, \n\t"InvoiceDate" DATETIME NOT NULL, \n\t"BillingAddress" NVARCHAR(70), \n\t"BillingCity" NVARCHAR(40), \n\t"BillingState" NVARCHAR(40), \n\t"BillingCountry" NVARCHAR(40), \n\t"BillingPostalCode" NVARCHAR(10), \n\t"Total" NUMERIC(10, 2) NOT NULL, \n\tPRIMARY KEY ("InvoiceId"), \n\tFOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")\n)\n\n/*\n3 rows from Invoice table:\nInvoiceId\tCustomerId\tInvoiceDate\tBillingAddress\tBillingCity\tBillingState\tBillingCountry\tBillingPostalCode\tTotal\n1\t2\t2009-01-01 00:00:00\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t1.98\n2\t4\t2009-01-02 00:00:00\tUllevålsveien 14\tOslo\tNone\tNorway\t0171\t3.96\n3\t8\t2009-01-03 00:00:00\tGrétrystraat 63\tBrussels\tNone\tBelgium\t1000\t5.94\n*/', name='sql_db_schema', id='bdce1dd3-30f6-4f37-b8ac-1ecb7edb4eaf', tool_call_id='call_crmN6plZmJp8WdulQenrTWb3')]
generate_query content="```sql\nSELECT e.FirstName, e.LastName, SUM(i.Total) AS TotalSales\nFROM Employee e\nJOIN Customer c ON e.EmployeeId = c.SupportRepId\nJOIN Invoice i ON c.CustomerId = i.CustomerId\nWHERE strftime('%Y', i.InvoiceDate) = '2009'\nGROUP BY e.EmployeeId\nORDER BY TotalSales DESC\nLIMIT 1;\n```" additional_kwargs={'refusal': None} response_metadata={'token_usage': {'completion_tokens': 81, 'prompt_tokens': 1019, 'total_tokens': 1100, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJvZWBNpqLIiTpSkNmvuf8q8wVpG', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None} id='run--7e6d1a50-1218-4622-bdea-07c455121793-0' usage_metadata={'input_tokens': 1019, 'output_tokens': 81, 'total_tokens': 1100, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}
================================== Ai Message ==================================


SELECT e.FirstName, e.LastName, SUM(i.Total) AS TotalSales
FROM Employee e
JOIN Customer c ON e.EmployeeId = c.SupportRepId
JOIN Invoice i ON c.CustomerId = i.CustomerId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY e.EmployeeId
ORDER BY TotalSales DESC
LIMIT 1;

##### CHECK QUERY #####
check_query SELECT e.FirstName, e.LastName, SUM(i.Total) AS TotalSales
FROM Employee e
JOIN Customer c ON e.EmployeeId = c.SupportRepId
JOIN Invoice i ON c.CustomerId = i.CustomerId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY e.EmployeeId
ORDER BY TotalSales DESC
LIMIT 1;
[('Steve', 'Johnson', 164.33999999999997)]
================================== Ai Message ==================================

[('Steve', 'Johnson', 164.33999999999997)]
##### ANSWER #####
context [('Steve', 'Johnson', 164.33999999999997)]
response content='2009년에 가장 많은 매출을 올린 영업 사원은 Steve Johnson입니다.' additional_kwargs={'refusal': None} response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 197, 'total_tokens': 216, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJvaRtkTvfThlIE21THJqipuoVxt', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None} id='run--048e5a27-2af8-4ae5-87e6-30f39ea28f0f-0' usage_metadata={'input_tokens': 197, 'output_tokens': 19, 'total_tokens': 216, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}
================================== Ai Message ==================================

2009년에 가장 많은 매출을 올린 영업 사원은 Steve Johnson입니다.
question = "DB 조회를 기반으로 2009년 가장 많은 양의 음반을 판매한 아티스트의 해당 앨범 판매 기간을 알려주세요. "

for step in graph.stream(
    {"messages": [{"role": "user", "content": question}]},
    stream_mode="values",
):
    step["messages"][-1].pretty_print()
================================ Human Message =================================

DB 조회를 기반으로 2009년 가장 많은 양의 음반을 판매한 아티스트의 해당 앨범 판매 기간을 알려주세요. 
================================== Ai Message ==================================
Tool Calls:
  sql_db_list_tables (call_NryktoBvtjFLYXz57zNGQFb9)
 Call ID: call_NryktoBvtjFLYXz57zNGQFb9
  Args:
    tool_input:
================================= Tool Message =================================
Name: sql_db_list_tables

Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
================================== Ai Message ==================================
Tool Calls:
  sql_db_schema (call_1tpMoRQC0PHuEkMVZ1XEFS2k)
 Call ID: call_1tpMoRQC0PHuEkMVZ1XEFS2k
  Args:
    table_names: Album, Artist, Invoice, InvoiceLine
================================= Tool Message =================================
Name: sql_db_schema


CREATE TABLE "Album" (
	"AlbumId" INTEGER NOT NULL, 
	"Title" NVARCHAR(160) NOT NULL, 
	"ArtistId" INTEGER NOT NULL, 
	PRIMARY KEY ("AlbumId"), 
	FOREIGN KEY("ArtistId") REFERENCES "Artist" ("ArtistId")
)

/*
3 rows from Album table:
AlbumId	Title	ArtistId
1	For Those About To Rock We Salute You	1
2	Balls to the Wall	2
3	Restless and Wild	2
*/


CREATE TABLE "Artist" (
	"ArtistId" INTEGER NOT NULL, 
	"Name" NVARCHAR(120), 
	PRIMARY KEY ("ArtistId")
)

/*
3 rows from Artist table:
ArtistId	Name
1	AC/DC
2	Accept
3	Aerosmith
*/


CREATE TABLE "Invoice" (
	"InvoiceId" INTEGER NOT NULL, 
	"CustomerId" INTEGER NOT NULL, 
	"InvoiceDate" DATETIME NOT NULL, 
	"BillingAddress" NVARCHAR(70), 
	"BillingCity" NVARCHAR(40), 
	"BillingState" NVARCHAR(40), 
	"BillingCountry" NVARCHAR(40), 
	"BillingPostalCode" NVARCHAR(10), 
	"Total" NUMERIC(10, 2) NOT NULL, 
	PRIMARY KEY ("InvoiceId"), 
	FOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")
)

/*
3 rows from Invoice table:
InvoiceId	CustomerId	InvoiceDate	BillingAddress	BillingCity	BillingState	BillingCountry	BillingPostalCode	Total
1	2	2009-01-01 00:00:00	Theodor-Heuss-Straße 34	Stuttgart	None	Germany	70174	1.98
2	4	2009-01-02 00:00:00	Ullevålsveien 14	Oslo	None	Norway	0171	3.96
3	8	2009-01-03 00:00:00	Grétrystraat 63	Brussels	None	Belgium	1000	5.94
*/


CREATE TABLE "InvoiceLine" (
	"InvoiceLineId" INTEGER NOT NULL, 
	"InvoiceId" INTEGER NOT NULL, 
	"TrackId" INTEGER NOT NULL, 
	"UnitPrice" NUMERIC(10, 2) NOT NULL, 
	"Quantity" INTEGER NOT NULL, 
	PRIMARY KEY ("InvoiceLineId"), 
	FOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"), 
	FOREIGN KEY("InvoiceId") REFERENCES "Invoice" ("InvoiceId")
)

/*
3 rows from InvoiceLine table:
InvoiceLineId	InvoiceId	TrackId	UnitPrice	Quantity
1	1	2	0.99	1
2	1	4	0.99	1
3	2	6	0.99	1
*/
##### GENERATE QUERY #####
[HumanMessage(content='DB 조회를 기반으로 2009년 가장 많은 양의 음반을 판매한 아티스트의 해당 앨범 판매 기간을 알려주세요. ', additional_kwargs={}, response_metadata={}, id='5cfd15d7-c67e-4e7e-86de-6370a06c7c05'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_NryktoBvtjFLYXz57zNGQFb9', 'function': {'arguments': '{"tool_input":""}', 'name': 'sql_db_list_tables'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 100, 'total_tokens': 117, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJwGq46PnwuhAanatA9eh8IdvRRx', 'service_tier': 'default', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--96c7c101-fd82-45fb-95eb-8d1f8eacfd9d-0', tool_calls=[{'name': 'sql_db_list_tables', 'args': {'tool_input': ''}, 'id': 'call_NryktoBvtjFLYXz57zNGQFb9', 'type': 'tool_call'}], usage_metadata={'input_tokens': 100, 'output_tokens': 17, 'total_tokens': 117, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), ToolMessage(content='Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track', name='sql_db_list_tables', id='52e417dc-70b0-4a0b-8c6c-89a8b4860f3c', tool_call_id='call_NryktoBvtjFLYXz57zNGQFb9'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_1tpMoRQC0PHuEkMVZ1XEFS2k', 'function': {'arguments': '{"table_names":"Album, Artist, Invoice, InvoiceLine"}', 'name': 'sql_db_schema'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 23, 'prompt_tokens': 202, 'total_tokens': 225, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJwH44usQyMQdPiuCw3vC2AyKAWA', 'service_tier': 'default', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--f846adad-7777-4579-98bf-13d943b78ec4-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Album, Artist, Invoice, InvoiceLine'}, 'id': 'call_1tpMoRQC0PHuEkMVZ1XEFS2k', 'type': 'tool_call'}], usage_metadata={'input_tokens': 202, 'output_tokens': 23, 'total_tokens': 225, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}), ToolMessage(content='\nCREATE TABLE "Album" (\n\t"AlbumId" INTEGER NOT NULL, \n\t"Title" NVARCHAR(160) NOT NULL, \n\t"ArtistId" INTEGER NOT NULL, \n\tPRIMARY KEY ("AlbumId"), \n\tFOREIGN KEY("ArtistId") REFERENCES "Artist" ("ArtistId")\n)\n\n/*\n3 rows from Album table:\nAlbumId\tTitle\tArtistId\n1\tFor Those About To Rock We Salute You\t1\n2\tBalls to the Wall\t2\n3\tRestless and Wild\t2\n*/\n\n\nCREATE TABLE "Artist" (\n\t"ArtistId" INTEGER NOT NULL, \n\t"Name" NVARCHAR(120), \n\tPRIMARY KEY ("ArtistId")\n)\n\n/*\n3 rows from Artist table:\nArtistId\tName\n1\tAC/DC\n2\tAccept\n3\tAerosmith\n*/\n\n\nCREATE TABLE "Invoice" (\n\t"InvoiceId" INTEGER NOT NULL, \n\t"CustomerId" INTEGER NOT NULL, \n\t"InvoiceDate" DATETIME NOT NULL, \n\t"BillingAddress" NVARCHAR(70), \n\t"BillingCity" NVARCHAR(40), \n\t"BillingState" NVARCHAR(40), \n\t"BillingCountry" NVARCHAR(40), \n\t"BillingPostalCode" NVARCHAR(10), \n\t"Total" NUMERIC(10, 2) NOT NULL, \n\tPRIMARY KEY ("InvoiceId"), \n\tFOREIGN KEY("CustomerId") REFERENCES "Customer" ("CustomerId")\n)\n\n/*\n3 rows from Invoice table:\nInvoiceId\tCustomerId\tInvoiceDate\tBillingAddress\tBillingCity\tBillingState\tBillingCountry\tBillingPostalCode\tTotal\n1\t2\t2009-01-01 00:00:00\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t1.98\n2\t4\t2009-01-02 00:00:00\tUllevålsveien 14\tOslo\tNone\tNorway\t0171\t3.96\n3\t8\t2009-01-03 00:00:00\tGrétrystraat 63\tBrussels\tNone\tBelgium\t1000\t5.94\n*/\n\n\nCREATE TABLE "InvoiceLine" (\n\t"InvoiceLineId" INTEGER NOT NULL, \n\t"InvoiceId" INTEGER NOT NULL, \n\t"TrackId" INTEGER NOT NULL, \n\t"UnitPrice" NUMERIC(10, 2) NOT NULL, \n\t"Quantity" INTEGER NOT NULL, \n\tPRIMARY KEY ("InvoiceLineId"), \n\tFOREIGN KEY("TrackId") REFERENCES "Track" ("TrackId"), \n\tFOREIGN KEY("InvoiceId") REFERENCES "Invoice" ("InvoiceId")\n)\n\n/*\n3 rows from InvoiceLine table:\nInvoiceLineId\tInvoiceId\tTrackId\tUnitPrice\tQuantity\n1\t1\t2\t0.99\t1\n2\t1\t4\t0.99\t1\n3\t2\t6\t0.99\t1\n*/', name='sql_db_schema', id='991a3fbc-a175-47fd-b699-264698545886', tool_call_id='call_1tpMoRQC0PHuEkMVZ1XEFS2k')]
generate_query content="```sql\nSELECT a.Name AS ArtistName, al.Title AS AlbumTitle, MIN(i.InvoiceDate) AS FirstSaleDate, MAX(i.InvoiceDate) AS LastSaleDate\nFROM Artist a\nJOIN Album al ON a.ArtistId = al.ArtistId\nJOIN InvoiceLine il ON il.TrackId IN (\n    SELECT TrackId FROM Track WHERE AlbumId = al.AlbumId\n)\nJOIN Invoice i ON i.InvoiceId = il.InvoiceId\nWHERE strftime('%Y', i.InvoiceDate) = '2009'\nGROUP BY a.ArtistId, al.AlbumId\nORDER BY SUM(il.Quantity) DESC\nLIMIT 1;\n```" additional_kwargs={'refusal': None} response_metadata={'token_usage': {'completion_tokens': 132, 'prompt_tokens': 879, 'total_tokens': 1011, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f33640a400', 'id': 'chatcmpl-CHJwI3rurRZROFlEN5AXibmuftDVI', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None} id='run--3cab18fe-7d29-4f21-b263-2d1eee4ccc73-0' usage_metadata={'input_tokens': 879, 'output_tokens': 132, 'total_tokens': 1011, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}
================================== Ai Message ==================================


SELECT a.Name AS ArtistName, al.Title AS AlbumTitle, MIN(i.InvoiceDate) AS FirstSaleDate, MAX(i.InvoiceDate) AS LastSaleDate
FROM Artist a
JOIN Album al ON a.ArtistId = al.ArtistId
JOIN InvoiceLine il ON il.TrackId IN (
    SELECT TrackId FROM Track WHERE AlbumId = al.AlbumId
)
JOIN Invoice i ON i.InvoiceId = il.InvoiceId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY a.ArtistId, al.AlbumId
ORDER BY SUM(il.Quantity) DESC
LIMIT 1;

##### CHECK QUERY #####
check_query SELECT a.Name AS ArtistName, al.Title AS AlbumTitle, MIN(i.InvoiceDate) AS FirstSaleDate, MAX(i.InvoiceDate) AS LastSaleDate
FROM Artist a
JOIN Album al ON a.ArtistId = al.ArtistId
JOIN InvoiceLine il ON il.TrackId IN (
    SELECT TrackId FROM Track WHERE AlbumId = al.AlbumId
)
JOIN Invoice i ON i.InvoiceId = il.InvoiceId
WHERE strftime('%Y', i.InvoiceDate) = '2009'
GROUP BY a.ArtistId, al.AlbumId
ORDER BY SUM(il.Quantity) DESC
LIMIT 1;
[('Chico Buarque', 'Minha Historia', '2009-01-19 00:00:00', '2009-03-09 00:00:00')]
================================== Ai Message ==================================

[('Chico Buarque', 'Minha Historia', '2009-01-19 00:00:00', '2009-03-09 00:00:00')]
##### ANSWER #####
context [('Chico Buarque', 'Minha Historia', '2009-01-19 00:00:00', '2009-03-09 00:00:00')]
response content='2009년에 가장 많은 양의 음반을 판매한 아티스트는 Chico Buarque이며, 해당 앨범 "Minha Historia"의 판매 기간은 2009년 1월 19일부터 2009년 3월 9일까지입니다.' additional_kwargs={'refusal': None} response_metadata={'token_usage': {'completion_tokens': 56, 'prompt_tokens': 285, 'total_tokens': 341, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_1827dd0c55', 'id': 'chatcmpl-CHJwLwqxhnAQe2TK24mk1E0vcBwLc', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None} id='run--f7039b36-795e-4f4c-9b89-9c0a5de5d7ec-0' usage_metadata={'input_tokens': 285, 'output_tokens': 56, 'total_tokens': 341, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}
================================== Ai Message ==================================

2009년에 가장 많은 양의 음반을 판매한 아티스트는 Chico Buarque이며, 해당 앨범 "Minha Historia"의 판매 기간은 2009년 1월 19일부터 2009년 3월 9일까지입니다.
profile
현서
The light shines in the darkness.
이전 포스트

벡터 데이터베이스

0개의 댓글