CodeGeeX4/repodemo/run.py

229 lines
8.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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? exhttps://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()