fictionChat/main.py at main · yuanlin-yu/fictionChat · 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
import crew
from crew import projectCrew
from dotenv import load_dotenv
import os
import threading
import time
import sys

# 加载环境变量
load_dotenv()
file_chunk_num = int(os.environ.get('FICTION_CHUNK_NUM'))
finish_id = 0

# 显示等待符号的函数
def show_waiting_symbol():
    symbols = ["/", "-", "\\", "|"]
    i = 0
    while not done:
        sys.stdout.write("\r" + f"{f'读小说中,已完成:{round(finish_id/file_chunk_num*100,2)}% ' if crew.user_input['initial'] else '思考中'} " + symbols[i % len(symbols)])
        sys.stdout.flush()
        time.sleep(0.1)
        i += 1

# 运行 Crew 的函数
def run(character, question='', source_id=0, initial=False):
    global done
    done = False

    # 输入参数
    inputs = {
        "character": character,
        "question": question,
        "source_id": source_id,
        "initial": initial
    }
    crew.user_input = inputs

    # 启动等待符号线程
    waiting_thread = threading.Thread(target=show_waiting_symbol)
    waiting_thread.start()

    # 运行crew
    result = projectCrew().crew().kickoff(inputs=inputs)

    done = True
    waiting_thread.join()
    sys.stdout.write("\r" + " " * 20 + "\r")
    return result

def run2(character, character2, scene='', source_id=0, initial=False):
    global done
    done = False

    # 输入参数
    inputs = {
        "character": character,
        "character2": character2,
        "scene": scene,
        "source_id": source_id,
        "initial": initial
    }
    crew.user_input = inputs

    # 启动等待符号线程
    waiting_thread = threading.Thread(target=show_waiting_symbol)
    waiting_thread.start()

    # 运行crew
    result = projectCrew().crew().kickoff(inputs=inputs)

    done = True
    waiting_thread.join()
    sys.stdout.write("\r" + " " * 20 + "\r")
    return result

fiction_name = os.environ.get('FICTION_NAME')
author_name = os.environ.get('FICTION_AUTHOR_NAME')
character = os.environ.get('FICTION_CHAT_CHARACTER')
character2 = os.environ.get('FICTION_CHAT_CHARACTER2')

if __name__ == "__main__":
    print(f"** 欢迎使用fictionChat, 在这里您可以与小说人物畅聊 **\n")
    print(f"** 小说：《{fiction_name}》，作者：{author_name} **\n")

    if crew.crew_mode == 'single_character':
        print(f"** 小说人物：{character} **\n")
        # 数据读取
        if crew.isFictionSplit:
            for i in range(file_chunk_num):
                run(character=character, source_id=i, initial=True)
                finish_id = i+1
        else:
            run(character=character, initial=True)
        print('** 数据初始化完成! **\n')
        print('** 现在可以开始对话了！请输入您的问题(输入"exit"退出): **\n')
        while True:
            print("- 问题：")
            question = input()
            if question == 'exit':
                break
            result = run(character, question)
            print(result.raw+'\n')
    else:
        print(f"** 小说人物：{character}、{character2} **\n")
        # 数据读取
        if crew.isFictionSplit:
            for i in range(file_chunk_num):
                run2(character=character, character2=character2, source_id=i, initial=True)
                finish_id = i+1
        else:
            run2(character=character, character2=character2, initial=True)
        print('** 数据初始化完成! **\n')
        print("* 请输入场景：")
        scene = input()
        result = run2(character, character2, scene)
        print(result.raw+'\n')