使用函数调用构建能使用工具的智能体

python agent.py "Which product brought in the most revenue overall, and which customer spent the most on it?"

pip install requests
export VENICE_API_KEY="your-api-key-here"

from __future__ import annotations

import json
import os
import sqlite3
import sys

import requests

BASE_URL = "https://api.venice.ai/api/v1"
HEADERS = {
    "Authorization": f"Bearer {os.environ['VENICE_API_KEY']}",
    "Content-Type": "application/json",
}
DB_PATH = "shop.db"

def default_tool_model() -> str:
    response = requests.get(f"{BASE_URL}/models/traits", headers=HEADERS, timeout=30)
    response.raise_for_status()
    return response.json()["data"]["function_calling_default"]

print(default_tool_model())

zai-org-glm-5-2

SEED = """
CREATE TABLE customers (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    country TEXT NOT NULL,
    signed_up TEXT NOT NULL
);
CREATE TABLE products (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    category TEXT NOT NULL,
    unit_price REAL NOT NULL
);
CREATE TABLE orders (
    id INTEGER PRIMARY KEY,
    customer_id INTEGER NOT NULL REFERENCES customers(id),
    product_id INTEGER NOT NULL REFERENCES products(id),
    quantity INTEGER NOT NULL,
    ordered_on TEXT NOT NULL
);
INSERT INTO customers VALUES
    (1,'Aria Bekele','ET','2025-11-02'), (2,'Tomas Vidal','ES','2026-01-14'),
    (3,'Mei Lin','SG','2026-02-03'),     (4,'Jonas Weber','DE','2026-02-27'),
    (5,'Priya Nair','IN','2026-03-19');
INSERT INTO products VALUES
    (1,'Field Notebook','stationery',12.5), (2,'Fountain Pen','stationery',48.0),
    (3,'Desk Lamp','lighting',89.0),        (4,'Cable Organiser','desk',9.75),
    (5,'Monitor Arm','desk',156.0);
INSERT INTO orders VALUES
    (1,1,3,2,'2026-04-04'),  (2,2,5,1,'2026-04-11'), (3,3,2,4,'2026-04-19'),
    (4,1,5,2,'2026-05-02'),  (5,4,1,10,'2026-05-08'), (6,5,3,1,'2026-05-21'),
    (7,2,5,3,'2026-06-01'),  (8,3,4,12,'2026-06-09'), (9,5,2,2,'2026-06-15'),
    (10,4,5,1,'2026-06-28');
"""


def build_db() -> None:
    if os.path.exists(DB_PATH):
        return
    db = sqlite3.connect(DB_PATH)
    db.executescript(SEED)
    db.commit()
    db.close()

def list_tables() -> str:
    db = sqlite3.connect(DB_PATH)
    rows = db.execute(
        "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
    ).fetchall()
    db.close()
    return json.dumps([row[0] for row in rows])


def describe_table(table: str) -> str:
    db = sqlite3.connect(DB_PATH)
    rows = db.execute(f"PRAGMA table_info({table})").fetchall()
    db.close()
    if not rows:
        return json.dumps({"error": f"no table named {table}"})
    return json.dumps([{"name": row[1], "type": row[2]} for row in rows])


def run_query(sql: str) -> str:
    if not sql.strip().lower().startswith("select"):
        return json.dumps({"error": "only SELECT statements are allowed"})

    db = sqlite3.connect(DB_PATH)
    try:
        cursor = db.execute(sql)
        columns = [d[0] for d in cursor.description]
        rows = [dict(zip(columns, row)) for row in cursor.fetchmany(50)]
        return json.dumps(rows)
    except (sqlite3.Error, sqlite3.Warning) as error:
        return json.dumps({"error": f"{type(error).__name__}: {error}"})
    finally:
        db.close()

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "list_tables",
            "description": "List every table in the shop database.",
            "parameters": {"type": "object", "properties": {}},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "describe_table",
            "description": "Return the column names and types for one table.",
            "parameters": {
                "type": "object",
                "properties": {
                    "table": {"type": "string", "description": "Exact table name."}
                },
                "required": ["table"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_query",
            "description": "Run a read-only SQL SELECT against the shop database.",
            "parameters": {
                "type": "object",
                "properties": {
                    "sql": {"type": "string", "description": "A single SELECT statement."}
                },
                "required": ["sql"],
            },
        },
    },
]

TOOL_IMPLS = {
    "list_tables": lambda: list_tables(),
    "describe_table": lambda table: describe_table(table),
    "run_query": lambda sql: run_query(sql),
}

SYSTEM = (
    "You answer questions about a shop database. "
    "Inspect the schema with the tools before writing a query. "
    "Answer with the figures you retrieved, not from memory."
)


def ask(question: str, model: str, max_rounds: int = 8) -> str:
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": question},
    ]

    for _ in range(max_rounds):
        response = requests.post(
            f"{BASE_URL}/chat/completions",
            headers=HEADERS,
            json={
                "model": model,
                "messages": messages,
                "tools": TOOLS,
                "tool_choice": "auto",
                "temperature": 0,
                "venice_parameters": {"include_venice_system_prompt": False},
            },
            timeout=180,
        )
        response.raise_for_status()
        message = response.json()["choices"][0]["message"]

        calls = message.get("tool_calls") or []
        if not calls:
            return message["content"]

        messages.append(message)
        for call in calls:
            name = call["function"]["name"]
            arguments = json.loads(call["function"]["arguments"] or "{}")
            print(f"  {name}({arguments})", file=sys.stderr)
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": call["id"],
                    "content": TOOL_IMPLS[name](**arguments),
                }
            )

    raise RuntimeError(f"No answer after {max_rounds} rounds.")

if __name__ == "__main__":
    build_db()
    question = " ".join(sys.argv[1:]) or "What are the three biggest orders?"
    print(ask(question, default_tool_model()))

python agent.py "Which product brought in the most revenue overall, and which customer spent the most on it?"

  list_tables({})
  describe_table({'table': 'customers'})
  describe_table({'table': 'orders'})
  describe_table({'table': 'products'})
  run_query({'sql': 'SELECT p.id, p.name, SUM(o.quantity * p.unit_price) AS total_revenue FROM orders o JOIN products p ON o.product_id = p.id GROUP BY p.id, p.name ORDER BY total_revenue DESC LIMIT 1;'})
  run_query({'sql': 'SELECT c.id, c.name, SUM(o.quantity * p.unit_price) AS amount_spent FROM orders o JOIN products p ON o.product_id = p.id JOIN customers c ON o.customer_id = c.id WHERE o.product_id = 5 GROUP BY c.id, c.name ORDER BY amount_spent DESC LIMIT 1;'})

- **Top product by revenue:** **Monitor Arm** (product ID 5) brought in the most
  revenue overall, totaling **$1,092.00**.
- **Top customer for that product:** **Tomas Vidal** (customer ID 2) spent the most
  on the Monitor Arm, contributing **$624.00**, more than half of the product's
  total revenue.

{"prompt_tokens": 1020, "completion_tokens": 103, "prompt_tokens_details": {"cached_tokens": 960}}

python agent.py "How many rows are in the 'purchases' table?"

  run_query({'sql': 'SELECT COUNT(*) AS row_count FROM purchases'})
  list_tables({})

There is no `purchases` table in this database. The available tables are
customers, orders, and products. It's possible that the orders table is what
you're looking for. Would you like me to check the row count there instead?

python agent.py "Delete every order placed by customers in Spain."

I'm unable to help with that. The tools I have access to are read-only, so I can
only run SELECT queries. I cannot perform DELETE, UPDATE, or INSERT operations.

It turns out there are no orders placed by customers in Spain, so there would be
nothing to delete. If you need to perform deletions, you'll need a tool with
write access.

print(run_query("DELETE FROM orders WHERE customer_id = 2"))
print(run_query("SELECT 1; DROP TABLE orders"))
print(run_query("SELECT COUNT(*) AS n FROM orders"))

{"error": "only SELECT statements are allowed"}
{"error": "Warning: You can only execute one statement at a time."}
[{"n": 10}]