weco 0.1.5__tar.gz → 0.1.6__tar.gz

{weco-0.1.5 → weco-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: weco
-Version: 0.1.5
+Version: 0.1.6
 Summary: A client facing API for interacting with the WeCo AI function builder service.
 Author-email: WeCo AI Team <dhruv@weco.ai>
 License: MIT

{weco-0.1.5 → weco-0.1.6}/examples/cookbook.ipynb

@@ -326,6 +326,13 @@
     "for key, value in query_response.items(): print(f\"{key}: {value}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## A/B Testing with Function Versions"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

{weco-0.1.5 → weco-0.1.6}/pyproject.toml

@@ -10,7 +10,7 @@ authors = [
 ]
 description = "A client facing API for interacting with the WeCo AI function builder service."
 readme = "README.md"
-version = "0.1.5"
+version = "0.1.6"
 license = {text = "MIT"}
 requires-python = ">=3.8"
 dependencies = ["asyncio", "httpx[http2]", "pillow"]

{weco-0.1.5 → weco-0.1.6}/tests/test_asynchronous.py

@@ -8,8 +8,9 @@ from weco import abuild, aquery
 @pytest.mark.asyncio
 async def test_abuild(text_evaluator, image_evaluator, text_and_image_evaluator):
     for evaluator in [text_evaluator, image_evaluator, text_and_image_evaluator]:
-        fn_name, fn_desc = await evaluator
+        fn_name, version_number, fn_desc = await evaluator
         assert isinstance(fn_name, str)
+        assert isinstance(version_number, int)
         assert isinstance(fn_desc, str)
 
 
@@ -23,16 +24,17 @@ async def assert_query_response(query_response):
 
 @pytest.fixture
 async def text_evaluator():
-    fn_name, fn_desc = await abuild(
-        task_description="Evaluate the sentiment of the given text. Provide a json object with 'sentiment' and 'explanation' keys."
+    fn_name, version_number, fn_desc = await abuild(
+        task_description="Evaluate the sentiment of the given text. Provide a json object with 'sentiment' and 'explanation' keys.",
+        multimodal=False
     )
-    return fn_name, fn_desc
+    return fn_name, version_number, fn_desc
 
 
 @pytest.mark.asyncio
 async def test_text_aquery(text_evaluator):
-    fn_name, _ = await text_evaluator
-    query_response = await aquery(fn_name=fn_name, text_input="I love this product!")
+    fn_name, version_number, _ = await text_evaluator
+    query_response = await aquery(fn_name=fn_name, version_number=version_number, text_input="I love this product!")
 
     await assert_query_response(query_response)
     assert set(query_response["output"].keys()) == {"sentiment", "explanation"}
@@ -40,17 +42,19 @@ async def test_text_aquery(text_evaluator):
 
 @pytest.fixture
 async def image_evaluator():
-    fn_name, fn_desc = await abuild(
-        task_description="Describe the contents of the given images. Provide a json object with 'description' and 'objects' keys."
+    fn_name, version_number, fn_desc = await abuild(
+        task_description="Describe the contents of the given images. Provide a json object with 'description' and 'objects' keys.",
+        multimodal=True
     )
-    return fn_name, fn_desc
+    return fn_name, version_number, fn_desc
 
 
 @pytest.mark.asyncio
 async def test_image_aquery(image_evaluator):
-    fn_name, _ = await image_evaluator
+    fn_name, version_number, _ = await image_evaluator
     query_response = await aquery(
         fn_name=fn_name,
+        version_number=version_number,
         images_input=[
             "https://www.integratedtreatmentservices.co.uk/wp-content/uploads/2013/12/Objects-of-Reference.jpg",
             "https://t4.ftcdn.net/jpg/05/70/90/23/360_F_570902339_kNj1reH40GFXakTy98EmfiZHci2xvUCS.jpg",
@@ -63,17 +67,19 @@ async def test_image_aquery(image_evaluator):
 
 @pytest.fixture
 async def text_and_image_evaluator():
-    fn_name, fn_desc = await abuild(
-        task_description="Evaluate, solve and arrive at a numerical answer for the image provided. Provide a json object with 'answer' and 'explanation' keys."
+    fn_name, version_number, fn_desc = await abuild(
+        task_description="Evaluate, solve and arrive at a numerical answer for the image provided. Provide a json object with 'answer' and 'explanation' keys.",
+        multimodal=True
     )
-    return fn_name, fn_desc
+    return fn_name, version_number, fn_desc
 
 
 @pytest.mark.asyncio
 async def test_text_and_image_aquery(text_and_image_evaluator):
-    fn_name, _ = await text_and_image_evaluator
+    fn_name, version_number, _ = await text_and_image_evaluator
     query_response = await aquery(
         fn_name=fn_name,
+        version_number=version_number,
         text_input="Find x and y.",
         images_input=[
             "https://i.ytimg.com/vi/cblHUeq3bkE/hq720.jpg?sqp=-oaymwEhCK4FEIIDSFryq4qpAxMIARUAAAAAGAElAADIQj0AgKJD&rs=AOn4CLAKn3piY91QRCBzRgnzAPf7MPrjDQ"

weco-0.1.6/tests/test_batching.py

@@ -0,0 +1,75 @@
+import pytest
+
+from weco import batch_query, build
+
+
+# Internally, these functions use the WecoAI client
+# therefore, we can test both the client and functional forms here
+@pytest.fixture
+def ml_task_evaluator():
+    fn_name, version_number, _ = build(
+        task_description="I want to evaluate the feasibility of a machine learning task. Give me a json object with three keys - 'feasibility', 'justification', and 'suggestions'.",
+        multimodal=False
+    )
+    return fn_name, version_number
+
+
+@pytest.fixture
+def ml_task_inputs():
+    return [
+        {"text_input": "I want to train a model to predict house prices using the Boston Housing dataset hosted on Kaggle."},
+        {"text_input": "I want to train a model to classify digits using the MNIST dataset hosted on Kaggle using a Google Colab notebook."},
+    ]
+
+
+@pytest.fixture
+def image_evaluator():
+    fn_name, version_number, _ = build(
+        task_description="Describe the contents of the given images. Provide a json object with 'description' and 'objects' keys.",
+        multimodal=True
+    )
+    return fn_name, version_number
+
+
+@pytest.fixture
+def image_inputs():
+    return [
+        {"images_input": ["https://www.integratedtreatmentservices.co.uk/wp-content/uploads/2013/12/Objects-of-Reference.jpg"]},
+        {"images_input": ["https://t4.ftcdn.net/jpg/05/70/90/23/360_F_570902339_kNj1reH40GFXakTy98EmfiZHci2xvUCS.jpg"]},
+    ]
+
+
+def test_batch_query_text(ml_task_evaluator, ml_task_inputs):
+    fn_name, version_number = ml_task_evaluator
+    batch_inputs = ml_task_inputs
+
+    query_responses = batch_query(fn_name=fn_name, version_number=version_number, batch_inputs=batch_inputs)
+
+    assert len(query_responses) == len(batch_inputs)
+
+    for query_response in query_responses:
+        assert isinstance(query_response["output"], dict)
+        assert isinstance(query_response["in_tokens"], int)
+        assert isinstance(query_response["out_tokens"], int)
+        assert isinstance(query_response["latency_ms"], float)
+
+        output = query_response["output"]
+        assert set(output.keys()) == {"feasibility", "justification", "suggestions"}
+
+
+def test_batch_query_image(image_evaluator, image_inputs):
+    fn_name, version_number = image_evaluator
+    batch_inputs = image_inputs
+
+    query_responses = batch_query(fn_name=fn_name, version_number=version_number, batch_inputs=batch_inputs)
+
+    assert len(query_responses) == len(batch_inputs)
+
+    for query_response in query_responses:
+        assert isinstance(query_response["output"], dict)
+        assert isinstance(query_response["in_tokens"], int)
+        assert isinstance(query_response["out_tokens"], int)
+        assert isinstance(query_response["latency_ms"], float)
+
+        output = query_response["output"]
+        assert set(output.keys()) == {"description", "objects"}

{weco-0.1.5 → weco-0.1.6}/tests/test_synchronous.py

@@ -7,8 +7,9 @@ from weco import build, query
 # therefore, we can test both the client and functional forms here
 def test_build(text_evaluator, image_evaluator, text_and_image_evaluator):
     for evaluator in [text_evaluator, image_evaluator, text_and_image_evaluator]:
-        fn_name, fn_desc = evaluator
+        fn_name, version_number, fn_desc = evaluator
         assert isinstance(fn_name, str)
+        assert isinstance(version_number, int)
         assert isinstance(fn_desc, str)
 
 
@@ -22,15 +23,16 @@ def assert_query_response(query_response):
 
 @pytest.fixture
 def text_evaluator():
-    fn_name, fn_desc = build(
-        task_description="Evaluate the sentiment of the given text. Provide a json object with 'sentiment' and 'explanation' keys."
+    fn_name, version_number, fn_desc = build(
+        task_description="Evaluate the sentiment of the given text. Provide a json object with 'sentiment' and 'explanation' keys.",
+        multimodal=False
     )
-    return fn_name, fn_desc
+    return fn_name, version_number, fn_desc
 
 
 def test_text_query(text_evaluator):
-    fn_name, _ = text_evaluator
-    query_response = query(fn_name=fn_name, text_input="I love this product!")
+    fn_name, version_number, _ = text_evaluator
+    query_response = query(fn_name=fn_name, version_number=version_number, text_input="I love this product!")
 
     assert_query_response(query_response)
     assert set(query_response["output"].keys()) == {"sentiment", "explanation"}
@@ -38,16 +40,18 @@ def test_text_query(text_evaluator):
 
 @pytest.fixture
 def image_evaluator():
-    fn_name, fn_desc = build(
-        task_description="Describe the contents of the given images. Provide a json object with 'description' and 'objects' keys."
+    fn_name, version_number, fn_desc = build(
+        task_description="Describe the contents of the given images. Provide a json object with 'description' and 'objects' keys.",
+        multimodal=True
     )
-    return fn_name, fn_desc
+    return fn_name, version_number, fn_desc
 
 
 def test_image_query(image_evaluator):
-    fn_name, _ = image_evaluator
+    fn_name, version_number, _ = image_evaluator
     query_response = query(
         fn_name=fn_name,
+        version_number=version_number,
         images_input=[
             "https://www.integratedtreatmentservices.co.uk/wp-content/uploads/2013/12/Objects-of-Reference.jpg",
             "https://t4.ftcdn.net/jpg/05/70/90/23/360_F_570902339_kNj1reH40GFXakTy98EmfiZHci2xvUCS.jpg",
@@ -60,16 +64,18 @@ def test_image_query(image_evaluator):
 
 @pytest.fixture
 def text_and_image_evaluator():
-    fn_name, fn_desc = build(
-        task_description="Evaluate, solve and arrive at a numerical answer for the image provided. Perform any additional things if instructed. Provide a json object with 'answer' and 'explanation' keys."
+    fn_name, version_number, fn_desc = build(
+        task_description="Evaluate, solve and arrive at a numerical answer for the image provided. Perform any additional things if instructed. Provide a json object with 'answer' and 'explanation' keys.",
+        multimodal=True
     )
-    return fn_name, fn_desc
+    return fn_name, version_number, fn_desc
 
 
 def test_text_and_image_query(text_and_image_evaluator):
-    fn_name, _ = text_and_image_evaluator
+    fn_name, version_number, _ = text_and_image_evaluator
     query_response = query(
         fn_name=fn_name,
+        version_number=version_number,
         text_input="Find x and y.",
         images_input=[
             "https://i.ytimg.com/vi/cblHUeq3bkE/hq720.jpg?sqp=-oaymwEhCK4FEIIDSFryq4qpAxMIARUAAAAAGAElAADIQj0AgKJD&rs=AOn4CLAKn3piY91QRCBzRgnzAPf7MPrjDQ"

{weco-0.1.5 → weco-0.1.6}/weco/client.py

@@ -65,11 +65,6 @@ class WecoAI:
         self.client = httpx.Client(http2=http2, timeout=timeout)
         self.async_client = httpx.AsyncClient(http2=http2, timeout=timeout)
 
-    def close(self):
-        """Close both synchronous and asynchronous clients."""
-        self.client.close()
-        asyncio.run(self.async_client.aclose())
-
     def _headers(self) -> Dict[str, str]:
         """Constructs the headers for the API requests."""
         return {
@@ -158,20 +153,24 @@ class WecoAI:
             "latency_ms": response["latency_ms"],
         }
 
-    def _build(self, task_description: str, is_async: bool) -> Union[Tuple[str, str], Coroutine[Any, Any, Tuple[str, str]]]:
+    def _build(self, task_description: str, multimodal: bool, is_async: bool) -> Union[Tuple[str, int, str], Coroutine[Any, Any, Tuple[str, int, str]]]:
         """Internal method to handle both synchronous and asynchronous build requests.
 
         Parameters
         ----------
         task_description : str
             A description of the task for which the function is being built.
+        
+        multimodal : bool
+            Whether the function is multimodal or not.
+        
         is_async : bool
             Whether to perform an asynchronous request.
 
         Returns
         -------
-        Union[tuple[str, str], Coroutine[Any, Any, tuple[str, str]]]
-            A tuple containing the name and description of the function, or a coroutine that returns such a tuple.
+        Union[tuple[str, int, str], Coroutine[Any, Any, tuple[str, int, str]]]
+            A tuple containing the name, version number and description of the function, or a coroutine that returns such a tuple.
 
         Raises
         ------
@@ -185,51 +184,57 @@ class WecoAI:
             raise ValueError(f"Task description must be less than {MAX_TEXT_LENGTH} characters.")
 
         endpoint = "build"
-        data = {"request": task_description}
+        data = {"request": task_description, "multimodal": multimodal}
         request = self._make_request(endpoint=endpoint, data=data, is_async=is_async)
 
         if is_async:
-
+            # return 0 for the version number
             async def _async_build():
                 response = await request
-                return response["name"], response["description"]
+                return response["function_name"], 0, response["description"]
 
             return _async_build()
         else:
             response = request  # the request has already been made and the response is available
-            return response["name"], response["description"]
+            return response["function_name"], 0, response["description"]
 
-    async def abuild(self, task_description: str) -> Tuple[str, str]:
+    async def abuild(self, task_description: str, multimodal: bool = False) -> Tuple[str, int, str]:
         """Asynchronously builds a specialized function given a task description.
 
         Parameters
         ----------
         task_description : str
             A description of the task for which the function is being built.
+        
+        multimodal : bool, optional
+            Whether the function is multimodal or not (default is False).
 
         Returns
         -------
         tuple[str, str]
-            A tuple containing the name and description of the function.
+            A tuple containing the name, version number and description of the function.
         """
-        return await self._build(task_description=task_description, is_async=True)
+        return await self._build(task_description=task_description, multimodal=multimodal, is_async=True)
 
-    def build(self, task_description: str) -> Tuple[str, str]:
+    def build(self, task_description: str, multimodal: bool = False) -> Tuple[str, int, str]:
         """Synchronously builds a specialized function given a task description.
 
         Parameters
         ----------
         task_description : str
             A description of the task for which the function is being built.
+        
+        multimodal : bool, optional
+            Whether the function is multimodal or not (default is False).
 
         Returns
         -------
         tuple[str, str]
-            A tuple containing the name and description of the function.
+            A tuple containing the name, version number and description of the function.
         """
-        return self._build(task_description=task_description, is_async=False)
+        return self._build(task_description=task_description, multimodal=multimodal, is_async=False)
 
-    def _upload_image(self, fn_name: str, upload_id: str, image_info: Dict[str, Any]) -> str:
+    def _upload_image(self, fn_name: str, version_number: int, upload_id: str, image_info: Dict[str, Any]) -> str:
         """
         Uploads an image to an S3 bucket and returns the URL of the uploaded image.
 
@@ -237,6 +242,8 @@ class WecoAI:
         ----------
         fn_name : str
             The name of the function for which the image is being uploaded.
+        version_number : int
+            The version number of the function for which the image is being uploaded.
         upload_id: str
             A unique identifier for the image upload.
         image_info : Dict[str, Any]
@@ -271,13 +278,12 @@ class WecoAI:
 
         # Request a presigned URL from the server
         endpoint = "upload_link"
-        request_data = {"fn_name": fn_name, "upload_id": upload_id, "file_type": file_type}
+        request_data = {"fn_name": fn_name, "version_number": version_number, "upload_id": upload_id, "file_type": file_type}
         # This needs to be a synchronous request since we need the presigned URL to upload the image
         response = self._make_request(endpoint=endpoint, data=request_data, is_async=False)
 
         # Upload the image to the S3 bucket
-        image_name = generate_random_base16_code()
-        files = {"file": (f"{image_name}.{file_type}", upload_data)}
+        files = {"file": upload_data}
         http_response = requests.post(response["url"], data=response["fields"], files=files)
         if http_response.status_code == 204:
             pass
@@ -372,7 +378,7 @@ class WecoAI:
         return image_info
 
     def _query(
-        self, is_async: bool, fn_name: str, text_input: Optional[str], images_input: Optional[List[str]]
+        self, is_async: bool, fn_name: str, version_number: Optional[int], text_input: Optional[str], images_input: Optional[List[str]]
     ) -> Union[Dict[str, Any], Coroutine[Any, Any, Dict[str, Any]]]:
         """Internal method to handle both synchronous and asynchronous query requests.
 
@@ -382,6 +388,8 @@ class WecoAI:
             Whether to perform an asynchronous request.
         fn_name : str
             The name of the function to query.
+        version_number : int, optional
+            The version number of the function to query.
         text_input : str, optional
             The text input to the function.
         images_input : List[str], optional
@@ -397,6 +405,8 @@ class WecoAI:
         ValueError
             If the input is invalid.
         """
+        warnings.warn("Setting the version number of the function is not yet supported. Currently, the first version of the function will be used i.e., version 0.")
+        version_number = 0
         # Validate the input
         image_info = self._validate_query(text_input=text_input, images_input=images_input)
 
@@ -405,14 +415,14 @@ class WecoAI:
         upload_id = generate_random_base16_code()
         for i, info in enumerate(image_info):
             if info["source"] == "url" or info["source"] == "base64" or info["source"] == "local":
-                url = self._upload_image(fn_name=fn_name, upload_id=upload_id, image_info=info)
+                url = self._upload_image(fn_name=fn_name, version_number=version_number, upload_id=upload_id, image_info=info)
             else:
                 raise ValueError(f"Image at index {i} must be a public URL or a path to a local image file.")
             image_urls.append(url)
 
         # Make the request
         endpoint = "query"
-        data = {"name": fn_name, "text": text_input, "images": image_urls}
+        data = {"name": fn_name, "text": text_input, "images": image_urls, "version_number": version_number}
         request = self._make_request(endpoint=endpoint, data=data, is_async=is_async)
 
         if is_async:
@@ -427,7 +437,7 @@ class WecoAI:
             return self._process_query_response(response=response)
 
     async def aquery(
-        self, fn_name: str, text_input: Optional[str] = "", images_input: Optional[List[str]] = []
+        self, fn_name: str, version_number: Optional[int] = -1, text_input: Optional[str] = "", images_input: Optional[List[str]] = []
     ) -> Dict[str, Any]:
         """Asynchronously queries a function with the given function ID and input.
 
@@ -435,6 +445,8 @@ class WecoAI:
         ----------
         fn_name : str
             The name of the function to query.
+        version_number : int, optional
+            The version number of the function to query. If not provided, the latest version will be used. Pass -1 to use the latest version.
         text_input : str, optional
             The text input to the function.
         images_input : List[str], optional
@@ -446,15 +458,17 @@ class WecoAI:
             A dictionary containing the output of the function, the number of input tokens, the number of output tokens,
             and the latency in milliseconds.
         """
-        return await self._query(fn_name=fn_name, text_input=text_input, images_input=images_input, is_async=True)
+        return await self._query(fn_name=fn_name, version_number=version_number, text_input=text_input, images_input=images_input, is_async=True)
 
-    def query(self, fn_name: str, text_input: Optional[str] = "", images_input: Optional[List[str]] = []) -> Dict[str, Any]:
+    def query(self, fn_name: str, version_number: Optional[int] = -1, text_input: Optional[str] = "", images_input: Optional[List[str]] = []) -> Dict[str, Any]:
         """Synchronously queries a function with the given function ID and input.
 
         Parameters
         ----------
         fn_name : str
             The name of the function to query.
+        version_number : int, optional
+            The version number of the function to query. If not provided, the latest version will be used. Pass -1 to use the latest version.
         text_input : str, optional
             The text input to the function.
         images_input : List[str], optional
@@ -466,49 +480,35 @@ class WecoAI:
                A dictionary containing the output of the function, the number of input tokens, the number of output tokens,
             and the latency in milliseconds.
         """
-        return self._query(fn_name=fn_name, text_input=text_input, images_input=images_input, is_async=False)
-
-    def batch_query(self, fn_names: Union[str, List[str]], batch_inputs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """Synchronously queries multiple functions using asynchronous calls internally.
+        return self._query(fn_name=fn_name, version_number=version_number, text_input=text_input, images_input=images_input, is_async=False)
 
-        This method uses the asynchronous queries to submit all queries concurrently
-        and waits for all responses to be received before returning the results.
+    def batch_query(self, fn_name: str, batch_inputs: List[Dict[str, Any]], version_number: Optional[int] = -1) -> List[Dict[str, Any]]:
+        """Batch queries a function version with a list of inputs.
 
         Parameters
         ----------
-        fn_name : Union[str, List[str]]
+        fn_name : str
             The name of the function or a list of function names to query.
-            Note that if a single function name is provided, it will be used for all queries.
-            If a list of function names is provided, the length must match the number of queries.
 
         batch_inputs : List[Dict[str, Any]]
             A list of inputs for the functions to query. The input must be a dictionary containing the data to be processed. e.g.,
             when providing for a text input, the dictionary should be {"text_input": "input text"}, for an image input, the dictionary should be {"images_input": ["url1", "url2", ...]}
             and for a combination of text and image inputs, the dictionary should be {"text_input": "input text", "images_input": ["url1", "url2", ...]}.
-            Note that the index of each input must correspond to the index of the function name when both inputs are lists.
+        
+        version_number : int, optional
+            The version number of the function to query. If not provided, the latest version will be used. Pass -1 to use the latest version.
 
         Returns
         -------
         List[Dict[str, Any]]
             A list of dictionaries, each containing the output of a function query,
             in the same order as the input queries.
-
-
-        Raises
-        ------
-        ValueError
-            If the number of function names (when provided as a list) does not match the number of inputs.
         """
-        if isinstance(fn_names, str):
-            fn_names = [fn_names] * len(batch_inputs)
-        elif len(fn_names) != len(batch_inputs):
-            raise ValueError("The number of function names must match the number of inputs.")
-
         async def run_queries():
-            tasks = [
-                self.aquery(fn_name=fn_name, **fn_input)  # unpack the input kwargs
-                for fn_name, fn_input in zip(fn_names, batch_inputs)
-            ]
+            tasks = list(map(
+                lambda fn_input: self.aquery(fn_name=fn_name, version_number=version_number, **fn_input),
+                batch_inputs
+            ))
             return await asyncio.gather(*tasks)
 
         return asyncio.run(run_queries())

{weco-0.1.5 → weco-0.1.6}/weco/functional.py

@@ -3,14 +3,15 @@ from typing import Any, Dict, List, Optional
 from .client import WecoAI
 
 
-# TODO: Implement the closing stuff for the client
-def build(task_description: str, api_key: str = None) -> tuple[str, str]:
+def build(task_description: str, multimodal: bool = False, api_key: str = None) -> tuple[str, int, str]:
     """Builds a specialized function synchronously given a task description.
 
     Parameters
     ----------
     task_description : str
         A description of the task for which the function is being built.
+    multimodal : bool, optional
+        A flag to indicate if the function should be multimodal. Default is False.
     api_key : str
         The API key for the WecoAI service. If not provided, the API key must be set using the environment variable - WECO_API_KEY.
 
@@ -20,32 +21,34 @@ def build(task_description: str, api_key: str = None) -> tuple[str, str]:
         A tuple containing the name and description of the function.
     """
     client = WecoAI(api_key=api_key)
-    response = client.build(task_description=task_description)
+    response = client.build(task_description=task_description, multimodal=multimodal)
     return response
 
 
-async def abuild(task_description: str, api_key: str = None) -> tuple[str, str]:
+async def abuild(task_description: str, multimodal: bool = False, api_key: str = None) -> tuple[str, int, str]:
     """Builds a specialized function asynchronously given a task description.
 
     Parameters
     ----------
     task_description : str
         A description of the task for which the function is being built.
+    multimodal : bool, optional
+        A flag to indicate if the function should be multimodal. Default is False.
     api_key : str
         The API key for the WecoAI service. If not provided, the API key must be set using the environment variable - WECO_API_KEY.
 
     Returns
     -------
     tuple[str, str]
-        A tuple containing the name and description of the function.
+        A tuple containing the name, version number and description of the function.
     """
     client = WecoAI(api_key=api_key)
-    response = await client.abuild(task_description=task_description)
+    response = await client.abuild(task_description=task_description, multimodal=multimodal)
     return response
 
 
 def query(
-    fn_name: str, text_input: Optional[str] = "", images_input: Optional[List[str]] = [], api_key: Optional[str] = None
+    fn_name: str, version_number: Optional[int] = -1, text_input: Optional[str] = "", images_input: Optional[List[str]] = [], api_key: Optional[str] = None
 ) -> Dict[str, Any]:
     """Queries a function synchronously with the given function ID and input.
 
@@ -53,6 +56,8 @@ def query(
     ----------
     fn_name : str
         The name of the function to query.
+    version_number : int, optional
+        The version number of the function to query. If not provided, the latest version is used. Default is -1 for the same behavior.
     text_input : str, optional
         The text input to the function.
     images_input : List[str], optional
@@ -67,12 +72,12 @@ def query(
         and the latency in milliseconds.
     """
     client = WecoAI(api_key=api_key)
-    response = client.query(fn_name=fn_name, text_input=text_input, images_input=images_input)
+    response = client.query(fn_name=fn_name, version_number=version_number, text_input=text_input, images_input=images_input)
     return response
 
 
 async def aquery(
-    fn_name: str, text_input: Optional[str] = "", images_input: Optional[List[str]] = [], api_key: Optional[str] = None
+    fn_name: str, version_number: Optional[int] = -1, text_input: Optional[str] = "", images_input: Optional[List[str]] = [], api_key: Optional[str] = None
 ) -> Dict[str, Any]:
     """Queries a function asynchronously with the given function ID and input.
 
@@ -80,6 +85,8 @@ async def aquery(
     ----------
     fn_name : str
         The name of the function to query.
+    version_number : int, optional
+        The version number of the function to query. If not provided, the latest version is used. Default is -1 for the same behavior.
     text_input : str, optional
         The text input to the function.
     images_input : List[str], optional
@@ -94,12 +101,12 @@ async def aquery(
         and the latency in milliseconds.
     """
     client = WecoAI(api_key=api_key)
-    response = await client.aquery(fn_name=fn_name, text_input=text_input, images_input=images_input)
+    response = await client.aquery(fn_name=fn_name, version_number=version_number, text_input=text_input, images_input=images_input)
     return response
 
 
 def batch_query(
-    fn_names: str | List[str], batch_inputs: List[Dict[str, Any]], api_key: Optional[str] = None
+    fn_name: str, batch_inputs: List[Dict[str, Any]], version_number: Optional[int] = -1, api_key: Optional[str] = None
 ) -> List[Dict[str, Any]]:
     """Synchronously queries multiple functions using asynchronous calls internally.
 
@@ -117,7 +124,9 @@ def batch_query(
        A list of inputs for the functions to query. The input must be a dictionary containing the data to be processed. e.g.,
        when providing for a text input, the dictionary should be {"text_input": "input text"}, for an image input, the dictionary should be {"images_input": ["url1", "url2", ...]}
        and for a combination of text and image inputs, the dictionary should be {"text_input": "input text", "images_input": ["url1", "url2", ...]}.
-       Note that the index of each input must correspond to the index of the function name when both inputs are lists.
+
+    version_number : int, optional
+        The version number of the function to query. If not provided, the latest version is used. Default is -1 for the same behavior.
 
     api_key : str, optional
         The API key for the WecoAI service. If not provided, the API key must be set using the environment variable - WECO_API_KEY.
@@ -129,5 +138,5 @@ def batch_query(
         in the same order as the input queries.
     """
     client = WecoAI(api_key=api_key)
-    responses = client.batch_query(fn_names=fn_names, batch_inputs=batch_inputs)
+    responses = client.batch_query(fn_name=fn_name, version_number=version_number, batch_inputs=batch_inputs)
     return responses

{weco-0.1.5 → weco-0.1.6}/weco.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: weco
-Version: 0.1.5
+Version: 0.1.6
 Summary: A client facing API for interacting with the WeCo AI function builder service.
 Author-email: WeCo AI Team <dhruv@weco.ai>
 License: MIT

weco-0.1.5/tests/test_batching.py

@@ -1,66 +0,0 @@
-import pytest
-
-from weco import batch_query, build
-
-
-# Internally, these functions use the WecoAI client
-# therefore, we can test both the client and functional forms here
-# We do NOT need to test for a list of function names here since interally
-# we cast the input to a list of function names to match the fn_input size
-@pytest.fixture
-def ml_task_evaluator():
-    fn_name, _ = build(
-        task_description="I want to evaluate the feasibility of a machine learning task. Give me a json object with three keys - 'feasibility', 'justification', and 'suggestions'."
-    )
-    return fn_name
-
-
-@pytest.fixture
-def ml_task_inputs():
-    return [
-        {"text_input": "I want to train a model to predict house prices using the Boston Housing dataset hosted on Kaggle."},
-        {
-            "text_input": "I want to train a model to classify digits using the MNIST dataset hosted on Kaggle using a Google Colab notebook."
-        },
-    ]
-
-
-@pytest.fixture
-def image_evaluator():
-    fn_name, _ = build(
-        task_description="Describe the contents of the given images. Provide a json object with 'description' and 'objects' keys."
-    )
-    return fn_name
-
-
-@pytest.fixture
-def image_inputs():
-    return [
-        {
-            "images_input": [
-                "https://www.integratedtreatmentservices.co.uk/wp-content/uploads/2013/12/Objects-of-Reference.jpg"
-            ]
-        },
-        {"images_input": ["https://t4.ftcdn.net/jpg/05/70/90/23/360_F_570902339_kNj1reH40GFXakTy98EmfiZHci2xvUCS.jpg"]},
-    ]
-
-
-def test_batch_query(ml_task_evaluator, ml_task_inputs, image_evaluator, image_inputs):
-    fn_names = [ml_task_evaluator, ml_task_evaluator, image_evaluator, image_evaluator]
-    batch_inputs = ml_task_inputs + image_inputs
-
-    query_responses = batch_query(fn_names=fn_names, batch_inputs=batch_inputs)
-
-    assert len(query_responses) == len(batch_inputs)
-
-    for i, query_response in enumerate(query_responses):
-        assert isinstance(query_response["output"], dict)
-        assert isinstance(query_response["in_tokens"], int)
-        assert isinstance(query_response["out_tokens"], int)
-        assert isinstance(query_response["latency_ms"], float)
-
-        output = query_response["output"]
-        if i < len(ml_task_inputs):
-            assert set(output.keys()) == {"feasibility", "justification", "suggestions"}
-        else:
-            assert set(output.keys()) == {"description", "objects"}