LangChain

from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from astapa import Astapa

astapa = Astapa()
session = None

@tool
def query_database(sql: str) -> str:
    """Execute SELECT query."""
    return str(session.execute(sql))

@tool
def update_database(sql: str) -> str:
    """Execute UPDATE/INSERT/DELETE. Staged for review."""
    session.execute(sql)
    return "Staged"

llm = ChatOpenAI(model="gpt-4")
tools = [query_database, update_database]
prompt = ChatPromptTemplate.from_messages([
    ("system", "Database assistant."),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent = create_openai_tools_agent(llm, tools, prompt)
executor = AgentExecutor(agent=agent, tools=tools)

with astapa.session(
    project_id="langchain-agent",
    database="mydb",
    user="postgres",
    password="pass",
) as sess:
    session = sess
    result = executor.invoke({"input": "Refund order #102"})
    print(result["output"])
    # auto-proposes on exit

session = astapa.get_session(session_id)
print(session.diff())
session.approve()  # or session.reject()