From 1c26afd335f9215661d194997115f554bdbf4feb Mon Sep 17 00:00:00 2001 From: Pierre Date: Mon, 10 Feb 2025 06:23:31 -0600 Subject: [PATCH 1/2] clean_docstring wrap the instructions from the docstring into clean_docstring to clean up the empty lines. --- workflowai/core/client/_fn_utils.py | 24 +++++++++- workflowai/core/client/_fn_utils_test.py | 56 ++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/workflowai/core/client/_fn_utils.py b/workflowai/core/client/_fn_utils.py index 6c29776..514347d 100644 --- a/workflowai/core/client/_fn_utils.py +++ b/workflowai/core/client/_fn_utils.py @@ -147,6 +147,28 @@ async def __call__(self, input: AgentInput, **kwargs: Unpack[RunParams[AgentOutp yield chunk.output +def clean_docstring(docstring: Optional[str]) -> str: + """Clean a docstring by removing empty lines at start/end and normalizing indentation.""" + if not docstring: + return "" + + # Split into lines and remove empty lines at start/end + lines = [line.rstrip() for line in docstring.split("\n")] + while lines and not lines[0].strip(): + lines.pop(0) + while lines and not lines[-1].strip(): + lines.pop() + + if not lines: + return "" + + # Find and remove common indentation + indent = min(len(line) - len(line.lstrip()) for line in lines if line.strip()) + lines = [line[indent:] if line.strip() else "" for line in lines] + + return "\n".join(lines) + + def wrap_run_template( client: Callable[[], APIClient], agent_id: str, @@ -165,7 +187,7 @@ def wrap_run_template( if not version and (fn.__doc__ or model): version = VersionProperties( - instructions=fn.__doc__, + instructions=clean_docstring(fn.__doc__), model=model, ) diff --git a/workflowai/core/client/_fn_utils_test.py b/workflowai/core/client/_fn_utils_test.py index c5cd37f..aa5dd4d 100644 --- a/workflowai/core/client/_fn_utils_test.py +++ b/workflowai/core/client/_fn_utils_test.py @@ -13,6 +13,7 @@ _RunnableStreamAgent, # pyright: ignore [reportPrivateUsage] _RunnableStreamOutputOnlyAgent, # pyright: ignore [reportPrivateUsage] agent_wrapper, + clean_docstring, extract_fn_spec, get_generic_args, is_async_iterator, @@ -113,3 +114,58 @@ async def test_fn_stream_output_only(self, mock_api_client: Mock): assert len(chunks) == 1 assert isinstance(chunks[0], HelloTaskOutput) assert chunks[0] == HelloTaskOutput(message="Hello, World!") + + +def test_clean_docstring(): + # Test empty docstring + assert clean_docstring("") == "" + assert clean_docstring(None) == "" + + # Test single line docstring + assert clean_docstring("Hello world") == "Hello world" + assert clean_docstring(" Hello world ") == "Hello world" + + # Test docstring with empty lines at start/end + assert clean_docstring(""" + + Hello world + + """) == "Hello world" + + # Test multi-line docstring with indentation + assert clean_docstring(""" + First line + Second line + Indented line + Last line + """) == "First line\nSecond line\n Indented line\nLast line" + + # Test docstring with empty lines in between + assert clean_docstring(""" + First line + + Second line + + Third line + """) == "First line\n\nSecond line\n\nThird line" + + # Test real-world example + expected = ( + "Find the capital city of the country where the input city is located.\n\n" + "Guidelines:\n" + "1. First identify the country where the input city is located\n" + "2. Then provide the capital city of that country\n" + "3. Include an interesting historical or cultural fact about the capital\n" + "4. Be accurate and precise with geographical information\n" + "5. If the input city is itself the capital, still provide the information" + ) + assert clean_docstring(""" + Find the capital city of the country where the input city is located. + + Guidelines: + 1. First identify the country where the input city is located + 2. Then provide the capital city of that country + 3. Include an interesting historical or cultural fact about the capital + 4. Be accurate and precise with geographical information + 5. If the input city is itself the capital, still provide the information + """) == expected From 7db7cf4f2e62dd8e842440a98b4bdcf9a8fd4005 Mon Sep 17 00:00:00 2001 From: Pierre Date: Mon, 10 Feb 2025 08:24:26 -0600 Subject: [PATCH 2/2] Update _fn_utils_test.py --- workflowai/core/client/_fn_utils_test.py | 62 +++++++++++++----------- 1 file changed, 33 insertions(+), 29 deletions(-) diff --git a/workflowai/core/client/_fn_utils_test.py b/workflowai/core/client/_fn_utils_test.py index aa5dd4d..58cfcac 100644 --- a/workflowai/core/client/_fn_utils_test.py +++ b/workflowai/core/client/_fn_utils_test.py @@ -1,4 +1,4 @@ -from typing import AsyncIterator +from typing import AsyncIterator, Union from unittest.mock import Mock import pytest @@ -116,50 +116,43 @@ async def test_fn_stream_output_only(self, mock_api_client: Mock): assert chunks[0] == HelloTaskOutput(message="Hello, World!") -def test_clean_docstring(): - # Test empty docstring - assert clean_docstring("") == "" - assert clean_docstring(None) == "" +@pytest.mark.parametrize( + ("value", "expected"), + [ + # Empty docstrings + ("", ""), + (None, ""), - # Test single line docstring - assert clean_docstring("Hello world") == "Hello world" - assert clean_docstring(" Hello world ") == "Hello world" + # Single line docstrings + ("Hello world", "Hello world"), + (" Hello world ", "Hello world"), - # Test docstring with empty lines at start/end - assert clean_docstring(""" + # Docstring with empty lines at start/end + (""" Hello world - """) == "Hello world" + """, "Hello world"), - # Test multi-line docstring with indentation - assert clean_docstring(""" + # Multi-line docstring with indentation + (""" First line Second line Indented line Last line - """) == "First line\nSecond line\n Indented line\nLast line" + """, "First line\nSecond line\n Indented line\nLast line"), - # Test docstring with empty lines in between - assert clean_docstring(""" + # Docstring with empty lines in between + (""" First line Second line Third line - """) == "First line\n\nSecond line\n\nThird line" + """, "First line\n\nSecond line\n\nThird line"), - # Test real-world example - expected = ( - "Find the capital city of the country where the input city is located.\n\n" - "Guidelines:\n" - "1. First identify the country where the input city is located\n" - "2. Then provide the capital city of that country\n" - "3. Include an interesting historical or cultural fact about the capital\n" - "4. Be accurate and precise with geographical information\n" - "5. If the input city is itself the capital, still provide the information" - ) - assert clean_docstring(""" + # Real-world example + (""" Find the capital city of the country where the input city is located. Guidelines: @@ -168,4 +161,15 @@ def test_clean_docstring(): 3. Include an interesting historical or cultural fact about the capital 4. Be accurate and precise with geographical information 5. If the input city is itself the capital, still provide the information - """) == expected + """, + "Find the capital city of the country where the input city is located.\n\n" + "Guidelines:\n" + "1. First identify the country where the input city is located\n" + "2. Then provide the capital city of that country\n" + "3. Include an interesting historical or cultural fact about the capital\n" + "4. Be accurate and precise with geographical information\n" + "5. If the input city is itself the capital, still provide the information"), + ], +) +def test_clean_docstring(value: Union[str, None], expected: str): + assert clean_docstring(value) == expected