274 lines
8.8 KiB
Python
274 lines
8.8 KiB
Python
import json
|
||
import re
|
||
from collections import OrderedDict
|
||
from collections.abc import Generator
|
||
from typing import Any
|
||
|
||
from dify_plugin import Tool
|
||
from dify_plugin.entities.model.llm import LLMModelConfig
|
||
from dify_plugin.entities.model.message import SystemPromptMessage, UserPromptMessage
|
||
from dify_plugin.entities.tool import ToolInvokeMessage
|
||
|
||
_SYSTEM_PROMPT = """You parse PDF table-of-contents text.
|
||
Return only valid JSON object, no markdown fences, no explanation.
|
||
Output schema:
|
||
{
|
||
"Chapter Name": {"start": 1, "end": 5},
|
||
"Another": {"start": 6, "end": 20}
|
||
}
|
||
Rules:
|
||
- start/end are integer printed page numbers from TOC.
|
||
- If end is unknown, use same value as start.
|
||
- Keep chapter names exactly as in TOC text.
|
||
"""
|
||
|
||
|
||
class PdfTocTool(Tool):
|
||
def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
|
||
toc_start = self._to_int(tool_parameters.get("toc_start"), None)
|
||
toc_end = self._to_int(tool_parameters.get("toc_end"), None)
|
||
toc_pages = (tool_parameters.get("toc_pages") or "").strip()
|
||
model_config = tool_parameters.get("model")
|
||
|
||
if toc_start is None or toc_end is None:
|
||
yield self.create_text_message("Error: toc_start and toc_end are required")
|
||
return
|
||
|
||
if not toc_pages:
|
||
yield self.create_text_message("Error: toc_pages text is empty")
|
||
return
|
||
|
||
cleaned = self._strip_index_lists(toc_pages)
|
||
|
||
# 1) deterministic parser first
|
||
catalog = self._parse_toc_lines(cleaned)
|
||
|
||
# 2) optional LLM fallback/enhance only when deterministic parser gives no result
|
||
llm_raw_output = ""
|
||
llm_error = None
|
||
if not catalog and model_config:
|
||
llm_catalog, llm_raw_output, llm_error = self._parse_with_llm(
|
||
toc_start=toc_start,
|
||
toc_end=toc_end,
|
||
toc_pages=cleaned,
|
||
model_config=model_config,
|
||
)
|
||
if llm_catalog:
|
||
catalog = self._normalize_catalog(llm_catalog)
|
||
|
||
result: dict[str, Any] = {
|
||
"toc_start": toc_start,
|
||
"toc_end": toc_end,
|
||
"catalog": catalog,
|
||
"meta": {
|
||
"catalog_size": len(catalog),
|
||
"parser": "rule" if catalog else "none",
|
||
},
|
||
}
|
||
|
||
if llm_raw_output:
|
||
result["meta"]["llm_used"] = True
|
||
if llm_error:
|
||
result["meta"]["llm_error"] = llm_error
|
||
|
||
# always return valid json text payload for downstream json.loads
|
||
yield self.create_text_message(json.dumps(result, ensure_ascii=False))
|
||
yield self.create_json_message(result)
|
||
|
||
def _parse_with_llm(
|
||
self,
|
||
toc_start: int,
|
||
toc_end: int,
|
||
toc_pages: str,
|
||
model_config: dict[str, Any],
|
||
) -> tuple[dict[str, Any] | None, str, str | None]:
|
||
user_content = (
|
||
f"TOC page index range: {toc_start}..{toc_end}\n\n"
|
||
f"TOC raw text:\n{toc_pages}"
|
||
)
|
||
response = self.session.model.llm.invoke(
|
||
model_config=LLMModelConfig(**model_config),
|
||
prompt_messages=[
|
||
SystemPromptMessage(content=_SYSTEM_PROMPT),
|
||
UserPromptMessage(content=user_content),
|
||
],
|
||
stream=False,
|
||
)
|
||
|
||
llm_text = ""
|
||
if hasattr(response, "message") and response.message:
|
||
content = response.message.content
|
||
if isinstance(content, str):
|
||
llm_text = content
|
||
elif isinstance(content, list):
|
||
llm_text = "".join(
|
||
item.data if hasattr(item, "data") else str(item) for item in content
|
||
)
|
||
|
||
parsed = self._extract_json_object(llm_text)
|
||
if parsed is None:
|
||
return None, llm_text, "Failed to parse LLM output as JSON"
|
||
if not isinstance(parsed, dict):
|
||
return None, llm_text, "LLM output JSON is not an object"
|
||
|
||
return parsed, llm_text, None
|
||
|
||
@staticmethod
|
||
def _strip_index_lists(text: str) -> str:
|
||
# Stop before common appendix lists that pollute TOC parsing.
|
||
pattern = re.compile(
|
||
r"^(List\s+of\s+Figures|List\s+of\s+Tables|图目录|表目录)",
|
||
re.IGNORECASE | re.MULTILINE,
|
||
)
|
||
m = pattern.search(text)
|
||
return text[: m.start()].rstrip() if m else text
|
||
|
||
def _parse_toc_lines(self, text: str) -> dict[str, dict[str, int]]:
|
||
"""Parse lines like:
|
||
1.2 Engine Overview ........ 35
|
||
Appendix A 120
|
||
"""
|
||
line_pattern = re.compile(
|
||
r"^\s*(?P<title>.+?)\s*(?:\.{2,}|\s)\s*(?P<page>\d{1,5})\s*$"
|
||
)
|
||
|
||
entries: list[tuple[str, int]] = []
|
||
for raw in text.splitlines():
|
||
line = raw.strip()
|
||
if not line or len(line) < 3:
|
||
continue
|
||
if re.fullmatch(r"\d+", line):
|
||
continue
|
||
|
||
m = line_pattern.match(line)
|
||
if not m:
|
||
continue
|
||
|
||
title = re.sub(r"\s+", " ", m.group("title")).strip("-_:： ")
|
||
page = self._to_int(m.group("page"), None)
|
||
if not title or page is None:
|
||
continue
|
||
|
||
# Skip obvious noise.
|
||
if len(title) <= 1 or title.lower() in {"page", "pages", "目录", "contents"}:
|
||
continue
|
||
|
||
entries.append((title, page))
|
||
|
||
if not entries:
|
||
return {}
|
||
|
||
# Deduplicate keeping earliest appearance.
|
||
dedup: OrderedDict[str, int] = OrderedDict()
|
||
for title, page in entries:
|
||
if title not in dedup:
|
||
dedup[title] = page
|
||
|
||
titles = list(dedup.keys())
|
||
pages = [dedup[t] for t in titles]
|
||
|
||
catalog: dict[str, dict[str, int]] = {}
|
||
for i, title in enumerate(titles):
|
||
start = pages[i]
|
||
if i + 1 < len(pages):
|
||
next_start = pages[i + 1]
|
||
end = max(start, next_start - 1)
|
||
else:
|
||
end = start
|
||
catalog[title] = {"start": int(start), "end": int(end)}
|
||
|
||
return catalog
|
||
|
||
def _normalize_catalog(self, raw: dict[str, Any]) -> dict[str, dict[str, int]]:
|
||
catalog: dict[str, dict[str, int]] = {}
|
||
source = raw.get("catalog") if isinstance(raw.get("catalog"), dict) else raw
|
||
if not isinstance(source, dict):
|
||
return catalog
|
||
|
||
for name, value in source.items():
|
||
if not isinstance(name, str) or not isinstance(value, dict):
|
||
continue
|
||
start = self._to_int(value.get("start"), None)
|
||
end = self._to_int(value.get("end"), start)
|
||
if start is None:
|
||
continue
|
||
if end is None:
|
||
end = start
|
||
catalog[name] = {"start": int(start), "end": int(max(start, end))}
|
||
return catalog
|
||
|
||
@staticmethod
|
||
def _extract_json_object(text: str) -> Any:
|
||
if not text:
|
||
return None
|
||
|
||
candidates: list[str] = []
|
||
|
||
code_blocks = re.findall(r"```(?:json)?\s*([\s\S]*?)\s*```", text, flags=re.IGNORECASE)
|
||
candidates.extend([c.strip() for c in code_blocks if c.strip()])
|
||
|
||
brace_candidate = PdfTocTool._extract_first_brace_object(text)
|
||
if brace_candidate:
|
||
candidates.append(brace_candidate)
|
||
|
||
candidates.append(text.strip())
|
||
|
||
for cand in candidates:
|
||
parsed = PdfTocTool._json_try_parse(cand)
|
||
if parsed is not None:
|
||
return parsed
|
||
return None
|
||
|
||
@staticmethod
|
||
def _extract_first_brace_object(text: str) -> str | None:
|
||
start = text.find("{")
|
||
if start < 0:
|
||
return None
|
||
|
||
depth = 0
|
||
in_str = False
|
||
escape = False
|
||
for i in range(start, len(text)):
|
||
ch = text[i]
|
||
if in_str:
|
||
if escape:
|
||
escape = False
|
||
elif ch == "\\":
|
||
escape = True
|
||
elif ch == '"':
|
||
in_str = False
|
||
continue
|
||
|
||
if ch == '"':
|
||
in_str = True
|
||
elif ch == "{":
|
||
depth += 1
|
||
elif ch == "}":
|
||
depth -= 1
|
||
if depth == 0:
|
||
return text[start : i + 1]
|
||
return None
|
||
|
||
@staticmethod
|
||
def _json_try_parse(text: str) -> Any:
|
||
try:
|
||
return json.loads(text)
|
||
except Exception:
|
||
pass
|
||
|
||
# Minimal repair: remove trailing commas before } or ]
|
||
repaired = re.sub(r",\s*([}\]])", r"\1", text)
|
||
try:
|
||
return json.loads(repaired)
|
||
except Exception:
|
||
return None
|
||
|
||
@staticmethod
|
||
def _to_int(value: Any, default: int | None) -> int | None:
|
||
try:
|
||
if value is None or value == "":
|
||
return default
|
||
return int(value)
|
||
except Exception:
|
||
return default
|