diff --git a/assemblyai/lemur.py b/assemblyai/lemur.py index c81ab34..d028f67 100644 --- a/assemblyai/lemur.py +++ b/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) """ diff --git a/assemblyai/transcriber.py b/assemblyai/transcriber.py index 72f83b3..82ca638 100644 --- a/assemblyai/transcriber.py +++ b/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: diff --git a/assemblyai/types.py b/assemblyai/types.py index f34243c..9f5bb37 100644 --- a/assemblyai/types.py +++ b/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, diff --git a/setup.py b/setup.py index 84e81c1..e7d5069 100644 --- a/setup.py +++ b/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="engineering.sdk@assemblyai.com", diff --git a/tests/unit/test_lemur.py b/tests/unit/test_lemur.py index d273112..6a67b62 100644 --- a/tests/unit/test_lemur.py +++ b/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. diff --git a/tests/unit/test_summarization.py b/tests/unit/test_summarization.py index 020487f..33966a1 100644 --- a/tests/unit/test_summarization.py +++ b/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 ), )