File size: 16,921 Bytes
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55f3e5a
3e8b0ee
 
 
 
55f3e5a
 
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55f3e5a
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55f3e5a
 
 
a4aaca9
55f3e5a
 
 
 
 
 
 
 
a4aaca9
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
 
 
 
3e8b0ee
 
 
a4aaca9
3e8b0ee
 
a4aaca9
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
 
cb9974f
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55f3e5a
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
 
 
 
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55f3e5a
 
 
 
 
3e8b0ee
a4aaca9
 
3e8b0ee
 
a4aaca9
 
 
3e8b0ee
a4aaca9
 
cb9974f
3e8b0ee
 
 
 
a4aaca9
3e8b0ee
55f3e5a
 
 
 
 
 
 
 
a4aaca9
 
 
 
 
 
 
 
 
 
 
 
3e8b0ee
 
a4aaca9
3e8b0ee
 
 
a4aaca9
 
 
 
3e8b0ee
 
 
a4aaca9
3e8b0ee
 
 
a4aaca9
3e8b0ee
55f3e5a
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
 
 
 
 
3e8b0ee
a4aaca9
3e8b0ee
a4aaca9
3e8b0ee
 
55f3e5a
 
 
 
 
 
3e8b0ee
55f3e5a
3e8b0ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4aaca9
3e8b0ee
 
 
 
 
 
 
55f3e5a
3e8b0ee
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
# -*- coding:utf-8 -*-
from __future__ import annotations
from typing import TYPE_CHECKING, List

import logging
import json
import os
import requests
import urllib3

from tqdm import tqdm
import colorama
from duckduckgo_search import ddg
import asyncio
import aiohttp


from modules.presets import *
from modules.llama_func import *
from modules.utils import *
from . import shared
from modules.config import retrieve_proxy

# logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s")

if TYPE_CHECKING:
    from typing import TypedDict

    class DataframeData(TypedDict):
        headers: List[str]
        data: List[List[str | int | bool]]


initial_prompt = "You are a helpful assistant."
HISTORY_DIR = "history"
TEMPLATES_DIR = "templates"

@shared.state.switching_api_key # 在不开启多账号模式的时候,这个装饰器不会起作用
def get_response(
    openai_api_key, system_prompt, history, temperature, top_p, stream, selected_model
):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {openai_api_key}",
    }

    history = [construct_system(system_prompt), *history]

    payload = {
        "model": selected_model,
        "messages": history,  # [{"role": "user", "content": f"{inputs}"}],
        "temperature": temperature,  # 1.0,
        "top_p": top_p,  # 1.0,
        "n": 1,
        "stream": stream,
        "presence_penalty": 0,
        "frequency_penalty": 0,
    }
    if stream:
        timeout = timeout_streaming
    else:
        timeout = timeout_all


    # 如果有自定义的api-host,使用自定义host发送请求,否则使用默认设置发送请求
    if shared.state.completion_url != COMPLETION_URL:
        logging.info(f"使用自定义API URL: {shared.state.completion_url}")

    with retrieve_proxy():
        response = requests.post(
            shared.state.completion_url,
            headers=headers,
            json=payload,
            stream=True,
            timeout=timeout,
        )

    return response


def stream_predict(
    openai_api_key,
    system_prompt,
    history,
    inputs,
    chatbot,
    all_token_counts,
    top_p,
    temperature,
    selected_model,
    fake_input=None,
    display_append=""
):
    def get_return_value():
        return chatbot, history, status_text, all_token_counts

    logging.info("实时回答模式")
    partial_words = ""
    counter = 0
    status_text = "开始实时传输回答……"
    history.append(construct_user(inputs))
    history.append(construct_assistant(""))
    if fake_input:
        chatbot.append((fake_input, ""))
    else:
        chatbot.append((inputs, ""))
    user_token_count = 0
    if fake_input is not None:
        input_token_count = count_token(construct_user(fake_input))
    else:
        input_token_count = count_token(construct_user(inputs))
    if len(all_token_counts) == 0:
        system_prompt_token_count = count_token(construct_system(system_prompt))
        user_token_count = (
            input_token_count + system_prompt_token_count
        )
    else:
        user_token_count = input_token_count
    all_token_counts.append(user_token_count)
    logging.info(f"输入token计数: {user_token_count}")
    yield get_return_value()
    try:
        response = get_response(
            openai_api_key,
            system_prompt,
            history,
            temperature,
            top_p,
            True,
            selected_model,
        )
    except requests.exceptions.ConnectTimeout:
        status_text = (
            standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
        )
        yield get_return_value()
        return
    except requests.exceptions.ReadTimeout:
        status_text = standard_error_msg + read_timeout_prompt + error_retrieve_prompt
        yield get_return_value()
        return

    yield get_return_value()
    error_json_str = ""

    if fake_input is not None:
        history[-2] = construct_user(fake_input)
    for chunk in tqdm(response.iter_lines()):
        if counter == 0:
            counter += 1
            continue
        counter += 1
        # check whether each line is non-empty
        if chunk:
            chunk = chunk.decode()
            chunklength = len(chunk)
            try:
                chunk = json.loads(chunk[6:])
            except json.JSONDecodeError:
                logging.info(chunk)
                error_json_str += chunk
                status_text = f"JSON解析错误。请重置对话。收到的内容: {error_json_str}"
                yield get_return_value()
                continue
            # decode each line as response data is in bytes
            if chunklength > 6 and "delta" in chunk["choices"][0]:
                finish_reason = chunk["choices"][0]["finish_reason"]
                status_text = construct_token_message(all_token_counts)
                if finish_reason == "stop":
                    yield get_return_value()
                    break
                try:
                    partial_words = (
                        partial_words + chunk["choices"][0]["delta"]["content"]
                    )
                except KeyError:
                    status_text = (
                        standard_error_msg
                        + "API回复中找不到内容。很可能是Token计数达到上限了。请重置对话。当前Token计数: "
                        + str(sum(all_token_counts))
                    )
                    yield get_return_value()
                    break
                history[-1] = construct_assistant(partial_words)
                chatbot[-1] = (chatbot[-1][0], partial_words+display_append)
                all_token_counts[-1] += 1
                yield get_return_value()


def predict_all(
    openai_api_key,
    system_prompt,
    history,
    inputs,
    chatbot,
    all_token_counts,
    top_p,
    temperature,
    selected_model,
    fake_input=None,
    display_append=""
):
    logging.info("一次性回答模式")
    history.append(construct_user(inputs))
    history.append(construct_assistant(""))
    if fake_input:
        chatbot.append((fake_input, ""))
    else:
        chatbot.append((inputs, ""))
    if fake_input is not None:
        all_token_counts.append(count_token(construct_user(fake_input)))
    else:
        all_token_counts.append(count_token(construct_user(inputs)))
    try:
        response = get_response(
            openai_api_key,
            system_prompt,
            history,
            temperature,
            top_p,
            False,
            selected_model,
        )
    except requests.exceptions.ConnectTimeout:
        status_text = (
            standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
        )
        return chatbot, history, status_text, all_token_counts
    except requests.exceptions.ProxyError:
        status_text = standard_error_msg + proxy_error_prompt + error_retrieve_prompt
        return chatbot, history, status_text, all_token_counts
    except requests.exceptions.SSLError:
        status_text = standard_error_msg + ssl_error_prompt + error_retrieve_prompt
        return chatbot, history, status_text, all_token_counts
    response = json.loads(response.text)
    if fake_input is not None:
        history[-2] = construct_user(fake_input)
    try:
        content = response["choices"][0]["message"]["content"]
        history[-1] = construct_assistant(content)
        chatbot[-1] = (chatbot[-1][0], content+display_append)
        total_token_count = response["usage"]["total_tokens"]
        if fake_input is not None:
            all_token_counts[-1] += count_token(construct_assistant(content))
        else:
            all_token_counts[-1] = total_token_count - sum(all_token_counts)
        status_text = construct_token_message(total_token_count)
        return chatbot, history, status_text, all_token_counts
    except KeyError:
        status_text = standard_error_msg + str(response)
        return chatbot, history, status_text, all_token_counts


def predict(
    openai_api_key,
    system_prompt,
    history,
    inputs,
    chatbot,
    all_token_counts,
    top_p,
    temperature,
    stream=False,
    selected_model=MODELS[0],
    use_websearch=False,
    files = None,
    reply_language="中文",
    should_check_token_count=True,
):  # repetition_penalty, top_k
    from llama_index.indices.vector_store.base_query import GPTVectorStoreIndexQuery
    from llama_index.indices.query.schema import QueryBundle
    from langchain.llms import OpenAIChat


    logging.info("输入为:" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
    if should_check_token_count:
        yield chatbot+[(inputs, "")], history, "开始生成回答……", all_token_counts
    if reply_language == "跟随问题语言(不稳定)":
        reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."
    old_inputs = None
    display_reference = []
    limited_context = False
    if files:
        limited_context = True
        old_inputs = inputs
        msg = "加载索引中……(这可能需要几分钟)"
        logging.info(msg)
        yield chatbot+[(inputs, "")], history, msg, all_token_counts
        index = construct_index(openai_api_key, file_src=files)
        msg = "索引构建完成,获取回答中……"
        logging.info(msg)
        yield chatbot+[(inputs, "")], history, msg, all_token_counts
        with retrieve_proxy():
            llm_predictor = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=selected_model))
            prompt_helper = PromptHelper(max_input_size = 4096, num_output = 5, max_chunk_overlap = 20, chunk_size_limit=600)
            from llama_index import ServiceContext
            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
            query_object = GPTVectorStoreIndexQuery(index.index_struct, service_context=service_context, similarity_top_k=5, vector_store=index._vector_store, docstore=index._docstore)
            query_bundle = QueryBundle(inputs)
            nodes = query_object.retrieve(query_bundle)
        reference_results = [n.node.text for n in nodes]
        reference_results = add_source_numbers(reference_results, use_source=False)
        display_reference = add_details(reference_results)
        display_reference = "\n\n" + "".join(display_reference)
        inputs = (
            replace_today(PROMPT_TEMPLATE)
            .replace("{query_str}", inputs)
            .replace("{context_str}", "\n\n".join(reference_results))
            .replace("{reply_language}", reply_language )
        )
    elif use_websearch:
        limited_context = True
        search_results = ddg(inputs, max_results=5)
        old_inputs = inputs
        reference_results = []
        for idx, result in enumerate(search_results):
            logging.info(f"搜索结果{idx + 1}{result}")
            domain_name = urllib3.util.parse_url(result["href"]).host
            reference_results.append([result["body"], result["href"]])
            display_reference.append(f"{idx+1}. [{domain_name}]({result['href']})\n")
        reference_results = add_source_numbers(reference_results)
        display_reference = "\n\n" + "".join(display_reference)
        inputs = (
            replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
            .replace("{query}", inputs)
            .replace("{web_results}", "\n\n".join(reference_results))
            .replace("{reply_language}", reply_language )
        )
    else:
        display_reference = ""

    if len(openai_api_key) == 0 and not shared.state.multi_api_key:
        status_text = standard_error_msg + no_apikey_msg
        logging.info(status_text)
        chatbot.append((inputs, ""))
        if len(history) == 0:
            history.append(construct_user(inputs))
            history.append("")
            all_token_counts.append(0)
        else:
            history[-2] = construct_user(inputs)
        yield chatbot+[(inputs, "")], history, status_text, all_token_counts
        return
    elif len(inputs.strip()) == 0:
        status_text = standard_error_msg + no_input_msg
        logging.info(status_text)
        yield chatbot+[(inputs, "")], history, status_text, all_token_counts
        return

    if stream:
        logging.info("使用流式传输")
        iter = stream_predict(
            openai_api_key,
            system_prompt,
            history,
            inputs,
            chatbot,
            all_token_counts,
            top_p,
            temperature,
            selected_model,
            fake_input=old_inputs,
            display_append=display_reference
        )
        for chatbot, history, status_text, all_token_counts in iter:
            if shared.state.interrupted:
                shared.state.recover()
                return
            yield chatbot, history, status_text, all_token_counts
    else:
        logging.info("不使用流式传输")
        chatbot, history, status_text, all_token_counts = predict_all(
            openai_api_key,
            system_prompt,
            history,
            inputs,
            chatbot,
            all_token_counts,
            top_p,
            temperature,
            selected_model,
            fake_input=old_inputs,
            display_append=display_reference
        )
        yield chatbot, history, status_text, all_token_counts

    logging.info(f"传输完毕。当前token计数为{all_token_counts}")
    if len(history) > 1 and history[-1]["content"] != inputs:
        logging.info(
            "回答为:"
            + colorama.Fore.BLUE
            + f"{history[-1]['content']}"
            + colorama.Style.RESET_ALL
        )

    if limited_context:
        history = history[-4:]
        all_token_counts = all_token_counts[-2:]
        yield chatbot, history, status_text, all_token_counts

    if stream:
        max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["streaming"]
    else:
        max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["all"]

    if sum(all_token_counts) > max_token and should_check_token_count:
        print(all_token_counts)
        count = 0
        while sum(all_token_counts) > max_token - 500 and sum(all_token_counts) > 0:
            count += 1
            del all_token_counts[0]
            del history[:2]
        logging.info(status_text)
        status_text = f"为了防止token超限,模型忘记了早期的 {count} 轮对话"
        yield chatbot, history, status_text, all_token_counts


def retry(
    openai_api_key,
    system_prompt,
    history,
    chatbot,
    token_count,
    top_p,
    temperature,
    stream=False,
    selected_model=MODELS[0],
    reply_language="中文",
):
    logging.info("重试中……")
    if len(history) == 0:
        yield chatbot, history, f"{standard_error_msg}上下文是空的", token_count
        return
    history.pop()
    inputs = history.pop()["content"]
    token_count.pop()
    iter = predict(
        openai_api_key,
        system_prompt,
        history,
        inputs,
        chatbot,
        token_count,
        top_p,
        temperature,
        stream=stream,
        selected_model=selected_model,
        reply_language=reply_language,
    )
    logging.info("重试中……")
    for x in iter:
        yield x
    logging.info("重试完毕")


def reduce_token_size(
    openai_api_key,
    system_prompt,
    history,
    chatbot,
    token_count,
    top_p,
    temperature,
    max_token_count,
    selected_model=MODELS[0],
    reply_language="中文",
):
    logging.info("开始减少token数量……")
    iter = predict(
        openai_api_key,
        system_prompt,
        history,
        summarize_prompt,
        chatbot,
        token_count,
        top_p,
        temperature,
        selected_model=selected_model,
        should_check_token_count=False,
        reply_language=reply_language,
    )
    logging.info(f"chatbot: {chatbot}")
    flag = False
    for chatbot, history, status_text, previous_token_count in iter:
        num_chat = find_n(previous_token_count, max_token_count)
        logging.info(f"previous_token_count: {previous_token_count}, keeping {num_chat} chats")
        if flag:
            chatbot = chatbot[:-1]
        flag = True
        history = history[-2*num_chat:] if num_chat > 0 else []
        token_count = previous_token_count[-num_chat:] if num_chat > 0 else []
        msg = f"保留了最近{num_chat}轮对话"
        yield chatbot, history, msg + "," + construct_token_message(
            token_count if len(token_count) > 0 else [0],
        ), token_count
    logging.info(msg)
    logging.info("减少token数量完毕")