出発時刻	到着時刻	種別
6:22	7:16	特急
6:35	7:28	特急
6:48	7:41	特急
7:01	7:54	特急
7:14	8:06	特急
7:25	8:18	特急
7:37	8:30	特急
7:48	8:42	特急
8:01	8:49	特急
8:15	8:56	特急
8:31	9:09	特急
8:42	9:36	特急
8:48	9:47	快速急行
9:01	9:54	特急
9:13	10:07	特急
9:24	10:11	快速特急
9:31	10:26	特急
9:43	10:39	特急
9:54	10:41	快速特急

llm = ChatOpenAI(model="gpt-4o-mini")


class State(BaseModel):
    question: str = Field(..., description="ユーザーの質問")
    query: str = Field("", description="LLMによって生成されたSQLクエリ")
    result: str = Field("", description="SQLクエリの実行結果")
    answer: str = Field("", description="ユーザーの質問に対する回答")

class QueryOutput(BaseModel):
    """SQLクエリを返すLLMの出力"""
    query: str = Field(..., description="モデルによって生成されたSQLクエリ")

def write_query(state: State):
    """ユーザーの質問からSQLクエリを生成してstateに書き込む"""
    query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")

    prompt = query_prompt_template.invoke(
        {
            "dialect": db.dialect,
            "top_k": 10,
            "table_info": db.get_table_info(),
            "input": state.question,
        }
    )
    structured_llm = llm.with_structured_output(QueryOutput)
    result = structured_llm.invoke(prompt)
    return {"query": result.query}

def write_result(state: State):
    """SQLを実行し、結果をstateに書き込む"""
    result = str(db.run(state.query)) # デフォルトだとinclude_columnsはFalseであり、結果は文字列で書かれたリストになる

    return {"result": result}


def write_answer(state: State):
    """resultから回答を生成してstateに書き込む"""
    template = """データ抽出結果から、ユーザーの質問に対する回答を生成してください。
    データ抽出結果: {result}
    ユーザーの質問: {question}
    """
    prompt = ChatPromptTemplate.from_template(template)
    chain = prompt | llm
    answer = chain.invoke({"result": state.result, "question": state.question})
    return {"answer":answer}

# ワークフローの定義
from langgraph.graph import StateGraph

workflow = StateGraph(State)

workflow.add_node(write_query)
workflow.add_node(write_result)
workflow.add_node(write_answer)

workflow.set_entry_point("write_query")
workflow.add_edge("write_query", "write_result")
workflow.add_edge("write_result", "write_answer")
workflow.set_finish_point("write_answer")

compiled = workflow.compile()

# 実行
initial_state = State(question="京阪電車で、8:00に出発します。一番早い電車は?")

result = compiled.invoke(initial_state)

print(result["answer"])

SELECT "出発時刻", "到着時刻", "種別" 
FROM public20250112."京阪電車_時刻表_土日" 
WHERE "出発時刻" > '08:00' 
ORDER BY "出発時刻" ASC 
LIMIT 1;

content='京阪電車の一番早い特急は、8:01に出発します。到着は8:49です。'

route_prompt = ChatPromptTemplate.from_template("""\
 質問に回答するために適切なRetrieverを選択してください。

 質問: {question}
 """)

class Route(str,Enum):
    rag_document = "stories includes dandadan, spy-family, and kingdom"
    web = "web"

class RouteOutput(BaseModel):
    route: Route

model = ChatOpenAI(model="gpt-4o-mini")

route_chain = (
    route_prompt
    | model.with_structured_output(RouteOutput)
    | (lambda x: x.route)
)

res = route_chain.invoke("「スパイファミリー」について教えてください")

print(res)