Add support for length parameter (#16)

* Add support for length parameter

* Add tests for length parameter

README.md

@@ -1,6 +1,6 @@
 # <img src="https://raw.githubusercontent.com/vsakkas/bard.py/master/images/logo.svg?token=GHSAT0AAAAAAB7MEK465TODCKRPHN3YQY54ZKGUN4Q" width="28px" /> Bard.py
 
-[![Latest Release](https://img.shields.io/github/v/release/vsakkas/bard.py.svg)](https://github.com/vsakkas/bard.py/releases/tag/v0.3.0)
+[![Latest Release](https://img.shields.io/github/v/release/vsakkas/bard.py.svg)](https://github.com/vsakkas/bard.py/releases/tag/v0.4.0)
 [![Python](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
 [![MIT License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/vsakkas/bard.py/blob/master/LICENSE)
 
@@ -12,7 +12,8 @@ Python Client for Bard, a Chat Based AI tool by Google.
 ## Features
 
 - Connect to Bard, Google's AI-powered personal assistant.
-- Ask questions and have a conversation in various tones.
+- Ask questions and have a conversation.
+- Improve responses by defining the conversation tone and length.
 - Use asyncio for efficient and non-blocking I/O operations.
 
 ## Requirements
@@ -155,6 +156,25 @@ The available options for the `tone` parameter are:
 > [!NOTE]
 > It is recommended to use the `tone` parameter on subsequent prompts and not in the first one. This is because this feature is typically used to change the previous response, rather than define the entire conversation tone.
 
+### Length
+
+You can set the response length when having a conversation with Bard:
+
+```python
+async with BardClient() as bard:
+    _ = await bard.ask("When was Bard released?")
+
+    response = await bard.ask("When was Bard released?", length="Short")
+    print(response)
+```
+
+The available options for the `length` parameter are:
+- `Short`
+- `Long`
+
+> [!NOTE]
+> It is recommended to use the `length` parameter on subsequent prompts and not in the first one. This is because this feature is typically used to change the previous response, rather than define the entire conversation length.
+
 
 ### Exceptions
 

bard/bard.py

@@ -5,7 +5,7 @@
 from aiohttp import ClientSession
 
 from bard.constants import BARD_STREAM_GENERATE_URL, BARD_URL, BARD_VERSION, HEADERS
-from bard.enums import ConversationTone
+from bard.enums import ConversationLength, ConversationTone
 from bard.exceptions import (
     AskException,
     CreateConversationException,
@@ -63,11 +63,29 @@ def _build_ask_parameters(self) -> dict:
             "rt": "c",
         }
 
-    def _build_ask_arguments(self, prompt: str, tone: str | None) -> dict:
+    def _build_ask_arguments(
+        self, prompt: str, tone: str | None, length: str | None
+    ) -> dict:
         conversation_arguments = None
+
+        tone_value = 0
+        length_value = 0
+
         if tone:
             tone_value = getattr(ConversationTone, tone.upper()).value
-            conversation_arguments = [0, [tone_value], None, None, None, None, []]
+        if length:
+            length_value = getattr(ConversationLength, length.upper()).value
+
+        if tone or length:
+            conversation_arguments = [
+                length_value,
+                [tone_value],
+                None,
+                None,
+                None,
+                None,
+                [],
+            ]
 
         request_data = [
             [prompt, 0, None, [], None, None, 0],
@@ -80,8 +98,8 @@ def _build_ask_arguments(self, prompt: str, tone: str | None) -> dict:
                 self.choice_id if self.choice_id else None,
                 [],
             ],
-            "",  # TODO: Find what this is
-            "",  # TODO: Find what this is
+            "",  # TODO: Find out what this is
+            "",  # TODO: Find out what this is
             conversation_arguments,
             [1],
             1,
@@ -96,9 +114,11 @@ def _build_ask_arguments(self, prompt: str, tone: str | None) -> dict:
             "at": self.snlm0e,
         }
 
-    async def _ask(self, prompt: str, tone: str | None = None) -> str | None:
+    async def _ask(
+        self, prompt: str, tone: str | None = None, length: str | None = None
+    ) -> str | None:
         parameters = self._build_ask_parameters()
-        arguments = self._build_ask_arguments(prompt, tone)
+        arguments = self._build_ask_arguments(prompt, tone, length)
 
         session = await self._get_session()
 
@@ -147,7 +167,9 @@ async def start_conversation(self) -> None:
 
             self.snlm0e = snlm0e_dict.group("value")
 
-    async def ask(self, prompt: str, tone: str | None = None) -> str:
+    async def ask(
+        self, prompt: str, tone: str | None = None, length: str | None = None
+    ) -> str:
         """
         Send a prompt to Bard and return the answer.
 
@@ -157,14 +179,18 @@ async def ask(self, prompt: str, tone: str | None = None) -> str:
             The prompt that needs to be sent to Bard.
         tone: str
             The tone that Bard will use in the next response. If no value is
-            given, it will use a default tone.
+            given, it will use a default tone which will depend on the provided
+            prompt.
+        length: str
+            The length that Bard will stick to in the next response. If no value is
+            given, the length will depend on the provided prompt.
 
         Returns
         -------
         str
             The response from Bard.
         """
-        response = await self._ask(prompt=prompt, tone=tone)
+        response = await self._ask(prompt=prompt, tone=tone, length=length)
         if not response:
             raise NoResponseException("No response was returned")
 

bard/enums.py

@@ -14,3 +14,16 @@ class ConversationTone(Enum):
     CASUAL = 2
     SIMPLE = 4
     PROFESSIONAL = 5
+
+
+class ConversationLength(Enum):
+    """
+    Bard conversation length. Supported options are:
+    - `Default`
+    - `Short`
+    - `Long`
+    """
+
+    DEFAULT = 0
+    SHORT = 1
+    LONG = 2

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bard-py"
-version = "0.3.0"
+version = "0.4.0"
 description = "Python Client for Bard."
 authors = ["vsakkas <[email protected]>"]
 license = "MIT"

tests/test_ask.py

@@ -88,6 +88,35 @@ async def test_ask_professional() -> bool:
         assert False, f"Unexpected response: {response}, match score: {score}"
 
 
+@pytest.mark.asyncio
+async def tesk_ask_short() -> bool:
+    expected_responses = [
+        "4.",
+        "You have 1 apple left today.",
+        "You have 4 apples left.",
+        "You have 4 apples today.",
+        "You still have 4 apples.",
+    ]
+
+    async with BardClient() as bard:
+        _ = await bard.ask(
+            "I have 4 apples today. I ate 3 apples yesterday. How many apples do I have today?"
+        )
+
+        response = await bard.ask(
+            "I have 4 apples today. I ate 3 apples yesterday. How many apples do I have today?",
+            length="Short",
+        )
+
+        score = 0
+        for expected_response in expected_responses:
+            score = fuzz.token_sort_ratio(response, expected_response)
+            if score >= 80:
+                return True
+
+        assert False, f"Unexpected response: {response}, match score: {score}"
+
+
 @pytest.mark.asyncio
 async def test_ask_multiple_prompts() -> None:
     async with BardClient() as bard: