Optimisations and fixes

pyrobbot/tokens.py

@@ -1,5 +1,6 @@
 """Management of token usage and costs for OpenAI API."""
 
+import contextlib
 import datetime
 import sqlite3
 from pathlib import Path
@@ -166,10 +167,9 @@ def get_n_tokens_from_msgs(messages: list[dict], model: str):
     """Returns the number of tokens used by a list of messages."""
     # Adapted from
     # <https://platform.openai.com/docs/guides/text-generation/managing-tokens>
-    try:
+    encoding = tiktoken.get_encoding("cl100k_base")
+    with contextlib.suppress(KeyError):
         encoding = tiktoken.encoding_for_model(model)
-    except KeyError:
-        encoding = tiktoken.get_encoding("cl100k_base")
 
     # OpenAI's original function was implemented for gpt-3.5-turbo-0613, but we'll use
     # it for all models for now. We are only interested in estimates, after all.

pyrobbot/voice_chat.py

@@ -206,7 +206,7 @@ def handle_update_audio_history(self, current_answer_audios_queue: queue.Queue):
                     merged_audio += new_audio
                     current_answer_audios_queue.task_done()
                     continue
-                if merged_audio.duration_seconds < min_audio_duration_seconds:
+                if merged_audio.duration_seconds <= min_audio_duration_seconds:
                     current_answer_audios_queue.task_done()
                     continue
 

tests/conftest.py

@@ -202,13 +202,9 @@ def _voice_chat_mockers(mocker, mock_wav_bytes_string):
     mock_google_tts_obj = type("mock_gTTS", (), {})
     mock_openai_tts_response = type("mock_openai_tts_response", (), {})
 
-    def _mock_write_to_fp(wav_buffer: io.BytesIO):
-        wav_buffer.write(mock_wav_bytes_string)
-
     def _mock_iter_bytes(*args, **kwargs):  # noqa: ARG001
         return [mock_wav_bytes_string]
 
-    mock_google_tts_obj.write_to_fp = _mock_write_to_fp
     mock_openai_tts_response.iter_bytes = _mock_iter_bytes
 
     mocker.patch(
@@ -234,3 +230,4 @@ def _mock_iter_bytes(*args, **kwargs):  # noqa: ARG001
     mocker.patch("webrtcvad.Vad.is_speech", return_value=False)
     mocker.patch("pygame.mixer.init")
     mocker.patch("chime.play_wav")
+    mocker.patch("chime.play_wav")

tests/smoke/test_app.py

@@ -1,18 +1,43 @@
+import contextlib
+
+import streamlit
+import streamlit_webrtc.component
+
 from pyrobbot.app import app
 
 
 def test_app(mocker, default_voice_chat_configs):
-    class MockAttrDict(dict):
+    class MockAttrDict(streamlit.runtime.state.session_state_proxy.SessionStateProxy):
         def __getattr__(self, attr):
             return self.get(attr, mocker.MagicMock())
 
-        def __setattr__(self, attr, value):
-            self[attr] = value
+        def __getitem__(self, key):
+            with contextlib.suppress(KeyError):
+                return super().__getitem__(key)
+            return mocker.MagicMock()
 
-    mocker.patch("streamlit.session_state", MockAttrDict())
+    mocker.patch.object(streamlit, "session_state", new=MockAttrDict())
+    mocker.patch.object(
+        streamlit.runtime.state.session_state_proxy,
+        "SessionStateProxy",
+        new=MockAttrDict,
+    )
     mocker.patch("streamlit.chat_input", return_value="foobar")
     mocker.patch(
         "pyrobbot.chat_configs.VoiceChatConfigs.from_file",
         return_value=default_voice_chat_configs,
     )
+    mocker.patch.object(
+        streamlit_webrtc.component,
+        "webrtc_streamer",
+        mocker.MagicMock(return_value=mocker.MagicMock()),
+    )
+
+    mocker.patch("streamlit.number_input", return_value=0)
+
+    mocker.patch(
+        "pyrobbot.chat_configs.VoiceChatConfigs.model_validate",
+        return_value=default_voice_chat_configs,
+    )
+
     app.run_app()