完成动态加载的重构

api/route.py

@@ -1,41 +0,0 @@
-from aiohttp import web
-from api.controller.ChatComplete import ChatComplete
-
-from api.controller.Hanzi import Hanzi
-from api.controller.Index import Index
-from api.controller.Kanji import Kanji
-from api.controller.Hanja import Hanja
-from api.controller.EmbeddingSearch import EmbeddingSearch
-
-def init(app: web.Application):
-    app.router.add_routes([
-        web.route('*', '/hanzi/pinyin/', Hanzi.hanziToPinyin),
-        web.route('*', '/hanzi/split/', Hanzi.splitHanzi),
-
-        web.route('*', '/kanji/romaji/', Kanji.kanji2romaji),
-
-        web.route('*', '/hanja/romaja/', Hanja.hanja2roma),
-
-        web.route('*', '/title/info', Index.update_title_info),
-        web.route('*', '/title/search', Index.search_title),
-        web.route('*', '/user/info', Index.get_user_info),
-        web.route('*', '/conversation/list', Index.get_conversation_list),
-        web.route('*', '/conversation/info', Index.get_conversation_info),
-        web.route('POST', '/conversation/remove', Index.remove_conversation),
-        web.route('DELETE', '/conversation/remove', Index.remove_conversation),
-        web.route('POST', '/conversation/set_pinned', Index.set_conversation_pinned),
-        web.route('POST', '/conversation/set_title', Index.set_conversation_title),
-
-        web.route('*', '/embedding_search/index_page', EmbeddingSearch.index_page),
-        web.route('*', '/embedding_search/search', EmbeddingSearch.search),
-        web.route('POST', '/sys/embedding_search/title/update', EmbeddingSearch.sys_update_title_info),
-
-        web.route('*', '/chatcomplete/conversation_chunk/list', ChatComplete.get_conversation_chunk_list),
-        web.route('*', '/chatcomplete/conversation_chunk/info', ChatComplete.get_conversation_chunk),
-        web.route('POST', '/chatcomplete/conversation/fork', ChatComplete.fork_conversation),
-        web.route('POST', '/chatcomplete/message', ChatComplete.start_chat_complete),
-        web.route('GET', '/chatcomplete/message/stream', ChatComplete.chat_complete_stream),
-        web.route('POST', '/chatcomplete/get_point_cost', ChatComplete.get_point_cost),
-        web.route('*', '/chatcomplete/persona/list', ChatComplete.get_persona_list),
-        web.route('*', '/chatcomplete/persona/info', ChatComplete.get_persona_info),
-    ])

init/server/base_api.py

@@ -0,0 +1,9 @@
+from aiohttp import web
+from utils.server import register_server_module
+
+from server.route.base_api import register_route
+
+def init(app: web.Application):
+    register_route(app)
+
+register_server_module("base_api", init)

init/server/chat_complete.py

@@ -0,0 +1,22 @@
+from aiohttp import web
+from utils.server import register_server_module
+
+import init.service.mw_api as _  # Init mediawiki api
+import init.service.database as _  # Init database
+import init.service.tiktoken as _  # Init tiktoken
+
+import toolbox_ui as _  # Init toolbox ui
+import embedding_search as _  # Init embedding search
+
+# Auto create database tables
+from server.model.chat_complete.conversation import ConversationChunkModel as _
+from server.model.chat_complete.bot_persona_category import BotPersonaCategoryModel as _
+from server.model.chat_complete.bot_persona import BotPersonaModel as _
+
+# Route
+from server.route.chat_complete import register_route
+
+def init(app: web.Application):
+    register_route(app)
+
+register_server_module("chat_complete", init)

init/server/embedding_search.py

@@ -0,0 +1,17 @@
+from aiohttp import web
+from utils.server import register_server_module
+
+import init.service.database as _  # Init database
+import init.service.tiktoken as _  # Init tiktoken
+
+# Auto create database tables
+from server.model.embedding_search.title_collection import TitleCollectionModel as _
+from server.model.embedding_search.title_index import TitleIndexModel as _
+
+# Route
+from server.route.embedding_search import register_route
+
+def init(app: web.Application):
+    register_route(app)
+
+register_server_module("embedding_search", init)

init/server/toolbox_ui.py

@@ -0,0 +1,16 @@
+from aiohttp import web
+from utils.server import register_server_module
+
+import init.service.database as _  # Init database
+
+# Auto create database tables
+from server.model.toolbox_ui.page_title import PageTitleModel as _
+from server.model.toolbox_ui.conversation import ConversationModel as _
+
+# Route
+from server.route.toolbox_ui import register_route
+
+def init(app: web.Application):
+    register_route(app)
+
+register_server_module("toolbox_ui", init)

init/service/database.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+from aiohttp import web
+from utils.server import register_server_module
+
+from server.model.base import BaseModel
+from service.database import DatabaseService
+
+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 close_database(app: web.Application):
+    dbs = await DatabaseService.create(app)
+    await dbs.close()
+
+def init(app: web.Application):
+    app.on_startup.append(init_database)
+    app.on_cleanup.append(close_database)
+
+register_server_module("database", init)

init/service/mw_api.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+from aiohttp import web
+from utils.server import register_server_module
+
+from lib.config import Config
+from service.mediawiki_api import MediaWikiApi
+
+
+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"])
+    )
+
+
+def init(app: web.Application):
+    app.on_startup.append(init_mw_api)
+
+
+register_server_module("mediawiki_api", init)

init/service/tiktoken.py

@@ -0,0 +1,16 @@
+from aiohttp import web
+from utils.server import register_server_module
+
+from service.tiktoken import TikTokenService
+
+
+async def init_tiktoken(app: web.Application):
+    await TikTokenService.create()
+    print("Tiktoken model loaded.")
+
+
+async def init(app: web.Application):
+    app.on_startup.append(init_tiktoken)
+
+
+register_server_module("tiktoken", init)

main.py

@@ -1,29 +1,17 @@
 import sys
 import traceback
 from lib.config import Config
+from utils.server import get_server_modules
 
 Config.load_config("config.toml")
 
 from utils.local import loop, noawait
 from aiohttp import web
 import utils.local as local
-import api.route
+import server.route as 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(
@@ -65,42 +53,6 @@ async def error_handler(request, handler):
         )
 
 
-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 close_database(app: web.Application):
-    dbs = await DatabaseService.create(app)
-    await dbs.close()
-
-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()
 
@@ -113,20 +65,29 @@ if __name__ == "__main__":
         ]
     )
 
-    if Config.get("database.host"):
-        app.on_startup.append(init_database)
-        app.on_cleanup.append(close_database)
+    import init.server.base_api as _
+
+    # Dynamically import modules
+    if Config.get("toolbox_ui.enable"):
+        import init.server.toolbox_ui as _
+
+    if Config.get("embedding.enable"):
+        import init.server.embedding_search as _
+
+    if Config.get("chatcomplete.enable"):
+        import init.server.chat_complete as _
 
-    if Config.get("mediawiki.api_endpoint"):
-        app.on_startup.append(init_mw_api)
+    # Initialize server modules
+    server_modules = get_server_modules()
+    print(server_modules)
 
-    if Config.get("chatcomplete.enabled"):
-        app.on_startup.append(init_tiktoken)
+    for server_module in server_modules:
+        server_module["init"](app)
 
     app.on_shutdown.append(stop_noawait_pool)
 
     app.router.add_route("*", "/", index)
-    api.route.init(app)
+    api_route.init(app)
 
     server_port = Config.get("port", 8144, int)
 

server/controller/ChatComplete.py

@@ -2,13 +2,13 @@ from __future__ import annotations
 import asyncio
 import time
 import traceback
-from api.controller.task.ChatCompleteTask import ChatCompleteTask
-from api.model.base import clone_model
-from api.model.chat_complete.bot_persona import BotPersonaHelper
-from api.model.toolkit_ui.conversation import ConversationHelper
+from server.controller.task.ChatCompleteTask import ChatCompleteTask
+from server.model.base import clone_model
+from server.model.chat_complete.bot_persona import BotPersonaHelper
+from server.model.toolbox_ui.conversation import ConversationHelper
 from utils.local import noawait
 from aiohttp import web
-from api.model.chat_complete.conversation import ConversationChunkHelper, ConversationModel, ConversationChunkModel
+from server.model.chat_complete.conversation import ConversationChunkHelper, ConversationModel, ConversationChunkModel
 from service.chat_complete import ChatCompleteQuestionTooLongException, ChatCompleteServiceResponse
 from service.database import DatabaseService
 from service.mediawiki_api import MediaWikiApi, MediaWikiPageNotFoundException, MediaWikiUserNoEnoughPointsException

server/controller/EmbeddingSearch.py

@@ -1,8 +1,8 @@
 import sys
 import traceback
 from aiohttp import web
-from api.model.embedding_search.title_collection import TitleCollectionHelper
-from api.model.embedding_search.title_index import TitleIndexHelper
+from server.model.embedding_search.title_collection import TitleCollectionHelper
+from server.model.embedding_search.title_index import TitleIndexHelper
 from service.database import DatabaseService
 from service.embedding_search import EmbeddingRunningException, EmbeddingSearchService
 from service.mediawiki_api import MediaWikiApi, MediaWikiApiException, MediaWikiPageNotFoundException

server/controller/ToolboxIndex.py

@@ -2,15 +2,15 @@ import sys
 import time
 import traceback
 from aiohttp import web
-from api.model.toolkit_ui.conversation import ConversationHelper
-from api.model.toolkit_ui.page_title import PageTitleHelper
+from server.model.toolbox_ui.conversation import ConversationHelper
+from server.model.toolbox_ui.page_title import PageTitleHelper
 from service.database import DatabaseService
 from service.event import EventService
 from service.mediawiki_api import MediaWikiApi, MediaWikiApiException, MediaWikiPageNotFoundException
 import utils.web
 
 
-class Index:
+class ToolboxIndex:
     @staticmethod
     @utils.web.token_auth
     async def update_title_info(request: web.Request):

server/model/chat_complete/bot_persona.py

@@ -2,13 +2,13 @@ from __future__ import annotations
 import math
 from typing import Optional
 import sqlalchemy
-from api.model.base import BaseHelper, BaseModel
+from server.model.base import BaseHelper, BaseModel
 
 import sqlalchemy
 from sqlalchemy import select, update
 from sqlalchemy.orm import mapped_column, relationship, load_only, Mapped
 
-from api.model.chat_complete.bot_persona_category import BotPersonaCategoryModel
+from server.model.chat_complete.bot_persona_category import BotPersonaCategoryModel
 from service.database import DatabaseService
 
 

server/model/chat_complete/bot_persona_category.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import sqlalchemy
-from api.model.base import BaseHelper, BaseModel
+from server.model.base import BaseHelper, BaseModel
 
 import sqlalchemy
 from sqlalchemy import select, update

server/model/chat_complete/conversation.py

@@ -5,8 +5,8 @@ import sqlalchemy
 from sqlalchemy import select, update
 from sqlalchemy.orm import mapped_column, relationship, Mapped
 
-from api.model.base import BaseHelper, BaseModel
-from api.model.toolkit_ui.conversation import ConversationModel
+from server.model.base import BaseHelper, BaseModel
+from server.model.toolbox_ui.conversation import ConversationModel
 from service.database import DatabaseService
 from service.event import EventService
 

server/model/embedding_search/page_index.py

@@ -3,7 +3,7 @@ import hashlib
 from typing import Optional, Type
 
 import asyncpg
-from api.model.base import BaseModel
+from server.model.base import BaseModel
 import numpy as np
 import sqlalchemy
 from lib.config import Config

server/model/embedding_search/title_collection.py

@@ -3,7 +3,7 @@ import sqlalchemy
 from sqlalchemy import select, update, delete
 from sqlalchemy.orm import mapped_column, Mapped
 
-from api.model.base import BaseHelper, BaseModel
+from server.model.base import BaseHelper, BaseModel
 from service.database import DatabaseService
 
 class TitleCollectionModel(BaseModel):

server/model/embedding_search/title_index.py

@@ -9,7 +9,7 @@ from sqlalchemy.orm import mapped_column, relationship, Mapped, deferred, defer
 from sqlalchemy.ext.asyncio import AsyncEngine
 
 from lib.config import Config
-from api.model.base import BaseHelper, BaseModel
+from server.model.base import BaseHelper, BaseModel
 from service.database import DatabaseService
 
 embedding_vector_size = Config.get("embedding.vector_size", 1536, int)

server/model/toolbox_ui/conversation.py

@@ -6,13 +6,13 @@ import sqlalchemy
 from sqlalchemy import update
 from sqlalchemy.orm import mapped_column, relationship, Mapped
 
-from api.model.base import BaseHelper, BaseModel
-from api.model.toolkit_ui.page_title import PageTitleModel
+from server.model.base import BaseHelper, BaseModel
+from server.model.toolbox_ui.page_title import PageTitleModel
 from service.database import DatabaseService
 
 
 class ConversationModel(BaseModel):
-    __tablename__ = "toolkit_ui_conversation"
+    __tablename__ = "toolbox_ui_conversation"
 
     id: Mapped[int] = mapped_column(sqlalchemy.Integer, primary_key=True, autoincrement=True)
     user_id: Mapped[int] = mapped_column(sqlalchemy.Integer, index=True)
@@ -20,7 +20,7 @@ class ConversationModel(BaseModel):
     title: Mapped[str] = mapped_column(sqlalchemy.String(255), nullable=True)
     thumbnail: Mapped[str] = mapped_column(sqlalchemy.Text(), nullable=True)
     description: Mapped[str] = mapped_column(sqlalchemy.Text(), nullable=True)
-    page_id: Mapped[int] = mapped_column(sqlalchemy.ForeignKey("toolkit_ui_page_title.page_id"), index=True, nullable=True)
+    page_id: Mapped[int] = mapped_column(sqlalchemy.ForeignKey("toolbox_ui_page_title.page_id"), index=True, nullable=True)
     rev_id: Mapped[int] = mapped_column(sqlalchemy.Integer, nullable=True)
     updated_at: Mapped[int] = mapped_column(sqlalchemy.BigInteger, index=True)
     pinned: Mapped[bool] = mapped_column(

server/model/toolbox_ui/page_title.py

@@ -6,12 +6,12 @@ import sqlalchemy
 from sqlalchemy import select, update
 from sqlalchemy.orm import mapped_column, Mapped
 
-from api.model.base import BaseHelper, BaseModel
+from server.model.base import BaseHelper, BaseModel
 from service.database import DatabaseService
 
 
 class PageTitleModel(BaseModel):
-    __tablename__ = "toolkit_ui_page_title"
+    __tablename__ = "toolbox_ui_page_title"
 
     id: Mapped[int] = mapped_column(
         sqlalchemy.Integer, primary_key=True, autoincrement=True)

server/route/base_api.py

@@ -0,0 +1,15 @@
+from aiohttp import web
+
+from server.controller.Hanja import Hanja
+from server.controller.Hanzi import Hanzi
+from server.controller.Kanji import Kanji
+
+def register_route(app: web.Application):
+    app.router.add_routes([
+        web.route('*', '/hanzi/pinyin/', Hanzi.hanziToPinyin),
+        web.route('*', '/hanzi/split/', Hanzi.splitHanzi),
+
+        web.route('*', '/kanji/romaji/', Kanji.kanji2romaji),
+
+        web.route('*', '/hanja/romaja/', Hanja.hanja2roma),
+    ])

server/route/chat_complete.py

@@ -0,0 +1,15 @@
+from aiohttp import web
+
+from server.controller.ChatComplete import ChatComplete
+
+def register_route(app: web.Application):
+    app.router.add_routes([
+        web.route('*', '/chatcomplete/conversation_chunk/list', ChatComplete.get_conversation_chunk_list),
+        web.route('*', '/chatcomplete/conversation_chunk/info', ChatComplete.get_conversation_chunk),
+        web.route('POST', '/chatcomplete/conversation/fork', ChatComplete.fork_conversation),
+        web.route('POST', '/chatcomplete/message', ChatComplete.start_chat_complete),
+        web.route('GET', '/chatcomplete/message/stream', ChatComplete.chat_complete_stream),
+        web.route('POST', '/chatcomplete/get_point_cost', ChatComplete.get_point_cost),
+        web.route('*', '/chatcomplete/persona/list', ChatComplete.get_persona_list),
+        web.route('*', '/chatcomplete/persona/info', ChatComplete.get_persona_info),
+    ])

server/route/embedding_search.py

@@ -0,0 +1,10 @@
+from aiohttp import web
+
+from server.controller.EmbeddingSearch import EmbeddingSearch
+
+def register_route(app: web.Application):
+    app.router.add_routes([
+        web.route('*', '/embedding_search/index_page', EmbeddingSearch.index_page),
+        web.route('*', '/embedding_search/search', EmbeddingSearch.search),
+        web.route('POST', '/sys/embedding_search/title/update', EmbeddingSearch.sys_update_title_info),
+    ])

server/route/toolbox_ui.py

@@ -0,0 +1,16 @@
+from aiohttp import web
+
+from server.controller.ToolboxIndex import ToolboxIndex
+
+def register_route(app: web.Application):
+    app.router.add_routes([
+        web.route('*', '/title/info', ToolboxIndex.update_title_info),
+        web.route('*', '/title/search', ToolboxIndex.search_title),
+        web.route('*', '/user/info', ToolboxIndex.get_user_info),
+        web.route('*', '/conversation/list', ToolboxIndex.get_conversation_list),
+        web.route('*', '/conversation/info', ToolboxIndex.get_conversation_info),
+        web.route('POST', '/conversation/remove', ToolboxIndex.remove_conversation),
+        web.route('DELETE', '/conversation/remove', ToolboxIndex.remove_conversation),
+        web.route('POST', '/conversation/set_pinned', ToolboxIndex.set_conversation_pinned),
+        web.route('POST', '/conversation/set_title', ToolboxIndex.set_conversation_title),
+    ])

service/chat_complete.py

@@ -3,13 +3,13 @@ import time
 import traceback
 from typing import Optional, Tuple, TypedDict
 
-from api.model.chat_complete.bot_persona import BotPersonaHelper
-from api.model.chat_complete.conversation import (
+from server.model.chat_complete.bot_persona import BotPersonaHelper
+from server.model.chat_complete.conversation import (
     ConversationChunkHelper,
     ConversationChunkModel,
 )
 import sys
-from api.model.toolkit_ui.conversation import ConversationHelper, ConversationModel
+from server.model.toolbox_ui.conversation import ConversationHelper, ConversationModel
 
 from lib.config import Config
 import utils.config, utils.web

service/embedding_search.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 from typing import Optional, TypedDict
 
 import sqlalchemy
-from api.model.embedding_search.title_collection import (
+from server.model.embedding_search.title_collection import (
     TitleCollectionHelper,
     TitleCollectionModel,
 )
-from api.model.embedding_search.title_index import TitleIndexHelper, TitleIndexModel
-from api.model.embedding_search.page_index import PageIndexHelper
+from server.model.embedding_search.title_index import TitleIndexHelper, TitleIndexModel
+from server.model.embedding_search.page_index import PageIndexHelper
 from service.database import DatabaseService
 from service.mediawiki_api import MediaWikiApi
 from service.openai_api import OpenAIApi

test/create_token.py

@@ -2,7 +2,7 @@ import asyncio
 import base
 
 from sqlalchemy import select
-from api.model.embedding_search.title_index import TitleIndexModel
+from server.model.embedding_search.title_index import TitleIndexModel
 import utils.local as local
 from service.database import DatabaseService
 

test/title_index.py

@@ -2,7 +2,7 @@ import asyncio
 import base
 
 from sqlalchemy import select
-from api.model.embedding_search.title_index import TitleIndexModel
+from server.model.embedding_search.title_index import TitleIndexModel
 import utils.local as local
 from service.database import DatabaseService
 

utils/server.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+from typing import TypedDict
+
+class ServerModuleItem(TypedDict):
+    name: str
+    init: callable
+
+_SERVER_MODULES: list[ServerModuleItem] = []
+
+def register_server_module(name: str, init_method: callable):
+    global _SERVER_MODULES
+
+    for module in _SERVER_MODULES:
+        if module["name"] == name:
+            return
+        
+    _SERVER_MODULES.append({
+        "name": name,
+        "init": init_method
+    })
+
+def get_server_modules() -> list[ServerModuleItem]:
+    return _SERVER_MODULES