mirror of
https://gitee.com/270580156/weiyu.git
synced 2026-05-14 19:27:53 +00:00
147 lines
3.9 KiB
Python
147 lines
3.9 KiB
Python
'''
|
||
Author: jackning 270580156@qq.com
|
||
Date: 2024-08-29 18:21:14
|
||
LastEditors: jackning 270580156@qq.com
|
||
LastEditTime: 2024-08-31 10:18:58
|
||
Description: bytedesk.com https://github.com/Bytedesk/bytedesk
|
||
Please be aware of the BSL license restrictions before installing Bytedesk IM –
|
||
selling, reselling, or hosting Bytedesk IM as a service is a breach of the terms and automatically terminates your rights under the license.
|
||
仅支持企业内部员工自用，严禁私自用于销售、二次销售或者部署SaaS方式销售
|
||
Business Source License 1.1: https://github.com/Bytedesk/bytedesk/blob/main/LICENSE
|
||
contact: 270580156@qq.com
|
||
技术/商务联系：270580156@qq.com
|
||
Copyright (c) 2024 by bytedesk.com, All Rights Reserved.
|
||
'''
|
||
import json
|
||
import logging
|
||
from typing import List
|
||
import redis
|
||
from app.config import get_settings
|
||
from app.consts import DELETE_FILE_ERROR, DELETE_FILE_SUCCESS, PARSE_FILE_ERROR, PARSE_FILE_SUCCESS, ANSWER, ANSWER_FINISHED, pubsubChannel
|
||
#
|
||
redisClient = redis.Redis(host=get_settings().REDIS_HOST,
|
||
password=get_settings().REDIS_PASSWORD,
|
||
port=get_settings().REDIS_PORT,
|
||
decode_responses=True)
|
||
#
|
||
|
||
|
||
def publishParseFileSuccess(fileUid: str, docIds: List[str]) -> None:
|
||
content = json.dumps({
|
||
"fileUid": fileUid,
|
||
"docIds": docIds
|
||
})
|
||
message = json.dumps({
|
||
"type": PARSE_FILE_SUCCESS,
|
||
"content": content,
|
||
}, ensure_ascii=False)
|
||
defaultPublish(content=message)
|
||
|
||
|
||
def publishParseFileError(fileUid: str, errorMsg: str) -> None:
|
||
content = json.dumps({
|
||
"fileUid": fileUid,
|
||
"errorMsg": errorMsg,
|
||
})
|
||
message = json.dumps({
|
||
"type": PARSE_FILE_ERROR,
|
||
"content": content,
|
||
}, ensure_ascii=False)
|
||
defaultPublish(content=message)
|
||
|
||
|
||
def publishDeleteFileSuccess(fileUid: str) -> None:
|
||
content = json.dumps({
|
||
"fileUid": fileUid,
|
||
})
|
||
message = json.dumps({
|
||
"type": DELETE_FILE_SUCCESS,
|
||
"content": content,
|
||
}, ensure_ascii=False)
|
||
defaultPublish(content=message)
|
||
|
||
|
||
def publishDeleteFileError(fileUid: str, errorMsg: str) -> None:
|
||
content = json.dumps({
|
||
"fileUid": fileUid,
|
||
"errorMsg": errorMsg,
|
||
})
|
||
message = json.dumps({
|
||
"type": DELETE_FILE_ERROR,
|
||
"content": content,
|
||
}, ensure_ascii=False)
|
||
defaultPublish(content=message)
|
||
|
||
|
||
def publishAnswerMessage(
|
||
id: int,
|
||
uid: str,
|
||
threadTopic: str,
|
||
kbUid: str,
|
||
question: str,
|
||
answer: str,
|
||
model: str,
|
||
created: int) -> None:
|
||
#
|
||
content = json.dumps({
|
||
"id": id,
|
||
"uid": uid,
|
||
"threadTopic": threadTopic,
|
||
"kbUid": kbUid,
|
||
"question": question,
|
||
"answer": answer,
|
||
"model": model,
|
||
"created": created
|
||
})
|
||
message = json.dumps({
|
||
"type": ANSWER,
|
||
"content": content,
|
||
}, ensure_ascii=False)
|
||
defaultPublish(content=message)
|
||
return
|
||
|
||
|
||
def publishAnswerFinished(
|
||
id: int,
|
||
uid: str,
|
||
threadTopic: str,
|
||
kbUid: str,
|
||
question: str,
|
||
answer: str,
|
||
model: str,
|
||
created: int,
|
||
promptTokens: str,
|
||
completionTokens: str,
|
||
totalTokens: str) -> None:
|
||
#
|
||
content = json.dumps({
|
||
"id": id,
|
||
"uid": uid,
|
||
"threadTopic": threadTopic,
|
||
"kbUid": kbUid,
|
||
"question": question,
|
||
"answer": answer,
|
||
"model": model,
|
||
"created": created,
|
||
"promptTokens": promptTokens,
|
||
"completionTokens": completionTokens,
|
||
"totalTokens": totalTokens
|
||
})
|
||
message = json.dumps({
|
||
"type": ANSWER_FINISHED,
|
||
"content": content,
|
||
}, ensure_ascii=False)
|
||
defaultPublish(content=message)
|
||
return
|
||
|
||
#
|
||
|
||
|
||
def defaultPublish(content: str) -> None:
|
||
publish(pubsubChannel, content)
|
||
|
||
|
||
def publish(channel, message):
|
||
# logging.info(f'publish {message} to channel: {channel}')
|
||
redisClient.publish(channel, message)
|