Coverage for src/semware/services/search.py: 80%

1"""Search service for semantic search operations.""" 

2 

3import time 

4 

5from loguru import logger 

6 

7from ..models.requests import SearchResponse, SimilaritySearchRequest, TopKSearchRequest 

8from .embedding import embedding_service 

9from .vectordb import vectordb 

10 

11 

12class SearchService: 

13 """Service for performing semantic search operations.""" 

14 

15 def __init__(self): 

16 """Initialize the search service.""" 

17 pass 

18 

19 def similarity_search( 

20 self, table_name: str, request: SimilaritySearchRequest 

21 ) -> SearchResponse: 

22 """Perform similarity-based search. 

23 

24 Args: 

25 table_name: Name of the table to search 

26 request: Search request parameters 

27 

28 Returns: 

29 Search response with results 

30 """ 

31 start_time = time.time() 

32 

33 try: 

34 logger.info( 

35 f"Similarity search in table '{table_name}' with threshold {request.threshold}" 

36 ) 

37 

38 # Validate table exists 

39 vectordb.get_table_schema(table_name) 

40 

41 # Generate query embedding 

42 logger.debug("Generating query embedding") 

43 query_embedding = embedding_service.generate_query_embedding(request.query) 

44 

45 # Perform search 

46 results = vectordb.similarity_search( 

47 table_name=table_name, 

48 query_embedding=query_embedding, 

49 threshold=request.threshold, 

50 limit=request.limit, 

51 ) 

52 

53 # Sort results by similarity score (descending) 

54 results.sort(key=lambda x: x.similarity_score, reverse=True) 

55 

56 search_time_ms = (time.time() - start_time) * 1000 

57 

58 logger.info( 

59 f"Similarity search completed in {search_time_ms:.2f}ms, found {len(results)} results" 

60 ) 

61 

62 return SearchResponse( 

63 query=request.query, 

64 results=results, 

65 total_results=len(results), 

66 search_time_ms=search_time_ms, 

67 ) 

68 

69 except ValueError as e: 

70 logger.error(f"Similarity search validation error: {e}") 

71 raise 

72 except Exception as e: 

73 logger.exception(f"Similarity search failed: {e}") 

74 raise 

75 

76 def top_k_search( 

77 self, table_name: str, request: TopKSearchRequest 

78 ) -> SearchResponse: 

79 """Perform top-k search. 

80 

81 Args: 

82 table_name: Name of the table to search 

83 request: Search request parameters 

84 

85 Returns: 

86 Search response with results 

87 """ 

88 start_time = time.time() 

89 

90 try: 

91 logger.info(f"Top-k search in table '{table_name}' for k={request.k}") 

92 

93 # Validate table exists 

94 vectordb.get_table_schema(table_name) 

95 

96 # Generate query embedding 

97 logger.debug("Generating query embedding") 

98 query_embedding = embedding_service.generate_query_embedding(request.query) 

99 

100 # Perform search 

101 results = vectordb.top_k_search( 

102 table_name=table_name, query_embedding=query_embedding, k=request.k 

103 ) 

104 

105 # Results are already sorted by similarity (descending) from LanceDB 

106 

107 search_time_ms = (time.time() - start_time) * 1000 

108 

109 logger.info( 

110 f"Top-k search completed in {search_time_ms:.2f}ms, found {len(results)} results" 

111 ) 

112 

113 return SearchResponse( 

114 query=request.query, 

115 results=results, 

116 total_results=len(results), 

117 search_time_ms=search_time_ms, 

118 ) 

119 

120 except ValueError as e: 

121 logger.error(f"Top-k search validation error: {e}") 

122 raise 

123 except Exception as e: 

124 logger.exception(f"Top-k search failed: {e}") 

125 raise 

126 

127 

128# Global search service instance 

129search_service = SearchService()