| 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.""" |
| |
| 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'. |
| """ |
| |
| cleaned = re.sub(r"[^a-zA-Z0-9\s]", "", task_description.lower()) |
|
|
| |
| words = cleaned.split() |
| selected_words = words[:max_words] if words else ["script"] |
|
|
| |
| 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 |
|
|
