Skip to content

Commit 6ed31bf

Browse files
committed
fix receiver
1 parent 0a7ebfb commit 6ed31bf

File tree

10 files changed

+513
-570
lines changed

10 files changed

+513
-570
lines changed

app/receiver/discord/__init__.py

+220
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
# -*- coding: utf-8 -*-
2+
# @Time : 2023/10/18 下午10:46
3+
# @Author : sudoskys
4+
# @File : __init__.py.py
5+
# @Software: PyCharm
6+
import atexit
7+
from typing import List
8+
9+
import hikari
10+
from hikari.impl import ProxySettings
11+
from loguru import logger
12+
13+
from app.middleware.llm_task import OpenaiMiddleware
14+
from app.receiver import function
15+
from app.receiver.receiver_client import BaseReceiver, BaseSender
16+
from app.setting.discord import BotSetting
17+
from llmkira.kv_manager.file import File
18+
from llmkira.openai import OpenAIResult
19+
from llmkira.openai.cell import Message
20+
from llmkira.sdk.utils import sync
21+
from llmkira.task import Task, TaskHeader
22+
23+
__receiver__ = "discord_hikari"
24+
25+
from llmkira.task.schema import Location, EventMessage
26+
27+
discord_rest: hikari.RESTApp = hikari.RESTApp(
28+
proxy_settings=ProxySettings(url=BotSetting.proxy_address)
29+
)
30+
31+
32+
class DiscordSender(BaseSender):
33+
"""
34+
平台路由
35+
"""
36+
37+
def __init__(self):
38+
self.bot = None
39+
40+
def acquire(self):
41+
self.bot: hikari.impl.RESTClientImpl = discord_rest.acquire(
42+
token=BotSetting.token, token_type=hikari.TokenType.BOT
43+
)
44+
45+
async def file_forward(self, receiver: Location, file_list: List[File]):
46+
for file_obj in file_list:
47+
# DATA
48+
file_data: bytes = await file_obj.download_file()
49+
if not file_data:
50+
logger.error(f"file not found {receiver.user_id}")
51+
continue
52+
file_warp = hikari.files.Bytes(file_data, file_obj.file_name, mimetype=None)
53+
if file_obj.file_name.endswith((".jpg", ".png")):
54+
async with self.bot as client:
55+
client: hikari.impl.RESTClientImpl
56+
_reply = None
57+
if receiver.thread_id != receiver.chat_id:
58+
_reply = await client.fetch_message(
59+
channel=int(receiver.thread_id),
60+
message=int(receiver.message_id),
61+
)
62+
await client.create_message(
63+
channel=int(receiver.thread_id),
64+
embed=hikari.EmbedImage(
65+
resource=file_warp,
66+
),
67+
reply=_reply,
68+
)
69+
else:
70+
async with self.bot as client:
71+
client: hikari.impl.RESTClientImpl
72+
_reply = None
73+
if receiver.thread_id != receiver.chat_id:
74+
_reply = await client.fetch_message(
75+
channel=int(receiver.thread_id),
76+
message=int(receiver.message_id),
77+
)
78+
await client.create_message(
79+
channel=int(receiver.thread_id),
80+
attachment=file_warp,
81+
reply=_reply,
82+
)
83+
84+
async def forward(self, receiver: Location, message: List[EventMessage]):
85+
"""
86+
插件专用转发,是Task通用类型
87+
"""
88+
for item in message:
89+
await self.file_forward(receiver=receiver, file_list=item.files)
90+
async with self.bot as client:
91+
client: hikari.impl.RESTClientImpl
92+
_reply = None
93+
if receiver.thread_id != receiver.chat_id:
94+
_reply = await client.fetch_message(
95+
channel=int(receiver.thread_id),
96+
message=int(receiver.message_id)
97+
if receiver.message_id
98+
else None,
99+
)
100+
await client.create_message(
101+
channel=int(receiver.thread_id) if receiver.thread_id else None,
102+
content=item.text,
103+
reply=_reply,
104+
)
105+
106+
async def reply(
107+
self, receiver: Location, messages: List[Message], reply_to_message: bool = True
108+
):
109+
"""
110+
模型直转发,Message是Openai的类型
111+
"""
112+
for item in messages:
113+
# raw_message = await self.loop_turn_from_openai(platform_name=__receiver__, message=item, locate=receiver)
114+
event_message = EventMessage.from_openai_message(
115+
message=item, locate=receiver
116+
)
117+
await self.file_forward(receiver=receiver, file_list=event_message.files)
118+
if not event_message.text:
119+
continue
120+
async with self.bot as client:
121+
client: hikari.impl.RESTClientImpl
122+
_reply = None
123+
if receiver.thread_id != receiver.chat_id:
124+
_reply = await client.fetch_message(
125+
channel=int(receiver.thread_id),
126+
message=int(receiver.message_id)
127+
if receiver.message_id
128+
else None,
129+
)
130+
await client.create_message(
131+
channel=int(receiver.thread_id),
132+
content=event_message.text,
133+
reply=_reply,
134+
)
135+
return logger.trace("reply message")
136+
137+
async def error(self, receiver: Location, text):
138+
async with self.bot as client:
139+
client: hikari.impl.RESTClientImpl
140+
_reply = None
141+
if receiver.thread_id != receiver.chat_id:
142+
_reply = await client.fetch_message(
143+
channel=int(receiver.thread_id),
144+
message=int(receiver.message_id) if receiver.message_id else None,
145+
)
146+
await client.create_message(
147+
channel=int(receiver.thread_id), content=text, reply=_reply
148+
)
149+
150+
async def function(
151+
self,
152+
receiver: Location,
153+
task: TaskHeader,
154+
llm: OpenaiMiddleware,
155+
llm_result: OpenAIResult,
156+
):
157+
tool_calls = llm_result.default_message.tool_calls
158+
certify_needed_map = await self.push_task_create_message(
159+
llm_result=llm_result, receiver=receiver, tool_calls=tool_calls
160+
)
161+
new_receiver = task.receiver.model_copy()
162+
new_receiver.platform = __receiver__
163+
"""更新接收者为当前平台,便于创建的函数消息能返回到正确的客户端"""
164+
165+
new_sign = task.task_sign.update_tool_calls(
166+
tool_calls=tool_calls,
167+
certify_needed_map=certify_needed_map,
168+
)
169+
"""克隆元数据为当前平台"""
170+
171+
logger.debug(
172+
f"Sender.function:Create a new task {function.__receiver__} to {new_receiver} // {new_sign}"
173+
)
174+
await Task.create_and_send(
175+
queue_name=function.__receiver__,
176+
task=TaskHeader.from_function(
177+
llm_response=llm_result,
178+
task_sign=new_sign,
179+
receiver=new_receiver,
180+
message=task.message,
181+
),
182+
)
183+
"""发送到函数处理接收器"""
184+
185+
186+
__sender__ = DiscordSender()
187+
188+
189+
class DiscordReceiver(BaseReceiver):
190+
"""
191+
receive message from telegram
192+
"""
193+
194+
async def discord(self):
195+
self.set_core(sender=__sender__, task=Task(queue=__receiver__))
196+
if not BotSetting.available:
197+
logger.warning("Receiver Runtime:Discord Setting empty")
198+
return None
199+
await discord_rest.start()
200+
__sender__.acquire()
201+
try:
202+
logger.success("Receiver Runtime:Discord start")
203+
await self.task.consuming_task(self.on_message)
204+
except KeyboardInterrupt:
205+
logger.warning("Discord Receiver shutdown")
206+
except Exception as e:
207+
logger.exception(e)
208+
raise e
209+
finally:
210+
await discord_rest.close()
211+
212+
213+
if BotSetting.available:
214+
215+
@atexit.register
216+
def __clean():
217+
try:
218+
sync(discord_rest.close())
219+
except Exception as e:
220+
logger.warning(f"Discord Receiver cleaning failed when exiting \n{e}")

app/receiver/function.py

+2
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,8 @@ async def run_pending_task(task: TaskHeader, pending_task: ToolCall):
181181
)
182182
# Get Env
183183
env_all = await EnvManager(user_id=task.receiver.uid).read_env()
184+
if not env_all:
185+
env_all = {}
184186
env_map = {}
185187
for require in _tool_obj.env_list:
186188
env_map[require] = env_all.get(require, None)

0 commit comments

Comments
 (0)