complete restructure
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user