130 lines
4.5 KiB
Python
130 lines
4.5 KiB
Python
"""Chat service: RAG context + tool-capped LLM completion + persistence + reply."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vibe_bot import llm_client
|
|
from vibe_bot.config import CHAT_MODEL, MAX_COMPLETION_TOKENS
|
|
from vibe_bot.prompts import build_system_prompt, get_user_info
|
|
from vibe_bot.textutil import split_message
|
|
|
|
if TYPE_CHECKING:
|
|
from discord.ext.commands import Bot, Context
|
|
|
|
from vibe_bot.database import ChatDatabase
|
|
from vibe_bot.llm_client import ToolRegistry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ChatService:
|
|
"""Handles one custom-bot chat turn: context -> LLM (with tools) -> persist -> reply."""
|
|
|
|
def __init__(self, db: ChatDatabase, registry: ToolRegistry) -> None:
|
|
self._db = db
|
|
self._registry = registry
|
|
|
|
async def handle(
|
|
self,
|
|
ctx: Context[Bot],
|
|
*,
|
|
bot_name: str,
|
|
message: str,
|
|
system_prompt: str,
|
|
response_prefix: str,
|
|
) -> None:
|
|
"""Run a single chat turn for ``bot_name`` and send the reply.
|
|
|
|
Args:
|
|
ctx: The Discord command context.
|
|
bot_name: The name of the custom bot.
|
|
message: The user message to process.
|
|
system_prompt: The base system prompt (personality) for the bot.
|
|
response_prefix: The prefix message sent before the reply.
|
|
|
|
"""
|
|
await ctx.send(f"{bot_name} is searching its databanks for {message[:50]}...")
|
|
|
|
# Get conversation context using RAG (SQLite + embedding HTTP call).
|
|
context = await asyncio.to_thread(
|
|
self._db.get_conversation_context,
|
|
user_id=str(ctx.author.id),
|
|
current_message=message,
|
|
max_context=5,
|
|
)
|
|
|
|
prompts: list[dict[str, str]] = [{"role": "user", "content": message}]
|
|
if context:
|
|
prompts = context + prompts
|
|
|
|
logger.info(
|
|
"chat: bot=%s user=%s context_msgs=%d",
|
|
bot_name,
|
|
ctx.author.id,
|
|
len(context),
|
|
)
|
|
|
|
system_prompt_edit = build_system_prompt(
|
|
system_prompt, get_user_info(ctx.author)
|
|
)
|
|
|
|
tools = self._registry.to_openai_tools()
|
|
|
|
def tool_executor(tool_name: str, tool_args: dict[str, str]) -> str:
|
|
"""Dispatch a tool call through the registry."""
|
|
return self._registry.execute(tool_name, tool_args, channel=ctx.channel)
|
|
|
|
async def tool_call_notifier(tool_name: str, tool_args: dict[str, str]) -> None:
|
|
"""Send a notification message when a tool is called."""
|
|
if tool_name == "get_channel_members":
|
|
await ctx.send(f"{bot_name} is looking at the channel members...")
|
|
|
|
try:
|
|
bot_response = await llm_client.chat_completion_with_tools(
|
|
system_prompt=system_prompt_edit,
|
|
prompts=prompts,
|
|
tools=tools,
|
|
tool_executor=tool_executor,
|
|
tool_call_notifier=tool_call_notifier,
|
|
model=CHAT_MODEL,
|
|
max_tokens=MAX_COMPLETION_TOKENS,
|
|
)
|
|
|
|
# Store both the user message and the bot response in the database.
|
|
await asyncio.to_thread(
|
|
self._db.add_message,
|
|
message_id=f"{ctx.message.id}",
|
|
user_id=str(ctx.author.id),
|
|
username=ctx.author.name,
|
|
content=f"User: {message}",
|
|
bot_name=bot_name,
|
|
channel_id=str(ctx.channel.id),
|
|
guild_id=str(ctx.guild.id) if ctx.guild else None,
|
|
)
|
|
|
|
if ctx.bot.user is not None:
|
|
await asyncio.to_thread(
|
|
self._db.add_message,
|
|
message_id=f"{ctx.message.id}_response",
|
|
user_id=str(ctx.bot.user.id),
|
|
username=ctx.bot.user.name,
|
|
content=bot_response,
|
|
bot_name=bot_name,
|
|
channel_id=str(ctx.channel.id),
|
|
guild_id=str(ctx.guild.id) if ctx.guild else None,
|
|
role="assistant",
|
|
embed=False,
|
|
)
|
|
|
|
# Send the response back to the chat.
|
|
await ctx.send(response_prefix)
|
|
for send_chunk in split_message(bot_response, 1000):
|
|
await ctx.send(send_chunk)
|
|
|
|
except Exception:
|
|
logger.exception("Error in handle_chat")
|
|
await ctx.send("An error occurred while processing your request.")
|