From 33d4cef00b266099a8d7e6d7eb53987f76d4b96b Mon Sep 17 00:00:00 2001
From: Dhanji Prasanna <dhanji@squareup.com>
Date: Sat, 6 Sep 2025 14:49:33 +1000
Subject: [PATCH] Difficulty rating

---
 crates/g3-core/src/lib.rs      | 20 +++++++++++---------
 crates/g3-execution/src/lib.rs | 26 +++++++++++++++++++++++---
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/crates/g3-core/src/lib.rs b/crates/g3-core/src/lib.rs
index 9257220..f09e409 100644
--- a/crates/g3-core/src/lib.rs
+++ b/crates/g3-core/src/lib.rs
@@ -5,9 +5,9 @@ use g3_providers::{CompletionRequest, Message, MessageRole, ProviderRegistry};
 use serde::{Deserialize, Serialize};
 use std::path::Path;
 use std::time::{Duration, Instant};
+use tokio_util::sync::CancellationToken;
 use tracing::field::debug;
 use tracing::info;
-use tokio_util::sync::CancellationToken;
 
 #[derive(Debug, Clone)]
 pub struct ContextWindow {
@@ -97,8 +97,8 @@ impl Agent {
         // Initialize context window with configured max context length
         let context_window = ContextWindow::new(config.agent.max_context_length as u32);
 
-        Ok(Self { 
-            providers, 
+        Ok(Self {
+            providers,
             config,
             context_window,
         })
@@ -173,14 +173,15 @@ impl Agent {
         let provider = self.providers.get(None)?;
 
         let system_prompt = format!(
-            "You are G3, a code-first AI agent. Your goal is to solve problems by writing and executing code autonomously.
+            "You are G3, a code-first AI agent. Your goal is to solve problems by writing code that completes the desired task.
 
 When given a task:
 1. Analyze what needs to be done
-2. Choose the most appropriate programming language{}
-3. Include any necessary imports/dependencies
-4. Add error handling where appropriate
-5. EXECUTE the code immediately to solve the user's problem
+2. Rate the difficulty of the task from 1 (easy, file operations) to 10 (difficult, build complex applications like Firefox)
+3. Choose the most appropriate programming language{}
+4. Include any necessary imports/dependencies
+5. Add error handling where appropriate
+6. Generate code to complete the task, or ask for more details, but no other output
 
 Prefer these languages:
 - Bash/Shell: File operations, system administration, simple tasks
@@ -189,7 +190,8 @@ Prefer these languages:
 
 Only use Rust/Python when you need libraries or complex logic that bash can't handle easily.
 
-Format your response as:
+Format your code response in markdown backticks as follows:
+difficulty rating: [X]
 ```[language]
 [code]
 ```
diff --git a/crates/g3-execution/src/lib.rs b/crates/g3-execution/src/lib.rs
index f18945d..ccb6f25 100644
--- a/crates/g3-execution/src/lib.rs
+++ b/crates/g3-execution/src/lib.rs
@@ -90,24 +90,44 @@ impl CodeExecutor {
     
     /// Extract code blocks from markdown-formatted text
     fn extract_code_blocks(&self, text: &str) -> Result<Vec<(String, String)>> {
-        let re = Regex::new(r"(?s)```(\w+)?\n(.*?)```")?;
         let mut blocks = Vec::new();
         
         debug!("Extracting code blocks from text: {}", text);
         
-        for cap in re.captures_iter(text) {
+        // Pattern 1: Standard markdown format ```language\ncode```
+        let markdown_re = Regex::new(r"(?s)```(\w+)?\n(.*?)```")?;
+        for cap in markdown_re.captures_iter(text) {
             let language = cap.get(1)
                 .map(|m| m.as_str().to_lowercase())
                 .unwrap_or_else(|| "bash".to_string()); // Default to bash
             let code = cap.get(2).map(|m| m.as_str()).unwrap_or("").trim();
             
-            debug!("Found code block - language: '{}', code: '{}'", language, code);
+            debug!("Found markdown code block - language: '{}', code: '{}'", language, code);
             
             if !code.is_empty() {
                 blocks.push((language, code.to_string()));
             }
         }
         
+        // Pattern 2: Bracket format [Language]code[/Language]
+        let bracket_re = Regex::new(r"(?s)\[(\w+)\]\s*(.*?)\s*\[/(\w+)\]")?;
+        for cap in bracket_re.captures_iter(text) {
+            let open_lang = cap.get(1).map(|m| m.as_str()).unwrap_or("");
+            let close_lang = cap.get(3).map(|m| m.as_str()).unwrap_or("");
+            
+            // Only match if opening and closing tags are the same (case insensitive)
+            if open_lang.to_lowercase() == close_lang.to_lowercase() {
+                let language = open_lang.to_lowercase();
+                let code = cap.get(2).map(|m| m.as_str()).unwrap_or("").trim();
+                
+                debug!("Found bracket code block - language: '{}', code: '{}'", language, code);
+                
+                if !code.is_empty() {
+                    blocks.push((language, code.to_string()));
+                }
+            }
+        }
+        
         debug!("Total code blocks found: {}", blocks.len());
         Ok(blocks)
     }