| 227 | |
| 228 | ################ INSERT METHODS ################ |
| 229 | async def upsert(self, data: dict[str, dict]): |
| 230 | left_data = {k: v for k, v in data.items() if k not in self._data} |
| 231 | self._data.update(left_data) |
| 232 | # print(self._data) |
| 233 | # values = [] |
| 234 | if self.namespace == "text_chunks": |
| 235 | list_data = [ |
| 236 | { |
| 237 | "__id__": k, |
| 238 | **{k1: v1 for k1, v1 in v.items()}, |
| 239 | } |
| 240 | for k, v in data.items() |
| 241 | ] |
| 242 | contents = [v["content"] for v in data.values()] |
| 243 | batches = [ |
| 244 | contents[i : i + self._max_batch_size] |
| 245 | for i in range(0, len(contents), self._max_batch_size) |
| 246 | ] |
| 247 | embeddings_list = await asyncio.gather( |
| 248 | *[self.embedding_func(batch) for batch in batches] |
| 249 | ) |
| 250 | embeddings = np.concatenate(embeddings_list) |
| 251 | for i, d in enumerate(list_data): |
| 252 | d["__vector__"] = embeddings[i] |
| 253 | # print(list_data) |
| 254 | for item in list_data: |
| 255 | merge_sql = SQL_TEMPLATES["merge_chunk"] |
| 256 | data = { |
| 257 | "check_id": item["__id__"], |
| 258 | "id": item["__id__"], |
| 259 | "content": item["content"], |
| 260 | "workspace": self.db.workspace, |
| 261 | "tokens": item["tokens"], |
| 262 | "chunk_order_index": item["chunk_order_index"], |
| 263 | "full_doc_id": item["full_doc_id"], |
| 264 | "content_vector": item["__vector__"], |
| 265 | } |
| 266 | # print(merge_sql) |
| 267 | await self.db.execute(merge_sql, data) |
| 268 | |
| 269 | if self.namespace == "full_docs": |
| 270 | for k, v in self._data.items(): |
| 271 | # values.clear() |
| 272 | merge_sql = SQL_TEMPLATES["merge_doc_full"] |
| 273 | data = { |
| 274 | "check_id": k, |
| 275 | "id": k, |
| 276 | "content": v["content"], |
| 277 | "workspace": self.db.workspace, |
| 278 | } |
| 279 | # print(merge_sql) |
| 280 | await self.db.execute(merge_sql, data) |
| 281 | return left_data |
| 282 | |
| 283 | async def index_done_callback(self): |
| 284 | if self.namespace in ["full_docs", "text_chunks"]: |