diff --git a/api/controller/ChatComplete.py b/api/controller/ChatComplete.py index 150ac1e..a751d1e 100644 --- a/api/controller/ChatComplete.py +++ b/api/controller/ChatComplete.py @@ -60,7 +60,7 @@ class ChatCompleteTask: self.transatcion_id: Optional[str] = None self.point_cost: int = 0 if not self.is_system: - usage_res = await self.mwapi.chat_complete_start_transaction(self.user_id, "chatcomplete", + usage_res = await self.mwapi.ai_toolbox_start_transaction(self.user_id, "chatcomplete", question_tokens, extract_limit) self.transatcion_id = usage_res["transaction_id"] self.point_cost = usage_res["point_cost"] @@ -109,7 +109,7 @@ class ChatCompleteTask: self.result = chat_res if self.transatcion_id: - await self.mwapi.chat_complete_end_transaction(self.transatcion_id, chat_res["total_tokens"]) + await self.mwapi.ai_toolbox_end_transaction(self.transatcion_id, chat_res["total_tokens"]) await self._on_finished() except Exception as e: @@ -119,7 +119,7 @@ class ChatCompleteTask: traceback.print_exc() if self.transatcion_id: - await self.mwapi.chat_complete_cancel_transaction(self.transatcion_id, error=err_msg) + await self.mwapi.ai_toolbox_cancel_transaction(self.transatcion_id, error=err_msg) await self._on_error(e) finally: @@ -282,7 +282,7 @@ class ChatComplete: tokens = await tiktoken.get_tokens(question) try: - res = await mwapi.chat_complete_get_point_cost(user_id, "chatcomplete", tokens, extract_limit) + res = await mwapi.ai_toolbox_get_point_cost(user_id, "chatcomplete", tokens, extract_limit) return await utils.web.api_response(1, { "point_cost": res["point_cost"], "tokens": tokens, diff --git a/api/controller/EmbeddingSearch.py b/api/controller/EmbeddingSearch.py index ff75a15..ba55def 100644 --- a/api/controller/EmbeddingSearch.py +++ b/api/controller/EmbeddingSearch.py @@ -31,7 +31,7 @@ class EmbeddingSearch: if await embedding_search.should_update_page_index(): if request.get("caller") == "user": user_id = request.get("user") - usage_res = await mwapi.chat_complete_start_transaction(user_id, "embeddingpage") + usage_res = await mwapi.ai_toolbox_start_transaction(user_id, "embeddingpage") transatcion_id = usage_res.get("transaction_id") await embedding_search.prepare_update_index() @@ -52,7 +52,7 @@ class EmbeddingSearch: }) if transatcion_id: - await mwapi.chat_complete_end_transaction(transatcion_id, token_usage) + await mwapi.ai_toolbox_end_transaction(transatcion_id, token_usage) else: await ws.send_json({ 'event': 'done', @@ -71,7 +71,7 @@ class EmbeddingSearch: }, }) if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) except MediaWikiApiException as e: error_msg = "MediaWiki API error: %s" % str(e) print(error_msg, file=sys.stderr) @@ -86,7 +86,7 @@ class EmbeddingSearch: }, }) if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) except EmbeddingRunningException: error_msg = "Page index is running now" await ws.send_json({ @@ -98,7 +98,7 @@ class EmbeddingSearch: }, }) if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) except Exception as e: error_msg = str(e) print(error_msg, file=sys.stderr) @@ -112,7 +112,7 @@ class EmbeddingSearch: } }) if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) finally: await ws.close() else: @@ -123,7 +123,7 @@ class EmbeddingSearch: if await embedding_search.should_update_page_index(): if request.get("caller") == "user": user_id = request.get("user") - usage_res = await mwapi.chat_complete_start_transaction(user_id, "embeddingpage") + usage_res = await mwapi.ai_toolbox_start_transaction(user_id, "embeddingpage") transatcion_id = usage_res.get("transaction_id") await embedding_search.prepare_update_index() @@ -131,7 +131,7 @@ class EmbeddingSearch: token_usage = await embedding_search.update_page_index() if transatcion_id: - result = await mwapi.chat_complete_end_transaction(transatcion_id, token_usage) + result = await mwapi.ai_toolbox_end_transaction(transatcion_id, token_usage) return await utils.web.api_response(1, {"data_indexed": True}) else: @@ -139,7 +139,7 @@ class EmbeddingSearch: except MediaWikiPageNotFoundException: error_msg = "Page \"%s\" not found." % page_title if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) return await utils.web.api_response(-2, error={ "code": "page-not-found", @@ -153,7 +153,7 @@ class EmbeddingSearch: traceback.print_exc() if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) return await utils.web.api_response(-3, error={ "code": "mediawiki-api-error", @@ -164,7 +164,7 @@ class EmbeddingSearch: error_msg = "Page index is running now" if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) return await utils.web.api_response(-4, error={ "code": "page-index-running", @@ -177,7 +177,7 @@ class EmbeddingSearch: traceback.print_exc() if transatcion_id: - await mwapi.chat_complete_cancel_transaction(transatcion_id, error_msg) + await mwapi.ai_toolbox_cancel_transaction(transatcion_id, error_msg) return await utils.web.api_response(-1, error={ "code": "internal-server-error", diff --git a/api/controller/Index.py b/api/controller/Index.py index 54f8ced..7f0051f 100644 --- a/api/controller/Index.py +++ b/api/controller/Index.py @@ -79,6 +79,22 @@ class Index: "code": "internal-server-error", "message": error_msg }, request=request, http_status=500) + + @staticmethod + async def search_title(request: web.Request): + params = await utils.web.get_param(request, { + "kw": { + "required": True, + } + }) + + keyword = params.get("kw") + mwapi = MediaWikiApi.create() + + search_result = await mwapi.search_title(keyword) + return await utils.web.api_response(1, { + "titles": search_result + }, request=request) @staticmethod @utils.web.token_auth @@ -89,7 +105,7 @@ class Index: "type": int }, "title": { - "required": True, + "required": False, }, "module": { "required": False @@ -105,14 +121,18 @@ class Index: module = params.get("module") db = await DatabaseService.create(request.app) - async with PageTitleHelper(db) as page_title_helper, ConversationHelper(db) as conversation_helper: - page_id = await page_title_helper.get_page_id_by_title(page_title) - if page_id is None: - return await utils.web.api_response(-2, error={ - "code": "page-not-found", - "message": "Page not found.", - }, request=request, http_status=404) + page_id = None + if page_title is not None: + async with PageTitleHelper(db) as page_title_helper: + page_id = await page_title_helper.get_page_id_by_title(page_title) + if page_id is None: + return await utils.web.api_response(-2, error={ + "code": "page-not-found", + "message": "Page not found.", + }, request=request, http_status=404) + + async with ConversationHelper(db) as conversation_helper: conversation_list = await conversation_helper.get_conversation_list(user_id, module=module, page_id=page_id) conversation_result = [] @@ -121,6 +141,7 @@ class Index: conversation_result.append({ "id": result.id, "module": result.module, + "page_title": result.page_info.title if result.page_info is not None else None, "title": result.title, "description": result.description, "thumbnail": result.thumbnail, @@ -319,7 +340,7 @@ class Index: mwapi = MediaWikiApi.create() try: - user_info = await mwapi.chat_complete_user_info(user_id) + user_info = await mwapi.ai_toolbox_get_user_info(user_id) return await utils.web.api_response(1, user_info, request=request) except MediaWikiPageNotFoundException as e: return await utils.web.api_response(-2, error={ diff --git a/api/model/toolkit_ui/conversation.py b/api/model/toolkit_ui/conversation.py index 9c262a2..0f5d0f1 100644 --- a/api/model/toolkit_ui/conversation.py +++ b/api/model/toolkit_ui/conversation.py @@ -4,9 +4,10 @@ import time from typing import List, Optional import sqlalchemy from sqlalchemy import update -from sqlalchemy.orm import mapped_column, Mapped +from sqlalchemy.orm import mapped_column, relationship, Mapped from api.model.base import BaseModel +from api.model.toolkit_ui.page_title import PageTitleModel from service.database import DatabaseService @@ -19,14 +20,15 @@ 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.Integer, index=True, nullable=True) + page_id: Mapped[int] = mapped_column(sqlalchemy.ForeignKey("toolkit_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( sqlalchemy.Boolean, default=False, index=True) extra: Mapped[dict] = mapped_column(sqlalchemy.JSON, default={}) + page_info: Mapped[PageTitleModel] = relationship("PageTitleModel", lazy="joined") + class ConversationHelper: def __init__(self, dbs: DatabaseService): diff --git a/api/route.py b/api/route.py index 07565ae..dfb3b54 100644 --- a/api/route.py +++ b/api/route.py @@ -17,6 +17,7 @@ def init(app: web.Application): 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), diff --git a/service/mediawiki_api.py b/service/mediawiki_api.py index 4cbbbc0..f7c45c9 100644 --- a/service/mediawiki_api.py +++ b/service/mediawiki_api.py @@ -125,7 +125,7 @@ class MediaWikiApi: return ret - async def is_logged_in(self,): + async def is_logged_in(self): async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: params = { "action": "query", @@ -157,6 +157,10 @@ class MediaWikiApi: return data["query"]["tokens"][token_type + "token"] async def robot_login(self, username: str, password: str): + if await self.is_logged_in(): + self.login_time = time.time() + return True + token = await self.get_token("login") async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: post_data = { @@ -173,7 +177,7 @@ class MediaWikiApi: raise MediaWikiApiException(data["error"]["info"], data["error"]["code"]) if "result" not in data["login"] or data["login"]["result"] != "Success": - raise MediaWikiApiException("Login failed") + raise MediaWikiApiException("Login failed", data["login"]["result"]) self.login_time = time.time() self.login_identity = { @@ -191,12 +195,24 @@ class MediaWikiApi: print("刷新MW机器人账号登录状态") return await self.robot_login(self.login_identity["username"], self.login_identity["password"]) - async def chat_complete_user_info(self, user_id: int): + async def search_title(self, keyword: str) -> list[str]: + async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: + params = { + "action": "opensearch", + "search": keyword, + "namespace": 0, + "format": "json", + } + async with session.get(self.api_url, params=params, proxy=config.REQUEST_PROXY) as resp: + data = await resp.json() + return data[1] + + async def ai_toolbox_get_user_info(self, user_id: int): await self.refresh_login() async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: params = { - "action": "chatcompletebot", + "action": "aitoolboxbot", "method": "userinfo", "userid": user_id, "format": "json", @@ -210,14 +226,14 @@ class MediaWikiApi: else: raise MediaWikiApiException(data["error"]["info"], data["error"]["code"]) - return data["chatcompletebot"]["userinfo"] + return data["aitoolboxbot"]["userinfo"] - async def chat_complete_get_point_cost(self, user_id: int, user_action: str, tokens: Optional[int] = None, extractlines: Optional[int] = None) -> ChatCompleteGetPointUsageResponse: + async def ai_toolbox_get_point_cost(self, user_id: int, user_action: str, tokens: Optional[int] = None, extractlines: Optional[int] = None) -> ChatCompleteGetPointUsageResponse: await self.refresh_login() async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: post_data = { - "action": "chatcompletebot", + "action": "aitoolboxbot", "method": "reportusage", "step": "check", "userid": int(user_id) if user_id is not None else None, @@ -235,15 +251,15 @@ class MediaWikiApi: print(data) raise MediaWikiApiException(data["error"]["info"], data["error"]["code"]) - point_cost = int(data["chatcompletebot"]["reportusage"]["pointcost"] or 0) + point_cost = int(data["aitoolboxbot"]["reportusage"]["pointcost"] or 0) return ChatCompleteGetPointUsageResponse(point_cost=point_cost) - async def chat_complete_start_transaction(self, user_id: int, user_action: str, tokens: Optional[int] = None, extractlines: Optional[int] = None) -> ChatCompleteReportUsageResponse: + async def ai_toolbox_start_transaction(self, user_id: int, user_action: str, tokens: Optional[int] = None, extractlines: Optional[int] = None) -> ChatCompleteReportUsageResponse: await self.refresh_login() async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: post_data = { - "action": "chatcompletebot", + "action": "aitoolboxbot", "method": "reportusage", "step": "start", "userid": int(user_id) if user_id is not None else None, @@ -263,17 +279,17 @@ class MediaWikiApi: else: raise MediaWikiApiException(data["error"]["info"], data["error"]["code"]) - point_cost = int(data["chatcompletebot"]["reportusage"]["pointcost"] or 0) + point_cost = int(data["aitoolboxbot"]["reportusage"]["pointcost"] or 0) return ChatCompleteReportUsageResponse(point_cost=point_cost, - transaction_id=data["chatcompletebot"]["reportusage"]["transactionid"]) + transaction_id=data["aitoolboxbot"]["reportusage"]["transactionid"]) - async def chat_complete_end_transaction(self, transaction_id: str, tokens: Optional[int] = None): + async def ai_toolbox_end_transaction(self, transaction_id: str, tokens: Optional[int] = None): await self.refresh_login() try: async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: post_data = { - "action": "chatcompletebot", + "action": "aitoolboxbot", "method": "reportusage", "step": "end", "transactionid": transaction_id, @@ -288,17 +304,17 @@ class MediaWikiApi: if "error" in data: raise MediaWikiApiException(data["error"]["info"], data["error"]["code"]) - return data["chatcompletebot"]["reportusage"]["success"] + return data["aitoolboxbot"]["reportusage"]["success"] except Exception as e: print(e, file=sys.stderr) - async def chat_complete_cancel_transaction(self, transaction_id: str, error: Optional[str] = None): + async def ai_toolbox_cancel_transaction(self, transaction_id: str, error: Optional[str] = None): await self.refresh_login() try: async with aiohttp.ClientSession(cookie_jar=self.cookie_jar) as session: post_data = { - "action": "chatcompletebot", + "action": "aitoolboxbot", "method": "reportusage", "step": "cancel", "transactionid": transaction_id, @@ -313,6 +329,6 @@ class MediaWikiApi: if "error" in data: raise MediaWikiApiException(data["error"]["info"], data["error"]["code"]) - return data["chatcompletebot"]["reportusage"]["success"] + return data["aitoolboxbot"]["reportusage"]["success"] except Exception as e: print(e, file=sys.stderr) \ No newline at end of file