quant-python-ai/main.py at main · Joe-GPT/quant-python-ai · 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
import sys

from dotenv import load_dotenv
from prompt_toolkit import PromptSession
from prompt_toolkit.history import InMemoryHistory
from rich.console import Console
from rich.panel import Panel
from rich.theme import Theme

from agent.quant_python_agent import QuantPythonAgent

theme = Theme(
    {
        "info": "cyan",
        "warning": "yellow",
        "stage": "bold magenta",
    }
)
console = Console(theme=theme)

BANNER = r"""
  ___                    _     ____        _   _                      _    ___
 / _ \ _   _  __ _ _ __ | |_  |  _ \ _   _| |_| |__   ___  _ __      / \  |_ _|
| | | | | | |/ _` | '_ \| __| | |_) | | | | __| '_ \ / _ \| '_ \    / _ \  | |
| |_| | |_| | (_| | | | | |_  |  __/| |_| | |_| | | | (_) | | | |  / ___ \ | |
 \__\_\\__,_|\__,_|_| |_|\__| |_|    \__, |\__|_| |_|\___/|_| |_| /_/   \_\___|
                                      |___/
"""


def main():
    load_dotenv()
    console.print(Panel(BANNER, style="bold cyan", subtitle="v0.1.0"))
    console.print("[dim]輸入投資研究任務，或輸入 /help 查看指令、/quit 離開[/dim]\n")

    agent = QuantPythonAgent(console)
    session = PromptSession(history=InMemoryHistory())

    while True:
        try:
            query = session.prompt(">> ")
        except (KeyboardInterrupt, EOFError):
            console.print("\n[dim]Bye![/dim]")
            sys.exit(0)

        query = query.strip()
        if not query:
            continue

        if query == "/quit":
            console.print("[dim]Bye![/dim]")
            break

        if query == "/help":
            console.print(
                Panel(
                    "[bold]/models[/bold]  - 列出可用模型（範例清單）\n"
                    "[bold]/model[/bold] [dim]<provider:model>[/dim] - 切換模型（例如 openai:gpt-4o-mini）\n"
                    "[bold]/quit[/bold]   - 離開程式\n"
                    "[bold]/help[/bold]   - 顯示此說明\n\n"
                    "直接輸入任務即可開始研究，例如:\n"
                    '  "整理台積電近期重大事件與財務重點，並判斷市場情緒"\n'
                    '  "比較 台積電 和 聯發科 的近期新聞情緒與財務指標差異"',
                    title="Help",
                )
            )
            continue

        if query == "/models":
            agent.list_models()
            continue

        if query.startswith("/model"):
            parts = query.split(maxsplit=1)
            if len(parts) < 2:
                console.print("[dim]Usage: /model <id>  (see /models for options)[/dim]")
            else:
                agent.set_model(parts[1])
            continue

        agent.run_mission(query)
        console.print()


if __name__ == "__main__":
    main()