feat(knowledge-base): add enabled flag for soft enable/disable - #8068
feat(knowledge-base): add enabled flag for soft enable/disable#8068YonganZhang wants to merge 3 commits into
enabled flag for soft enable/disable#8068Conversation
## Use case Currently the only way to "turn off" a knowledge base is to delete it, which loses all indexed documents and embeddings. Users frequently want to temporarily disable a KB (e.g. testing, debugging, A/B comparison) without re-indexing later. ## Change Three small additions to the existing schema: 1. `KnowledgeBase` model gets an `enabled: bool = True` field. 2. `kb_db_sqlite.migrate_to_v2()` adds the column to existing DBs via `ALTER TABLE ... ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1`, wrapped in a try/except so re-running the migration is a no-op. 3. `RetrievalManager.invalidate_sparse_cache(kb_id)` exposes a small helper so the BM25 cache can be cleared when a KB is enabled/disabled or its docs change. Useful for the upcoming dashboard toggle and any other consumer that mutates KB state. ## Backwards compatibility - New rows default to `enabled=True` → no behavior change for existing KBs. - Existing rows get `1` (true) via the migration default → no behavior change. - Pure additive: no existing field renamed or removed. ## Diff size 16 lines across 3 files. Pure model + migration + helper — no API contract change in this PR. ## Follow-up A separate PR can wire this up in the dashboard (toggle button + filter `enabled=true` in retrieval). Posting this as a small, self-contained schema change first to keep review focused.
dosubot
Bot
added
size:S
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- In
migrate_to_v2, catching a broadExceptionand silently passing can mask real migration issues; consider catching a more specific DB exception (e.g., the 'duplicate column' error) or at least logging unexpected errors while still making re-runs idempotent.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `migrate_to_v2`, catching a broad `Exception` and silently passing can mask real migration issues; consider catching a more specific DB exception (e.g., the 'duplicate column' error) or at least logging unexpected errors while still making re-runs idempotent.
## Individual Comments
### Comment 1
<location path="astrbot/core/knowledge_base/kb_db_sqlite.py" line_range="170-178" />
<code_context>
await session.commit()
+ async def migrate_to_v2(self) -> None:
+ """Add enabled column to knowledge_bases table."""
+ async with self.get_db() as session:
+ try:
+ await session.execute(
+ text("ALTER TABLE knowledge_bases ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1")
+ )
+ await session.commit()
+ except Exception:
+ pass # Column already exists
+
</code_context>
<issue_to_address>
**issue (bug_risk):** Avoid swallowing all exceptions during migration without logging or narrowing the failure mode
Catching a bare `Exception` and then `pass`-ing will hide any migration failure (SQL errors, permission issues, etc.), potentially leaving the DB in an inconsistent state and making problems hard to trace.
Instead, either:
- Catch only the specific error that indicates the column already exists, or
- Log the exception and special-case the "column already exists" scenario.
This keeps the migration idempotent without silently masking real issues.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an enabled field to the knowledge base model and adds a corresponding SQLite migration. It also implements a cache invalidation method for sparse retrievers. However, several critical issues were identified: the migration method is defined but never called, the new cache invalidation logic references a non-existent method in the SparseRetriever class which will lead to runtime errors, and the migration's error handling is overly broad and should be more specific or include logging.
| async def migrate_to_v2(self) -> None: | ||
| """Add enabled column to knowledge_bases table.""" | ||
| async with self.get_db() as session: | ||
| try: | ||
| await session.execute( | ||
| text("ALTER TABLE knowledge_bases ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1") | ||
| ) | ||
| await session.commit() | ||
| except Exception: | ||
| pass # Column already exists |
There was a problem hiding this comment.
The migrate_to_v2 method is defined but never called. To ensure existing databases are updated with the new enabled column, you must invoke this method during the initialization process. Based on the current implementation of KnowledgeBaseManager._init_kb_database in astrbot/core/knowledge_base/kb_mgr.py, a call to await self.kb_db.migrate_to_v2() should be added after migrate_to_v1().
| def invalidate_sparse_cache(self, kb_id: str) -> None: | ||
| """清除指定 KB 的 BM25 缓存""" | ||
| self.sparse_retriever.invalidate_cache(kb_id) |
There was a problem hiding this comment.
The invalidate_sparse_cache method calls self.sparse_retriever.invalidate_cache(kb_id), but the invalidate_cache method is not defined in the SparseRetriever class (see astrbot/core/knowledge_base/retrieval/sparse_retriever.py). This will cause an AttributeError at runtime. Furthermore, the current SparseRetriever implementation does not seem to utilize its _index_cache for BM25 indexing, as it rebuilds the index on every call to _retrieve_with_bm25. You should implement the caching logic and the corresponding invalidation method in SparseRetriever before exposing this helper. Additionally, ensure this new functionality is covered by unit tests.
References
- New functionality, such as handling attachments, should be accompanied by corresponding unit tests.
| except Exception: | ||
| pass # Column already exists |
There was a problem hiding this comment.
Catching a broad Exception and silently passing is generally discouraged as it can hide unexpected errors. While this is intended to handle the case where the enabled column already exists (since SQLite's ALTER TABLE doesn't support IF NOT EXISTS), it is safer to catch a more specific exception (e.g., sqlalchemy.exc.OperationalError) and verify the error message. At the very least, consider logging a warning if an exception occurs.
…ache + narrow except Three review findings from `@gemini-code-assist` and `@sourcery-ai`: 1. HIGH: `migrate_to_v2` was defined but never called. Wire it into `KnowledgeBaseDBSqlite.initialize()` right after the `engine.begin()` PRAGMA setup block commits, so every existing database picks up the new `enabled` column on startup. Idempotent — the method swallows the "duplicate column" path quietly. 2. HIGH: `RetrievalManager.invalidate_sparse_cache` called `self.sparse_retriever.invalidate_cache(kb_id)` but `SparseRetriever` had no such method, so any call to invalidate sparse cache would raise AttributeError at runtime. Add `SparseRetriever.invalidate_cache(kb_id)` that pops the entry from the internal `_index_cache` dict (no-op if absent). 3. MEDIUM: `migrate_to_v2` caught bare `Exception` and `pass`'d, hiding every failure (FS / permission / SQL). Narrow the swallow: only the SQLite "duplicate column" / "already exists" idempotency path returns silently; every other exception is logged at debug so a real schema problem is no longer invisible. Keeps the migration idempotent without becoming an error sink.
dosubot
Bot
added
size:M
|
Addressed the bot review findings in commit ff8f65f: HIGH bugs (both shipped in this update):
MEDIUM (also fixed):
|
There was a problem hiding this comment.
Code Review
This pull request introduces an enabled field to the knowledge base system, including a SQLite schema migration and cache invalidation logic for the sparse retriever. Feedback highlights that the migration error handling should re-raise unexpected exceptions to prevent the application from running in a broken state. Additionally, the migration execution order is currently inconsistent, as migrate_to_v2 is called before migrate_to_v1, and unit tests for the new functionality are recommended.
| except Exception as e: | ||
| msg = str(e).lower() | ||
| if "duplicate column" in msg or "already exists" in msg: | ||
| # Column already present from a prior migration run — expected. | ||
| return | ||
| logger.debug(f"知识库 v2 迁移跳过 (非重复列错误): {e!r}") |
There was a problem hiding this comment.
Unexpected errors during migration are currently logged at debug level and swallowed, which allows the application to start with an incomplete schema. Since the enabled column is defined as NOT NULL in the model, this will likely lead to runtime errors when the field is accessed. It is safer to log unexpected errors at error level and re-raise the exception to prevent the application from running in a broken state.
| except Exception as e: | |
| msg = str(e).lower() | |
| if "duplicate column" in msg or "already exists" in msg: | |
| # Column already present from a prior migration run — expected. | |
| return | |
| logger.debug(f"知识库 v2 迁移跳过 (非重复列错误): {e!r}") | |
| except Exception as e: | |
| msg = str(e).lower() | |
| if "duplicate column" in msg or "already exists" in msg: | |
| # Column already present from a prior migration run — expected. | |
| return | |
| logger.error(f"知识库 v2 迁移失败 (非重复列错误): {e!r}") | |
| raise |
| # Idempotent and re-entrant — safe to call on every startup; the | ||
| # method swallows the "duplicate column" path quietly and logs anything | ||
| # else at debug. See migrate_to_v2 docstring. | ||
| await self.migrate_to_v2() |
There was a problem hiding this comment.
The migration migrate_to_v2 is called here within initialize, but migrate_to_v1 is called externally in KnowledgeBaseManager._init_kb_database (line 61 of kb_mgr.py). This results in migrations running out of order (v2 before v1). It is recommended to consolidate migration calls within initialize or ensure they are called sequentially in the manager to maintain a clear and predictable schema evolution path. Additionally, as this is new functionality, please ensure it is accompanied by corresponding unit tests.
References
- New functionality should be accompanied by corresponding unit tests.
…d migration errors Two new findings from @gemini-code-assist's re-review: 1. ORDER BUG: previous commit invoked migrate_to_v2 inside KnowledgeBaseDBSqlite.initialize() right after the PRAGMA block, but migrate_to_v1 is invoked from kb_mgr.py:61 (the manager layer, not the DB initializer). That meant v2 ran *before* v1 on every startup — wrong order, and migrate_to_v1 creates indices that v2 ought to build on top of. Fix: remove the migrate_to_v2 call from initialize(), and add it immediately after the existing self.kb_db.migrate_to_v1() in kb_mgr.py so v1 → v2 always run in declaration order. 2. SWALLOWED ERRORS: previous commit caught everything-not-"duplicate column" and only logger.debug'd. As gemini-code-assist pointed out, a real schema failure (locked DB, missing table, permission error) should NOT silently leave the application running with the wrong schema. Switch the non-idempotent path to logger.error + raise so the manager surfaces the broken state on startup instead of running with stale schema.
|
Addressed both new findings from 1. Migration order (v2 was running before v1) ✅ 2. Re-raise unexpected migration errors ✅ 3. Unit tests — agree these would help; the existing
|
There was a problem hiding this comment.
Code Review
This pull request introduces a database migration to add an 'enabled' column to the knowledge base table, updates the corresponding data model, and adds cache invalidation methods for the sparse retriever. Feedback focuses on improving the migration logic by using idiomatic transaction handling, fixing docstring inconsistencies, and narrowing exception types. Additionally, it was noted that the cache invalidation logic currently acts as dead code because the cache is not utilized during retrieval, indicating a need for the actual caching implementation to resolve performance bottlenecks.
| async def migrate_to_v2(self) -> None: | ||
| """Add enabled column to knowledge_bases table. | ||
|
|
||
| SQLite has no IF NOT EXISTS for ALTER TABLE ADD COLUMN, so the | ||
| re-run case raises OperationalError("duplicate column name"). That | ||
| specific failure is the expected idempotent path; anything else | ||
| is logged at debug so a real schema problem is not lost. | ||
| """ | ||
| async with self.get_db() as session: | ||
| try: | ||
| await session.execute( | ||
| text("ALTER TABLE knowledge_bases ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1") | ||
| ) | ||
| await session.commit() | ||
| except Exception as e: | ||
| msg = str(e).lower() | ||
| if "duplicate column" in msg or "already exists" in msg: | ||
| # Column already present from a prior migration run — expected. | ||
| return | ||
| # Real schema failure — let it propagate so the manager surfaces | ||
| # a broken-knowledge-base state on startup instead of silently | ||
| # running with the old schema. | ||
| logger.error(f"知识库 v2 迁移失败: {e!r}") | ||
| raise |
There was a problem hiding this comment.
The migrate_to_v2 method has a few minor issues:
- Inconsistency: It uses manual
try/exceptandsession.commit()instead of theasync with session.begin():pattern used inmigrate_to_v1. Usingsession.begin()is more idiomatic as it handles transactions automatically and ensures consistency across the codebase. - Docstring Discrepancy: The docstring states that non-expected failures are logged at
debug, but the implementation useslogger.error. Given that a migration failure is a significant event that re-raises the exception,logger.erroris appropriate, but the docstring should be updated to reflect this. - Broad Exception: Catching
Exceptionis generally discouraged. While SQLite error handling can be driver-specific, it's better to catch specific SQLAlchemy exceptions if possible. - Unit Tests: As this is new functionality, please ensure it is accompanied by corresponding unit tests.
async def migrate_to_v2(self) -> None:
"""Add enabled column to knowledge_bases table.
SQLite has no IF NOT EXISTS for ALTER TABLE ADD COLUMN, so the
re-run case raises OperationalError("duplicate column name"). That
specific failure is the expected idempotent path; anything else
is logged so a real schema problem is not lost.
"""
async with self.get_db() as session:
try:
async with session.begin():
await session.execute(
text("ALTER TABLE knowledge_bases ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1")
)
except Exception as e:
msg = str(e).lower()
if "duplicate column" in msg or "already exists" in msg:
# Column already present from a prior migration run — expected.
return
# Real schema failure — let it propagate so the manager surfaces
# a broken-knowledge-base state on startup instead of silently
# running with the old schema.
logger.error(f"知识库 v2 迁移失败: {e!r}")
raiseReferences
- New functionality, such as handling attachments, should be accompanied by corresponding unit tests.
| def invalidate_cache(self, kb_id: str) -> None: | ||
| """清除指定 KB 的 BM25 索引缓存。 | ||
|
|
||
| 当 KB 的 enabled 状态切换或内容增删改时调用,确保下次检索使用 | ||
| 最新的索引。pop 用 default=None 保证未缓存的 kb_id 调用安全 (no-op)。 | ||
| """ | ||
| self._index_cache.pop(kb_id, None) |
There was a problem hiding this comment.
The _index_cache dictionary (initialized in __init__) is modified here but never actually read in retrieve or _retrieve_with_bm25. This makes the invalidate_cache method dead code and indicates that the BM25 index is still being rebuilt on every call to _retrieve_with_bm25, which is a performance bottleneck. You should implement the caching logic to store and reuse the BM25Okapi instances per kb_id to make this invalidation logic functional.
a4c4a7d to
9bd38ca
Compare
1 participant
Use case
Currently the only way to "turn off" a knowledge base is to delete it, which loses all indexed documents and embeddings. Users frequently want to temporarily disable a KB (e.g. testing, debugging, A/B comparison) without paying the re-indexing cost later.
Change
Three small additions to the existing schema:
KnowledgeBasemodel gets anenabled: bool = Truefield.kb_db_sqlite.migrate_to_v2()adds the column to existing DBs viaALTER TABLE ... ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1, wrapped intry/exceptso re-running the migration is a no-op.RetrievalManager.invalidate_sparse_cache(kb_id)exposes a small helper to clear the BM25 cache when a KB is enabled/disabled or its docs change. Useful for the dashboard toggle and any other consumer that mutates KB state.Backwards compatibility
enabled=True→ no behavior change for existing KBs.1(true) via the migration default → no behavior change.Diff size
16 lines across 3 files. Pure model + migration + helper — no API contract change in this PR.
Follow-up
A separate PR can wire this up in the dashboard (toggle button + filter
enabled=truein retrieval). Posting this as a small, self-contained schema change first to keep review focused.Summary by Sourcery
Introduce a soft enable/disable mechanism for knowledge bases and wire supporting schema migration and cache invalidation helper methods.
New Features:
Enhancements: