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, DEFAULT_CHAT_MODEL
from .ollama_manager import get_ollama_manager
from . import tools as _tools

# Cap on tool-call round-trips before the final answer — stops a confused small
# model from looping forever.
MAX_TOOL_STEPS = 5


# -------------------------
# 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 DEFAULT_CHAT_MODEL
    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)


async def _run_tool_loop(manager, messages, model, tool_schemas, temperature, num_gpu):
    """Let the model call read-only tools before the final streamed answer.

    Mutates `messages` IN PLACE, appending the assistant tool-call turns and
    their `role:"tool"` results, and yields a `__status__<tool>` sentinel before
    each tool runs (surfaced to the UI as a "running tool" indicator).
    Non-streamed — tool calls arrive as whole messages. Degrades to an untouched
    `messages` if the model can't do tool calling.

    ponytail: the turn that finally returns content is thrown away and the answer
    is re-generated by the streaming turn (one wasted call). Simpler than
    streaming a maybe-already-complete message; revisit if latency matters.
    """
    for _ in range(MAX_TOOL_STEPS):
        msg = await manager.chat(
            messages=messages, model=model, stream=False,
            temperature=temperature, num_gpu=num_gpu, tools=tool_schemas,
        )
        if not isinstance(msg, dict):
            break  # None/error or no tool support -> fall back to plain stream
        calls = msg.get("tool_calls")
        if not calls:
            break
        messages.append(msg)
        for c in calls:
            fn = c.get("function", {})
            name = fn.get("name", "")
            yield f"__status__{name}"
            result = await _tools.dispatch(name, fn.get("arguments"))
            messages.append({"role": "tool", "content": result})


# -------------------------
# 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 DEFAULT_CHAT_MODEL
    temperature = metadata.get("temperature")
    num_gpu = metadata.get("num_gpu")
    num_ctx = metadata.get("num_ctx")
    think = metadata.get("think", False)

    # 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"]})
    user_msg: Dict[str, Any] = {"role": "user", "content": user_message}
    images = metadata.get("images")  # base64 strings (no data: prefix) for vision models
    if images:
        user_msg["images"] = images
    messages.append(user_msg)

    # Tool-using playbooks: run read-only tool calls, then stream the final answer
    # with their results already in the messages array.
    tool_schemas = metadata.get("tools")
    if tool_schemas:
        try:
            async for status in _run_tool_loop(manager, messages, model, tool_schemas, temperature, num_gpu):
                yield status
        except Exception:
            _logger.exception("tool loop failed; streaming without tools")

    _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, think=think, num_ctx=num_ctx)
        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