agent: add brave_search & fetch_page tools + move to examples/agent/tools/

examples/agent/README.md

@@ -48,8 +48,9 @@
 
   ```bash
   docker run -p 8088:8088 -w /src -v $PWD/examples/agent:/src \
+    --env BRAVE_SEARCH_API_KEY=$BRAVE_SEARCH_API_KEY \
     --rm -it ghcr.io/astral-sh/uv:python3.12-alpine \
-    uv run fastify.py --port 8088 tools.py
+    uv run fastify.py --port 8088 tools
   ```
 
   > [!WARNING]
@@ -58,9 +59,14 @@
 - Run the agent with a given goal:
 
   ```bash
-  uv run examples/agent/run.py \
-    --tool-endpoint http://localhost:8088 \
-    --goal "What is the sum of 2535 squared and 32222000403?"
+  uv run examples/agent/run.py --tools http://localhost:8088 \
+    "What is the sum of 2535 squared and 32222000403?"
+
+  uv run examples/agent/run.py --tools http://localhost:8088 \
+    "What is the best BBQ join in Laguna Beach?"
+
+  uv run examples/agent/run.py --tools http://localhost:8088 \
+    "Search for, fetch and summarize the homepage of llama.cpp"
   ```
 
 ## TODO

examples/agent/fastify.py

@@ -1,14 +1,17 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
+#     "aiohttp",
 #     "fastapi",
-#     "uvicorn",
-#     "typer",
+#     "html2text",
 #     "ipython",
+#     "pyppeteer",
+#     "typer",
+#     "uvicorn",
 # ]
 # ///
 '''
-    Binds the functions of a python script as a FastAPI server.
+    Discovers and binds python script functions as a FastAPI server.
 '''
 import os
 import sys
@@ -45,7 +48,7 @@ def _load_module(f: str):
 def main(files: List[str], host: str = '0.0.0.0', port: int = 8000):
     app = fastapi.FastAPI()
 
-    for f in files:
+    def load_python(f):
         print(f'Binding functions from {f}')
         module = _load_module(f)
         for k in dir(module):
@@ -69,7 +72,15 @@ def main(files: List[str], host: str = '0.0.0.0', port: int = 8000):
             except Exception as e:
                 print(f'WARNING:    Failed to bind /{k}\n\t{e}')
 
-    print(f'INFO:     CWD = {os.getcwd()}')
+    for f in files:
+        if os.path.isdir(f):
+            for root, _, files in os.walk(f):
+                for file in files:
+                    if file.endswith('.py'):
+                        load_python(os.path.join(root, file))
+        else:
+            load_python(f)
+
     uvicorn.run(app, host=host, port=port)
 
 

examples/agent/run.py

@@ -136,16 +136,16 @@ def typer_async_workaround():
 
 @typer_async_workaround()
 async def main(
-    goal: Annotated[str, typer.Option()],
+    goal: str,
     api_key: str = '<unset>',
-    tool_endpoint: Optional[list[str]] = None,
+    tools: Optional[list[str]] = None,
     max_iterations: Optional[int] = 10,
     verbose: bool = False,
     endpoint: str = "http://localhost:8080/v1/",
 ):
     client = AsyncOpenAI(api_key=api_key, base_url=endpoint)
 
-    tool_map, tools = await discover_tools(tool_endpoint or [], verbose)
+    tool_map, tools = await discover_tools(tools or [], verbose)
 
     sys.stdout.write(f'🛠️  {", ".join(tool_map.keys())}\n')
 

examples/agent/tools/fetch.py

@@ -0,0 +1,58 @@
+import aiohttp
+import sys
+from typing import Optional
+
+from pydantic import BaseModel
+import html2text
+
+
+class FetchResult(BaseModel):
+    content: Optional[str] = None
+    markdown: Optional[str] = None
+    error: Optional[str] = None
+
+
+async def fetch_page(url: str) -> FetchResult:
+    '''
+        Fetch a web page (convert it to markdown if possible).
+    '''
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as res:
+                res.raise_for_status()
+                content = await res.text()
+    except aiohttp.ClientError as e:
+        return FetchResult(error=str(e))
+
+    # NOTE: Pyppeteer doesn't work great in docker, short of installing a bunch of dependencies
+    # from pyppeteer import launch
+    # from pyppeteer.errors import TimeoutError, NetworkError
+    # browser = await launch()
+    # try:
+    #     page = await browser.newPage()
+    #     response = await page.goto(url)
+
+    #     if not response.ok:
+    #         return FetchResult(error=f"HTTP {response.status} {response.statusText}")
+
+    #     content=await page.content()
+    # except TimeoutError:
+    #     return FetchResult(error="Page load timed out")
+    # except NetworkError:
+    #     return FetchResult(error="Network error occurred")
+    # except Exception as e:
+    #     return FetchResult(error=str(e))
+    # finally:
+    #     await browser.close()
+
+    try:
+        h = html2text.HTML2Text()
+        h.ignore_links = False
+        h.ignore_images = False
+        h.ignore_emphasis = False
+        markdown = h.handle(content)
+        return FetchResult(markdown=markdown)
+    except Exception as e:
+        print(f'Failed to convert HTML of {url} to markdown: {e}', file=sys.stderr)
+        return FetchResult(content=content)

examples/agent/tools/python.py

@@ -0,0 +1,28 @@
+from IPython.core.interactiveshell import InteractiveShell
+from io import StringIO
+import sys
+
+
+def python(code: str) -> str:
+    """
+    Execute Python code in a siloed environment using IPython and returns the output.
+
+    Parameters:
+        code (str): The Python code to execute.
+
+    Returns:
+        str: The output of the executed code.
+    """
+    shell = InteractiveShell()
+
+    old_stdout = sys.stdout
+    sys.stdout = out = StringIO()
+
+    try:
+        shell.run_cell(code)
+    except Exception as e:
+        return f"An error occurred: {e}"
+    finally:
+        sys.stdout = old_stdout
+
+    return out.getvalue()

examples/agent/tools/search.py

@@ -0,0 +1,72 @@
+import aiohttp
+import itertools
+import json
+import os
+import sys
+from typing import Dict, List
+import urllib.parse
+
+
+def _extract_values(keys, obj):
+    values = {}
+    for k in keys:
+        v = obj.get(k)
+        if v is not None:
+            values[k] = v
+    return values
+
+
+# Let's keep this tool aligned w/ llama_stack.providers.impls.meta_reference.agents.tools.builtin.BraveSearch
+# (see https://github.com/meta-llama/llama-stack/blob/main/llama_stack/providers/impls/meta_reference/agents/tools/builtin.py)
+_result_keys_by_type = {
+    "web":       ("type", "title", "url", "description", "date", "extra_snippets"),
+    "videos":    ("type", "title", "url", "description", "date"),
+    "news":      ("type", "title", "url", "description"),
+    "infobox":   ("type", "title", "url", "description", "long_desc"),
+    "locations": ("type", "title", "url", "description", "coordinates", "postal_address", "contact", "rating", "distance", "zoom_level"),
+    "faq":       ("type", "title", "url", "question", "answer"),
+}
+
+
+async def brave_search(query: str, max_results: int = 10) -> List[Dict]:
+    """
+    Search the Brave Search API for the specified query.
+
+    Parameters:
+        query (str): The query to search for.
+        max_results (int): The maximum number of results to return (defaults to 10)
+
+    Returns:
+        List[Dict]: The search results.
+    """
+
+    url = f"https://api.search.brave.com/res/v1/web/search?q={urllib.parse.quote(query)}"
+    headers = {
+        'Accept': 'application/json',
+        'Accept-Encoding': 'gzip',
+        'X-Subscription-Token': os.environ['BRAVE_SEARCH_API_KEY'],
+    }
+
+    def extract_results(search_response):
+        for m in search_response['mixed']['main']:
+            result_type = m['type']
+            keys = _result_keys_by_type.get(result_type)
+            if keys is None:
+                print(f'[brave_search] Unknown result type: {result_type}', file=sys.stderr)
+                continue
+
+            results_of_type = search_response[result_type]["results"]
+            if (idx := m.get("index")) is not None:
+                yield _extract_values(keys, results_of_type[idx])
+            elif m["all"]:
+                for r in results_of_type:
+                    yield _extract_values(keys, r)
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(url, headers=headers) as res:
+            res.raise_for_status()
+            response = await res.json()
+
+            results = list(itertools.islice(extract_results(response), max_results))
+            print(json.dumps(dict(query=query, response=response, results=results), indent=2))
+            return results

examples/agent/tools/wait.py

@@ -1,16 +1,9 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#     "ipython",
-# ]
-# ///
+import asyncio
 import datetime
 from pydantic import BaseModel
 import sys
-import time
 from typing import Optional
 
-
 class Duration(BaseModel):
     seconds: Optional[int] = None
     minutes: Optional[int] = None
@@ -34,7 +27,7 @@ class Duration(BaseModel):
         ])
 
     @property
-    def get_total_seconds(self) -> int:
+    def get_total_seconds(self) -> float:
         return sum([
             self.seconds or 0,
             (self.minutes or 0)*60,
@@ -44,23 +37,18 @@ class Duration(BaseModel):
             (self.years or 0)*31536000,
         ])
 
-
 class WaitForDuration(BaseModel):
     duration: Duration
 
-    def __call__(self):
+    async def __call__(self):
         sys.stderr.write(f"Waiting for {self.duration}...\n")
-        time.sleep(self.duration.get_total_seconds)
+        await asyncio.sleep(self.duration.get_total_seconds)
 
-
-def wait_for_duration(duration: Duration) -> None:
+async def wait_for_duration(duration: Duration) -> None:
     'Wait for a certain amount of time before continuing.'
+    await asyncio.sleep(duration.get_total_seconds)
 
-    # sys.stderr.write(f"Waiting for {duration}...\n")
-    time.sleep(duration.get_total_seconds)
-
-
-def wait_for_date(target_date: datetime.date) -> None:
+async def wait_for_date(target_date: datetime.date) -> None:
     f'''
         Wait until a specific date is reached before continuing.
         Today's date is {datetime.date.today()}
@@ -75,34 +63,4 @@ def wait_for_date(target_date: datetime.date) -> None:
 
     days, seconds = time_diff.days, time_diff.seconds
 
-    # sys.stderr.write(f"Waiting for {days} days and {seconds} seconds until {target_date}...\n")
-    time.sleep(days * 86400 + seconds)
-
-
-def python(code: str) -> str:
-    """
-    Executes Python code in a siloed environment using IPython and returns the output.
-
-    Parameters:
-        code (str): The Python code to execute.
-
-    Returns:
-        str: The output of the executed code.
-    """
-    from IPython.core.interactiveshell import InteractiveShell
-    from io import StringIO
-    import sys
-
-    shell = InteractiveShell()
-
-    old_stdout = sys.stdout
-    sys.stdout = out = StringIO()
-
-    try:
-        shell.run_cell(code)
-    except Exception as e:
-        return f"An error occurred: {e}"
-    finally:
-        sys.stdout = old_stdout
-
-    return out.getvalue()
+    await asyncio.sleep(days * 86400 + seconds)