mirror of
https://github.com/JasonYANG170/CodeGeeX4.git
synced 2024-11-23 12:16:33 +00:00
229 lines
8.4 KiB
Python
229 lines
8.4 KiB
Python
import chainlit as cl
|
||
from chainlit.input_widget import Slider,Switch
|
||
import json
|
||
import re
|
||
|
||
from llm.api.codegeex4 import codegeex4
|
||
from prompts.base_prompt import (
|
||
get_cur_base_user_prompt,
|
||
build_message_list,
|
||
tools_choose_prompt,
|
||
tools_input_prompt
|
||
)
|
||
from utils.bingsearch import bing_search_prompt
|
||
from utils.tools import unzip_file, get_project_files_with_content,clone_repo,is_valid_json
|
||
|
||
def tools_choose_agent(input_text):
|
||
tools_prompt = tools_choose_prompt+tools_input_prompt.format(input_text=input_text)
|
||
message_list = build_message_list(tools_prompt)
|
||
judge_tmp = codegeex4(
|
||
messages_list=message_list,
|
||
temperature=0.2,
|
||
top_p=0.95,
|
||
)
|
||
judge_context = ""
|
||
for part in judge_tmp:
|
||
judge_context += part
|
||
attempt = 1
|
||
max_attempts = 10
|
||
while not is_valid_json(judge_context) and attempt <= max_attempts:
|
||
judge_tmp = codegeex4(
|
||
messages_list=message_list,
|
||
temperature=0.2,
|
||
top_p=0.95,
|
||
)
|
||
judge_context = ""
|
||
for part in judge_tmp:
|
||
judge_context += part
|
||
attempt += 1
|
||
match = re.search(r'\{.*\}', judge_context, re.DOTALL)
|
||
if match:
|
||
dict_str = match.group()
|
||
response = json.loads(dict_str)
|
||
else:
|
||
response = json.loads(judge_context)
|
||
tool_name = response["tool"]["name"]
|
||
return tool_name
|
||
|
||
@cl.set_chat_profiles
|
||
async def chat_profile():
|
||
return [
|
||
cl.ChatProfile(
|
||
name="联网聊天",
|
||
markdown_description="聊天demo：支持多轮对话。支持联网回答用户问题。默认联网，如不联网在输入框左边关闭联网功能。",
|
||
starters=[
|
||
cl.Starter(
|
||
label="请你用python写一个快速排序。",
|
||
message="请你用python写一个快速排序。",
|
||
),
|
||
cl.Starter(
|
||
label="请你介绍一下自己。",
|
||
message="请你介绍一下自己。",
|
||
),
|
||
cl.Starter(
|
||
label="用 Python 编写一个脚本来自动发送每日电子邮件报告，并指导我如何进行设置。",
|
||
message="用 Python 编写一个脚本来自动发送每日电子邮件报告，并指导我如何进行设置。",
|
||
),
|
||
cl.Starter(
|
||
label="我是一个python初学者，请你告诉我怎么才能学好python。",
|
||
message="我是一个python初学者，请你告诉我怎么才能学好python。",
|
||
),
|
||
],
|
||
),
|
||
cl.ChatProfile(
|
||
name="项目问答",
|
||
markdown_description="项目级能力demo：支持上传本地zip压缩包项目，支持输入GitHub链接项目，可以进行项目问答和对项目进行修改。",
|
||
),
|
||
]
|
||
|
||
@cl.on_settings_update
|
||
async def setup_agent(settings):
|
||
temperature = settings["temperature"]
|
||
top_p = settings["top_p"]
|
||
is_online = settings["is_online"]
|
||
cl.user_session.set("temperature", temperature)
|
||
cl.user_session.set("top_p", top_p)
|
||
cl.user_session.set("is_online", is_online)
|
||
@cl.on_chat_start
|
||
async def start():
|
||
settings = await cl.ChatSettings(
|
||
[
|
||
Slider(
|
||
id="temperature",
|
||
label="CodeGeeX4 - Temperature",
|
||
initial=0.2,
|
||
min=0,
|
||
max=1,
|
||
step=0.1,
|
||
),
|
||
Slider(
|
||
id="top_p",
|
||
label="CodeGeeX4 - top_p",
|
||
initial=0.95,
|
||
min=0,
|
||
max=1,
|
||
step=0.1,
|
||
),
|
||
Switch(
|
||
id="is_online",
|
||
label="CodeGeeX4 - is_online",
|
||
initial=True
|
||
),
|
||
]
|
||
).send()
|
||
temperature = settings["temperature"]
|
||
top_p = settings["top_p"]
|
||
is_online = settings["is_online"]
|
||
cl.user_session.set("temperature", temperature)
|
||
cl.user_session.set("top_p", top_p)
|
||
cl.user_session.set("is_online", is_online)
|
||
cl.user_session.set("message_history", [])
|
||
chat_profile = cl.user_session.get("chat_profile")
|
||
extract_dir = "repodata"
|
||
if chat_profile == "项目问答":
|
||
res = await cl.AskActionMessage(
|
||
content="请选择项目上传方式",
|
||
actions=[
|
||
cl.Action(name="zip", value="zip", label="本地上传zip文件"),
|
||
cl.Action(name="url", value="url", label="上传GitHub链接"),
|
||
],
|
||
).send()
|
||
if res.get("value") == "url":
|
||
repo_path =None
|
||
while repo_path == None:
|
||
res = await cl.AskUserMessage(content="请你在下面消息框中提供GitHub仓库URL? ex：https://github.com/THUDM/CodeGeeX4", timeout=3600).send()
|
||
if res:
|
||
repo_path = clone_repo(res['output'],extract_dir)
|
||
if repo_path is None:
|
||
await cl.Message(
|
||
content=f"您的github链接无法正常下载，请检查项目链接或github网络连通情况。",
|
||
).send()
|
||
|
||
files_list = get_project_files_with_content(repo_path)
|
||
cl.user_session.set("project_index", files_list)
|
||
if len(files_list) > 0:
|
||
await cl.Message(
|
||
content=f"已成功上传，您可以开始对项目进行提问！",
|
||
).send()
|
||
elif res.get("value") == "zip":
|
||
files = None
|
||
while files == None:
|
||
files = await cl.AskFileMessage(
|
||
content="请上传项目zip压缩文件!",
|
||
accept={"application/zip": [".zip"]},
|
||
max_size_mb=50,
|
||
).send()
|
||
|
||
text_file = files[0]
|
||
extracted_path = unzip_file(text_file.path, extract_dir)
|
||
files_list = get_project_files_with_content(extracted_path)
|
||
cl.user_session.set("project_index", files_list)
|
||
if len(files_list) > 0:
|
||
await cl.Message(
|
||
content=f"已成功上传，您可以开始对项目进行提问！",
|
||
).send()
|
||
|
||
|
||
@cl.step(type="tool")
|
||
async def bing_search_tool(search_text):
|
||
current_step = cl.context.current_step
|
||
# Simulate a running task
|
||
current_step.input = search_text
|
||
|
||
prompt_tmp = bing_search_prompt(search_text)
|
||
current_step.output = prompt_tmp
|
||
|
||
return prompt_tmp
|
||
|
||
@cl.on_message
|
||
async def main(message: cl.Message):
|
||
chat_profile = cl.user_session.get("chat_profile")
|
||
message_history = cl.user_session.get("message_history")
|
||
|
||
tool_name = tools_choose_agent(message.content)
|
||
is_online = cl.user_session.get("is_online")
|
||
|
||
if chat_profile == "联网聊天":
|
||
if "online_query" in tool_name and is_online:
|
||
prompt_tmp = await bing_search_tool(message.content)
|
||
message_history.append({"role": "tool", "content": prompt_tmp})
|
||
message_history.append({"role": "user", "content": message.content})
|
||
prompt_content = get_cur_base_user_prompt(message_history=message_history)
|
||
|
||
elif chat_profile == "项目问答":
|
||
message_history.append({"role": "user", "content": message.content})
|
||
project_index = cl.user_session.get("project_index")
|
||
index_prompt = ""
|
||
index_tmp = """###PATH:{path}\n{code}\n"""
|
||
for index in project_index:
|
||
index_prompt += index_tmp.format(path=index["path"], code=index["content"])
|
||
if len(tool_name)>0:
|
||
prompt_content = get_cur_base_user_prompt(
|
||
message_history=message_history,
|
||
index_prompt=index_prompt,
|
||
judge_context=tool_name[0],
|
||
)
|
||
|
||
else:
|
||
prompt_content = get_cur_base_user_prompt(message_history=message_history)
|
||
|
||
|
||
|
||
msg = cl.Message(content="")
|
||
await msg.send()
|
||
temperature = cl.user_session.get("temperature")
|
||
top_p = cl.user_session.get("top_p")
|
||
|
||
|
||
if len(prompt_content) / 4 < 120000:
|
||
stream = codegeex4(prompt_content, temperature=temperature, top_p=top_p)
|
||
|
||
for part in stream:
|
||
if token := (part or " "):
|
||
await msg.stream_token(token)
|
||
else:
|
||
await msg.stream_token("项目太大了，请换小一点的项目。")
|
||
|
||
message_history.append({"role": "assistant", "content": msg.content})
|
||
await msg.update()
|