149 lines
4.1 KiB
Python
149 lines
4.1 KiB
Python
"""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()
|