瀏覽代碼

Fix: too long context during KG issue. (#5723)

### What problem does this PR solve?

#5088

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
tags/v0.17.1
Kevin Hu 8 月之前
父節點
當前提交
9fc7174612
沒有連結到貢獻者的電子郵件帳戶。
共有 2 個檔案被更改,包括 4 行新增2 行删除
  1. 1
    1
      api/db/services/document_service.py
  2. 3
    1
      graphrag/general/extractor.py

+ 1
- 1
api/db/services/document_service.py 查看文件

from api.db.db_utils import bulk_insert_into_db from api.db.db_utils import bulk_insert_into_db
from api import settings from api import settings
from api.utils import current_timestamp, get_format_time, get_uuid from api.utils import current_timestamp, get_format_time, get_uuid
from graphrag.general.mind_map_extractor import MindMapExtractor
from rag.settings import SVR_QUEUE_NAME from rag.settings import SVR_QUEUE_NAME
from rag.utils.storage_factory import STORAGE_IMPL from rag.utils.storage_factory import STORAGE_IMPL
from rag.nlp import search, rag_tokenizer from rag.nlp import search, rag_tokenizer
cks = [c for c in docs if c["doc_id"] == doc_id] cks = [c for c in docs if c["doc_id"] == doc_id]


if parser_ids[doc_id] != ParserType.PICTURE.value: if parser_ids[doc_id] != ParserType.PICTURE.value:
from graphrag.general.mind_map_extractor import MindMapExtractor
mindmap = MindMapExtractor(llm_bdl) mindmap = MindMapExtractor(llm_bdl)
try: try:
mind_map = trio.run(mindmap, [c["content_with_weight"] for c in docs if c["doc_id"] == doc_id]) mind_map = trio.run(mindmap, [c["content_with_weight"] for c in docs if c["doc_id"] == doc_id])

+ 3
- 1
graphrag/general/extractor.py 查看文件

from graphrag.utils import get_llm_cache, set_llm_cache, handle_single_entity_extraction, \ from graphrag.utils import get_llm_cache, set_llm_cache, handle_single_entity_extraction, \
handle_single_relationship_extraction, split_string_by_multi_markers, flat_uniq_list, chat_limiter handle_single_relationship_extraction, split_string_by_multi_markers, flat_uniq_list, chat_limiter
from rag.llm.chat_model import Base as CompletionLLM from rag.llm.chat_model import Base as CompletionLLM
from rag.prompts import message_fit_in
from rag.utils import truncate from rag.utils import truncate


GRAPH_FIELD_SEP = "<SEP>" GRAPH_FIELD_SEP = "<SEP>"
response = get_llm_cache(self._llm.llm_name, system, hist, conf) response = get_llm_cache(self._llm.llm_name, system, hist, conf)
if response: if response:
return response return response
response = self._llm.chat(system, hist, conf)
_, system_msg = message_fit_in([{"role": "system", "content": system}], int(self._llm.max_length * 0.97))
response = self._llm.chat(system_msg[0]["content"], hist, conf)
response = re.sub(r"<think>.*</think>", "", response, flags=re.DOTALL) response = re.sub(r"<think>.*</think>", "", response, flags=re.DOTALL)
if response.find("**ERROR**") >= 0: if response.find("**ERROR**") >= 0:
raise Exception(response) raise Exception(response)

Loading…
取消
儲存