albert-einstein-09/codedark
3
1"""2CodeDark Tool Implementations3 4Tools available to agents:5- run_python: Execute Python/pandas code in sandboxed environment6- read_notes: Read all saved notes from current episode7- save_note: Save a note for later recall8- clarify: Ask clarifying question (max 2 per episode)9- submit_answer: Submit final answer (ends episode)10"""11 12import re13import ast14from typing import Any, Dict, List, Optional, Tuple15 16import pandas as pd17import numpy as np18 19 20# Safe builtins for sandboxed code execution21SAFE_BUILTINS = {22 "len": len,23 "sum": sum,24 "min": min,25 "max": max,26 "abs": abs,27 "round": round,28 "sorted": sorted,29 "range": range,30 "int": int,31 "float": float,32 "str": str,33 "bool": bool,34 "list": list,35 "dict": dict,36 "set": set,37 "tuple": tuple,38 "enumerate": enumerate,39 "zip": zip,40 "True": True,41 "False": False,42 "None": None,43 "print": print,44 "type": type,45 "isinstance": isinstance,46 "map": map,47 "filter": filter,48 "any": any,49 "all": all,50 "hasattr": hasattr,51 "getattr": getattr,52 "repr": repr,53 "locals": locals,54 "globals": globals,55 "dir": dir,56 "vars": vars,57 "reversed": reversed,58 "slice": slice,59 "format": format,60 "Exception": Exception,61 "ValueError": ValueError,62 "TypeError": TypeError,63 "KeyError": KeyError,64 "IndexError": IndexError,65 "AttributeError": AttributeError,66}67 68 69def run_python(70 code: str, df: pd.DataFrame, max_output_chars: int = 20071) -> Tuple[str, str, int]:72 """Execute Python code in sandboxed environment.73 74 Args:75 code: Python code to execute76 df: DataFrame available as 'df' in execution context77 max_output_chars: Maximum characters for output truncation78 79 Returns:80 Tuple of (stdout, stderr, exit_code)81 """82 if df is None:83 return "", "Error: No dataframe loaded", 184 85 local_vars = {86 "pd": pd,87 "np": np,88 "df": df.copy(),89 }90 91 try:92 exec(code, {"__builtins__": SAFE_BUILTINS}, local_vars)93 result = local_vars.get("result")94 95 if result is None:96 return (97 "",98 "Error: No 'result' variable set. Store your result in 'result'.",99 1,100 )101 102 # Format output with truncation103 if isinstance(result, pd.DataFrame):104 preview = result.head(3).to_string()105 elif isinstance(result, pd.Series):106 preview = result.head(5).to_string()107 else:108 preview = str(result)109 110 # Truncate if needed111 if len(preview) > max_output_chars:112 preview = preview[:max_output_chars] + "..."113 114 return f"run_python Result:\n{preview}", "", 0115 116 except Exception as e:117 return "", f"run_python Error: {e}", 1118 119 120def read_notes(notes: List[str]) -> Tuple[str, str, int]:121 """Read all saved notes.122 123 Args:124 notes: List of saved notes125 126 Returns:127 Tuple of (stdout, stderr, exit_code)128 """129 if not notes:130 return "No notes saved yet.", "", 0131 132 notes_list = "\n".join(f"- {n}" for n in notes)133 return f"Saved notes:\n{notes_list}", "", 0134 135 136def save_note(content: str, notes: List[str]) -> Tuple[str, str, int]:137 """Save a note to persistent memory.138 139 Args:140 content: Note content to save141 notes: List to append note to (modified in place)142 143 Returns:144 Tuple of (stdout, stderr, exit_code)145 """146 content = content.strip()147 if not content:148 return "", "Error: Empty note content", 1149 150 notes.append(content)151 notes_list = "\n".join(f"- {n}" for n in notes)152 return f"Note saved.\n\nAll notes:\n{notes_list}", "", 0153 154 155def clarify(156 question: str,157 clarify_count: int,158 max_clarifications: int,159 ambiguities: Optional[List[str]] = None,160 answer_type: str = "scalar",161) -> Tuple[str, str, int, int]:162 """Ask a clarifying question about the task.163 164 Args:165 question: The clarifying question166 clarify_count: Current number of clarifications used167 max_clarifications: Maximum allowed clarifications168 ambiguities: List of known ambiguities from task metadata169 answer_type: Expected answer type ("scalar", "list", etc.)170 171 Returns:172 Tuple of (stdout, stderr, exit_code, new_clarify_count)173 """174 if clarify_count >= max_clarifications:175 return (176 "",177 f"Error: Maximum {max_clarifications} clarifications per episode. Please proceed with your best interpretation.",178 1,179 clarify_count,180 )181 182 question_lower = question.lower()183 ambiguities = ambiguities or []184 185 # Build clarification responses from task metadata186 clarifications = {}187 188 for amb in ambiguities:189 amb_lower = amb.lower()190 if (191 "percentile" in amb_lower192 or "inclusive" in amb_lower193 or "exclusive" in amb_lower194 ):195 clarifications["percentile"] = (196 "Use >= for 'top X%' (inclusive of threshold) and <= for 'bottom X%'."197 )198 if "rate" in amb_lower or "percentage" in amb_lower:199 clarifications["rate"] = (200 "Express rates as percentages 0-100, rounded to 2 decimal places."201 )202 if (203 "positive" in amb_lower204 or "success" in amb_lower205 or "target" in amb_lower206 or "y=1" in amb_lower207 ):208 clarifications["target"] = "Subscription/success means y=1 in the dataset."209 if "boundary" in amb_lower:210 clarifications["boundary"] = (211 "Include boundary values (>=, <=) when filtering."212 )213 214 # Add format clarifications from answer type215 if answer_type == "scalar":216 clarifications["format"] = (217 "Return a single numeric value, rounded to 2 decimal places."218 )219 elif answer_type == "list":220 clarifications["format"] = (221 "Return as a list/DataFrame with the specified columns."222 )223 224 # Try to match question to a clarification225 response = None226 for key, value in clarifications.items():227 if key in question_lower or any(word in question_lower for word in key.split()):228 response = value229 break230 231 if response:232 return f"Clarification: {response}", "", 0, clarify_count + 1233 else:234 return (235 "Clarification: Please proceed with your best interpretation based on standard data analysis conventions.",236 "",237 0,238 clarify_count + 1,239 )240 241 242def submit_answer(answer_str: str) -> Tuple[str, str, int, Any]:243 """Submit final answer.244 245 Args:246 answer_str: Answer string to parse and submit247 248 Returns:249 Tuple of (stdout, stderr, exit_code, parsed_answer)250 """251 answer_str = answer_str.strip().rstrip("%").strip()252 253 if not answer_str:254 return "", "Error: Empty answer", 1, None255 256 # Try to parse as structured data first (list/dict)257 try:258 answer = ast.literal_eval(answer_str)259 except (ValueError, SyntaxError):260 # Fall back to numeric parsing261 try:262 answer = float(answer_str)263 except ValueError:264 answer = answer_str265 266 return "[SUBMITTED]", "", 0, answer267 268 269def parse_tool_call(args: str, tool_name: str) -> Tuple[Optional[str], Optional[str]]:270 """Parse tool-specific arguments from args string.271 272 Args:273 args: Raw args string274 tool_name: Name of the tool being called275 276 Returns:277 Tuple of (parsed_content, error_message)278 """279 if tool_name == "run_python":280 # Extract code from <code></code> tags281 match = re.search(r"<code>(.*?)</code>", args, re.DOTALL)282 if not match:283 return None, "No <code> tag found. Use: <code>your_code</code>"284 return match.group(1).strip(), None285 286 elif tool_name == "clarify":287 # Extract question from <question></question> tags288 match = re.search(r"<question>(.*?)</question>", args, re.DOTALL)289 if not match:290 return (291 None,292 "No <question> tag found. Use: <question>your question</question>",293 )294 return match.group(1).strip(), None295 296 elif tool_name == "submit_answer":297 # Extract answer from <answer></answer> tags298 match = re.search(r"<answer>(.*?)</answer>", args, re.DOTALL)299 if not match:300 return None, "No <answer> tag found. Use: <answer>value</answer>"301 return match.group(1).strip(), None302 303 elif tool_name in ("read_notes", "save_note"):304 # These take raw args305 return args.strip(), None306 307 else:308 return None, f"Unknown tool: {tool_name}"309