mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-06 12:50:15 +08:00
Compare commits
6 Commits
refactor/b
...
fix/8414
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dcb714ec8c | ||
|
|
0db7fc9b39 | ||
|
|
556903c135 | ||
|
|
bdc32bb78c | ||
|
|
c70a1924fe | ||
|
|
6ae103a24f |
@@ -3,7 +3,6 @@ from typing import Any
|
||||
from mcp.types import CallToolResult
|
||||
|
||||
from astrbot.core.agent.hooks import BaseAgentRunHooks
|
||||
from astrbot.core.agent.message import Message
|
||||
from astrbot.core.agent.run_context import ContextWrapper
|
||||
from astrbot.core.agent.tool import FunctionTool
|
||||
from astrbot.core.astr_agent_context import AstrAgentContext
|
||||
@@ -70,37 +69,6 @@ class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]):
|
||||
tool_result,
|
||||
)
|
||||
|
||||
# special handle web_search_tavily
|
||||
platform_name = run_context.context.event.get_platform_name()
|
||||
if (
|
||||
platform_name == "webchat"
|
||||
and tool.name
|
||||
in [
|
||||
"web_search_baidu",
|
||||
"web_search_tavily",
|
||||
"web_search_bocha",
|
||||
"web_search_brave",
|
||||
]
|
||||
and len(run_context.messages) > 0
|
||||
and tool_result
|
||||
and len(tool_result.content)
|
||||
):
|
||||
# inject system prompt
|
||||
first_part = run_context.messages[0]
|
||||
if (
|
||||
isinstance(first_part, Message)
|
||||
and first_part.role == "system"
|
||||
and first_part.content
|
||||
and isinstance(first_part.content, str)
|
||||
):
|
||||
# we assume system part is str
|
||||
first_part.content += (
|
||||
"Always cite web search results you rely on. "
|
||||
"Index is a unique identifier for each search result. "
|
||||
"Use the exact citation format <ref>index</ref> (e.g. <ref>abcd.3</ref>) "
|
||||
"after the sentence that uses the information. Do not invent citations."
|
||||
)
|
||||
|
||||
|
||||
class EmptyAgentHooks(BaseAgentRunHooks[AstrAgentContext]):
|
||||
pass
|
||||
|
||||
@@ -115,6 +115,20 @@ from astrbot.core.utils.quoted_message_parser import (
|
||||
from astrbot.core.utils.string_utils import normalize_and_dedupe_strings
|
||||
|
||||
LLM_ERROR_MESSAGE_EXTRA_KEY = "_llm_error_message"
|
||||
WEB_SEARCH_CITATION_TOOL_NAMES = frozenset(
|
||||
{
|
||||
"web_search_baidu",
|
||||
"web_search_tavily",
|
||||
"web_search_bocha",
|
||||
"web_search_brave",
|
||||
}
|
||||
)
|
||||
WEB_SEARCH_CITATION_PROMPT = (
|
||||
"Always cite web search results you rely on. "
|
||||
"Index is a unique identifier for each search result. "
|
||||
"Use the exact citation format <ref>index</ref> (e.g. <ref>abcd.3</ref>) "
|
||||
"after the sentence that uses the information. Do not invent citations."
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -1149,6 +1163,23 @@ async def _apply_web_search_tools(
|
||||
req.func_tool.add_tool(tool_mgr.get_builtin_tool(BaiduWebSearchTool))
|
||||
|
||||
|
||||
def _apply_web_search_citation_prompt(
|
||||
event: AstrMessageEvent,
|
||||
req: ProviderRequest,
|
||||
) -> None:
|
||||
if event.get_platform_name() != "webchat" or not req.func_tool:
|
||||
return
|
||||
|
||||
if not any(req.func_tool.get_tool(name) for name in WEB_SEARCH_CITATION_TOOL_NAMES):
|
||||
return
|
||||
|
||||
system_prompt = req.system_prompt or ""
|
||||
if WEB_SEARCH_CITATION_PROMPT in system_prompt:
|
||||
return
|
||||
|
||||
req.system_prompt = f"{system_prompt}\n{WEB_SEARCH_CITATION_PROMPT}\n"
|
||||
|
||||
|
||||
def _get_compress_provider(
|
||||
config: MainAgentBuildConfig,
|
||||
plugin_context: Context,
|
||||
@@ -1520,6 +1551,8 @@ async def build_main_agent(
|
||||
if action_type == "live":
|
||||
req.system_prompt += f"\n{LIVE_MODE_SYSTEM_PROMPT}\n"
|
||||
|
||||
_apply_web_search_citation_prompt(event, req)
|
||||
|
||||
reset_coro = agent_runner.reset(
|
||||
provider=provider,
|
||||
request=req,
|
||||
|
||||
@@ -5,10 +5,7 @@ from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
|
||||
from deprecated import deprecated
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.engine import make_url
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.pool import NullPool
|
||||
|
||||
from astrbot.core.db.po import (
|
||||
ApiKey,
|
||||
@@ -32,19 +29,6 @@ from astrbot.core.db.po import (
|
||||
)
|
||||
|
||||
|
||||
def _configure_sqlite_connection(dbapi_connection, connection_record) -> None:
|
||||
cursor = dbapi_connection.cursor()
|
||||
try:
|
||||
cursor.execute("PRAGMA journal_mode=WAL")
|
||||
cursor.execute("PRAGMA synchronous=NORMAL")
|
||||
cursor.execute("PRAGMA cache_size=20000")
|
||||
cursor.execute("PRAGMA temp_store=MEMORY")
|
||||
cursor.execute("PRAGMA mmap_size=134217728")
|
||||
cursor.execute("PRAGMA optimize")
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseDatabase(abc.ABC):
|
||||
"""数据库基类"""
|
||||
@@ -57,29 +41,14 @@ class BaseDatabase(abc.ABC):
|
||||
# second write is attempted. Setting timeout=30 tells SQLite to
|
||||
# wait up to 30 s for the lock, which is enough to ride out brief
|
||||
# write bursts from concurrent agent/metrics/session operations.
|
||||
db_url = make_url(self.DATABASE_URL)
|
||||
is_sqlite = db_url.get_backend_name() == "sqlite"
|
||||
is_sqlite = "sqlite" in self.DATABASE_URL
|
||||
connect_args = {"timeout": 30} if is_sqlite else {}
|
||||
engine_kwargs = {
|
||||
"echo": False,
|
||||
"future": True,
|
||||
"connect_args": connect_args,
|
||||
}
|
||||
if is_sqlite:
|
||||
# Keep SQLite async engines off SQLAlchemy's default async queue
|
||||
# pool so packaged runtimes don't depend on dialect-specific pool
|
||||
# event support.
|
||||
engine_kwargs["poolclass"] = NullPool
|
||||
self.engine = create_async_engine(
|
||||
self.DATABASE_URL,
|
||||
**engine_kwargs,
|
||||
echo=False,
|
||||
future=True,
|
||||
connect_args=connect_args,
|
||||
)
|
||||
if is_sqlite:
|
||||
event.listen(
|
||||
self.engine.sync_engine,
|
||||
"connect",
|
||||
_configure_sqlite_connection,
|
||||
)
|
||||
self.AsyncSessionLocal = async_sessionmaker(
|
||||
self.engine,
|
||||
class_=AsyncSession,
|
||||
|
||||
@@ -53,6 +53,7 @@ class SQLiteDatabase(BaseDatabase):
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.run_sync(SQLModel.metadata.create_all)
|
||||
await conn.execute(text("PRAGMA journal_mode=WAL"))
|
||||
await conn.execute(text("PRAGMA busy_timeout=30000"))
|
||||
await conn.execute(text("PRAGMA synchronous=NORMAL"))
|
||||
await conn.execute(text("PRAGMA cache_size=20000"))
|
||||
await conn.execute(text("PRAGMA temp_store=MEMORY"))
|
||||
|
||||
@@ -5,11 +5,8 @@ from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import Column, Text, bindparam
|
||||
from sqlalchemy.dialects import sqlite
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy.pool import NullPool
|
||||
from sqlalchemy.schema import CreateTable
|
||||
from sqlmodel import Field, MetaData, SQLModel, col, func, select, text
|
||||
|
||||
from astrbot.core import logger
|
||||
@@ -63,7 +60,8 @@ class DocumentStorage:
|
||||
"""Initialize the SQLite database and create the documents table if it doesn't exist."""
|
||||
await self.connect()
|
||||
async with self.engine.begin() as conn: # type: ignore
|
||||
await self._ensure_documents_table(conn)
|
||||
# Create tables using SQLModel
|
||||
await conn.run_sync(BaseDocModel.metadata.create_all)
|
||||
|
||||
try:
|
||||
await conn.execute(
|
||||
@@ -93,59 +91,15 @@ class DocumentStorage:
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
await conn.execute(
|
||||
text(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_doc_id_unique ON documents(doc_id)",
|
||||
),
|
||||
)
|
||||
|
||||
await self._initialize_fts5(conn)
|
||||
await conn.commit()
|
||||
|
||||
async def _ensure_documents_table(self, executor) -> None:
|
||||
"""Create the document table from the SQLModel definition."""
|
||||
result = await executor.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT 1
|
||||
FROM sqlite_master
|
||||
WHERE type='table' AND name=:table_name
|
||||
LIMIT 1
|
||||
""",
|
||||
),
|
||||
{"table_name": Document.__tablename__},
|
||||
)
|
||||
if result.scalar_one_or_none() is not None:
|
||||
await self._ensure_doc_id_unique_index(executor)
|
||||
return
|
||||
|
||||
create_table = CreateTable(Document.__table__, if_not_exists=True) # type: ignore[attr-defined]
|
||||
|
||||
await executor.execute(
|
||||
text(str(create_table.compile(dialect=sqlite.dialect())))
|
||||
)
|
||||
await self._ensure_doc_id_unique_index(executor)
|
||||
|
||||
async def _ensure_doc_id_unique_index(self, executor) -> None:
|
||||
duplicate_result = await executor.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT doc_id
|
||||
FROM documents
|
||||
GROUP BY doc_id
|
||||
HAVING COUNT(*) > 1
|
||||
LIMIT 1
|
||||
""",
|
||||
),
|
||||
)
|
||||
if duplicate_result.scalar_one_or_none() is not None:
|
||||
logger.warning(
|
||||
"Skipping documents.doc_id unique index migration because duplicate "
|
||||
f"doc_id values already exist in {self.db_path}.",
|
||||
)
|
||||
return
|
||||
|
||||
await executor.execute(
|
||||
text(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS "
|
||||
"idx_documents_doc_id_unique ON documents(doc_id)",
|
||||
),
|
||||
)
|
||||
|
||||
async def _initialize_fts5(self, executor) -> None:
|
||||
try:
|
||||
await self._create_fts5_table(executor, if_not_exists=True)
|
||||
@@ -249,7 +203,6 @@ class DocumentStorage:
|
||||
self.DATABASE_URL,
|
||||
echo=False,
|
||||
future=True,
|
||||
poolclass=NullPool,
|
||||
)
|
||||
self.async_session_maker = sessionmaker(
|
||||
self.engine, # type: ignore
|
||||
|
||||
@@ -33,6 +33,8 @@ class EventBus:
|
||||
# abconf uuid -> scheduler
|
||||
self.pipeline_scheduler_mapping = pipeline_scheduler_mapping
|
||||
self.astrbot_config_mgr = astrbot_config_mgr
|
||||
# 持有正在执行的 pipeline 任务的强引用, 防止 task 在 pending 状态被 GC 回收
|
||||
self._pending_tasks: set[asyncio.Task] = set()
|
||||
|
||||
async def dispatch(self) -> None:
|
||||
while True:
|
||||
@@ -47,7 +49,18 @@ class EventBus:
|
||||
f"PipelineScheduler not found for id: {conf_id}, event ignored."
|
||||
)
|
||||
continue
|
||||
asyncio.create_task(scheduler.execute(event))
|
||||
task = asyncio.create_task(scheduler.execute(event))
|
||||
self._pending_tasks.add(task)
|
||||
task.add_done_callback(self._on_task_done)
|
||||
|
||||
def _on_task_done(self, task: asyncio.Task) -> None:
|
||||
"""pipeline 任务结束回调: 移除强引用并暴露未捕获的异常"""
|
||||
self._pending_tasks.discard(task)
|
||||
if task.cancelled():
|
||||
return
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
logger.error("pipeline 任务执行异常", exc_info=exc)
|
||||
|
||||
def _print_event(self, event: AstrMessageEvent, conf_name: str) -> None:
|
||||
"""用于记录事件信息
|
||||
|
||||
@@ -2,9 +2,8 @@ from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import delete, event, func, select, text, update
|
||||
from sqlalchemy import delete, func, select, text, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.pool import NullPool
|
||||
from sqlmodel import col, desc
|
||||
|
||||
from astrbot.core import logger
|
||||
@@ -20,19 +19,6 @@ if TYPE_CHECKING:
|
||||
from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
|
||||
|
||||
|
||||
def _configure_sqlite_connection(dbapi_connection, connection_record) -> None:
|
||||
cursor = dbapi_connection.cursor()
|
||||
try:
|
||||
cursor.execute("PRAGMA journal_mode=WAL")
|
||||
cursor.execute("PRAGMA synchronous=NORMAL")
|
||||
cursor.execute("PRAGMA cache_size=20000")
|
||||
cursor.execute("PRAGMA temp_store=MEMORY")
|
||||
cursor.execute("PRAGMA mmap_size=134217728")
|
||||
cursor.execute("PRAGMA optimize")
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
class KBSQLiteDatabase:
|
||||
def __init__(self, db_path: str | None = None) -> None:
|
||||
"""初始化知识库数据库
|
||||
@@ -54,12 +40,8 @@ class KBSQLiteDatabase:
|
||||
self.engine = create_async_engine(
|
||||
self.DATABASE_URL,
|
||||
echo=False,
|
||||
poolclass=NullPool,
|
||||
)
|
||||
event.listen(
|
||||
self.engine.sync_engine,
|
||||
"connect",
|
||||
_configure_sqlite_connection,
|
||||
pool_pre_ping=True,
|
||||
pool_recycle=3600,
|
||||
)
|
||||
|
||||
# 创建会话工厂
|
||||
|
||||
@@ -5,8 +5,6 @@ import base64
|
||||
from collections.abc import AsyncGenerator
|
||||
from dataclasses import replace
|
||||
|
||||
from sqlalchemy.exc import OperationalError
|
||||
|
||||
from astrbot.core import db_helper, logger
|
||||
from astrbot.core.agent.message import (
|
||||
CheckpointData,
|
||||
@@ -521,15 +519,6 @@ class InternalAgentSubStage(Stage):
|
||||
BLOCKED = {"dGZid2h2d3IuY2xvdWQuc2VhbG9zLmlv", "a291cmljaGF0"}
|
||||
decoded_blocked = [base64.b64decode(b).decode("utf-8") for b in BLOCKED]
|
||||
|
||||
PROVIDER_STATS_SQLITE_LOCK_RETRY_ATTEMPTS = 3
|
||||
PROVIDER_STATS_SQLITE_LOCK_RETRY_BASE_DELAY = 0.2
|
||||
|
||||
|
||||
def _is_sqlite_database_locked_error(exc: OperationalError) -> bool:
|
||||
raw = getattr(exc, "orig", exc)
|
||||
message = str(raw).lower()
|
||||
return "database" in message and "locked" in message
|
||||
|
||||
|
||||
async def _record_internal_agent_stats(
|
||||
event: AstrMessageEvent,
|
||||
@@ -560,35 +549,15 @@ async def _record_internal_agent_stats(
|
||||
status = "error"
|
||||
else:
|
||||
status = "completed"
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
|
||||
await db_helper.insert_provider_stat(
|
||||
umo=event.unified_msg_origin,
|
||||
conversation_id=conversation_id,
|
||||
provider_id=provider_config.get("id", "") or provider.meta().id,
|
||||
provider_model=provider.get_model(),
|
||||
status=status,
|
||||
stats=stats.to_dict(),
|
||||
agent_type="internal",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Persist provider stats failed: %s", e, exc_info=True)
|
||||
return
|
||||
|
||||
for attempt in range(PROVIDER_STATS_SQLITE_LOCK_RETRY_ATTEMPTS):
|
||||
last_attempt = attempt == PROVIDER_STATS_SQLITE_LOCK_RETRY_ATTEMPTS - 1
|
||||
try:
|
||||
await db_helper.insert_provider_stat(
|
||||
umo=event.unified_msg_origin,
|
||||
conversation_id=conversation_id,
|
||||
provider_id=provider_config.get("id", "") or provider.meta().id,
|
||||
provider_model=provider.get_model(),
|
||||
status=status,
|
||||
stats=stats.to_dict(),
|
||||
agent_type="internal",
|
||||
)
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except OperationalError as e:
|
||||
if _is_sqlite_database_locked_error(e) and not last_attempt:
|
||||
await asyncio.sleep(
|
||||
PROVIDER_STATS_SQLITE_LOCK_RETRY_BASE_DELAY * (2**attempt)
|
||||
)
|
||||
continue
|
||||
logger.warning("Persist provider stats failed: %s", e, exc_info=True)
|
||||
break
|
||||
except Exception as e:
|
||||
logger.warning("Persist provider stats failed: %s", e, exc_info=True)
|
||||
break
|
||||
|
||||
1
dashboard/pnpm-lock.yaml
generated
1
dashboard/pnpm-lock.yaml
generated
@@ -998,6 +998,7 @@ packages:
|
||||
|
||||
'@ungap/structured-clone@1.3.0':
|
||||
resolution: {integrity: sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g==}
|
||||
deprecated: Potential CWE-502 - Update to 1.3.1 or higher
|
||||
|
||||
'@vitejs/plugin-vue@5.2.4':
|
||||
resolution: {integrity: sha512-7Yx/SXSOcQq5HiiV3orevHUFn+pmMB4cgbEkDYgnkUWb0WfeQ/wa2yFv6D5ICiCQOVpjA7vYDXrC7AGO8yjDHA==}
|
||||
|
||||
3
dashboard/pnpm-workspace.yaml
Normal file
3
dashboard/pnpm-workspace.yaml
Normal file
@@ -0,0 +1,3 @@
|
||||
allowBuilds:
|
||||
esbuild: true
|
||||
vue-demi: true
|
||||
@@ -27,6 +27,7 @@ function onSubmit() {
|
||||
<v-text-field
|
||||
:model-value="props.username"
|
||||
:label="t('username')"
|
||||
autocomplete="username"
|
||||
class="mb-6 input-field"
|
||||
required
|
||||
hide-details="auto"
|
||||
@@ -40,6 +41,7 @@ function onSubmit() {
|
||||
<v-text-field
|
||||
:model-value="props.password"
|
||||
:label="t('password')"
|
||||
autocomplete="current-password"
|
||||
required
|
||||
variant="outlined"
|
||||
hide-details="auto"
|
||||
|
||||
3
main.py
3
main.py
@@ -54,9 +54,6 @@ def check_env() -> None:
|
||||
|
||||
site_packages_path = get_astrbot_site_packages_path()
|
||||
if not is_packaged_desktop_runtime() and site_packages_path not in sys.path:
|
||||
# Packaged desktop runtime keeps shared plugin dependencies out of the
|
||||
# global import path so bundled core libraries don't mix with user-
|
||||
# installed wheels from ~/.astrbot/data/site-packages.
|
||||
sys.path.append(site_packages_path)
|
||||
|
||||
os.makedirs(get_astrbot_config_path(), exist_ok=True)
|
||||
|
||||
@@ -476,6 +476,46 @@ class TestBuiltinToolInjection:
|
||||
assert req.func_tool.get_tool("web_search_firecrawl") is search_tool
|
||||
assert req.func_tool.get_tool("firecrawl_extract_web_page") is extract_tool
|
||||
|
||||
def test_apply_web_search_citation_prompt_for_webchat(self, mock_event):
|
||||
module = ama
|
||||
req = ProviderRequest(system_prompt="base")
|
||||
search_tool = MagicMock(spec=FunctionTool)
|
||||
search_tool.name = "web_search_tavily"
|
||||
req.func_tool = ToolSet()
|
||||
req.func_tool.add_tool(search_tool)
|
||||
mock_event.get_platform_name.return_value = "webchat"
|
||||
|
||||
module._apply_web_search_citation_prompt(mock_event, req)
|
||||
|
||||
assert module.WEB_SEARCH_CITATION_PROMPT in req.system_prompt
|
||||
|
||||
def test_apply_web_search_citation_prompt_is_idempotent(self, mock_event):
|
||||
module = ama
|
||||
req = ProviderRequest(system_prompt="")
|
||||
search_tool = MagicMock(spec=FunctionTool)
|
||||
search_tool.name = "web_search_tavily"
|
||||
req.func_tool = ToolSet()
|
||||
req.func_tool.add_tool(search_tool)
|
||||
mock_event.get_platform_name.return_value = "webchat"
|
||||
|
||||
module._apply_web_search_citation_prompt(mock_event, req)
|
||||
module._apply_web_search_citation_prompt(mock_event, req)
|
||||
|
||||
assert req.system_prompt.count(module.WEB_SEARCH_CITATION_PROMPT) == 1
|
||||
|
||||
def test_apply_web_search_citation_prompt_requires_webchat(self, mock_event):
|
||||
module = ama
|
||||
req = ProviderRequest(system_prompt="")
|
||||
search_tool = MagicMock(spec=FunctionTool)
|
||||
search_tool.name = "web_search_tavily"
|
||||
req.func_tool = ToolSet()
|
||||
req.func_tool.add_tool(search_tool)
|
||||
mock_event.get_platform_name.return_value = "test_platform"
|
||||
|
||||
module._apply_web_search_citation_prompt(mock_event, req)
|
||||
|
||||
assert module.WEB_SEARCH_CITATION_PROMPT not in req.system_prompt
|
||||
|
||||
def test_proactive_cron_job_tools_uses_builtin_tool_manager(self, mock_context):
|
||||
"""Test cron tool injection through the builtin tool manager."""
|
||||
module = ama
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage
|
||||
|
||||
@@ -102,38 +101,3 @@ async def test_document_storage_fts_recovers_from_legacy_non_fts_table(tmp_path)
|
||||
assert [result["doc_id"] for result in results] == ["legacy-fix"]
|
||||
|
||||
await storage.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_storage_adds_unique_doc_id_index_to_existing_table(tmp_path):
|
||||
db_path = tmp_path / "doc.db"
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE documents (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
doc_id VARCHAR NOT NULL,
|
||||
text VARCHAR NOT NULL,
|
||||
metadata TEXT,
|
||||
created_at DATETIME,
|
||||
updated_at DATETIME
|
||||
)
|
||||
""",
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO documents (doc_id, text) VALUES ('legacy-chunk', 'legacy text')"
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
storage = DocumentStorage(str(db_path))
|
||||
await storage.initialize()
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
await storage.insert_document(
|
||||
doc_id="legacy-chunk",
|
||||
text="duplicate text",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
await storage.close()
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlmodel import select
|
||||
|
||||
from astrbot.core.agent.response import AgentStats
|
||||
@@ -65,143 +63,3 @@ async def test_record_internal_agent_stats_persists_provider_stat(
|
||||
assert record.start_time == 100.0
|
||||
assert record.end_time == 108.5
|
||||
assert record.time_to_first_token == 0.6
|
||||
|
||||
|
||||
def _provider_stats_recording_args():
|
||||
event = SimpleNamespace(unified_msg_origin="webchat:FriendMessage:session-42")
|
||||
req = ProviderRequest(conversation=SimpleNamespace(cid="conv-123"))
|
||||
provider = SimpleNamespace(
|
||||
provider_config={"id": "provider-1"},
|
||||
meta=lambda: SimpleNamespace(id="provider-1", type="openai"),
|
||||
get_model=lambda: "gpt-4.1",
|
||||
)
|
||||
agent_runner = SimpleNamespace(
|
||||
provider=provider,
|
||||
stats=AgentStats(),
|
||||
was_aborted=lambda: False,
|
||||
)
|
||||
return event, req, agent_runner, SimpleNamespace(role="assistant")
|
||||
|
||||
|
||||
def _provider_stats_operational_error(message: str) -> OperationalError:
|
||||
return OperationalError("insert into provider_stats", {}, Exception(message))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"lock_message",
|
||||
["database is locked", "database table is locked"],
|
||||
)
|
||||
async def test_record_internal_agent_stats_retries_transient_database_locks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
lock_message: str,
|
||||
):
|
||||
attempts = 0
|
||||
|
||||
class LockedOnceDb:
|
||||
async def insert_provider_stat(self, **kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
if attempts == 1:
|
||||
raise _provider_stats_operational_error(lock_message)
|
||||
return SimpleNamespace(**kwargs)
|
||||
|
||||
monkeypatch.setattr(internal, "db_helper", LockedOnceDb())
|
||||
|
||||
async def no_sleep(delay: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(internal.asyncio, "sleep", no_sleep)
|
||||
|
||||
await internal._record_internal_agent_stats(
|
||||
*_provider_stats_recording_args(),
|
||||
)
|
||||
|
||||
assert attempts == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_internal_agent_stats_logs_after_exhausting_database_lock_retries(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
attempts = 0
|
||||
sleep_delays = []
|
||||
warnings = []
|
||||
|
||||
class AlwaysLockedDb:
|
||||
async def insert_provider_stat(self, **kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
raise _provider_stats_operational_error("database is locked")
|
||||
|
||||
monkeypatch.setattr(internal, "db_helper", AlwaysLockedDb())
|
||||
|
||||
async def record_sleep(delay: float) -> None:
|
||||
sleep_delays.append(delay)
|
||||
|
||||
monkeypatch.setattr(internal.asyncio, "sleep", record_sleep)
|
||||
monkeypatch.setattr(
|
||||
internal.logger,
|
||||
"warning",
|
||||
lambda *args, **kwargs: warnings.append((args, kwargs)),
|
||||
)
|
||||
|
||||
await internal._record_internal_agent_stats(*_provider_stats_recording_args())
|
||||
|
||||
assert attempts == internal.PROVIDER_STATS_SQLITE_LOCK_RETRY_ATTEMPTS
|
||||
base_delay = internal.PROVIDER_STATS_SQLITE_LOCK_RETRY_BASE_DELAY
|
||||
expected_sleep_delays = [
|
||||
base_delay * (2**attempt)
|
||||
for attempt in range(internal.PROVIDER_STATS_SQLITE_LOCK_RETRY_ATTEMPTS - 1)
|
||||
]
|
||||
assert sleep_delays == expected_sleep_delays
|
||||
assert len(warnings) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_internal_agent_stats_does_not_retry_other_operational_errors(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
attempts = 0
|
||||
warnings = []
|
||||
|
||||
class FailingDb:
|
||||
async def insert_provider_stat(self, **kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
raise _provider_stats_operational_error("no such table: provider_stats")
|
||||
|
||||
monkeypatch.setattr(internal, "db_helper", FailingDb())
|
||||
monkeypatch.setattr(
|
||||
internal.logger,
|
||||
"warning",
|
||||
lambda *args, **kwargs: warnings.append((args, kwargs)),
|
||||
)
|
||||
|
||||
await internal._record_internal_agent_stats(*_provider_stats_recording_args())
|
||||
|
||||
assert attempts == 1
|
||||
assert len(warnings) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_internal_agent_stats_propagates_cancelled_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
warnings = []
|
||||
|
||||
class CancellingDb:
|
||||
async def insert_provider_stat(self, **kwargs):
|
||||
raise asyncio.CancelledError
|
||||
|
||||
monkeypatch.setattr(internal, "db_helper", CancellingDb())
|
||||
monkeypatch.setattr(
|
||||
internal.logger,
|
||||
"warning",
|
||||
lambda *args, **kwargs: warnings.append((args, kwargs)),
|
||||
)
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await internal._record_internal_agent_stats(*_provider_stats_recording_args())
|
||||
|
||||
assert warnings == []
|
||||
|
||||
Reference in New Issue
Block a user