OpenAI

from openai import OpenAI
from astapa import Astapa

openai = OpenAI()
astapa = Astapa()

tools = [
    {"type": "function", "function": {
        "name": "query_database",
        "description": "Execute a SELECT query",
        "parameters": {"type": "object", "properties": {
            "sql": {"type": "string"}
        }, "required": ["sql"]}
    }},
    {"type": "function", "function": {
        "name": "update_database",
        "description": "Execute UPDATE/INSERT/DELETE",
        "parameters": {"type": "object", "properties": {
            "sql": {"type": "string"}
        }, "required": ["sql"]}
    }}
]

with astapa.session(
    project_id="openai-agent",
    database="mydb",
    user="postgres",
    password="pass",
) as session:

    def handle_tool_call(name: str, args: dict) -> str:
        if name == "query_database":
            return str(session.execute(args["sql"]))
        elif name == "update_database":
            session.execute(args["sql"])
            return "Staged for review"
        return "Unknown"

    response = openai.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "Refund order #102"}],
        tools=tools,
    )

    for tc in response.choices[0].message.tool_calls or []:
        print(handle_tool_call(tc.function.name, eval(tc.function.arguments)))

    print("Staged:", session.diff())
    # auto-proposes on exit

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