complete restructure
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""SQLite storage layer: connection, schema, message store, RAG, custom bots."""
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Custom bot configuration store (create, read, list, delete)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
||||
from vibe_bot.config import DB_PATH
|
||||
from vibe_bot.db.connection import connect
|
||||
from vibe_bot.db.schema import initialize_custom_bots_table
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CustomBotManager:
|
||||
"""Manages custom bot configurations stored in SQLite database."""
|
||||
|
||||
def __init__(self, db_path: str = DB_PATH) -> None:
|
||||
"""Initialize the custom bot manager.
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file.
|
||||
|
||||
"""
|
||||
self.db_path = db_path
|
||||
self._initialize_custom_bots_table()
|
||||
|
||||
def _initialize_custom_bots_table(self) -> None:
|
||||
"""Initialize the custom bots table in SQLite."""
|
||||
initialize_custom_bots_table(self.db_path)
|
||||
|
||||
def create_custom_bot(
|
||||
self,
|
||||
bot_name: str,
|
||||
system_prompt: str,
|
||||
created_by: str,
|
||||
) -> str | bool:
|
||||
"""Create a custom bot configuration.
|
||||
|
||||
Returns "created" if the name was new, "replaced" if a bot with that
|
||||
name already existed (the namespace is shared), or False on error.
|
||||
"""
|
||||
conn = connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
name = bot_name.lower()
|
||||
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT 1 FROM custom_bots WHERE bot_name = ?",
|
||||
(name,),
|
||||
)
|
||||
exists = cursor.fetchone() is not None
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO custom_bots
|
||||
(bot_name, system_prompt, created_by, is_active)
|
||||
VALUES (?, ?, ?, 1)
|
||||
""",
|
||||
(name, system_prompt, created_by),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
except Exception:
|
||||
logger.exception("Error creating custom bot")
|
||||
conn.rollback()
|
||||
return False
|
||||
else:
|
||||
return "replaced" if exists else "created"
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_custom_bot(self, bot_name: str) -> tuple[str, str, str, datetime] | None:
|
||||
"""Get a custom bot configuration by name."""
|
||||
conn = connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT bot_name, system_prompt, created_by, created_at
|
||||
FROM custom_bots
|
||||
WHERE bot_name = ? AND is_active = 1
|
||||
""",
|
||||
(bot_name.lower(),),
|
||||
)
|
||||
|
||||
result = cursor.fetchone()
|
||||
conn.close()
|
||||
|
||||
if result is None:
|
||||
return None
|
||||
return (result[0], result[1], result[2], result[3])
|
||||
|
||||
def list_custom_bots(
|
||||
self,
|
||||
user_id: str | None = None,
|
||||
) -> list[tuple[str, str, str]]:
|
||||
"""List all custom bots, optionally filtered by creator."""
|
||||
conn = connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
if user_id:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT bot_name, system_prompt, created_by
|
||||
FROM custom_bots
|
||||
WHERE is_active = 1 AND created_by = ?
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
(user_id,),
|
||||
)
|
||||
else:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT bot_name, system_prompt, created_by
|
||||
FROM custom_bots
|
||||
WHERE is_active = 1
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
)
|
||||
|
||||
bots = cursor.fetchall()
|
||||
conn.close()
|
||||
|
||||
return bots
|
||||
|
||||
def delete_custom_bot(self, bot_name: str) -> bool:
|
||||
"""Delete a custom bot configuration."""
|
||||
conn = connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
DELETE FROM custom_bots
|
||||
WHERE bot_name = ?
|
||||
""",
|
||||
(bot_name.lower(),),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
except Exception:
|
||||
logger.exception("Error deleting custom bot")
|
||||
conn.rollback()
|
||||
return False
|
||||
else:
|
||||
return cursor.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
@@ -0,0 +1,36 @@
|
||||
"""SQLite connection plumbing shared by the database layer."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
|
||||
# Per-connection busy timeout (ms) so concurrent writers wait instead of
|
||||
# failing with "database is locked".
|
||||
SQLITE_BUSY_TIMEOUT_MS = 5000
|
||||
|
||||
|
||||
def _parse_timestamp(value: str | bytes) -> datetime:
|
||||
"""Decode a stored TIMESTAMP value into a naive datetime."""
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode("utf-8")
|
||||
return datetime.fromisoformat(value)
|
||||
|
||||
|
||||
sqlite3.register_converter("TIMESTAMP", _parse_timestamp)
|
||||
|
||||
|
||||
def connect(db_path: str) -> sqlite3.Connection:
|
||||
"""Open a SQLite connection configured for concurrent access.
|
||||
|
||||
WAL journaling is persistent (set once per database file); the busy
|
||||
timeout is per-connection, so it is applied on every connection here.
|
||||
``PARSE_DECLTYPES`` plus the registered ``TIMESTAMP`` converter decode
|
||||
declared ``TIMESTAMP`` columns into ``datetime`` objects instead of
|
||||
raw strings.
|
||||
|
||||
"""
|
||||
conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute(f"PRAGMA busy_timeout={SQLITE_BUSY_TIMEOUT_MS}")
|
||||
return conn
|
||||
@@ -0,0 +1,306 @@
|
||||
"""Chat message store: persistence, cleanup, and recency queries."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sqlite3
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vibe_bot import llm_client
|
||||
from vibe_bot.config import (
|
||||
DB_PATH,
|
||||
EMBEDDING_ENDPOINT,
|
||||
EMBEDDING_ENDPOINT_KEY,
|
||||
EMBEDDING_MODEL,
|
||||
MAX_HISTORY_MESSAGES,
|
||||
SIMILARITY_THRESHOLD,
|
||||
TOP_K_RESULTS,
|
||||
)
|
||||
from vibe_bot.db.connection import connect
|
||||
from vibe_bot.db.schema import initialize_chat_tables
|
||||
from vibe_bot.db.search import (
|
||||
get_bot_history,
|
||||
get_user_history,
|
||||
search_similar_messages,
|
||||
)
|
||||
from vibe_bot.db.timing import (
|
||||
get_image_generation_time_estimate,
|
||||
record_image_generation_time,
|
||||
)
|
||||
from vibe_bot.db.vectors import (
|
||||
bytes_to_vector,
|
||||
cosine_similarity,
|
||||
vector_to_bytes,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChatDatabase:
|
||||
"""SQLite store for chat history, embedding-backed RAG, and image timing."""
|
||||
|
||||
def __init__(self, db_path: str = DB_PATH) -> None:
|
||||
"""Initialize the database connection.
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file.
|
||||
|
||||
"""
|
||||
logger.info("Initializing ChatDatabase with path: %s", db_path)
|
||||
self.db_path = db_path
|
||||
initialize_chat_tables(db_path)
|
||||
|
||||
def _vector_to_bytes(self, vector: list[float]) -> bytes:
|
||||
"""Convert vector to bytes for SQLite storage."""
|
||||
return vector_to_bytes(vector)
|
||||
|
||||
def _bytes_to_vector(self, blob: bytes) -> np.ndarray:
|
||||
"""Convert bytes back to a vector."""
|
||||
return bytes_to_vector(blob)
|
||||
|
||||
def _calculate_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
|
||||
"""Calculate cosine similarity between two vectors."""
|
||||
return cosine_similarity(vec1, vec2)
|
||||
|
||||
def add_message(
|
||||
self,
|
||||
*,
|
||||
message_id: str,
|
||||
user_id: str,
|
||||
username: str,
|
||||
content: str,
|
||||
bot_name: str | None = None,
|
||||
channel_id: str | None = None,
|
||||
guild_id: str | None = None,
|
||||
role: str = "user",
|
||||
embed: bool = True,
|
||||
) -> bool:
|
||||
"""Add a message to the database, optionally storing its embedding.
|
||||
|
||||
Args:
|
||||
role: Either "user" (a human message) or "assistant" (a bot
|
||||
response). Used to scope RAG retrieval instead of matching a
|
||||
hard-coded bot username.
|
||||
embed: Whether to generate and store an embedding for the message.
|
||||
Response rows pass False: only user rows feed RAG retrieval.
|
||||
|
||||
"""
|
||||
logger.debug("Adding message %s from user %s", message_id, user_id)
|
||||
conn = connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
logger.debug(
|
||||
"Inserting message into chat_messages table: message_id=%s",
|
||||
message_id,
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO chat_messages
|
||||
(message_id, user_id, username, content, bot_name, channel_id,
|
||||
guild_id, role)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
message_id,
|
||||
user_id,
|
||||
username,
|
||||
content,
|
||||
bot_name,
|
||||
channel_id,
|
||||
guild_id,
|
||||
role,
|
||||
),
|
||||
)
|
||||
logger.debug("Message %s inserted into chat_messages table", message_id)
|
||||
|
||||
if embed:
|
||||
logger.debug("Generating embedding for message %s", message_id)
|
||||
embedding = llm_client.embedding(
|
||||
content,
|
||||
model=EMBEDDING_MODEL,
|
||||
url=EMBEDDING_ENDPOINT,
|
||||
api_key=EMBEDDING_ENDPOINT_KEY,
|
||||
)
|
||||
if embedding:
|
||||
logger.debug(
|
||||
"Embedding generated successfully for message %s, "
|
||||
"storing in database",
|
||||
message_id,
|
||||
)
|
||||
vector = np.array(embedding, dtype=np.float32)
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT OR REPLACE INTO message_embeddings
|
||||
(message_id, embedding, norm)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(message_id, vector.tobytes(), float(np.linalg.norm(vector))),
|
||||
)
|
||||
logger.debug(
|
||||
"Embedding stored in message_embeddings table for message %s",
|
||||
message_id,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Failed to generate embedding for message %s, "
|
||||
"skipping embedding storage",
|
||||
message_id,
|
||||
)
|
||||
|
||||
logger.debug("Checking if cleanup of old messages is needed")
|
||||
self._cleanup_old_messages(cursor)
|
||||
|
||||
conn.commit()
|
||||
except Exception:
|
||||
logger.exception("Error adding message %s", message_id)
|
||||
conn.rollback()
|
||||
return False
|
||||
else:
|
||||
logger.debug("Successfully added message %s to database", message_id)
|
||||
return True
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _cleanup_old_messages(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Remove old messages to stay within the limit.
|
||||
|
||||
The rows to delete are captured up front. Deriving the embedding
|
||||
message_ids from a fresh subquery *after* the chat_messages delete
|
||||
would select the next-oldest live rows instead of the ones just
|
||||
removed, orphaning the deleted rows' embeddings and deleting the
|
||||
embeddings of rows that should survive.
|
||||
|
||||
"""
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT COUNT(*) FROM chat_messages
|
||||
""",
|
||||
)
|
||||
count = cursor.fetchone()[0]
|
||||
|
||||
if count <= MAX_HISTORY_MESSAGES:
|
||||
return
|
||||
|
||||
excess = count - MAX_HISTORY_MESSAGES
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT id, message_id FROM chat_messages
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT ?
|
||||
""",
|
||||
(excess,),
|
||||
)
|
||||
oldest = cursor.fetchall()
|
||||
if not oldest:
|
||||
return
|
||||
|
||||
row_ids = [row[0] for row in oldest]
|
||||
# Include each row's `_response` companion so a deleted user message
|
||||
# also sheds its response embedding (and vice versa).
|
||||
message_ids: list[str] = []
|
||||
for _id, message_id in oldest:
|
||||
message_ids.append(message_id)
|
||||
message_ids.append(f"{message_id}_response")
|
||||
|
||||
id_placeholders = ", ".join("?" for _ in row_ids)
|
||||
cursor.execute(
|
||||
f"DELETE FROM chat_messages WHERE id IN ({id_placeholders})",
|
||||
row_ids,
|
||||
)
|
||||
mid_placeholders = ", ".join("?" for _ in message_ids)
|
||||
cursor.execute(
|
||||
f"DELETE FROM message_embeddings WHERE message_id IN ({mid_placeholders})",
|
||||
message_ids,
|
||||
)
|
||||
|
||||
def search_similar_messages(
|
||||
self,
|
||||
query: str,
|
||||
top_k: int = TOP_K_RESULTS,
|
||||
min_similarity: float = SIMILARITY_THRESHOLD,
|
||||
) -> list[tuple[str, str, float]]:
|
||||
"""Search for messages similar to the query using embeddings."""
|
||||
return search_similar_messages(
|
||||
self.db_path,
|
||||
query,
|
||||
top_k=top_k,
|
||||
min_similarity=min_similarity,
|
||||
)
|
||||
|
||||
def get_bot_history(self, bot_name: str, limit: int = 20) -> list[tuple[str, str]]:
|
||||
"""Get message history for a specific custom bot.
|
||||
|
||||
Args:
|
||||
bot_name: The name of the custom bot.
|
||||
limit: Maximum number of messages to retrieve.
|
||||
|
||||
Returns:
|
||||
List of (user_message, bot_response) tuples.
|
||||
|
||||
"""
|
||||
return get_bot_history(self.db_path, bot_name, limit)
|
||||
|
||||
def get_user_history(self, user_id: str, limit: int = 20) -> list[tuple[str, str]]:
|
||||
"""Get message history for a specific user."""
|
||||
return get_user_history(self.db_path, user_id, limit)
|
||||
|
||||
def get_conversation_context(
|
||||
self,
|
||||
user_id: str,
|
||||
current_message: str,
|
||||
max_context: int = 5,
|
||||
) -> list[dict[str, str]]:
|
||||
"""Get relevant conversation context for RAG."""
|
||||
recent_messages = get_user_history(self.db_path, user_id, limit=max_context * 2)
|
||||
|
||||
similar_messages = search_similar_messages(
|
||||
self.db_path,
|
||||
current_message,
|
||||
top_k=max_context,
|
||||
)
|
||||
|
||||
context_parts: list[dict[str, str]] = []
|
||||
|
||||
for user_message, bot_message in recent_messages:
|
||||
context_parts.append({"role": "assistant", "content": bot_message})
|
||||
context_parts.append({"role": "user", "content": user_message})
|
||||
|
||||
for user_message, bot_message, _similarity in similar_messages:
|
||||
context_parts.append({"role": "assistant", "content": bot_message})
|
||||
context_parts.append({"role": "user", "content": user_message})
|
||||
|
||||
# Conversation history needs to be delivered in "newest context last" order
|
||||
context_parts.reverse()
|
||||
return context_parts
|
||||
|
||||
def clear_all_messages(self) -> None:
|
||||
"""Clear all messages and embeddings from the database."""
|
||||
conn = connect(self.db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("DELETE FROM message_embeddings")
|
||||
cursor.execute("DELETE FROM chat_messages")
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def record_image_generation_time(self, duration_seconds: float) -> bool:
|
||||
"""Record how long an image generation took.
|
||||
|
||||
Args:
|
||||
duration_seconds: Wall-clock seconds the generation took.
|
||||
|
||||
"""
|
||||
return record_image_generation_time(self.db_path, duration_seconds)
|
||||
|
||||
def get_image_generation_time_estimate(self) -> float | None:
|
||||
"""Get a moving-average estimate of image generation time.
|
||||
|
||||
Returns:
|
||||
The average duration in seconds over the most recent generations,
|
||||
or None if there is no generation history yet.
|
||||
|
||||
"""
|
||||
return get_image_generation_time_estimate(self.db_path)
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Schema creation and column migrations for the chat and custom-bot tables."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sqlite3
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vibe_bot.db.connection import connect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def initialize_chat_tables(db_path: str) -> None:
|
||||
"""Create (and migrate) the chat history and embedding tables."""
|
||||
logger.info("Initializing SQLite database at %s", db_path)
|
||||
conn = connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
logger.info("Creating chat_messages table if not exists")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS 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
|
||||
)
|
||||
""",
|
||||
)
|
||||
logger.info("chat_messages table initialized successfully")
|
||||
_migrate_chat_messages(cursor)
|
||||
|
||||
logger.info("Creating message_embeddings table if not exists")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS message_embeddings (
|
||||
message_id TEXT PRIMARY KEY,
|
||||
embedding BLOB,
|
||||
norm REAL,
|
||||
FOREIGN KEY (message_id) REFERENCES chat_messages(message_id)
|
||||
)
|
||||
""",
|
||||
)
|
||||
logger.info("message_embeddings table initialized successfully")
|
||||
_migrate_message_embeddings(cursor)
|
||||
|
||||
logger.info("Creating idx_timestamp index if not exists")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS idx_timestamp ON chat_messages(timestamp)
|
||||
""",
|
||||
)
|
||||
logger.info("idx_timestamp index created successfully")
|
||||
|
||||
logger.info("Creating idx_user_id index if not exists")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS idx_user_id ON chat_messages(user_id)
|
||||
""",
|
||||
)
|
||||
logger.info("idx_user_id index created successfully")
|
||||
|
||||
logger.info("Creating image_generation_times table if not exists")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS image_generation_times (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
duration_seconds REAL NOT NULL,
|
||||
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""",
|
||||
)
|
||||
logger.info("image_generation_times table initialized successfully")
|
||||
|
||||
conn.commit()
|
||||
logger.info("Database initialization completed successfully")
|
||||
conn.close()
|
||||
|
||||
|
||||
def _migrate_chat_messages(cursor: sqlite3.Cursor) -> None:
|
||||
"""Add the bot_name and role columns to pre-existing databases."""
|
||||
logger.info("Checking for chat_messages column migrations")
|
||||
cursor.execute("PRAGMA table_info(chat_messages)")
|
||||
columns = {row[1] for row in cursor.fetchall()}
|
||||
|
||||
if "bot_name" not in columns:
|
||||
logger.info("Adding bot_name column to chat_messages table")
|
||||
cursor.execute("ALTER TABLE chat_messages ADD COLUMN bot_name TEXT")
|
||||
logger.info("bot_name column added successfully")
|
||||
|
||||
# role replaces the old convention of identifying bot responses by a
|
||||
# hard-coded bot username.
|
||||
if "role" not in columns:
|
||||
logger.info("Adding role column to chat_messages table")
|
||||
cursor.execute("ALTER TABLE chat_messages ADD COLUMN role TEXT")
|
||||
cursor.execute(
|
||||
"UPDATE chat_messages SET role = 'assistant' "
|
||||
"WHERE message_id LIKE '%_response' AND role IS NULL",
|
||||
)
|
||||
cursor.execute(
|
||||
"UPDATE chat_messages SET role = 'user' WHERE role IS NULL",
|
||||
)
|
||||
logger.info("role column added and backfilled")
|
||||
|
||||
|
||||
# Backfill in batches so a large legacy table does not build one huge
|
||||
# executemany parameter list in memory.
|
||||
NORM_BACKFILL_BATCH = 500
|
||||
|
||||
|
||||
def _migrate_message_embeddings(cursor: sqlite3.Cursor) -> None:
|
||||
"""Add the norm column to pre-existing databases and backfill it.
|
||||
|
||||
The norm is the L2 norm of the stored float32 blob, so search can score
|
||||
candidates with one matrix multiply and no per-vector renormalization.
|
||||
"""
|
||||
logger.info("Checking for message_embeddings column migrations")
|
||||
cursor.execute("PRAGMA table_info(message_embeddings)")
|
||||
columns = {row[1] for row in cursor.fetchall()}
|
||||
|
||||
if "norm" in columns:
|
||||
return
|
||||
|
||||
logger.info("Adding norm column to message_embeddings table")
|
||||
cursor.execute("ALTER TABLE message_embeddings ADD COLUMN norm REAL")
|
||||
|
||||
cursor.execute(
|
||||
"SELECT message_id, embedding FROM message_embeddings "
|
||||
"WHERE embedding IS NOT NULL",
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
for start in range(0, len(rows), NORM_BACKFILL_BATCH):
|
||||
cursor.executemany(
|
||||
"UPDATE message_embeddings SET norm = ? WHERE message_id = ?",
|
||||
[
|
||||
(
|
||||
float(np.linalg.norm(np.frombuffer(blob, dtype=np.float32))),
|
||||
message_id,
|
||||
)
|
||||
for message_id, blob in rows[start : start + NORM_BACKFILL_BATCH]
|
||||
],
|
||||
)
|
||||
logger.info("norm column added and backfilled for %d rows", len(rows))
|
||||
|
||||
|
||||
def initialize_custom_bots_table(db_path: str) -> None:
|
||||
"""Create the custom bots table in SQLite."""
|
||||
conn = connect(db_path)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS custom_bots (
|
||||
bot_name TEXT PRIMARY KEY,
|
||||
system_prompt TEXT NOT NULL,
|
||||
created_by TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
is_active INTEGER DEFAULT 1
|
||||
)
|
||||
""",
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
@@ -0,0 +1,204 @@
|
||||
"""RAG retrieval: similarity search over user messages and history lookups."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vibe_bot import llm_client
|
||||
from vibe_bot.config import (
|
||||
EMBEDDING_ENDPOINT,
|
||||
EMBEDDING_ENDPOINT_KEY,
|
||||
EMBEDDING_MODEL,
|
||||
SIMILARITY_THRESHOLD,
|
||||
TOP_K_RESULTS,
|
||||
)
|
||||
from vibe_bot.db.connection import connect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def search_similar_messages(
|
||||
db_path: str,
|
||||
query: str,
|
||||
top_k: int = TOP_K_RESULTS,
|
||||
min_similarity: float = SIMILARITY_THRESHOLD,
|
||||
) -> list[tuple[str, str, float]]:
|
||||
"""Search for messages similar to the query using embeddings.
|
||||
|
||||
A single JOIN pulls every user row, its stored embedding, the stored L2
|
||||
norm, and its ``_response`` companion. Similarities are one matrix
|
||||
multiply over the stored norms — no per-vector renormalization. Rows
|
||||
with a missing or zero norm score 0 instead of dividing by zero.
|
||||
"""
|
||||
query_embedding = llm_client.embedding(
|
||||
text=query,
|
||||
model=EMBEDDING_MODEL,
|
||||
url=EMBEDDING_ENDPOINT,
|
||||
api_key=EMBEDDING_ENDPOINT_KEY,
|
||||
)
|
||||
if not query_embedding:
|
||||
return []
|
||||
|
||||
query_vector = np.array(query_embedding, dtype=np.float32)
|
||||
query_norm = float(np.linalg.norm(query_vector))
|
||||
if query_norm == 0:
|
||||
return []
|
||||
|
||||
conn = connect(db_path)
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT cm.content, r.content, me.embedding, me.norm
|
||||
FROM chat_messages cm
|
||||
JOIN message_embeddings me ON me.message_id = cm.message_id
|
||||
LEFT JOIN chat_messages r ON r.message_id = cm.message_id || '_response'
|
||||
WHERE cm.role = 'user'
|
||||
""",
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
if not rows:
|
||||
return []
|
||||
|
||||
n_rows = len(rows)
|
||||
blobs = [embedding_blob for _c, _r, embedding_blob, _n in rows]
|
||||
dim = len(blobs[0]) // 4
|
||||
if sum(len(blob) for blob in blobs) == n_rows * dim * 4:
|
||||
vectors = np.frombuffer(b"".join(blobs), dtype=np.float32).reshape(n_rows, dim)
|
||||
else:
|
||||
# Mixed blob lengths (e.g. EMBEDDING_MODEL changed mid-life) can't be
|
||||
# batched into one reshape; reconstruct per row, zero-padded (or
|
||||
# truncated) to the query dim so the single matrix multiply still works.
|
||||
vectors = np.zeros((n_rows, query_vector.size), dtype=np.float32)
|
||||
for i, blob in enumerate(blobs):
|
||||
row = np.frombuffer(blob, dtype=np.float32)
|
||||
k = min(row.size, query_vector.size)
|
||||
vectors[i, :k] = row[:k]
|
||||
norms = np.array(
|
||||
[
|
||||
stored_norm if stored_norm is not None else 0.0
|
||||
for _c, _r, _b, stored_norm in rows
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
safe_norms = np.where(norms > 0, norms, 1.0)
|
||||
similarities = vectors @ query_vector / (safe_norms * query_norm)
|
||||
similarities = np.where(norms > 0, similarities, 0.0)
|
||||
|
||||
results: list[tuple[str, str, float]] = []
|
||||
for (content, response, _blob, _norm), similarity in zip(
|
||||
rows, similarities, strict=True
|
||||
):
|
||||
if response is None or similarity < min_similarity:
|
||||
continue
|
||||
results.append((str(content), str(response), float(similarity)))
|
||||
|
||||
results.sort(key=lambda item: item[2], reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
|
||||
def get_bot_history(
|
||||
db_path: str, bot_name: str, limit: int = 20
|
||||
) -> list[tuple[str, str]]:
|
||||
"""Get message history for a specific custom bot.
|
||||
|
||||
Args:
|
||||
bot_name: The name of the custom bot.
|
||||
limit: Maximum number of messages to retrieve.
|
||||
|
||||
Returns:
|
||||
List of (user_message, bot_response) tuples.
|
||||
|
||||
"""
|
||||
conn = connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
logger.debug(
|
||||
"Fetching last %d messages for bot %r",
|
||||
limit,
|
||||
bot_name,
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT message_id, content
|
||||
FROM chat_messages
|
||||
WHERE bot_name = ? AND message_id NOT LIKE '%%_response'
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(bot_name, limit),
|
||||
)
|
||||
|
||||
conversations: list[tuple[str, str]] = []
|
||||
try:
|
||||
for message_id, msg_content in cursor.fetchall():
|
||||
logger.debug("Finding response for message_id=%s", message_id)
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT content
|
||||
FROM chat_messages
|
||||
WHERE message_id = ?
|
||||
ORDER BY timestamp DESC
|
||||
""",
|
||||
(f"{message_id}_response",),
|
||||
)
|
||||
response_row = cursor.fetchone()
|
||||
if response_row:
|
||||
logger.debug("Found response for message_id=%s", message_id)
|
||||
conversations.append((str(msg_content), str(response_row[0])))
|
||||
else:
|
||||
logger.debug("No response found")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
return conversations
|
||||
|
||||
|
||||
def get_user_history(
|
||||
db_path: str, user_id: str, limit: int = 20
|
||||
) -> list[tuple[str, str]]:
|
||||
"""Get message history for a specific user."""
|
||||
conn = connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
logger.debug("Fetching last %d user messages", limit)
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT message_id, content
|
||||
FROM chat_messages
|
||||
WHERE user_id = ? AND role = 'user'
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(user_id, limit),
|
||||
)
|
||||
|
||||
# Format is [user message, bot response]
|
||||
conversations: list[tuple[str, str]] = []
|
||||
try:
|
||||
for message_id, msg_content in cursor.fetchall():
|
||||
logger.debug("Finding response for message_id=%s", message_id)
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT content
|
||||
FROM chat_messages
|
||||
WHERE message_id = ?
|
||||
ORDER BY timestamp DESC
|
||||
""",
|
||||
(f"{message_id}_response",),
|
||||
)
|
||||
response_row = cursor.fetchone()
|
||||
if response_row:
|
||||
logger.debug("Found response for message_id=%s", message_id)
|
||||
conversations.append((str(msg_content), str(response_row[0])))
|
||||
else:
|
||||
logger.debug("No response found")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
return conversations
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Image-generation timing statistics (moving-average estimate)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from vibe_bot.db.connection import connect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Moving-average window (most recent generations) for the time estimate.
|
||||
IMAGE_GEN_TIME_WINDOW = 10
|
||||
# Maximum number of generation times retained in the database.
|
||||
IMAGE_GEN_TIME_LIMIT = 100
|
||||
|
||||
|
||||
def record_image_generation_time(db_path: str, duration_seconds: float) -> bool:
|
||||
"""Record how long an image generation took.
|
||||
|
||||
Args:
|
||||
duration_seconds: Wall-clock seconds the generation took.
|
||||
|
||||
"""
|
||||
logger.debug("Recording image generation time: %.2fs", duration_seconds)
|
||||
conn = connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO image_generation_times (duration_seconds)
|
||||
VALUES (?)
|
||||
""",
|
||||
(duration_seconds,),
|
||||
)
|
||||
|
||||
# Cap the table so it doesn't grow unbounded.
|
||||
cursor.execute(
|
||||
"""
|
||||
DELETE FROM image_generation_times
|
||||
WHERE id NOT IN (
|
||||
SELECT id FROM image_generation_times
|
||||
ORDER BY id DESC
|
||||
LIMIT ?
|
||||
)
|
||||
""",
|
||||
(IMAGE_GEN_TIME_LIMIT,),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
except Exception:
|
||||
logger.exception("Error recording image generation time")
|
||||
conn.rollback()
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_image_generation_time_estimate(db_path: str) -> float | None:
|
||||
"""Get a moving-average estimate of image generation time.
|
||||
|
||||
Returns:
|
||||
The average duration in seconds over the most recent generations,
|
||||
or None if there is no generation history yet.
|
||||
|
||||
"""
|
||||
logger.debug("Computing moving-average image generation time estimate")
|
||||
conn = connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT AVG(duration_seconds)
|
||||
FROM (
|
||||
SELECT duration_seconds
|
||||
FROM image_generation_times
|
||||
ORDER BY id DESC
|
||||
LIMIT ?
|
||||
)
|
||||
""",
|
||||
(IMAGE_GEN_TIME_WINDOW,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
except Exception:
|
||||
logger.exception("Error reading image generation times")
|
||||
return None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
if row is None or row[0] is None:
|
||||
return None
|
||||
return float(row[0])
|
||||
@@ -0,0 +1,42 @@
|
||||
"""float32 embedding (de)serialization and cosine similarity."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def vector_to_bytes(vector: list[float]) -> bytes:
|
||||
"""Convert a vector to bytes for SQLite storage."""
|
||||
logger.debug("Converting vector (length: %d) to bytes", len(vector))
|
||||
result = np.array(vector, dtype=np.float32).tobytes()
|
||||
logger.debug("Vector converted to %d bytes", len(result))
|
||||
return result
|
||||
|
||||
|
||||
def bytes_to_vector(blob: bytes) -> np.ndarray:
|
||||
"""Convert bytes back to a vector."""
|
||||
logger.debug("Converting %d bytes back to vector", len(blob))
|
||||
result = np.frombuffer(blob, dtype=np.float32)
|
||||
logger.debug("Vector reconstructed with %d dimensions", len(result))
|
||||
return result
|
||||
|
||||
|
||||
def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
|
||||
"""Calculate cosine similarity between two vectors."""
|
||||
vec1 = vec1.flatten()
|
||||
vec2 = vec2.flatten()
|
||||
logger.debug(
|
||||
"Calculating cosine similarity between vectors of dimension %d",
|
||||
len(vec1),
|
||||
)
|
||||
norm1 = np.linalg.norm(vec1)
|
||||
norm2 = np.linalg.norm(vec2)
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
result = float(np.dot(vec1, vec2) / (norm1 * norm2))
|
||||
logger.debug("Similarity calculated: %.4f", result)
|
||||
return result
|
||||
Reference in New Issue
Block a user