ollama.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  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 = {
  26. "query": query,
  27. }
  28. try:
  29. response = requests.post(f"{url}/api/web_search", headers=headers, json=payload)
  30. response.raise_for_status()
  31. data = response.json()
  32. results = data.get("results", [])
  33. log.info(f"Found {len(results)} results")
  34. return [
  35. SearchResult(
  36. link=result.get("url", ""),
  37. title=result.get("title", ""),
  38. snippet=result.get("content", ""),
  39. )
  40. for result in results
  41. ]
  42. except Exception as e:
  43. log.error(f"Error searching Ollama: {e}")
  44. return []