hendrik/tui/app.py

import curses
import threading
from datetime import datetime
import config
from .render import init_colors, draw
from .input import handle_key
from .agent import log, WELCOME_ART
from services.session_manager_neo import NeoSessionManager, NeoSession


class HendrikTUI:
    def __init__(self, llm_client, tools_definition, TOOLS, TOOL_HANDLERS,
                 build_system_prompt, agent_max_iterations):
        self.llm = llm_client
        self.tools_def = tools_definition
        self.TOOLS = TOOLS
        self.TOOL_HANDLERS = TOOL_HANDLERS
        self.build_system_prompt = build_system_prompt
        self.agent_max_iterations = agent_max_iterations

        self.messages = None
        self.log = []
        self.input_buffer = [""]
        self.input_line = 0
        self.input_col = 0
        self.scroll = 0
        self.processing = False
        self.running = True
        self.h, self.w = 0, 0

        self.agent_thread: threading.Thread | None = None
        self.agent_done = threading.Event()

        self.session_mgr = NeoSessionManager()
        self.current_session: NeoSession | None = None

    def switch_model(self, item: dict):
        self.llm.base_url = item["base_url"]
        self.llm.model = item["model"]
        self.llm.api_key = item["api_key"]
        if self.current_session:
            self.session_mgr.update_model_info(
                self.current_session.doc_id, self._model_info()
            )

    def _model_info(self) -> dict:
        return {
            "provider": config.resolve_provider(self.llm.base_url, self.llm.model),
            "base_url": self.llm.base_url,
            "model": self.llm.model,
        }

    def new_session(self):
        name = f"Session {datetime.now().strftime('%Y-%m-%d %H:%M')}"
        self.current_session = self.session_mgr.create(name, self._model_info())
        self.messages = [{"role": "system", "content": self.build_system_prompt(
            tools_definition=self.tools_def,
            character=config.AGENT_CHARACTER or None,
            skills=config.AGENT_SKILLS.split(",") if config.AGENT_SKILLS else None,
        )}]
        self.log.clear()
        self.scroll = 0
        log(self, "welcome", WELCOME_ART)

    def switch_session(self, doc_id: int):
        session = self.session_mgr.get(doc_id)
        if not session:
            return
        self.current_session = session
        self.messages = list(session.messages)
        self.log.clear()
        self.scroll = 0
        log(self, "welcome", WELCOME_ART)
        for i, msg in enumerate(session.messages):
            if msg["role"] == "user":
                if i > 0:
                    log(self, "sep", "")
                log(self, "user", msg["content"])
                mi = msg.get("model_info")
                if mi:
                    self.log[-1]["model_info"] = (mi.get("provider"), mi.get("model"))
            elif msg["role"] == "assistant":
                next_role = session.messages[i + 1]["role"] if i + 1 < len(session.messages) else None
                if next_role in ("user", None):
                    log(self, "ai", msg["content"])
        if session.messages and session.messages[-1]["role"] == "assistant":
            log(self, "sep", "")

    def run(self):
        try:
            curses.wrapper(self._main)
        except KeyboardInterrupt:
            pass

    def _main(self, stdscr):
        curses.use_default_colors()
        init_colors()
        stdscr.keypad(True)
        stdscr.refresh()

        self.messages = [{"role": "system", "content": self.build_system_prompt(
            tools_definition=self.tools_def,
            character=config.AGENT_CHARACTER or None,
            skills=config.AGENT_SKILLS.split(",") if config.AGENT_SKILLS else None,
        )}]
        log(self, "welcome", WELCOME_ART)

        while self.running:
            self.h, self.w = stdscr.getmaxyx()
            if self.h < 14 or self.w < 40:
                stdscr.erase()
                stdscr.addstr(0, 0, "Terminal too small (min 40x14)")
                stdscr.refresh()
                stdscr.getch()
                continue

            draw(self, stdscr)
            curses.curs_set(2)

            if self.processing:
                stdscr.timeout(100)
            else:
                stdscr.timeout(-1)

            try:
                key = stdscr.getch()
            except KeyboardInterrupt:
                break

            handle_key(self, stdscr, key)

            if self.agent_done.is_set():
                self.agent_thread.join()
                self.agent_done.clear()
                self.processing = False
                self.agent_thread = None