2024-07-07 01:58:23 +08:00
|
|
|
import os
|
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
|
from core.history import HistoryManager
|
|
|
|
|
from config import Config
|
2024-07-07 17:21:07 +08:00
|
|
|
from models import select_model
|
2024-07-07 01:58:23 +08:00
|
|
|
|
|
|
|
|
load_dotenv()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
|
config = Config("config/base.yaml")
|
2024-07-07 17:21:07 +08:00
|
|
|
model = select_model(config)
|
2024-07-07 01:58:23 +08:00
|
|
|
|
2024-07-07 17:21:07 +08:00
|
|
|
print(f"[{config.model_provider}:{config.get('model_name', 'default')}] Type 'exit' to quit")
|
2024-07-07 01:58:23 +08:00
|
|
|
|
|
|
|
|
history_manager = HistoryManager()
|
|
|
|
|
while True:
|
|
|
|
|
message = input("\nUser: ")
|
|
|
|
|
if message == "exit":
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
messages = history_manager.add_user(message)
|
|
|
|
|
response = model.predict(messages, stream=config.stream)
|
|
|
|
|
|
|
|
|
|
if config.stream:
|
|
|
|
|
content = ""
|
|
|
|
|
print(f"AI: ", end='', flush=True)
|
|
|
|
|
for chunk in response:
|
|
|
|
|
content += chunk.content
|
|
|
|
|
print(f"{chunk.content}", end='', flush=True)
|
|
|
|
|
print()
|
|
|
|
|
else:
|
|
|
|
|
content = response.content
|
|
|
|
|
print(f"AI: {content}")
|
|
|
|
|
|
|
|
|
|
history_manager.add_ai(content)
|