Add token options

Author: ddrscott

README.md

@@ -27,6 +27,7 @@ export ANTHROPIC_API_KEY=your_api_key_here
 - 🧪 Powerful file filtering to focus the AI's attention on specific scrolls of code
 - 📚 Handles file output with proper directory creation spells
 - 📥 Supports reading enchantments from stdin
+- 🖼️ Includes image attachment capabilities for visual context
 
 ## 🧙 Installation
 
@@ -85,9 +86,18 @@ Include specific scrolls:
 wiz prompt -f main.py -f utils.py "How can I make these files more efficient?"
 ```
 
+Add images for visual context:
+
+```bash
+wiz prompt -i screenshot.png "What UI improvements would you suggest based on this screenshot?"
+```
+
 Options:
 - `-f, --file`: Specify files to include (can be used multiple times)
+- `-i, --image`: Include image files as context (can be used multiple times)
 - `-o, --output`: Location to write response (default: `.response.md`)
+- `-m, --max-tokens`: Maximum tokens for the response (default: 60000)
+- `-t, --thinking-tokens`: Maximum tokens for Claude's thinking process (default: 16000)
 
 The wizard's response will be saved to `.response.md` by default, and a copy of the full messages including system prompt will be saved to `.messages.md`.
 
@@ -132,6 +142,18 @@ Conjure a new feature:
 wiz prompt "Add a progress bar to the file processing function"
 ```
 
+Get feedback on UI design with an image:
+
+```bash
+wiz prompt -i design.png -f styles.css "How can I improve this layout?"
+```
+
+Adjust token limits for complex analysis:
+
+```bash
+wiz prompt -m 80000 -t 20000 "Perform a comprehensive security audit of this codebase"
+```
+
 Then cast the spell to apply the changes:
 
 ```bash

wiz.py

@@ -118,7 +118,8 @@ def parse_attachment(attachment):
                 },
             }
 
-def reply(question, files=None, attachments=None):
+def reply(question, files=None, attachments=None, max_tokens=60000, thinking_tokens=16000):
+    attachments = attachments or []  # Ensure attachments is a list
 
     if files:
         console.print(f"Using specified files: {', '.join(files)}")
@@ -131,6 +132,7 @@ def reply(question, files=None, attachments=None):
     table.add_column("File", style="cyan")
     table.add_column("Size", justify="right", style="green")
     table.add_column("Status", style="yellow")
+    table.add_column("Type", style="magenta")  # Add type column to differentiate files and attachments
 
     body = [f"Help me with following files: {', '.join(file_list)}"]
 
@@ -154,12 +156,25 @@ def reply(question, files=None, attachments=None):
 
                     # Add to table
                     size_str = f"{file_size / 1024:.1f} KB" if file_size > 1024 else f"{file_size} bytes"
-                    table.add_row(file, size_str, "✓ Read")
+                    table.add_row(file, size_str, "✓ Read", "Text File")
             except Exception as e:
-                table.add_row(file, "N/A", str(e))
+                table.add_row(file, "N/A", str(e), "Error")
 
             progress.update(task, advance=1)
 
+    # Add attachments to the table
+    for attachment in attachments:
+        try:
+            if attachment.startswith("http"):
+                table.add_row(attachment, "URL", "✓ Included", "Remote Image")
+            else:
+                file_size = os.path.getsize(attachment)
+                size_str = f"{file_size / 1024:.1f} KB" if file_size > 1024 else f"{file_size} bytes"
+                ext = os.path.splitext(attachment)[1][1:].upper() or "Unknown"
+                table.add_row(attachment, size_str, "✓ Included", f"Image ({ext})")
+        except Exception as e:
+            table.add_row(attachment, "N/A", f"Error: {str(e)}", "Image Error")
+
     # Show summary table
     console.print(table)
 
@@ -169,7 +184,6 @@ def reply(question, files=None, attachments=None):
 **Reminder**
 - wrap resulting code between `[{TAG}]` and `[/{TAG}]` tags!!!
 """
-    attachments = attachments or []
     images = [
         parse_attachment(att)
         for att in attachments
@@ -185,7 +199,7 @@ def reply(question, files=None, attachments=None):
             ]
         }
     ]
-    open('.messages.md', 'w').write(system + "\n---\n" + body)
+    open('.messages.md', 'w').write(system + "\n---\\n" + body)
 
     console.print("[bold yellow]Question:[/bold yellow]")
     console.print(question)
@@ -200,13 +214,13 @@ def reply(question, files=None, attachments=None):
     # Stream response through a Live display
     with client.messages.stream(
         model="claude-3-7-sonnet-20250219",
-        max_tokens=60000,
+        max_tokens=max_tokens,
         temperature=1,
         system=system,
         messages=messages, # type: ignore
         thinking={
             "type": "enabled",
-            "budget_tokens": 16000
+            "budget_tokens": thinking_tokens,
         }
     ) as stream:
         for event in stream:
@@ -291,12 +305,14 @@ def cli():
 @click.option('--file', '-f', help='Files to include in the question', multiple=True)
 @click.option('--image', '-i', help='Image to include in the question', multiple=True)
 @click.option('--output', '-o', help='location write response without thoughts', default='.response.md')
-def prompt(question_text, file, output, image):
+@click.option('--max-tokens', '-m', help='Max tokens for the response', default=60000)
+@click.option('--thinking-tokens', '-t', help='Max tokens for the thinking', default=16000)
+def prompt(question_text, file, output, image, max_tokens, thinking_tokens):
     question = ' '.join(question_text)
 
     if question:
         try:
-            response = reply(question, files=file, attachments=image)
+            response = reply(question, files=file, attachments=image, max_tokens=max_tokens)
             with open(output, 'w') as f:
                 f.write(response)
             console.print(f"[bold green]Output written to {escape(output)}[/bold green]")