File size: 4,298 Bytes
b2c86fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | import re
from biomni.llm import get_llm
class base_agent:
def __init__(self, llm="claude-3-haiku-20240307", cheap_llm=None, tools=None, temperature=0.7):
self.tools = tools
self.llm = get_llm(llm, temperature)
if cheap_llm is None:
self.cheap_llm = llm
else:
self.cheap_llm = cheap_llm
def configure(self):
pass
def go(self, input):
pass
class FunctionGenerator(base_agent):
"""Agent that generates executable Python code scripts given a task description."""
def __init__(self, llm="claude-3-7-sonnet-20250219", cheap_llm=None, temperature=0.7):
"""Initialize the PaperTaskExtractor agent.
Args:
llm (str): The LLM model to use
cheap_llm (str, optional): A cheaper LLM for simpler tasks
"""
super().__init__(llm, cheap_llm, temperature)
self.log = []
self.configure()
def configure(self):
"""Configure the agent with appropriate prompts."""
# Prompt for Python code generation
self.system_prompt = """You are a senior Python engineer. Generate robust, idiomatic Python code that solves the user's task. Requirements:
1. Output ONLY Python code, ideally inside a single triple-backtick code block.
2. Include minimal inline comments and a small docstring.
3. Add a `main()` and an `if __name__ == '__main__':` guard when appropriate.
4. Avoid external dependencies unless necessary; if used, show `pip` installs in comments.
5. Do not include prose before or after the code.
6. When applicable, prioritize the use of codes on public repositories, such as HuggingFace or Github
Generate Python codes for the following task:
{task}
"""
def _generate_code(self, task_description: str) -> str:
"""Generate codes given a task description.
Args:
task_description (str): task descriptions (possibly generated from previous steps)
Returns:
str: generated code string
"""
prompt = self.system_prompt.format(task=task_description)
message = self.llm.invoke(prompt)
return message.content
def _generate_script_filename(self, task_description: str, max_words: int = 6) -> str:
"""
Generate a safe, meaningful Python script filename from a task description.
Poised for update: may ask the agent to suggest meaningful names.
Parameters:
-----------
task_description (str): task descriptions (possibly generated from previous steps)
max_words : int
Maximum number of words to include in the filename.
Returns:
--------
str
A lowercase, hyphen-free, safe filename ending in '.py'.
"""
# Lowercase and remove non-alphanumeric (allow spaces for splitting)
cleaned = re.sub(r"[^a-zA-Z0-9\s]", "", task_description.lower())
# Tokenize and select top words
words = cleaned.split()
selected_words = words[:max_words] if words else ["script"]
# Join with underscores
base_name = "_".join(selected_words)
return f"{base_name}.py"
def go(self, task_description: str):
"""Implement the inherited function to get the tasks done.
Args:
task_description (str): task descriptions (possibly generated from previous steps)
Returns:
tuple: (script_filename, results) where script_filename is a generated name for script file and results is the generated codes
"""
self.log = []
self.log.append(
(
"user",
"Generate Python codes given a task description",
)
)
script_filename = self._generate_script_filename(task_description)
results = self._generate_code(task_description)
return script_filename, self._extract_code_block(results)
def _extract_code_block(self, s: str) -> str:
"""
Extract the first fenced code block (``` or ```python) from s.
"""
m = re.search(r"```(?:python)?\s*(.+?)\s*```", s, flags=re.DOTALL | re.IGNORECASE)
return m.group(1).strip() if m else None
|