Dynavera/apps/onboarding/consumers/chat.py

from apps.onboarding.consumers.base import BaseOnboardingConsumer, LogType

__all__ = ["OnboardingChatConsumer"]

class OnboardingChatConsumer(BaseOnboardingConsumer):
    """
    Route: /ws/onboarding/chat/<config_uuid>/
    """

    config_uuid: str 

    def parse_extra(self):
        self.config_uuid = self.scope["url_route"]["kwargs"].get("config_uuid")

    async def action_message(self, data: dict):
        user_query = data.get("query")
        if not user_query:
            return await self.send_error("Missing 'query' field in payload.")
        max_tokens = self.parse_max_tokens(data.get("max_tokens"))
        config = await self.get_config_for_user(self.config_uuid)
        if config is None:
            await self.send_error("Forbidden or Invalid Config UUID")
            return
        response = await self.stream_llm(
            config,
            user_query,
            max_tokens=max_tokens or 1024,
            system_prompt_suffix="Respond in plain text only. Do not use markdown formatting, bullet points, headers, bold, italics, or code blocks.",
        )
        await self.send_log(LogType.COMPLETED, "Inference complete.", {"response": response or ""})