chat/engine.py

"""Chat engine — manages conversations with streaming support."""

import json
import logging
from typing import Optional, AsyncIterator

from config import config
from database.repositories import ConversationRepo, MessageRepo
from agent.synapse_agent import SynapseAgent
from memory.manager import MemoryManager

logger = logging.getLogger("synapse.chat")


class ChatEngine:
    """Manages chat sessions with the Synapse agent."""

    def __init__(self):
        self._active_sessions: dict[str, SynapseAgent] = {}

    def get_agent(self, user_id: str = "guest",
                  conversation_id: str = None,
                  mode: str = "agent",
                  personality_prompt: str = None) -> SynapseAgent:
        key = f"{user_id}:{conversation_id or 'new'}"
        if key not in self._active_sessions:
            self._active_sessions[key] = SynapseAgent(
                user_id=user_id,
                conversation_id=conversation_id,
                mode=mode,
                personality_prompt=personality_prompt,
            )
        return self._active_sessions[key]

    async def send_message(self, content: str, user_id: str = "guest",
                           conversation_id: str = None,
                           mode: str = "agent",
                           personality_prompt: str = None,
                           files: list[dict] = None) -> dict:
        if not conversation_id:
            conv = await ConversationRepo.create(user_id=user_id)
            conversation_id = conv["id"]

        await MessageRepo.create(
            conversation_id=conversation_id,
            role="user",
            content=content,
            metadata={"files": files} if files else None,
        )

        history = await MessageRepo.get_last_n(conversation_id, n=20)
        conversation_messages = [
            {"role": m["role"], "content": m["content"]} for m in history
        ]

        agent = self.get_agent(user_id, conversation_id, mode, personality_prompt)
        result = await agent.process(content, conversation_messages, mode)

        await MessageRepo.create(
            conversation_id=conversation_id,
            role="assistant",
            content=result.get("content", ""),
            reasoning=json.dumps(result.get("reasoning")) if result.get("reasoning") else None,
            tool_calls=json.dumps(result.get("tools_used")) if result.get("tools_used") else None,
            model=result.get("model"),
        )

        msg_count = len(await MessageRepo.get_conversation_messages(conversation_id))
        if msg_count == 2:
            auto_title = content[:60] + ("..." if len(content) > 60 else "")
            await ConversationRepo.update_title(conversation_id, auto_title)

        result["conversation_id"] = conversation_id
        return result

    async def stream_message(self, content: str, user_id: str = "guest",
                             conversation_id: str = None,
                             mode: str = "agent",
                             personality_prompt: str = None) -> AsyncIterator[str]:
        if not conversation_id:
            conv = await ConversationRepo.create(user_id=user_id)
            conversation_id = conv["id"]
            yield json.dumps({"event": "conversation_created", "id": conversation_id})

        await MessageRepo.create(
            conversation_id=conversation_id,
            role="user",
            content=content,
        )

        history = await MessageRepo.get_last_n(conversation_id, n=20)
        conversation_messages = [
            {"role": m["role"], "content": m["content"]} for m in history
        ]

        agent = self.get_agent(user_id, conversation_id, mode, personality_prompt)

        from models.router import router
        tool_defs = agent.tool_executor.get_tool_definitions()
        system_msg = agent._build_system_message(tool_defs)
        all_messages = [system_msg] + conversation_messages

        try:
            stream = await router.chat(
                messages=all_messages,
                temperature=config.groq.temperature,
                max_tokens=config.groq.max_tokens,
                stream=True,
            )

            full_content = ""
            async for chunk in self._process_stream(stream):
                full_content += chunk
                yield json.dumps({"event": "token", "content": chunk})

            await MessageRepo.create(
                conversation_id=conversation_id,
                role="assistant",
                content=full_content,
            )

            yield json.dumps({"event": "done", "conversation_id": conversation_id})
        except Exception as e:
            logger.error(f"Stream error: {e}")
            yield json.dumps({"event": "error", "message": str(e)})

    async def _process_stream(self, stream) -> AsyncIterator[str]:
        try:
            async for chunk in stream:
                if hasattr(chunk, "choices") and chunk.choices:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, "content") and delta.content:
                        yield delta.content
                elif isinstance(chunk, dict):
                    choices = chunk.get("choices", [])
                    if choices:
                        delta = choices[0].get("delta", {})
                        content = delta.get("content", "")
                        if content:
                            yield content
        except Exception as e:
            logger.error(f"Stream processing error: {e}")
            yield f"\n[Stream error: {e}]"

    async def get_conversation_history(self, conversation_id: str) -> list[dict]:
        return await MessageRepo.get_conversation_messages(conversation_id)

    async def create_conversation(self, user_id: str, title: str = "New Chat",
                                  mode: str = "agent",
                                  personality_id: str = "synapse",
                                  folder_id: str = None) -> dict:
        return await ConversationRepo.create(
            user_id=user_id,
            title=title,
            personality_id=personality_id,
            mode=mode,
            folder_id=folder_id,
        )

    async def get_user_conversations(self, user_id: str,
                                     folder_id: str = None) -> list[dict]:
        return await ConversationRepo.get_user_conversations(user_id, folder_id)

    async def delete_conversation(self, conv_id: str) -> bool:
        return await ConversationRepo.delete(conv_id)

    async def update_title(self, conv_id: str, title: str) -> None:
        await ConversationRepo.update_title(conv_id, title)

    async def toggle_pin(self, conv_id: str) -> None:
        await ConversationRepo.toggle_pin(conv_id)

    async def toggle_bookmark(self, conv_id: str) -> None:
        await ConversationRepo.toggle_bookmark(conv_id)

    async def search_conversations(self, user_id: str, query: str) -> list[dict]:
        return await ConversationRepo.search(user_id, query)

    async def search_messages(self, user_id: str, query: str) -> list[dict]:
        return await MessageRepo.search(user_id, query)


chat_engine = ChatEngine()