handle multiple codeblocks (#15)

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

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