scripts/quant.py

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

import re
import sys
import traceback

import duckdb
import numpy as np
import ollama
import pandas as pd

DB = duckdb.connect()
PARQUET_FILE = "data/alpaca_merged.parquet"


def execute_backtest(python_code: str) -> str:
    """Executes the generated python code string natively."""
    import io
    import sys as native_sys

    exec_globals = {"db": DB, "parquet_file": PARQUET_FILE, "pd": pd, "np": np}

    old_stdout = native_sys.stdout
    redirected_output = io.StringIO()
    native_sys.stdout = redirected_output

    try:
        exec(python_code, exec_globals)
        native_sys.stdout = old_stdout
        captured_text = redirected_output.getvalue().strip()
        return captured_text if captured_text else "Success! Code executed but did not print any output metrics."
    except Exception:
        native_sys.stdout = old_stdout
        return f"Execution Error:\n{traceback.format_exc()}"


def run_quant_search(user_request: str):
    print(f"📈 Initiating Direct Alpha Search for: '{user_request}'\n")

    try:
        schema_info = DB.execute(f"DESCRIBE SELECT * FROM read_parquet('{PARQUET_FILE}');").df().to_string(index=False)
    except Exception as e:
        sys.exit(f"❌ Could not open file at '{PARQUET_FILE}': {e}")

    system_prompt = f"""
You are an expert algorithmic quant researcher writing code compatible with modern Pandas 3.0+ (Copy-on-Write permanently active).
Your goal is to write a single standalone Python script block that reads market data, runs a backtest, and PRINTS the final return metrics.

[DATA ENVIRONMENT]
- DuckDB connection object: `db`
- Target file path string variable: `parquet_file`
- Dataset Column Schemas:
{schema_info}

[STRICT MODERN CODING LAWS]
1. Load data via: df = db.execute(f"SELECT * FROM read_parquet('{{parquet_file}}') ORDER BY timestamp;").df()
2. Convert dates cleanly: df['timestamp'] = pd.to_datetime(df['timestamp'])
3. CRITICAL: Pandas Copy-on-Write is permanently enabled. Chained assignment (e.g., df['signal'][window:] = ...) will throw a fatal error.
4. To set values based on conditions, you MUST use `.loc` in a single step (e.g., `df.loc[condition, 'column_name'] = value`) or use `np.where()`. Never assign directly onto a sliced dataframe view.
5. Use high-speed vectorized pandas/numpy computations. No slow row-based loops.
6. You do NOT need to import pandas or numpy. They are pre-injected.
7. You MUST use 'print()' statements at the end to output parameters and strategy return values.
8. Output ONLY the code inside a markdown block: ```python ... ```
"""

    response = ollama.generate(
        model="qwen2.5-coder:7b", system=system_prompt, prompt=user_request, options={"temperature": 0.0}
    )

    raw_content = response["response"]
    code_match = re.search(r"```python\n([\s\S]*?)\n```", raw_content)
    extracted_code = code_match.group(1).strip() if code_match else raw_content.strip()

    print("🖥️ Generated Backtest Engine Code:")
    print("-" * 60)
    print(extracted_code)
    print("-" * 60 + "\n")

    print("⚙️ Running matrix simulation against Parquet data...")
    report = execute_backtest(extracted_code)

    print("\n📊 Strategy Execution Report Output:")
    print("=" * 60)
    print(report)
    print("=" * 60)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit("Usage: uv run scripts/quant.py 'strategy description'")
    # Correctly grab the text payload from command line
    run_quant_search(sys.argv[1])