weco 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

weco/client.py

@@ -1,9 +1,24 @@
 import asyncio
+import base64
 import os
 import warnings
-from typing import Any, Callable, Coroutine, Dict, List, Tuple
+from io import BytesIO
+from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple, Union
 
 import httpx
+import requests
+from httpx import HTTPStatusError
+from PIL import Image
+
+from .constants import MAX_IMAGE_SIZE_MB, MAX_IMAGE_UPLOADS, MAX_TEXT_LENGTH, SUPPORTED_IMAGE_EXTENSIONS
+from .utils import (
+    generate_random_base16_code,
+    get_image_size,
+    is_base64_image,
+    is_local_image,
+    is_public_url_image,
+    preprocess_image,
+)
 
 
 class WecoAI:
@@ -17,7 +32,7 @@ class WecoAI:
         The API key used for authentication.
     """
 
-    def __init__(self, api_key: str = None, timeout: float = 30.0) -> None:
+    def __init__(self, api_key: str = None, timeout: float = 120.0, http2: bool = True) -> None:
         """Initializes the WecoAI client with the provided API key and base URL.
 
         Parameters
@@ -25,6 +40,12 @@ class WecoAI:
         api_key : str, optional
             The API key used for authentication. If not provided, the client will attempt to read it from the environment variable - WECO_API_KEY.
 
+        timeout : float, optional
+            The timeout for the HTTP requests in seconds (default is 30.0).
+
+        http2 : bool, optional
+            Whether to use HTTP/2 protocol for the HTTP requests (default is True).
+
         Raises
         ------
         ValueError
@@ -37,30 +58,17 @@ class WecoAI:
             except KeyError:
                 raise ValueError("WECO_API_KEY must be passed to client or set as an environment variable")
         self.api_key = api_key
-
-        self.base_url = "https://function-builder.vercel.app"
-
+        self.http2 = http2
+        self.timeout = timeout
+        self.base_url = "https://function.api.weco.ai"
         # Setup clients
-        self.client = httpx.Client(http2=False, timeout=timeout)
-        self.async_client = httpx.AsyncClient(http2=False, timeout=timeout)
-
-    def __del__(self):
-        """Closes the HTTP clients when the WecoAI instance is deleted."""
-        try:
-            self.client.close()
-            if not self.async_client.is_closed:
-                try:
-                    loop = asyncio.get_event_loop()
-                    if loop.is_running():
-                        loop.create_task(self.async_client.aclose())
-                    else:
-                        loop.run_until_complete(self.async_client.aclose())
-                except RuntimeError:
-                    # If the event loop is closed, we can't do anything about it
-                    pass
-        except AttributeError:
-            # If the client is not initialized, we can't do anything about it
-            pass
+        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."""
@@ -92,22 +100,38 @@ class WecoAI:
         if is_async:
 
             async def _request():
-                response = await self.async_client.post(url, json=data, headers=headers)
-                response.raise_for_status()
-                return response.json()
+                try:
+                    response = await self.async_client.post(url, json=data, headers=headers)
+                    response.raise_for_status()
+                    return response.json()
+                except HTTPStatusError as e:
+                    # Handle HTTP errors (4xx and 5xx status codes)
+                    error_message = f"HTTP error occurred: {e.response.status_code} - {e.response.text}"
+                    raise ValueError(error_message) from e
+                except Exception as e:
+                    # Handle other exceptions
+                    raise ValueError(f"An error occurred: {str(e)}") from e
 
             return _request()
         else:
 
             def _request():
-                response = self.client.post(url, json=data, headers=headers)
-                response.raise_for_status()
-                return response.json()
+                try:
+                    response = self.client.post(url, json=data, headers=headers)
+                    response.raise_for_status()
+                    return response.json()
+                except HTTPStatusError as e:
+                    # Handle HTTP errors (4xx and 5xx status codes)
+                    error_message = f"HTTP error occurred: {e.response.status_code} - {e.response.text}"
+                    raise ValueError(error_message) from e
+                except Exception as e:
+                    # Handle other exceptions
+                    raise ValueError(f"An error occurred: {str(e)}") from e
 
             return _request()
 
-    def _process_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
-        """Processes the API response and handles warnings.
+    def _process_query_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
+        """Processes the query response and handles warnings.
 
         Parameters
         ----------
@@ -134,7 +158,7 @@ class WecoAI:
             "latency_ms": response["latency_ms"],
         }
 
-    def _build(self, task_description: str, is_async: bool) -> Tuple[str, str] | Coroutine[Any, Any, Tuple[str, str]]:
+    def _build(self, task_description: str, is_async: bool) -> Union[Tuple[str, str], Coroutine[Any, Any, Tuple[str, str]]]:
         """Internal method to handle both synchronous and asynchronous build requests.
 
         Parameters
@@ -146,9 +170,20 @@ class WecoAI:
 
         Returns
         -------
-        tuple[str, str] | Coroutine[Any, Any, tuple[str, str]]
+        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.
+
+        Raises
+        ------
+        ValueError
+            If the task description is empty or exceeds the maximum length.
         """
+        # Validate the input
+        if len(task_description) == 0:
+            raise ValueError("Task description must be provided.")
+        if len(task_description) > MAX_TEXT_LENGTH:
+            raise ValueError(f"Task description must be less than {MAX_TEXT_LENGTH} characters.")
+
         endpoint = "build"
         data = {"request": task_description}
         request = self._make_request(endpoint=endpoint, data=data, is_async=is_async)
@@ -194,47 +229,216 @@ class WecoAI:
         """
         return self._build(task_description=task_description, is_async=False)
 
-    def _query(self, fn_name: str, fn_input: str, is_async: bool) -> Dict[str, Any] | Coroutine[Any, Any, Dict[str, Any]]:
-        """Internal method to handle both synchronous and asynchronous query requests.
+    def _upload_image(self, fn_name: str, upload_id: str, image_info: Dict[str, Any]) -> str:
+        """
+        Uploads an image to an S3 bucket and returns the URL of the uploaded image.
 
         Parameters
         ----------
         fn_name : str
-            The name of the function to query.
-        fn_input : str
-            The input to the function.
+            The name 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]
+            A dictionary containing the image metadata.
+
+        Returns
+        -------
+        str
+            The URL of the uploaded image.
+        """
+
+        if image_info["source"] == "base64":
+            _, base64_info = is_base64_image(maybe_base64=image_info["image"])
+            img_data = base64.b64decode(base64_info["encoding"])
+        elif image_info["source"] == "url":
+            response = requests.get(image_info["image"])
+            response.raise_for_status()
+            img_data = response.content
+        elif image_info["source"] == "local":
+            with open(image_info["image"], "rb") as f:
+                img_data = f.read()
+        else:
+            raise ValueError("Invalid image input")
+
+        # Preprocess the image
+        img = Image.open(BytesIO(img_data))
+        file_type = image_info["file_type"]
+        processed_img, file_type = preprocess_image(image=img, file_type=file_type)
+        upload_data = BytesIO()
+        processed_img.save(upload_data, format=file_type)
+        upload_data = upload_data.getvalue()
+
+        # Request a presigned URL from the server
+        endpoint = "upload_link"
+        request_data = {"fn_name": fn_name, "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)}
+        http_response = requests.post(response["url"], data=response["fields"], files=files)
+        if http_response.status_code == 204:
+            pass
+        else:
+            raise ValueError("Image upload failed")
+
+        # Return the URL of the uploaded image
+        upload_link = f"{response['url']}{response['fields']['key']}"
+        return upload_link
+
+    def _validate_query(self, text_input: str, images_input: List[str]) -> List[Dict[str, Any]]:
+        """
+        Validate the input for the query method.
+
+        Parameters
+        ----------
+        text_input : str
+            The text input to the function.
+        images_input : List[str]
+            A list of image URLs or images encoded in base64 with their metadata to be sent as input to the function.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            A list of dictionaries containing the image metadata.
+
+        Raises
+        ------
+        ValueError
+            If the input is invalid.
+        """
+        if not isinstance(text_input, str) or not isinstance(images_input, list):
+            raise ValueError("Text input must be a string and images input must be a list of strings.")
+        for image in images_input:
+            if not isinstance(image, str):
+                raise ValueError("Images input must be a list of strings.")
+
+        # Assert that either text or images or both must be provded
+        if len(text_input) == 0 and len(images_input) == 0:
+            raise ValueError("Either text or images or both must be provided as input.")
+
+        # Check if the text input is within the limit
+        if len(text_input) > MAX_TEXT_LENGTH:
+            raise ValueError(f"Text input must be less than {MAX_TEXT_LENGTH} characters.")
+
+        # Check if the images input is within the limit
+        if len(images_input) > MAX_IMAGE_UPLOADS:
+            raise ValueError(f"Number of images must be less than {MAX_IMAGE_UPLOADS}.")
+
+        # Check if input is an valid image
+        image_info = []
+        for image in images_input:
+            is_base64, base64_info = is_base64_image(maybe_base64=image)
+            if is_base64:
+                file_type = base64_info["media_type"].split("/")[1]
+
+            is_public_url = is_public_url_image(maybe_url_image=image)
+            if is_public_url:
+                response = requests.get(image)
+                response.raise_for_status()
+                file_type = response.headers["content-type"].split("/")[1]
+
+            is_local = is_local_image(maybe_local_image=image)
+            if is_local:
+                file_type = os.path.splitext(image)[1][1:]
+
+            if not (is_base64 or is_public_url or is_local):
+                raise ValueError("Images must be local paths, public URLs or base64 encoded strings.")
+
+            # Determine the source of image
+            if is_base64:
+                source = "base64"
+            elif is_public_url:
+                source = "url"
+            elif is_local:
+                source = "local"
+
+            # Check if the image type is supported
+            file_type = file_type.lower()
+            if file_type not in SUPPORTED_IMAGE_EXTENSIONS:
+                raise ValueError(
+                    f"Image file type {file_type} is not supported. Supported types are {SUPPORTED_IMAGE_EXTENSIONS}."
+                )
+
+            # Check if the image size is within the limit
+            size = get_image_size(image=image, source=source)
+            if size > MAX_IMAGE_SIZE_MB:
+                raise ValueError(f"Individual image sizes must be less than {MAX_IMAGE_SIZE_MB} MB each.")
+
+            image_info.append({"image": image, "file_type": file_type, "size": size, "source": source})
+
+        return image_info
+
+    def _query(
+        self, is_async: bool, fn_name: str, 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.
+
+        Parameters
+        ----------
         is_async : bool
             Whether to perform an asynchronous request.
+        fn_name : str
+            The name of the function to query.
+        text_input : str, optional
+            The text input to the function.
+        images_input : List[str], optional
+            A list of image URLs or images encoded in base64 with their metadata to be sent as input to the function.
 
         Returns
         -------
-        dict | Coroutine[Any, Any, dict]
+        Union[Dict[str, Any], Coroutine[Any, Any, dict]]
             A dictionary containing the query results, or a coroutine that returns such a dictionary.
+
+        Raises
+        ------
+        ValueError
+            If the input is invalid.
         """
+        # Validate the input
+        image_info = self._validate_query(text_input=text_input, images_input=images_input)
+
+        # Create links for all images that are not public URLs and upload images
+        image_urls = []
+        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)
+            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, "user_message": fn_input}
+        data = {"name": fn_name, "text": text_input, "images": image_urls}
         request = self._make_request(endpoint=endpoint, data=data, is_async=is_async)
 
         if is_async:
 
             async def _async_query():
                 response = await request
-                return self._process_response(response=response)
+                return self._process_query_response(response=response)
 
             return _async_query()
         else:
             response = request  # the request has already been made and the response is available
-            return self._process_response(response=response)
+            return self._process_query_response(response=response)
 
-    async def aquery(self, fn_name: str, fn_input: str) -> Dict[str, Any]:
+    async def aquery(
+        self, fn_name: str, text_input: Optional[str] = "", images_input: Optional[List[str]] = []
+    ) -> Dict[str, Any]:
         """Asynchronously queries a function with the given function ID and input.
 
         Parameters
         ----------
         fn_name : str
             The name of the function to query.
-        fn_input : str
-            The input to the function.
+        text_input : str, optional
+            The text input to the function.
+        images_input : List[str], optional
+            A list of image URLs or images encoded in base64 with their metadata to be sent as input to the function.
 
         Returns
         -------
@@ -242,17 +446,19 @@ 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, fn_input=fn_input, is_async=True)
+        return await self._query(fn_name=fn_name, text_input=text_input, images_input=images_input, is_async=True)
 
-    def query(self, fn_name: str, fn_input: str) -> Dict[str, Any]:
+    def query(self, fn_name: str, 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.
-        fn_input : str
-            The input to the function.
+        text_input : str, optional
+            The text input to the function.
+        images_input : List[str], optional
+            A list of image URLs or images encoded in base64 with their metadata to be sent as input to the function.
 
         Returns
         -------
@@ -260,9 +466,9 @@ 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, fn_input=fn_input, is_async=False)
+        return self._query(fn_name=fn_name, text_input=text_input, images_input=images_input, is_async=False)
 
-    def batch_query(self, fn_names: str | List[str], batch_inputs: List[str]) -> List[Dict[str, Any]]:
+    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.
 
         This method uses the asynchronous queries to submit all queries concurrently
@@ -270,20 +476,28 @@ class WecoAI:
 
         Parameters
         ----------
-        fn_name : str | List[str]
+        fn_name : Union[str, List[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[str]
-            A list of inputs for the functions to query.
-            Note that the index of each input must correspond to the index of the function name.
+        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.
 
         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)
@@ -291,7 +505,10 @@ class WecoAI:
             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=fn_input) for fn_name, fn_input in zip(fn_names, batch_inputs)]
+            tasks = [
+                self.aquery(fn_name=fn_name, **fn_input)  # unpack the input kwargs
+                for fn_name, fn_input in zip(fn_names, batch_inputs)
+            ]
             return await asyncio.gather(*tasks)
 
         return asyncio.run(run_queries())

weco/constants.py

@@ -0,0 +1,4 @@
+MAX_TEXT_LENGTH = 10000
+MAX_IMAGE_UPLOADS = 5
+MAX_IMAGE_SIZE_MB = 10  # 10 MB
+SUPPORTED_IMAGE_EXTENSIONS = ["png", "jpeg", "jpg", "webp", "gif"]

weco/functional.py

@@ -3,6 +3,7 @@ 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]:
     """Builds a specialized function synchronously given a task description.
 
@@ -43,15 +44,19 @@ async def abuild(task_description: str, api_key: str = None) -> tuple[str, str]:
     return response
 
 
-def query(fn_name: str, fn_input: str, api_key: Optional[str] = None) -> Dict[str, Any]:
+def query(
+    fn_name: str, 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.
 
     Parameters
     ----------
     fn_name : str
         The name of the function to query.
-    fn_input : str
-        The input to the function.
+    text_input : str, optional
+        The text input to the function.
+    images_input : List[str], optional
+        A list of image URLs or base64 encoded images to be used as input to the function.
     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.
 
@@ -62,19 +67,23 @@ def query(fn_name: str, fn_input: str, api_key: Optional[str] = None) -> Dict[st
         and the latency in milliseconds.
     """
     client = WecoAI(api_key=api_key)
-    response = client.query(fn_name=fn_name, fn_input=fn_input)
+    response = client.query(fn_name=fn_name, text_input=text_input, images_input=images_input)
     return response
 
 
-async def aquery(fn_name: str, fn_input: str, api_key: Optional[str] = None) -> Dict[str, Any]:
+async def aquery(
+    fn_name: str, 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.
 
     Parameters
     ----------
     fn_name : str
         The name of the function to query.
-    fn_input : str
-        The input to the function.
+    text_input : str, optional
+        The text input to the function.
+    images_input : List[str], optional
+        A list of image URLs to be used as input to the function.
     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.
 
@@ -85,11 +94,13 @@ async def aquery(fn_name: str, fn_input: str, api_key: Optional[str] = None) ->
         and the latency in milliseconds.
     """
     client = WecoAI(api_key=api_key)
-    response = await client.aquery(fn_name=fn_name, fn_input=fn_input)
+    response = await client.aquery(fn_name=fn_name, text_input=text_input, images_input=images_input)
     return response
 
 
-def batch_query(fn_names: str | List[str], batch_inputs: List[str], api_key: Optional[str] = None) -> List[Dict[str, Any]]:
+def batch_query(
+    fn_names: str | List[str], batch_inputs: List[Dict[str, Any]], api_key: Optional[str] = None
+) -> List[Dict[str, Any]]:
     """Synchronously queries multiple functions using asynchronous calls internally.
 
     This method uses the asynchronous queries to submit all queries concurrently
@@ -103,8 +114,10 @@ def batch_query(fn_names: str | List[str], batch_inputs: List[str], api_key: Opt
         If a list of function names is provided, the length must match the number of queries.
 
     batch_inputs : List[str]
-        A list of inputs for the functions to query.
-        Note that the index of each input must correspond to the index of the function name.
+       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.
 
     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.

weco/utils.py

@@ -0,0 +1,180 @@
+import base64
+import os
+import random
+import re
+import string
+from io import BytesIO
+from typing import Dict, Optional, Tuple
+from urllib.parse import urlparse
+
+import requests
+from PIL import Image
+
+
+def is_local_image(maybe_local_image: str) -> bool:
+    """
+    Check if the file is a local image.
+
+    Parameters
+    ----------
+    maybe_local_image : str
+        The file path.
+
+    Returns
+    -------
+    bool
+        True if the file is a local image, False otherwise.
+    """
+    if not os.path.exists(maybe_local_image):  # Check if the file exists
+        return False
+
+    try:  # Check if the file is an image
+        Image.open(maybe_local_image)
+    except IOError:
+        return False
+
+    return True
+
+
+def is_base64_image(maybe_base64: str) -> Tuple[bool, Optional[Dict[str, str]]]:
+    """
+    Check if the image is a base64 encoded image and if so, extract the image information from the encoded data or URL provided.
+
+    Parameters
+    ----------
+    data : str
+        The image data or URL.
+
+    Returns
+    -------
+    Tuple[bool, Optional[Dict[str, str]]]
+    """
+    pattern = r"data:(?P<media_type>[\w/]+);(?P<source_type>\w+),(?P<encoding>.*)"
+    match = re.match(pattern, maybe_base64)
+    if match:
+        return True, match.groupdict()
+
+    return False, None
+
+
+def is_public_url_image(maybe_url_image: str) -> bool:
+    """
+    Check if the string is a publicly accessible URL
+
+    Parameters
+    ----------
+    maybe_url_image : str
+        The URL to check.
+
+    Returns
+    -------
+    bool
+        True if the URL is publicly accessible, False otherwise.
+    """
+    try:
+        # Check if it is a valid URL
+        if not urlparse(maybe_url_image).scheme:
+            return False
+
+        # Check if the URL is publicly accessible
+        response = requests.head(maybe_url_image)
+        if response.status_code != 200:
+            return False
+
+        # Check if the URL is an image
+        content_type = response.headers.get("content-type")
+        if not content_type:
+            return False
+        if not content_type.startswith("image"):
+            return False
+    except Exception:
+        return False
+
+    return True
+
+
+def get_image_size(image: str, source: str) -> float:
+    """
+    Get the size of the image in MB.
+
+    Parameters
+    ----------
+    image : str
+        The image data or URL.
+
+    source : str
+        The source of the image. It can be 'base64', 'url', or 'local'.
+
+    Returns
+    -------
+    float
+        The size of the image in MB.
+
+    Raises
+    ------
+    ValueError
+        If the image is not a valid image.
+    """
+    if source == "base64":
+        _, base64_info = is_base64_image(maybe_base64=image)
+        img_data = base64.b64decode(base64_info["encoding"])
+    elif source == "url":
+        response = requests.get(image)
+        response.raise_for_status()
+        img_data = response.content
+    elif source == "local":
+        with open(image, "rb") as f:
+            img_data = f.read()
+    else:
+        raise ValueError("Invalid image input")
+
+    img = Image.open(BytesIO(img_data))
+    img_byte_arr = BytesIO()
+    img.save(img_byte_arr, format=img.format)
+    return img_byte_arr.tell() / 1000000  # MB
+
+
+def preprocess_image(image: Image, file_type: str) -> Tuple:
+    """
+    Preprocess the image by converting it to RGB if it has an alpha channel.
+
+    Parameters
+    ----------
+    image : Image
+        The image to preprocess.
+    file_type : str
+        The file type of the image.
+
+    Returns
+    -------
+    Image
+        The preprocessed image.
+    file_type : str
+        The file type of the image.
+    """
+    # Do not rescale or resize. Only do this if latency becomes an issue.
+    # Remove the alpha channel for PNG and WEBP images if it exists.
+    if image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info):
+        image = image.convert("RGB")
+
+    # If the image file type is JPG, convert to JPEG for PIL compatibility.
+    if file_type == "jpg":
+        file_type = "jpeg"
+    return image, file_type
+
+
+def generate_random_base16_code(length: int = 5):
+    """
+    Generate a random base16 code.
+
+    Parameters
+    ----------
+    length : int
+        The length of the code.
+
+    Returns
+    -------
+    str
+        The random base16 code.
+    """
+    return "".join(random.choices(string.hexdigits, k=length))

{weco-0.1.3.dist-info → weco-0.1.4.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.1
 Name: weco
-Version: 0.1.3
+Version: 0.1.4
 Summary: A client facing API for interacting with the WeCo AI function builder service.
 Author-email: WeCo AI Team <dhruv@weco.ai>
 License: MIT
 Project-URL: Homepage, https://github.com/WecoAI/weco-python
-Keywords: AI,LLM,machine learning,data science,function builder
+Keywords: AI,LLM,machine learning,data science,function builder,AI function
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Classifier: License :: OSI Approved :: MIT License
@@ -14,21 +14,31 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: asyncio
 Requires-Dist: httpx[http2]
+Requires-Dist: pillow
 Provides-Extra: dev
-Requires-Dist: build ; extra == 'dev'
-Requires-Dist: setuptools-scm ; extra == 'dev'
 Requires-Dist: flake8 ; extra == 'dev'
 Requires-Dist: flake8-pyproject ; extra == 'dev'
 Requires-Dist: black ; extra == 'dev'
 Requires-Dist: isort ; extra == 'dev'
+Requires-Dist: pytest-asyncio ; extra == 'dev'
+Requires-Dist: pytest-xdist ; extra == 'dev'
+Requires-Dist: build ; extra == 'dev'
+Requires-Dist: setuptools-scm ; extra == 'dev'
 
 <div align="center" style="display: flex; align-items: center; justify-content: center;">
   <img src="assets/weco.svg" alt="WeCo AI" style="height: 50px; margin-right: 10px;">
   <a href="https://git.io/typing-svg"><img src="https://readme-typing-svg.demolab.com?font=Georgia&size=32&duration=4000&pause=400&color=FD4578&vCenter=true&multiline=false&width=200&height=50&lines=WeCo+Client" alt="Typing SVG" /></a>
 </div>
 
+![Python](https://img.shields.io/badge/Python-3.10.14-blue)
+[![License](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
+
+<!-- TODO: Update examples -->
 # $f$(👷‍♂️)
 
+<a href="https://colab.research.google.com/github/WecoAI/weco-python/blob/main/examples/cookbook.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" width=110 height=20/></a>
+<a target="_blank" href="https://lightning.ai/new?repo_url=https%3A%2F%2Fgithub.com%2FWecoAI%2Fweco-python%2Fblob%2Fmain%2Fexamples%2Fcookbook.ipynb"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open in Studio" width=100 height=20/></a>
+
 A client facing API for interacting with the [WeCo AI](https://www.weco.ai/) function builder [service](https://weco-app.vercel.app/function)!
 
 
@@ -47,6 +57,7 @@ pip install weco
 - The **query** function allows you to test and use the newly created function in your own code.
 - We offer asynchronous versions of the above clients.
 - We provide a **batch_query** functions that allows users to batch functions for various inputs as well as multiple inputs for the same function in a query. This is helpful to make a large number of queries more efficiently.
+- We also offer multimodality capabilities. You can now query our client with both **language** AND **vision** inputs!
 
 We provide both services in two ways:
 - `weco.WecoAI` client to be used when you want to maintain the same client service across a portion of code. This is better for dense service usage.
@@ -63,18 +74,48 @@ export WECO_API_KEY=<YOUR_WECO_API_KEY>
 ## Example
 
 We create a function on the [web console](https://weco-app.vercel.app/function) for the following task:
-> "I want to evaluate the feasibility of a machine learning task. Give me a json object with three keys - 'feasibility', 'justification', and 'suggestions'."
+> "Analyze a business idea and provide a structured evaluation. Output a JSON with 'viability_score' (0-100), 'strengths' (list), 'weaknesses' (list), and 'next_steps' (list)."
 
 Now, you're ready to query this function anywhere in your code!
 
 ```python
 from weco import query
 response = query(
-    fn_name=fn_name,
-    fn_input="I want to train a model to predict house prices using the Boston Housing dataset hosted on Kaggle.",
+    fn_name="BusinessIdeaAnalyzer-XYZ123",  # Replace with your actual function name
+    text_input="A subscription service for personalized, AI-generated bedtime stories for children."
 )
 ```
 
 For more examples and an advanced user guide, check out our function builder [cookbook](examples/cookbook.ipynb).
 
 ## Happy building $f$(👷‍♂️)!
+
+## Contributing
+
+We value your contributions! If you believe you can help to improve our package enabling people to build AI with AI, please contribute!
+
+Use the following steps as a guideline to help you make contributions:
+
+1. Download and install package from source:
+   ```bash
+   git clone https://github.com/WecoAI/weco-python.git
+   cd weco-python
+   pip install -e ".[dev]"
+   ```
+
+2. Create a new branch for your feature or bugfix:
+   ```bash
+   git checkout -b feature/your-feature-name
+   ```
+
+3. Make your changes and run tests to ensure everything is working:
+   
+   > **Tests can be expensive to run as they make LLM requests with the API key being used so it is the developers best interests to write small and simple tests that adds coverage for a large portion of the package.**
+   
+   ```bash
+   pytest -n auto tests
+   ```
+
+4. Commit and push your changes, then open a PR for us to view 😁
+
+Please ensure your code follows our style guidelines (Numpy docstrings) and includes appropriate tests. We appreciate your contributions!

weco-0.1.4.dist-info/RECORD

@@ -0,0 +1,10 @@
+weco/__init__.py,sha256=qiKpnrm6t0n0bpAtXEKJO1Yz2xYXnJJRZBWt-cH7DdU,168
+weco/client.py,sha256=MGwLxV0pgZ_dighM5tvE32VbM1lspppTq8EmdZnXCfE,20608
+weco/constants.py,sha256=zIVqyFh6JQVWkPpeh7ABUvTmJ9CjQk_4LPlz4pEO0Uk,145
+weco/functional.py,sha256=MpyFaREjoPBWqY0iLgsN0JovL6cyCyNf5SzYI0i1b5w,5328
+weco/utils.py,sha256=UUSw6ocqWdlSmIXVcH66DAL4NuLU2rFOyviD8aTWsv0,4371
+weco-0.1.4.dist-info/LICENSE,sha256=NvpxfBuSajszAczWBGKxhHe4gsvil1H63zmu8xXZdL0,1064
+weco-0.1.4.dist-info/METADATA,sha256=YonF54xTkK71fPqBSuRmRcyqQDRh059LKxUpVvrZ94w,5957
+weco-0.1.4.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+weco-0.1.4.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.1.4.dist-info/RECORD,,

{weco-0.1.3.dist-info → weco-0.1.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.1.1)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

weco-0.1.3.dist-info/RECORD

@@ -1,8 +0,0 @@
-weco/__init__.py,sha256=qiKpnrm6t0n0bpAtXEKJO1Yz2xYXnJJRZBWt-cH7DdU,168
-weco/client.py,sha256=RZ6wsFohWhcJSz8HIGTIfVpXozUE_bYK7K037GSlZdA,10991
-weco/functional.py,sha256=oZljY8enxw8nPKP1TGBc8AxERt3Y9ukpV1_40vIJ8oE,4375
-weco-0.1.3.dist-info/LICENSE,sha256=NvpxfBuSajszAczWBGKxhHe4gsvil1H63zmu8xXZdL0,1064
-weco-0.1.3.dist-info/METADATA,sha256=vqzoh9iOLs7ZbSVjQwZz1jiLmSdEITITT0A7FWAdz0c,3876
-weco-0.1.3.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
-weco-0.1.3.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.1.3.dist-info/RECORD,,