fix(python/sdk): Add missing context parameter for LeMUR task (#6586)

GitOrigin-RevId: 37d05a8a56e2a45f031e5bcbdd4735e855da3e01

Author: ploeber

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "0.33.0"
+__version__ = "0.34.0"

assemblyai/lemur.py

@@ -103,6 +103,7 @@ def action_items(
     def task(
         self,
         prompt: str,
+        context: Optional[Union[str, Dict[str, Any]]],
         final_model: Optional[types.LemurModel],
         max_output_size: Optional[int],
         timeout: Optional[float],
@@ -114,6 +115,7 @@ def task(
             request=types.LemurTaskRequest(
                 sources=self._sources,
                 prompt=prompt,
+                context=context,
                 final_model=final_model,
                 max_output_size=max_output_size,
                 temperature=temperature,
@@ -438,6 +440,7 @@ def action_items_async(
     def task(
         self,
         prompt: str,
+        context: Optional[Union[str, Dict[str, Any]]] = None,
         final_model: Optional[types.LemurModel] = None,
         max_output_size: Optional[int] = None,
         timeout: Optional[float] = None,
@@ -451,6 +454,7 @@ def task(
 
         Args:
             prompt: The prompt to use for this task.
+            context: An optional context on the transcript.
             final_model: The model that is used for the final prompt after compression is performed.
             max_output_size: Max output size in tokens
             timeout: The timeout in seconds to wait for the task.
@@ -462,6 +466,7 @@ def task(
 
         return self._impl.task(
             prompt=prompt,
+            context=context,
             final_model=final_model,
             max_output_size=max_output_size,
             timeout=timeout,
@@ -472,6 +477,7 @@ def task(
     def task_async(
         self,
         prompt: str,
+        context: Optional[Union[str, Dict[str, Any]]] = None,
         final_model: Optional[types.LemurModel] = None,
         max_output_size: Optional[int] = None,
         timeout: Optional[float] = None,
@@ -485,6 +491,7 @@ def task_async(
 
         Args:
             prompt: The prompt to use for this task.
+            context: An optional context on the transcript.
             final_model: The model that is used for the final prompt after compression is performed.
             max_output_size: Max output size in tokens
             timeout: The timeout in seconds to wait for the task.
@@ -497,6 +504,7 @@ def task_async(
         return self._executor.submit(
             self._impl.task,
             prompt=prompt,
+            context=context,
             final_model=final_model,
             max_output_size=max_output_size,
             timeout=timeout,

assemblyai/types.py

@@ -2068,6 +2068,7 @@ class LemurStringResponse(BaseLemurResponse):
 
 
 class LemurTaskRequest(BaseLemurRequest):
+    context: Optional[Union[str, Dict[str, Any]]]
     prompt: str
 
 

tests/unit/test_lemur.py

@@ -477,7 +477,9 @@ def test_lemur_task_succeeds_transcript(httpx_mock: HTTPXMock):
     lemur = aai.Lemur(
         sources=[aai.LemurSource(transcript)],
     )
-    result = lemur.task(prompt="Create action items of the meeting")
+    result = lemur.task(
+        prompt="Create action items of the meeting", context="An important meeting"
+    )
 
     # check the response
     assert isinstance(result, aai.LemurTaskResponse)
@@ -559,6 +561,7 @@ def test_lemur_task_succeeds(final_model, httpx_mock: HTTPXMock):
     result = lemur.task(
         final_model=final_model,
         prompt="Create action items of the meeting",
+        context="An important meeting",
         input_text="Test test",
     )