improve_prompt.py

improve_prompt.py

# /// script
# dependencies = [
#   "claude-agent-sdk",
# ]
# ///

from claude_agent_sdk import (
    ClaudeSDKClient,
    ClaudeAgentOptions,
    AssistantMessage,
    UserMessage,
    SystemMessage,
    ResultMessage,
    PermissionResultAllow,
    ToolPermissionContext,
    query,
)

import json
import os
import signal
import sys
from collections import defaultdict

signal.signal(signal.SIGPIPE, signal.SIG_DFL)




async def auto_approve(
    tool_name: str, tool_input: dict[str, Any], context: ToolPermissionContext
) -> PermissionResultAllow:
    """Auto-approve all tool requests."""
    return PermissionResultAllow()


def extract_session(entries: list[dict]) -> str:
    output = []
    for obj in entries:
        message = obj.get("message")
        if not message:
            continue

        content = message.get("content")

        if isinstance(content, str) and content.strip():
            output.append(f"[USER]\n{content}\n")

        elif isinstance(content, list):
            for item in content:
                item_type = item.get("type")

                if item_type == "thinking":
                    thinking = item.get("thinking", "")
                    if thinking:
                        output.append(f"[THINKING]\n{thinking}\n")

                elif item_type == "text":
                    text = item.get("text", "")
                    if text:
                        output.append(f"[ASSISTANT]\n{text}\n")

                elif item_type == "tool_use" and item.get("name") == "Bash":
                    command = item.get("input", {}).get("command", "")
                    if command:
                        output.append(f"[BASH]\n{command}\n")

    return "\n".join(output)


async def run_query(prompt: str, session_id: Optional[str] = None):
    options = ClaudeAgentOptions(
        can_use_tool=auto_approve,
        system_prompt={"type": "preset", "preset": "claude_code"},
        setting_sources=["user", "project", "local"],  # Load all settings
        resume=session_id,
    )
    async with ClaudeSDKClient(options) as client:
        await client.query(prompt)
        async for message in client.receive_messages():
            if message.get("type") == "result":
                print(message.get("content"))

import argparse

def load_session(filepath: str) -> list[dict]:
    entries = []
    with open(filepath) as f:
        for line in f:
            if line.strip():
                entries.append(json.loads(line))
    return entries


def export_session(session_id: str):
    selected_file = os.path.join(projects_dir, session_id)
    entries = load_session(selected_file)
    content = extract_session(entries)

    output_path = "ctx.txt"
    with open(output_path, "w") as f:
        f.write(content)

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--session-id", type=str, required=True)
    args = parser.parse_args()
    projects_dir = get_projects_dir()
    session_id = args.session_id
    # step 0: first export
    export_session(session_id)
    step_1_prompt = "write a prompt for how you would guide a single opus agent and tell him to read ctx.txt and guide him on how to proceed"
    # step 1: run query
    asyncio.run(run_query(step_1_prompt, session_id))
    
    # step 2: export
    export_session(session_id)
    
    # step 3: iterate
    step_3_prompt = "improve the prompt at the end of ctx.txt.  (assume line numbers are accurate). you can save time for the agent by doing some of the prep work the prompt asks etc"
    for i in range(4):
        asyncio.run(run_query(step_3_prompt, None))