cathub/LocalChatGPT
0
1# -*- coding:utf-8 -*-2from __future__ import annotations3from typing import TYPE_CHECKING, List4 5import logging6import json7import os8import requests9import urllib310 11from tqdm import tqdm12import colorama13from duckduckgo_search import ddg14import asyncio15import aiohttp16from llama_index.indices.query.vector_store import GPTVectorStoreIndexQuery17from llama_index.indices.query.schema import QueryBundle18from langchain.llms import OpenAIChat19 20from modules.presets import *21from modules.llama_func import *22from modules.utils import *23import modules.shared as shared24 25# logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s")26 27if TYPE_CHECKING:28 from typing import TypedDict29 30 class DataframeData(TypedDict):31 headers: List[str]32 data: List[List[str | int | bool]]33 34 35initial_prompt = "You are a helpful assistant."36HISTORY_DIR = "history"37TEMPLATES_DIR = "templates"38 39def get_response(40 openai_api_key, system_prompt, history, temperature, top_p, stream, selected_model41):42 headers = {43 "Content-Type": "application/json",44 "Authorization": f"Bearer {openai_api_key}",45 }46 47 history = [construct_system(system_prompt), *history]48 49 payload = {50 "model": selected_model,51 "messages": history, # [{"role": "user", "content": f"{inputs}"}],52 "temperature": temperature, # 1.0,53 "top_p": top_p, # 1.0,54 "n": 1,55 "stream": stream,56 "presence_penalty": 0,57 "frequency_penalty": 0,58 }59 if stream:60 timeout = timeout_streaming61 else:62 timeout = timeout_all63 64 proxies = get_proxies()65 66 # 如果有自定义的api-url，使用自定义url发送请求，否则使用默认设置发送请求67 if shared.state.api_url != API_URL:68 logging.info(f"使用自定义API URL: {shared.state.api_url}")69 70 response = requests.post(71 shared.state.api_url,72 headers=headers,73 json=payload,74 stream=True,75 timeout=timeout,76 proxies=proxies,77 )78 79 return response80 81 82def stream_predict(83 openai_api_key,84 system_prompt,85 history,86 inputs,87 chatbot,88 all_token_counts,89 top_p,90 temperature,91 selected_model,92 fake_input=None,93 display_append=""94):95 def get_return_value():96 return chatbot, history, status_text, all_token_counts97 98 logging.info("实时回答模式")99 partial_words = ""100 counter = 0101 status_text = "开始实时传输回答……"102 history.append(construct_user(inputs))103 history.append(construct_assistant(""))104 if fake_input:105 chatbot.append((fake_input, ""))106 else:107 chatbot.append((inputs, ""))108 user_token_count = 0109 if fake_input is not None:110 input_token_count = count_token(construct_user(fake_input))111 else:112 input_token_count = count_token(construct_user(inputs))113 if len(all_token_counts) == 0:114 system_prompt_token_count = count_token(construct_system(system_prompt))115 user_token_count = (116 input_token_count + system_prompt_token_count117 )118 else:119 user_token_count = input_token_count120 all_token_counts.append(user_token_count)121 logging.info(f"输入token计数: {user_token_count}")122 yield get_return_value()123 try:124 response = get_response(125 openai_api_key,126 system_prompt,127 history,128 temperature,129 top_p,130 True,131 selected_model,132 )133 except requests.exceptions.ConnectTimeout:134 status_text = (135 standard_error_msg + connection_timeout_prompt + error_retrieve_prompt136 )137 yield get_return_value()138 return139 except requests.exceptions.ReadTimeout:140 status_text = standard_error_msg + read_timeout_prompt + error_retrieve_prompt141 yield get_return_value()142 return143 144 yield get_return_value()145 error_json_str = ""146 147 if fake_input is not None:148 history[-2] = construct_user(fake_input)149 for chunk in response.iter_lines():150 if counter == 0:151 counter += 1152 continue153 counter += 1154 # check whether each line is non-empty155 if chunk:156 chunk = chunk.decode()157 chunklength = len(chunk)158 try:159 chunk = json.loads(chunk[6:])160 except json.JSONDecodeError:161 logging.info(chunk)162 error_json_str += chunk163 status_text = f"JSON解析错误。请重置对话。收到的内容: {error_json_str}"164 yield get_return_value()165 continue166 # decode each line as response data is in bytes167 if chunklength > 6 and "delta" in chunk["choices"][0]:168 finish_reason = chunk["choices"][0]["finish_reason"]169 status_text = construct_token_message(170 sum(all_token_counts), stream=True171 )172 if finish_reason == "stop":173 yield get_return_value()174 break175 try:176 partial_words = (177 partial_words + chunk["choices"][0]["delta"]["content"]178 )179 except KeyError:180 status_text = (181 standard_error_msg182 + "API回复中找不到内容。很可能是Token计数达到上限了。请重置对话。当前Token计数: "183 + str(sum(all_token_counts))184 )185 yield get_return_value()186 break187 history[-1] = construct_assistant(partial_words)188 chatbot[-1] = (chatbot[-1][0], partial_words+display_append)189 all_token_counts[-1] += 1190 yield get_return_value()191 192 193def predict_all(194 openai_api_key,195 system_prompt,196 history,197 inputs,198 chatbot,199 all_token_counts,200 top_p,201 temperature,202 selected_model,203 fake_input=None,204 display_append=""205):206 logging.info("一次性回答模式")207 history.append(construct_user(inputs))208 history.append(construct_assistant(""))209 if fake_input:210 chatbot.append((fake_input, ""))211 else:212 chatbot.append((inputs, ""))213 if fake_input is not None:214 all_token_counts.append(count_token(construct_user(fake_input)))215 else:216 all_token_counts.append(count_token(construct_user(inputs)))217 try:218 response = get_response(219 openai_api_key,220 system_prompt,221 history,222 temperature,223 top_p,224 False,225 selected_model,226 )227 except requests.exceptions.ConnectTimeout:228 status_text = (229 standard_error_msg + connection_timeout_prompt + error_retrieve_prompt230 )231 return chatbot, history, status_text, all_token_counts232 except requests.exceptions.ProxyError:233 status_text = standard_error_msg + proxy_error_prompt + error_retrieve_prompt234 return chatbot, history, status_text, all_token_counts235 except requests.exceptions.SSLError:236 status_text = standard_error_msg + ssl_error_prompt + error_retrieve_prompt237 return chatbot, history, status_text, all_token_counts238 response = json.loads(response.text)239 if fake_input is not None:240 history[-2] = construct_user(fake_input)241 try:242 content = response["choices"][0]["message"]["content"]243 history[-1] = construct_assistant(content)244 chatbot[-1] = (chatbot[-1][0], content+display_append)245 total_token_count = response["usage"]["total_tokens"]246 if fake_input is not None:247 all_token_counts[-1] += count_token(construct_assistant(content))248 else:249 all_token_counts[-1] = total_token_count - sum(all_token_counts)250 status_text = construct_token_message(total_token_count)251 return chatbot, history, status_text, all_token_counts252 except KeyError:253 status_text = standard_error_msg + str(response)254 return chatbot, history, status_text, all_token_counts255 256def is_repeated_string(s):257 n = len(s)258 for i in range(1, n // 2 + 1):259 if n % i == 0:260 sub = s[:i]261 if sub * (n // i) == s:262 return True263 return False264 265def predict(266 openai_api_key,267 system_prompt,268 history,269 inputs,270 chatbot,271 all_token_counts,272 top_p,273 temperature,274 stream=False,275 selected_model=MODELS[0],276 use_websearch=False,277 files = None,278 reply_language="中文",279 should_check_token_count=True,280): # repetition_penalty, top_k281 logging.info("输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)282 if is_repeated_string(inputs):283 print("================== 有人来浪费了 ======================")284 yield chatbot+[(inputs, "🖕️🖕️🖕️🖕️🖕️看不起你")], history, "🖕️🖕️🖕️🖕️🖕️🖕️", all_token_counts285 return286 if should_check_token_count:287 yield chatbot+[(inputs, "")], history, "开始生成回答……", all_token_counts288 if reply_language == "跟随问题语言（不稳定）":289 reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."290 old_inputs = None291 display_reference = []292 limited_context = False293 if files:294 limited_context = True295 old_inputs = inputs296 msg = "加载索引中……（这可能需要几分钟）"297 logging.info(msg)298 yield chatbot+[(inputs, "")], history, msg, all_token_counts299 index = construct_index(openai_api_key, file_src=files)300 msg = "索引构建完成，获取回答中……"301 logging.info(msg)302 yield chatbot+[(inputs, "")], history, msg, all_token_counts303 llm_predictor = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=selected_model))304 prompt_helper = PromptHelper(max_input_size = 4096, num_output = 5, max_chunk_overlap = 20, chunk_size_limit=600)305 service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)306 query_object = GPTVectorStoreIndexQuery(index.index_struct, service_context=service_context, similarity_top_k=5, vector_store=index._vector_store, docstore=index._docstore)307 query_bundle = QueryBundle(inputs)308 nodes = query_object.retrieve(query_bundle)309 reference_results = [n.node.text for n in nodes]310 reference_results = add_source_numbers(reference_results, use_source=False)311 display_reference = add_details(reference_results)312 display_reference = "\n\n" + "".join(display_reference)313 inputs = (314 replace_today(PROMPT_TEMPLATE)315 .replace("{query_str}", inputs)316 .replace("{context_str}", "\n\n".join(reference_results))317 .replace("{reply_language}", reply_language )318 )319 elif use_websearch:320 limited_context = True321 search_results = ddg(inputs, max_results=5)322 old_inputs = inputs323 reference_results = []324 for idx, result in enumerate(search_results):325 logging.info(f"搜索结果{idx + 1}：{result}")326 domain_name = urllib3.util.parse_url(result["href"]).host327 reference_results.append([result["body"], result["href"]])328 display_reference.append(f"{idx+1}. [{domain_name}]({result['href']})\n")329 reference_results = add_source_numbers(reference_results)330 display_reference = "\n\n" + "".join(display_reference)331 inputs = (332 replace_today(WEBSEARCH_PTOMPT_TEMPLATE)333 .replace("{query}", inputs)334 .replace("{web_results}", "\n\n".join(reference_results))335 .replace("{reply_language}", reply_language )336 )337 else:338 display_reference = ""339 340 if len(openai_api_key) != 51:341 status_text = standard_error_msg + no_apikey_msg342 logging.info(status_text)343 chatbot.append((inputs, ""))344 if len(history) == 0:345 history.append(construct_user(inputs))346 history.append("")347 all_token_counts.append(0)348 else:349 history[-2] = construct_user(inputs)350 yield chatbot+[(inputs, "")], history, status_text, all_token_counts351 return352 elif len(inputs.strip()) == 0:353 status_text = standard_error_msg + no_input_msg354 logging.info(status_text)355 yield chatbot+[(inputs, "")], history, status_text, all_token_counts356 return357 358 if stream:359 logging.info("使用流式传输")360 iter = stream_predict(361 openai_api_key,362 system_prompt,363 history,364 inputs,365 chatbot,366 all_token_counts,367 top_p,368 temperature,369 selected_model,370 fake_input=old_inputs,371 display_append=display_reference372 )373 for chatbot, history, status_text, all_token_counts in iter:374 if shared.state.interrupted:375 shared.state.recover()376 return377 yield chatbot, history, status_text, all_token_counts378 else:379 logging.info("不使用流式传输")380 chatbot, history, status_text, all_token_counts = predict_all(381 openai_api_key,382 system_prompt,383 history,384 inputs,385 chatbot,386 all_token_counts,387 top_p,388 temperature,389 selected_model,390 fake_input=old_inputs,391 display_append=display_reference392 )393 yield chatbot, history, status_text, all_token_counts394 395 logging.info(f"传输完毕。当前token计数为{all_token_counts}")396 if len(history) > 1 and history[-1]["content"] != inputs:397 logging.info(398 "回答为："399 + colorama.Fore.BLUE400 + f"{history[-1]['content']}"401 + colorama.Style.RESET_ALL402 )403 404 if limited_context:405 history = history[-4:]406 all_token_counts = all_token_counts[-2:]407 yield chatbot, history, status_text, all_token_counts408 409 if stream:410 max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["streaming"]411 else:412 max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["all"]413 414 if sum(all_token_counts) > max_token and should_check_token_count:415 status_text = f"精简token中{all_token_counts}/{max_token}"416 logging.info(status_text)417 yield chatbot, history, status_text, all_token_counts418 iter = reduce_token_size(419 openai_api_key,420 system_prompt,421 history,422 chatbot,423 all_token_counts,424 top_p,425 temperature,426 max_token//2,427 selected_model=selected_model,428 )429 for chatbot, history, status_text, all_token_counts in iter:430 status_text = f"Token 达到上限，已自动降低Token计数至 {status_text}"431 yield chatbot, history, status_text, all_token_counts432 433 434def retry(435 openai_api_key,436 system_prompt,437 history,438 chatbot,439 token_count,440 top_p,441 temperature,442 stream=False,443 selected_model=MODELS[0],444 reply_language="中文",445):446 logging.info("重试中……")447 if len(history) == 0:448 yield chatbot, history, f"{standard_error_msg}上下文是空的", token_count449 return450 history.pop()451 inputs = history.pop()["content"]452 token_count.pop()453 iter = predict(454 openai_api_key,455 system_prompt,456 history,457 inputs,458 chatbot,459 token_count,460 top_p,461 temperature,462 stream=stream,463 selected_model=selected_model,464 reply_language=reply_language,465 )466 logging.info("重试中……")467 for x in iter:468 yield x469 logging.info("重试完毕")470 471 472def reduce_token_size(473 openai_api_key,474 system_prompt,475 history,476 chatbot,477 token_count,478 top_p,479 temperature,480 max_token_count,481 selected_model=MODELS[0],482 reply_language="中文",483):484 logging.info("开始减少token数量……")485 iter = predict(486 openai_api_key,487 system_prompt,488 history,489 summarize_prompt,490 chatbot,491 token_count,492 top_p,493 temperature,494 selected_model=selected_model,495 should_check_token_count=False,496 reply_language=reply_language,497 )498 logging.info(f"chatbot: {chatbot}")499 flag = False500 for chatbot, history, status_text, previous_token_count in iter:501 num_chat = find_n(previous_token_count, max_token_count)502 logging.info(f"previous_token_count: {previous_token_count}, keeping {num_chat} chats")503 if flag:504 chatbot = chatbot[:-1]505 flag = True506 history = history[-2*num_chat:] if num_chat > 0 else []507 token_count = previous_token_count[-num_chat:] if num_chat > 0 else []508 msg = f"保留了最近{num_chat}轮对话"509 yield chatbot, history, msg + "，" + construct_token_message(510 sum(token_count) if len(token_count) > 0 else 0,511 ), token_count512 logging.info(msg)513 logging.info("减少token数量完毕")514