feat(mistral): add mistral support (#43)

Co-authored-by: Patrick Loeber <[email protected]>
Co-authored-by: Niels Swimberghe <[email protected]>
Co-authored-by: Robert McHardy <[email protected]>
Co-authored-by: Robert McHardy <[email protected]>
Co-authored-by: Martin Schweiger <[email protected]>

Author: 6 people

assemblyai/lemur.py

@@ -1,8 +1,7 @@
 from typing import Any, Dict, List, Optional, Union
 
-from . import api
+from . import api, types
 from . import client as _client
-from . import types
 
 
 class _LemurImpl:
@@ -173,10 +172,11 @@ def question(
         Args:
             questions: One or a list of questions to ask.
             context: The context which is shared among all questions. This can be a string or a dictionary.
-            final_model: The model that is used for the final prompt after compression is performed (options: "basic" and "default").
+            final_model: The model that is used for the final prompt after compression is performed (options: "basic", "default", and "assemblyai/mistral-7b").
             max_output_size: Max output size in tokens
             timeout: The timeout in seconds to wait for the answer(s).
             temperature: Change how deterministic the response is, with 0 being the most deterministic and 1 being the least deterministic.
+            input_text: Custom formatted transcript data. Use instead of transcript_ids.
 
         Returns: One or a list of answer objects.
         """
@@ -214,10 +214,11 @@ def summarize(
         Args:
             context: An optional context on the transcript.
             answer_format: The format on how the summary shall be summarized.
-            final_model: The model that is used for the final prompt after compression is performed (options: "basic" and "default").
+            final_model: The model that is used for the final prompt after compression is performed (options: "basic", "default", and "assemblyai/mistral-7b").
             max_output_size: Max output size in tokens
             timeout: The timeout in seconds to wait for the summary.
             temperature: Change how deterministic the response is, with 0 being the most deterministic and 1 being the least deterministic.
+            input_text: Custom formatted transcript data. Use instead of transcript_ids.
 
         Returns: The summary as a string.
         """
@@ -253,10 +254,11 @@ def action_items(
         Args:
             context: An optional context on the transcript.
             answer_format: The preferred format for the result action items.
-            final_model: The model that is used for the final prompt after compression is performed (options: "basic" and "default").
+            final_model: The model that is used for the final prompt after compression is performed (options: "basic", "default", and "assemblyai/mistral-7b").
             max_output_size: Max output size in tokens
             timeout: The timeout in seconds to wait for the action items response.
             temperature: Change how deterministic the response is, with 0 being the most deterministic and 1 being the least deterministic.
+            input_text: Custom formatted transcript data. Use instead of transcript_ids.
 
         Returns: The action items as a string.
         """
@@ -287,10 +289,11 @@ def task(
 
         Args:
             prompt: The prompt to use for this task.
-            final_model: The model that is used for the final prompt after compression is performed (options: "basic" and "default").
+            final_model: The model that is used for the final prompt after compression is performed (options: "basic", "default", and "assemblyai/mistral-7b").
             max_output_size: Max output size in tokens
             timeout: The timeout in seconds to wait for the task.
             temperature: Change how deterministic the response is, with 0 being the most deterministic and 1 being the least deterministic.
+            input_text: Custom formatted transcript data. Use instead of transcript_ids.
 
         Returns: A response to a question or task submitted via custom prompt (with source transcripts or other sources taken into the context)
         """

assemblyai/transcriber.py

@@ -27,9 +27,8 @@
 from typing_extensions import Self
 from websockets.sync.client import connect as websocket_connect
 
-from . import api
+from . import api, lemur, types
 from . import client as _client
-from . import lemur, types
 
 
 class _TranscriptImpl:

assemblyai/types.py

@@ -1247,6 +1247,7 @@ class Word(BaseModel):
     start: int
     end: int
     confidence: float
+    speaker: Optional[str]
 
 
 class UtteranceWord(Word):
@@ -1382,6 +1383,10 @@ class RedactedAudioResponse(BaseModel):
 
 class Sentence(Word):
     words: List[Word]
+    start: int
+    end: int
+    confidence: int
+    speaker: Optional[str]
 
 
 class SentencesResponse(BaseModel):
@@ -1392,6 +1397,10 @@ class SentencesResponse(BaseModel):
 
 class Paragraph(Word):
     words: List[Word]
+    start: int
+    end: int
+    confidence: int
+    text: str
 
 
 class ParagraphsResponse(BaseModel):
@@ -1695,7 +1704,7 @@ def from_lemur_source(cls, source: LemurSource) -> Self:
 
 class LemurModel(str, Enum):
     """
-    LeMUR features two model modes, Basic and Default, that allow you to configure your request
+    LeMUR features three model modes, Basic, Default and Mistral 7B, that allow you to configure your request
     to suit your needs. These options tell LeMUR whether to use the more advanced Default model or
     the cheaper, faster, but simplified Basic model. The implicit setting is Default when no option
     is explicitly passed in.
@@ -1720,6 +1729,11 @@ class LemurModel(str, Enum):
     for complex/subjective tasks where answers require more nuance to be effective.
     """
 
+    mistral7b = "assemblyai/mistral-7b"
+    """
+    Mistral 7B is an open source model that works well for summarization and answering questions.
+    """
+
 
 class LemurQuestionAnswer(BaseModel):
     """
@@ -1921,7 +1935,7 @@ class RealtimeTranscript(BaseModel):
     text: str
     "The transcript for your audio"
 
-    words: List[Word]
+    words: List[RealtimeWord]
     """
     An array of objects, with the information for each word in the transcription text.
     Will include the `start`/`end` time (in milliseconds) of the word, the `confidence` score of the word,

setup.py

@@ -7,7 +7,7 @@
 
 setup(
     name="assemblyai",
-    version="0.20.0",
+    version="0.20.1",
     description="AssemblyAI Python SDK",
     author="AssemblyAI",
     author_email="[email protected]",

tests/unit/test_lemur.py

@@ -514,6 +514,40 @@ def test_lemur_task_succeeds_input_text(httpx_mock: HTTPXMock):
     assert len(httpx_mock.get_requests()) == 1
 
 
+def test_lemur_task_succeeds_mistral(httpx_mock: HTTPXMock):
+    """
+    Tests whether creating a task request succeeds with mistral.
+    """
+
+    # create a mock response of a LemurSummaryResponse
+    mock_lemur_task_response = factories.generate_dict_factory(
+        factories.LemurTaskResponse
+    )()
+
+    # mock the specific endpoints
+    httpx_mock.add_response(
+        url=f"{aai.settings.base_url}{ENDPOINT_LEMUR}/task",
+        status_code=httpx.codes.OK,
+        method="POST",
+        json=mock_lemur_task_response,
+    )
+    # test input_text input
+    lemur = aai.Lemur()
+    result = lemur.task(
+        final_model=aai.LemurModel.mistral7b,
+        prompt="Create action items of the meeting",
+        input_text="Test test",
+    )
+
+    # check the response
+    assert isinstance(result, aai.LemurTaskResponse)
+
+    assert result.response == mock_lemur_task_response["response"]
+
+    # check whether we mocked everything
+    assert len(httpx_mock.get_requests()) == 1
+
+
 def test_lemur_ask_coach_fails(httpx_mock: HTTPXMock):
     """
     Tests whether creating a task request fails.

tests/unit/test_summarization.py

@@ -26,7 +26,8 @@ def test_summarization_fails_without_required_field(
             httpx_mock,
             {},
             config=aai.TranscriptionConfig(
-                summarization=True, **{required_field: False}  # type: ignore
+                summarization=True,
+                **{required_field: False},  # type: ignore
             ),
         )