services/vibe-bot
2
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity

complete restructure

Browse Source
This commit is contained in:
ducoterra
2026-08-19 13:16:43 -04:00
parent d7b6f28cbd
commit f87e1d51ef
60 changed files with 8176 additions and 5143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vibe_bot/tests/test_database.py
+774 -44
View File
@@ -2,17 +2,40 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
if TYPE_CHECKING:
import sqlite3
from vibe_bot.database import ChatDatabase
def _recent_messages(
db_path: str,
limit: int,
) -> list[tuple[str, str, str, datetime]]:
"""Read the newest rows straight from the database (test-side helper)."""
from vibe_bot.db.connection import connect
conn = connect(db_path)
try:
rows = conn.execute(
"SELECT message_id, username, content, timestamp "
"FROM chat_messages ORDER BY timestamp DESC LIMIT ?",
(limit,),
).fetchall()
finally:
conn.close()
return [
(str(row[0]), str(row[1]), str(row[2]), cast(datetime, row[3])) for row in rows
]
def test_vector_to_bytes(chat_db: ChatDatabase) -> None:
"""Test converting a vector to bytes and back."""
vector: list[float] = [0.1, 0.2, 0.3, 0.4]
@@ -55,6 +78,13 @@ def test_calculate_similarity_negative(chat_db: ChatDatabase) -> None:
assert similarity == pytest.approx(-1.0, abs=1e-6)
def test_calculate_similarity_zero_norm(chat_db: ChatDatabase) -> None:
"""A zero vector has no direction, so its similarity is 0."""
zero = np.zeros(3, dtype=np.float32)
other = np.array([1.0, 0.0, 0.0], dtype=np.float32)
assert chat_db._calculate_similarity(zero, other) == 0.0
def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
"""Test adding a message to the database."""
result = chat_db.add_message(
@@ -67,7 +97,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
)
assert result is True
messages = chat_db.get_recent_messages(limit=10)
messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 1
assert messages[0][0] == "msg-1"
assert messages[0][1] == "testuser"
@@ -76,7 +106,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
def test_add_message_no_embedding(chat_db: ChatDatabase) -> None:
"""Test adding a message when embedding generation fails."""
with patch("vibe_bot.llama_wrapper.embedding", return_value=None):
with patch("vibe_bot.llm_client.embedding", return_value=None):
result = chat_db.add_message(
message_id="msg-no-embed",
user_id="user-1",
@@ -106,14 +136,14 @@ def test_add_message_duplicate(
content="Second content",
)
messages = chat_db.get_recent_messages(limit=10)
messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 1
assert messages[0][2] == "Second content"
def test_add_message_failure(chat_db: ChatDatabase) -> None:
"""Test that add_message returns False on database error."""
with patch.object(chat_db, "_vector_to_bytes", side_effect=Exception("fail")):
with patch("vibe_bot.db.messages.connect", return_value=_broken_connection()):
result = chat_db.add_message(
message_id="msg-fail",
user_id="user-1",
@@ -123,11 +153,304 @@ def test_add_message_failure(chat_db: ChatDatabase) -> None:
assert result is False
def test_get_recent_messages(
def _embedding_row_count(db_path: str) -> int:
"""Count the rows stored in message_embeddings."""
import sqlite3
conn = sqlite3.connect(db_path)
count = conn.execute("SELECT COUNT(*) FROM message_embeddings").fetchone()[0]
conn.close()
return int(count)
def test_add_message_embed_false_skips_embedding(chat_db: ChatDatabase) -> None:
"""embed=False neither calls the embedding API nor stores a row."""
with patch("vibe_bot.llm_client.embedding") as mock_embedding:
result = chat_db.add_message(
message_id="msg-assist",
user_id="bot-1",
username="some-bot",
content="assistant reply",
role="assistant",
embed=False,
)
assert result is True
mock_embedding.assert_not_called()
assert _embedding_row_count(chat_db.db_path) == 0
def test_add_message_stores_embedding_by_default(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Test retrieving recent messages."""
"""The default embed=True still stores exactly one embedding row."""
assert chat_db.add_message(
message_id="msg-embed",
user_id="user-1",
username="testuser",
content="default embed",
)
assert _embedding_row_count(chat_db.db_path) == 1
def test_add_message_stores_embedding_norm(chat_db: ChatDatabase) -> None:
"""add_message stores the L2 norm of the float32-encoded embedding."""
import sqlite3
vector: list[float] = [0.6, 0.8]
with patch("vibe_bot.llm_client.embedding", return_value=vector):
assert chat_db.add_message(
message_id="norm-1",
user_id="u1",
username="alice",
content="normed message",
)
conn = sqlite3.connect(chat_db.db_path)
blob, norm = conn.execute(
"SELECT embedding, norm FROM message_embeddings WHERE message_id = 'norm-1'"
).fetchone()
conn.close()
stored = np.frombuffer(blob, dtype=np.float32)
assert float(norm) == pytest.approx(float(np.linalg.norm(stored)), abs=1e-6)
def test_cleanup_old_messages_no_orphaned_embeddings(chat_db: ChatDatabase) -> None:
"""Deleting the oldest rows must delete their embeddings, not the next ones.
Regression test for the bug where the embedding cleanup re-queried
``chat_messages`` *after* the message delete, so it stripped embeddings
from the next-oldest live rows while leaving the deleted rows' embeddings
behind as orphans.
"""
import sqlite3
with patch("vibe_bot.db.messages.MAX_HISTORY_MESSAGES", 5):
# Seed 7 rows with distinct ascending timestamps and an embedding for
# each, so the "oldest" ordering is deterministic.
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
for i in range(1, 8):
ts = f"2024-01-0{i} 00:00:00"
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, bot_name, timestamp) "
"VALUES (?, ?, ?, ?, ?, ?)",
(f"m{i}", "u1", "alice", f"content {i}", "bot", ts),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding) "
"VALUES (?, ?)",
(f"m{i}", chat_db._vector_to_bytes([0.1, 0.2, 0.3])),
)
conn.commit()
conn.close()
# The 8th insert pushes the count to 8 (> 5), so add_message's
# cleanup must delete exactly the 3 oldest rows and their embeddings.
assert chat_db.add_message(
message_id="m8",
user_id="u1",
username="alice",
content="content 8",
)
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
cursor.execute("SELECT message_id FROM chat_messages")
live = {row[0] for row in cursor.fetchall()}
cursor.execute("SELECT message_id FROM message_embeddings")
embedded = {row[0] for row in cursor.fetchall()}
conn.close()
# The three oldest rows are gone; the rest (incl. the new one) remain.
assert live == {"m4", "m5", "m6", "m7", "m8"}
# No orphaned embeddings and no stripped survivors: the embedding set
# exactly matches the live message rows.
assert embedded == live
def test_role_scopes_history_and_search(chat_db: ChatDatabase) -> None:
"""get_user_history excludes responses; search matches only user rows."""
chat_db.add_message(
message_id="r-1",
user_id="u1",
username="alice",
content="User asks about the weather",
role="user",
)
chat_db.add_message(
message_id="r-1_response",
user_id="bot",
username="some-bot",
content="Bot answers the weather",
role="assistant",
)
# get_user_history returns only the user row (paired with its response).
conversations = chat_db.get_user_history("u1")
assert len(conversations) == 1
assert conversations[0][0] == "User asks about the weather"
assert conversations[0][1] == "Bot answers the weather"
# search_similar_messages only considers user rows, never responses.
results = chat_db.search_similar_messages(
"User asks about the weather", top_k=5, min_similarity=0.0
)
assert len(results) == 1
assert results[0][0] == "User asks about the weather"
assert results[0][1] == "Bot answers the weather"
def test_role_migration_backfills_legacy_rows(
temp_db_path: str,
) -> None:
"""Legacy rows (no role column) are backfilled when ChatDatabase inits."""
import sqlite3
from vibe_bot.database import ChatDatabase
# Create the legacy schema (no role column) with a user row and its
# response row inserted directly.
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT)"
)
cursor.execute(
"INSERT INTO chat_messages (message_id, user_id, username, content) "
"VALUES ('legacy-1', 'u1', 'alice', 'old question')"
)
cursor.execute(
"INSERT INTO chat_messages (message_id, user_id, username, content) "
"VALUES ('legacy-1_response', 'bot', 'old-bot', 'old answer')"
)
conn.commit()
conn.close()
# Initializing ChatDatabase should add and backfill the role column.
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("SELECT message_id, role FROM chat_messages ORDER BY message_id")
roles = {row[0]: row[1] for row in cursor.fetchall()}
conn.close()
assert roles["legacy-1"] == "user"
assert roles["legacy-1_response"] == "assistant"
def test_bot_name_migration_adds_column(
temp_db_path: str,
) -> None:
"""A pre-bot_name schema gets the column added when ChatDatabase inits."""
import sqlite3
from vibe_bot.database import ChatDatabase
conn = sqlite3.connect(temp_db_path)
conn.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT, role TEXT)"
)
conn.commit()
conn.close()
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(chat_messages)")
columns = {row[1] for row in cursor.fetchall()}
conn.close()
assert "bot_name" in columns
def _seed_pre_norm_db(db_path: str) -> list[tuple[str, str, list[float]]]:
"""Create a pre-norm-schema database (no norm column) with seeded rows."""
import sqlite3
rows: list[tuple[str, str, list[float]]] = [
("n-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-5", "ask about nothing", [0.0] * 8),
]
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT, role TEXT)"
)
cursor.execute(
"CREATE TABLE message_embeddings ("
"message_id TEXT PRIMARY KEY, embedding BLOB,"
"FOREIGN KEY (message_id) REFERENCES chat_messages(message_id))"
)
for i, (message_id, content, vector) in enumerate(rows):
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'u1', 'alice', ?, 'user', ?)",
(message_id, content, f"2024-01-0{i + 1} 00:00:00"),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', ?)",
(
f"{message_id}_response",
f"response {i + 1}",
f"2024-01-0{i + 1} 00:00:01",
),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding) VALUES (?, ?)",
(message_id, np.array(vector, dtype=np.float32).tobytes()),
)
conn.commit()
conn.close()
return rows
def test_norm_migration_backfills_stored_norms(temp_db_path: str) -> None:
"""A pre-norm database gets the norm column added and backfilled on init."""
import sqlite3
from vibe_bot.database import ChatDatabase
rows = _seed_pre_norm_db(temp_db_path)
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
stored = dict(conn.execute("SELECT message_id, norm FROM message_embeddings"))
conn.close()
for message_id, _content, vector in rows:
expected = float(np.linalg.norm(np.array(vector, dtype=np.float32)))
assert stored[message_id] == pytest.approx(expected, abs=1e-5)
def test_recent_messages_desc_order(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Newest-first ordering of stored messages."""
chat_db.add_message(
message_id="msg-1",
user_id="u1",
@@ -147,17 +470,17 @@ def test_get_recent_messages(
content="Third",
)
messages = chat_db.get_recent_messages(limit=2)
messages = _recent_messages(chat_db.db_path, 2)
assert len(messages) == 2
assert messages[0][2] == "Third"
assert messages[1][2] == "Second"
def test_get_recent_messages_limit(
def test_recent_messages_limit(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Test that get_recent_messages respects the limit."""
"""The newest-rows query respects the limit."""
for i in range(5):
chat_db.add_message(
message_id=f"msg-{i}",
@@ -166,10 +489,25 @@ def test_get_recent_messages_limit(
content=f"Message {i}",
)
messages = chat_db.get_recent_messages(limit=3)
messages = _recent_messages(chat_db.db_path, 3)
assert len(messages) == 3
def test_recent_messages_returns_datetime(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""The timestamp column comes back as a real datetime, not a string."""
chat_db.add_message(
message_id="dt-1",
user_id="u1",
username="alice",
content="fresh message",
)
messages = _recent_messages(chat_db.db_path, 1)
assert isinstance(messages[0][3], datetime)
def test_clear_all_messages(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
@@ -190,7 +528,7 @@ def test_clear_all_messages(
chat_db.clear_all_messages()
messages = chat_db.get_recent_messages(limit=10)
messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 0
@@ -264,6 +602,25 @@ def test_image_generation_estimate_after_capping(
assert estimate == pytest.approx(99.5)
def _broken_connection() -> MagicMock:
"""A mock connection whose first cursor.execute raises."""
fake_conn = MagicMock()
fake_conn.cursor.return_value.execute.side_effect = Exception("db error")
return fake_conn
def test_record_image_generation_time_failure(chat_db: ChatDatabase) -> None:
"""A database error while recording yields False, not an exception."""
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
assert chat_db.record_image_generation_time(1.0) is False
def test_image_generation_estimate_failure(chat_db: ChatDatabase) -> None:
"""A database error while reading yields None, not an exception."""
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
assert chat_db.get_image_generation_time_estimate() is None
def test_get_user_history(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
@@ -320,6 +677,39 @@ def test_get_user_history_excludes_bot(
assert len(conversations) == 0
def test_get_bot_history(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""get_bot_history pairs user messages with responses for one bot."""
chat_db.add_message(
message_id="bh-1",
user_id="u1",
username="alice",
content="bot question",
bot_name="alfred",
)
chat_db.add_message(
message_id="bh-1_response",
user_id="bot-1",
username="some-bot",
content="bot answer",
bot_name="alfred",
role="assistant",
embed=False,
)
chat_db.add_message(
message_id="bh-2",
user_id="u1",
username="alice",
content="unanswered question",
bot_name="alfred",
)
history = chat_db.get_bot_history("alfred")
assert history == [("bot question", "bot answer")]
def test_get_conversation_context(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
@@ -349,21 +739,333 @@ def test_get_conversation_context_empty(chat_db: ChatDatabase) -> None:
assert context == []
def _query_vector() -> list[float]:
"""A 8-dim query vector along the first axis."""
return [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
def _seed_search_rows(chat_db: ChatDatabase) -> list[tuple[str, str, list[float]]]:
"""Seed user/response rows with known embeddings (plus exclusion traps)."""
import sqlite3
rows: list[tuple[str, str, list[float]]] = [
("s-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
]
def stored_norm(vector: list[float]) -> float:
return float(np.linalg.norm(np.array(vector, dtype=np.float32)))
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
for i, (message_id, content, vector) in enumerate(rows):
timestamp = f"2024-01-0{i + 1} 00:00:00"
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, ?, ?, ?, 'user', ?)",
(message_id, "u1", "alice", content, timestamp),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, ?, ?, ?, 'assistant', ?)",
(
f"{message_id}_response",
"bot-1",
"some-bot",
f"response {i + 1}",
f"2024-01-0{i + 1} 00:00:01",
),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(message_id, chat_db._vector_to_bytes(vector), stored_norm(vector)),
)
# A user row with the top possible similarity but no response row: the
# JOIN must exclude it instead of returning a NULL response.
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES ('s-5', 'u1', 'alice', 'orphan question', 'user', "
"'2024-01-05 00:00:00')",
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(
"s-5",
chat_db._vector_to_bytes(_query_vector()),
stored_norm(_query_vector()),
),
)
# An assistant row carrying an embedding: the role filter must skip it.
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES ('s-6', 'bot-1', 'some-bot', 'assistant noise', 'assistant', "
"'2024-01-06 00:00:00')",
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(
"s-6",
chat_db._vector_to_bytes(_query_vector()),
stored_norm(_query_vector()),
),
)
conn.commit()
conn.close()
return rows
def test_search_matches_reference_topk_and_ordering(
chat_db: ChatDatabase,
) -> None:
"""The JOINed search matches the per-row reference: same top-k and order."""
rows = _seed_search_rows(chat_db)
query = _query_vector()
expected: list[tuple[str, str, float]] = []
query_arr = np.array(query, dtype=np.float32)
for i, (_message_id, content, vector) in enumerate(rows):
stored = np.frombuffer(chat_db._vector_to_bytes(vector), dtype=np.float32)
similarity = float(
np.dot(query_arr, stored)
/ (np.linalg.norm(query_arr) * np.linalg.norm(stored))
)
expected.append((content, f"response {i + 1}", similarity))
expected.sort(key=lambda item: item[2], reverse=True)
expected_top = expected[:3]
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = chat_db.search_similar_messages(
"query text", top_k=3, min_similarity=0.0
)
assert len(results) == 3
for (content, response, actual), (_ec, er, reference) in zip(
results, expected_top, strict=True
):
assert content == _ec
assert response == er
assert actual == pytest.approx(reference, abs=1e-5)
def test_search_with_stored_norms_matches_pre_norm_reference(
temp_db_path: str,
) -> None:
"""Stored-norm search is identical to per-vector renormalization.
Seeds a pre-norm database, migrates it, then asserts the stored-norm
search (top-k + ordering + similarities) matches a reference that
renormalizes every candidate vector inline — the pre-optimization
algorithm.
"""
from vibe_bot.database import ChatDatabase
rows = _seed_pre_norm_db(temp_db_path)
db = ChatDatabase(db_path=temp_db_path)
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=3, min_similarity=0.0)
query_arr = np.array(query, dtype=np.float32)
query_norm = float(np.linalg.norm(query_arr))
expected: list[tuple[str, str, float]] = []
for i, (_message_id, content, vector) in enumerate(rows):
stored = np.array(vector, dtype=np.float32)
stored_norm = float(np.linalg.norm(stored))
similarity = (
0.0
if stored_norm == 0
else float(np.dot(query_arr, stored) / (query_norm * stored_norm))
)
expected.append((content, f"response {i + 1}", similarity))
expected.sort(key=lambda item: item[2], reverse=True)
assert len(results) == 3
for (content, response, actual), (ec, er, reference) in zip(
results, expected[:3], strict=True
):
assert content == ec
assert response == er
assert actual == pytest.approx(reference, abs=1e-6)
def test_search_null_norm_row_scores_zero(temp_db_path: str) -> None:
"""A row whose norm was never written scores 0 instead of crashing."""
import sqlite3
from vibe_bot.database import ChatDatabase
_rows = _seed_pre_norm_db(temp_db_path)
db = ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
conn.execute("UPDATE message_embeddings SET norm = NULL WHERE message_id = 'n-1'")
conn.commit()
conn.close()
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
scores = {content: similarity for content, _response, similarity in results}
assert scores["ask about the sky"] == 0.0
# All four scored rows plus the zero-vector row (0.0 passes min_similarity=0.0).
assert len(results) == 5
def test_search_mixed_embedding_dims_falls_back_to_per_row(temp_db_path: str) -> None:
"""Rows with different blob lengths (mid-life model change) don't crash."""
import sqlite3
from vibe_bot.database import ChatDatabase
wide = [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
narrow = [0.5, 0.5, 0.0, 0.0]
db = ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
for message_id, content, vector in (
("m-1", "wide question", wide),
("m-2", "narrow question", narrow),
):
blob = np.array(vector, dtype=np.float32).tobytes()
stored_norm = float(np.linalg.norm(np.frombuffer(blob, dtype=np.float32)))
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'u1', 'alice', ?, 'user', '2024-01-01 00:00:00')",
(message_id, content),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', '2024-01-01 00:00:01')",
(f"{message_id}_response", f"{message_id} answer"),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(message_id, blob, stored_norm),
)
conn.commit()
conn.close()
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
scores = {content: similarity for content, _response, similarity in results}
# Both rows are returned: the 8-dim row scores exactly, the 4-dim row is
# zero-padded to the query dim and still scores sanely.
assert set(scores) == {"wide question", "narrow question"}
wide_norm = float(np.linalg.norm(np.array(wide, dtype=np.float32)))
narrow_norm = float(np.linalg.norm(np.array(narrow, dtype=np.float32)))
assert scores["wide question"] == pytest.approx(0.9 / wide_norm, abs=1e-5)
assert scores["narrow question"] == pytest.approx(0.5 / narrow_norm, abs=1e-5)
assert scores["wide question"] > scores["narrow question"]
def test_search_issues_single_select_per_call(chat_db: ChatDatabase) -> None:
"""search_similar_messages issues exactly one SELECT per call (no N+1)."""
import vibe_bot.db.search as db_search
from vibe_bot.db.connection import connect as realconnect
chat_db.add_message(
message_id="n-1",
user_id="u1",
username="alice",
content="one question",
)
chat_db.add_message(
message_id="n-1_response",
user_id="bot-1",
username="some-bot",
content="one answer",
role="assistant",
embed=False,
)
select_statements: list[str] = []
class _TracingCursor:
def __init__(self, cursor: sqlite3.Cursor) -> None:
self._cursor = cursor
def execute(self, sql: str, *args: Any) -> Any:
if sql.strip().upper().startswith("SELECT"):
select_statements.append(sql)
return self._cursor.execute(sql, *args)
def fetchall(self) -> Any:
return self._cursor.fetchall()
def fetchone(self) -> Any:
return self._cursor.fetchone()
class _TracingConnection:
def __init__(self, conn: sqlite3.Connection) -> None:
self._conn = conn
def cursor(self) -> _TracingCursor:
return _TracingCursor(self._conn.cursor())
def close(self) -> None:
self._conn.close()
def tracingconnect(db_path: str) -> _TracingConnection:
return _TracingConnection(realconnect(db_path))
with patch.object(db_search, "connect", side_effect=tracingconnect):
results = chat_db.search_similar_messages(
"one question", top_k=5, min_similarity=0.0
)
assert len(results) == 1
assert results[0][0] == "one question"
assert results[0][1] == "one answer"
assert len(select_statements) == 1
def test_search_empty_query_embedding_returns_empty(chat_db: ChatDatabase) -> None:
"""A failed query embedding yields no results."""
with patch("vibe_bot.llm_client.embedding", return_value=[]):
assert chat_db.search_similar_messages("anything") == []
def test_search_zero_query_vector_returns_empty(chat_db: ChatDatabase) -> None:
"""A zero-norm query vector yields no results (no division by zero)."""
with patch("vibe_bot.llm_client.embedding", return_value=[0.0] * 8):
assert chat_db.search_similar_messages("anything") == []
def test_custom_bot_create(custom_bot_manager: Any) -> None:
"""Test creating a custom bot."""
"""Test creating a custom bot returns "created" for a new name."""
result = custom_bot_manager.create_custom_bot(
bot_name="alfred",
system_prompt="You are a british butler",
created_by="user-123",
)
assert result is True
assert result == "created"
def test_custom_bot_create_duplicate(
custom_bot_manager: Any,
) -> None:
"""Test creating a duplicate custom bot replaces the old one."""
custom_bot_manager.create_custom_bot(
first = custom_bot_manager.create_custom_bot(
bot_name="alfred",
system_prompt="First personality",
created_by="user-1",
@@ -373,13 +1075,25 @@ def test_custom_bot_create_duplicate(
system_prompt="Second personality",
created_by="user-1",
)
assert result is True
assert first == "created"
assert result == "replaced"
bot = custom_bot_manager.get_custom_bot("alfred")
assert bot is not None
assert bot[1] == "Second personality"
def test_custom_bot_create_failure(custom_bot_manager: Any) -> None:
"""A database error while creating yields False, not an exception."""
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
result = custom_bot_manager.create_custom_bot(
bot_name="failbot",
system_prompt="a long enough personality",
created_by="user-1",
)
assert result is False
def test_custom_bot_create_case_insensitive(
custom_bot_manager: Any,
) -> None:
@@ -399,6 +1113,18 @@ def test_custom_bot_get_not_found(custom_bot_manager: Any) -> None:
assert result is None
def test_custom_bot_get_returns_datetime(custom_bot_manager: Any) -> None:
"""created_at comes back as a real datetime, not a string."""
custom_bot_manager.create_custom_bot(
bot_name="dtbot",
system_prompt="a long enough personality",
created_by="user-1",
)
result = custom_bot_manager.get_custom_bot("dtbot")
assert result is not None
assert isinstance(result[3], datetime)
def test_custom_bot_get_returns_correct_data(
custom_bot_manager: Any,
) -> None:
@@ -413,8 +1139,7 @@ def test_custom_bot_get_returns_correct_data(
assert result[0] == "testbot"
assert result[1] == "test prompt"
assert result[2] == "creator-1"
assert result[3] is not None
assert "20" in result[3]
assert isinstance(result[3], datetime)
def test_custom_bot_list_empty(custom_bot_manager: Any) -> None:
@@ -440,6 +1165,23 @@ def test_custom_bot_list(custom_bot_manager: Any) -> None:
assert len(bots) == 2
def test_custom_bot_list_by_creator(custom_bot_manager: Any) -> None:
"""list_custom_bots filters by creator when user_id is given."""
custom_bot_manager.create_custom_bot(
bot_name="bot-x",
system_prompt="prompt x",
created_by="user-1",
)
custom_bot_manager.create_custom_bot(
bot_name="bot-y",
system_prompt="prompt y",
created_by="user-2",
)
bots = custom_bot_manager.list_custom_bots(user_id="user-1")
assert [bot[0] for bot in bots] == ["bot-x"]
def test_custom_bot_delete(custom_bot_manager: Any) -> None:
"""Test deleting a custom bot."""
custom_bot_manager.create_custom_bot(
@@ -462,32 +1204,18 @@ def test_custom_bot_delete_nonexistent(
assert result is False
def test_custom_bot_deactivate(custom_bot_manager: Any) -> None:
"""Test deactivating a custom bot."""
custom_bot_manager.create_custom_bot(
bot_name="inactive-bot",
system_prompt="will be deactivated",
created_by="user-1",
)
result = custom_bot_manager.deactivate_custom_bot("inactive-bot")
assert result is True
bot = custom_bot_manager.get_custom_bot("inactive-bot")
assert bot is None
def test_custom_bot_deactivate_nonexistent(
custom_bot_manager: Any,
) -> None:
"""Test deactivating a non-existent bot returns False."""
result = custom_bot_manager.deactivate_custom_bot("nonexistent")
assert result is False
def test_custom_bot_delete_failure(custom_bot_manager: Any) -> None:
"""A database error while deleting yields False, not an exception."""
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
assert custom_bot_manager.delete_custom_bot("whatever") is False
def test_custom_bot_list_excludes_inactive(
custom_bot_manager: Any,
) -> None:
"""Test that list_custom_bots excludes deactivated bots."""
"""Test that list_custom_bots excludes bots with is_active = 0."""
import sqlite3
custom_bot_manager.create_custom_bot(
bot_name="active-bot",
system_prompt="stays active",
@@ -498,7 +1226,12 @@ def test_custom_bot_list_excludes_inactive(
system_prompt="should not appear",
created_by="user-1",
)
custom_bot_manager.deactivate_custom_bot("deactivated-bot")
conn = sqlite3.connect(custom_bot_manager.db_path)
conn.execute(
"UPDATE custom_bots SET is_active = 0 WHERE bot_name = 'deactivated-bot'"
)
conn.commit()
conn.close()
bots = custom_bot_manager.list_custom_bots()
assert len(bots) == 1
@@ -532,16 +1265,13 @@ def test_database_get_database_singleton(temp_db_path: str) -> None:
db2 = get_database()
assert db1 is db2
db1.client.close()
def test_database_init_creates_tables(temp_db_path: str) -> None:
"""Test that database initialization creates the expected tables."""
from vibe_bot.database import ChatDatabase, CustomBotManager
db = ChatDatabase(db_path=temp_db_path)
ChatDatabase(db_path=temp_db_path)
CustomBotManager(db_path=temp_db_path)
db.client.close()
import sqlite3