Coverage for /var/devmt/py/docp-dbi_1.1.0/docp_dbi/databases/chroma.py: 100%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4:Purpose: This module provides a localised wrapper and specialised 

5 functionality around the :class:`langchain_chroma.Chroma` 

6 class, for interacting with a ChromaDB database. 

7 

8:Platform: Linux/Windows | Python 3.11+ 

9:Developer: J Berendt 

10:Email: development@s3dev.uk 

11 

12:Comments: This module uses the :class:`langchain_chroma.Chroma` 

13 class, rather than the base ``chromadb`` library as 

14 langchain's implementation provides the ``add_texts`` method 

15 which supports **GPU processing** and parallelisation. This 

16 functionality is implemented and accessed through the 

17 :meth:`~ChromaDB.add_documents` method. 

18 

19""" 

20# pylint: disable=wrong-import-order 

21 

22from __future__ import annotations 

23 

24import chromadb 

25import os 

26from hashlib import md5 

27# langchain's Chroma is used rather than the base chromadb as it provides 

28# the add_texts method which support GPU processing and parallelisation. 

29from langchain_chroma import Chroma as _LangChroma 

30# locals 

31try: 

32 from ._embeddingfunctions import CustomEmbeddingFunction 

33except ImportError: 

34 from docp_dbi.databases._embeddingfunctions import CustomEmbeddingFunction 

35 

36 

37class ChromaDB(_LangChroma): 

38 r"""Wrapper class around the ``chromadb`` library. 

39 

40 Args: 

41 path (str): Path to the chromadb database's *directory*. 

42 collection (str): Collection name. 

43 embedding_model_path (str, optional): Path to a *local* embedding 

44 model *repo* of your choosing. Defaults to None. 

45 If not provided, the ``SentenceTransformers`` class be be 

46 allowed free reign to download or use the model named in the 

47 :attr:`DEFAULT_EMBEDDING_MODEL` attribute. 

48 repo_id (str, optional): HuggingFace repository ID for the 

49 requested embedding model. Defaults to None. 

50 offline (bool, optional): Remain offline. Use the local model 

51 embedding function model repo rather than obtaining one 

52 online. Defaults to False. 

53 

54 :Examples: 

55 

56 Load a new PDF document into ChromaDB:: 

57 

58 >>> from docp_dbi import ChromaDB 

59 >>> from docp_parsers import PDFParser 

60 >>> from langchain_text_splitters import RecursiveCharacterTextSplitter 

61 

62 # Parse the PDF document. 

63 >>> pdf = PDFParser(path='/path/to/documents/rag-pipelines-how-to.pdf') 

64 >>> pdf.extract_text() 

65 

66 # Setup a text splitter (for chunking the document). 

67 >>> splitter = RecursiveCharacterTextSplitter( 

68 ... separators=['\n\n\n', '\n\n', '\n', '.'], 

69 ... chunk_size=512, 

70 ... chunk_overlap=128 

71 ... ) 

72 # Split the document for storage. 

73 >>> docs = splitter.split_documents(pdf.doc.documents) 

74 

75 # Create a database interface, using an offline, local user-defined embedding model. 

76 >>> db = ChromaDB(path='/path/to/databases/chroma/', 

77 ... collection='test', 

78 ... embedding_model_path=('/path/to/models/sentence-transformers/' 

79 'all-MiniLM-L6-v2'), 

80 ... offline=True) 

81 # Embed and store the document chunks. 

82 >>> db.add_documents(documents=docs) 

83 

84 # Run your first query. 

85 >>> result = db.collection.query(query_texts=['How do I implement a RAG pipeline?']) 

86 

87 """ 

88 

89 # TODO: Move this to a config file in docp-core. 

90 DEFAULT_EMBEDDING_MODEL = 'sentence-transformers/all-MiniLM-L6-v2' 

91 

92 def __init__(self, 

93 path: str, 

94 collection: str, 

95 *, 

96 embedding_model_path: str=None, 

97 repo_id: str=None, 

98 offline: bool=False): 

99 """Chroma database class initialiser.""" 

100 self._path = os.path.realpath(path) 

101 self._cname = collection 

102 self._embpath = embedding_model_path 

103 self._repo_id = repo_id 

104 self._offline = offline 

105 self._oclient = None # Database 'client' object 

106 self._ocollection = None # Database 'collection' object. 

107 self._set_client() 

108 self._set_embedding_fn() 

109 self._set_collection() 

110 super().__init__(client=self._oclient, 

111 collection_name=self._cname, 

112 embedding_function=self._embfn, 

113 persist_directory=self._path) 

114 

115 def __repr__(self) -> str: 

116 """Define the information about this class to be displayed.""" 

117 metric = self.collection.__dict__['_model']['configuration_json']['hnsw']['space'] 

118 return (f'<class: {self.__class__.__name__}>\n' 

119 f'Collection: {self.collection.name}\n' 

120 f'Similarity metric: {metric}\n' 

121 f'Embedding function: {self.embedding_function.model}\n' 

122 f'Embedding model path: {self.embedding_function.model_path}') 

123 

124 @property 

125 def client(self) -> chromadb.api.client.Client: # nocover 

126 """Accessor to the :class:`chromadb.PersistentClient` class.""" 

127 return self._oclient 

128 

129 @property 

130 def collection(self) -> chromadb.api.models.Collection.Collection: # nocover 

131 """Accessor to the chromadb client's collection object.""" 

132 return self._ocollection 

133 

134 @property 

135 def embedding_function(self) -> CustomEmbeddingFunction: 

136 """Accessor to the embedding function used.""" 

137 return self._embfn 

138 

139 @property 

140 def path(self) -> str: 

141 """Accessor to the database's path.""" 

142 return self._path 

143 

144 def add_documents(self, documents: list[Document]) -> None: # noqa # pylint: disable=undefined-variable 

145 """Add multiple documents to the collection. 

146 

147 This method overrides the base class' ``add_documents`` method 

148 to enable local ID derivation. Knowing *how* the IDs are derived 

149 gives us greater understanding and querying ability of the 

150 documents in the database. Each ID is derived locally by the 

151 :meth:`_preproc` method from the file's basename, page number 

152 and page content. 

153 

154 Additionally, this method wraps the 

155 :func:`langchain_chroma.Chroma.add_texts` method which supports 

156 GPU processing and parallelisation. 

157 

158 Args: 

159 documents (list[Document]): A list of :class:`Document` 

160 objects. These objects can be of either type: 

161 

162 - :class:`langchain_core.documents.Document` 

163 - :class:`docp_core.objects.documentobject.Document` 

164 

165 """ 

166 # pylint: disable=arguments-differ 

167 if not isinstance(documents, list): 

168 documents = [documents] 

169 ids_, docs_, meta_ = self._preproc(docs=documents) 

170 self.add_texts(ids=ids_, texts=docs_, metadatas=meta_) 

171 

172 def show_all(self, include: list[str]=None) -> dict: 

173 """Return the entire contents of the collection. 

174 

175 This is an alias around ``.collection.get()``. 

176 

177 Args: 

178 include (list[str], optional): A list of what to include in 

179 the results. Can contain `"embeddings"`, `"metadatas"`, 

180 `"documents"`. IDs are always included. 

181 Defaults to `["metadatas", "documents"]`. 

182 

183 """ 

184 include = include or ['metadatas', 'documents'] 

185 return self._ocollection.get(include=include) 

186 

187 def _get_embedding_function_model(self) -> str: 

188 """Derive the path to the embedding function model. 

189 

190 .. note:: 

191 If the user has specified a *local* embedding model path, 

192 and that path exists, this model path is returned. 

193 

194 If the user provides a repository ID, this ID is passed into 

195 the :class:`SentenceTransformer` class to handle as 

196 appropriate, via the :class:`CustomEmbeddingFunction`. class 

197 

198 If the user does not provide either an embedding model path, 

199 nor a repository ID, the default embedding model value is 

200 used. 

201 

202 Raises: 

203 FileNotFoundError: If the requested embedding model path does 

204 not exist. 

205 

206 Returns: 

207 str: The path to the local embedding model, the model 

208 repository ID or the default embedding model ID. 

209 

210 """ 

211 # User has specified the path to the local model they want to use. 

212 if self._embpath: 

213 if os.path.exists(self._embpath): 

214 return self._embpath 

215 raise FileNotFoundError('The requested model path cannot be found:\n' 

216 f'-- {self._embpath}') 

217 # User has specified repo_id *only*. Lets SentenceTransformer decide how to handle. 

218 if self._repo_id and self._embpath is None: # nocover # N/A in testing env. 

219 return self._repo_id 

220 # User has *not* specified an embedding model path, nor a repo ID. 

221 return self.DEFAULT_EMBEDDING_MODEL # nocover # N/A in testing env. 

222 

223 @staticmethod 

224 def _preproc(docs: list[Document]) -> tuple: # noqa # pylint: disable=undefined-variable 

225 """Pre-process the document objects to create the IDs. 

226 

227 Parse the ``Document`` object into its parts for storage. 

228 Additionally, create the ID as a hash of the source document's 

229 basename, page number and content. 

230 

231 Args: 

232 docs (list[Document]): A list of Document objects which are 

233 used to derive the database ID, and split into the 

234 return values. 

235 

236 Returns: 

237 tuple: A tuple containing: 

238 

239 - ([ids], [texts], [metadatas]) 

240 

241 These values are passed into the :meth:`~add_texts` method 

242 for embedding and storage. 

243 

244 """ 

245 ids = [] 

246 txts = [] 

247 metas = [] 

248 for doc in docs: 

249 pc = doc.page_content 

250 m = doc.metadata 

251 pc_, src_ = map(str.encode, (pc, m['source'])) 

252 pg_ = str(m.get('pageno', 0)).zfill(4) 

253 id_ = f'id_{md5(src_).hexdigest()}_{pg_}_{md5(pc_).hexdigest()}' 

254 ids.append(id_) 

255 txts.append(pc) 

256 metas.append(m) 

257 return ids, txts, metas 

258 

259 def _set_client(self) -> None: 

260 """Set the database client object. 

261 

262 :Note: 

263 ChromaDB telemetry has been disabled. 

264 

265 """ 

266 # pylint: disable=no-member # Settings, PersistentClient 

267 settings = chromadb.Settings(anonymized_telemetry=False) 

268 self._oclient = chromadb.PersistentClient(path=self._path, settings=settings) 

269 

270 def _set_collection(self) -> None: 

271 """Set the database collection object. 

272 

273 If the collection does not already exist, it is created using 

274 *cosine* as a similarity metric. 

275 

276 """ 

277 self._ocollection = self._oclient.get_or_create_collection(self._cname, 

278 metadata={'hnsw:space': 

279 'cosine'}, 

280 embedding_function=self._embfn) 

281 

282 def _set_embedding_fn(self) -> None: 

283 """Set the embeddings function object.""" 

284 path = self._get_embedding_function_model() 

285 self._embfn = CustomEmbeddingFunction(embedding_model_path=path, 

286 local_files_only=self._offline)