from langchain_ibm.utilities.sql_database import WatsonxSQLDatabase

wx_sql_database = WatsonxSQLDatabase(connection_id=postgres_sql_connection_id,
                                     **credentials,
                                     schema="public"
                                     )

from langchain_ibm import ChatWatsonx

chat_llama = ChatWatsonx(
    watsonx_client=wx_sql_database.watsonx_client,
    model_id="mistralai/mistral-medium-2505",
    params={"max_new_tokens": 200},
)

from langchain_ibm.agent_toolkits.sql_toolkit import WatsonxSQLDatabaseToolkit


toolkit = WatsonxSQLDatabaseToolkit(db=wx_sql_database, llm=chat_llama)


from langchain import hub

prompt_template = hub.pull("langchain-ai/sql-agent-system-prompt")
system_message = prompt_template.format(dialect="PostgreSQL", top_k=5)

from langgraph.prebuilt import create_react_agent

agent_executor = create_react_agent(chat_llama, toolkit.get_tools(), prompt=system_message)


example_query = "What city in Japan has the largest population? How much it is?"

events = agent_executor.stream(
    {"messages": [("user", example_query)]},
    stream_mode="values",
)
for event in events:
    event["messages"][-1].pretty_print()