Fix bugs

Author: DhruvSrikanth

aifn/client.py

@@ -1,4 +1,5 @@
 import asyncio
+import nest_asyncio
 import base64
 import os
 import warnings
@@ -613,15 +614,15 @@ def _validate_batch_query(self, batch_inputs: List[Dict[str, Any]]):
             if not all(isinstance(img, str) for inp in batch_inputs for img in inp.get("images_input", [])):
                 raise ValueError("Images input must be a list of strings.")
 
-    def _batch_query(
+    async def _batch_query(
         self,
         fn_name: str,
         version: Optional[Union[str, int]],
         batch_inputs: List[Dict[str, Any]],
         return_reasoning: bool,
         strict: bool,
     ) -> List[NamedTuple]:
-        """Internal method to handle both synchronous and asynchronous batch query requests.
+        """Internal method to handle asynchronous batch query requests.
 
         Parameters
         ----------
@@ -644,21 +645,18 @@ def _batch_query(
 
         Returns
         -------
-        Union[List[NamedTuple], Coroutine[Any, Any, List[NamedTuple]]]
-            A list of NamedTuples, each containing the output and metadata of the response, or a coroutine that returns such a list.
+        List[NamedTuple]
+            A list of NamedTuples, each containing the output and metadata of the response.
         """
 
         # Validate the batch inputs
         self._validate_batch_query(batch_inputs=batch_inputs)
 
-        async def run_queries():
-            tasks = [
-                self.aquery(fn_name=fn_name, version=version, return_reasoning=return_reasoning, strict=strict, **fn_input)
-                for fn_input in batch_inputs
-            ]
-            return await asyncio.gather(*tasks)
-
-        return run_queries()
+        tasks = [
+            self.aquery(fn_name=fn_name, version=version, return_reasoning=return_reasoning, strict=strict, **fn_input)
+            for fn_input in batch_inputs
+        ]
+        return await asyncio.gather(*tasks)
 
     async def abatch_query(
         self,
@@ -732,8 +730,17 @@ def batch_query(
         List[NamedTuple]
             A list of NamedTuples, each containing the output and metadata of the response.
         """
-        return asyncio.run(
-            self._batch_query(
+        try:
+            # Check if an event loop is already running
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                nest_asyncio.apply()
+            task = loop.create_task(self._batch_query(
                 fn_name=fn_name, version=version, batch_inputs=batch_inputs, return_reasoning=return_reasoning, strict=strict
-            )
-        )
+            ))
+            return loop.run_until_complete(task)
+        except RuntimeError:
+            # If no event loop is running, use asyncio.run
+            return asyncio.run(self._batch_query(
+                fn_name=fn_name, version=version, batch_inputs=batch_inputs, return_reasoning=return_reasoning, strict=strict
+            ))

examples/cookbook.ipynb

@@ -38,11 +38,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "%env WECO_API_KEY=<YOUR_WECO_API_KEY>"
+    "import os\n",
+    "os.environ[\"WECO_API_KEY\"] = \"YOUR_WECO_API_KEY\""
    ]
   },
   {
@@ -214,13 +215,11 @@
    "source": [
     "task_evaluator = build(task_description=\"I want to know if AI can solve a problem for me, how easy it is to arrive at a solution and whether any helpful tips for me along the way. Help me understand this through - 'feasibility', 'justification', and 'suggestions'.\")\n",
     "\n",
-    "\n",
     "task1 = {\n",
     "    \"text_input\": \"I want to train a model to predict house prices using the Boston Housing dataset hosted on Kaggle.\"\n",
     "}\n",
     "task2 = {\n",
     "    \"text_input\": \"I want to train a model to classify digits using the MNIST dataset hosted on Kaggle using a Google Colab notebook. Attached is an example of what some of the digits would look like.\",\n",
-    "    \"images_input\": [\"https://machinelearningmastery.com/wp-content/uploads/2019/02/Plot-of-a-Subset-of-Images-from-the-MNIST-Dataset-1024x768.png\"]\n",
     "}\n",
     "responses = task_evaluator.batch([task1, task2])\n",
     "for response in responses:\n",
@@ -352,10 +351,16 @@
     "task_evaluator = build(task_description=\"I want to know if AI can solve a problem for me, how easy it is to arrive at a solution and whether any helpful tips for me along the way. Help me understand this through - 'feasibility', 'justification', and 'suggestions'.\")\n",
     "\n",
     "output, metadata = task_evaluator(\"I want to train a model to predict house prices using the Boston Housing dataset hosted on Kaggle.\", return_reasoning=True)\n",
-    "reasoning_steps = response[\"reasoning_steps\"]\n",
     "for key, value in output.items(): print(f\"{key}: {value}\")\n",
     "for i, step in enumerate(metadata[\"reasoning_steps\"]): print(f\"Step {i+1}: {step}\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

pyproject.toml

@@ -13,7 +13,7 @@ readme = "README.md"
 version = "0.2.0"
 license = {text = "MIT"}
 requires-python = ">=3.8"
-dependencies = ["requests", "asyncio", "httpx[http2]", "pillow"]
+dependencies = ["requests", "asyncio", "nest_asyncio", "httpx[http2]", "pillow"]
 keywords = ["AI", "LLM", "VLM", "AI functions", "Prompt Engineering", "NLP"]
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -33,6 +33,10 @@ packages = ["aifn"]
 
 [tool.setuptools_scm]
 
+# Test configuration
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope = "function"
+
 # Lint and format code
 [tool.ruff]
 fix = false