Coverage for src/app/repositories/faiss_repository.py: 100%

1from langchain_community.vectorstores import FAISS 

2from langchain_core.documents import Document 

3 

4from dependencies.init_vector_store import store_vector_store 

5 

6from entities.document_context_entity import DocumentContextEntity 

7from entities.query_entity import QueryEntity 

8from entities.file_chunk_entity import FileChunkEntity 

9 

10class FaissRepository: 

11 

12 def __init__(self, vectorstore: FAISS): 

13 self.vectorstore = vectorstore 

14 

15 def similarity_search(self, query: QueryEntity) -> list[DocumentContextEntity]: 

16 """ 

17 Performs a similarity search in the collection and returns the most relevant documents. 

18  

19 Args: 

20 query (QueryEntity): The query entity containing the search parameters. 

21  

22 Returns: 

23 list[DocumentContextEntity]: A list of the most relevant document context entities. 

24 """ 

25 

26 

27 

28 if not query.get_query().strip(): 

29 raise ValueError("Query cannot be empty") 

30 

31 try: 

32 result = self.vectorstore.similarity_search(query.get_query(), k=4) 

33 

34 context_list = [] 

35 

36 for context in result: 

37 context_list.append(DocumentContextEntity(context.page_content)) 

38 

39 return context_list 

40 

41 except Exception as e: 

42 return str(e) 

43 

44 

45 def load_chunks(self, chunks: list[FileChunkEntity]) -> str: 

46 """ 

47 Perform a load of the chunks into the vectorstore. 

48 

49 Args: 

50 chunks (list[FileChunkEntity]): The list of file chunk entities to load. 

51 

52 Returns: 

53 str: The number of chunks loaded. 

54 

55 Raises: 

56 Exception: If an error occurs during the load. 

57 """ 

58 if not chunks: 

59 raise ValueError("No chunks to load.") 

60 

61 try: 

62 result = [] 

63 

64 for chunk in chunks: 

65 doc = Document(page_content=chunk.get_chunk_content(), metadata={"metadata": chunk.get_metadata()}) 

66 result.append(self.vectorstore.add_documents([doc])) 

67 

68 store_vector_store(self.vectorstore) 

69 

70 return f"{len(result)} chunks loaded." 

71 

72 except Exception as e: 

73 return f"Error occurred during chunk loading: {str(e)}" 

74 

75 

76 

77 

78 

79