Merge pull request #7 from speechmatics/v0.0.6

Added client function calling support

Author: dumitrugutu

CHANGELOG.md

@@ -4,6 +4,18 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [0.0.6] - 2024-11-18
+
+### Added
+
+- `tools` parameter: Introduced in the `client.run()` function to enable custom tool functionality.
+- `ToolFunctionParam` class: Added for enhanced type-checking when building client functions.
+- New message types: `ToolInvoke` and `ToolResult` messages are now supported for handling function calling.
+
+### Changed
+
+- StartConversation message: Updated to include the `tools` parameter.
+
 ## [0.0.5] - 2024-11-13
 
 ### Added

requirements-dev.txt

@@ -1,6 +1,7 @@
-pytest==7.1.1
-pytest-mock==3.7.0
+pytest
+pytest-mock
 black==22.3.0
 ruff==0.0.280
 pre-commit==2.21.0
-pytest-cov==3.0.0
+pytest-cov
+pytest-asyncio

speechmatics_flow/client.py

@@ -5,10 +5,12 @@
 
 import asyncio
 import copy
+import inspect
 import json
 import logging
 import os
-from typing import List
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Optional
 
 import httpx
 import pyaudio
@@ -27,6 +29,7 @@
     Interaction,
     ConnectionSettings,
 )
+from speechmatics_flow.tool_function_param import ToolFunctionParam
 from speechmatics_flow.utils import read_in_chunks, json_utf8
 
 LOGGER = logging.getLogger(__name__)
@@ -60,6 +63,7 @@ def __init__(
         self.websocket = None
         self.conversation_config = None
         self.audio_settings = None
+        self.tools = None
 
         self.event_handlers = {x: [] for x in ServerMessageType}
         self.middlewares = {x: [] for x in ClientMessageType}
@@ -70,6 +74,7 @@ def __init__(
         self.conversation_ended_wait_timeout = 5
         self._session_needs_closing = False
         self._audio_buffer = None
+        self._executor = ThreadPoolExecutor()
 
         # The following asyncio fields are fully instantiated in
         # _init_synchronization_primitives
@@ -124,6 +129,8 @@ def _start_conversation(self):
             "audio_format": self.audio_settings.asdict(),
             "conversation_config": self.conversation_config.asdict(),
         }
+        if self.tools is not None:
+            msg["tools"] = self.tools
         self.session_running = True
         self._call_middleware(ClientMessageType.StartConversation, msg, False)
         LOGGER.debug(msg)
@@ -166,7 +173,7 @@ async def _wait_for_conversation_ended(self):
             self._conversation_ended.wait(), self.conversation_ended_wait_timeout
         )
 
-    async def _consumer(self, message, from_cli: False):
+    async def _consumer(self, message, from_cli=False):
         """
         Consumes messages and acts on them.
 
@@ -204,10 +211,18 @@ async def _consumer(self, message, from_cli: False):
 
         for handler in self.event_handlers[message_type]:
             try:
-                handler(copy.deepcopy(message))
+                if inspect.iscoroutinefunction(handler):
+                    await handler(copy.deepcopy(message))
+                else:
+                    loop = asyncio.get_event_loop()
+                    await loop.run_in_executor(
+                        self._executor, handler, copy.deepcopy(message)
+                    )
             except ForceEndSession:
                 LOGGER.warning("Session was ended forcefully by an event handler")
                 raise
+            except Exception as e:
+                LOGGER.error(f"Unhandled exception in {handler=}: {e=}")
 
         if message_type == ServerMessageType.ConversationStarted:
             self._flag_conversation_started()
@@ -262,7 +277,7 @@ async def _read_from_microphone(self):
             stream.close()
             _pyaudio.terminate()
 
-    async def _consumer_handler(self, from_cli: False):
+    async def _consumer_handler(self, from_cli=False):
         """
         Controls the consumer loop for handling messages from the server.
 
@@ -492,6 +507,7 @@ async def run(
         audio_settings: AudioSettings = AudioSettings(),
         conversation_config: ConversationConfig = None,
         from_cli: bool = False,
+        tools: Optional[List[ToolFunctionParam]] = None,
     ):
         """
         Begin a new recognition session.
@@ -508,13 +524,17 @@ async def run(
         :param conversation_config: Configuration for the conversation.
         :type conversation_config: models.ConversationConfig
 
+        :param tools: Optional list of tool functions.
+        :type tools: List[ToolFunctionParam]
+
         :raises Exception: Can raise any exception returned by the
             consumer/producer tasks.
         """
         self.client_seq_no = 0
         self.server_seq_no = 0
         self.conversation_config = conversation_config
         self.audio_settings = audio_settings
+        self.tools = tools
 
         await self._init_synchronization_primitives()
 

speechmatics_flow/models.py

@@ -94,6 +94,10 @@ class ClientMessageType(str, Enum):
     AudioEnded = "AudioEnded"
     """Indicates audio input has finished."""
 
+    ToolResult = "ToolResult"
+    """Client response to :py:attr:`ServerMessageType.ToolInvoke`, containing
+    the result of the function call."""
+
 
 class ServerMessageType(str, Enum):
     # pylint: disable=invalid-name
@@ -133,7 +137,7 @@ class ServerMessageType(str, Enum):
 
     AddAudio = "AddAudio"
     """Implicit name for all outbound binary messages. The client confirms
-    receipt by sending an :py:attr:`ServerMessageType.AudioReceived` message."""
+    receipt by sending an :py:attr:`ClientMessageType.AudioReceived` message."""
 
     audio = "audio"
     """Message contains binary data"""
@@ -148,6 +152,11 @@ class ServerMessageType(str, Enum):
     ConversationEnded = "ConversationEnded"
     """Message indicates the session ended."""
 
+    ToolInvoke = "ToolInvoke"
+    """Indicates invocation of a function call. The client responds by sending
+    an :py:attr:`ClientMessageType.ToolResult` message.
+    """
+
     Info = "Info"
     """Indicates a generic info message."""
 

speechmatics_flow/tool_function_param.py

@@ -0,0 +1,63 @@
+"""Define a set of typed dictionaries to represent a structured message
+for a client side function, using TypedDict for type enforcement.
+
+Example:
+    tool_function = ToolFunctionParam(
+        type="function",
+        function={
+            "name": "add_nums",
+            "description": "Adds two numbers",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "a": {"type": "int", "description": "First number"},
+                    "b": {"type": "int", "description": "Second number"}
+                },
+                "required": ["a", "b"]
+            }
+        }
+    )
+
+    # Convert to dictionary
+    dict(tool_function)
+"""
+
+import sys
+
+from typing import Literal, Optional, Dict, List, TypedDict
+
+if sys.version_info < (3, 11):
+    from typing_extensions import Required
+else:
+    from typing import Required
+
+
+class Property(TypedDict):
+    type: Required[str]
+    description: Required[str]
+
+
+class FunctionParam(TypedDict, total=False):
+    type: Required[str]
+    properties: Required[Dict[str, Property]]
+    required: Optional[List[str]]
+
+
+class FunctionDefinition(TypedDict, total=False):
+    name: Required[str]
+    """The name of the function to be called."""
+
+    description: Optional[str]
+    """The description of what the function does, used by the model to choose
+    when and how to call the function.
+    """
+
+    parameters: Optional[FunctionParam]
+    """The parameters of the function to be called."""
+
+
+class ToolFunctionParam(TypedDict):
+    type: Required[Literal["function"]]
+    """Currently, only 'function' is supported."""
+
+    function: Required[FunctionDefinition]

tests/test_client.py

@@ -0,0 +1,112 @@
+import json
+from typing import List, Dict, Optional
+
+import pytest
+from pytest import param
+
+from speechmatics_flow.client import WebsocketClient
+from speechmatics_flow.models import (
+    ServerMessageType,
+    ConnectionSettings,
+    ClientMessageType,
+    AudioSettings,
+    ConversationConfig,
+)
+from speechmatics_flow.tool_function_param import ToolFunctionParam
+
+TOOL_FUNCTION = dict(
+    ToolFunctionParam(
+        type="function",
+        function={
+            "name": "test_function",
+            "description": "test function to be called.",
+        },
+    )
+)
+
+
+@pytest.fixture
+def ws_client():
+    return WebsocketClient(
+        connection_settings=ConnectionSettings(url="ws://test"),
+    )
+
+
+@pytest.mark.parametrize(
+    "audio_format, conversation_config, tools, expected_start_message",
+    [
+        param(
+            AudioSettings(),
+            ConversationConfig(),
+            None,
+            {
+                "message": ClientMessageType.StartConversation.value,
+                "audio_format": AudioSettings().asdict(),
+                "conversation_config": ConversationConfig().asdict(),
+            },
+            id="with default values",
+        ),
+        param(
+            AudioSettings(),
+            ConversationConfig(),
+            [TOOL_FUNCTION],
+            {
+                "message": ClientMessageType.StartConversation.value,
+                "audio_format": AudioSettings().asdict(),
+                "conversation_config": ConversationConfig().asdict(),
+                "tools": [TOOL_FUNCTION],
+            },
+            id="with default values and tools",
+        ),
+    ],
+)
+def test_start_conversation(
+    ws_client: WebsocketClient,
+    audio_format: AudioSettings,
+    conversation_config: ConversationConfig,
+    tools: Optional[List[ToolFunctionParam]],
+    expected_start_message: Dict,
+):
+    handler_called = False
+
+    def handler(*_):
+        nonlocal handler_called
+        handler_called = True
+
+    ws_client.middlewares = {ClientMessageType.StartConversation: [handler]}
+    ws_client.audio_settings = audio_format
+    ws_client.conversation_config = conversation_config
+    ws_client.tools = tools
+    start_conversation_msg = ws_client._start_conversation()
+    assert start_conversation_msg == json.dumps(
+        expected_start_message
+    ), f"expected={start_conversation_msg}, got={expected_start_message}"
+    assert handler_called, "handler was not called"
+
+
+@pytest.mark.asyncio
+async def test_consumer_supports_sync_and_async_handlers(ws_client):
+    await ws_client._init_synchronization_primitives()
+
+    async_handler_called = False
+    sync_handler_called = False
+
+    async def async_handler(_):
+        nonlocal async_handler_called
+        async_handler_called = True
+
+    def sync_handler(_):
+        nonlocal sync_handler_called
+        sync_handler_called = True
+
+    # Add event handlers for a message type
+    ws_client.event_handlers = {
+        ServerMessageType.ConversationStarted: [async_handler, sync_handler],
+    }
+
+    message = json.dumps({"message": ServerMessageType.ConversationStarted})
+    await ws_client._consumer(message, from_cli=False)
+
+    # Check if both handlers were called
+    assert async_handler_called, "async handler was not called"
+    assert sync_handler_called, "sync handler was not called"

tests/test_tool_function_param.py

@@ -0,0 +1,47 @@
+from pytest import mark, param
+
+from speechmatics_flow.tool_function_param import ToolFunctionParam
+
+
+@mark.parametrize(
+    "params",
+    [
+        param(
+            {
+                "type": "function",
+                "function": {
+                    "name": "test_function",
+                    "description": "test_description",
+                },
+            },
+            id="function without optional params",
+        ),
+        param(
+            {
+                "type": "function",
+                "function": {
+                    "name": "test_function",
+                    "description": "test_description",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "a": {
+                                "type": "int",
+                                "description": "First number",
+                            },
+                            "b": {
+                                "type": "int",
+                                "description": "Second number",
+                            },
+                        },
+                        "required": ["a", "b"],
+                    },
+                },
+            },
+            id="function with optional params",
+        ),
+    ],
+)
+def test_websocket_function(params):
+    websocket_function = ToolFunctionParam(**params)
+    assert dict(websocket_function) == params