NexusOS/synapse/chat.py

from __future__ import annotations

import asyncio
import json as _json
import logging
import threading
from typing import AsyncGenerator, Dict, List, Optional, Any

from .nexus_config import settings
from .ollama_manager import get_ollama_manager


# -------------------------
# Logger setup
# -------------------------
_logger = logging.getLogger("nexus.chat")
_logger.setLevel(logging.INFO)
if not _logger.handlers:
    handler = logging.FileHandler(str(settings.chat_log)) if getattr(settings, "chat_log", None) else logging.StreamHandler()
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
    handler.setFormatter(formatter)
    _logger.addHandler(handler)


# -------------------------
# Synapse tracer (real-time prompt/token view for control panel)
# -------------------------
_synapse_lock = threading.Lock()
_synapse_fh = None

def _synapse_trace(text: str) -> None:
    global _synapse_fh
    try:
        log_path = getattr(settings, "chat_log", None)
        if not log_path:
            return
        with _synapse_lock:
            if _synapse_fh is None or _synapse_fh.closed:
                _synapse_fh = open(str(log_path), "a", buffering=1, encoding="utf-8")
            _synapse_fh.write(text)
            _synapse_fh.flush()
    except Exception:
        pass


# -------------------------
# Non-streaming generation
# -------------------------
async def generate_chat_response(
    user_message: str,
    metadata: Optional[Dict[str, Any]] = None,
    history: Optional[List[Dict[str, str]]] = None,
    timeout: Optional[float] = None,
) -> Dict[str, Any]:
    metadata = metadata or {}
    timeout = timeout or getattr(settings, "ollama_timeout", 120)

    manager = get_ollama_manager()
    system = metadata.get("system", "")
    model = metadata.get("model") or "mistral"
    temperature = metadata.get("temperature")
    num_gpu = metadata.get("num_gpu")

    messages: List[Dict[str, str]] = []
    if system:
        messages.append({"role": "system", "content": system})
    for msg in (history or []):
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})

    _logger.info("generate_chat_response: model=%s turns=%d timeout=%s", model, len(messages), timeout)

    sys_preview = (system or "")[:200].replace("\n", " ")
    _synapse_trace(f"\n── TURN [{model} | {len(messages)} msgs] {'─' * 30}\n")
    if system:
        _synapse_trace(f"SYS: {sys_preview}{'…' if len(system) > 200 else ''}\n")
    _synapse_trace(f"USR: {user_message}\n{'─' * 50}\n")

    try:
        result = await asyncio.wait_for(
            manager.chat(messages=messages, model=model, stream=False, temperature=temperature, num_gpu=num_gpu),
            timeout=timeout,
        )
        response_text = result if isinstance(result, str) else str(result)

        preview = response_text[:500].replace("\n", " ")
        _synapse_trace(f"{preview}{'…' if len(response_text) > 500 else ''}\n{'─' * 50}\n")
        _logger.info("generate_chat_response: completed model=%s", model)
        return {"response": response_text, "model": model, "metadata": metadata}

    except asyncio.TimeoutError:
        _logger.exception("generate_chat_response: timeout after %s seconds", timeout)
        raise
    except Exception:
        _logger.exception("generate_chat_response: unexpected error")
        raise


# -------------------------
# Async iterator timeout helper
# -------------------------
async def _aiter_with_timeout(aiterable, timeout: Optional[float]):
    if timeout is None or timeout <= 0:
        async for item in aiterable:
            yield item
        return

    aiter = aiterable.__aiter__()
    while True:
        try:
            item = await asyncio.wait_for(aiter.__anext__(), timeout=timeout)
            yield item
        except StopAsyncIteration:
            break


# -------------------------
# Normalizer for many return shapes
# -------------------------
async def _normalize_to_async_generator(maybe_iterable) -> AsyncGenerator[str, None]:
    # The sole caller passes manager.chat(stream=True) — an async-def call, i.e.
    # a coroutine that resolves to an async generator. Await it if needed, then
    # stream the tokens.
    result = await maybe_iterable if asyncio.iscoroutine(maybe_iterable) else maybe_iterable
    async for item in result:
        yield str(item)


# -------------------------
# Streaming implementation
# -------------------------
async def stream_chat_response(
    user_message: str,
    metadata: Optional[Dict[str, Any]] = None,
    history: Optional[List[Dict[str, str]]] = None,
    timeout: Optional[float] = None,
) -> AsyncGenerator[str, None]:
    metadata = metadata or {}
    timeout = timeout or getattr(settings, "ollama_timeout", 120)

    manager = get_ollama_manager()
    system = metadata.get("system", "")
    model = metadata.get("model") or "mistral"
    temperature = metadata.get("temperature")
    num_gpu = metadata.get("num_gpu")

    # Build messages array for /api/chat multi-turn format
    messages: List[Dict[str, str]] = []
    if system:
        messages.append({"role": "system", "content": system})
    for msg in (history or []):
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})

    _logger.info("stream_chat_response: starting stream (model=%s, turns=%d, timeout=%s)", model, len(messages), timeout)

    sys_preview = (system or "")[:200].replace("\n", " ")
    _synapse_trace(f"\n── TURN [{model} | {len(messages)} msgs] {'─' * 30}\n")
    if system:
        _synapse_trace(f"SYS: {sys_preview}{'…' if len(system) > 200 else ''}\n")
    _synapse_trace(f"USR: {user_message}\n{'─' * 50}\n")

    try:
        maybe_iter = manager.chat(messages=messages, model=model, stream=True, temperature=temperature, num_gpu=num_gpu)
        async_gen = _normalize_to_async_generator(maybe_iter)

        buffer_parts: list[str] = []
        buffer_len = 0
        FLUSH_THRESHOLD = 24

        async for piece in _aiter_with_timeout(async_gen, timeout):
            if piece is None:
                continue
            text = str(piece)
            if not text:
                continue

            # Pass stats sentinel through immediately, don't buffer it
            if text.startswith("__meta__"):
                if buffer_parts:
                    chunk = "".join(buffer_parts)
                    buffer_parts = []
                    buffer_len = 0
                    _synapse_trace(chunk.replace("\n", " ") + "\n")
                    yield chunk
                yield text
                continue

            buffer_parts.append(text)
            buffer_len += len(text)

            if buffer_len >= FLUSH_THRESHOLD or any(text.endswith(c) for c in (".", "!", "?", "\n")):
                chunk = "".join(buffer_parts)
                buffer_parts = []
                buffer_len = 0
                _synapse_trace(chunk.replace("\n", " ") + "\n")
                yield chunk

        if buffer_parts:
            chunk = "".join(buffer_parts)
            _synapse_trace(chunk.replace("\n", " ") + "\n")
            yield chunk

        _synapse_trace(f"{'─' * 50}\n")
        _logger.info("stream_chat_response: stream completed")

    except asyncio.TimeoutError:
        _logger.exception("stream_chat_response: timeout after %s seconds", timeout)
        raise
    except Exception:
        _logger.exception("stream_chat_response: unexpected error during streaming")
        raise