scripts/query.py

# /// script
# dependencies = [
#   "duckdb",
#   "ollama",
#   "pandas",
# ]
# ///

import json
import sys

import duckdb
import ollama

DB = duckdb.connect()
DB.execute("CREATE OR REPLACE VIEW market_data AS SELECT * FROM read_parquet('data/alpaca_merged.parquet');")


def get_schema():
    return DB.execute("DESCRIBE market_data;").df().to_string(index=False)


def execute_sql(sql_query: str):
    try:
        df = DB.execute(sql_query).df()
        return df.to_string(index=False) if not df.empty else "Warning: 0 rows returned. Check filter values."
    except Exception as e:
        return f"Database Error: {str(e)}. Query from 'market_data' table only."


FORMAT_SCHEMA = {
    "type": "object",
    "properties": {
        "action": {"type": "string", "enum": ["get_schema", "execute_sql", "final_answer"]},
        "sql_query": {"type": "string", "description": "The exact SQL string if action is execute_sql"},
        "final_output": {"type": "string", "description": "The final table result if action is final_answer"},
    },
    "required": ["action"],
}


def run_pipeline(user_prompt: str):
    print(f"🚀 Processing: '{user_prompt}'")

    messages = [
        {
            "role": "system",
            "content": "You are a database agent. First choose 'get_schema' to verify column names. Then use 'execute_sql' to query the 'market_data' table. Once you see the successful output from the data table, choose 'final_answer'.",
        },
        {"role": "user", "content": user_prompt},
    ]

    last_table_output = "No data retrieved."

    for _ in range(5):
        response = ollama.chat(
            model="qwen2.5-coder:7b", messages=messages, format=FORMAT_SCHEMA, options={"temperature": 0.0}
        )

        res = json.loads(response["message"]["content"])
        action = res.get("action")

        if action == "final_answer":
            # Safety fallback: If the model leaves 'final_output' blank or writes a sentence,
            # we print the actual dataframe string we captured from the tool.
            output = res.get("final_output")
            if not output or len(output) < 50:
                output = last_table_output
            print(f"\n📊 Final Answer:\n{output}")
            return

        if action == "get_schema":
            print("🛠️ Tool call: get_schema()")
            output = get_schema()
        elif action == "execute_sql":
            query = res.get("sql_query", "")
            print(f"🛠️ Tool call: execute_sql(\n{query}\n)")
            output = execute_sql(query)
            last_table_output = output  # Capture the raw data table right here

        messages.append({"role": "assistant", "content": response["message"]["content"]})
        messages.append({"role": "user", "content": f"Result: {output}"})


if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit("Usage: uv run query.py 'your query'")
    run_pipeline(sys.argv[1])  # Cleaned to pass the string payload correctly