Coverage for mindsdb / integrations / utilities / rag / loaders / vector_store_loader / MDBVectorStore.py: 45%

1from mindsdb_sql_parser.ast import Select, BinaryOperation, Identifier, Constant, Star 

2from mindsdb.integrations.libs.vectordatabase_handler import TableField 

3 

4from typing import Any, List, Optional 

5 

6from langchain_core.documents import Document 

7from langchain_core.embeddings import Embeddings 

8from langchain_core.vectorstores import VectorStore 

9 

10 

11class MDBVectorStore(VectorStore): 

12 

13 def __init__(self, kb_table) -> None: 

14 self.kb_table = kb_table 

15 

16 @property 

17 def embeddings(self) -> Optional[Embeddings]: 

18 return None 

19 

20 def similarity_search( 

21 self, 

22 query: str, 

23 k: int = 4, 

24 **kwargs: Any, 

25 ) -> List[Document]: 

26 

27 query = Select( 

28 targets=[Star()], 

29 where=BinaryOperation(op='=', args=[ 

30 Identifier(TableField.CONTENT.value), Constant(query) 

31 ]), 

32 limit=Constant(k), 

33 ) 

34 

35 df = self.kb_table.select_query(query) 

36 

37 docs = [] 

38 for _, row in df.iterrows(): 

39 metadata = row[TableField.METADATA.value] 

40 if metadata is None: 

41 metadata = {} 

42 docs.append(Document( 

43 page_content=row[TableField.CONTENT.value], 

44 metadata=metadata 

45 )) 

46 

47 return docs 

48 

49 def add_texts(self, *args, **kwargs) -> List[str]: 

50 raise NotImplementedError 

51 

52 @classmethod 

53 def from_texts(self, *args, **kwargs): 

54 raise NotImplementedError