pip install duckdb openai  # or anthropic, or deepseek

import duckdb

# Create an in-memory database (or persistent: duckdb.connect('agent.duckdb'))
con = duckdb.connect()

# Load sample data — DuckDB reads CSV/Parquet/JSON directly
con.execute("""
    CREATE TABLE sales AS
    SELECT * FROM read_csv_auto('ecommerce_10m.csv');
""")

# Check structure
schema = con.execute("DESCRIBE sales").fetchdf()
print(schema)

import json
from openai import OpenAI

client = OpenAI(api_key="your-key-here")

def ask_agent(question: str) -> str:
    """Ask a natural language question, get DuckDB-powered insights."""
    
    # Step 1: Get database schema for LLM context
    schema_info = con.execute("""
        SELECT table_name, column_name, data_type 
        FROM duckdb_columns()
        ORDER BY table_name, column_name
    """).fetchdf().to_string()
    
    # Step 2: LLM generates DuckDB SQL
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{
            "role": "system",
            "content": f"""You are a DuckDB SQL expert. 
Database schema:\n{schema_info}\n
Convert the user's question to DuckDB SQL.
Return ONLY valid DuckDB SQL, no explanations.
Use DuckDB-specific syntax when beneficial:
- read_csv_auto, read_parquet for external data
- LIST, UNNEST, STRUCT for nested data
- QUALIFY for window function filtering"""
        }, {
            "role": "user",
            "content": question
        }]
    )
    
    sql = response.choices[0].message.content.strip()
    sql = sql.replace("```sql", "").replace("```", "").strip()
    
    # Step 3: Execute on DuckDB
    try:
        result = con.execute(sql).fetchdf()
    except Exception as e:
        # Retry with error feedback
        return f"SQL Error: {e}\nGenerated SQL: {sql}"
    
    # Step 4: LLM summarizes results
    summary = client.chat.completions.create(
        model="gpt-4o",
        messages=[{
            "role": "system",
            "content": "Summarize the data analysis results for the user. Be concise and highlight key insights."
        }, {
            "role": "user",
            "content": f"Question: {question}\n\nResults:\n{result.head(20).to_string()}"
        }]
    )
    
    return f"```sql\n{sql}\n```\n\n{summary.choices[0].message.content}"

print(ask_agent("What are our top 3 products by revenue in electronics?"))

SELECT product, SUM(amount) as revenue
FROM sales
WHERE category = 'Electronics'
GROUP BY product
ORDER BY revenue DESC
LIMIT 3;

import json

TOOLS = [{
    "type": "function",
    "function": {
        "name": "query_duckdb",
        "description": "Execute a DuckDB SQL query and return results as JSON",
        "parameters": {
            "type": "object",
            "properties": {
                "sql": {
                    "type": "string",
                    "description": "The DuckDB SQL query to execute"
                }
            },
            "required": ["sql"]
        }
    }
}, {
    "type": "function",
    "function": {
        "name": "describe_table",
        "description": "Get column names and types for a table",
        "parameters": {
            "type": "object",
            "properties": {
                "table_name": {
                    "type": "string",
                    "description": "Name of the table to describe"
                }
            },
            "required": ["table_name"]
        }
    }
}]

def execute_tool(name: str, args: dict):
    if name == "query_duckdb":
        df = con.execute(args["sql"]).fetchdf()
        return df.head(100).to_json(orient="records")
    elif name == "describe_table":
        df = con.execute(f"DESCRIBE {args['table_name']}").fetchdf()
        return df.to_json(orient="records")

def agent_with_tools(question: str, max_steps: int = 5):
    messages = [
        {"role": "system", "content": "You are a DuckDB data analyst. Use the available tools to answer questions."},
        {"role": "user", "content": question}
    ]
    
    for step in range(max_steps):
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=TOOLS,
            tool_choice="auto"
        )
        
        msg = response.choices[0].message
        
        # No tool call → final answer
        if not msg.tool_calls:
            return msg.content
        
        # Execute each tool call
        for tc in msg.tool_calls:
            args = json.loads(tc.function.arguments)
            try:
                result = execute_tool(tc.function.name, args)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result
                })
            except Exception as e:
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": f"Error: {e}"
                })
        
        messages.append(msg)
    
    return "Max steps reached"

result = agent_with_tools(
    "Compare month-over-month revenue growth for each category. "
    "Show which categories are declining and suggest why."
)
print(result)

# duckdb_mcp_server.py
from mcp.server import Server
import duckdb

app = Server("duckdb-agent")
con = duckdb.connect(":memory:")

@app.tool()
def query(sql: str) -> str:
    """Execute DuckDB SQL and return results as text."""
    return con.execute(sql).fetchdf().to_string()

@app.tool()
def load_csv(path: str, table: str = "data") -> str:
    """Load a CSV file into DuckDB as a new table."""
    con.execute(f"CREATE TABLE {table} AS SELECT * FROM read_csv_auto('{path}')")
    info = con.execute(f"SELECT COUNT(*) AS rows, COUNT(DISTINCT column_name) AS cols FROM duckdb_columns() WHERE table_name = '{table}'").fetchdf()
    return f"Loaded {info['rows'][0]} rows, {info['cols'][0]} columns"

if __name__ == "__main__":
    app.run()

{
  "mcpServers": {
    "duckdb": {
      "command": "python",
      "args": ["duckdb_mcp_server.py"],
      "env": {}
    }
  }
}