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

weco/__init__.py

@@ -1,8 +1,4 @@
 from .client import WecoAI
-from .functional import build, query
+from .functional import abuild, aquery, batch_query, build, query
 
-__all__ = [
-    "WecoAI",
-    "build",
-    "query",
-]
+__all__ = ["WecoAI", "build", "abuild", "query", "aquery", "batch_query"]

weco/client.py

@@ -1,19 +1,29 @@
+import asyncio
 import os
-from typing import Any, Dict
+import warnings
+from typing import Any, Callable, Coroutine, Dict, List, Tuple
 
-import requests
+import httpx
 
 
 class WecoAI:
-    def __init__(self, api_key: str = None) -> None:
+    """A client for the WecoAI function builder API that allows users to build and query specialized functions built by LLMs.
+    The user must simply provide a task description to build a function, and then query the function with an input to get the result they need.
+    Our client supports both synchronous and asynchronous request paradigms and uses HTTP/2 for faster communication with the API.
+
+    Attributes
+    ----------
+    api_key : str
+        The API key used for authentication.
+    """
+
+    def __init__(self, api_key: str = None, timeout: float = 30.0) -> None:
         """Initializes the WecoAI client with the provided API key and base URL.
 
         Parameters
         ----------
-        api_key : str
+        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.
-        base_url : str
-            The base URL of the WecoAI API.
 
         Raises
         ------
@@ -25,29 +35,42 @@ class WecoAI:
             try:
                 api_key = os.environ["WECO_API_KEY"]
             except KeyError:
-                raise ValueError(
-                    "WECO_API_KEY must be passed to client or set as an environment variable"
-                )
+                raise ValueError("WECO_API_KEY must be passed to client or set as an environment variable")
         self.api_key = api_key
 
-        # base URL
         self.base_url = "https://function-builder.vercel.app"
 
-    def _headers(self) -> Dict[str, str]:
-        """Constructs the headers for the API requests.
+        # Setup clients
+        self.client = httpx.Client(http2=False, timeout=timeout)
+        self.async_client = httpx.AsyncClient(http2=False, timeout=timeout)
 
-        Returns
-        -------
-        dict
-            A dictionary containing the headers.
-        """
+    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
+
+    def _headers(self) -> Dict[str, str]:
+        """Constructs the headers for the API requests."""
         return {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
         }
 
-    def _post(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Makes a POST request to the specified endpoint with the provided data.
+    def _make_request(self, endpoint: str, data: Dict[str, Any], is_async: bool = False) -> Callable:
+        """Creates a callable for making either synchronous or asynchronous requests.
 
         Parameters
         ----------
@@ -55,42 +78,156 @@ class WecoAI:
             The API endpoint to which the request will be made.
         data : dict
             The data to be sent in the request body.
+        is_async : bool, optional
+            Whether to create an asynchronous request (default is False).
+
+        Returns
+        -------
+        Callable
+            A callable that performs the HTTP request.
+        """
+        url = f"{self.base_url}/{endpoint}"
+        headers = self._headers()
+
+        if is_async:
+
+            async def _request():
+                response = await self.async_client.post(url, json=data, headers=headers)
+                response.raise_for_status()
+                return response.json()
+
+            return _request()
+        else:
+
+            def _request():
+                response = self.client.post(url, json=data, headers=headers)
+                response.raise_for_status()
+                return response.json()
+
+            return _request()
+
+    def _process_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
+        """Processes the API response and handles warnings.
+
+        Parameters
+        ----------
+        response : dict
+            The raw API response.
 
         Returns
         -------
         dict
-            The response from the server as a dictionary.
+            A processed dictionary containing the output, token counts, and latency.
 
         Raises
         ------
-        requests.HTTPError
-            If the request fails.
+        UserWarning
+            If there are any warnings in the API response.
         """
-        url = f"{self.base_url}/{endpoint}"
-        response = requests.post(url, json=data, headers=self._headers())
-        response.raise_for_status()
-        return response.json()
+        for _warning in response.get("warnings", []):
+            warnings.warn(_warning)
 
-    def build(self, task_description: str) -> tuple[str, str]:
-        """Builds a specialized function given a task description.
+        return {
+            "output": response["response"],
+            "in_tokens": response["num_input_tokens"],
+            "out_tokens": response["num_output_tokens"],
+            "latency_ms": response["latency_ms"],
+        }
+
+    def _build(self, task_description: str, is_async: bool) -> Tuple[str, str] | Coroutine[Any, Any, Tuple[str, 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.
+        is_async : bool
+            Whether to perform an asynchronous request.
 
         Returns
         -------
-        tuple[str, str]
-            A tuple containing the name and description of the function.
+        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.
         """
         endpoint = "build"
         data = {"request": task_description}
-        response = self._post(endpoint, data)
-        return response["name"], response["description"]
+        request = self._make_request(endpoint=endpoint, data=data, is_async=is_async)
 
-    def query(self, fn_name: str, fn_input: str) -> Dict[str, Any]:
-        """Queries a function with the given function ID and input.
+        if is_async:
+
+            async def _async_build():
+                response = await request
+                return response["name"], response["description"]
+
+            return _async_build()
+        else:
+            response = request  # the request has already been made and the response is available
+            return response["name"], response["description"]
+
+    async def abuild(self, task_description: str) -> Tuple[str, 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.
+
+        Returns
+        -------
+        tuple[str, str]
+            A tuple containing the name and description of the function.
+        """
+        return await self._build(task_description=task_description, is_async=True)
+
+    def build(self, task_description: str) -> Tuple[str, 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.
+
+        Returns
+        -------
+        tuple[str, str]
+            A tuple containing the name and description of the function.
+        """
+        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.
+
+        Parameters
+        ----------
+        fn_name : str
+            The name of the function to query.
+        fn_input : str
+            The input to the function.
+        is_async : bool
+            Whether to perform an asynchronous request.
+
+        Returns
+        -------
+        dict | Coroutine[Any, Any, dict]
+            A dictionary containing the query results, or a coroutine that returns such a dictionary.
+        """
+        endpoint = "query"
+        data = {"name": fn_name, "user_message": fn_input}
+        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 _async_query()
+        else:
+            response = request  # the request has already been made and the response is available
+            return self._process_response(response=response)
+
+    async def aquery(self, fn_name: str, fn_input: str) -> Dict[str, Any]:
+        """Asynchronously queries a function with the given function ID and input.
 
         Parameters
         ----------
@@ -105,17 +242,56 @@ 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)
 
-        endpoint = "query"
-        data = {
-            "name": fn_name,
-            "user_message": fn_input,
-        }
-        response = self._post(endpoint, data)
-        result = {
-            "output": response["response"],
-            "in_tokens": response["num_input_tokens"],
-            "out_tokens": response["num_output_tokens"],
-            "latency_ms": response["latency_ms"],
-        }
-        return result
+    def query(self, fn_name: str, fn_input: 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.
+
+        Returns
+        -------
+        dict
+               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)
+
+    def batch_query(self, fn_names: str | List[str], batch_inputs: List[str]) -> List[Dict[str, Any]]:
+        """Synchronously queries multiple functions using asynchronous calls internally.
+
+        This method uses the asynchronous queries to submit all queries concurrently
+        and waits for all responses to be received before returning the results.
+
+        Parameters
+        ----------
+        fn_name : 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.
+
+        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.
+        """
+        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=fn_input) for fn_name, fn_input in zip(fn_names, batch_inputs)]
+            return await asyncio.gather(*tasks)
+
+        return asyncio.run(run_queries())

weco/functional.py

@@ -1,10 +1,10 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from .client import WecoAI
 
 
 def build(task_description: str, api_key: str = None) -> tuple[str, str]:
-    """Builds a specialized function given a task description.
+    """Builds a specialized function synchronously given a task description.
 
     Parameters
     ----------
@@ -23,8 +23,28 @@ def build(task_description: str, api_key: str = None) -> tuple[str, str]:
     return response
 
 
+async def abuild(task_description: str, api_key: str = None) -> tuple[str, 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.
+    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.
+    """
+    client = WecoAI(api_key=api_key)
+    response = await client.abuild(task_description=task_description)
+    return response
+
+
 def query(fn_name: str, fn_input: str, api_key: Optional[str] = None) -> Dict[str, Any]:
-    """Queries a function with the given function ID and input.
+    """Queries a function synchronously with the given function ID and input.
 
     Parameters
     ----------
@@ -44,3 +64,57 @@ def query(fn_name: str, fn_input: str, api_key: Optional[str] = None) -> Dict[st
     client = WecoAI(api_key=api_key)
     response = client.query(fn_name=fn_name, fn_input=fn_input)
     return response
+
+
+async def aquery(fn_name: str, fn_input: 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.
+    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
+    -------
+    dict
+        A dictionary containing the output of the function, the number of input tokens, the number of output tokens,
+        and the latency in milliseconds.
+    """
+    client = WecoAI(api_key=api_key)
+    response = await client.aquery(fn_name=fn_name, fn_input=fn_input)
+    return response
+
+
+def batch_query(fn_names: str | List[str], batch_inputs: List[str], 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
+    and waits for all responses to be received before returning the results.
+
+    Parameters
+    ----------
+    fn_name : 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.
+
+    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.
+
+    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.
+    """
+    client = WecoAI(api_key=api_key)
+    responses = client.batch_query(fn_names=fn_names, batch_inputs=batch_inputs)
+    return responses

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

@@ -1,26 +1,33 @@
 Metadata-Version: 2.1
 Name: weco
-Version: 0.1.1
+Version: 0.1.3
 Summary: A client facing API for interacting with the WeCo AI function builder service.
-Home-page: https://github.com/WecoAI/weco
-Author: ['WeCo AI Team']
-Author-email: dhruv@weco.ai
+Author-email: WeCo AI Team <dhruv@weco.ai>
 License: MIT
-Keywords: artificial intelligence,machine learning,data science,function builder,LLM
+Project-URL: Homepage, https://github.com/WecoAI/weco-python
+Keywords: AI,LLM,machine learning,data science,function builder
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: MIT License
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: requests
-
-[![Typing SVG](https://readme-typing-svg.demolab.com?font=Georgia&size=32&duration=4000&pause=400&color=FD4578&vCenter=true&multiline=false&width=750&height=100&lines=WeCo:+Function+Builder+Client;)](https://git.io/typing-svg)
-
-[![License](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
-![Python](https://img.shields.io/badge/Python-3.10.14-blue)
-
-
-# WeCo $f$(👷‍♂️)
+Requires-Dist: asyncio
+Requires-Dist: httpx[http2]
+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'
+
+<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>
+
+# $f$(👷‍♂️)
 
 A client facing API for interacting with the [WeCo AI](https://www.weco.ai/) function builder [service](https://weco-app.vercel.app/function)!
 
@@ -36,18 +43,20 @@ pip install weco
 
 ## Features
 
-- The **build** function enables quick and easy prototyping of new functions via LLMs through just natural language. We encourage users to do this through our [web console](https://weco-app.vercel.app/function) for maximum control and ease of use, however, you can also do this through our API as shown in [here](examples/).
+- The **build** function enables quick and easy prototyping of new functions via LLMs through just natural language. We encourage users to do this through our [web console](https://weco-app.vercel.app/function) for maximum control and ease of use, however, you can also do this through our API as shown in [here](examples/cookbook.ipynb).
 - 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 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. An example is shown [here](examples/example_client.py).
-- `weco.query` and `weco.build` to be used when you only require sparse usage. An example is provided [here](examples/example_functional.py).
+- `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.
+- `weco.query` and `weco.build` to be used when you only require sparse usage.
 
 ## Usage
 
 When using the WeCo API, you will need to set the API key:
 You can find/setup your API key [here](https://weco-app.vercel.app/account) by navigating to the API key tab. Once you have your API key, you may pass it to the `weco` client using the `api_key` argument input or set it as an environment variable such as:
-```
+```bash
 export WECO_API_KEY=<YOUR_WECO_API_KEY>
 ```
 
@@ -66,4 +75,6 @@ response = query(
 )
 ```
 
-## Enjoy $f$(👷‍♂️)!
+For more examples and an advanced user guide, check out our function builder [cookbook](examples/cookbook.ipynb).
+
+## Happy building $f$(👷‍♂️)!

weco-0.1.3.dist-info/RECORD

@@ -0,0 +1,8 @@
+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,,

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

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

weco-0.1.1.dist-info/RECORD

@@ -1,8 +0,0 @@
-weco/__init__.py,sha256=Xl80uFblDNbhIxL3_BQgYc_SninCeOeF0shnzPCrxrw,119
-weco/client.py,sha256=YOUgWJwLKlp9GR2u40pmnjaBtmDjLtdP_pUuvDAMCJQ,3744
-weco/functional.py,sha256=AnQ0PaZKr2K9a7wxXfVPJTTJLgrn1niZDQctdkj5vFA,1501
-weco-0.1.1.dist-info/LICENSE,sha256=NvpxfBuSajszAczWBGKxhHe4gsvil1H63zmu8xXZdL0,1064
-weco-0.1.1.dist-info/METADATA,sha256=ieRckEKjQ-shxHf5VPcNCHENd6CUpJz3i43-EeLSnMw,3214
-weco-0.1.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-weco-0.1.1.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.1.1.dist-info/RECORD,,