Merge branch 'main' into examples

i-am-bee · Mar 3, 2025 · bb4cc5d · bb4cc5d
2 parents a65a518 + a6fe493
commit bb4cc5d
Show file tree

Hide file tree

Showing 7 changed files with 47 additions and 35 deletions.
diff --git a/python/beeai_framework/tools/search/duckduckgo.py b/python/beeai_framework/tools/search/duckduckgo.py
@@ -59,7 +59,7 @@ def __init__(self, max_results: int = 10, safe_search: str = DuckDuckGoSearchTyp
             creator=self,
         )
 
-    def _run(self, input: DuckDuckGoSearchToolInput, _: Any | None = None) -> DuckDuckGoSearchToolOutput:
+    async def _run(self, input: DuckDuckGoSearchToolInput, _: Any | None = None) -> DuckDuckGoSearchToolOutput:
         try:
             results = DDGS().text(input.query, max_results=self.max_results, safesearch=self.safe_search)
             search_results: list[SearchToolResult] = [

diff --git a/python/beeai_framework/tools/search/wikipedia.py b/python/beeai_framework/tools/search/wikipedia.py
@@ -61,7 +61,7 @@ def get_section_titles(self, sections: wikipediaapi.WikipediaPage.sections) -> s
             titles.append(section.title)
         return ",".join(str(title) for title in titles)
 
-    def _run(self, input: WikipediaToolInput, _: Any | None = None) -> WikipediaToolOutput:
+    async def _run(self, input: WikipediaToolInput, _: Any | None = None) -> WikipediaToolOutput:
         page_py = self.client.page(input.query)
 
         if not page_py.exists():

diff --git a/python/beeai_framework/tools/tool.py b/python/beeai_framework/tools/tool.py
@@ -84,7 +84,7 @@ def input_schema(self) -> type[T]:
         pass
 
     @abstractmethod
-    def _run(self, input: Any, options: dict[str, Any] | None = None) -> Any:
+    async def _run(self, input: Any, options: dict[str, Any] | None = None) -> Any:
         pass
 
     def validate_input(self, input: T | dict[str, Any]) -> T:
@@ -113,7 +113,7 @@ async def executor(_: RetryableContext) -> Any:
                     nonlocal error_propagated
                     error_propagated = False
                     await context.emitter.emit("start", meta)
-                    return self._run(validated_input, options)
+                    return await self._run(validated_input, options)
 
                 async def on_error(error: Exception, _: RetryableContext) -> None:
                     nonlocal error_propagated
@@ -202,9 +202,12 @@ def __init__(self, options: dict[str, Any] | None = None) -> None:
                 creator=self,
             )
 
-        def _run(self, tool_in: Any, _: dict[str, Any] | None = None) -> None:
+        async def _run(self, tool_in: Any, _: dict[str, Any] | None = None) -> None:
             tool_input_dict = tool_in.model_dump()
-            return tool_function(**tool_input_dict)
+            if inspect.iscoroutinefunction(tool_function):
+                return await tool_function(**tool_input_dict)
+            else:
+                return tool_function(**tool_input_dict)
 
     f_tool = FunctionTool()
     return f_tool
diff --git a/python/beeai_framework/tools/weather/openmeteo.py b/python/beeai_framework/tools/weather/openmeteo.py
@@ -19,6 +19,7 @@
 from typing import Any, Literal
 from urllib.parse import urlencode
 
+import httpx
 import requests
 from pydantic import BaseModel, Field
 
@@ -134,12 +135,14 @@ def _trim_date(date_str: str) -> str:
         params["temperature_unit"] = input.temperature_unit
         return params
 
-    def _run(self, input: OpenMeteoToolInput, options: Any = None) -> StringToolOutput:
+    async def _run(self, input: OpenMeteoToolInput, options: Any = None) -> StringToolOutput:
         params = urlencode(self.get_params(input), doseq=True)
         logger.debug(f"Using OpenMeteo URL: https://api.open-meteo.com/v1/forecast?{params}")
-        response = requests.get(
-            f"https://api.open-meteo.com/v1/forecast?{params}",
-            headers={"Content-Type": "application/json", "Accept": "application/json"},
-        )
-        response.raise_for_status()
-        return StringToolOutput(json.dumps(response.json()))
+
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"https://api.open-meteo.com/v1/forecast?{params}",
+                headers={"Content-Type": "application/json", "Accept": "application/json"},
+            )
+            response.raise_for_status()
+            return StringToolOutput(json.dumps(response.json()))
diff --git a/python/docs/tools.md b/python/docs/tools.md
@@ -378,7 +378,7 @@ class RiddleTool(Tool[RiddleToolInput]):
             creator=self,
         )
 
-    def _run(self, input: RiddleToolInput, _: Any | None = None) -> None:
+    async def _run(self, input: RiddleToolInput, _: Any | None = None) -> None:
         index = input.riddle_number % (len(self.data))
         riddle = self.data[index]
         return riddle
@@ -419,7 +419,7 @@ import asyncio
 import sys
 from typing import Any
 
-import requests
+import httpx
 from pydantic import BaseModel, Field
 
 from beeai_framework.emitter.emitter import Emitter
@@ -453,7 +453,7 @@ class OpenLibraryTool(Tool[OpenLibraryToolInput]):
             creator=self,
         )
 
-    def _run(self, tool_input: OpenLibraryToolInput, _: Any | None = None) -> OpenLibraryToolResult:
+    async def _run(self, tool_input: OpenLibraryToolInput, _: Any | None = None) -> OpenLibraryToolResult:
         key = ""
         value = ""
         input_vars = vars(tool_input)
@@ -465,17 +465,20 @@ class OpenLibraryTool(Tool[OpenLibraryToolInput]):
         else:
             raise ToolInputValidationError("All input values in OpenLibraryToolInput were empty.") from None
 
-        response = requests.get(
-            f"https://openlibrary.org/api/books?bibkeys={key}:{value}&jsmcd=data&format=json",
-            headers={"Content-Type": "application/json", "Accept": "application/json"},
-        )
-
-        response.raise_for_status()
+        json_output = {}
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"https://openlibrary.org/api/books?bibkeys={key}:{value}&jsmcd=data&format=json",
+                headers={"Content-Type": "application/json", "Accept": "application/json"},
+            )
+            response.raise_for_status()
 
-        json_output = response.json()[f"{key}:{value}"]
+            json_output = response.json()[f"{key}:{value}"]
 
         return OpenLibraryToolResult(
-            preview_url=json_output["preview_url"], info_url=json_output["info_url"], bib_key=json_output["bib_key"]
+            preview_url=json_output.get("preview_url"),
+            info_url=json_output.get("info_url"),
+            bib_key=json_output.get("bib_key"),
         )
 
 

diff --git a/python/examples/tools/custom/base.py b/python/examples/tools/custom/base.py
@@ -36,7 +36,7 @@ def __init__(self, options: dict[str, Any] | None = None) -> None:
             creator=self,
         )
 
-    def _run(self, input: RiddleToolInput, _: Any | None = None) -> None:
+    async def _run(self, input: RiddleToolInput, _: Any | None = None) -> None:
         index = input.riddle_number % (len(self.data))
         riddle = self.data[index]
         return riddle

diff --git a/python/examples/tools/custom/openlibrary.py b/python/examples/tools/custom/openlibrary.py
@@ -2,7 +2,7 @@
 import sys
 from typing import Any
 
-import requests
+import httpx
 from pydantic import BaseModel, Field
 
 from beeai_framework.emitter.emitter import Emitter
@@ -36,7 +36,7 @@ def __init__(self, options: dict[str, Any] | None = None) -> None:
             creator=self,
         )
 
-    def _run(self, tool_input: OpenLibraryToolInput, _: Any | None = None) -> OpenLibraryToolResult:
+    async def _run(self, tool_input: OpenLibraryToolInput, _: Any | None = None) -> OpenLibraryToolResult:
         key = ""
         value = ""
         input_vars = vars(tool_input)
@@ -48,17 +48,20 @@ def _run(self, tool_input: OpenLibraryToolInput, _: Any | None = None) -> OpenLi
         else:
             raise ToolInputValidationError("All input values in OpenLibraryToolInput were empty.") from None
 
-        response = requests.get(
-            f"https://openlibrary.org/api/books?bibkeys={key}:{value}&jsmcd=data&format=json",
-            headers={"Content-Type": "application/json", "Accept": "application/json"},
-        )
-
-        response.raise_for_status()
+        json_output = {}
+        async with httpx.AsyncClient() as client:
+            response = await client.get(
+                f"https://openlibrary.org/api/books?bibkeys={key}:{value}&jsmcd=data&format=json",
+                headers={"Content-Type": "application/json", "Accept": "application/json"},
+            )
+            response.raise_for_status()
 
-        json_output = response.json()[f"{key}:{value}"]
+            json_output = response.json()[f"{key}:{value}"]
 
         return OpenLibraryToolResult(
-            preview_url=json_output["preview_url"], info_url=json_output["info_url"], bib_key=json_output["bib_key"]
+            preview_url=json_output.get("preview_url"),
+            info_url=json_output.get("info_url"),
+            bib_key=json_output.get("bib_key"),
         )