1
0

ollama.py 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. import logging
  2. from dataclasses import dataclass
  3. from typing import Optional
  4. import requests
  5. from open_webui.env import SRC_LOG_LEVELS
  6. from open_webui.retrieval.web.main import SearchResult
  7. log = logging.getLogger(__name__)
  8. log.setLevel(SRC_LOG_LEVELS["RAG"])
  9. def search_ollama_cloud(
  10. url: str,
  11. api_key: str,
  12. query: str,
  13. count: int,
  14. filter_list: Optional[list[str]] = None,
  15. ) -> list[SearchResult]:
  16. """Search using Ollama Search API and return the results as a list of SearchResult objects.
  17. Args:
  18. api_key (str): A Ollama Search API key
  19. query (str): The query to search for
  20. count (int): Number of results to return
  21. filter_list (Optional[list[str]]): List of domains to filter results by
  22. """
  23. log.info(f"Searching with Ollama for query: {query}")
  24. headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
  25. payload = {"query": query, "max_results": count}
  26. try:
  27. response = requests.post(f"{url}/api/web_search", headers=headers, json=payload)
  28. response.raise_for_status()
  29. data = response.json()
  30. results = data.get("results", [])
  31. log.info(f"Found {len(results)} results")
  32. return [
  33. SearchResult(
  34. link=result.get("url", ""),
  35. title=result.get("title", ""),
  36. snippet=result.get("content", ""),
  37. )
  38. for result in results
  39. ]
  40. except Exception as e:
  41. log.error(f"Error searching Ollama: {e}")
  42. return []