Coverage for baibu/chat/search.py: 100%

1"""Web search providers for the ``internet_search`` tool. 

2 

3No provider is bundled. A deployment writes a subclass of 

4:class:`SearchProvider` for the search service it uses and names it in 

5``CHAT_SEARCH_PROVIDER``. :class:`MockSearchProvider` returns made-up results 

6for development and tests. 

7""" 

8 

9from dataclasses import asdict 

10from dataclasses import dataclass 

11from functools import cache 

12 

13from django.conf import settings 

14from django.utils.module_loading import import_string 

15 

16 

17class SearchError(Exception): 

18 """The search failed. The message is stored with the tool invocation.""" 

19 

20 

21@dataclass(frozen=True) 

22class SearchResult: 

23 title: str 

24 url: str 

25 snippet: str = "" 

26 

27 def as_dict(self) -> dict: 

28 return asdict(self) 

29 

30 

31class SearchProvider: 

32 name = "provider" 

33 

34 def search(self, query: str, *, max_results: int) -> list[SearchResult]: 

35 raise NotImplementedError 

36 

37 

38class MockSearchProvider(SearchProvider): 

39 """Deterministic results under example.org. A query containing 

40 ``[search-fail]`` fails; ``[search-empty]`` finds nothing.""" 

41 

42 name = "mock" 

43 

44 def search(self, query: str, *, max_results: int) -> list[SearchResult]: 

45 if "[search-fail]" in query: 

46 msg = "Mock search failure." 

47 raise SearchError(msg) 

48 if "[search-empty]" in query: 

49 return [] 

50 return [ 

51 SearchResult( 

52 title=f"Result {n} for {query}", 

53 url=f"https://example.org/search/{n}", 

54 snippet=f"Made-up text about {query}, number {n}.", 

55 ) 

56 for n in range(1, max_results + 1) 

57 ] 

58 

59 

60@cache 

61def _provider_class(path: str) -> type[SearchProvider]: 

62 return import_string(path) 

63 

64 

65def get_provider() -> SearchProvider | None: 

66 path = settings.CHAT_SEARCH_PROVIDER 

67 return _provider_class(path)() if path else None