ai-agent/src/ai_agent/cli.py at 34ee2f410477b85c5a5935daa1d750af6518cb4b · Imaging-Plaza/ai-agent · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# src/ai_agent/cli.py
from __future__ import annotations

import argparse
import os
import sys
import threading
import time
from dotenv import load_dotenv
import logging

load_dotenv()
log = logging.getLogger("ai_agent.cli")

from ai_agent.catalog.sync import sync_once


def _ui_funcs():
    # Lazy import avoids loading agent/model modules for non-UI commands.
    from ai_agent.ui.app import (
        get_pipeline,
        refresh_ui_docs_from_index,
        launch,
        ensure_logging_initialized,
    )

    return get_pipeline, refresh_ui_docs_from_index, launch, ensure_logging_initialized


# --------------------------- catalog background refresher ---------------------------
def _background_refresh():
    """If SYNC_EVERY_HOURS > 0, refresh in the background while UI runs."""
    hours = float(os.getenv("SYNC_EVERY_HOURS", "0") or 0)

    if hours <= 0:
        log.info("[auto-refresh] disabled")
        return

    def _loop():
        # Startup already performs one sync in run_chat(); wait one full interval
        # before the first background refresh to avoid duplicate work.
        interval_s = max(60.0, hours * 3600.0)
        time.sleep(interval_s)

        while True:
            try:
                res = sync_once()
                log.info(
                    "[auto-refresh] %s → %s",
                    res.get("count", "?"),
                    res.get("jsonl_path"),
                )

                get_pipeline, refresh_ui_docs_from_index, _, _ = _ui_funcs()
                pipe = get_pipeline()

                if res.get("changed"):
                    ok = pipe.reload_index()
                    if ok:
                        log.info("[auto-refresh] reloaded FAISS index")
                        refresh_ui_docs_from_index()
                    else:
                        log.warning(
                            "[auto-refresh] reload failed; serving previous index"
                        )
                else:
                    log.info("[auto-refresh] catalog unchanged; FAISS not touched")
            except Exception:
                log.exception("[auto-refresh] error")
            try:
                time.sleep(interval_s)
            except Exception:
                time.sleep(3600.0)

    t = threading.Thread(target=_loop, daemon=True)
    t.start()


# --------------------------- custom tasks ---------------------------
def run_chat():
    """Launch the chat-based UI."""
    try:
        _, _, _, ensure_logging_initialized = _ui_funcs()
        ensure_logging_initialized()

        res = sync_once()
        log.info("[startup-sync] %s → %s", res.get("count", "?"), res.get("jsonl_path"))

        get_pipeline, refresh_ui_docs_from_index, launch, _ = _ui_funcs()

        # Initialize pipeline
        pipe = get_pipeline()

        if res.get("changed"):
            ok = pipe.reload_index()
            if ok:
                log.info("[startup-refresh] reloaded FAISS index")
                refresh_ui_docs_from_index()
            else:
                log.warning("[startup-refresh] reload failed; serving previous index")
        else:
            log.info(
                "[startup-refresh] catalog unchanged; keeping existing FAISS index"
            )
    except Exception:
        log.exception("[startup-sync] failed")

    _background_refresh()

    try:
        _, _, launch, _ = _ui_funcs()
        launch()
    except Exception:
        log.exception("[chat-launch] failed")
        raise


def run_sync():
    try:
        r = sync_once()
        log.info("[sync] %s → %s", r.get("count", "?"), r.get("jsonl_path"))
    except Exception:
        log.exception("[sync] failed")
        raise


# --------------------------- main entry ---------------------------
def main():
    p = argparse.ArgumentParser(description="AI Agent CLI")
    p.add_argument(
        "mode",
        choices=["chat", "sync"],
        help="'chat' launches the chat UI; 'sync' runs one catalog refresh.",
    )
    args = p.parse_args()

    if args.mode == "chat":
        run_chat()
    elif args.mode == "sync":
        run_sync()
    else:
        p.print_help()
        sys.exit(f"Unsupported mode: {args.mode}")


if __name__ == "__main__":
    main()