feat: Maestro - add data sources support (#282)

fix: add data sources support

* fix: typedDict import

* fix: add missing param

* fix: add test

* fix: remove skip

* fix: fix test

Author: MichalAI21

ai21/clients/common/maestro/run.py

@@ -12,6 +12,7 @@
     DEFAULT_RUN_POLL_TIMEOUT,
     Requirement,
     Budget,
+    OutputOptions,
 )
 from ai21.types import NOT_GIVEN, NotGiven
 from ai21.utils.typing import remove_not_given
@@ -30,6 +31,7 @@ def _create_body(
         context: Dict[str, Any] | NotGiven,
         requirements: List[Requirement] | NotGiven,
         budget: Budget | NotGiven,
+        include: List[OutputOptions] | NotGiven,
         **kwargs,
     ) -> dict:
         return remove_not_given(
@@ -41,6 +43,7 @@ def _create_body(
                 "context": context,
                 "requirements": requirements,
                 "budget": budget,
+                "include": include,
                 **kwargs,
             }
         )
@@ -56,6 +59,7 @@ def create(
         context: Dict[str, Any] | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
+        include: List[OutputOptions] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> RunResponse:
         pass
@@ -79,6 +83,7 @@ def create_and_poll(
         context: Dict[str, Any] | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
+        include: List[OutputOptions] | NotGiven = NOT_GIVEN,
         poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
         poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
         **kwargs,

ai21/clients/studio/resources/maestro/run.py

@@ -16,6 +16,7 @@
     DEFAULT_RUN_POLL_TIMEOUT,
     Requirement,
     Budget,
+    OutputOptions,
 )
 from ai21.types import NotGiven, NOT_GIVEN
 
@@ -31,6 +32,7 @@ def create(
         context: Dict[str, Any] | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
+        include: List[OutputOptions] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> RunResponse:
         body = self._create_body(
@@ -41,6 +43,7 @@ def create(
             context=context,
             requirements=requirements,
             budget=budget,
+            include=include,
             **kwargs,
         )
 
@@ -76,6 +79,7 @@ def create_and_poll(
         context: Dict[str, Any] | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
+        include: List[OutputOptions] | NotGiven = NOT_GIVEN,
         poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
         poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
         **kwargs,
@@ -88,6 +92,7 @@ def create_and_poll(
             context=context,
             requirements=requirements,
             budget=budget,
+            include=include,
             **kwargs,
         )
 
@@ -105,6 +110,7 @@ async def create(
         context: Dict[str, Any] | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
+        include: List[OutputOptions] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> RunResponse:
         body = self._create_body(
@@ -115,6 +121,7 @@ async def create(
             context=context,
             requirements=requirements,
             budget=budget,
+            include=include,
             **kwargs,
         )
 
@@ -150,6 +157,7 @@ async def create_and_poll(
         context: Dict[str, Any] | NotGiven = NOT_GIVEN,
         requirements: List[Requirement] | NotGiven = NOT_GIVEN,
         budget: Budget | NotGiven = NOT_GIVEN,
+        include: List[OutputOptions] | NotGiven = NOT_GIVEN,
         poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
         poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
         **kwargs,
@@ -162,6 +170,7 @@ async def create_and_poll(
             context=context,
             requirements=requirements,
             budget=budget,
+            include=include,
             **kwargs,
         )
 

ai21/models/__init__.py

@@ -8,8 +8,16 @@
     ConversationalRagSource,
 )
 from ai21.models.responses.file_response import FileResponse
-from ai21.models.maestro.run import Requirement, Budget, Tool, ToolResources
-
+from ai21.models.maestro.run import (
+    Requirement,
+    Budget,
+    Tool,
+    ToolResources,
+    DataSources,
+    FileSearchResult,
+    WebSearchResult,
+    OutputOptions,
+)
 
 __all__ = [
     "ChatMessage",
@@ -26,4 +34,8 @@
     "Budget",
     "Tool",
     "ToolResources",
+    "DataSources",
+    "FileSearchResult",
+    "WebSearchResult",
+    "OutputOptions",
 ]

ai21/models/maestro/run.py

@@ -1,5 +1,5 @@
-from typing import TypedDict, Literal, List, Optional, Any, Set, Dict, Type, Union
-
+from typing import Literal, List, Optional, Any, Set, Dict, Type, Union
+from typing_extensions import TypedDict
 from pydantic import BaseModel
 
 from ai21.models.ai21_base_model import AI21BaseModel
@@ -8,6 +8,7 @@
 Role = Literal["user", "assistant"]
 RunStatus = Literal["completed", "failed", "in_progress", "requires_action"]
 ToolType = Literal["file_search", "web_search"]
+OutputOptions = Literal["data_sources"]
 PrimitiveTypes = Union[Type[str], Type[int], Type[float], Type[bool]]
 PrimitiveLists = Type[List[PrimitiveTypes]]
 OutputType = Union[Type[BaseModel], PrimitiveTypes, Dict[str, Any]]
@@ -40,12 +41,32 @@ class ToolResources(TypedDict, total=False):
     web_search: Optional[WebSearchToolResource]
 
 
-class Requirement(TypedDict):
+class Requirement(TypedDict, total=False):
     name: str
     description: str
 
 
+class FileSearchResult(TypedDict, total=False):
+    text: Optional[str]
+    file_id: str
+    file_name: str
+    score: float
+    order: int
+
+
+class WebSearchResult(TypedDict, total=False):
+    text: str
+    url: str
+    score: float
+
+
+class DataSources(TypedDict, total=False):
+    file_search: Optional[List[FileSearchResult]]
+    web_search: Optional[List[WebSearchResult]]
+
+
 class RunResponse(AI21BaseModel):
     id: str
     status: RunStatus
     result: Any
+    data_sources: Optional[DataSources] = None

tests/integration_tests/clients/studio/conftest.py

@@ -20,7 +20,7 @@ def _wait_for_file_to_process(client: AI21Client, file_id: str, timeout: float =
             return
 
         elapsed_time = time.time() - start_time
-        time.sleep(0.5)
+        time.sleep(2)
 
     raise TimeoutError(f"Timeout: {timeout} seconds passed. File processing not completed")
 

tests/integration_tests/clients/studio/test_maestro.py

@@ -0,0 +1,16 @@
+import pytest
+
+from ai21 import AsyncAI21Client
+
+
+@pytest.mark.asyncio
+async def test_maestro__when_upload__should_return_data_sources():  # file_in_library: str):
+    client = AsyncAI21Client()
+    result = await client.beta.maestro.runs.create_and_poll(
+        input="When did Einstein receive a Nobel Prize?", tools=[{"type": "file_search"}], include=["data_sources"]
+    )
+    assert result.status == "completed", "Expected 'completed' status"
+    assert result.result, "Expected a non-empty answer"
+    assert result.data_sources, "Expected data sources"
+    assert len(result.data_sources["file_search"]) > 0, "Expected at least one file search data source"
+    assert result.data_sources.get("web_search") is None, "Expected no web search data sources"