You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

78 lines
2.7 KiB
Python

from config import Config
Config.load_config("config.toml")
from local import loop, noawait
from aiohttp import web
import local
import api.route
import utils.web
from service.database import DatabaseService
from service.mediawiki_api import MediaWikiApi
# Auto create Table
from api.model.base import BaseModel
from api.model.toolkit_ui.page_title import PageTitleModel as _
from api.model.toolkit_ui.conversation import ConversationModel as _
from api.model.chat_complete.conversation import ConversationChunkModel as _
from api.model.chat_complete.bot_persona_category import BotPersonaCategoryModel as _
from api.model.chat_complete.bot_persona import BotPersonaModel as _
from api.model.embedding_search.title_collection import TitleCollectionModel as _
from api.model.embedding_search.title_index import TitleIndexModel as _
from service.tiktoken import TikTokenService
async def index(request: web.Request):
return await utils.web.api_response(1, data={"message": "Isekai toolkit API"}, request=request)
async def init_mw_api(app: web.Application):
mw_api = MediaWikiApi.create()
bot_username = Config.get("mediawiki.bot_username", "", str)
bot_password = Config.get("mediawiki.bot_password", "", str)
if bot_username and bot_password:
try:
await mw_api.robot_login(bot_username, bot_password)
except Exception as e:
print("Cannot login to Robot account, please check config.")
site_meta = await mw_api.get_site_meta()
print("Connected to Wiki %s, Robot username: %s" % (site_meta["sitename"], site_meta["user"]))
async def init_database(app: web.Application):
dbs = await DatabaseService.create(app)
print("Database connected.")
async with dbs.engine.begin() as conn:
await conn.run_sync(BaseModel.metadata.create_all)
async def init_tiktoken(app: web.Application):
await TikTokenService.create()
print("Tiktoken model loaded.")
async def stop_noawait_pool(app: web.Application):
await noawait.end()
if __name__ == '__main__':
local.debug = Config.get("server.debug", False, bool)
app = web.Application()
if Config.get("database.host"):
app.on_startup.append(init_database)
if Config.get("mediawiki.api_endpoint"):
app.on_startup.append(init_mw_api)
if Config.get("chatcomplete.enabled"):
app.on_startup.append(init_tiktoken)
app.on_shutdown.append(stop_noawait_pool)
app.router.add_route('*', '/', index)
api.route.init(app)
server_port = Config.get("port", 8144, int)
web.run_app(app, host='0.0.0.0', port=server_port, loop=loop)