Skip to content

Commit 234f782

Browse files
authored
feat(kb): add list_docs and kb_exists method (#152)
* fix: knowledgebase add method & list_chunks * fix: knowledgebase add method * feat: knowledge exists and list_docs * fix: fix list_docs empty * fix: fix bug * fix: viking delete as docs level * fix: delete method with adapter bug
1 parent a918d4d commit 234f782

File tree

3 files changed

+421
-50
lines changed

3 files changed

+421
-50
lines changed

veadk/database/database_adapter.py

Lines changed: 125 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,26 @@ def __init__(self, client):
2828

2929
self.client: RedisDatabase = client
3030

31-
def add(self, data: list[str], index: str):
31+
def index_exists(self, index: str) -> bool:
32+
"""
33+
Check if the index (key) exists in Redis.
34+
35+
Args:
36+
index: The Redis key to check
37+
38+
Returns:
39+
bool: True if the key exists, False otherwise
40+
"""
41+
try:
42+
# Use Redis EXISTS command to check if key exists
43+
return bool(self.client._client.exists(index))
44+
except Exception as e:
45+
logger.error(
46+
f"Failed to check if key exists in Redis: index={index} error={e}"
47+
)
48+
return False
49+
50+
def add(self, data: list[str], index: str, **kwargs):
3251
logger.debug(f"Adding documents to Redis database: index={index}")
3352

3453
try:
@@ -78,7 +97,7 @@ def delete_doc(self, index: str, id: str) -> bool:
7897
)
7998
return False
8099

81-
def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
100+
def list_chunks(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
82101
logger.debug(f"Listing documents from Redis database: index={index}")
83102
try:
84103
# Get all documents from Redis
@@ -99,6 +118,24 @@ def __init__(self, client):
99118

100119
self.client: MysqlDatabase = client
101120

121+
def index_exists(self, index: str) -> bool:
122+
"""
123+
Check if the table (index) exists in MySQL database.
124+
125+
Args:
126+
index: The table name to check
127+
128+
Returns:
129+
bool: True if the table exists, False otherwise
130+
"""
131+
try:
132+
return self.client.table_exists(index)
133+
except Exception as e:
134+
logger.error(
135+
f"Failed to check if table exists in MySQL: index={index} error={e}"
136+
)
137+
return False
138+
102139
def create_table(self, table_name: str):
103140
logger.debug(f"Creating table for SQL database: table_name={table_name}")
104141

@@ -111,7 +148,7 @@ def create_table(self, table_name: str):
111148
"""
112149
self.client.add(sql)
113150

114-
def add(self, data: list[str], index: str):
151+
def add(self, data: list[str], index: str, **kwargs):
115152
logger.debug(
116153
f"Adding documents to SQL database: table_name={index} data_len={len(data)}"
117154
)
@@ -188,6 +225,25 @@ def __init__(self, client):
188225

189226
self.client: OpenSearchVectorDatabase = client
190227

228+
def index_exists(self, index: str) -> bool:
229+
"""
230+
Check if the collection (index) exists in OpenSearch.
231+
232+
Args:
233+
index: The collection name to check
234+
235+
Returns:
236+
bool: True if the collection exists, False otherwise
237+
"""
238+
try:
239+
self._validate_index(index)
240+
return self.client.collection_exists(index)
241+
except Exception as e:
242+
logger.error(
243+
f"Failed to check if collection exists in OpenSearch: index={index} error={e}"
244+
)
245+
return False
246+
191247
def _validate_index(self, index: str):
192248
"""
193249
Verify whether the string conforms to the naming rules of index_name in OpenSearch.
@@ -203,7 +259,7 @@ def _validate_index(self, index: str):
203259
"The index name does not conform to the naming rules of OpenSearch"
204260
)
205261

206-
def add(self, data: list[str], index: str):
262+
def add(self, data: list[str], index: str, **kwargs):
207263
self._validate_index(index)
208264

209265
logger.debug(
@@ -247,7 +303,7 @@ def delete_doc(self, index: str, id: str) -> bool:
247303
)
248304
return False
249305

250-
def list_docs(self, index: str, offset: int = 0, limit: int = 1000) -> list[dict]:
306+
def list_chunks(self, index: str, offset: int = 0, limit: int = 1000) -> list[dict]:
251307
self._validate_index(index)
252308
logger.debug(f"Listing documents from vector database: index={index}")
253309
return self.client.list_docs(collection_name=index, offset=offset, limit=limit)
@@ -259,6 +315,25 @@ def __init__(self, client):
259315

260316
self.client: VikingDatabase = client
261317

318+
def index_exists(self, index: str) -> bool:
319+
"""
320+
Check if the collection (index) exists in VikingDB.
321+
322+
Args:
323+
index: The collection name to check
324+
325+
Returns:
326+
bool: True if the collection exists, False otherwise
327+
"""
328+
try:
329+
self._validate_index(index)
330+
return self.client.collection_exists(index)
331+
except Exception as e:
332+
logger.error(
333+
f"Failed to check if collection exists in VikingDB: index={index} error={e}"
334+
)
335+
return False
336+
262337
def _validate_index(self, index: str):
263338
"""
264339
Only English letters, numbers, and underscores (_) are allowed.
@@ -322,6 +397,13 @@ def delete_doc(self, index: str, id: str) -> bool:
322397
logger.debug(f"Deleting documents from vector database: index={index} id={id}")
323398
return self.client.delete_by_id(collection_name=index, id=id)
324399

400+
def list_chunks(self, index: str, offset: int, limit: int) -> list[dict]:
401+
self._validate_index(index)
402+
logger.debug(f"Listing documents from vector database: index={index}")
403+
return self.client.list_chunks(
404+
collection_name=index, offset=offset, limit=limit
405+
)
406+
325407
def list_docs(self, index: str, offset: int, limit: int) -> list[dict]:
326408
self._validate_index(index)
327409
logger.debug(f"Listing documents from vector database: index={index}")
@@ -334,6 +416,25 @@ def __init__(self, client):
334416

335417
self.client: VikingMemoryDatabase = client
336418

419+
def index_exists(self, index: str) -> bool:
420+
"""
421+
Check if the collection (index) exists in VikingMemoryDB.
422+
423+
Note:
424+
VikingMemoryDatabase does not support checking if a collection exists.
425+
This method always returns False.
426+
427+
Args:
428+
index: The collection name to check
429+
430+
Returns:
431+
bool: Always returns False as VikingMemoryDatabase does not support this functionality
432+
"""
433+
logger.warning(
434+
"VikingMemoryDatabase does not support checking if a collection exists"
435+
)
436+
raise NotImplementedError("VikingMemoryDatabase does not support index_exists")
437+
337438
def _validate_index(self, index: str):
338439
if not (
339440
isinstance(index, str)
@@ -371,7 +472,7 @@ def delete(self, index: str) -> bool:
371472
def delete_docs(self, index: str, ids: list[int]):
372473
raise NotImplementedError("VikingMemoryDatabase does not support delete_docs")
373474

374-
def list_docs(self, index: str):
475+
def list_chunks(self, index: str):
375476
raise NotImplementedError("VikingMemoryDatabase does not support list_docs")
376477

377478

@@ -381,6 +482,23 @@ def __init__(self, client):
381482

382483
self.client: LocalDataBase = client
383484

485+
def index_exists(self, index: str) -> bool:
486+
"""
487+
Check if the index exists in LocalDataBase.
488+
489+
Note:
490+
LocalDataBase does not support checking if an index exists.
491+
This method always returns False.
492+
493+
Args:
494+
index: The index name to check (not used in LocalDataBase)
495+
496+
Returns:
497+
bool: Always returns False as LocalDataBase does not support this functionality
498+
"""
499+
logger.warning("LocalDataBase does not support checking if an index exists")
500+
return True
501+
384502
def add(self, data: list[str], **kwargs):
385503
self.client.add(data)
386504

@@ -393,7 +511,7 @@ def delete(self, index: str) -> bool:
393511
def delete_doc(self, index: str, id: str) -> bool:
394512
return self.client.delete_doc(id)
395513

396-
def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
514+
def list_chunks(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
397515
return self.client.list_docs(offset=offset, limit=limit)
398516

399517

0 commit comments

Comments
 (0)