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..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 @@ -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,62 @@ 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!") + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + # Empty docstrings + ("", ""), + (None, ""), + + # Single line docstrings + ("Hello world", "Hello world"), + (" Hello world ", "Hello world"), + + # Docstring with empty lines at start/end + (""" + + Hello world + + """, "Hello world"), + + # Multi-line docstring with indentation + (""" + First line + Second line + Indented line + Last line + """, "First line\nSecond line\n Indented line\nLast line"), + + # Docstring with empty lines in between + (""" + First line + + Second line + + Third line + """, "First line\n\nSecond line\n\nThird line"), + + # Real-world example + (""" + 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 + """, + "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