hendrik/tui/agent.py

import json
import threading
from datetime import datetime
import config
from scripts import ntro

def _add_msg(app, role, content, **kwargs):
    msg = {"role": role, "content": content}
    msg.update(kwargs)
    app.messages.append(msg)
    if app.current_session:
        app.session_mgr.add_message(
            app.current_session.doc_id, role, content, **kwargs
        )

WELCOME_ART = """
,--.  ,--.,------.,--.  ,--.,------.  ,------. ,--.,--. ,--. 
|  '--'  ||  .---'|  ,'.|  ||  .-.  \\ |  .--. '|  ||  .'   / 
|  .--.  ||  `--, |  |' '  ||  |  \\  :|  '--'.'|  ||  .   '  
|  |  |  ||  `---.|  | `   ||  '--'  /|  |\\  \\ |  ||  |\\   \\ 
`--'  `--'`------'`--'  `--'`-------' `--' '--'`--'`--' '--' 

"""
"""
 __   __  _______  __    _  ______   ______    ___   ___   _ 
|  | |  ||       ||  |  | ||      | |    _ |  |   | |   | | |
|  |_|  ||    ___||   |_| ||  _    ||   | ||  |   | |   |_| |
|       ||   |___ |       || | |   ||   |_||_ |   | |      _|
|       ||    ___||  _    || |_|   ||    __  ||   | |     |_ 
|   _   ||   |___ | | |   ||       ||   |  | ||   | |    _  |
|__| |__||_______||_|  |__||______| |___|  |_||___| |___| |_|
"""
"""
╔══════════════════════════════════════════╗
║                                          ║
║     /\\_/\\                                ║
║    ( o.o )      HENDRIK                  ║
║     > ^ <                                ║
║    (     )      AI Agent                 ║
║     (___)                                ║
║                                          ║
╚══════════════════════════════════════════╝
"""

def log(app, role, text):
    app.log.append({
        "role": role,
        "text": text,
        "time": datetime.now().strftime("%H:%M"),
    })


def submit(app, stdscr):
    query = "\n".join(app.input_buffer).strip()
    if not query:
        return

    log(app, "user", query)
    model_info = (
        config.resolve_provider(app.llm.base_url, app.llm.model),
        app.llm.model
    )
    if app.log:
        app.log[-1]["model_info"] = model_info
    app.input_buffer = [""]
    app.input_line = 0
    app.input_col = 0
    app.scroll = 999999
    app.processing = True

    if app.current_session is None:
        app.current_session = app.session_mgr.create(
            f"Session {datetime.now().strftime('%Y-%m-%d %H:%M')}",
            app._model_info(),
        )

    _add_msg(app, "user", query, model_info=app._model_info())

    app.agent_done.clear()
    app.agent_thread = threading.Thread(
        target=_agent_loop,
        args=(app,),
        daemon=True,
    )
    app.agent_thread.start()


def _agent_loop(app):
    stamp = ntro.start()

    for step in range(app.agent_max_iterations):
        stamp_step = ntro.start()
        log(app, "system", f"  step {step + 1} \u2014 Thinking...")
        app.scroll = 999999

        # Streaming response
        stream_buffer = []
        placeholder_marker = None

        def on_stream_chunk(chunk):
            nonlocal placeholder_marker
            # Buat placeholder di chunk pertama saja
            if placeholder_marker is None:
                placeholder_marker = f"_stream_placeholder_{step}"
                log(app, "ai", placeholder_marker)
            stream_buffer.append(chunk)
            current_text = ''.join(stream_buffer)
            # Update placeholder secara real-time
            for i in range(len(app.log) - 1, -1, -1):
                # Cari placeholder berdasarkan content aslinya (atau apa yang sudah diupdate)
                # Karena placeholder text berubah seiring streaming, kita harus teliti
                if app.log[i].get('role') == 'ai' and (app.log[i].get('text') == placeholder_marker or (i > 0 and app.log[i-1].get('role') == 'system' and 'Thinking' in app.log[i-1].get('text', ''))):
                    app.log[i]['text'] = current_text
                    break
            app.scroll = 999999

        response = app.llm.chat(app.messages, tools=app.TOOLS, on_stream_chunk=on_stream_chunk)

        # Hapus "Thinking..." log
        for i in range(len(app.log) - 1, -1, -1):
            if app.log[i].get('role') == 'system' and 'Thinking' in app.log[i].get('text', ''):
                app.log.pop(i)
                break

        if response.warning:
            log(app, "system", f"  {response.warning}")

        # Cek apakah ada tool_calls
        has_tool_calls = bool(response.tool_calls)

        if has_tool_calls:
            # Hapus placeholder jika ada tool_calls (cari semua AI log yang mungkin placeholder)
            if placeholder_marker is not None:
                for i in range(len(app.log) - 1, -1, -1):
                    # Jika ini adalah log AI yang muncul tepat setelah "Thinking..." atau contains placeholder marker
                    if app.log[i].get('role') == 'ai':
                        text = app.log[i].get('text', '')
                        # Hapus jika text-nya adalah placeholder atau jika itu adalah entry AI yang baru saja kita buat untuk streaming
                        if placeholder_marker in text or text == "" or text == "...":
                            app.log.pop(i)
            
            _add_msg(app, "assistant", response.content, tool_calls=response.tool_calls)

            # Log tool_calls dengan label AI
            for tc in response.tool_calls:
                tname = tc["function"]["name"]
                targs = tc["function"]["arguments"]
                log(app, "tool_call", json.dumps({
                    "name": tname,
                    "arguments": targs,
                }))
                app.scroll = 999999
                execute_tool(app, tc)

            # Log content AI setelah tools (jika ada)
            if response.content and response.content.strip():
                log(app, "ai", response.content)
        else:
            if response.content:
                _add_msg(app, "assistant", response.content)
                # Placeholder sudah terupdate via streaming, jangan log lagi
            log(app, "sep", "")
            ntro.end(stamp)
            app.agent_done.set()
            return
        ntro.end(stamp_step)

    log(app, "error", "Max iterations reached without final answer.")
    _add_msg(app, "assistant", "Max iterations reached without final answer.")
    ntro.end(stamp)
    app.agent_done.set()


def execute_tool(app, tool_call):
    tname = tool_call["function"]["name"]
    targs = json.loads(tool_call["function"]["arguments"])
    handler = app.TOOL_HANDLERS.get(tname)
    if not handler:
        result = f"Tool {tname} not found"
    else:
        try:
            if tname == "search_code":
                result = handler(
                    pattern=targs["pattern"],
                    search_type=targs["search_type"],
                    path=targs.get("path", "."),
                )
            elif tname == "git_operation":
                result = handler(args=targs["args"])
            else:
                result = handler(**targs)
        except Exception as e:
            result = f"Error executing tool: {str(e)}"

    _add_msg(app, "tool", str(result), tool_call_id=tool_call["id"])