From f8cca3a4e458e1f416893d65cb7c18fa19f48a70 Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Sat, 5 Apr 2025 12:24:48 -0400 Subject: [PATCH] handle multiple codeblocks (#15) --- langgraph_codeact/__init__.py | 10 +- langgraph_codeact/utils.py | 61 ++++++++++++ tests/test_utils.py | 175 ++++++++++++++++++++++++++++++++++ 3 files changed, 241 insertions(+), 5 deletions(-) create mode 100644 langgraph_codeact/utils.py create mode 100644 tests/test_utils.py diff --git a/langgraph_codeact/__init__.py b/langgraph_codeact/__init__.py index c55900f..4d25ab5 100644 --- a/langgraph_codeact/__init__.py +++ b/langgraph_codeact/__init__.py @@ -7,6 +7,8 @@ from langchain_core.tools import tool as create_tool from langgraph.graph import END, START, MessagesState, StateGraph from langgraph.types import Command +from langgraph_codeact.utils import extract_and_combine_codeblocks + class CodeActState(MessagesState): """State for CodeAct agent.""" @@ -75,11 +77,9 @@ def create_codeact( def call_model(state: CodeActState) -> Command: messages = [{"role": "system", "content": prompt}] + state["messages"] response = model.invoke(messages) - if "```" in response.content: - # get content between fences - code = response.content.split("```")[1] - # remove first line, which is the language or empty string - code = "\n".join(code.splitlines()[1:]) + # Extract and combine all code blocks + code = extract_and_combine_codeblocks(response.content) + if code: return Command(goto="sandbox", update={"messages": [response], "script": code}) else: # no code block, end the loop and respond to the user diff --git a/langgraph_codeact/utils.py b/langgraph_codeact/utils.py new file mode 100644 index 0000000..3ca07a9 --- /dev/null +++ b/langgraph_codeact/utils.py @@ -0,0 +1,61 @@ +import re + +BACKTICK_PATTERN = r"(?:^|\n)```(.*?)(?:```(?:\n|$))" + + +def extract_and_combine_codeblocks(text: str) -> str: + """ + Extracts all codeblocks from a text string and combines them into a single code string. + + Args: + text: A string containing zero or more codeblocks, where each codeblock is + surrounded by triple backticks (```). + + Returns: + A string containing the combined code from all codeblocks, with each codeblock + separated by a newline. + + Example: + text = '''Here's some code: + + ```python + print('hello') + ``` + And more: + + ``` + print('world') + ```''' + + result = extract_and_combine_codeblocks(text) + + Result: + + print('hello') + + print('world') + """ + # Find all code blocks in the text using regex + # Pattern matches anything between triple backticks, with or without a language identifier + code_blocks = re.findall(BACKTICK_PATTERN, text, re.DOTALL) + + if not code_blocks: + return "" + + # Process each codeblock + processed_blocks = [] + for block in code_blocks: + # Strip leading and trailing whitespace + block = block.strip() + + # If the first line looks like a language identifier, remove it + lines = block.split("\n") + if lines and (not lines[0].strip() or " " not in lines[0].strip()): + # First line is empty or likely a language identifier (no spaces) + block = "\n".join(lines[1:]) + + processed_blocks.append(block) + + # Combine all codeblocks with newlines between them + combined_code = "\n\n".join(processed_blocks) + return combined_code diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..a94241c --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,175 @@ +from langgraph_codeact.utils import extract_and_combine_codeblocks + + +def test_empty_text(): + """Test when the input text has no codeblocks.""" + text = "This is a text without any code blocks." + result = extract_and_combine_codeblocks(text) + assert result == "" + + +def test_single_codeblock_no_language(): + """Test extracting a single codeblock without language identifier.""" + text = """Here is a code block: +``` +print("Hello, world!") +x = 10 +``` +End of the code.""" + + expected = """\ +print("Hello, world!") +x = 10\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected + + +def test_single_codeblock_with_language(): + """Test extracting a single codeblock with language identifier.""" + text = """Here is a code block: +```python +print("Hello, world!") +x = 10 +``` +End of the code.""" + + expected = """\ +print("Hello, world!") +x = 10\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected + + +def test_multiple_codeblocks(): + """Test extracting and combining multiple codeblocks.""" + text = """Here's the first code block: +```python +def hello(): + print("Hello!") +``` + +And here's the second one: +```python +result = 42 +print(f"The answer is {result}") +```""" + + expected = """\ +def hello(): + print("Hello!") + +result = 42 +print(f"The answer is {result}")\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected + + +def test_multiple_codeblocks_mixed(): + """Test codeblocks with a mix of language identifiers / no identifiers.""" + text = """Different language identifiers: +```python +x = 10 +``` + +```python +y = 20 +``` + +``` +z = 30 +```""" + + expected = """\ +x = 10 + +y = 20 + +z = 30\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected + + +def test_empty_codeblock(): + """Test an empty codeblock.""" + text = "Empty block: `````` should be ignored." + result = extract_and_combine_codeblocks(text) + assert result == "" + + +def test_language_with_spaces(): + """Test a codeblock with a language identifier containing spaces.""" + text = """Here is code with a more unusual language tag: +```python code +x = 10 +y = 20 +```""" + + # The first line shouldn't be removed since it contains spaces + expected = """\ +python code +x = 10 +y = 20\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected + + +def test_with_nested_backticks(): + """Test with nested backticks inside the code block.""" + text = """Code with nested backticks: +``` +def example(): + code = "```nested```" + return code +```""" + + expected = """\ +def example(): + code = "```nested```" + return code\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected + + +def test_realistic_example(): + """Test with a realistic example similar to the one provided in the user query.""" + text = """First, I'll find where the baseball lands when hit by the batter. Then, I'll calculate where the ball lands after being thrown by the outfielder. + +```python +# Constants +g = 9.81 # acceleration due to gravity +v0_batter = 45.847 # initial velocity +angle_batter_deg = 23.474 # angle in degrees + +print(f"The ball lands {distance:.2f} meters away") +``` + +Now, let's calculate the second trajectory: + +``` +# Outfielder's throw +v0_outfielder = 24.12 # initial velocity +distance_2 = v0_outfielder * 2 # simplified calculation +print(f"Final position: {distance_2:.2f} meters") +```""" + + expected = """\ +# Constants +g = 9.81 # acceleration due to gravity +v0_batter = 45.847 # initial velocity +angle_batter_deg = 23.474 # angle in degrees + +print(f"The ball lands {distance:.2f} meters away") + +# Outfielder's throw +v0_outfielder = 24.12 # initial velocity +distance_2 = v0_outfielder * 2 # simplified calculation +print(f"Final position: {distance_2:.2f} meters")\ +""" + result = extract_and_combine_codeblocks(text) + assert result == expected