feat(mistral): add mistral support to the sdk (#3157)

GitOrigin-RevId: 67e94fa0b2f00c8343d3c19988c3c8a682575cdd

Author: jhazenaai

assemblyai/lemur.py

@@ -172,7 +172,7 @@ 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.
@@ -214,7 +214,7 @@ 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.
@@ -254,7 +254,7 @@ 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.
@@ -289,7 +289,7 @@ 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.

assemblyai/types.py

@@ -1704,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.
@@ -1729,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):
     """

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.