2024-01-15 08:46:22 +08:00
|
|
|
#
|
2025-03-26 19:33:14 +08:00
|
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
2024-01-15 08:46:22 +08:00
|
|
|
#
|
|
|
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
|
# you may not use this file except in compliance with the License.
|
|
|
|
|
# You may obtain a copy of the License at
|
|
|
|
|
#
|
|
|
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
|
#
|
|
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
|
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
|
# See the License for the specific language governing permissions and
|
|
|
|
|
# limitations under the License.
|
|
|
|
|
#
|
2025-03-26 19:33:14 +08:00
|
|
|
import asyncio
|
|
|
|
|
import json
|
|
|
|
|
import logging
|
|
|
|
|
import os
|
2025-03-22 23:07:03 +08:00
|
|
|
import random
|
2025-07-30 19:41:09 +08:00
|
|
|
import re
|
2025-03-26 19:33:14 +08:00
|
|
|
import time
|
2023-12-25 19:05:59 +08:00
|
|
|
from abc import ABC
|
2025-06-11 17:20:12 +08:00
|
|
|
from copy import deepcopy
|
2025-05-16 16:32:19 +08:00
|
|
|
from typing import Any, Protocol
|
2025-06-03 14:18:40 +08:00
|
|
|
from urllib.parse import urljoin
|
2025-03-26 19:33:14 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
import json_repair
|
2025-08-12 10:59:20 +08:00
|
|
|
import litellm
|
2024-02-27 14:57:34 +08:00
|
|
|
import openai
|
2025-03-26 19:33:14 +08:00
|
|
|
import requests
|
|
|
|
|
from openai import OpenAI
|
|
|
|
|
from openai.lib.azure import AzureOpenAI
|
2025-07-30 19:41:09 +08:00
|
|
|
from strenum import StrEnum
|
2025-03-26 19:33:14 +08:00
|
|
|
from zhipuai import ZhipuAI
|
|
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
from rag.llm import FACTORY_DEFAULT_BASE_URL, LITELLM_PROVIDER_PREFIX, SupportedLiteLLMProvider
|
2024-12-05 13:28:42 +08:00
|
|
|
from rag.nlp import is_chinese, is_english
|
2025-09-22 17:17:06 +08:00
|
|
|
from rag.utils import num_tokens_from_string, total_token_count_from_response
|
2025-03-22 23:07:03 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
|
2025-03-22 23:07:03 +08:00
|
|
|
# Error message constants
|
2025-07-30 19:41:09 +08:00
|
|
|
class LLMErrorCode(StrEnum):
|
|
|
|
|
ERROR_RATE_LIMIT = "RATE_LIMIT_EXCEEDED"
|
|
|
|
|
ERROR_AUTHENTICATION = "AUTH_ERROR"
|
|
|
|
|
ERROR_INVALID_REQUEST = "INVALID_REQUEST"
|
|
|
|
|
ERROR_SERVER = "SERVER_ERROR"
|
|
|
|
|
ERROR_TIMEOUT = "TIMEOUT"
|
|
|
|
|
ERROR_CONNECTION = "CONNECTION_ERROR"
|
|
|
|
|
ERROR_MODEL = "MODEL_ERROR"
|
|
|
|
|
ERROR_MAX_ROUNDS = "ERROR_MAX_ROUNDS"
|
|
|
|
|
ERROR_CONTENT_FILTER = "CONTENT_FILTERED"
|
|
|
|
|
ERROR_QUOTA = "QUOTA_EXCEEDED"
|
|
|
|
|
ERROR_MAX_RETRIES = "MAX_RETRIES_EXCEEDED"
|
|
|
|
|
ERROR_GENERIC = "GENERIC_ERROR"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ReActMode(StrEnum):
|
|
|
|
|
FUNCTION_CALL = "function_call"
|
|
|
|
|
REACT = "react"
|
2023-12-28 13:50:13 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
ERROR_PREFIX = "**ERROR**"
|
2025-03-12 19:40:54 +08:00
|
|
|
LENGTH_NOTIFICATION_CN = "······\n由于大模型的上下文窗口大小限制,回答已经被大模型截断。"
|
2025-03-13 14:43:24 +08:00
|
|
|
LENGTH_NOTIFICATION_EN = "...\nThe answer is truncated by your chosen LLM due to its limitation on context length."
|
2024-09-27 13:17:25 +08:00
|
|
|
|
2025-01-22 19:43:14 +08:00
|
|
|
|
2025-05-16 16:32:19 +08:00
|
|
|
class ToolCallSession(Protocol):
|
|
|
|
|
def tool_call(self, name: str, arguments: dict[str, Any]) -> str: ...
|
|
|
|
|
|
|
|
|
|
|
2023-12-25 19:05:59 +08:00
|
|
|
class Base(ABC):
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url, **kwargs):
|
2025-03-26 19:33:14 +08:00
|
|
|
timeout = int(os.environ.get("LM_TIMEOUT_SECONDS", 600))
|
2024-10-21 12:11:08 +08:00
|
|
|
self.client = OpenAI(api_key=key, base_url=base_url, timeout=timeout)
|
2024-01-22 19:51:38 +08:00
|
|
|
self.model_name = model_name
|
2025-03-22 23:07:03 +08:00
|
|
|
# Configure retry parameters
|
2025-06-11 17:20:12 +08:00
|
|
|
self.max_retries = kwargs.get("max_retries", int(os.environ.get("LLM_MAX_RETRIES", 5)))
|
|
|
|
|
self.base_delay = kwargs.get("retry_interval", float(os.environ.get("LLM_BASE_DELAY", 2.0)))
|
2025-06-12 12:31:10 +08:00
|
|
|
self.max_rounds = kwargs.get("max_rounds", 5)
|
2025-04-08 16:09:03 +08:00
|
|
|
self.is_tools = False
|
2025-06-23 17:45:35 +08:00
|
|
|
self.tools = []
|
|
|
|
|
self.toolcall_sessions = {}
|
2025-03-22 23:07:03 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
def _get_delay(self):
|
2025-03-22 23:07:03 +08:00
|
|
|
"""Calculate retry delay time"""
|
2025-07-30 19:41:09 +08:00
|
|
|
return self.base_delay * random.uniform(10, 150)
|
2025-03-26 19:33:14 +08:00
|
|
|
|
2025-03-22 23:07:03 +08:00
|
|
|
def _classify_error(self, error):
|
|
|
|
|
"""Classify error based on error message content"""
|
|
|
|
|
error_str = str(error).lower()
|
2025-03-26 19:33:14 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
keywords_mapping = [
|
|
|
|
|
(["quota", "capacity", "credit", "billing", "balance", "欠费"], LLMErrorCode.ERROR_QUOTA),
|
|
|
|
|
(["rate limit", "429", "tpm limit", "too many requests", "requests per minute"], LLMErrorCode.ERROR_RATE_LIMIT),
|
|
|
|
|
(["auth", "key", "apikey", "401", "forbidden", "permission"], LLMErrorCode.ERROR_AUTHENTICATION),
|
|
|
|
|
(["invalid", "bad request", "400", "format", "malformed", "parameter"], LLMErrorCode.ERROR_INVALID_REQUEST),
|
|
|
|
|
(["server", "503", "502", "504", "500", "unavailable"], LLMErrorCode.ERROR_SERVER),
|
|
|
|
|
(["timeout", "timed out"], LLMErrorCode.ERROR_TIMEOUT),
|
|
|
|
|
(["connect", "network", "unreachable", "dns"], LLMErrorCode.ERROR_CONNECTION),
|
|
|
|
|
(["filter", "content", "policy", "blocked", "safety", "inappropriate"], LLMErrorCode.ERROR_CONTENT_FILTER),
|
|
|
|
|
(["model", "not found", "does not exist", "not available"], LLMErrorCode.ERROR_MODEL),
|
|
|
|
|
(["max rounds"], LLMErrorCode.ERROR_MODEL),
|
|
|
|
|
]
|
|
|
|
|
for words, code in keywords_mapping:
|
|
|
|
|
if re.search("({})".format("|".join(words)), error_str):
|
|
|
|
|
return code
|
|
|
|
|
|
|
|
|
|
return LLMErrorCode.ERROR_GENERIC
|
2023-12-25 19:05:59 +08:00
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
|
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2025-08-22 19:33:09 +08:00
|
|
|
|
|
|
|
|
allowed_conf = {
|
|
|
|
|
"temperature",
|
|
|
|
|
"max_completion_tokens",
|
|
|
|
|
"top_p",
|
|
|
|
|
"stream",
|
|
|
|
|
"stream_options",
|
|
|
|
|
"stop",
|
|
|
|
|
"n",
|
|
|
|
|
"presence_penalty",
|
|
|
|
|
"frequency_penalty",
|
|
|
|
|
"functions",
|
|
|
|
|
"function_call",
|
|
|
|
|
"logit_bias",
|
|
|
|
|
"user",
|
|
|
|
|
"response_format",
|
|
|
|
|
"seed",
|
|
|
|
|
"tools",
|
|
|
|
|
"tool_choice",
|
|
|
|
|
"logprobs",
|
|
|
|
|
"top_logprobs",
|
2025-10-10 17:07:55 +08:00
|
|
|
"extra_headers"
|
2025-08-22 19:33:09 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
gen_conf = {k: v for k, v in gen_conf.items() if k in allowed_conf}
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
return gen_conf
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat(self, history, gen_conf, **kwargs):
|
|
|
|
|
logging.info("[HISTORY]" + json.dumps(history, ensure_ascii=False, indent=2))
|
2025-10-10 16:38:04 +08:00
|
|
|
if self.model_name.lower().find("qwq") >= 0:
|
|
|
|
|
logging.info(f"[INFO] {self.model_name} detected as reasoning model, using _chat_streamly")
|
|
|
|
|
|
|
|
|
|
final_ans = ""
|
2025-10-10 13:18:24 +02:00
|
|
|
tol_token = 0
|
2025-10-10 16:38:04 +08:00
|
|
|
for delta, tol in self._chat_streamly(history, gen_conf, with_reasoning=False, **kwargs):
|
|
|
|
|
if delta.startswith("<think>") or delta.endswith("</think>"):
|
|
|
|
|
continue
|
|
|
|
|
final_ans += delta
|
|
|
|
|
tol_token = tol
|
|
|
|
|
|
|
|
|
|
if len(final_ans.strip()) == 0:
|
|
|
|
|
final_ans = "**ERROR**: Empty response from reasoning model"
|
|
|
|
|
|
|
|
|
|
return final_ans.strip(), tol_token
|
2025-10-10 13:18:24 +02:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
if self.model_name.lower().find("qwen3") >= 0:
|
2025-07-30 19:41:09 +08:00
|
|
|
kwargs["extra_body"] = {"enable_thinking": False}
|
2025-09-25 14:11:09 +08:00
|
|
|
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, **gen_conf, **kwargs)
|
2025-06-11 17:20:12 +08:00
|
|
|
|
2025-09-29 14:49:45 +08:00
|
|
|
if not response.choices or not response.choices[0].message or not response.choices[0].message.content:
|
2025-06-11 17:20:12 +08:00
|
|
|
return "", 0
|
|
|
|
|
ans = response.choices[0].message.content.strip()
|
|
|
|
|
if response.choices[0].finish_reason == "length":
|
2025-07-30 19:41:09 +08:00
|
|
|
ans = self._length_stop(ans)
|
2025-10-22 12:25:31 +08:00
|
|
|
return ans, total_token_count_from_response(response)
|
2025-06-11 17:20:12 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat_streamly(self, history, gen_conf, **kwargs):
|
|
|
|
|
logging.info("[HISTORY STREAMLY]" + json.dumps(history, ensure_ascii=False, indent=4))
|
|
|
|
|
reasoning_start = False
|
2025-09-25 14:11:09 +08:00
|
|
|
|
2025-09-08 19:00:52 +08:00
|
|
|
if kwargs.get("stop") or "stop" in gen_conf:
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, stream=True, **gen_conf, stop=kwargs.get("stop"))
|
2025-09-08 19:00:52 +08:00
|
|
|
else:
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, stream=True, **gen_conf)
|
2025-09-25 14:11:09 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
for resp in response:
|
|
|
|
|
if not resp.choices:
|
|
|
|
|
continue
|
|
|
|
|
if not resp.choices[0].delta.content:
|
|
|
|
|
resp.choices[0].delta.content = ""
|
|
|
|
|
if kwargs.get("with_reasoning", True) and hasattr(resp.choices[0].delta, "reasoning_content") and resp.choices[0].delta.reasoning_content:
|
|
|
|
|
ans = ""
|
|
|
|
|
if not reasoning_start:
|
|
|
|
|
reasoning_start = True
|
|
|
|
|
ans = "<think>"
|
|
|
|
|
ans += resp.choices[0].delta.reasoning_content + "</think>"
|
|
|
|
|
else:
|
|
|
|
|
reasoning_start = False
|
|
|
|
|
ans = resp.choices[0].delta.content
|
|
|
|
|
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-07-30 19:41:09 +08:00
|
|
|
if not tol:
|
|
|
|
|
tol = num_tokens_from_string(resp.choices[0].delta.content)
|
|
|
|
|
|
|
|
|
|
if resp.choices[0].finish_reason == "length":
|
|
|
|
|
if is_chinese(ans):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
|
|
|
|
yield ans, tol
|
|
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
def _length_stop(self, ans):
|
|
|
|
|
if is_chinese([ans]):
|
|
|
|
|
return ans + LENGTH_NOTIFICATION_CN
|
|
|
|
|
return ans + LENGTH_NOTIFICATION_EN
|
|
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
@property
|
|
|
|
|
def _retryable_errors(self) -> set[str]:
|
|
|
|
|
return {
|
|
|
|
|
LLMErrorCode.ERROR_RATE_LIMIT,
|
|
|
|
|
LLMErrorCode.ERROR_SERVER,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
def _should_retry(self, error_code: str) -> bool:
|
|
|
|
|
return error_code in self._retryable_errors
|
|
|
|
|
|
|
|
|
|
def _exceptions(self, e, attempt) -> str | None:
|
2025-07-16 13:47:38 +08:00
|
|
|
logging.exception("OpenAI chat_with_tools")
|
2025-06-12 12:31:10 +08:00
|
|
|
# Classify the error
|
|
|
|
|
error_code = self._classify_error(e)
|
2025-07-30 19:41:09 +08:00
|
|
|
if attempt == self.max_retries:
|
|
|
|
|
error_code = LLMErrorCode.ERROR_MAX_RETRIES
|
2025-06-12 12:31:10 +08:00
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
if self._should_retry(error_code):
|
|
|
|
|
delay = self._get_delay()
|
|
|
|
|
logging.warning(f"Error: {error_code}. Retrying in {delay:.2f} seconds... (Attempt {attempt + 1}/{self.max_retries})")
|
|
|
|
|
time.sleep(delay)
|
|
|
|
|
return None
|
2025-06-12 12:31:10 +08:00
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
return f"{ERROR_PREFIX}: {error_code} - {str(e)}"
|
2025-07-30 19:41:09 +08:00
|
|
|
|
2025-06-27 19:28:41 +08:00
|
|
|
def _verbose_tool_use(self, name, args, res):
|
2025-08-12 10:59:20 +08:00
|
|
|
return "<tool_call>" + json.dumps({"name": name, "args": args, "result": res}, ensure_ascii=False, indent=2) + "</tool_call>"
|
2025-06-27 19:28:41 +08:00
|
|
|
|
|
|
|
|
def _append_history(self, hist, tool_call, tool_res):
|
|
|
|
|
hist.append(
|
|
|
|
|
{
|
|
|
|
|
"role": "assistant",
|
|
|
|
|
"tool_calls": [
|
|
|
|
|
{
|
|
|
|
|
"index": tool_call.index,
|
|
|
|
|
"id": tool_call.id,
|
|
|
|
|
"function": {
|
|
|
|
|
"name": tool_call.function.name,
|
|
|
|
|
"arguments": tool_call.function.arguments,
|
|
|
|
|
},
|
|
|
|
|
"type": "function",
|
|
|
|
|
},
|
|
|
|
|
],
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
try:
|
|
|
|
|
if isinstance(tool_res, dict):
|
|
|
|
|
tool_res = json.dumps(tool_res, ensure_ascii=False)
|
|
|
|
|
finally:
|
|
|
|
|
hist.append({"role": "tool", "tool_call_id": tool_call.id, "content": str(tool_res)})
|
|
|
|
|
return hist
|
|
|
|
|
|
2025-04-08 16:09:03 +08:00
|
|
|
def bind_tools(self, toolcall_session, tools):
|
|
|
|
|
if not (toolcall_session and tools):
|
|
|
|
|
return
|
|
|
|
|
self.is_tools = True
|
2025-07-30 19:41:09 +08:00
|
|
|
self.toolcall_session = toolcall_session
|
|
|
|
|
self.tools = tools
|
2025-06-23 17:45:35 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
def chat_with_tools(self, system: str, history: list, gen_conf: dict = {}):
|
2025-06-27 12:10:53 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2025-04-08 16:09:03 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
|
|
|
|
|
|
|
|
|
ans = ""
|
|
|
|
|
tk_count = 0
|
2025-06-12 12:31:10 +08:00
|
|
|
hist = deepcopy(history)
|
2025-04-08 16:09:03 +08:00
|
|
|
# Implement exponential backoff retry strategy
|
2025-07-03 19:05:31 +08:00
|
|
|
for attempt in range(self.max_retries + 1):
|
2025-06-12 12:31:10 +08:00
|
|
|
history = hist
|
2025-06-27 19:28:41 +08:00
|
|
|
try:
|
2025-08-12 10:59:20 +08:00
|
|
|
for _ in range(self.max_rounds + 1):
|
2025-07-30 19:41:09 +08:00
|
|
|
logging.info(f"{self.tools=}")
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, tools=self.tools, tool_choice="auto", **gen_conf)
|
2025-10-22 12:25:31 +08:00
|
|
|
tk_count += total_token_count_from_response(response)
|
2025-06-27 19:28:41 +08:00
|
|
|
if any([not response.choices, not response.choices[0].message]):
|
|
|
|
|
raise Exception(f"500 response structure error. Response: {response}")
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
if not hasattr(response.choices[0].message, "tool_calls") or not response.choices[0].message.tool_calls:
|
|
|
|
|
if hasattr(response.choices[0].message, "reasoning_content") and response.choices[0].message.reasoning_content:
|
|
|
|
|
ans += "<think>" + response.choices[0].message.reasoning_content + "</think>"
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
ans += response.choices[0].message.content
|
|
|
|
|
if response.choices[0].finish_reason == "length":
|
|
|
|
|
ans = self._length_stop(ans)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
return ans, tk_count
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
for tool_call in response.choices[0].message.tool_calls:
|
2025-07-30 19:41:09 +08:00
|
|
|
logging.info(f"Response {tool_call=}")
|
2025-06-12 12:31:10 +08:00
|
|
|
name = tool_call.function.name
|
|
|
|
|
try:
|
|
|
|
|
args = json_repair.loads(tool_call.function.arguments)
|
2025-07-30 19:41:09 +08:00
|
|
|
tool_response = self.toolcall_session.tool_call(name, args)
|
2025-06-27 19:28:41 +08:00
|
|
|
history = self._append_history(history, tool_call, tool_response)
|
|
|
|
|
ans += self._verbose_tool_use(name, args, tool_response)
|
2025-06-12 12:31:10 +08:00
|
|
|
except Exception as e:
|
2025-06-27 19:28:41 +08:00
|
|
|
logging.exception(msg=f"Wrong JSON argument format in LLM tool call response: {tool_call}")
|
2025-06-12 12:31:10 +08:00
|
|
|
history.append({"role": "tool", "tool_call_id": tool_call.id, "content": f"Tool call error: \n{tool_call}\nException:\n" + str(e)})
|
2025-06-27 19:28:41 +08:00
|
|
|
ans += self._verbose_tool_use(name, {}, str(e))
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
logging.warning(f"Exceed max rounds: {self.max_rounds}")
|
2025-07-30 19:41:09 +08:00
|
|
|
history.append({"role": "user", "content": f"Exceed max rounds: {self.max_rounds}"})
|
|
|
|
|
response, token_count = self._chat(history, gen_conf)
|
|
|
|
|
ans += response
|
|
|
|
|
tk_count += token_count
|
|
|
|
|
return ans, tk_count
|
2025-06-27 19:28:41 +08:00
|
|
|
except Exception as e:
|
|
|
|
|
e = self._exceptions(e, attempt)
|
|
|
|
|
if e:
|
|
|
|
|
return e, tk_count
|
2025-07-30 19:41:09 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
assert False, "Shouldn't be here."
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat(self, system, history, gen_conf={}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-03-27 11:33:46 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2025-06-11 17:20:12 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2025-03-22 23:07:03 +08:00
|
|
|
|
|
|
|
|
# Implement exponential backoff retry strategy
|
2025-06-23 15:59:25 +08:00
|
|
|
for attempt in range(self.max_retries + 1):
|
2025-03-22 23:07:03 +08:00
|
|
|
try:
|
2025-07-30 19:41:09 +08:00
|
|
|
return self._chat(history, gen_conf, **kwargs)
|
2025-03-22 23:07:03 +08:00
|
|
|
except Exception as e:
|
2025-06-12 12:31:10 +08:00
|
|
|
e = self._exceptions(e, attempt)
|
|
|
|
|
if e:
|
|
|
|
|
return e, 0
|
|
|
|
|
assert False, "Shouldn't be here."
|
2023-12-25 19:05:59 +08:00
|
|
|
|
2025-04-08 16:09:03 +08:00
|
|
|
def _wrap_toolcall_message(self, stream):
|
|
|
|
|
final_tool_calls = {}
|
|
|
|
|
|
|
|
|
|
for chunk in stream:
|
|
|
|
|
for tool_call in chunk.choices[0].delta.tool_calls or []:
|
|
|
|
|
index = tool_call.index
|
|
|
|
|
|
|
|
|
|
if index not in final_tool_calls:
|
|
|
|
|
final_tool_calls[index] = tool_call
|
|
|
|
|
|
|
|
|
|
final_tool_calls[index].function.arguments += tool_call.function.arguments
|
|
|
|
|
|
|
|
|
|
return final_tool_calls
|
|
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
def chat_streamly_with_tools(self, system: str, history: list, gen_conf: dict = {}):
|
2025-06-27 19:28:41 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2025-04-08 16:09:03 +08:00
|
|
|
tools = self.tools
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2025-04-08 16:09:03 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
|
|
|
|
|
|
|
|
|
total_tokens = 0
|
2025-06-12 12:31:10 +08:00
|
|
|
hist = deepcopy(history)
|
|
|
|
|
# Implement exponential backoff retry strategy
|
2025-06-23 15:59:25 +08:00
|
|
|
for attempt in range(self.max_retries + 1):
|
2025-06-12 12:31:10 +08:00
|
|
|
history = hist
|
2025-06-27 19:28:41 +08:00
|
|
|
try:
|
2025-08-12 10:59:20 +08:00
|
|
|
for _ in range(self.max_rounds + 1):
|
2025-06-27 19:28:41 +08:00
|
|
|
reasoning_start = False
|
2025-07-30 19:41:09 +08:00
|
|
|
logging.info(f"{tools=}")
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, stream=True, tools=tools, tool_choice="auto", **gen_conf)
|
2025-06-12 12:31:10 +08:00
|
|
|
final_tool_calls = {}
|
|
|
|
|
answer = ""
|
|
|
|
|
for resp in response:
|
|
|
|
|
if resp.choices[0].delta.tool_calls:
|
|
|
|
|
for tool_call in resp.choices[0].delta.tool_calls or []:
|
|
|
|
|
index = tool_call.index
|
|
|
|
|
|
|
|
|
|
if index not in final_tool_calls:
|
2025-06-27 19:28:41 +08:00
|
|
|
if not tool_call.function.arguments:
|
|
|
|
|
tool_call.function.arguments = ""
|
2025-06-12 12:31:10 +08:00
|
|
|
final_tool_calls[index] = tool_call
|
|
|
|
|
else:
|
2025-06-27 19:28:41 +08:00
|
|
|
final_tool_calls[index].function.arguments += tool_call.function.arguments if tool_call.function.arguments else ""
|
2025-04-08 16:09:03 +08:00
|
|
|
continue
|
2025-06-12 12:31:10 +08:00
|
|
|
|
|
|
|
|
if any([not resp.choices, not resp.choices[0].delta, not hasattr(resp.choices[0].delta, "content")]):
|
|
|
|
|
raise Exception("500 response structure error.")
|
|
|
|
|
|
2025-04-08 16:09:03 +08:00
|
|
|
if not resp.choices[0].delta.content:
|
|
|
|
|
resp.choices[0].delta.content = ""
|
2025-06-12 12:31:10 +08:00
|
|
|
|
2025-04-08 16:09:03 +08:00
|
|
|
if hasattr(resp.choices[0].delta, "reasoning_content") and resp.choices[0].delta.reasoning_content:
|
|
|
|
|
ans = ""
|
|
|
|
|
if not reasoning_start:
|
|
|
|
|
reasoning_start = True
|
|
|
|
|
ans = "<think>"
|
|
|
|
|
ans += resp.choices[0].delta.reasoning_content + "</think>"
|
2025-06-12 12:31:10 +08:00
|
|
|
yield ans
|
2025-04-08 16:09:03 +08:00
|
|
|
else:
|
|
|
|
|
reasoning_start = False
|
2025-06-12 12:31:10 +08:00
|
|
|
answer += resp.choices[0].delta.content
|
|
|
|
|
yield resp.choices[0].delta.content
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-04-08 16:09:03 +08:00
|
|
|
if not tol:
|
|
|
|
|
total_tokens += num_tokens_from_string(resp.choices[0].delta.content)
|
|
|
|
|
else:
|
2025-09-05 19:17:21 +08:00
|
|
|
total_tokens = tol
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
finish_reason = resp.choices[0].finish_reason if hasattr(resp.choices[0], "finish_reason") else ""
|
|
|
|
|
if finish_reason == "length":
|
|
|
|
|
yield self._length_stop("")
|
|
|
|
|
|
|
|
|
|
if answer:
|
|
|
|
|
yield total_tokens
|
|
|
|
|
return
|
2025-05-16 16:32:19 +08:00
|
|
|
|
2025-06-12 12:31:10 +08:00
|
|
|
for tool_call in final_tool_calls.values():
|
|
|
|
|
name = tool_call.function.name
|
|
|
|
|
try:
|
|
|
|
|
args = json_repair.loads(tool_call.function.arguments)
|
2025-07-30 19:41:09 +08:00
|
|
|
yield self._verbose_tool_use(name, args, "Begin to call...")
|
|
|
|
|
tool_response = self.toolcall_session.tool_call(name, args)
|
2025-06-27 19:28:41 +08:00
|
|
|
history = self._append_history(history, tool_call, tool_response)
|
|
|
|
|
yield self._verbose_tool_use(name, args, tool_response)
|
2025-06-12 12:31:10 +08:00
|
|
|
except Exception as e:
|
|
|
|
|
logging.exception(msg=f"Wrong JSON argument format in LLM tool call response: {tool_call}")
|
|
|
|
|
history.append({"role": "tool", "tool_call_id": tool_call.id, "content": f"Tool call error: \n{tool_call}\nException:\n" + str(e)})
|
2025-06-27 19:28:41 +08:00
|
|
|
yield self._verbose_tool_use(name, {}, str(e))
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
logging.warning(f"Exceed max rounds: {self.max_rounds}")
|
2025-07-30 19:41:09 +08:00
|
|
|
history.append({"role": "user", "content": f"Exceed max rounds: {self.max_rounds}"})
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, stream=True, **gen_conf)
|
2025-07-30 19:41:09 +08:00
|
|
|
for resp in response:
|
|
|
|
|
if any([not resp.choices, not resp.choices[0].delta, not hasattr(resp.choices[0].delta, "content")]):
|
|
|
|
|
raise Exception("500 response structure error.")
|
|
|
|
|
if not resp.choices[0].delta.content:
|
|
|
|
|
resp.choices[0].delta.content = ""
|
|
|
|
|
continue
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-07-30 19:41:09 +08:00
|
|
|
if not tol:
|
|
|
|
|
total_tokens += num_tokens_from_string(resp.choices[0].delta.content)
|
|
|
|
|
else:
|
2025-09-05 19:17:21 +08:00
|
|
|
total_tokens = tol
|
2025-07-30 19:41:09 +08:00
|
|
|
answer += resp.choices[0].delta.content
|
|
|
|
|
yield resp.choices[0].delta.content
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
|
|
|
|
return
|
|
|
|
|
|
2025-06-27 19:28:41 +08:00
|
|
|
except Exception as e:
|
|
|
|
|
e = self._exceptions(e, attempt)
|
|
|
|
|
if e:
|
2025-07-30 19:41:09 +08:00
|
|
|
yield e
|
2025-06-27 19:28:41 +08:00
|
|
|
yield total_tokens
|
|
|
|
|
return
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
assert False, "Shouldn't be here."
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf: dict = {}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-05-16 20:14:53 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2025-06-27 19:28:41 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2024-05-16 20:14:53 +08:00
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
2025-07-30 19:41:09 +08:00
|
|
|
for delta_ans, tol in self._chat_streamly(history, gen_conf, **kwargs):
|
|
|
|
|
yield delta_ans
|
|
|
|
|
total_tokens += tol
|
2024-05-16 20:14:53 +08:00
|
|
|
except openai.APIError as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
|
|
|
|
|
Dynamic Context Window Size for Ollama Chat (#6582)
# Dynamic Context Window Size for Ollama Chat
## Problem Statement
Previously, the Ollama chat implementation used a fixed context window
size of 32768 tokens. This caused two main issues:
1. Performance degradation due to unnecessarily large context windows
for small conversations
2. Potential business logic failures when using smaller fixed sizes
(e.g., 2048 tokens)
## Solution
Implemented a dynamic context window size calculation that:
1. Uses a base context size of 8192 tokens
2. Applies a 1.2x buffer ratio to the total token count
3. Adds multiples of 8192 tokens based on the buffered token count
4. Implements a smart context size update strategy
## Implementation Details
### Token Counting Logic
```python
def count_tokens(text):
"""Calculate token count for text"""
# Simple calculation: 1 token per ASCII character
# 2 tokens for non-ASCII characters (Chinese, Japanese, Korean, etc.)
total = 0
for char in text:
if ord(char) < 128: # ASCII characters
total += 1
else: # Non-ASCII characters
total += 2
return total
```
### Dynamic Context Calculation
```python
def _calculate_dynamic_ctx(self, history):
"""Calculate dynamic context window size"""
# Calculate total tokens for all messages
total_tokens = 0
for message in history:
content = message.get("content", "")
content_tokens = count_tokens(content)
role_tokens = 4 # Role marker token overhead
total_tokens += content_tokens + role_tokens
# Apply 1.2x buffer ratio
total_tokens_with_buffer = int(total_tokens * 1.2)
# Calculate context size in multiples of 8192
if total_tokens_with_buffer <= 8192:
ctx_size = 8192
else:
ctx_multiplier = (total_tokens_with_buffer // 8192) + 1
ctx_size = ctx_multiplier * 8192
return ctx_size
```
### Integration in Chat Method
```python
def chat(self, system, history, gen_conf):
if system:
history.insert(0, {"role": "system", "content": system})
if "max_tokens" in gen_conf:
del gen_conf["max_tokens"]
try:
# Calculate new context size
new_ctx_size = self._calculate_dynamic_ctx(history)
# Prepare options with context size
options = {
"num_ctx": new_ctx_size
}
# Add other generation options
if "temperature" in gen_conf:
options["temperature"] = gen_conf["temperature"]
if "max_tokens" in gen_conf:
options["num_predict"] = gen_conf["max_tokens"]
if "top_p" in gen_conf:
options["top_p"] = gen_conf["top_p"]
if "presence_penalty" in gen_conf:
options["presence_penalty"] = gen_conf["presence_penalty"]
if "frequency_penalty" in gen_conf:
options["frequency_penalty"] = gen_conf["frequency_penalty"]
# Make API call with dynamic context size
response = self.client.chat(
model=self.model_name,
messages=history,
options=options,
keep_alive=60
)
return response["message"]["content"].strip(), response.get("eval_count", 0) + response.get("prompt_eval_count", 0)
except Exception as e:
return "**ERROR**: " + str(e), 0
```
## Benefits
1. **Improved Performance**: Uses appropriate context windows based on
conversation length
2. **Better Resource Utilization**: Context window size scales with
content
3. **Maintained Compatibility**: Works with existing business logic
4. **Predictable Scaling**: Context growth in 8192-token increments
5. **Smart Updates**: Context size updates are optimized to reduce
unnecessary model reloads
## Future Considerations
1. Fine-tune buffer ratio based on usage patterns
2. Add monitoring for context window utilization
3. Consider language-specific token counting optimizations
4. Implement adaptive threshold based on conversation patterns
5. Add metrics for context size update frequency
---------
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
2025-03-28 12:38:27 +08:00
|
|
|
def _calculate_dynamic_ctx(self, history):
|
|
|
|
|
"""Calculate dynamic context window size"""
|
2025-04-08 16:09:03 +08:00
|
|
|
|
Dynamic Context Window Size for Ollama Chat (#6582)
# Dynamic Context Window Size for Ollama Chat
## Problem Statement
Previously, the Ollama chat implementation used a fixed context window
size of 32768 tokens. This caused two main issues:
1. Performance degradation due to unnecessarily large context windows
for small conversations
2. Potential business logic failures when using smaller fixed sizes
(e.g., 2048 tokens)
## Solution
Implemented a dynamic context window size calculation that:
1. Uses a base context size of 8192 tokens
2. Applies a 1.2x buffer ratio to the total token count
3. Adds multiples of 8192 tokens based on the buffered token count
4. Implements a smart context size update strategy
## Implementation Details
### Token Counting Logic
```python
def count_tokens(text):
"""Calculate token count for text"""
# Simple calculation: 1 token per ASCII character
# 2 tokens for non-ASCII characters (Chinese, Japanese, Korean, etc.)
total = 0
for char in text:
if ord(char) < 128: # ASCII characters
total += 1
else: # Non-ASCII characters
total += 2
return total
```
### Dynamic Context Calculation
```python
def _calculate_dynamic_ctx(self, history):
"""Calculate dynamic context window size"""
# Calculate total tokens for all messages
total_tokens = 0
for message in history:
content = message.get("content", "")
content_tokens = count_tokens(content)
role_tokens = 4 # Role marker token overhead
total_tokens += content_tokens + role_tokens
# Apply 1.2x buffer ratio
total_tokens_with_buffer = int(total_tokens * 1.2)
# Calculate context size in multiples of 8192
if total_tokens_with_buffer <= 8192:
ctx_size = 8192
else:
ctx_multiplier = (total_tokens_with_buffer // 8192) + 1
ctx_size = ctx_multiplier * 8192
return ctx_size
```
### Integration in Chat Method
```python
def chat(self, system, history, gen_conf):
if system:
history.insert(0, {"role": "system", "content": system})
if "max_tokens" in gen_conf:
del gen_conf["max_tokens"]
try:
# Calculate new context size
new_ctx_size = self._calculate_dynamic_ctx(history)
# Prepare options with context size
options = {
"num_ctx": new_ctx_size
}
# Add other generation options
if "temperature" in gen_conf:
options["temperature"] = gen_conf["temperature"]
if "max_tokens" in gen_conf:
options["num_predict"] = gen_conf["max_tokens"]
if "top_p" in gen_conf:
options["top_p"] = gen_conf["top_p"]
if "presence_penalty" in gen_conf:
options["presence_penalty"] = gen_conf["presence_penalty"]
if "frequency_penalty" in gen_conf:
options["frequency_penalty"] = gen_conf["frequency_penalty"]
# Make API call with dynamic context size
response = self.client.chat(
model=self.model_name,
messages=history,
options=options,
keep_alive=60
)
return response["message"]["content"].strip(), response.get("eval_count", 0) + response.get("prompt_eval_count", 0)
except Exception as e:
return "**ERROR**: " + str(e), 0
```
## Benefits
1. **Improved Performance**: Uses appropriate context windows based on
conversation length
2. **Better Resource Utilization**: Context window size scales with
content
3. **Maintained Compatibility**: Works with existing business logic
4. **Predictable Scaling**: Context growth in 8192-token increments
5. **Smart Updates**: Context size updates are optimized to reduce
unnecessary model reloads
## Future Considerations
1. Fine-tune buffer ratio based on usage patterns
2. Add monitoring for context window utilization
3. Consider language-specific token counting optimizations
4. Implement adaptive threshold based on conversation patterns
5. Add metrics for context size update frequency
---------
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
2025-03-28 12:38:27 +08:00
|
|
|
def count_tokens(text):
|
|
|
|
|
"""Calculate token count for text"""
|
|
|
|
|
# Simple calculation: 1 token per ASCII character
|
|
|
|
|
# 2 tokens for non-ASCII characters (Chinese, Japanese, Korean, etc.)
|
|
|
|
|
total = 0
|
|
|
|
|
for char in text:
|
|
|
|
|
if ord(char) < 128: # ASCII characters
|
|
|
|
|
total += 1
|
|
|
|
|
else: # Non-ASCII characters (Chinese, Japanese, Korean, etc.)
|
|
|
|
|
total += 2
|
|
|
|
|
return total
|
|
|
|
|
|
|
|
|
|
# Calculate total tokens for all messages
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
for message in history:
|
|
|
|
|
content = message.get("content", "")
|
|
|
|
|
# Calculate content tokens
|
|
|
|
|
content_tokens = count_tokens(content)
|
|
|
|
|
# Add role marker token overhead
|
|
|
|
|
role_tokens = 4
|
|
|
|
|
total_tokens += content_tokens + role_tokens
|
|
|
|
|
|
|
|
|
|
# Apply 1.2x buffer ratio
|
|
|
|
|
total_tokens_with_buffer = int(total_tokens * 1.2)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
Dynamic Context Window Size for Ollama Chat (#6582)
# Dynamic Context Window Size for Ollama Chat
## Problem Statement
Previously, the Ollama chat implementation used a fixed context window
size of 32768 tokens. This caused two main issues:
1. Performance degradation due to unnecessarily large context windows
for small conversations
2. Potential business logic failures when using smaller fixed sizes
(e.g., 2048 tokens)
## Solution
Implemented a dynamic context window size calculation that:
1. Uses a base context size of 8192 tokens
2. Applies a 1.2x buffer ratio to the total token count
3. Adds multiples of 8192 tokens based on the buffered token count
4. Implements a smart context size update strategy
## Implementation Details
### Token Counting Logic
```python
def count_tokens(text):
"""Calculate token count for text"""
# Simple calculation: 1 token per ASCII character
# 2 tokens for non-ASCII characters (Chinese, Japanese, Korean, etc.)
total = 0
for char in text:
if ord(char) < 128: # ASCII characters
total += 1
else: # Non-ASCII characters
total += 2
return total
```
### Dynamic Context Calculation
```python
def _calculate_dynamic_ctx(self, history):
"""Calculate dynamic context window size"""
# Calculate total tokens for all messages
total_tokens = 0
for message in history:
content = message.get("content", "")
content_tokens = count_tokens(content)
role_tokens = 4 # Role marker token overhead
total_tokens += content_tokens + role_tokens
# Apply 1.2x buffer ratio
total_tokens_with_buffer = int(total_tokens * 1.2)
# Calculate context size in multiples of 8192
if total_tokens_with_buffer <= 8192:
ctx_size = 8192
else:
ctx_multiplier = (total_tokens_with_buffer // 8192) + 1
ctx_size = ctx_multiplier * 8192
return ctx_size
```
### Integration in Chat Method
```python
def chat(self, system, history, gen_conf):
if system:
history.insert(0, {"role": "system", "content": system})
if "max_tokens" in gen_conf:
del gen_conf["max_tokens"]
try:
# Calculate new context size
new_ctx_size = self._calculate_dynamic_ctx(history)
# Prepare options with context size
options = {
"num_ctx": new_ctx_size
}
# Add other generation options
if "temperature" in gen_conf:
options["temperature"] = gen_conf["temperature"]
if "max_tokens" in gen_conf:
options["num_predict"] = gen_conf["max_tokens"]
if "top_p" in gen_conf:
options["top_p"] = gen_conf["top_p"]
if "presence_penalty" in gen_conf:
options["presence_penalty"] = gen_conf["presence_penalty"]
if "frequency_penalty" in gen_conf:
options["frequency_penalty"] = gen_conf["frequency_penalty"]
# Make API call with dynamic context size
response = self.client.chat(
model=self.model_name,
messages=history,
options=options,
keep_alive=60
)
return response["message"]["content"].strip(), response.get("eval_count", 0) + response.get("prompt_eval_count", 0)
except Exception as e:
return "**ERROR**: " + str(e), 0
```
## Benefits
1. **Improved Performance**: Uses appropriate context windows based on
conversation length
2. **Better Resource Utilization**: Context window size scales with
content
3. **Maintained Compatibility**: Works with existing business logic
4. **Predictable Scaling**: Context growth in 8192-token increments
5. **Smart Updates**: Context size updates are optimized to reduce
unnecessary model reloads
## Future Considerations
1. Fine-tune buffer ratio based on usage patterns
2. Add monitoring for context window utilization
3. Consider language-specific token counting optimizations
4. Implement adaptive threshold based on conversation patterns
5. Add metrics for context size update frequency
---------
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
2025-03-28 12:38:27 +08:00
|
|
|
if total_tokens_with_buffer <= 8192:
|
|
|
|
|
ctx_size = 8192
|
|
|
|
|
else:
|
|
|
|
|
ctx_multiplier = (total_tokens_with_buffer // 8192) + 1
|
|
|
|
|
ctx_size = ctx_multiplier * 8192
|
2025-04-08 16:09:03 +08:00
|
|
|
|
Dynamic Context Window Size for Ollama Chat (#6582)
# Dynamic Context Window Size for Ollama Chat
## Problem Statement
Previously, the Ollama chat implementation used a fixed context window
size of 32768 tokens. This caused two main issues:
1. Performance degradation due to unnecessarily large context windows
for small conversations
2. Potential business logic failures when using smaller fixed sizes
(e.g., 2048 tokens)
## Solution
Implemented a dynamic context window size calculation that:
1. Uses a base context size of 8192 tokens
2. Applies a 1.2x buffer ratio to the total token count
3. Adds multiples of 8192 tokens based on the buffered token count
4. Implements a smart context size update strategy
## Implementation Details
### Token Counting Logic
```python
def count_tokens(text):
"""Calculate token count for text"""
# Simple calculation: 1 token per ASCII character
# 2 tokens for non-ASCII characters (Chinese, Japanese, Korean, etc.)
total = 0
for char in text:
if ord(char) < 128: # ASCII characters
total += 1
else: # Non-ASCII characters
total += 2
return total
```
### Dynamic Context Calculation
```python
def _calculate_dynamic_ctx(self, history):
"""Calculate dynamic context window size"""
# Calculate total tokens for all messages
total_tokens = 0
for message in history:
content = message.get("content", "")
content_tokens = count_tokens(content)
role_tokens = 4 # Role marker token overhead
total_tokens += content_tokens + role_tokens
# Apply 1.2x buffer ratio
total_tokens_with_buffer = int(total_tokens * 1.2)
# Calculate context size in multiples of 8192
if total_tokens_with_buffer <= 8192:
ctx_size = 8192
else:
ctx_multiplier = (total_tokens_with_buffer // 8192) + 1
ctx_size = ctx_multiplier * 8192
return ctx_size
```
### Integration in Chat Method
```python
def chat(self, system, history, gen_conf):
if system:
history.insert(0, {"role": "system", "content": system})
if "max_tokens" in gen_conf:
del gen_conf["max_tokens"]
try:
# Calculate new context size
new_ctx_size = self._calculate_dynamic_ctx(history)
# Prepare options with context size
options = {
"num_ctx": new_ctx_size
}
# Add other generation options
if "temperature" in gen_conf:
options["temperature"] = gen_conf["temperature"]
if "max_tokens" in gen_conf:
options["num_predict"] = gen_conf["max_tokens"]
if "top_p" in gen_conf:
options["top_p"] = gen_conf["top_p"]
if "presence_penalty" in gen_conf:
options["presence_penalty"] = gen_conf["presence_penalty"]
if "frequency_penalty" in gen_conf:
options["frequency_penalty"] = gen_conf["frequency_penalty"]
# Make API call with dynamic context size
response = self.client.chat(
model=self.model_name,
messages=history,
options=options,
keep_alive=60
)
return response["message"]["content"].strip(), response.get("eval_count", 0) + response.get("prompt_eval_count", 0)
except Exception as e:
return "**ERROR**: " + str(e), 0
```
## Benefits
1. **Improved Performance**: Uses appropriate context windows based on
conversation length
2. **Better Resource Utilization**: Context window size scales with
content
3. **Maintained Compatibility**: Works with existing business logic
4. **Predictable Scaling**: Context growth in 8192-token increments
5. **Smart Updates**: Context size updates are optimized to reduce
unnecessary model reloads
## Future Considerations
1. Fine-tune buffer ratio based on usage patterns
2. Add monitoring for context window utilization
3. Consider language-specific token counting optimizations
4. Implement adaptive threshold based on conversation patterns
5. Add metrics for context size update frequency
---------
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
2025-03-28 12:38:27 +08:00
|
|
|
return ctx_size
|
2023-12-25 19:05:59 +08:00
|
|
|
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-05-08 10:30:02 +08:00
|
|
|
class GptTurbo(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "OpenAI"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name="gpt-3.5-turbo", base_url="https://api.openai.com/v1", **kwargs):
|
2024-12-08 14:21:12 +08:00
|
|
|
if not base_url:
|
|
|
|
|
base_url = "https://api.openai.com/v1"
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2024-05-08 10:30:02 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class XinferenceChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "Xinference"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key=None, model_name="", base_url="", **kwargs):
|
2024-07-25 10:23:35 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("Local llm url cannot be None")
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin(base_url, "v1")
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2024-10-29 10:42:45 +08:00
|
|
|
|
|
|
|
|
|
2024-10-11 14:45:48 +08:00
|
|
|
class HuggingFaceChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "HuggingFace"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key=None, model_name="", base_url="", **kwargs):
|
2024-10-11 14:45:48 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("Local llm url cannot be None")
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin(base_url, "v1")
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name.split("___")[0], base_url, **kwargs)
|
2024-05-08 10:30:02 +08:00
|
|
|
|
2024-10-29 10:42:45 +08:00
|
|
|
|
2025-02-24 10:12:20 +08:00
|
|
|
class ModelScopeChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "ModelScope"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key=None, model_name="", base_url="", **kwargs):
|
2025-02-24 10:12:20 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("Local llm url cannot be None")
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin(base_url, "v1")
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name.split("___")[0], base_url, **kwargs)
|
2025-02-24 10:12:20 +08:00
|
|
|
|
|
|
|
|
|
2024-07-04 09:57:16 +08:00
|
|
|
class AzureChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "Azure-OpenAI"
|
|
|
|
|
|
2025-06-23 15:59:25 +08:00
|
|
|
def __init__(self, key, model_name, base_url, **kwargs):
|
2025-03-26 19:33:14 +08:00
|
|
|
api_key = json.loads(key).get("api_key", "")
|
|
|
|
|
api_version = json.loads(key).get("api_version", "2024-02-01")
|
2025-06-23 15:59:25 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
|
|
|
|
self.client = AzureOpenAI(api_key=api_key, azure_endpoint=base_url, api_version=api_version)
|
2024-07-04 09:57:16 +08:00
|
|
|
self.model_name = model_name
|
|
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
@property
|
|
|
|
|
def _retryable_errors(self) -> set[str]:
|
|
|
|
|
return {
|
|
|
|
|
LLMErrorCode.ERROR_RATE_LIMIT,
|
|
|
|
|
LLMErrorCode.ERROR_SERVER,
|
|
|
|
|
LLMErrorCode.ERROR_QUOTA,
|
|
|
|
|
}
|
|
|
|
|
|
2024-03-14 19:45:29 +08:00
|
|
|
|
2024-05-28 09:09:37 +08:00
|
|
|
class BaiChuanChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "BaiChuan"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name="Baichuan3-Turbo", base_url="https://api.baichuan-ai.com/v1", **kwargs):
|
2024-05-28 09:09:37 +08:00
|
|
|
if not base_url:
|
|
|
|
|
base_url = "https://api.baichuan-ai.com/v1"
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2024-05-28 09:09:37 +08:00
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _format_params(params):
|
|
|
|
|
return {
|
|
|
|
|
"temperature": params.get("temperature", 0.3),
|
|
|
|
|
"top_p": params.get("top_p", 0.85),
|
|
|
|
|
}
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
|
|
|
|
return {
|
|
|
|
|
"temperature": gen_conf.get("temperature", 0.3),
|
|
|
|
|
"top_p": gen_conf.get("top_p", 0.85),
|
|
|
|
|
}
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat(self, history, gen_conf={}, **kwargs):
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(
|
2025-06-11 17:20:12 +08:00
|
|
|
model=self.model_name,
|
|
|
|
|
messages=history,
|
|
|
|
|
extra_body={"tools": [{"type": "web_search", "web_search": {"enable": True, "search_mode": "performance_first"}}]},
|
|
|
|
|
**gen_conf,
|
|
|
|
|
)
|
|
|
|
|
ans = response.choices[0].message.content.strip()
|
|
|
|
|
if response.choices[0].finish_reason == "length":
|
|
|
|
|
if is_chinese([ans]):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
2025-10-22 12:25:31 +08:00
|
|
|
return ans, total_token_count_from_response(response)
|
2024-05-28 09:09:37 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-05-28 09:09:37 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2024-05-28 09:09:37 +08:00
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(
|
2024-05-28 09:09:37 +08:00
|
|
|
model=self.model_name,
|
|
|
|
|
messages=history,
|
2025-03-26 19:33:14 +08:00
|
|
|
extra_body={"tools": [{"type": "web_search", "web_search": {"enable": True, "search_mode": "performance_first"}}]},
|
2024-05-28 09:09:37 +08:00
|
|
|
stream=True,
|
2025-03-26 19:33:14 +08:00
|
|
|
**self._format_params(gen_conf),
|
|
|
|
|
)
|
2024-05-28 09:09:37 +08:00
|
|
|
for resp in response:
|
2024-12-08 14:21:12 +08:00
|
|
|
if not resp.choices:
|
|
|
|
|
continue
|
2024-05-28 09:09:37 +08:00
|
|
|
if not resp.choices[0].delta.content:
|
2024-10-08 18:27:04 +08:00
|
|
|
resp.choices[0].delta.content = ""
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = resp.choices[0].delta.content
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-01-26 13:54:26 +08:00
|
|
|
if not tol:
|
|
|
|
|
total_tokens += num_tokens_from_string(resp.choices[0].delta.content)
|
|
|
|
|
else:
|
|
|
|
|
total_tokens = tol
|
2024-05-28 09:09:37 +08:00
|
|
|
if resp.choices[0].finish_reason == "length":
|
2024-12-04 09:34:49 +08:00
|
|
|
if is_chinese([ans]):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
2024-05-28 09:09:37 +08:00
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
|
|
|
|
|
|
|
|
|
|
2024-02-08 17:01:01 +08:00
|
|
|
class ZhipuChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "ZHIPU-AI"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name="glm-3-turbo", base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-02-08 17:01:01 +08:00
|
|
|
self.client = ZhipuAI(api_key=key)
|
|
|
|
|
self.model_name = model_name
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2025-09-25 10:28:03 +08:00
|
|
|
gen_conf = self._clean_conf_plealty(gen_conf)
|
|
|
|
|
return gen_conf
|
2025-09-25 14:11:09 +08:00
|
|
|
|
2025-09-25 10:28:03 +08:00
|
|
|
def _clean_conf_plealty(self, gen_conf):
|
2025-06-11 17:20:12 +08:00
|
|
|
if "presence_penalty" in gen_conf:
|
|
|
|
|
del gen_conf["presence_penalty"]
|
|
|
|
|
if "frequency_penalty" in gen_conf:
|
|
|
|
|
del gen_conf["frequency_penalty"]
|
|
|
|
|
return gen_conf
|
2024-03-12 11:57:08 +08:00
|
|
|
|
2025-05-16 16:32:19 +08:00
|
|
|
def chat_with_tools(self, system: str, history: list, gen_conf: dict):
|
2025-09-25 10:28:03 +08:00
|
|
|
gen_conf = self._clean_conf_plealty(gen_conf)
|
2025-05-16 16:32:19 +08:00
|
|
|
|
|
|
|
|
return super().chat_with_tools(system, history, gen_conf)
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-05-16 20:14:53 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2025-09-25 10:28:03 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2024-05-16 20:14:53 +08:00
|
|
|
ans = ""
|
2024-05-30 16:18:15 +08:00
|
|
|
tk_count = 0
|
2024-05-16 20:14:53 +08:00
|
|
|
try:
|
2025-07-30 19:41:09 +08:00
|
|
|
logging.info(json.dumps(history, ensure_ascii=False, indent=2))
|
2025-09-23 16:06:12 +08:00
|
|
|
response = self.client.chat.completions.create(model=self.model_name, messages=history, stream=True, **gen_conf)
|
2024-05-16 20:14:53 +08:00
|
|
|
for resp in response:
|
2024-12-08 14:21:12 +08:00
|
|
|
if not resp.choices[0].delta.content:
|
|
|
|
|
continue
|
2024-05-16 20:14:53 +08:00
|
|
|
delta = resp.choices[0].delta.content
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = delta
|
2024-05-17 17:07:33 +08:00
|
|
|
if resp.choices[0].finish_reason == "length":
|
2024-12-04 09:34:49 +08:00
|
|
|
if is_chinese(ans):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
2025-10-22 12:25:31 +08:00
|
|
|
tk_count = total_token_count_from_response(resp)
|
2024-12-08 14:21:12 +08:00
|
|
|
if resp.choices[0].finish_reason == "stop":
|
2025-10-22 12:25:31 +08:00
|
|
|
tk_count = total_token_count_from_response(resp)
|
2024-05-16 20:14:53 +08:00
|
|
|
yield ans
|
|
|
|
|
except Exception as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield tk_count
|
|
|
|
|
|
2025-05-16 16:32:19 +08:00
|
|
|
def chat_streamly_with_tools(self, system: str, history: list, gen_conf: dict):
|
2025-09-25 10:28:03 +08:00
|
|
|
gen_conf = self._clean_conf_plealty(gen_conf)
|
2025-05-16 16:32:19 +08:00
|
|
|
return super().chat_streamly_with_tools(system, history, gen_conf)
|
|
|
|
|
|
2024-03-27 11:33:46 +08:00
|
|
|
|
2024-07-19 15:50:28 +08:00
|
|
|
class LocalAIChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "LocalAI"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-07-25 10:23:35 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("Local llm url cannot be None")
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin(base_url, "v1")
|
2024-08-07 18:10:42 +08:00
|
|
|
self.client = OpenAI(api_key="empty", base_url=base_url)
|
2024-07-19 15:50:28 +08:00
|
|
|
self.model_name = model_name.split("___")[0]
|
|
|
|
|
|
|
|
|
|
|
2024-05-20 12:23:51 +08:00
|
|
|
class LocalLLM(Base):
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2024-07-30 14:07:00 +08:00
|
|
|
from jina import Client
|
2025-06-23 15:59:25 +08:00
|
|
|
|
2024-07-30 14:07:00 +08:00
|
|
|
self.client = Client(port=12345, protocol="grpc", asyncio=True)
|
2024-05-20 12:40:59 +08:00
|
|
|
|
2024-07-30 14:07:00 +08:00
|
|
|
def _prepare_prompt(self, system, history, gen_conf):
|
2024-11-29 14:52:27 +08:00
|
|
|
from rag.svr.jina_server import Prompt
|
2025-03-26 19:33:14 +08:00
|
|
|
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-05-20 12:40:59 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2024-07-30 14:07:00 +08:00
|
|
|
return Prompt(message=history, gen_conf=gen_conf)
|
|
|
|
|
|
|
|
|
|
def _stream_response(self, endpoint, prompt):
|
2024-11-29 14:52:27 +08:00
|
|
|
from rag.svr.jina_server import Generation
|
2025-03-26 19:33:14 +08:00
|
|
|
|
2024-05-20 12:40:59 +08:00
|
|
|
answer = ""
|
|
|
|
|
try:
|
2025-03-26 19:33:14 +08:00
|
|
|
res = self.client.stream_doc(on=endpoint, inputs=prompt, return_type=Generation)
|
2024-07-30 14:07:00 +08:00
|
|
|
loop = asyncio.get_event_loop()
|
|
|
|
|
try:
|
|
|
|
|
while True:
|
|
|
|
|
answer = loop.run_until_complete(res.__anext__()).text
|
|
|
|
|
yield answer
|
|
|
|
|
except StopAsyncIteration:
|
|
|
|
|
pass
|
2024-05-20 12:40:59 +08:00
|
|
|
except Exception as e:
|
|
|
|
|
yield answer + "\n**ERROR**: " + str(e)
|
2024-07-30 14:07:00 +08:00
|
|
|
yield num_tokens_from_string(answer)
|
2024-05-20 12:40:59 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat(self, system, history, gen_conf={}, **kwargs):
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2024-07-30 14:07:00 +08:00
|
|
|
prompt = self._prepare_prompt(system, history, gen_conf)
|
|
|
|
|
chat_gen = self._stream_response("/chat", prompt)
|
|
|
|
|
ans = next(chat_gen)
|
|
|
|
|
total_tokens = next(chat_gen)
|
|
|
|
|
return ans, total_tokens
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2024-07-30 14:07:00 +08:00
|
|
|
prompt = self._prepare_prompt(system, history, gen_conf)
|
|
|
|
|
return self._stream_response("/stream", prompt)
|
2024-05-23 11:15:29 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class VolcEngineChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "VolcEngine"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url="https://ark.cn-beijing.volces.com/api/v3", **kwargs):
|
2024-05-23 11:15:29 +08:00
|
|
|
"""
|
|
|
|
|
Since do not want to modify the original database fields, and the VolcEngine authentication method is quite special,
|
2024-08-26 13:34:29 +08:00
|
|
|
Assemble ark_api_key, ep_id into api_key, store it as a dictionary type, and parse it for use
|
2024-05-23 11:15:29 +08:00
|
|
|
model_name is for display only
|
|
|
|
|
"""
|
2025-03-26 19:33:14 +08:00
|
|
|
base_url = base_url if base_url else "https://ark.cn-beijing.volces.com/api/v3"
|
|
|
|
|
ark_api_key = json.loads(key).get("ark_api_key", "")
|
|
|
|
|
model_name = json.loads(key).get("ep_id", "") + json.loads(key).get("endpoint_id", "")
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(ark_api_key, model_name, base_url, **kwargs)
|
2024-05-31 16:38:53 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class MiniMaxChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "MiniMax"
|
|
|
|
|
|
2025-06-23 15:59:25 +08:00
|
|
|
def __init__(self, key, model_name, base_url="https://api.minimax.chat/v1/text/chatcompletion_v2", **kwargs):
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-05-31 16:38:53 +08:00
|
|
|
if not base_url:
|
2024-07-17 15:32:51 +08:00
|
|
|
base_url = "https://api.minimax.chat/v1/text/chatcompletion_v2"
|
|
|
|
|
self.base_url = base_url
|
|
|
|
|
self.model_name = model_name
|
|
|
|
|
self.api_key = key
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
2024-07-17 15:32:51 +08:00
|
|
|
for k in list(gen_conf.keys()):
|
|
|
|
|
if k not in ["temperature", "top_p", "max_tokens"]:
|
|
|
|
|
del gen_conf[k]
|
2025-06-11 17:20:12 +08:00
|
|
|
return gen_conf
|
|
|
|
|
|
|
|
|
|
def _chat(self, history, gen_conf):
|
2024-07-17 15:32:51 +08:00
|
|
|
headers = {
|
|
|
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
|
|
|
"Content-Type": "application/json",
|
|
|
|
|
}
|
2025-03-26 19:33:14 +08:00
|
|
|
payload = json.dumps({"model": self.model_name, "messages": history, **gen_conf})
|
2025-06-11 17:20:12 +08:00
|
|
|
response = requests.request("POST", url=self.base_url, headers=headers, data=payload)
|
|
|
|
|
response = response.json()
|
|
|
|
|
ans = response["choices"][0]["message"]["content"].strip()
|
|
|
|
|
if response["choices"][0]["finish_reason"] == "length":
|
|
|
|
|
if is_chinese(ans):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
2025-10-22 12:25:31 +08:00
|
|
|
return ans, total_token_count_from_response(response)
|
2024-07-17 15:32:51 +08:00
|
|
|
|
|
|
|
|
def chat_streamly(self, system, history, gen_conf):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-07-17 15:32:51 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2025-03-06 11:29:40 +08:00
|
|
|
for k in list(gen_conf.keys()):
|
|
|
|
|
if k not in ["temperature", "top_p", "max_tokens"]:
|
|
|
|
|
del gen_conf[k]
|
2024-07-17 15:32:51 +08:00
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
|
|
|
|
headers = {
|
|
|
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
|
|
|
"Content-Type": "application/json",
|
|
|
|
|
}
|
|
|
|
|
payload = json.dumps(
|
|
|
|
|
{
|
|
|
|
|
"model": self.model_name,
|
|
|
|
|
"messages": history,
|
|
|
|
|
"stream": True,
|
|
|
|
|
**gen_conf,
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
response = requests.request(
|
|
|
|
|
"POST",
|
|
|
|
|
url=self.base_url,
|
|
|
|
|
headers=headers,
|
|
|
|
|
data=payload,
|
|
|
|
|
)
|
|
|
|
|
for resp in response.text.split("\n\n")[:-1]:
|
|
|
|
|
resp = json.loads(resp[6:])
|
2024-07-29 19:35:16 +08:00
|
|
|
text = ""
|
|
|
|
|
if "choices" in resp and "delta" in resp["choices"][0]:
|
2024-07-17 15:32:51 +08:00
|
|
|
text = resp["choices"][0]["delta"]["content"]
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = text
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-01-26 13:54:26 +08:00
|
|
|
if not tol:
|
|
|
|
|
total_tokens += num_tokens_from_string(text)
|
|
|
|
|
else:
|
|
|
|
|
total_tokens = tol
|
2024-07-17 15:32:51 +08:00
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
2024-06-14 11:32:58 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class MistralChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "Mistral"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-06-14 11:32:58 +08:00
|
|
|
from mistralai.client import MistralClient
|
2025-03-26 19:33:14 +08:00
|
|
|
|
2024-06-14 11:32:58 +08:00
|
|
|
self.client = MistralClient(api_key=key)
|
|
|
|
|
self.model_name = model_name
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
2024-06-14 11:32:58 +08:00
|
|
|
for k in list(gen_conf.keys()):
|
|
|
|
|
if k not in ["temperature", "top_p", "max_tokens"]:
|
|
|
|
|
del gen_conf[k]
|
2025-06-11 17:20:12 +08:00
|
|
|
return gen_conf
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat(self, history, gen_conf={}, **kwargs):
|
2025-09-24 10:49:34 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2025-06-11 17:20:12 +08:00
|
|
|
response = self.client.chat(model=self.model_name, messages=history, **gen_conf)
|
|
|
|
|
ans = response.choices[0].message.content
|
|
|
|
|
if response.choices[0].finish_reason == "length":
|
|
|
|
|
if is_chinese(ans):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
2025-10-22 12:25:31 +08:00
|
|
|
return ans, total_token_count_from_response(response)
|
2024-06-14 11:32:58 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-06-14 11:32:58 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
2025-09-24 10:49:34 +08:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2024-06-14 11:32:58 +08:00
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
2025-07-30 19:41:09 +08:00
|
|
|
response = self.client.chat_stream(model=self.model_name, messages=history, **gen_conf, **kwargs)
|
2024-06-14 11:32:58 +08:00
|
|
|
for resp in response:
|
2024-12-08 14:21:12 +08:00
|
|
|
if not resp.choices or not resp.choices[0].delta.content:
|
|
|
|
|
continue
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = resp.choices[0].delta.content
|
2024-06-14 11:32:58 +08:00
|
|
|
total_tokens += 1
|
|
|
|
|
if resp.choices[0].finish_reason == "length":
|
2024-12-04 09:34:49 +08:00
|
|
|
if is_chinese(ans):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
2024-06-14 11:32:58 +08:00
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except openai.APIError as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
2024-07-08 09:37:34 +08:00
|
|
|
|
|
|
|
|
|
2024-07-24 12:46:43 +08:00
|
|
|
class LmStudioChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "LM-Studio"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url, **kwargs):
|
2024-07-24 12:46:43 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("Local llm url cannot be None")
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin(base_url, "v1")
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2024-08-06 16:20:21 +08:00
|
|
|
self.client = OpenAI(api_key="lm-studio", base_url=base_url)
|
2024-07-24 12:46:43 +08:00
|
|
|
self.model_name = model_name
|
2024-08-06 16:20:21 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class OpenAI_APIChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = ["VLLM", "OpenAI-API-Compatible"]
|
|
|
|
|
|
2025-08-07 08:45:37 +07:00
|
|
|
def __init__(self, key, model_name, base_url, **kwargs):
|
2024-08-06 16:20:21 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("url cannot be None")
|
|
|
|
|
model_name = model_name.split("___")[0]
|
2025-08-07 08:45:37 +07:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2024-08-07 18:40:51 +08:00
|
|
|
|
|
|
|
|
|
2024-08-08 12:09:50 +08:00
|
|
|
class LeptonAIChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "LeptonAI"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
2024-08-08 12:09:50 +08:00
|
|
|
if not base_url:
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin("https://" + model_name + ".lepton.run", "api/v1")
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2024-08-12 10:11:50 +08:00
|
|
|
|
|
|
|
|
|
2024-08-19 10:36:57 +08:00
|
|
|
class ReplicateChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "Replicate"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-08-19 10:36:57 +08:00
|
|
|
from replicate.client import Client
|
|
|
|
|
|
|
|
|
|
self.model_name = model_name
|
|
|
|
|
self.client = Client(api_token=key)
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat(self, history, gen_conf={}, **kwargs):
|
2025-06-11 17:20:12 +08:00
|
|
|
system = history[0]["content"] if history and history[0]["role"] == "system" else ""
|
|
|
|
|
prompt = "\n".join([item["role"] + ":" + item["content"] for item in history[-5:] if item["role"] != "system"])
|
|
|
|
|
response = self.client.run(
|
|
|
|
|
self.model_name,
|
|
|
|
|
input={"system_prompt": system, "prompt": prompt, **gen_conf},
|
|
|
|
|
)
|
|
|
|
|
ans = "".join(response)
|
|
|
|
|
return ans, num_tokens_from_string(ans)
|
2024-08-19 10:36:57 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
2024-08-19 10:36:57 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
2025-03-06 11:29:40 +08:00
|
|
|
del gen_conf["max_tokens"]
|
2025-03-26 19:33:14 +08:00
|
|
|
prompt = "\n".join([item["role"] + ":" + item["content"] for item in history[-5:]])
|
2024-08-19 10:36:57 +08:00
|
|
|
ans = ""
|
|
|
|
|
try:
|
|
|
|
|
response = self.client.run(
|
|
|
|
|
self.model_name,
|
2025-06-11 17:20:12 +08:00
|
|
|
input={"system_prompt": system, "prompt": prompt, **gen_conf},
|
2024-08-19 10:36:57 +08:00
|
|
|
)
|
|
|
|
|
for resp in response:
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = resp
|
2024-08-19 10:36:57 +08:00
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield num_tokens_from_string(ans)
|
2024-08-20 15:27:13 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class HunyuanChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "Tencent Hunyuan"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-08-20 15:27:13 +08:00
|
|
|
from tencentcloud.common import credential
|
|
|
|
|
from tencentcloud.hunyuan.v20230901 import hunyuan_client
|
|
|
|
|
|
|
|
|
|
key = json.loads(key)
|
|
|
|
|
sid = key.get("hunyuan_sid", "")
|
|
|
|
|
sk = key.get("hunyuan_sk", "")
|
|
|
|
|
cred = credential.Credential(sid, sk)
|
|
|
|
|
self.model_name = model_name
|
|
|
|
|
self.client = hunyuan_client.HunyuanClient(cred, "")
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
2024-08-20 15:27:13 +08:00
|
|
|
_gen_conf = {}
|
|
|
|
|
if "temperature" in gen_conf:
|
|
|
|
|
_gen_conf["Temperature"] = gen_conf["temperature"]
|
|
|
|
|
if "top_p" in gen_conf:
|
|
|
|
|
_gen_conf["TopP"] = gen_conf["top_p"]
|
2025-06-11 17:20:12 +08:00
|
|
|
return _gen_conf
|
|
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat(self, history, gen_conf={}, **kwargs):
|
2025-06-11 17:20:12 +08:00
|
|
|
from tencentcloud.hunyuan.v20230901 import models
|
2024-08-20 15:27:13 +08:00
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
hist = [{k.capitalize(): v for k, v in item.items()} for item in history]
|
2024-08-20 15:27:13 +08:00
|
|
|
req = models.ChatCompletionsRequest()
|
2025-06-11 17:20:12 +08:00
|
|
|
params = {"Model": self.model_name, "Messages": hist, **gen_conf}
|
2024-08-20 15:27:13 +08:00
|
|
|
req.from_json_string(json.dumps(params))
|
2025-06-11 17:20:12 +08:00
|
|
|
response = self.client.ChatCompletions(req)
|
|
|
|
|
ans = response.Choices[0].Message.Content
|
|
|
|
|
return ans, response.Usage.TotalTokens
|
2024-08-20 15:27:13 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
2024-08-20 15:27:13 +08:00
|
|
|
from tencentcloud.common.exception.tencent_cloud_sdk_exception import (
|
|
|
|
|
TencentCloudSDKException,
|
|
|
|
|
)
|
2025-03-26 19:33:14 +08:00
|
|
|
from tencentcloud.hunyuan.v20230901 import models
|
2024-08-20 16:56:42 +08:00
|
|
|
|
2024-08-20 15:27:13 +08:00
|
|
|
_gen_conf = {}
|
2024-10-08 18:27:04 +08:00
|
|
|
_history = [{k.capitalize(): v for k, v in item.items()} for item in history]
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2024-08-20 15:27:13 +08:00
|
|
|
_history.insert(0, {"Role": "system", "Content": system})
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2024-08-20 15:27:13 +08:00
|
|
|
if "temperature" in gen_conf:
|
|
|
|
|
_gen_conf["Temperature"] = gen_conf["temperature"]
|
|
|
|
|
if "top_p" in gen_conf:
|
|
|
|
|
_gen_conf["TopP"] = gen_conf["top_p"]
|
|
|
|
|
req = models.ChatCompletionsRequest()
|
|
|
|
|
params = {
|
|
|
|
|
"Model": self.model_name,
|
|
|
|
|
"Messages": _history,
|
|
|
|
|
"Stream": True,
|
|
|
|
|
**_gen_conf,
|
|
|
|
|
}
|
|
|
|
|
req.from_json_string(json.dumps(params))
|
|
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
|
|
|
|
response = self.client.ChatCompletions(req)
|
|
|
|
|
for resp in response:
|
|
|
|
|
resp = json.loads(resp["data"])
|
|
|
|
|
if not resp["Choices"] or not resp["Choices"][0]["Delta"]["Content"]:
|
|
|
|
|
continue
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = resp["Choices"][0]["Delta"]["Content"]
|
2024-08-20 15:27:13 +08:00
|
|
|
total_tokens += 1
|
|
|
|
|
|
|
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except TencentCloudSDKException as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
2024-08-20 16:56:42 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class SparkChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "XunFei Spark"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url="https://spark-api-open.xf-yun.com/v1", **kwargs):
|
2024-08-20 16:56:42 +08:00
|
|
|
if not base_url:
|
|
|
|
|
base_url = "https://spark-api-open.xf-yun.com/v1"
|
|
|
|
|
model2version = {
|
|
|
|
|
"Spark-Max": "generalv3.5",
|
|
|
|
|
"Spark-Lite": "general",
|
|
|
|
|
"Spark-Pro": "generalv3",
|
|
|
|
|
"Spark-Pro-128K": "pro-128k",
|
|
|
|
|
"Spark-4.0-Ultra": "4.0Ultra",
|
|
|
|
|
}
|
2024-11-20 12:16:36 +08:00
|
|
|
version2model = {v: k for k, v in model2version.items()}
|
|
|
|
|
assert model_name in model2version or model_name in version2model, f"The given model name is not supported yet. Support: {list(model2version.keys())}"
|
|
|
|
|
if model_name in model2version:
|
|
|
|
|
model_version = model2version[model_name]
|
2024-12-08 14:21:12 +08:00
|
|
|
else:
|
|
|
|
|
model_version = model_name
|
2025-06-11 17:20:12 +08:00
|
|
|
super().__init__(key, model_version, base_url, **kwargs)
|
2024-08-22 16:45:15 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class BaiduYiyanChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "BaiduYiyan"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-08-22 16:45:15 +08:00
|
|
|
import qianfan
|
2024-08-29 13:30:06 +08:00
|
|
|
|
2024-08-22 16:45:15 +08:00
|
|
|
key = json.loads(key)
|
2024-10-08 18:27:04 +08:00
|
|
|
ak = key.get("yiyan_ak", "")
|
|
|
|
|
sk = key.get("yiyan_sk", "")
|
|
|
|
|
self.client = qianfan.ChatCompletion(ak=ak, sk=sk)
|
2024-08-22 16:45:15 +08:00
|
|
|
self.model_name = model_name.lower()
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
2025-03-26 19:33:14 +08:00
|
|
|
gen_conf["penalty_score"] = ((gen_conf.get("presence_penalty", 0) + gen_conf.get("frequency_penalty", 0)) / 2) + 1
|
2024-08-22 16:45:15 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
2025-03-06 11:29:40 +08:00
|
|
|
del gen_conf["max_tokens"]
|
2025-06-11 17:20:12 +08:00
|
|
|
return gen_conf
|
2024-08-29 13:30:06 +08:00
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _chat(self, history, gen_conf):
|
|
|
|
|
system = history[0]["content"] if history and history[0]["role"] == "system" else ""
|
|
|
|
|
response = self.client.do(model=self.model_name, messages=[h for h in history if h["role"] != "system"], system=system, **gen_conf).body
|
|
|
|
|
ans = response["result"]
|
2025-10-22 12:25:31 +08:00
|
|
|
return ans, total_token_count_from_response(response)
|
2024-08-22 16:45:15 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
2025-03-26 19:33:14 +08:00
|
|
|
gen_conf["penalty_score"] = ((gen_conf.get("presence_penalty", 0) + gen_conf.get("frequency_penalty", 0)) / 2) + 1
|
2024-08-22 16:45:15 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
2025-03-06 11:29:40 +08:00
|
|
|
del gen_conf["max_tokens"]
|
2024-08-22 16:45:15 +08:00
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
2024-08-29 13:30:06 +08:00
|
|
|
|
2024-08-22 16:45:15 +08:00
|
|
|
try:
|
2025-06-11 17:20:12 +08:00
|
|
|
response = self.client.do(model=self.model_name, messages=history, system=system, stream=True, **gen_conf)
|
2024-08-22 16:45:15 +08:00
|
|
|
for resp in response:
|
|
|
|
|
resp = resp.body
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = resp["result"]
|
2025-10-22 12:25:31 +08:00
|
|
|
total_tokens = total_token_count_from_response(resp)
|
2024-08-22 16:45:15 +08:00
|
|
|
|
|
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
return ans + "\n**ERROR**: " + str(e), 0
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
2024-08-29 13:30:06 +08:00
|
|
|
|
|
|
|
|
|
2024-09-02 12:06:41 +08:00
|
|
|
class GoogleChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "Google Cloud"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
super().__init__(key, model_name, base_url=base_url, **kwargs)
|
2025-04-08 16:09:03 +08:00
|
|
|
|
2024-09-02 12:06:41 +08:00
|
|
|
import base64
|
|
|
|
|
|
2025-03-26 19:33:14 +08:00
|
|
|
from google.oauth2 import service_account
|
|
|
|
|
|
2025-02-07 12:00:19 +08:00
|
|
|
key = json.loads(key)
|
2025-03-26 19:33:14 +08:00
|
|
|
access_token = json.loads(base64.b64decode(key.get("google_service_account_key", "")))
|
2024-09-02 12:06:41 +08:00
|
|
|
project_id = key.get("google_project_id", "")
|
|
|
|
|
region = key.get("google_region", "")
|
|
|
|
|
|
|
|
|
|
scopes = ["https://www.googleapis.com/auth/cloud-platform"]
|
|
|
|
|
self.model_name = model_name
|
|
|
|
|
|
|
|
|
|
if "claude" in self.model_name:
|
|
|
|
|
from anthropic import AnthropicVertex
|
|
|
|
|
from google.auth.transport.requests import Request
|
|
|
|
|
|
|
|
|
|
if access_token:
|
2025-03-26 19:33:14 +08:00
|
|
|
credits = service_account.Credentials.from_service_account_info(access_token, scopes=scopes)
|
2024-09-02 12:06:41 +08:00
|
|
|
request = Request()
|
|
|
|
|
credits.refresh(request)
|
|
|
|
|
token = credits.token
|
2025-03-26 19:33:14 +08:00
|
|
|
self.client = AnthropicVertex(region=region, project_id=project_id, access_token=token)
|
2024-09-02 12:06:41 +08:00
|
|
|
else:
|
|
|
|
|
self.client = AnthropicVertex(region=region, project_id=project_id)
|
|
|
|
|
else:
|
2025-10-15 08:54:20 +02:00
|
|
|
from google import genai
|
2024-09-02 12:06:41 +08:00
|
|
|
|
|
|
|
|
if access_token:
|
2025-10-15 08:54:20 +02:00
|
|
|
credits = service_account.Credentials.from_service_account_info(access_token, scopes=scopes)
|
|
|
|
|
self.client = genai.Client(vertexai=True, project=project_id, location=region, credentials=credits)
|
2024-09-02 12:06:41 +08:00
|
|
|
else:
|
2025-10-15 08:54:20 +02:00
|
|
|
self.client = genai.Client(vertexai=True, project=project_id, location=region)
|
2024-09-02 12:06:41 +08:00
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def _clean_conf(self, gen_conf):
|
2024-09-02 12:06:41 +08:00
|
|
|
if "claude" in self.model_name:
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2024-09-02 12:06:41 +08:00
|
|
|
else:
|
|
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
gen_conf["max_output_tokens"] = gen_conf["max_tokens"]
|
2025-10-10 13:18:24 +02:00
|
|
|
del gen_conf["max_tokens"]
|
2024-09-02 12:06:41 +08:00
|
|
|
for k in list(gen_conf.keys()):
|
|
|
|
|
if k not in ["temperature", "top_p", "max_output_tokens"]:
|
|
|
|
|
del gen_conf[k]
|
2025-06-11 17:20:12 +08:00
|
|
|
return gen_conf
|
2024-09-02 12:06:41 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def _chat(self, history, gen_conf={}, **kwargs):
|
2025-06-11 17:20:12 +08:00
|
|
|
system = history[0]["content"] if history and history[0]["role"] == "system" else ""
|
2025-10-15 08:54:20 +02:00
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
if "claude" in self.model_name:
|
2025-10-15 08:54:20 +02:00
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
2025-06-11 17:20:12 +08:00
|
|
|
response = self.client.messages.create(
|
|
|
|
|
model=self.model_name,
|
|
|
|
|
messages=[h for h in history if h["role"] != "system"],
|
|
|
|
|
system=system,
|
|
|
|
|
stream=False,
|
|
|
|
|
**gen_conf,
|
|
|
|
|
).json()
|
|
|
|
|
ans = response["content"][0]["text"]
|
|
|
|
|
if response["stop_reason"] == "max_tokens":
|
|
|
|
|
ans += "...\nFor the content length reason, it stopped, continue?" if is_english([ans]) else "······\n由于长度的原因,回答被截断了,要继续吗?"
|
|
|
|
|
return (
|
|
|
|
|
ans,
|
|
|
|
|
response["usage"]["input_tokens"] + response["usage"]["output_tokens"],
|
|
|
|
|
)
|
2024-09-02 12:06:41 +08:00
|
|
|
|
2025-10-15 08:54:20 +02:00
|
|
|
# Gemini models with google-genai SDK
|
|
|
|
|
# Set default thinking_budget=0 if not specified
|
|
|
|
|
if "thinking_budget" not in gen_conf:
|
|
|
|
|
gen_conf["thinking_budget"] = 0
|
|
|
|
|
|
|
|
|
|
thinking_budget = gen_conf.pop("thinking_budget", 0)
|
|
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
|
|
|
|
|
|
|
|
|
# Build GenerateContentConfig
|
|
|
|
|
try:
|
|
|
|
|
from google.genai.types import GenerateContentConfig, ThinkingConfig, Content, Part
|
|
|
|
|
except ImportError as e:
|
|
|
|
|
logging.error(f"[GoogleChat] Failed to import google-genai: {e}. Please install: pip install google-genai>=1.41.0")
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
config_dict = {}
|
|
|
|
|
if system:
|
|
|
|
|
config_dict["system_instruction"] = system
|
|
|
|
|
if "temperature" in gen_conf:
|
|
|
|
|
config_dict["temperature"] = gen_conf["temperature"]
|
|
|
|
|
if "top_p" in gen_conf:
|
|
|
|
|
config_dict["top_p"] = gen_conf["top_p"]
|
|
|
|
|
if "max_output_tokens" in gen_conf:
|
|
|
|
|
config_dict["max_output_tokens"] = gen_conf["max_output_tokens"]
|
|
|
|
|
|
|
|
|
|
# Add ThinkingConfig
|
|
|
|
|
config_dict["thinking_config"] = ThinkingConfig(thinking_budget=thinking_budget)
|
|
|
|
|
|
|
|
|
|
config = GenerateContentConfig(**config_dict)
|
|
|
|
|
|
|
|
|
|
# Convert history to google-genai Content format
|
|
|
|
|
contents = []
|
2025-06-11 17:20:12 +08:00
|
|
|
for item in history:
|
|
|
|
|
if item["role"] == "system":
|
|
|
|
|
continue
|
2025-10-15 08:54:20 +02:00
|
|
|
# google-genai uses 'model' instead of 'assistant'
|
|
|
|
|
role = "model" if item["role"] == "assistant" else item["role"]
|
|
|
|
|
content = Content(
|
|
|
|
|
role=role,
|
|
|
|
|
parts=[Part(text=item["content"])]
|
|
|
|
|
)
|
|
|
|
|
contents.append(content)
|
|
|
|
|
|
|
|
|
|
response = self.client.models.generate_content(
|
|
|
|
|
model=self.model_name,
|
|
|
|
|
contents=contents,
|
|
|
|
|
config=config
|
|
|
|
|
)
|
2025-06-11 17:20:12 +08:00
|
|
|
|
|
|
|
|
ans = response.text
|
2025-10-15 08:54:20 +02:00
|
|
|
# Get token count from response
|
|
|
|
|
try:
|
|
|
|
|
total_tokens = response.usage_metadata.total_token_count
|
|
|
|
|
except Exception:
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
|
|
|
|
|
return ans, total_tokens
|
2025-06-11 17:20:12 +08:00
|
|
|
|
2025-07-30 19:41:09 +08:00
|
|
|
def chat_streamly(self, system, history, gen_conf={}, **kwargs):
|
2024-09-02 12:06:41 +08:00
|
|
|
if "claude" in self.model_name:
|
2025-03-06 11:29:40 +08:00
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
2024-09-02 12:06:41 +08:00
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
|
|
|
|
response = self.client.messages.create(
|
|
|
|
|
model=self.model_name,
|
|
|
|
|
messages=history,
|
2025-06-11 17:20:12 +08:00
|
|
|
system=system,
|
2024-09-02 12:06:41 +08:00
|
|
|
stream=True,
|
|
|
|
|
**gen_conf,
|
|
|
|
|
)
|
|
|
|
|
for res in response.iter_lines():
|
|
|
|
|
res = res.decode("utf-8")
|
|
|
|
|
if "content_block_delta" in res and "data" in res:
|
|
|
|
|
text = json.loads(res[6:])["delta"]["text"]
|
2025-03-26 19:33:14 +08:00
|
|
|
ans = text
|
2024-09-02 12:06:41 +08:00
|
|
|
total_tokens += num_tokens_from_string(text)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
|
|
|
|
else:
|
2025-10-15 08:54:20 +02:00
|
|
|
# Gemini models with google-genai SDK
|
|
|
|
|
ans = ""
|
2025-10-10 13:18:24 +02:00
|
|
|
total_tokens = 0
|
2025-10-15 08:54:20 +02:00
|
|
|
|
|
|
|
|
# Set default thinking_budget=0 if not specified
|
|
|
|
|
if "thinking_budget" not in gen_conf:
|
|
|
|
|
gen_conf["thinking_budget"] = 0
|
|
|
|
|
|
|
|
|
|
thinking_budget = gen_conf.pop("thinking_budget", 0)
|
|
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
|
|
|
|
|
|
|
|
|
# Build GenerateContentConfig
|
|
|
|
|
try:
|
|
|
|
|
from google.genai.types import GenerateContentConfig, ThinkingConfig, Content, Part
|
|
|
|
|
except ImportError as e:
|
|
|
|
|
logging.error(f"[GoogleChat] Failed to import google-genai: {e}. Please install: pip install google-genai>=1.41.0")
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
config_dict = {}
|
|
|
|
|
if system:
|
|
|
|
|
config_dict["system_instruction"] = system
|
|
|
|
|
if "temperature" in gen_conf:
|
|
|
|
|
config_dict["temperature"] = gen_conf["temperature"]
|
|
|
|
|
if "top_p" in gen_conf:
|
|
|
|
|
config_dict["top_p"] = gen_conf["top_p"]
|
|
|
|
|
if "max_output_tokens" in gen_conf:
|
|
|
|
|
config_dict["max_output_tokens"] = gen_conf["max_output_tokens"]
|
|
|
|
|
|
|
|
|
|
# Add ThinkingConfig
|
|
|
|
|
config_dict["thinking_config"] = ThinkingConfig(thinking_budget=thinking_budget)
|
|
|
|
|
|
|
|
|
|
config = GenerateContentConfig(**config_dict)
|
|
|
|
|
|
|
|
|
|
# Convert history to google-genai Content format
|
|
|
|
|
contents = []
|
2024-09-02 12:06:41 +08:00
|
|
|
for item in history:
|
2025-10-15 08:54:20 +02:00
|
|
|
# google-genai uses 'model' instead of 'assistant'
|
|
|
|
|
role = "model" if item["role"] == "assistant" else item["role"]
|
|
|
|
|
content = Content(
|
|
|
|
|
role=role,
|
|
|
|
|
parts=[Part(text=item["content"])]
|
|
|
|
|
)
|
|
|
|
|
contents.append(content)
|
|
|
|
|
|
2024-09-02 12:06:41 +08:00
|
|
|
try:
|
2025-10-15 08:54:20 +02:00
|
|
|
for chunk in self.client.models.generate_content_stream(
|
|
|
|
|
model=self.model_name,
|
|
|
|
|
contents=contents,
|
|
|
|
|
config=config
|
|
|
|
|
):
|
|
|
|
|
text = chunk.text
|
|
|
|
|
ans = text
|
|
|
|
|
total_tokens += num_tokens_from_string(text)
|
2024-09-02 12:06:41 +08:00
|
|
|
yield ans
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
2025-10-10 13:18:24 +02:00
|
|
|
yield total_tokens
|
2025-01-15 14:15:58 +08:00
|
|
|
|
2025-03-06 11:29:40 +08:00
|
|
|
|
2025-01-15 14:15:58 +08:00
|
|
|
class GPUStackChat(Base):
|
2025-07-03 19:05:31 +08:00
|
|
|
_FACTORY_NAME = "GPUStack"
|
|
|
|
|
|
2025-06-11 17:20:12 +08:00
|
|
|
def __init__(self, key=None, model_name="", base_url="", **kwargs):
|
2025-01-15 14:15:58 +08:00
|
|
|
if not base_url:
|
|
|
|
|
raise ValueError("Local llm url cannot be None")
|
2025-06-03 14:18:40 +08:00
|
|
|
base_url = urljoin(base_url, "v1")
|
2025-07-03 19:05:31 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
2025-07-23 18:10:35 +08:00
|
|
|
|
2025-07-31 14:48:30 +08:00
|
|
|
|
2025-09-11 17:25:31 +08:00
|
|
|
class TokenPonyChat(Base):
|
|
|
|
|
_FACTORY_NAME = "TokenPony"
|
|
|
|
|
|
|
|
|
|
def __init__(self, key, model_name, base_url="https://ragflow.vip-api.tokenpony.cn/v1", **kwargs):
|
|
|
|
|
if not base_url:
|
|
|
|
|
base_url = "https://ragflow.vip-api.tokenpony.cn/v1"
|
2025-10-28 10:04:41 +08:00
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
|
|
|
|
|
2025-09-11 17:25:31 +08:00
|
|
|
|
2025-10-09 11:14:49 +08:00
|
|
|
class DeerAPIChat(Base):
|
|
|
|
|
_FACTORY_NAME = "DeerAPI"
|
|
|
|
|
|
|
|
|
|
def __init__(self, key, model_name, base_url="https://api.deerapi.com/v1", **kwargs):
|
|
|
|
|
if not base_url:
|
|
|
|
|
base_url = "https://api.deerapi.com/v1"
|
|
|
|
|
super().__init__(key, model_name, base_url, **kwargs)
|
|
|
|
|
|
2025-09-08 19:00:52 +08:00
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
class LiteLLMBase(ABC):
|
2025-09-18 17:16:59 +08:00
|
|
|
_FACTORY_NAME = [
|
|
|
|
|
"Tongyi-Qianwen",
|
|
|
|
|
"Bedrock",
|
|
|
|
|
"Moonshot",
|
|
|
|
|
"xAI",
|
|
|
|
|
"DeepInfra",
|
|
|
|
|
"Groq",
|
|
|
|
|
"Cohere",
|
|
|
|
|
"Gemini",
|
|
|
|
|
"DeepSeek",
|
|
|
|
|
"NVIDIA",
|
|
|
|
|
"TogetherAI",
|
|
|
|
|
"Anthropic",
|
|
|
|
|
"Ollama",
|
|
|
|
|
"Meituan",
|
|
|
|
|
"CometAPI",
|
|
|
|
|
"SILICONFLOW",
|
|
|
|
|
"OpenRouter",
|
|
|
|
|
"StepFun",
|
|
|
|
|
"PPIO",
|
|
|
|
|
"PerfXCloud",
|
|
|
|
|
"Upstage",
|
|
|
|
|
"NovitaAI",
|
|
|
|
|
"01.AI",
|
|
|
|
|
"GiteeAI",
|
|
|
|
|
"302.AI",
|
|
|
|
|
]
|
|
|
|
|
|
2025-08-12 10:59:20 +08:00
|
|
|
def __init__(self, key, model_name, base_url=None, **kwargs):
|
|
|
|
|
self.timeout = int(os.environ.get("LM_TIMEOUT_SECONDS", 600))
|
|
|
|
|
self.provider = kwargs.get("provider", "")
|
|
|
|
|
self.prefix = LITELLM_PROVIDER_PREFIX.get(self.provider, "")
|
|
|
|
|
self.model_name = f"{self.prefix}{model_name}"
|
|
|
|
|
self.api_key = key
|
2025-09-18 17:16:59 +08:00
|
|
|
self.base_url = (base_url or FACTORY_DEFAULT_BASE_URL.get(self.provider, "")).rstrip("/")
|
2025-08-12 10:59:20 +08:00
|
|
|
# Configure retry parameters
|
|
|
|
|
self.max_retries = kwargs.get("max_retries", int(os.environ.get("LLM_MAX_RETRIES", 5)))
|
|
|
|
|
self.base_delay = kwargs.get("retry_interval", float(os.environ.get("LLM_BASE_DELAY", 2.0)))
|
|
|
|
|
self.max_rounds = kwargs.get("max_rounds", 5)
|
|
|
|
|
self.is_tools = False
|
|
|
|
|
self.tools = []
|
|
|
|
|
self.toolcall_sessions = {}
|
|
|
|
|
|
|
|
|
|
# Factory specific fields
|
|
|
|
|
if self.provider == SupportedLiteLLMProvider.Bedrock:
|
|
|
|
|
self.bedrock_ak = json.loads(key).get("bedrock_ak", "")
|
|
|
|
|
self.bedrock_sk = json.loads(key).get("bedrock_sk", "")
|
|
|
|
|
self.bedrock_region = json.loads(key).get("bedrock_region", "")
|
2025-10-16 09:39:59 +08:00
|
|
|
elif self.provider == SupportedLiteLLMProvider.OpenRouter:
|
|
|
|
|
self.api_key = json.loads(key).get("api_key", "")
|
|
|
|
|
self.provider_order = json.loads(key).get("provider_order", "")
|
2025-08-12 10:59:20 +08:00
|
|
|
|
|
|
|
|
def _get_delay(self):
|
|
|
|
|
"""Calculate retry delay time"""
|
|
|
|
|
return self.base_delay * random.uniform(10, 150)
|
|
|
|
|
|
|
|
|
|
def _classify_error(self, error):
|
|
|
|
|
"""Classify error based on error message content"""
|
|
|
|
|
error_str = str(error).lower()
|
|
|
|
|
|
|
|
|
|
keywords_mapping = [
|
|
|
|
|
(["quota", "capacity", "credit", "billing", "balance", "欠费"], LLMErrorCode.ERROR_QUOTA),
|
|
|
|
|
(["rate limit", "429", "tpm limit", "too many requests", "requests per minute"], LLMErrorCode.ERROR_RATE_LIMIT),
|
|
|
|
|
(["auth", "key", "apikey", "401", "forbidden", "permission"], LLMErrorCode.ERROR_AUTHENTICATION),
|
|
|
|
|
(["invalid", "bad request", "400", "format", "malformed", "parameter"], LLMErrorCode.ERROR_INVALID_REQUEST),
|
|
|
|
|
(["server", "503", "502", "504", "500", "unavailable"], LLMErrorCode.ERROR_SERVER),
|
|
|
|
|
(["timeout", "timed out"], LLMErrorCode.ERROR_TIMEOUT),
|
|
|
|
|
(["connect", "network", "unreachable", "dns"], LLMErrorCode.ERROR_CONNECTION),
|
|
|
|
|
(["filter", "content", "policy", "blocked", "safety", "inappropriate"], LLMErrorCode.ERROR_CONTENT_FILTER),
|
|
|
|
|
(["model", "not found", "does not exist", "not available"], LLMErrorCode.ERROR_MODEL),
|
|
|
|
|
(["max rounds"], LLMErrorCode.ERROR_MODEL),
|
|
|
|
|
]
|
|
|
|
|
for words, code in keywords_mapping:
|
|
|
|
|
if re.search("({})".format("|".join(words)), error_str):
|
|
|
|
|
return code
|
|
|
|
|
|
|
|
|
|
return LLMErrorCode.ERROR_GENERIC
|
|
|
|
|
|
|
|
|
|
def _clean_conf(self, gen_conf):
|
|
|
|
|
if "max_tokens" in gen_conf:
|
|
|
|
|
del gen_conf["max_tokens"]
|
|
|
|
|
return gen_conf
|
|
|
|
|
|
|
|
|
|
def _chat(self, history, gen_conf, **kwargs):
|
|
|
|
|
logging.info("[HISTORY]" + json.dumps(history, ensure_ascii=False, indent=2))
|
|
|
|
|
if self.model_name.lower().find("qwen3") >= 0:
|
|
|
|
|
kwargs["extra_body"] = {"enable_thinking": False}
|
|
|
|
|
|
2025-08-12 15:54:30 +08:00
|
|
|
completion_args = self._construct_completion_args(history=history, stream=False, tools=False, **gen_conf)
|
2025-08-12 10:59:20 +08:00
|
|
|
response = litellm.completion(
|
|
|
|
|
**completion_args,
|
|
|
|
|
drop_params=True,
|
|
|
|
|
timeout=self.timeout,
|
|
|
|
|
)
|
2025-09-23 16:06:12 +08:00
|
|
|
# response = self.client.chat.completions.create(model=self.model_name, messages=history, **gen_conf, **kwargs)
|
2025-08-12 10:59:20 +08:00
|
|
|
if any([not response.choices, not response.choices[0].message, not response.choices[0].message.content]):
|
|
|
|
|
return "", 0
|
|
|
|
|
ans = response.choices[0].message.content.strip()
|
|
|
|
|
if response.choices[0].finish_reason == "length":
|
|
|
|
|
ans = self._length_stop(ans)
|
|
|
|
|
|
2025-10-22 12:25:31 +08:00
|
|
|
return ans, total_token_count_from_response(response)
|
2025-08-12 10:59:20 +08:00
|
|
|
|
|
|
|
|
def _chat_streamly(self, history, gen_conf, **kwargs):
|
|
|
|
|
logging.info("[HISTORY STREAMLY]" + json.dumps(history, ensure_ascii=False, indent=4))
|
|
|
|
|
reasoning_start = False
|
|
|
|
|
|
2025-08-12 15:54:30 +08:00
|
|
|
completion_args = self._construct_completion_args(history=history, stream=True, tools=False, **gen_conf)
|
2025-08-12 10:59:20 +08:00
|
|
|
stop = kwargs.get("stop")
|
|
|
|
|
if stop:
|
|
|
|
|
completion_args["stop"] = stop
|
|
|
|
|
response = litellm.completion(
|
|
|
|
|
**completion_args,
|
|
|
|
|
drop_params=True,
|
|
|
|
|
timeout=self.timeout,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
for resp in response:
|
|
|
|
|
if not hasattr(resp, "choices") or not resp.choices:
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
delta = resp.choices[0].delta
|
|
|
|
|
if not hasattr(delta, "content") or delta.content is None:
|
|
|
|
|
delta.content = ""
|
|
|
|
|
|
|
|
|
|
if kwargs.get("with_reasoning", True) and hasattr(delta, "reasoning_content") and delta.reasoning_content:
|
|
|
|
|
ans = ""
|
|
|
|
|
if not reasoning_start:
|
|
|
|
|
reasoning_start = True
|
|
|
|
|
ans = "<think>"
|
|
|
|
|
ans += delta.reasoning_content + "</think>"
|
|
|
|
|
else:
|
|
|
|
|
reasoning_start = False
|
|
|
|
|
ans = delta.content
|
|
|
|
|
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-08-12 10:59:20 +08:00
|
|
|
if not tol:
|
|
|
|
|
tol = num_tokens_from_string(delta.content)
|
|
|
|
|
|
|
|
|
|
finish_reason = resp.choices[0].finish_reason if hasattr(resp.choices[0], "finish_reason") else ""
|
|
|
|
|
if finish_reason == "length":
|
|
|
|
|
if is_chinese(ans):
|
|
|
|
|
ans += LENGTH_NOTIFICATION_CN
|
|
|
|
|
else:
|
|
|
|
|
ans += LENGTH_NOTIFICATION_EN
|
|
|
|
|
|
|
|
|
|
yield ans, tol
|
|
|
|
|
|
|
|
|
|
def _length_stop(self, ans):
|
|
|
|
|
if is_chinese([ans]):
|
|
|
|
|
return ans + LENGTH_NOTIFICATION_CN
|
|
|
|
|
return ans + LENGTH_NOTIFICATION_EN
|
|
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
@property
|
|
|
|
|
def _retryable_errors(self) -> set[str]:
|
|
|
|
|
return {
|
|
|
|
|
LLMErrorCode.ERROR_RATE_LIMIT,
|
|
|
|
|
LLMErrorCode.ERROR_SERVER,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
def _should_retry(self, error_code: str) -> bool:
|
|
|
|
|
return error_code in self._retryable_errors
|
|
|
|
|
|
|
|
|
|
def _exceptions(self, e, attempt) -> str | None:
|
2025-08-12 10:59:20 +08:00
|
|
|
logging.exception("OpenAI chat_with_tools")
|
|
|
|
|
# Classify the error
|
|
|
|
|
error_code = self._classify_error(e)
|
|
|
|
|
if attempt == self.max_retries:
|
|
|
|
|
error_code = LLMErrorCode.ERROR_MAX_RETRIES
|
|
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
if self._should_retry(error_code):
|
|
|
|
|
delay = self._get_delay()
|
|
|
|
|
logging.warning(f"Error: {error_code}. Retrying in {delay:.2f} seconds... (Attempt {attempt + 1}/{self.max_retries})")
|
|
|
|
|
time.sleep(delay)
|
|
|
|
|
return None
|
2025-08-12 10:59:20 +08:00
|
|
|
|
2025-09-23 06:19:28 +02:00
|
|
|
return f"{ERROR_PREFIX}: {error_code} - {str(e)}"
|
2025-08-12 10:59:20 +08:00
|
|
|
|
|
|
|
|
def _verbose_tool_use(self, name, args, res):
|
|
|
|
|
return "<tool_call>" + json.dumps({"name": name, "args": args, "result": res}, ensure_ascii=False, indent=2) + "</tool_call>"
|
|
|
|
|
|
|
|
|
|
def _append_history(self, hist, tool_call, tool_res):
|
|
|
|
|
hist.append(
|
|
|
|
|
{
|
|
|
|
|
"role": "assistant",
|
|
|
|
|
"tool_calls": [
|
|
|
|
|
{
|
|
|
|
|
"index": tool_call.index,
|
|
|
|
|
"id": tool_call.id,
|
|
|
|
|
"function": {
|
|
|
|
|
"name": tool_call.function.name,
|
|
|
|
|
"arguments": tool_call.function.arguments,
|
|
|
|
|
},
|
|
|
|
|
"type": "function",
|
|
|
|
|
},
|
|
|
|
|
],
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
try:
|
|
|
|
|
if isinstance(tool_res, dict):
|
|
|
|
|
tool_res = json.dumps(tool_res, ensure_ascii=False)
|
|
|
|
|
finally:
|
|
|
|
|
hist.append({"role": "tool", "tool_call_id": tool_call.id, "content": str(tool_res)})
|
|
|
|
|
return hist
|
|
|
|
|
|
|
|
|
|
def bind_tools(self, toolcall_session, tools):
|
|
|
|
|
if not (toolcall_session and tools):
|
|
|
|
|
return
|
|
|
|
|
self.is_tools = True
|
|
|
|
|
self.toolcall_session = toolcall_session
|
|
|
|
|
self.tools = tools
|
|
|
|
|
|
2025-08-12 15:54:30 +08:00
|
|
|
def _construct_completion_args(self, history, stream: bool, tools: bool, **kwargs):
|
2025-08-12 10:59:20 +08:00
|
|
|
completion_args = {
|
|
|
|
|
"model": self.model_name,
|
|
|
|
|
"messages": history,
|
|
|
|
|
"api_key": self.api_key,
|
2025-09-03 13:31:43 +08:00
|
|
|
"num_retries": self.max_retries,
|
2025-08-12 10:59:20 +08:00
|
|
|
**kwargs,
|
|
|
|
|
}
|
2025-08-12 15:54:30 +08:00
|
|
|
if stream:
|
|
|
|
|
completion_args.update(
|
|
|
|
|
{
|
|
|
|
|
"stream": stream,
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
if tools and self.tools:
|
|
|
|
|
completion_args.update(
|
|
|
|
|
{
|
|
|
|
|
"tools": self.tools,
|
|
|
|
|
"tool_choice": "auto",
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
if self.provider in FACTORY_DEFAULT_BASE_URL:
|
2025-08-12 10:59:20 +08:00
|
|
|
completion_args.update({"api_base": self.base_url})
|
|
|
|
|
elif self.provider == SupportedLiteLLMProvider.Bedrock:
|
|
|
|
|
completion_args.pop("api_key", None)
|
|
|
|
|
completion_args.pop("api_base", None)
|
|
|
|
|
completion_args.update(
|
|
|
|
|
{
|
|
|
|
|
"aws_access_key_id": self.bedrock_ak,
|
|
|
|
|
"aws_secret_access_key": self.bedrock_sk,
|
|
|
|
|
"aws_region_name": self.bedrock_region,
|
|
|
|
|
}
|
|
|
|
|
)
|
2025-10-16 09:39:59 +08:00
|
|
|
|
|
|
|
|
if self.provider == SupportedLiteLLMProvider.OpenRouter:
|
|
|
|
|
if self.provider_order:
|
|
|
|
|
def _to_order_list(x):
|
|
|
|
|
if x is None:
|
|
|
|
|
return []
|
|
|
|
|
if isinstance(x, str):
|
|
|
|
|
return [s.strip() for s in x.split(",") if s.strip()]
|
|
|
|
|
if isinstance(x, (list, tuple)):
|
|
|
|
|
return [str(s).strip() for s in x if str(s).strip()]
|
|
|
|
|
return []
|
|
|
|
|
extra_body = {}
|
|
|
|
|
provider_cfg = {}
|
|
|
|
|
provider_order = _to_order_list(self.provider_order)
|
|
|
|
|
provider_cfg["order"] = provider_order
|
|
|
|
|
provider_cfg["allow_fallbacks"] = False
|
|
|
|
|
extra_body["provider"] = provider_cfg
|
|
|
|
|
completion_args.update({"extra_body": extra_body})
|
2025-08-12 10:59:20 +08:00
|
|
|
return completion_args
|
|
|
|
|
|
|
|
|
|
def chat_with_tools(self, system: str, history: list, gen_conf: dict = {}):
|
|
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2025-08-12 10:59:20 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
|
|
|
|
|
|
|
|
|
ans = ""
|
|
|
|
|
tk_count = 0
|
|
|
|
|
hist = deepcopy(history)
|
|
|
|
|
|
|
|
|
|
# Implement exponential backoff retry strategy
|
|
|
|
|
for attempt in range(self.max_retries + 1):
|
|
|
|
|
history = deepcopy(hist) # deepcopy is required here
|
|
|
|
|
try:
|
|
|
|
|
for _ in range(self.max_rounds + 1):
|
|
|
|
|
logging.info(f"{self.tools=}")
|
|
|
|
|
|
2025-08-12 15:54:30 +08:00
|
|
|
completion_args = self._construct_completion_args(history=history, stream=False, tools=True, **gen_conf)
|
2025-08-12 10:59:20 +08:00
|
|
|
response = litellm.completion(
|
|
|
|
|
**completion_args,
|
|
|
|
|
drop_params=True,
|
|
|
|
|
timeout=self.timeout,
|
|
|
|
|
)
|
|
|
|
|
|
2025-10-22 12:25:31 +08:00
|
|
|
tk_count += total_token_count_from_response(response)
|
2025-08-12 10:59:20 +08:00
|
|
|
|
|
|
|
|
if not hasattr(response, "choices") or not response.choices or not response.choices[0].message:
|
|
|
|
|
raise Exception(f"500 response structure error. Response: {response}")
|
|
|
|
|
|
|
|
|
|
message = response.choices[0].message
|
|
|
|
|
|
|
|
|
|
if not hasattr(message, "tool_calls") or not message.tool_calls:
|
|
|
|
|
if hasattr(message, "reasoning_content") and message.reasoning_content:
|
|
|
|
|
ans += f"<think>{message.reasoning_content}</think>"
|
|
|
|
|
ans += message.content or ""
|
|
|
|
|
if response.choices[0].finish_reason == "length":
|
|
|
|
|
ans = self._length_stop(ans)
|
|
|
|
|
return ans, tk_count
|
|
|
|
|
|
|
|
|
|
for tool_call in message.tool_calls:
|
|
|
|
|
logging.info(f"Response {tool_call=}")
|
|
|
|
|
name = tool_call.function.name
|
|
|
|
|
try:
|
|
|
|
|
args = json_repair.loads(tool_call.function.arguments)
|
|
|
|
|
tool_response = self.toolcall_session.tool_call(name, args)
|
|
|
|
|
history = self._append_history(history, tool_call, tool_response)
|
|
|
|
|
ans += self._verbose_tool_use(name, args, tool_response)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logging.exception(msg=f"Wrong JSON argument format in LLM tool call response: {tool_call}")
|
|
|
|
|
history.append({"role": "tool", "tool_call_id": tool_call.id, "content": f"Tool call error: \n{tool_call}\nException:\n" + str(e)})
|
|
|
|
|
ans += self._verbose_tool_use(name, {}, str(e))
|
|
|
|
|
|
|
|
|
|
logging.warning(f"Exceed max rounds: {self.max_rounds}")
|
|
|
|
|
history.append({"role": "user", "content": f"Exceed max rounds: {self.max_rounds}"})
|
|
|
|
|
|
|
|
|
|
response, token_count = self._chat(history, gen_conf)
|
|
|
|
|
ans += response
|
|
|
|
|
tk_count += token_count
|
|
|
|
|
return ans, tk_count
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
e = self._exceptions(e, attempt)
|
|
|
|
|
if e:
|
|
|
|
|
return e, tk_count
|
|
|
|
|
|
|
|
|
|
assert False, "Shouldn't be here."
|
|
|
|
|
|
|
|
|
|
def chat(self, system, history, gen_conf={}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2025-08-12 10:59:20 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
|
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
|
|
|
|
|
|
|
|
|
# Implement exponential backoff retry strategy
|
|
|
|
|
for attempt in range(self.max_retries + 1):
|
|
|
|
|
try:
|
|
|
|
|
response = self._chat(history, gen_conf, **kwargs)
|
|
|
|
|
return response
|
|
|
|
|
except Exception as e:
|
|
|
|
|
e = self._exceptions(e, attempt)
|
|
|
|
|
if e:
|
|
|
|
|
return e, 0
|
|
|
|
|
assert False, "Shouldn't be here."
|
|
|
|
|
|
|
|
|
|
def _wrap_toolcall_message(self, stream):
|
|
|
|
|
final_tool_calls = {}
|
|
|
|
|
|
|
|
|
|
for chunk in stream:
|
|
|
|
|
for tool_call in chunk.choices[0].delta.tool_calls or []:
|
|
|
|
|
index = tool_call.index
|
|
|
|
|
|
|
|
|
|
if index not in final_tool_calls:
|
|
|
|
|
final_tool_calls[index] = tool_call
|
|
|
|
|
|
|
|
|
|
final_tool_calls[index].function.arguments += tool_call.function.arguments
|
|
|
|
|
|
|
|
|
|
return final_tool_calls
|
|
|
|
|
|
|
|
|
|
def chat_streamly_with_tools(self, system: str, history: list, gen_conf: dict = {}):
|
|
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
|
|
|
|
tools = self.tools
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2025-08-12 10:59:20 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
|
|
|
|
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
hist = deepcopy(history)
|
|
|
|
|
|
|
|
|
|
# Implement exponential backoff retry strategy
|
|
|
|
|
for attempt in range(self.max_retries + 1):
|
|
|
|
|
history = deepcopy(hist) # deepcopy is required here
|
|
|
|
|
try:
|
|
|
|
|
for _ in range(self.max_rounds + 1):
|
|
|
|
|
reasoning_start = False
|
|
|
|
|
logging.info(f"{tools=}")
|
|
|
|
|
|
2025-08-12 15:54:30 +08:00
|
|
|
completion_args = self._construct_completion_args(history=history, stream=True, tools=True, **gen_conf)
|
2025-08-12 10:59:20 +08:00
|
|
|
response = litellm.completion(
|
|
|
|
|
**completion_args,
|
|
|
|
|
drop_params=True,
|
|
|
|
|
timeout=self.timeout,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
final_tool_calls = {}
|
|
|
|
|
answer = ""
|
|
|
|
|
|
|
|
|
|
for resp in response:
|
|
|
|
|
if not hasattr(resp, "choices") or not resp.choices:
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
delta = resp.choices[0].delta
|
|
|
|
|
|
|
|
|
|
if hasattr(delta, "tool_calls") and delta.tool_calls:
|
|
|
|
|
for tool_call in delta.tool_calls:
|
|
|
|
|
index = tool_call.index
|
|
|
|
|
if index not in final_tool_calls:
|
|
|
|
|
if not tool_call.function.arguments:
|
|
|
|
|
tool_call.function.arguments = ""
|
|
|
|
|
final_tool_calls[index] = tool_call
|
|
|
|
|
else:
|
|
|
|
|
final_tool_calls[index].function.arguments += tool_call.function.arguments or ""
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
if not hasattr(delta, "content") or delta.content is None:
|
|
|
|
|
delta.content = ""
|
|
|
|
|
|
|
|
|
|
if hasattr(delta, "reasoning_content") and delta.reasoning_content:
|
|
|
|
|
ans = ""
|
|
|
|
|
if not reasoning_start:
|
|
|
|
|
reasoning_start = True
|
|
|
|
|
ans = "<think>"
|
|
|
|
|
ans += delta.reasoning_content + "</think>"
|
|
|
|
|
yield ans
|
|
|
|
|
else:
|
|
|
|
|
reasoning_start = False
|
|
|
|
|
answer += delta.content
|
|
|
|
|
yield delta.content
|
|
|
|
|
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-08-12 10:59:20 +08:00
|
|
|
if not tol:
|
|
|
|
|
total_tokens += num_tokens_from_string(delta.content)
|
|
|
|
|
else:
|
|
|
|
|
total_tokens += tol
|
|
|
|
|
|
|
|
|
|
finish_reason = getattr(resp.choices[0], "finish_reason", "")
|
|
|
|
|
if finish_reason == "length":
|
|
|
|
|
yield self._length_stop("")
|
|
|
|
|
|
|
|
|
|
if answer:
|
|
|
|
|
yield total_tokens
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
for tool_call in final_tool_calls.values():
|
|
|
|
|
name = tool_call.function.name
|
|
|
|
|
try:
|
|
|
|
|
args = json_repair.loads(tool_call.function.arguments)
|
|
|
|
|
yield self._verbose_tool_use(name, args, "Begin to call...")
|
|
|
|
|
tool_response = self.toolcall_session.tool_call(name, args)
|
|
|
|
|
history = self._append_history(history, tool_call, tool_response)
|
|
|
|
|
yield self._verbose_tool_use(name, args, tool_response)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
logging.exception(msg=f"Wrong JSON argument format in LLM tool call response: {tool_call}")
|
|
|
|
|
history.append(
|
|
|
|
|
{
|
|
|
|
|
"role": "tool",
|
|
|
|
|
"tool_call_id": tool_call.id,
|
|
|
|
|
"content": f"Tool call error: \n{tool_call}\nException:\n{str(e)}",
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
yield self._verbose_tool_use(name, {}, str(e))
|
|
|
|
|
|
|
|
|
|
logging.warning(f"Exceed max rounds: {self.max_rounds}")
|
|
|
|
|
history.append({"role": "user", "content": f"Exceed max rounds: {self.max_rounds}"})
|
|
|
|
|
|
2025-08-12 15:54:30 +08:00
|
|
|
completion_args = self._construct_completion_args(history=history, stream=True, tools=True, **gen_conf)
|
2025-08-12 10:59:20 +08:00
|
|
|
response = litellm.completion(
|
|
|
|
|
**completion_args,
|
|
|
|
|
drop_params=True,
|
|
|
|
|
timeout=self.timeout,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
for resp in response:
|
|
|
|
|
if not hasattr(resp, "choices") or not resp.choices:
|
|
|
|
|
continue
|
|
|
|
|
delta = resp.choices[0].delta
|
|
|
|
|
if not hasattr(delta, "content") or delta.content is None:
|
|
|
|
|
continue
|
2025-10-22 12:25:31 +08:00
|
|
|
tol = total_token_count_from_response(resp)
|
2025-08-12 10:59:20 +08:00
|
|
|
if not tol:
|
|
|
|
|
total_tokens += num_tokens_from_string(delta.content)
|
|
|
|
|
else:
|
|
|
|
|
total_tokens += tol
|
|
|
|
|
yield delta.content
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
e = self._exceptions(e, attempt)
|
|
|
|
|
if e:
|
|
|
|
|
yield e
|
|
|
|
|
yield total_tokens
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
assert False, "Shouldn't be here."
|
|
|
|
|
|
|
|
|
|
def chat_streamly(self, system, history, gen_conf: dict = {}, **kwargs):
|
Refa: revise the implementation of LightRAG and enable response caching (#9828)
### What problem does this PR solve?
This revision performed a comprehensive check on LightRAG to ensure the
correctness of its implementation. It **did not involve** Entity
Resolution and Community Reports Generation. There is an example using
default entity types and the General chunking method, which shows good
results in both time and effectiveness. Moreover, response caching is
enabled for resuming failed tasks.
[The-Necklace.pdf](https://github.com/user-attachments/files/22042432/The-Necklace.pdf)
After:

```bash
Begin at:
Fri, 29 Aug 2025 16:48:03 GMT
Duration:
222.31 s
Progress:
16:48:04 Task has been received.
16:48:06 Page(1~7): Start to parse.
16:48:06 Page(1~7): OCR started
16:48:08 Page(1~7): OCR finished (1.89s)
16:48:11 Page(1~7): Layout analysis (3.72s)
16:48:11 Page(1~7): Table analysis (0.00s)
16:48:11 Page(1~7): Text merged (0.00s)
16:48:11 Page(1~7): Finish parsing.
16:48:12 Page(1~7): Generate 7 chunks
16:48:12 Page(1~7): Embedding chunks (0.29s)
16:48:12 Page(1~7): Indexing done (0.04s). Task done (7.84s)
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ...
16:48:17 Start processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin...
16:49:30 Completed processing for f421fb06849e11f0bdd32724b93a52b2: She had no dresses, no je... after 1 gleanings, 21985 tokens.
16:49:30 Entities extraction of chunk 3 1/7 done, 12 nodes, 13 edges, 21985 tokens.
16:49:40 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Finally, she replied, hes... after 1 gleanings, 22584 tokens.
16:49:40 Entities extraction of chunk 5 2/7 done, 19 nodes, 19 edges, 22584 tokens.
16:50:02 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Then she asked, hesitatin... after 1 gleanings, 24610 tokens.
16:50:02 Entities extraction of chunk 0 3/7 done, 16 nodes, 28 edges, 24610 tokens.
16:50:03 Completed processing for f421fb06849e11f0bdd32724b93a52b2: And this life lasted ten ... after 1 gleanings, 24031 tokens.
16:50:04 Entities extraction of chunk 1 4/7 done, 24 nodes, 22 edges, 24031 tokens.
16:50:14 Completed processing for f421fb06849e11f0bdd32724b93a52b2: So they begged the jewell... after 1 gleanings, 24635 tokens.
16:50:14 Entities extraction of chunk 6 5/7 done, 27 nodes, 26 edges, 24635 tokens.
16:50:29 Completed processing for f421fb06849e11f0bdd32724b93a52b2: Her husband, already half... after 1 gleanings, 25758 tokens.
16:50:29 Entities extraction of chunk 2 6/7 done, 25 nodes, 35 edges, 25758 tokens.
16:51:35 Completed processing for f421fb06849e11f0bdd32724b93a52b2: The Necklace By Guy de Ma... after 1 gleanings, 27491 tokens.
16:51:35 Entities extraction of chunk 4 7/7 done, 39 nodes, 37 edges, 27491 tokens.
16:51:35 Entities and relationships extraction done, 147 nodes, 177 edges, 171094 tokens, 198.58s.
16:51:35 Entities merging done, 0.01s.
16:51:35 Relationships merging done, 0.01s.
16:51:35 ignored 7 relations due to missing entities.
16:51:35 generated subgraph for doc f421fb06849e11f0bdd32724b93a52b2 in 198.68 seconds.
16:51:35 run_graphrag f421fb06849e11f0bdd32724b93a52b2 graphrag_task_lock acquired
16:51:35 set_graph removed 0 nodes and 0 edges from index in 0.00s.
16:51:35 Get embedding of nodes: 9/147
16:51:35 Get embedding of nodes: 109/147
16:51:37 Get embedding of edges: 9/170
16:51:37 Get embedding of edges: 109/170
16:51:40 set_graph converted graph change to 319 chunks in 4.21s.
16:51:40 Insert chunks: 4/319
16:51:40 Insert chunks: 104/319
16:51:40 Insert chunks: 204/319
16:51:40 Insert chunks: 304/319
16:51:40 set_graph added/updated 147 nodes and 170 edges from index in 0.53s.
16:51:40 merging subgraph for doc f421fb06849e11f0bdd32724b93a52b2 into the global graph done in 4.79 seconds.
16:51:40 Knowledge Graph done (204.29s)
```
Before:

```bash
Begin at:
Fri, 29 Aug 2025 17:00:47 GMT
processDuration:
173.38 s
Progress:
17:00:49 Task has been received.
17:00:51 Page(1~7): Start to parse.
17:00:51 Page(1~7): OCR started
17:00:53 Page(1~7): OCR finished (1.82s)
17:00:57 Page(1~7): Layout analysis (3.64s)
17:00:57 Page(1~7): Table analysis (0.00s)
17:00:57 Page(1~7): Text merged (0.00s)
17:00:57 Page(1~7): Finish parsing.
17:00:57 Page(1~7): Generate 7 chunks
17:00:57 Page(1~7): Embedding chunks (0.31s)
17:00:57 Page(1~7): Indexing done (0.03s). Task done (7.88s)
17:00:57 created task graphrag
17:01:00 Task has been received.
17:02:17 Entities extraction of chunk 1 1/7 done, 9 nodes, 9 edges, 10654 tokens.
17:02:31 Entities extraction of chunk 2 2/7 done, 12 nodes, 13 edges, 11066 tokens.
17:02:33 Entities extraction of chunk 4 3/7 done, 9 nodes, 10 edges, 10433 tokens.
17:02:42 Entities extraction of chunk 5 4/7 done, 11 nodes, 14 edges, 11290 tokens.
17:02:52 Entities extraction of chunk 6 5/7 done, 13 nodes, 15 edges, 11039 tokens.
17:02:55 Entities extraction of chunk 3 6/7 done, 14 nodes, 13 edges, 11466 tokens.
17:03:32 Entities extraction of chunk 0 7/7 done, 19 nodes, 18 edges, 13107 tokens.
17:03:32 Entities and relationships extraction done, 71 nodes, 89 edges, 79055 tokens, 149.66s.
17:03:32 Entities merging done, 0.01s.
17:03:32 Relationships merging done, 0.01s.
17:03:32 ignored 1 relations due to missing entities.
17:03:32 generated subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 in 149.69 seconds.
17:03:32 run_graphrag b1d9d3b6848711f0aacd7ddc0714c4d3 graphrag_task_lock acquired
17:03:32 set_graph removed 0 nodes and 0 edges from index in 0.00s.
17:03:32 Get embedding of nodes: 9/71
17:03:33 Get embedding of edges: 9/88
17:03:34 set_graph converted graph change to 161 chunks in 2.27s.
17:03:34 Insert chunks: 4/161
17:03:34 Insert chunks: 104/161
17:03:34 set_graph added/updated 71 nodes and 88 edges from index in 0.28s.
17:03:34 merging subgraph for doc b1d9d3b6848711f0aacd7ddc0714c4d3 into the global graph done in 2.60 seconds.
17:03:34 Knowledge Graph done (153.18s)
```
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement
2025-08-29 17:58:36 +08:00
|
|
|
if system and history and history[0].get("role") != "system":
|
2025-08-12 10:59:20 +08:00
|
|
|
history.insert(0, {"role": "system", "content": system})
|
|
|
|
|
gen_conf = self._clean_conf(gen_conf)
|
|
|
|
|
ans = ""
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
try:
|
|
|
|
|
for delta_ans, tol in self._chat_streamly(history, gen_conf, **kwargs):
|
|
|
|
|
yield delta_ans
|
|
|
|
|
total_tokens += tol
|
|
|
|
|
except openai.APIError as e:
|
|
|
|
|
yield ans + "\n**ERROR**: " + str(e)
|
|
|
|
|
|
|
|
|
|
yield total_tokens
|
|
|
|
|
|
|
|
|
|
def _calculate_dynamic_ctx(self, history):
|
|
|
|
|
"""Calculate dynamic context window size"""
|
|
|
|
|
|
|
|
|
|
def count_tokens(text):
|
|
|
|
|
"""Calculate token count for text"""
|
|
|
|
|
# Simple calculation: 1 token per ASCII character
|
|
|
|
|
# 2 tokens for non-ASCII characters (Chinese, Japanese, Korean, etc.)
|
|
|
|
|
total = 0
|
|
|
|
|
for char in text:
|
|
|
|
|
if ord(char) < 128: # ASCII characters
|
|
|
|
|
total += 1
|
|
|
|
|
else: # Non-ASCII characters (Chinese, Japanese, Korean, etc.)
|
|
|
|
|
total += 2
|
|
|
|
|
return total
|
|
|
|
|
|
|
|
|
|
# Calculate total tokens for all messages
|
|
|
|
|
total_tokens = 0
|
|
|
|
|
for message in history:
|
|
|
|
|
content = message.get("content", "")
|
|
|
|
|
# Calculate content tokens
|
|
|
|
|
content_tokens = count_tokens(content)
|
|
|
|
|
# Add role marker token overhead
|
|
|
|
|
role_tokens = 4
|
|
|
|
|
total_tokens += content_tokens + role_tokens
|
|
|
|
|
|
|
|
|
|
# Apply 1.2x buffer ratio
|
|
|
|
|
total_tokens_with_buffer = int(total_tokens * 1.2)
|
|
|
|
|
|
|
|
|
|
if total_tokens_with_buffer <= 8192:
|
|
|
|
|
ctx_size = 8192
|
|
|
|
|
else:
|
|
|
|
|
ctx_multiplier = (total_tokens_with_buffer // 8192) + 1
|
|
|
|
|
ctx_size = ctx_multiplier * 8192
|
|
|
|
|
|
|
|
|
|
return ctx_size
|