requirements.txt

@@ -1,5 +1,4 @@
 ollama
-httpx
 trilium-py
 gitpython
 PyGithub

src/ai_generators/tools/ollama_web_search_tool.py

@@ -4,25 +4,17 @@ Custom CrewAI tool that wraps Ollama's native web search API.
 This tool allows CrewAI agents to perform web searches using an Ollama
 subscription instead of third-party services like Serper or EXA.
 
-Uses direct HTTP requests via httpx with explicit Authorization: Bearer
-header to ensure the OLLAMA_API_KEY is properly passed to the Ollama cloud
-API endpoint (https://ollama.com/api/web_search).
-
 Requires:
-- httpx library: pip install httpx (already a transitive dependency of crewai)
+- Ollama Python library: pip install ollama
 - OLLAMA_API_KEY environment variable set with your Ollama API key
-
-Reference: https://docs.ollama.com/capabilities/web-search
 """
 
 import os
 
-import httpx
+import ollama
 from crewai.tools import BaseTool
 from pydantic import BaseModel, Field
 
-OLLAMA_WEB_SEARCH_URL = "https://ollama.com/api/web_search"
-
 
 class OllamaWebSearchInput(BaseModel):
     """Input schema for OllamaWebSearchTool."""
@@ -50,10 +42,6 @@ class OllamaWebSearchTool(BaseTool):
     The tool requires an Ollama subscription and the OLLAMA_API_KEY environment
     variable to be set.
 
-    Authentication is handled by sending the OLLAMA_API_KEY as a Bearer token
-    in the Authorization header, as documented at:
-    https://docs.ollama.com/capabilities/web-search
-
     Example usage:
         from ai_generators.tools.ollama_web_search_tool import OllamaWebSearchTool
 
@@ -77,9 +65,6 @@ class OllamaWebSearchTool(BaseTool):
         """
         Execute a web search and return formatted results.
 
-        Makes a POST request to https://ollama.com/api/web_search with the
-        OLLAMA_API_KEY as a Bearer token in the Authorization header.
-
         Args:
             query: The search query string
             max_results: Maximum number of results to return (1-10)
@@ -87,40 +72,16 @@ class OllamaWebSearchTool(BaseTool):
         Returns:
             Formatted string with search results, each containing title, URL, and content
         """
-        api_key = os.environ.get("OLLAMA_API_KEY")
-        if not api_key:
-            return (
-                "Error: OLLAMA_API_KEY environment variable is not set. "
-                "Please set your Ollama API key."
-            )
-
         try:
-            response = httpx.post(
-                OLLAMA_WEB_SEARCH_URL,
-                json={"query": query, "max_results": max_results},
-                headers={
-                    "Authorization": f"Bearer {api_key}",
-                    "Content-Type": "application/json",
-                },
-                timeout=30.0,
-            )
+            # Ensure API key is set
+            if not os.environ.get("OLLAMA_API_KEY"):
+                return "Error: OLLAMA_API_KEY environment variable is not set. Please set your Ollama API key."
 
-            # Raise for HTTP errors so we can catch them with specific messages
-            if response.status_code == 401:
-                return (
-                    "Authentication error: OLLAMA_API_KEY was rejected. "
-                    "Your key may be invalid or expired. Please verify it at "
-                    "https://ollama.com/settings/keys"
-                )
-            if response.status_code == 429:
-                return (
-                    "Rate limit exceeded: Too many search requests. "
-                    "Please wait a moment and try again."
-                )
-            response.raise_for_status()
+            # Perform the web search
+            response = ollama.web_search(query=query, max_results=max_results)
 
-            data = response.json()
-            results = data.get("results", [])
+            # Extract and format results
+            results = response.get("results", [])
 
             if not results:
                 return f"No search results found for query: '{query}'"
@@ -137,20 +98,27 @@ class OllamaWebSearchTool(BaseTool):
 
             return "\n".join(formatted_results)
 
-        except httpx.TimeoutException:
+        except Exception as exc:
+            return f"Error performing web search: {exc}"
+
+    def _handle_exception(self, exc: Exception) -> str:
+        """Handle exceptions gracefully and return a user-friendly error message."""
+        error_message = str(exc)
+
+        # Check for common error types
+        if "authentication" in error_message.lower() or "401" in error_message:
             return (
-                "Timeout error: The web search request timed out. "
-                "Please try again with a simpler query."
+                "Authentication error: Your OLLAMA_API_KEY may be invalid or expired. "
+                "Please check your API key and ensure it's set correctly in the environment."
             )
-        except httpx.ConnectError:
+        elif "rate limit" in error_message.lower() or "429" in error_message:
+            return "Rate limit exceeded: Too many search requests. Please wait a moment and try again."
+        elif (
+            "network" in error_message.lower() or "connection" in error_message.lower()
+        ):
             return (
                 "Network error: Unable to connect to Ollama's web search service. "
                 "Please check your internet connection and try again."
             )
-        except httpx.HTTPStatusError as exc:
-            return (
-                f"HTTP error {exc.response.status_code} from Ollama web search API: "
-                f"{exc.response.text}"
-            )
-        except Exception as exc:
-            return f"Error performing web search: {exc}"
+        else:
+            return f"Search failed: {error_message}"