feat: response_parser both sync and async (#432)

Author: kzamlynska

packages/ragbits-core/CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
+- Allow Prompt class to accept the asynchronous response_parser. Change the signature of parse_response method.
 
 ## 0.11.0 (2025-03-25)
 - Add HybridSearchVectorStore which can aggregate results from multiple VectorStores (#412)

packages/ragbits-core/src/ragbits/core/llms/base.py

@@ -163,7 +163,7 @@ async def generate(
         with trace(model_name=self.model_name, prompt=prompt, options=repr(options)) as outputs:
             raw_response = await self.generate_raw(prompt, options=options)
             if isinstance(prompt, BasePromptWithParser):
-                response = prompt.parse_response(raw_response["response"])
+                response = await prompt.parse_response(raw_response["response"])
             else:
                 response = cast(OutputT, raw_response["response"])
             raw_response["response"] = response
@@ -225,7 +225,7 @@ async def generate_with_metadata(
             response = await self.generate_raw(prompt, options=options)
             content = response.pop("response")
             if isinstance(prompt, BasePromptWithParser):
-                content = prompt.parse_response(content)
+                content = await prompt.parse_response(content)
             outputs.response = LLMResponseWithMetadata[type(content)](  # type: ignore
                 content=content,
                 metadata=response,

packages/ragbits-core/src/ragbits/core/prompt/base.py

@@ -54,7 +54,7 @@ class BasePromptWithParser(Generic[OutputT], BasePrompt, metaclass=ABCMeta):
     """
 
     @abstractmethod
-    def parse_response(self, response: str) -> OutputT:
+    async def parse_response(self, response: str) -> OutputT:
         """
         Parse the response from the LLM to the desired output type.
 

packages/ragbits-core/src/ragbits/core/prompt/prompt.py

@@ -1,8 +1,9 @@
+import asyncio
 import base64
 import imghdr
 import textwrap
 from abc import ABCMeta
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from typing import Any, Generic, cast, get_args, get_origin, overload
 
 from jinja2 import Environment, Template, meta
@@ -34,7 +35,7 @@ class Prompt(Generic[InputT, OutputT], BasePromptWithParser[OutputT], metaclass=
 
     # function that parses the response from the LLM to specific output type
     # if not provided, the class tries to set it automatically based on the output type
-    response_parser: Callable[[str], OutputT]
+    response_parser: Callable[[str], OutputT | Awaitable[OutputT]]
 
     # Automatically set in __init_subclass__
     input_type: type[InputT] | None
@@ -98,7 +99,7 @@ def _format_message(cls, message: str) -> str:
         return textwrap.dedent(message).strip()
 
     @classmethod
-    def _detect_response_parser(cls) -> Callable[[str], OutputT]:
+    def _detect_response_parser(cls) -> Callable[[str], OutputT | Awaitable[OutputT]]:
         if hasattr(cls, "response_parser") and cls.response_parser is not None:
             return cls.response_parser
         if issubclass(cls.output_type, BaseModel):
@@ -265,7 +266,7 @@ def json_mode(self) -> bool:
         """
         return issubclass(self.output_type, BaseModel)
 
-    def parse_response(self, response: str) -> OutputT:
+    async def parse_response(self, response: str) -> OutputT:
         """
         Parse the response from the LLM to the desired output type.
 
@@ -278,7 +279,11 @@ def parse_response(self, response: str) -> OutputT:
         Raises:
             ResponseParsingError: If the response cannot be parsed.
         """
-        return self.response_parser(response)
+        if asyncio.iscoroutinefunction(self.response_parser):
+            result = await self.response_parser(response)
+        else:
+            result = self.response_parser(response)
+        return result
 
     @classmethod
     def to_promptfoo(cls, config: dict[str, Any]) -> ChatFormat:

packages/ragbits-core/tests/unit/llms/test_base.py

@@ -28,7 +28,7 @@ def chat(self) -> ChatFormat:
         return [{"role": "user", "content": self._content}]
 
     @staticmethod
-    def parse_response(response: str) -> CustomOutputType:
+    async def parse_response(response: str) -> CustomOutputType:
         return CustomOutputType(message=response)
 
 

packages/ragbits-core/tests/unit/llms/test_litellm.py

@@ -57,7 +57,7 @@ def chat(self) -> ChatFormat:
         return [{"content": self.message, "role": "user"}]
 
     @staticmethod
-    def parse_response(response: str) -> int:
+    async def parse_response(response: str) -> int:
         """
         Parser for the prompt.
 

packages/ragbits-core/tests/unit/prompts/test_parsers.py

@@ -8,7 +8,7 @@
 from .test_prompt import _PromptOutput
 
 
-def test_prompt_with_str_output():
+async def test_prompt_with_str_output():
     """Test a prompt with a string output."""
 
     class TestPrompt(Prompt[None, str]):  # pylint: disable=unused-variable
@@ -17,10 +17,10 @@ class TestPrompt(Prompt[None, str]):  # pylint: disable=unused-variable
         user_prompt = "Hello"
 
     prompt = TestPrompt()
-    assert prompt.parse_response("Hi") == "Hi"
+    assert await prompt.parse_response("Hi") == "Hi"
 
 
-def test_prompt_with_int_output():
+async def test_prompt_with_int_output():
     """Test a prompt with an int output."""
 
     class TestPrompt(Prompt[None, int]):  # pylint: disable=unused-variable
@@ -29,13 +29,13 @@ class TestPrompt(Prompt[None, int]):  # pylint: disable=unused-variable
         user_prompt = "Hello"
 
     prompt = TestPrompt()
-    assert prompt.parse_response("1") == 1
+    assert await prompt.parse_response("1") == 1
 
     with pytest.raises(ResponseParsingError):
-        prompt.parse_response("a")
+        await prompt.parse_response("a")
 
 
-def test_prompt_with_model_output():
+async def test_prompt_with_model_output():
     """Test a prompt with a model output."""
 
     class TestPrompt(Prompt[None, _PromptOutput]):  # pylint: disable=unused-variable
@@ -44,15 +44,15 @@ class TestPrompt(Prompt[None, _PromptOutput]):  # pylint: disable=unused-variabl
         user_prompt = "Hello"
 
     prompt = TestPrompt()
-    assert prompt.parse_response('{"song_title": "Hello", "song_lyrics": "World"}') == _PromptOutput(
+    assert await prompt.parse_response('{"song_title": "Hello", "song_lyrics": "World"}') == _PromptOutput(
         song_title="Hello", song_lyrics="World"
     )
 
     with pytest.raises(ResponseParsingError):
-        prompt.parse_response('{"song_title": "Hello"}')
+        await prompt.parse_response('{"song_title": "Hello"}')
 
 
-def test_prompt_with_float_output():
+async def test_prompt_with_float_output():
     """Test a prompt with a float output."""
 
     class TestPrompt(Prompt[None, float]):  # pylint: disable=unused-variable
@@ -61,13 +61,13 @@ class TestPrompt(Prompt[None, float]):  # pylint: disable=unused-variable
         user_prompt = "Hello"
 
     prompt = TestPrompt()
-    assert prompt.parse_response("1.0") == 1.0
+    assert await prompt.parse_response("1.0") == 1.0
 
     with pytest.raises(ResponseParsingError):
-        prompt.parse_response("a")
+        await prompt.parse_response("a")
 
 
-def test_prompt_with_bool_output():
+async def test_prompt_with_bool_output():
     """Test a prompt with a bool output."""
 
     class TestPrompt(Prompt[None, bool]):  # pylint: disable=unused-variable
@@ -76,14 +76,14 @@ class TestPrompt(Prompt[None, bool]):  # pylint: disable=unused-variable
         user_prompt = "Hello"
 
     prompt = TestPrompt()
-    assert prompt.parse_response("true") is True
-    assert prompt.parse_response("false") is False
+    assert await prompt.parse_response("true") is True
+    assert await prompt.parse_response("false") is False
 
     with pytest.raises(ResponseParsingError):
-        prompt.parse_response("a")
+        await prompt.parse_response("a")
 
 
-def test_prompt_with_int_and_custom_parser():
+async def test_prompt_with_int_and_custom_parser():
     """Test a prompt with an int output and a custom parser."""
 
     class TestPrompt(Prompt[None, int]):  # pylint: disable=unused-variable
@@ -111,10 +111,10 @@ def response_parser(response: str) -> int:
             raise ResponseParsingError("Could not parse response")
 
     prompt = TestPrompt()
-    assert prompt.parse_response("abcd k2") == 2
+    assert await prompt.parse_response("abcd k2") == 2
 
     with pytest.raises(ResponseParsingError):
-        prompt.parse_response("a")
+        await prompt.parse_response("a")
 
 
 def test_prompt_with_unknown_output_and_no_parser():
@@ -127,7 +127,7 @@ class TestPrompt(Prompt[None, list]):  # pylint: disable=unused-variable
             user_prompt = "Hello"
 
 
-def test_prompt_with_unknown_output_and_custom_parser():
+async def test_prompt_with_unknown_output_and_custom_parser():
     """Test a prompt with an output type that doesn't have a default parser but has a custom parser."""
 
     class TestPrompt(Prompt[None, list]):  # pylint: disable=unused-variable
@@ -152,5 +152,5 @@ def response_parser(response: str) -> list:
             return response.split()
 
     prompt = TestPrompt()
-    assert prompt.parse_response("Hello World") == ["Hello", "World"]
-    assert prompt.parse_response("Hello") == ["Hello"]
+    assert await prompt.parse_response("Hello World") == ["Hello", "World"]
+    assert await prompt.parse_response("Hello") == ["Hello"]

packages/ragbits-core/tests/unit/prompts/test_prompt.py

@@ -77,7 +77,7 @@ class TestPromptSystem(Prompt):  # pylint: disable=unused-variable
             system_prompt = "Hello, {{ name }}"
 
 
-def test_raises_when_unknow_user_template_variable():
+def test_raises_when_unknown_user_template_variable():
     """Test that a ValueError is raised when an unknown template variable is provided."""
     with pytest.raises(ValueError):
 
@@ -87,7 +87,7 @@ class TestPromptUser(Prompt[_PromptInput, str]):  # pylint: disable=unused-varia
             user_prompt = "Hello, {{ foo }}"
 
 
-def test_raises_when_unknow_system_template_variable():
+def test_raises_when_unknown_system_template_variable():
     """Test that a ValueError is raised when an unknown template variable is provided."""
     with pytest.raises(ValueError):
 
@@ -515,3 +515,25 @@ class TestPrompt(Prompt[_PromptInput, str]):  # pylint: disable=unused-variable
         {"role": "assistant", "content": "Why do I know all the words?"},
         {"role": "user", "content": "Theme for the song is rock."},
     ]
+
+
+async def test_response_parser():
+    class TestPrompt(Prompt):
+        user_prompt = "Hello AI"
+
+    async def async_parser(response: str) -> str:
+        return response.upper()
+
+    def sync_parser(response: str) -> str:
+        return response.lower()
+
+    test_prompt = TestPrompt()
+
+    resp = "Hello Human"
+    test_prompt.response_parser = async_parser
+    resp_async = await test_prompt.parse_response(resp)
+    assert resp_async == "HELLO HUMAN"
+
+    test_prompt.response_parser = sync_parser
+    resp_sync = await test_prompt.parse_response(resp)
+    assert resp_sync == "hello human"