ordering, not expanded python-docx cells).
-
- Returns:
- HTML string representation of the table with merged cells
- """
-
- NS = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
-
- def get_cell_vmerge(cell):
- """Get vMerge status: 'restart', 'continue', or None"""
- tc = cell._tc
- tcPr = tc.find(".//w:tcPr", namespaces=NS)
- if tcPr is not None:
- vMerge = tcPr.find(".//w:vMerge", namespaces=NS)
- if vMerge is not None:
- val = vMerge.get(
- "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val"
- )
- return (
- val if val else "continue"
- ) # If no val attribute, it's a continuation
- return None
-
- n_rows = len(table.rows)
- if n_rows == 0:
- return ""
-
- # Build grid: track unique cells and their positions
- # grid[row][col] = (cell_tc_id, cell, is_new_cell)
- # We use id(cell._tc) as unique identifier for cells
-
- grid = []
- for row_idx, row in enumerate(table.rows):
- row_data = []
- prev_tc_id = None
- for cell in row.cells:
- tc_id = id(cell._tc)
- is_new = tc_id != prev_tc_id
- row_data.append((tc_id, cell, is_new))
- prev_tc_id = tc_id
- grid.append(row_data)
-
- # Rows may have different cell counts due to complex merges;
- # use the maximum for grid allocation, per-row length for access.
- n_cols = max(len(r) for r in grid) if grid else 0
-
- # Calculate colspan for each cell (count consecutive cells with same _tc)
- colspan_grid = [[0] * n_cols for _ in range(n_rows)]
-
- for row_idx in range(n_rows):
- row_len = len(grid[row_idx])
- col_idx = 0
- while col_idx < row_len:
- tc_id = grid[row_idx][col_idx][0]
- span = 1
- while (
- col_idx + span < row_len and grid[row_idx][col_idx + span][0] == tc_id
- ):
- span += 1
- colspan_grid[row_idx][col_idx] = span
- col_idx += span
-
- # Calculate rowspan for cells with vMerge='restart'
- rowspan_grid = [[1] * n_cols for _ in range(n_rows)]
-
- for col_idx in range(n_cols):
- row_idx = 0
- while row_idx < n_rows:
- if col_idx >= len(grid[row_idx]):
- row_idx += 1
- continue
- cell = grid[row_idx][col_idx][1]
- vmerge = get_cell_vmerge(cell)
-
- if vmerge == "restart":
- # Count how many 'continue' cells follow
- span = 1
- while row_idx + span < n_rows:
- if col_idx >= len(grid[row_idx + span]):
- break
- next_cell = grid[row_idx + span][col_idx][1]
- next_vmerge = get_cell_vmerge(next_cell)
- if next_vmerge == "continue":
- span += 1
- else:
- break
- rowspan_grid[row_idx][col_idx] = span
- row_idx += span
- elif vmerge == "continue":
- # This cell is part of a vertical merge, mark as 0 (skip)
- rowspan_grid[row_idx][col_idx] = 0
- row_idx += 1
- else:
- row_idx += 1
-
- # Build HTML
- html_parts = [""]
-
- for row_idx in range(n_rows):
- html_parts.append("")
- col_idx = 0
- unique_col_idx = 0 # Tracks unique tc index per row (matches XML order)
-
- while col_idx < len(grid[row_idx]):
- tc_id, cell, is_new = grid[row_idx][col_idx]
-
- # Skip if this cell is a horizontal continuation
- if not is_new:
- col_idx += 1
- continue
-
- # Skip if this cell is a vertical continuation
- rowspan = rowspan_grid[row_idx][col_idx]
- if rowspan == 0:
- unique_col_idx += 1
- col_idx += 1
- continue
-
- colspan = colspan_grid[row_idx][col_idx]
-
- # Build cell content
- if cell.tables:
- # Nested table
- content = "".join(
- table2html(nested_table) for nested_table in cell.tables
- )
- else:
- content = cell.text.strip().replace("\n", "
")
-
- # Append image descriptions if available
- if cell_image_map:
- img_desc = cell_image_map.get((row_idx, unique_col_idx))
- if img_desc:
- content += f"
{img_desc}"
-
- # Build attributes
- attrs = []
- if colspan > 1:
- attrs.append(f'colspan="{colspan}"')
- if rowspan > 1:
- attrs.append(f'rowspan="{rowspan}"')
-
- attr_str = " " + " ".join(attrs) if attrs else ""
- html_parts.append(f"{content}")
-
- unique_col_idx += 1
- col_idx += colspan
-
- html_parts.append("
")
-
- html_parts.append("
")
- return "".join(html_parts)
-
-
-def render_multiindex_thead(columns: pd.MultiIndex, escape: bool = False) -> str:
- """
- Convert MultiIndex columns to HTML thead with colspan/rowspan.
-
- This function generates a proper multi-row structure where:
- - Horizontally adjacent identical values are merged with colspan
- - Vertically repeated values are merged with rowspan
-
- Args:
- columns: pandas MultiIndex representing the column headers
- escape: Whether to HTML-escape the cell values
-
- Returns:
- HTML string for the element
- """
- import html as html_lib
-
- n_levels = columns.nlevels
- n_cols = len(columns)
-
- # Build a 2D grid of values [level][col]
- grid = []
- for level in range(n_levels):
- row = [columns.get_level_values(level)[col] for col in range(n_cols)]
- grid.append(row)
-
- # Calculate colspan for each cell (horizontal merging)
- # colspan[level][col] = number of columns this cell spans
- colspan = [[1] * n_cols for _ in range(n_levels)]
-
- for level in range(n_levels):
- col = 0
- while col < n_cols:
- span = 1
- while col + span < n_cols and grid[level][col] == grid[level][col + span]:
- # Check if the parent cells also match (for correct hierarchical merging)
- parent_match = True
- for parent_level in range(level):
- if grid[parent_level][col] != grid[parent_level][col + span]:
- parent_match = False
- break
- if parent_match:
- span += 1
- else:
- break
- colspan[level][col] = span
- col += span
-
- # Calculate rowspan for each cell (vertical merging)
- # A cell has rowspan > 1 if all cells in the same column below have the same value
- # AND if they would have the same colspan
- rowspan = [[1] * n_cols for _ in range(n_levels)]
-
- for col in range(n_cols):
- level = 0
- while level < n_levels:
- span = 1
- # Check if cells below have the same value AND same colspan
- while level + span < n_levels:
- if (
- grid[level][col] == grid[level + span][col]
- and colspan[level][col] == colspan[level + span][col]
- ):
- span += 1
- else:
- break
- rowspan[level][col] = span
- level += span
-
- # Build HTML rows
- # Track which cells are "covered" by rowspan from above
- covered = [[False] * n_cols for _ in range(n_levels)]
-
- html_parts = [""]
-
- for level in range(n_levels):
- html_parts.append('')
- col = 0
- while col < n_cols:
- if covered[level][col]:
- # This cell is covered by a rowspan from above, skip it
- col += 1
- continue
-
- # Get cell value
- val = grid[level][col]
- val_str = str(val) if val is not None else ""
- if escape:
- val_str = html_lib.escape(val_str)
-
- # Get spans
- cs = colspan[level][col]
- rs = rowspan[level][col]
-
- # Mark covered cells
- for r_offset in range(rs):
- for c_offset in range(cs):
- if r_offset > 0 or c_offset > 0:
- if level + r_offset < n_levels and col + c_offset < n_cols:
- covered[level + r_offset][col + c_offset] = True
-
- # Build th element with attributes
- attrs = []
- if cs > 1:
- attrs.append(f'colspan="{cs}"')
- if rs > 1:
- attrs.append(f'rowspan="{rs}"')
-
- attr_str = " " + " ".join(attrs) if attrs else ""
- html_parts.append(f"{val_str}")
-
- col += cs
-
- html_parts.append("
")
-
- html_parts.append("")
- return "".join(html_parts)
-
-
-def render_tbody_with_row_headers(
- tb_df: pd.DataFrame,
- row_header_cols: int = 0,
- na_rep: str = "—",
- escape: bool = False,
-) -> str:
- """
- Render DataFrame body with support for row headers and cell merging.
-
- This function generates a proper structure where:
- - Row header columns use instead of |
- - Horizontally adjacent identical values in row headers are merged with colspan
- - Vertically adjacent identical values in row headers are merged with rowspan
- - Merging respects hierarchical structure
-
- Args:
- tb_df: DataFrame to render
- row_header_cols: Number of leftmost columns to render as |
- na_rep: String representation for NaN values
- escape: Whether to HTML-escape values
-
- Returns:
- HTML string for the | element
- """
- import html as html_lib
-
- if row_header_cols <= 0:
- # No row headers - simple rendering without merging
- html_parts = [""]
- for _, row in tb_df.iterrows():
- html_parts.append("")
- for val in row:
- if pd.isna(val):
- val_str = na_rep
- else:
- val_str = str(val)
- if escape:
- val_str = html_lib.escape(val_str)
- html_parts.append(f"| {val_str} | ")
- html_parts.append("
")
- html_parts.append("")
- return "".join(html_parts)
-
- n_rows = len(tb_df)
- n_cols = len(tb_df.columns)
-
- if n_rows == 0:
- return ""
-
- # Build 2D grid of values for row header columns
- # grid[row_idx][col_idx] = value
- grid = []
- for row_idx in range(n_rows):
- row_values = []
- for col_idx in range(row_header_cols):
- val = tb_df.iloc[row_idx, col_idx]
- if pd.isna(val):
- val = na_rep
- else:
- val = str(val)
- row_values.append(val)
- grid.append(row_values)
-
- # Calculate colspan for each cell (horizontal merging within same row)
- # colspan[row_idx][col_idx] = number of columns this cell spans
- colspan = [[1] * row_header_cols for _ in range(n_rows)]
-
- for row_idx in range(n_rows):
- col_idx = 0
- while col_idx < row_header_cols:
- span = 1
- while (
- col_idx + span < row_header_cols
- and grid[row_idx][col_idx] == grid[row_idx][col_idx + span]
- ):
- span += 1
- colspan[row_idx][col_idx] = span
- col_idx += span
-
- # Calculate rowspan for each cell (vertical merging)
- # Only calculate rowspan for cells that start a colspan group
- # rowspan[row_idx][col_idx] = number of rows this cell spans
- rowspan = [[1] * row_header_cols for _ in range(n_rows)]
-
- col_idx = 0
- while col_idx < row_header_cols:
- row_idx = 0
- while row_idx < n_rows:
- # Only process cells that start a colspan group (not covered by colspan from left)
- if col_idx > 0 and grid[row_idx][col_idx] == grid[row_idx][col_idx - 1]:
- row_idx += 1
- continue
-
- current_colspan = colspan[row_idx][col_idx]
- span = 1
-
- while row_idx + span < n_rows:
- # Check if the value matches
- if grid[row_idx][col_idx] != grid[row_idx + span][col_idx]:
- break
- # Check if colspan in the next row also matches
- if colspan[row_idx + span][col_idx] != current_colspan:
- break
- # Check if all parent columns (to the left) also have same rowspan behavior
- parent_match = True
- for parent_col in range(col_idx):
- if grid[row_idx][parent_col] != grid[row_idx + span][parent_col]:
- parent_match = False
- break
- if parent_match:
- span += 1
- else:
- break
-
- rowspan[row_idx][col_idx] = span
- row_idx += span
- col_idx += 1
-
- # Track which cells are covered by rowspan from above or colspan from left
- covered = [[False] * row_header_cols for _ in range(n_rows)]
-
- # Mark cells covered by colspan (horizontal)
- for row_idx in range(n_rows):
- col_idx = 0
- while col_idx < row_header_cols:
- cs = colspan[row_idx][col_idx]
- for offset in range(1, cs):
- if col_idx + offset < row_header_cols:
- covered[row_idx][col_idx + offset] = True
- col_idx += cs
-
- # Mark cells covered by rowspan (vertical)
- for row_idx in range(n_rows):
- for col_idx in range(row_header_cols):
- if covered[row_idx][col_idx]:
- continue # Skip cells already covered by colspan
- rs = rowspan[row_idx][col_idx]
- for offset in range(1, rs):
- if row_idx + offset < n_rows:
- # Mark all cells in the rowspan as covered
- cs = colspan[row_idx][col_idx]
- for c_offset in range(cs):
- if col_idx + c_offset < row_header_cols:
- covered[row_idx + offset][col_idx + c_offset] = True
-
- # Build HTML
- html_parts = [""]
-
- for row_idx in range(n_rows):
- html_parts.append("")
-
- # Render row header columns with rowspan/colspan
- for col_idx in range(row_header_cols):
- if covered[row_idx][col_idx]:
- # This cell is covered by a rowspan/colspan, skip it
- continue
-
- val_str = grid[row_idx][col_idx]
- if escape:
- val_str = html_lib.escape(val_str)
-
- rs = rowspan[row_idx][col_idx]
- cs = colspan[row_idx][col_idx]
-
- attrs = []
- if rs > 1:
- attrs.append(f'rowspan="{rs}"')
- if cs > 1:
- attrs.append(f'colspan="{cs}"')
-
- attr_str = " " + " ".join(attrs) if attrs else ""
- html_parts.append(f'| {val_str} | ')
-
- # Render data columns
- for col_idx in range(row_header_cols, n_cols):
- val = tb_df.iloc[row_idx, col_idx]
- if pd.isna(val):
- val_str = na_rep
- else:
- val_str = str(val)
- if escape:
- val_str = html_lib.escape(val_str)
- html_parts.append(f"{val_str} | ")
-
- html_parts.append("
")
-
- html_parts.append("")
- return "".join(html_parts)
-
-
-def df2html(
- tb_df: pd.DataFrame,
- *,
- index: bool = False,
- classes: Union[str, List[str], None] = "table table-striped",
- na_rep: str = "—",
- escape: bool = False,
- row_header_cols: int = 0,
-) -> str:
- """Convert DataFrame to HTML table.
-
- Supports:
- - MultiIndex columns with proper colspan/rowspan merging
- - Row headers (leftmost columns rendered as )
- - Custom CSS classes and NA representation
-
- Args:
- tb_df: DataFrame to convert
- index: Whether to include the DataFrame index (not commonly used)
- classes: CSS classes to add to the table
- na_rep: String representation for NaN values
- escape: Whether to HTML-escape values
- row_header_cols: Number of leftmost columns to render as row headers ( | ).
- These columns will use | instead of | .
-
- Returns:
- HTML table string
- """
- class_str = (
- classes if isinstance(classes, str) else " ".join(classes) if classes else ""
- )
-
- # Check if columns are MultiIndex - use advanced rendering
- if isinstance(tb_df.columns, pd.MultiIndex):
- # Use specialized rendering for MultiIndex columns
- thead_html = render_multiindex_thead(tb_df.columns, escape=escape)
- tbody_html = render_tbody_with_row_headers(
- tb_df, row_header_cols, na_rep, escape
- )
- return f''
-
- # Simple columns case
- if row_header_cols <= 0:
- # Use default pandas to_html for simple case without row headers
- table_html = tb_df.to_html(
- index=index,
- na_rep=na_rep,
- classes=classes,
- escape=escape,
- border=0,
- justify="center",
- )
- return table_html.replace("\n", "")
-
- # Simple columns with row headers - custom rendering
- import html as html_lib
-
- html_parts = [f'']
-
- # Build thead
- html_parts.append("")
- html_parts.append('')
- for col in tb_df.columns:
- col_str = str(col) if col is not None else ""
- if escape:
- col_str = html_lib.escape(col_str)
- html_parts.append(f"| {col_str} | ")
- html_parts.append(" ")
- html_parts.append("")
-
- # Build tbody with row headers
- tbody_html = render_tbody_with_row_headers(tb_df, row_header_cols, na_rep, escape)
- html_parts.append(tbody_html)
-
- html_parts.append(" ")
-
- return "".join(html_parts)
diff --git a/apps/worker/app/services/document_parser/identifiers.py b/apps/worker/app/services/document_parser/identifiers.py
new file mode 100644
index 000000000..c8d177c3a
--- /dev/null
+++ b/apps/worker/app/services/document_parser/identifiers.py
@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import uuid
+from datetime import datetime
+
+
+def gen_str_codes(input_string: str) -> str:
+ """Generate a UUID5 code from a string."""
+ return str(uuid.uuid5(uuid.NAMESPACE_DNS, input_string))
+
+
+def get_str_time() -> str:
+ """Get the current time as a string."""
+ return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py
index dc8ffb541..cc9be0ab1 100755
--- a/apps/worker/app/services/document_parser/image_parser.py
+++ b/apps/worker/app/services/document_parser/image_parser.py
@@ -8,11 +8,9 @@
from pathlib import Path
import pandas as pd
-from app.services.common.kb_utils import (
- gen_str_codes,
- get_str_time,
- process_dup_paths_df,
-)
+from app.services.document_parser.dataframe_helpers import process_dup_paths_df
+from app.services.document_parser.identifiers import gen_str_codes, get_str_time
+from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder
from loguru import logger
from PIL import Image
@@ -25,7 +23,7 @@
from shared.services.ai.prompt_service import build_prompt
from shared.services.ai.response_process_service import eval_response
from shared.utils.chunk_refs import build_chunk_ref
-from shared.utils.CommonHelperSync import is_remote, load_file_bytes
+from shared.utils.file_loading import is_remote, load_file_bytes
from shared.utils.file_utils import path_handle
from shared.utils.OpenAICompatibleClientSync import (
OpenAICompatibleClientSync,
@@ -227,7 +225,6 @@ def parse_image(
relative_root=None,
):
split_char = settings.SPLIT_CHAR or "/"
- df_list = []
time_stamp = get_str_time()
os.makedirs(output_dir, exist_ok=True)
img_dir = os.path.join(output_dir, "images")
@@ -358,23 +355,19 @@ def parse_image(
)
img_ref = build_chunk_ref(relative_img_path)
img_bottom_content = f"{img_ref}\nImage Content:\n{image_content}"
- df_list.append(
- [
- img_bottom_content,
- relative_img_path,
- "image",
- len(img_bottom_content),
- "",
- image_summary,
- temp_uid,
- "",
- "",
- time_stamp,
- "",
- ]
+ rows_builder = ParsedRowsBuilder()
+ rows_builder.append(
+ ParsedRow(
+ content=img_bottom_content,
+ path=relative_img_path,
+ type="image",
+ summary=image_summary,
+ know_id=temp_uid,
+ addtime=time_stamp,
+ )
)
- img_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(","))
+ img_df = rows_builder.to_dataframe()
img_df = process_dup_paths_df(img_df)
return img_df
diff --git a/apps/worker/app/services/document_parser/inline_asset.py b/apps/worker/app/services/document_parser/inline_asset.py
new file mode 100644
index 000000000..4cb16b531
--- /dev/null
+++ b/apps/worker/app/services/document_parser/inline_asset.py
@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from app.services.document_parser.parser_rows import ParsedRow
+
+
+def build_image_asset_row(
+ *,
+ content: str,
+ relative_path: str,
+ summary: str,
+ know_id: str,
+ addtime: str,
+ page_nums: str = "",
+) -> ParsedRow:
+ return ParsedRow(
+ content=content,
+ path=relative_path,
+ type="image",
+ keywords="",
+ summary=summary,
+ know_id=know_id,
+ tokens="",
+ connectto="",
+ addtime=addtime,
+ page_nums=page_nums,
+ )
+
+
+def build_table_asset_row(
+ *,
+ content: str,
+ relative_path: str,
+ summary: str,
+ keywords: str,
+ know_id: str,
+ addtime: str,
+ page_nums: str = "",
+) -> ParsedRow:
+ return ParsedRow(
+ content=content,
+ path=relative_path,
+ type="table",
+ keywords=keywords,
+ summary=summary,
+ know_id=know_id,
+ tokens="",
+ connectto="",
+ addtime=addtime,
+ page_nums=page_nums,
+ )
diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py
index f308e7f5a..b0490d576 100755
--- a/apps/worker/app/services/document_parser/layout_parser.py
+++ b/apps/worker/app/services/document_parser/layout_parser.py
@@ -1,33 +1,34 @@
# pyright: reportArgumentType=false, reportAssignmentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalSubscript=false
import os
-import re
-import unicodedata
-from collections import Counter, defaultdict
import gevent
import pandas as pd
-from app.services.common.kb_utils import (
- count_cn_en,
- truncate_text_by_tokens,
+from app.services.document_parser.heading_candidates import (
+ filter_document_headings,
+ filter_markdown_headings,
+ postprocess_headings,
+)
+from app.services.document_parser.heading_llm_executor import (
+ build_level_mapping,
+ execute_level_mapping,
+ execute_llm_heading_hierarchy,
+)
+from app.services.document_parser.heading_tree import (
+ build_tree_from_dataframe as build_heading_tree_from_dataframe,
+)
+from app.services.document_parser.heading_tree import (
+ remove_isolated_nodes as remove_isolated_heading_nodes,
+)
+from app.services.document_parser.heading_tree import (
+ tree_to_dataframe as heading_tree_to_dataframe,
)
from app.services.document_parser.stage_profiler import stage_timer
-from app.services.document_parser.table_parser import df2md
-from docx.oxml.ns import qn
+from app.services.document_parser.table_text_parser import df2md
from gevent.pool import Pool as GeventPool
-try:
- from markitdown import MarkItDown
-except ImportError:
- # Fall back to a pass-through shim when markitdown is unavailable.
- class MarkItDown:
- def convert(self, content):
- return content
-
-
from loguru import logger
from shared.core.config import settings
-from shared.core.exceptions.domain_exceptions import WorkerHandlingException
# TaskRedis dependency is removed, use Redis directly to track
from shared.services.ai.prompt_service import build_prompt
@@ -71,1069 +72,15 @@ def save_intermediate_csv(df: pd.DataFrame, output_dir: str, filename: str):
def build_tree_from_dataframe(df):
- """
- develop json tree from dataframe
-
- Args:
- df: DataFrame, including id, heading, level columns
-
- Returns:
- tree: pure nested dict structure
- node_to_id: map from tree node to id (use unique node key)
- id_to_row: map from id to original row data
- """
- headings = df[df["level"] > -1].copy()
-
- node_to_id = {} # {(tree_node_key, parent_path): id}
- id_to_node_info = {} # {id: (tree_node_key, parent_path)}
- id_to_row = {}
- root = {}
- stack = [(0, root, "ROOT", "")]
-
- for _, row in headings.iterrows():
- heading_txt = row["heading"]
- row_id = int(row["id"])
- level = int(row["level"])
-
- # record id to row mapping
- id_to_row[row_id] = row.to_dict()
-
- # find suitable parent node
- while len(stack) > 1 and stack[-1][0] >= level:
- stack.pop()
-
- # get parent node info
- parent_level, parent_dict, parent_heading, parent_path = stack[-1]
-
- # create unique key for tree node: if there are duplicate headings under the same parent, add ID suffix
- tree_node_key = heading_txt
-
- if tree_node_key in parent_dict:
- tree_node_key = f"{heading_txt}#{row_id}"
-
- # build mapping: use (tree_node_key, parent_path) as key
- node_key = (tree_node_key, parent_path)
- node_to_id[node_key] = row_id
- id_to_node_info[row_id] = node_key
-
- parent_dict[tree_node_key] = {}
- current_path = (
- f"{parent_path}/{tree_node_key}" if parent_path else tree_node_key
- )
- stack.append((level, parent_dict[tree_node_key], tree_node_key, current_path))
- return root, node_to_id, id_to_row
+ return build_heading_tree_from_dataframe(df)
def tree_to_dataframe(tree, node_to_id, original_df):
- """
- convert processed tree structure back to dataframe
-
- Args:
- tree: processed pure nested dict structure
- node_to_id: map from node to id {(tree_node_key, parent_path): id}
- original_df: original dataframe
-
- Returns:
- updated_df: updated dataframe
- """
-
- # extract all retained headings from tree
- def extract_headings(node_dict, current_level=1, parent_path=""):
- """recursively extract all headings and their new levels"""
- results = []
- for tree_node_key, children in node_dict.items():
- # use (tree_node_key, parent_path) as key to find ID
- node_key = (tree_node_key, parent_path)
- row_id = node_to_id.get(node_key, -1)
-
- if row_id >= 0:
- # extract original heading from tree_node_key (remove possible ID suffix)
- original_heading = (
- tree_node_key.split("#")[0]
- if "#" in tree_node_key
- else tree_node_key
- )
-
- results.append(
- {
- "id": row_id,
- "heading": original_heading,
- "level": current_level,
- "tree_key": tree_node_key,
- "parent_path": parent_path,
- }
- )
- # recursively process child nodes
- if isinstance(children, dict) and children:
- current_path = (
- f"{parent_path}/{tree_node_key}"
- if parent_path
- else tree_node_key
- )
- results.extend(
- extract_headings(children, current_level + 1, current_path)
- )
- return results
-
- preserved_headings = extract_headings(tree)
- preserved_ids = set([h["id"] for h in preserved_headings])
-
- updated_df = original_df.copy()
- removed_count = 0
- level_changed_count = 0
-
- for idx, row in original_df.iterrows():
- row_id = int(row["id"])
- old_level = int(row["level"]) if row["level"] not in [-2, "nan", -1] else -1
-
- if old_level > -1:
- if row_id in preserved_ids:
- new_level = next(
- (h["level"] for h in preserved_headings if h["id"] == row_id),
- old_level,
- )
- updated_df.at[idx, "level"] = new_level
- if new_level != old_level:
- level_changed_count += 1
- else:
- updated_df.at[idx, "level"] = -1
- removed_count += 1
-
- logger.debug(
- f"Tree changed: removed headings={removed_count}, level changed={level_changed_count}, preserved headings={len(preserved_ids)}"
- )
- return updated_df
+ return heading_tree_to_dataframe(tree, node_to_id, original_df)
def remove_isolated_nodes(tree):
- """
- rules: if a heading has only one child heading, and the child heading has no further child headings,
- then delete this isolated child heading
-
- Args:
- tree: pure nested dict structure, format as {heading: {child_heading: {...}}}
-
- Returns:
- processed_tree: processed tree structure
- """
-
- def recursive_check_and_remove(node_dict, parent_path=""):
- if not isinstance(node_dict, dict):
- return node_dict
-
- result_dict = {}
-
- for heading, children in node_dict.items():
- if isinstance(children, dict) and len(children) == 1:
- child_heading = list(children.keys())[0]
- grandchildren = children[child_heading]
-
- if not grandchildren or (
- isinstance(grandchildren, dict) and len(grandchildren) == 0
- ):
- result_dict[heading] = {}
- logger.debug(
- f"remove isolated heading: {parent_path}/{heading}/{child_heading}"
- )
- else:
- processed_children = recursive_check_and_remove(
- children, f"{parent_path}/{heading}" if parent_path else heading
- )
- result_dict[heading] = processed_children
- elif isinstance(children, dict) and children:
- processed_children = recursive_check_and_remove(
- children, f"{parent_path}/{heading}" if parent_path else heading
- )
- result_dict[heading] = processed_children
- else:
- result_dict[heading] = children
-
- return result_dict
-
- processed_tree = recursive_check_and_remove(tree)
- return processed_tree
-
-
-# def if_no_pos_code(reason_str: str) -> bool:
-# """
-# Check whether all pos_code values are zero.
-# reason format: "POS [0, 0, ...] NEG [...]"
-# """
-# if not reason_str or not isinstance(reason_str, str):
-# return True
-
-# pos_match = re.search(r'POS\s*\[([^\]]*)\]', reason_str)
-# if not pos_match:
-# return True
-# pos_content = pos_match.group(1)
-# try:
-# nums = [int(x.strip()) for x in pos_content.split(',') if x.strip()]
-# return all(x == 0 for x in nums)
-# except:
-# return True
-
-
-# ==================== Level Mapping Functions ====================
-
-
-def build_level_mapping(df, origin_lvls, mode="max"):
- df = df.copy()
- df["origin_level"] = origin_lvls
-
- mapping = df.groupby("reason")["level"].apply(list).to_dict()
-
- processed_mapping = {}
- for reason, lvls in mapping.items():
- positive_lvls = [lvl for lvl in lvls if lvl > -1]
- counts = Counter(lvls)
-
- if not positive_lvls:
- mapped_lvl = -1
- elif mode == "max":
- mapped_lvl = max(positive_lvls)
- elif mode == "freq":
- mapped_lvl = counts.most_common(1)[0][0]
- else:
- raise WorkerHandlingException(
- internal_message=f"wrong input mode: {mode}. Must be 'max' or 'freq'"
- )
-
- processed_mapping[reason] = {
- "lvls": lvls,
- "positive_lvls": positive_lvls,
- "freqs": dict(counts),
- "mapped_lvl": mapped_lvl,
- }
- return df, processed_mapping
-
-
-def execute_level_mapping(df: pd.DataFrame, mapping: dict) -> pd.DataFrame:
- def map_row(row):
- reason = row["reason"]
- if reason in mapping:
- return mapping[reason]["mapped_lvl"]
- return row["level"]
-
- df = df.copy()
- origin_est_lvls = df["level"].tolist()
- df["level"] = df.apply(map_row, axis=1)
- df["origin_level"] = origin_est_lvls
- return df
-
-
-def extract_non_neg_code(reason_str: str) -> str:
- """
- Extract the non-NEG code from reason_str (strip only NEG part, preserve META)
-
- Example: "POS [1, 0, 0] NEG [0, 0, 0] META [1, 2, 1]" -> "POS [1, 0, 0] META [1, 2, 1]"
- Example: "3# AND POS [1, 0] NEG [0, 0]" -> "3# AND POS [1, 0]"
- Example: "3# AND POS [1, 0] NEG [0, 0] META [1, 1, 0]" -> "3# AND POS [1, 0] META [1, 1, 0]"
- """
- if not reason_str or not isinstance(reason_str, str):
- return ""
- neg_match = re.search(r"\s*NEG\s*\[[^\]]*\]", reason_str)
- if neg_match:
- # Remove only the NEG [...] part, keep everything before and after
- before_neg = reason_str[: neg_match.start()]
- after_neg = reason_str[neg_match.end() :]
- return (before_neg + after_neg).strip()
- return reason_str.strip()
-
-
-def build_non_neg_mapping(lvl_mapping: dict) -> dict:
- """
- Build non-NEG code mapping from complete lvl_mapping, select by highest frequency
-
- Args:
- lvl_mapping: reason -> level mapping
-
- Returns:
- non_neg_mapping: {non_neg_code: mapped_lvl}
- """
- # collect all levels for each non_neg_code
- non_neg_levels = {}
- for reason, info in lvl_mapping.items():
- non_neg_code = extract_non_neg_code(reason)
- mapped_lvl = info.get("mapped_lvl", -1)
- if non_neg_code:
- if non_neg_code not in non_neg_levels:
- non_neg_levels[non_neg_code] = []
- non_neg_levels[non_neg_code].append(mapped_lvl)
-
- # select by highest frequency
- non_neg_mapping = {}
- for non_neg_code, levels in non_neg_levels.items():
- positive_levels = [lvl for lvl in levels if lvl > -1]
- if positive_levels:
- level_counts = Counter(positive_levels)
- most_common_level = level_counts.most_common(1)[0][0]
- non_neg_mapping[non_neg_code] = most_common_level
- else:
- non_neg_mapping[non_neg_code] = -1
-
- return non_neg_mapping
-
-
-def handle_unseen_codes(
- df: pd.DataFrame,
- level_dfs: list,
- lvl_mapping: dict,
- output_dir: str = None,
- window_half_size: int = 10,
- strategy: str = "double_mapping",
-) -> dict:
- """
- Handle unseen codes with configurable strategy
-
- Args:
- df: original complete DataFrame
- level_dfs: segment DataFrames
- lvl_mapping: existing level mapping
- output_dir: output directory (optional, only used for window_llm strategy)
- window_half_size: window half size (how many rows above and below)
- strategy: "double_mapping" or "window_llm"
- - double_mapping: use non-neg code fallback (fast, no LLM call)
- - window_llm: create windows for LLM to judge (slower, more accurate)
-
- Returns:
- updated lvl_mapping
- """
-
- def extract_reason_signature(reason: str) -> str:
- """Extract reason signature"""
- return reason.strip() if reason else ""
-
- def has_neg_signal(reason_str: str) -> bool:
- """Check if NEG signal exists (any value >= 1)"""
- if not reason_str or not isinstance(reason_str, str):
- return False
- neg_match = re.search(r"NEG\s*\[([^\]]*)\]", reason_str)
- if not neg_match:
- return False
- neg_content = neg_match.group(1)
- try:
- nums = [int(x.strip()) for x in neg_content.split(",") if x.strip()]
- return any(x >= 1 for x in nums)
- except Exception:
- return False
-
- def build_context_window(
- target_idx: int, known_codes_set: set, total_rows: int, half_size: int = 10
- ) -> dict:
- """
- Build context window for unseen codes
- 1. window size: half_size
- 2. window should contain at least one known code
- """
- min_start = max(0, target_idx - half_size)
- min_end = min(total_rows - 1, target_idx + half_size)
-
- start_idx = min_start
- end_idx = min_end
-
- found_known_above = False
- found_known_below = False
- known_positions = []
-
- # check above
- for i in range(start_idx, target_idx):
- reason = df.iloc[i].get("reason", "")
- sig = extract_reason_signature(reason)
- if sig in known_codes_set:
- found_known_above = True
- known_positions.append(i)
-
- # check below
- for i in range(target_idx + 1, end_idx + 1):
- reason = df.iloc[i].get("reason", "")
- sig = extract_reason_signature(reason)
- if sig in known_codes_set:
- found_known_below = True
- known_positions.append(i)
-
- # expand above if needed
- if not found_known_above and min_start > 0:
- search_idx = min_start - 1
- while search_idx >= 0:
- reason = df.iloc[search_idx].get("reason", "")
- sig = extract_reason_signature(reason)
- if sig in known_codes_set:
- found_known_above = True
- known_positions.append(search_idx)
- start_idx = search_idx
- break
- search_idx -= 1
-
- # expand below if needed
- if not found_known_below and min_end < total_rows - 1:
- search_idx = min_end + 1
- while search_idx < total_rows:
- reason = df.iloc[search_idx].get("reason", "")
- sig = extract_reason_signature(reason)
- if sig in known_codes_set:
- found_known_below = True
- known_positions.append(search_idx)
- end_idx = search_idx
- break
- search_idx += 1
-
- return {
- "start": start_idx,
- "end": end_idx,
- "found_known": found_known_above or found_known_below,
- "known_positions": known_positions,
- }
-
- # build non-neg mapping
- non_neg_mapping = build_non_neg_mapping(lvl_mapping)
-
- # get known codes
- known_codes = set(lvl_mapping.keys())
-
- # record all codes from all segments. Placeholder rows (reason ==
- # PLACEHOLDER_REASON) are injected by _compact_for_llm and are never real
- # heading candidates, so they must be skipped here — otherwise they would
- # show up as an "unseen code" and fall through to NO_MATCH_FALLBACK, adding
- # harmless but noisy warnings to the log.
- all_codes_in_full = {}
- for seg_idx, seg_df in enumerate(level_dfs):
- for _, row in seg_df.iterrows():
- reason = row.get("reason", "")
- sig = extract_reason_signature(reason)
- if not sig or sig == PLACEHOLDER_REASON:
- continue
- if sig not in all_codes_in_full:
- all_codes_in_full[sig] = {
- "first_seg": seg_idx,
- "first_id": row.get("id", 0),
- "reason": reason,
- }
-
- # find unseen codes
- unseen_codes = {}
- unseen_neg_filtered = {}
- for sig, info in all_codes_in_full.items():
- if sig in known_codes:
- continue
- if has_neg_signal(info["reason"]):
- unseen_neg_filtered[sig] = info
- else:
- unseen_codes[sig] = info
-
- logger.info(
- f"Unseen codes total: {len(unseen_codes) + len(unseen_neg_filtered)}, NEG filtered: {len(unseen_neg_filtered)}, to process: {len(unseen_codes)}"
- )
-
- # if neg signal, map to -1
- for sig in unseen_neg_filtered:
- lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NEG_FILTERED"}
-
- # handle remaining unseen_codes based on strategy
- if unseen_codes:
- if strategy == "double_mapping":
- # Strategy 1: use non-neg code fallback
- fallback_success = 0
- fallback_failed = 0
- failed_codes = []
- for sig, info in unseen_codes.items():
- non_neg_code = extract_non_neg_code(sig)
- if non_neg_code in non_neg_mapping:
- mapped_level = non_neg_mapping[non_neg_code]
- lvl_mapping[sig] = {
- "mapped_lvl": mapped_level,
- "note": f"NON_NEG_FALLBACK from '{non_neg_code}'",
- }
- fallback_success += 1
- else:
- lvl_mapping[sig] = {"mapped_lvl": -1, "note": "NO_MATCH_FALLBACK"}
- fallback_failed += 1
- failed_codes.append(
- f"'{non_neg_code}' (from '{sig[:60]}...')"
- if len(sig) > 60
- else f"'{non_neg_code}' (from '{sig}')"
- )
-
- logger.info(
- f"Double mapping result: success={fallback_success}, failed={fallback_failed}"
- )
- if failed_codes:
- logger.warning(
- f"Failed codes (non_neg not in mapping): {failed_codes[:5]}{'...' if len(failed_codes) > 5 else ''}"
- )
-
- elif strategy == "window_llm" and output_dir:
- # Strategy 2: create windows for LLM to judge
- total_rows = len(df)
- windows = []
- for sig, info in unseen_codes.items():
- first_id = info["first_id"]
- first_seg = info["first_seg"]
- df_indices = df.index[df["id"] == first_id].tolist()
- if df_indices:
- first_df_idx = df_indices[0]
- window_info = build_context_window(
- first_df_idx, known_codes, total_rows, window_half_size
- )
- windows.append(
- {
- "code": sig,
- "first_id": first_id,
- "first_seg": first_seg,
- "start": window_info["start"],
- "end": window_info["end"],
- "found_known": window_info["found_known"],
- }
- )
-
- # merge windows
- sorted_windows = sorted(windows, key=lambda x: x["start"])
- merged_windows = []
- current_window = None
-
- for w in sorted_windows:
- if current_window is None:
- current_window = {
- "start": w["start"],
- "end": w["end"],
- "codes": [w["code"]],
- "segments": [w["first_seg"]],
- }
- elif w["start"] <= current_window["end"]:
- current_window["end"] = max(current_window["end"], w["end"])
- current_window["codes"].append(w["code"])
- current_window["segments"].append(w["first_seg"])
- else:
- merged_windows.append(current_window)
- current_window = {
- "start": w["start"],
- "end": w["end"],
- "codes": [w["code"]],
- "segments": [w["first_seg"]],
- }
-
- if current_window:
- merged_windows.append(current_window)
-
- # save windows
- windows_dir = os.path.join(output_dir, "merged_windows")
- os.makedirs(windows_dir, exist_ok=True)
-
- unseen_codes_set = set(unseen_codes.keys())
- unseen_neg_set = set(unseen_neg_filtered.keys())
-
- for i, mw in enumerate(merged_windows):
- window_df = df.iloc[mw["start"] : mw["end"] + 1].copy()
-
- def get_code_status(row):
- reason = row.get("reason", "")
- sig = extract_reason_signature(reason)
- if not sig:
- return ""
- if sig in unseen_codes_set:
- return "★ UNSEEN_TARGET"
- elif sig in unseen_neg_set:
- return "NEG→-1"
- elif sig in known_codes:
- return "KNOWN"
- else:
- return ""
-
- window_df["code_status"] = window_df.apply(get_code_status, axis=1)
- window_path = os.path.join(
- windows_dir,
- f"window_{i + 1:02d}_rows_{mw['start']}-{mw['end']}.csv",
- )
- window_df.to_csv(window_path, index=False, encoding="utf-8-sig")
-
- logger.debug(
- f"Window LLM: {len(merged_windows)} windows created in {windows_dir}"
- )
- # TODO: use llm to assign level based on window data
-
- return lvl_mapping
-
-
-def detect_outlines_md(line):
- pos_code = judge_by_conditions(line)
- any(x > 0 for x in pos_code)
-
-
-def get_max_lvl(code_str: str):
- match = re.search(r"\[([^]]+)]", code_str)
- if not match:
- return "Sure"
-
- nums = [int(x.strip()) for x in match.group(1).split(",")]
- max_val = int(max(nums))
- return max_val if max_val > 1 else -2 # -2 = "Not Sure" sentinel (int-safe)
-
-
-PLACEHOLDER_REASON = "__PLACEHOLDER__"
-
-
-def _compact_for_llm(df: pd.DataFrame) -> pd.DataFrame:
- """Collapse consecutive ``level == -1`` rows into a single placeholder row.
-
- Rows with ``level >= 1`` (heading candidates) and ``level == -2`` ("Not Sure")
- are preserved verbatim so the LLM can still judge them. Each run of
- consecutive ``-1`` rows becomes one placeholder row whose:
-
- id = "start-end" (always a range; "N-N" when the run is one row)
- heading = "[N BODY LINES]" where N is the run length
- level = "-"
- reason = ``PLACEHOLDER_REASON``
-
- The id is ALWAYS a hyphenated string, even for single-row runs, so that
- ``int(id)`` fails for every placeholder. This lets downstream code identify
- placeholders structurally (non-integer id) without depending on ``reason``
- or length heuristics.
- """
- if df is None or len(df) == 0:
- return pd.DataFrame(columns=["id", "heading", "level", "reason"])
-
- rows = []
- i = 0
- n = len(df)
- while i < n:
- lvl_raw = df.iloc[i]["level"]
- try:
- lvl_int = int(lvl_raw)
- except (TypeError, ValueError):
- lvl_int = None
-
- if lvl_int == -1:
- j = i
- while j < n:
- try:
- nxt_lvl = int(df.iloc[j]["level"])
- except (TypeError, ValueError):
- break
- if nxt_lvl != -1:
- break
- j += 1
- start_id = int(df.iloc[i]["id"])
- end_id = int(df.iloc[j - 1]["id"])
- run = j - i
- rows.append(
- {
- "id": f"{start_id}-{end_id}",
- "heading": f"[{run} BODY LINES]",
- "level": "-",
- "reason": PLACEHOLDER_REASON,
- }
- )
- i = j
- else:
- r = df.iloc[i]
- rows.append(
- {
- "id": int(r["id"]),
- "heading": str(r["heading"]),
- "level": (
- int(lvl_int)
- if lvl_int is not None and lvl_int != -2
- else "Not Sure"
- ),
- "reason": str(r.get("reason", "") or ""),
- }
- )
- i += 1
-
- return pd.DataFrame(rows, columns=["id", "heading", "level", "reason"])
-
-
-def heading_tb_transfer(df, threshold=3000, max_start=50, max_end=10):
- raw_headings = df["heading"].tolist()
- df["heading"] = df["heading"].apply(
- lambda x: truncate_text_by_tokens(x, max_start, max_end)
- )
-
- sub_dfs = []
- current_rows = []
- current_len = 0
- for _, row in df.iterrows():
- row_filtered = row.drop(labels=["reason"], errors="ignore")
- row_len = sum(count_cn_en(str(v)) for v in row_filtered.values)
-
- if current_len + row_len > threshold and current_rows:
- sub_dfs.append(pd.DataFrame(current_rows, columns=df.columns))
- current_rows = [row.tolist()]
- current_len = row_len
- else:
- current_rows.append(row.tolist())
- current_len += row_len
-
- if current_rows:
- sub_dfs.append(pd.DataFrame(current_rows, columns=df.columns))
- return sub_dfs, raw_headings
-
-
-def judge_by_conditions(text, scope=20, return_detail=False, CN_SPECIAL_IDX=12):
- """
- judge level features as one-hot embeddings for texts
-
- Args:
- text: input text
- scope: text scope for judging
- return_detail: whether to return detailed information (including unit type)
- CN_SPECIAL_IDX: index of special Chinese number
-
- Returns:
- if return_detail=False: return pos_triggered_code list
- if return_detail=True: return (pos_triggered_code, detail_info) tuple
- where detail_info is a dictionary containing additional information, such as Chinese unit type
- """
- text = text.replace("\u3000", " ")
- text = unicodedata.normalize("NFKC", text)[:scope]
-
- # ========== English Numbering ==========
- regex_en_num_dots = r"^\d+(?:\s*\.\s*\d+)+(?![、，。！？；：])(?=\s|$|\w|[一-龥])"
- regex_en_num_dun = r"^\d、\s{0,4}(?=\S|$)" # 1、xxx
- regex_en_num_dots_dun = r"^\d+(?:\.\d+)*、\s*(?=[A-Za-z一-龥])"
- regex_en_num_single_dot = r"^\d+\.(?!\d)\s{0,4}(?=\S)" # 1.xxx
- regex_en_num_space = r"^[0-9]{1,2}\s{1,8}(?=\S)" # 1 xxx
- # ========== Chinese Numbering ==========
- regex_cn_num_dun = r"^[一二三四五六七八九十百千万]+、\s{0,4}(?=\S|$)"
- regex_cn_num_mix = (
- r"^[一二三四五六七八九十百千万]+(?:\s*\.[一二三四五六七八九十百千万\d]+)+"
- )
- regex_cn_num_plain = r"^[一二三四五六七八九十百千万]+(?=\s|$)"
- # ========== English Bracketing ==========
- regex_en_brac_paren = r"^[\(\（]\s*\d+(?:\.\d+)*(?!\.0)\s*[\)\）]"
- regex_en_brac_right = r"^\d+(?:\.\d+)*(?!\.0)\s*[\)\）]"
- # ========== Chinese Bracketing ==========
- regex_cn_brac_paren = r"^[\(\（]\s*[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]"
- regex_cn_brac_right = r"^[一二三四五六七八九十百千万]+(?:\.[一二三四五六七八九十百千万\d]+)*\s*[\)\）]"
- # ========== Chinese Special ==========
- regex_cn_special = r"^第[一二三四五六七八九十百千万\d]+(?:\.[一二三四五六七八九十百千万\d]+)*(章|节|条|部分|款|目|项|编|篇|卷|辑)?(?=$|\s|[A-Za-z0-9\u4e00-\u9fa5])"
- # ========== English Letter Numbering ==========
- regex_letter_dot = r"^[A-Za-z](?:\.\d+)*[\.、](?=\s*\S)"
- regex_letter_brac_paren = r"^[\(\（]\s*[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]"
- regex_letter_brac_right = r"^[A-Za-z](?:\.\d+)*(?!\.0)\s*[\)\）]"
- # ========== Appendix ==========
- regex_appendix = r"^((附件|附录|附表|附图)|(?i:appendix))[\s_\-—]{0,4}(?:\[)?[一二三四五六七八九十A-Za-z\d]"
-
- pos_regex_conditions = [
- # English Numbering
- regex_en_num_dots,
- regex_en_num_dun,
- regex_en_num_single_dot,
- regex_en_num_space,
- regex_en_num_dots_dun,
- # Chinese Numbering
- regex_cn_num_dun,
- regex_cn_num_mix,
- regex_cn_num_plain,
- # English Bracketing
- regex_en_brac_paren,
- regex_en_brac_right,
- # Chinese Bracketing
- regex_cn_brac_paren,
- regex_cn_brac_right,
- # Chinese Special
- regex_cn_special,
- # English Letter Numbering
- regex_letter_dot,
- regex_letter_brac_paren,
- regex_letter_brac_right,
- # Appendix
- regex_appendix,
- ]
-
- pos_triggered_code = []
- reason_suffix_parts = []
-
- for idx, regex in enumerate(pos_regex_conditions):
- match = re.match(regex, text)
- if match:
- matched_text = match.group(0)
- symbols = ".-"
- count_ = sum(matched_text.count(s) for s in symbols) + 1
-
- # Special handling for Chinese chapter/section/item markers.
- if idx == CN_SPECIAL_IDX and return_detail:
- unit_match = re.search(
- r"(章|节|条|部分|款|目|项|编|篇|卷|辑)", matched_text
- )
- if unit_match:
- unit = unit_match.group(1)
- reason_suffix_parts.append(f"[CN:{unit}]")
- pos_triggered_code.append(count_)
- else:
- pos_triggered_code.append(0)
-
- if return_detail:
- detail_info = {
- "reason_suffix": (
- " ".join(reason_suffix_parts) if reason_suffix_parts else ""
- )
- }
- if detail_info["reason_suffix"]:
- detail_info["reason_suffix"] = " " + detail_info["reason_suffix"]
- return pos_triggered_code, detail_info
- return pos_triggered_code
-
-
-def remove_by_conditions(text, include_punc=False):
- neg_condition_num = r"^\d{3,}"
- neg_condition_zero = r"^0\.\d+[\u4e00-\u9fa5A-Za-z\S]*" # 0.2xxx
- neg_decimal_only = r"^\d*\.\d+$" # 0.2 .23
- neg_condition_http = (
- r"(?i)(^https?://\S+|^www\.\S+|^P\.S|^\b\d{0,2}\s*(?:a\.m|p\.m)\b)"
- )
- # LaTeX: wrapped ($..\cmd..$) OR bare commands (\times, \mathrm, etc.)
- neg_condition_latex = (
- r"(?:"
- r"\$[^$]*\\[A-Za-z]+(?:\s*\{[^{}]*\})?[^$]*\$" # wrapped: $...\cmd...$
- r"|"
- r"\\(?:times|div|cdot|pm|mp|leq|geq|neq|approx|equiv|sim|infty"
- r"|sum|prod|int|sqrt|frac|mathrm|mathbf|mathit|mathcal"
- r"|text(?:bf|it|rm)?|alpha|beta|gamma|delta|epsilon|theta"
- r"|lambda|mu|sigma|pi|omega|partial|nabla"
- r"|left|right|begin|end|overline|underline|hat|vec|tilde)\b"
- r")"
- )
- # Number immediately followed by measurement unit (e.g. 25.40mm, 100kPa)
- neg_condition_unit = (
- r"^\d+\.?\d*\s{0,2}"
- r"(?:mm|cm|km|nm|μm|inch(?:es)?|ft|yd|mi"
- r"|kg|mg|μg|lb|oz"
- r"|kPa|MPa|GPa|Pa|psi|bar"
- r"|°[CFK]"
- r"|Hz|kHz|MHz|GHz"
- r"|mol|mL|dL|dB|Nm|kN|MN|kW|MW|GW|hp|rpm|cc|cal|kcal)\b"
- )
- neg_condition_punc_mid = r"[。！；].+"
- neg_condition_punc_end = r"[.,;，。；]$"
-
- neg_conditions = [
- neg_condition_num,
- neg_condition_http,
- neg_condition_latex,
- neg_condition_zero,
- neg_decimal_only,
- neg_condition_punc_mid,
- neg_condition_unit,
- ]
-
- neg_triggered_code = []
- for regex in neg_conditions:
- match = re.search(regex, text)
- neg_triggered_code.append(1 if match else 0)
-
- if include_punc:
- match = re.search(neg_condition_punc_end, text)
- neg_triggered_code.append(1 if match else 0)
- else:
- neg_triggered_code.append(0)
-
- return neg_triggered_code
-
-
-def md_heading_match(line, as_is=True):
- """handle markdown headings, considering # < ! [...."""
- match = re.match(r"^\s*(#+)\s*(.*)$", line)
- if match:
- level = len(match.group(1)) # count the number of '#'
- if as_is: # determine if remove the '#'
- return line, level
- else:
- return line.lstrip("#").strip(), level
- else:
- return line, -1
-
-
-def filter_md_headings(md_lines, num_pos=17, num_neg=7, layout_json_path=None):
- """filter candidate headings for .md
-
- Args:
- md_lines: list of markdown lines
- num_pos: number of positive conditions
- num_neg: number of negative conditions
- layout_json_path: optional path to layout.json for META features (size ranking)
- """
- # Create MetadataContext if layout_json_path is provided
- meta_ctx = None
- if layout_json_path:
- try:
- from .metadata_extractor import MetadataContext
-
- meta_ctx = MetadataContext(md_lines, layout_json_path)
- except Exception as e:
- logger.warning(f"Failed to create MetadataContext: {e}")
-
- raw_candidates = []
- for i, line in enumerate(md_lines):
- line = line.strip()
- if not line:
- continue
-
- if (
- ("" in line) # annotation line
- or line.startswith("|") # table line
- or line.startswith("")
- or "![" in line
- and "](" in line # image line
- ):
- est_lvl = -1
- zero_pos_code = [0] * num_pos
- zero_neg_code = [0] * num_neg
- str_lvl = f"POS {zero_pos_code} NEG {zero_neg_code}"
- if meta_ctx:
- str_lvl += " META [0, 0, 0]"
- line = "Figure/Image"
- else:
- line_clean, hash_lvl = md_heading_match(
- line, as_is=False
- ) # detect "#" in .md lines
-
- # NEW: detect and strip full-line bold markers (e.g. **3.4 Title** -> 3.4 Title)
- from .metadata_extractor import detect_and_strip_md_bold
-
- line_clean_stripped, is_full_bold = detect_and_strip_md_bold(line_clean)
-
- # Use stripped text for POS/NEG analysis (fixes '**3.4' -> '3.4' issue)
- pos_code, detail_info = judge_by_conditions(
- line_clean_stripped, return_detail=True
- )
- neg_code = remove_by_conditions(line_clean_stripped)
-
- if any(x > 0 for x in neg_code):
- code_lvl = -1
- code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}"
-
- elif any(x > 0 for x in pos_code) and all(x == 0 for x in neg_code):
- code_lvl = get_max_lvl(str(pos_code))
- code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}"
-
- else:
- code_lvl = -1
- code_str = f"POS {pos_code} NEG {neg_code}"
-
- # Add META suffix with bold dimension
- if meta_ctx:
- size_rank, occurrence = meta_ctx.get_meta_for_line(line_clean)
- is_bold_int = 1 if is_full_bold else 0
- code_str += meta_ctx.format_meta_suffix(
- size_rank, occurrence, is_bold_int
- )
- else:
- # Even without layout.json, output bold info in META
- if is_full_bold:
- code_str += " META [0, 0, 1]"
-
- if hash_lvl <= 0:
- est_lvl = code_lvl
- str_lvl = code_str
- else:
- if isinstance(code_lvl, int):
- est_lvl = max(
- hash_lvl, code_lvl
- ) # current miner tend to produce fewer #s
- else:
- est_lvl = code_lvl # code_lvl could be not sure
- str_lvl = f"{hash_lvl}# AND {code_str}"
- raw_candidates.append((i, line, est_lvl, str_lvl))
-
- preds_df = pd.DataFrame(
- raw_candidates, columns=["id", "heading", "level", "reason"], index=None
- )
- return preds_df
-
-
-def filter_doc_headings(titles_material, enable_regx=True, enable_style_check=False):
- """filter candidate headings for docx"""
-
- def find_docstyle(para_):
- try:
- style_name = para_.style.name
- except Exception:
- style_name = "normal"
- if style_name.startswith("Heading") or style_name.startswith("标题"):
- try:
- outline_level = int(style_name.split(" ")[1])
- except Exception:
- outline_level = -2 # "Not Sure" sentinel
- return outline_level
- else:
- return None
-
- def find_otsetting(para_):
- ppr = para_._element.find(qn("w:pPr"))
- if ppr is not None:
- plvl = ppr.find(qn("w:outlineLvl"))
- else:
- return None
-
- if plvl is not None:
- outline_level = int(plvl.get(qn("w:val"))) + 1
- return outline_level
- else:
- return None
-
- def find_bold(para_):
- if para_.runs and all(run.bold for run in para_.runs if run.text.strip()):
- return True
- else:
- return None
-
- raw_candidates = []
- logger.debug(
- "Filtering docx heading candidates... total_items={}", len(titles_material)
- )
- for ele_id, para, text in titles_material:
- str_lvl = ""
- est_lvl = None
- style_lvl = find_docstyle(para)
- setting_lvl = find_otsetting(para)
-
- # 1. check .docx style settings
- if style_lvl is not None:
- est_lvl = style_lvl
- str_lvl = f"style-{style_lvl}"
-
- # 2. check .docx paragraph numbering settings
- elif setting_lvl is not None:
- est_lvl = setting_lvl
- str_lvl = f"outline-{setting_lvl}"
-
- # 3. detect bold (unconditionally, encode as META dimension)
- is_bold = 1 if find_bold(para) else 0
-
- # 4. proceed condition judge
- if enable_regx:
- pos_code, detail_info = judge_by_conditions(text, return_detail=True)
- neg_code = remove_by_conditions(text)
-
- if any(x > 0 for x in neg_code):
- code_lvl = -1
- code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}"
- elif any(x > 0 for x in pos_code) and all(x == 0 for x in neg_code):
- code_lvl = get_max_lvl(str(pos_code))
- code_str = f"POS {pos_code}{detail_info.get('reason_suffix', '')} NEG {neg_code}"
- else:
- code_lvl = -1
- code_str = f"POS {pos_code} NEG {neg_code}"
-
- # Append bold as META dimension (DOCX has no layout.json, so only bold)
- if is_bold:
- code_str += f" META [0, 0, {is_bold}]"
-
- if est_lvl is None:
- est_lvl = code_lvl
- str_lvl = code_str
- else:
- str_lvl = f"{str_lvl} AND {code_str}"
- raw_candidates.append((ele_id, text, est_lvl, str_lvl))
-
- preds_df = pd.DataFrame(
- raw_candidates, columns=["id", "heading", "level", "reason"], index=None
- )
-
- # initial merge isolated and short texts
- preds_df = postprocess_headings(preds_df, task="merge_continuous")
- preds_df = postprocess_headings(preds_df, task="merge_short")
- return preds_df
+ return remove_isolated_heading_nodes(tree)
def format_toc_context_for_llm(toc_context) -> str:
@@ -1437,11 +384,11 @@ def pred_titles(
)
if doc_type == "pptx":
- raw_preds = filter_md_headings(infos)
+ raw_preds = filter_markdown_headings(infos)
elif doc_type == "md":
- raw_preds = filter_md_headings(infos, layout_json_path=layout_json_path)
+ raw_preds = filter_markdown_headings(infos, layout_json_path=layout_json_path)
elif doc_type == "docx":
- raw_preds = filter_doc_headings(infos, enable_regx)
+ raw_preds = filter_document_headings(infos, enable_regex=enable_regx)
else:
raw_preds = pd.DataFrame(columns=["id", "heading", "level", "reason"])
@@ -1709,406 +656,23 @@ def est_hierarchies_llm(
raw_preds: raw data
prompt_limt: prompt character limit
toc_hierarchies: TOC hierarchies
- max_len: maximum heading length (passed through to heading_tb_transfer)
+ max_len: maximum heading length for executor chunk preparation
max_depth: maximum hierarchy depth
model_name: LLM model name
output_dir: output directory, used to save intermediate results CSV
csv_suffix: suffix for intermediate CSV filenames
"""
model_name = _resolve_hierarchy_model_name(model_name)
- if len(raw_preds) == 0:
- return pd.DataFrame(columns=["id", "heading", "level", "reason"])
-
- compact_enabled = os.environ.get(
- "KB_LAYOUT_LLM_COMPACT_INPUT", "true"
- ).strip().lower() in ("true", "1", "yes", "on")
- preds_for_llm = _compact_for_llm(raw_preds) if compact_enabled else raw_preds.copy()
- if compact_enabled:
- placeholder_count = int(preds_for_llm["reason"].eq(PLACEHOLDER_REASON).sum())
- logger.info(
- f"smart parse => compact input: {len(raw_preds)} → {len(preds_for_llm)} rows "
- f"({placeholder_count} placeholder groups)"
- )
-
- # Short-circuit: if there are no heading candidates to judge (all rows were
- # collapsed into placeholders, or raw_preds contains only level==-1 rows
- # with compaction disabled), skip the LLM entirely and return raw_preds
- # with all levels set to -1.
- non_placeholder = (
- preds_for_llm[preds_for_llm["reason"].astype(str) != PLACEHOLDER_REASON]
- if compact_enabled
- else preds_for_llm
- )
- if len(non_placeholder) == 0:
- logger.info(
- "smart parse => no heading candidates, skipping LLM hierarchy detection"
- )
- fallback = raw_preds.copy()[["id", "heading", "level", "reason"]]
- fallback["level"] = -1
- return fallback.sort_values("id").reset_index(drop=True)
-
- level_dfs, _raw_headings = heading_tb_transfer(
- preds_for_llm, threshold=prompt_limt, max_start=max_len, max_end=5
+ return execute_llm_heading_hierarchy(
+ raw_preds=raw_preds,
+ prompt_limt=prompt_limt,
+ hierarchy_judge=hiearchy_llm,
+ fallback_hierarchy=est_hierarchies_naive,
+ save_intermediate_csv=save_intermediate_csv,
+ toc_hierarchies=toc_hierarchies,
+ max_len=max_len,
+ max_depth=max_depth,
+ model_name=model_name,
+ output_dir=output_dir,
+ csv_suffix=csv_suffix,
)
- chunk_sizes = [len(d) for d in level_dfs]
- logger.info(
- f"smart parse => {len(level_dfs)} chunk(s) | rows per chunk: {chunk_sizes} | "
- f"threshold={prompt_limt} | max_start={max_len}"
- )
-
- # Pick the first chunk that actually contains heading candidates. When
- # compaction is enabled a small prompt_limt may push a placeholder-only
- # chunk to index 0 — using it would waste an LLM call and produce an empty
- # mapping. Placeholder chunks that precede the chosen one contribute no
- # reason-code signal (their ids map to -1 anyway).
- basic_idx = 0
- for idx, chunk in enumerate(level_dfs):
- if (chunk["reason"].astype(str) != PLACEHOLDER_REASON).any():
- basic_idx = idx
- break
- basic_df = level_dfs[basic_idx]
- if basic_idx != 0:
- logger.info(
- f"smart parse => promoted chunk {basic_idx} as basic_df "
- f"(chunks 0..{basic_idx - 1} contain only placeholders)"
- )
- full_preds = None
- try:
- with stage_timer(
- "heading.hierarchy_llm",
- chunk_count=len(level_dfs),
- base_chunk_rows=len(basic_df),
- compact_enabled=compact_enabled,
- source_row_count=len(raw_preds),
- model_name=model_name,
- ):
- logger.debug("🚀 smart parse => interpreting hierarchy patterns...")
- df4llm = basic_df.drop(columns=["reason"]).copy()
- from .metadata_extractor import clean_md_text_for_llm
-
- # Keep formatting signals in `reason` / preliminary `level`, but let the LLM
- # judge hierarchy from the semantic heading text instead of raw markdown markers.
- df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm)
- logger.debug(f"DataFrame transformation completed, rows: {len(df4llm)}")
-
- layout_res = hiearchy_llm(
- df4llm, model_name, max_depth, toc_hierarchies, task="eval-headings"
- )
-
- # Build base_preds by aligning on basic_df["id"]: we always have one
- # row per chunk-0 row in the rendered output, regardless of how many
- # entries the LLM actually returned. Missing ids -> level=-1.
- layout_level_by_id = {}
- if isinstance(layout_res, list):
- for item in layout_res:
- if isinstance(item, dict) and "id" in item and "level" in item:
- layout_level_by_id[item["id"]] = item["level"]
-
- def _level_for(rid):
- if rid in layout_level_by_id:
- return layout_level_by_id[rid]
- try:
- return layout_level_by_id.get(int(rid), -1)
- except (TypeError, ValueError):
- return -1
-
- base_preds = (
- basic_df[["id", "heading", "reason"]].copy().reset_index(drop=True)
- )
- base_preds.insert(2, "level", base_preds["id"].map(_level_for))
-
- # Save base_preds as preds_3 (reflects what the LLM saw, compact or not)
- save_intermediate_csv(
- base_preds, output_dir, f"preds_3_llm_base{csv_suffix}"
- )
-
- # Collect {int_id -> level} from chunk-0 LLM output. Placeholder rows
- # have non-integer ids and are skipped.
- llm_levels = {}
- for _, row in base_preds.iterrows():
- rid = row["id"]
- if isinstance(rid, bool):
- continue
- if isinstance(rid, int):
- llm_levels[rid] = row["level"]
-
- if len(level_dfs) > 1:
- # Multi-chunk: build reason-code mapping from chunk-0 candidates and
- # apply it to chunks 1..N to infer levels for headings beyond chunk 0.
- # Placeholder rows are excluded from both the mapping source and the
- # per-chunk application — they always map to level=-1 in the final df.
- placeholder_mask_base = base_preds["reason"].eq(PLACEHOLDER_REASON)
- figure_mask_base = base_preds["heading"].eq("Figure/Image")
- exclude_mask_base = placeholder_mask_base | figure_mask_base
- base_preds_for_mapping = base_preds[~exclude_mask_base].copy()
- base_origin_for_mapping = basic_df.loc[
- ~exclude_mask_base.values, "level"
- ].tolist()
-
- base_preds_for_mapping, lvl_mapping = build_level_mapping(
- base_preds_for_mapping, base_origin_for_mapping, mode="freq"
- )
- logger.debug(
- f"mapping development finished: {len(lvl_mapping)} rules "
- f"(placeholders and Figure/Image excluded)"
- )
-
- logger.debug(
- f"mapping dataframe to levels across {len(level_dfs)} chunks..."
- )
- lvl_mapping = handle_unseen_codes(
- preds_for_llm, level_dfs, lvl_mapping, output_dir
- )
-
- for level_df in level_dfs:
- placeholder_mask_chunk = level_df["reason"].eq(PLACEHOLDER_REASON)
- figure_mask_chunk = level_df["heading"].eq("Figure/Image")
- exclude_mask_chunk = placeholder_mask_chunk | figure_mask_chunk
- non_excluded = level_df[~exclude_mask_chunk].copy()
- if not non_excluded.empty:
- non_excluded = execute_level_mapping(non_excluded, lvl_mapping)
- for _, row in non_excluded.iterrows():
- rid = row["id"]
- if isinstance(rid, bool):
- continue
- if isinstance(rid, int):
- # Mapping may override chunk-0 LLM decisions when
- # two rows share the same reason; accept that (the
- # mapping is by construction the "representative"
- # level for each reason-code).
- llm_levels[rid] = row["level"]
- logger.info(
- f"multi-chunk mapping produced {len(llm_levels)} id→level entries"
- )
- else:
- logger.info(
- "single chunk — skipping reason-code mapping, using LLM output directly"
- )
-
- # Expand back onto the original raw_preds: heading candidates take
- # the LLM/mapping-assigned level; everything else is body text (-1).
- full_preds = raw_preds.copy()
- full_preds = full_preds[["id", "heading", "level", "reason"]]
-
- def _resolve_level(rid):
- try:
- int_id = int(rid)
- except (TypeError, ValueError):
- return -1
- lvl = llm_levels.get(int_id, -1)
- try:
- return int(lvl)
- except (TypeError, ValueError):
- return -1
-
- full_preds["level"] = full_preds["id"].map(_resolve_level).astype(int)
-
- save_intermediate_csv(
- full_preds, output_dir, f"preds_4_llm_final{csv_suffix}"
- )
-
- except Exception as e:
- logger.warning(f"LLM-based parsing fails due to {e}, using non-llm pipeline...")
- full_preds = est_hierarchies_naive(raw_preds.copy())
- return full_preds
-
-
-def collapse_recursive(df, task, indices, merge_th=3, checked_pairs=None, depth=0):
- """recursive collapse"""
- if checked_pairs is None:
- checked_pairs = set()
-
- if len(indices) < 2:
- return
-
- for k in range(len(indices) - 1):
- i, j = indices[k], indices[k + 1]
- if (i, j) in checked_pairs:
- continue
- checked_pairs.add((i, j))
-
- between = df.loc[i + 1 : j - 1]
- i_txt = df.at[i, "heading"].strip()
- j_txt = df.at[j, "heading"].strip()
-
- if task == "merge_short" and len(between) > 0:
- between_lens = [count_cn_en(c) for c in between["heading"].tolist()]
- between_lvls = [bl for bl in between["level"].tolist()]
- i_half_len = int(count_cn_en(i_txt) / 2)
- too_short = sum(between_lens) <= merge_th or sum(between_lens) < i_half_len
-
- if too_short and all(
- bl == -1 for bl in between_lvls
- ): # only non-headings can be merged
- logger.debug(
- f"⚠️ too short between {i}=>{i_txt[:15]} and {j}=>{j_txt[:15]} => merge to {i}"
- )
- between_txts = [
- str(r["heading"]).strip()
- for _, r in between.iterrows()
- if isinstance(r.get("heading"), str) and r["heading"].strip()
- ]
-
- if between_txts:
- joined_txt = "\n".join(between_txts)
- df.at[i, "heading"] = f"{i_txt} {joined_txt}"
-
- for idx in between.index:
- df.at[idx, "level"] = -1
- df.at[idx, "reason"] = f"Merged into {i}"
- logger.debug(f"\tmerged texts: {joined_txt[:50]}...")
-
- elif task == "collapse" and len(between) == 0:
- logger.debug(
- f"⚠️ Empty between i={i_txt[:15]}, j={j_txt[:15]} => set i.level=-1, j.level=Not Sure"
- )
- df.at[i, "level"] = -2 # "Not Sure" sentinel (int-safe)
- df.at[j, "level"] = -2 # "Not Sure" sentinel (int-safe)
-
- # ========== get subgroups for recursive tasks ==========
- sub_between = between[between["level"] != -1]
- code2sub = defaultdict(list)
- for idx, row in sub_between.iterrows():
- level = row["level"]
- reason = row["reason"]
- if level != -1:
- code2sub[(level, reason)].append(idx)
-
- for _, sub_indices in code2sub.items():
- collapse_recursive(
- df, task, sub_indices, merge_th, checked_pairs, depth + 1
- )
-
-
-def postprocess_headings(df, task, max_depth=-1):
- """postprocess headings"""
- if task == "judge_negs":
- for i, row in df.iterrows():
- neg_code = remove_by_conditions(row["heading"], include_punc=True)
- if any(x > 0 for x in neg_code):
- current_code = str(df.loc[i, "reason"])
-
- neg_match = re.search(r"(.*NEG\s*)\[[^\]]*\](.*)", current_code)
- if neg_match:
- update_code = f"{neg_match.group(1)}{neg_code}{neg_match.group(2)}"
- else:
- update_code = f"{current_code} NEG {neg_code}"
-
- df.loc[i, "level"] = -1
- df.loc[i, "reason"] = update_code
- return df
-
- elif task == "merge_continuous":
- denoised_rows = []
- punc_pattern = re.compile(r'[.,!?;:，。！？；：）】〕｝〉》’”"]$')
-
- i = 0
- while i < len(df):
- row = df.iloc[i]
- current_content = str(row["heading"]).strip()
- current_level = row["level"]
-
- j = i + 1
- while j < len(df):
- next_row = df.iloc[j]
- next_content = str(next_row["heading"]).strip()
- next_level = next_row["level"]
-
- # Skip merge if ID is not continuous (indicates table/image was skipped in between)
- expected_id = row["id"] + (j - i)
- if next_row["id"] != expected_id:
- break
-
- # both current and next rows are not heading & current row has no punctuation -> merge
- current_not_punc = not punc_pattern.search(current_content[-2:])
- if (current_level == -1 and next_level == -1) and current_not_punc:
- current_content += " " + next_content
- j += 1
- else:
- break
-
- merge_row = row.copy()
- merge_row["heading"] = current_content
- denoised_rows.append(tuple(merge_row))
- i = j
- return pd.DataFrame(denoised_rows, columns=["id", "heading", "level", "reason"])
-
- elif task == "merge_short" or task == "collapse":
- group2indices = defaultdict(list)
- for idx, row in df.iterrows():
- level = row["level"]
- reason = row["reason"]
- if level != -1:
- group2indices[(level, reason)].append(idx)
-
- checked_pairs = set()
- for _, indices in group2indices.items():
- collapse_recursive(
- df, task, indices, merge_th=3, checked_pairs=checked_pairs, depth=0
- )
-
- if task == "merge_short":
- drop_between = df.index[
- df["reason"].astype(str).str.startswith("Merged into", na=False)
- ].tolist()
- if drop_between:
- logger.debug(
- f"🛠️ Delete rows labeled as merged into, total {len(drop_between)} rows"
- )
- df.drop(drop_between, inplace=True)
- df.reset_index(drop=True, inplace=True)
- return df
-
- else:
- return None
-
-
-# def parse_outline_hier(markdown_text):
-# lines = markdown_text.strip().splitlines()
-# stack = []
-# root = []
-# for line in lines:
-# line = line.replace('markdown', '') # handle possible unexpected outputs
-# if not line.strip():
-# continue
-
-# stripped = line.lstrip()
-# indent = len(line) - len(stripped)
-# match = re.match(r"[-*+] (.+)", stripped)
-# if not match:
-# continue
-
-# title = match.group(1).strip()
-# node = {"chapter": title, "children": [], 'serial': 1}
-# level = indent // 2 # Two spaces per level, adjustable if needed.
-# if level == 0:
-# node['serial'] = len(root)+1
-# root.append(node)
-# stack = [(level, node)]
-# else:
-# while stack and stack[-1][0] >= level:
-# stack.pop()
-# if stack:
-# parent = stack[-1][1]
-# node['serial'] = len(parent['children']) + 1
-# parent["children"].append(node)
-# stack.append((level, node))
-# return root
-
-
-# def outline_to_markdown(nodes, level=0, path=""):
-# rows = []
-# def traverse(node_list, level, path_prefix):
-# for node in node_list:
-# split_char = settings.SPLIT_CHAR or "/"
-# current_path = f"{path_prefix} {split_char} {node['chapter']}" if path_prefix else node['chapter']
-# rows.append({
-# "path": current_path,
-# "title": node["chapter"],
-# "thoughts": node.get("thoughts", "").strip(),
-# "level": level
-# })
-# if node.get("children"):
-# traverse(node["children"], level + 1, current_path)
-# traverse(nodes, level, path)
-# return pd.DataFrame(rows)
diff --git a/apps/worker/app/services/document_parser/markdown_deferred_summary.py b/apps/worker/app/services/document_parser/markdown_deferred_summary.py
new file mode 100644
index 000000000..270303e90
--- /dev/null
+++ b/apps/worker/app/services/document_parser/markdown_deferred_summary.py
@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from typing import Literal, TypeGuard
+
+import gevent
+from app.services.document_parser.markdown_deferred_task import (
+ ImageDeferredSummaryTask,
+ MarkdownDeferredSummaryTask,
+ TableDeferredSummaryTask,
+ TextDeferredSummaryTask,
+)
+from app.services.document_parser.image_parser import _get_vision_client, ask_image
+from app.services.document_parser.stage_profiler import stage_timer
+from app.services.document_parser.table_text_parser import sanitize_table_name_from_header
+from app.services.document_parser.txt_parser import (
+ extract_title_keywords_summary,
+ split_title_summary,
+)
+from gevent.pool import Pool as GeventPool
+from loguru import logger
+
+from shared.core.config import settings
+from shared.utils.chunk_refs import build_chunk_ref
+from shared.utils.file_utils import path_handle
+
+DeferredResult = (
+ tuple[
+ int,
+ Literal["image", "table", "text"],
+ tuple[str | None, str | None] | tuple[str, str, str] | tuple[str, str],
+ ]
+)
+ImageSummaryResult = tuple[str | None, str | None]
+TableSummaryResult = tuple[str, str, str]
+TextSummaryResult = tuple[str, str]
+
+
+@dataclass(frozen=True)
+class MarkdownDeferredSummaryInput:
+ rows: list[list[str | int]]
+ tasks: list[MarkdownDeferredSummaryTask]
+ output_dir: str
+ summary_len: int = 1500
+
+
+def apply_markdown_deferred_summaries(
+ deferred_input: MarkdownDeferredSummaryInput,
+) -> None:
+ if not deferred_input.tasks:
+ return
+
+ image_task_count = sum(
+ 1 for task in deferred_input.tasks if isinstance(task, ImageDeferredSummaryTask)
+ )
+ table_task_count = sum(
+ 1 for task in deferred_input.tasks if isinstance(task, TableDeferredSummaryTask)
+ )
+ text_task_count = sum(
+ 1 for task in deferred_input.tasks if isinstance(task, TextDeferredSummaryTask)
+ )
+ logger.info(
+ f"Running {len(deferred_input.tasks)} deferred summary LLM calls in parallel"
+ )
+ max_concurrent = getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 8)
+
+ with stage_timer(
+ "md.deferred_summaries",
+ total_tasks=len(deferred_input.tasks),
+ image_tasks=image_task_count,
+ table_tasks=table_task_count,
+ text_tasks=text_task_count,
+ max_concurrent=min(max_concurrent, len(deferred_input.tasks)),
+ ):
+ results = _run_deferred_summary_tasks(deferred_input, max_concurrent)
+ _apply_deferred_summary_results(deferred_input, results)
+
+ logger.info(f"Completed {len(deferred_input.tasks)} deferred summary LLM calls")
+
+
+def replace_chunk_ref_in_rows(
+ rows: list[list[str | int]], old_path: str, new_path: str
+) -> None:
+ old_ref = build_chunk_ref(old_path)
+ new_ref = build_chunk_ref(new_path)
+ if not old_ref or old_ref == new_ref:
+ return
+
+ for row in rows:
+ if len(row) > 0 and isinstance(row[0], str):
+ row[0] = row[0].replace(old_ref, new_ref)
+ if len(row) > 1 and row[1] == old_path:
+ row[1] = new_path
+ if len(row) > 2 and isinstance(row[2], str):
+ row[2] = row[2].replace(old_ref, new_ref)
+ if len(row) > 8 and isinstance(row[8], str):
+ row[8] = row[8].replace(old_ref, new_ref)
+
+
+def _run_deferred_summary_tasks(
+ deferred_input: MarkdownDeferredSummaryInput,
+ max_concurrent: int,
+) -> list[DeferredResult | None]:
+ pool = GeventPool(size=min(max_concurrent, len(deferred_input.tasks)))
+ greenlets = [
+ pool.spawn(_run_deferred_summary_task, task, deferred_input)
+ for task in deferred_input.tasks
+ ]
+ gevent.joinall(greenlets)
+ return [greenlet.value for greenlet in greenlets]
+
+
+def _run_deferred_summary_task(
+ task: MarkdownDeferredSummaryTask,
+ deferred_input: MarkdownDeferredSummaryInput,
+) -> DeferredResult | None:
+ try:
+ if isinstance(task, ImageDeferredSummaryTask):
+ client = _get_vision_client()
+ # TODO: Risk of missing text content if MinerU outputted a pure text image.
+ # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image.
+ llm_resp = ask_image(
+ client, deferred_input.output_dir, paths_=[task.relative_path]
+ )
+ if llm_resp:
+ img_title, img_summary = split_title_summary(llm_resp)
+ else:
+ img_title, img_summary = None, None
+ return task.row_index, "image", (img_title, img_summary)
+
+ if isinstance(task, TableDeferredSummaryTask):
+ title, keywords, summary = extract_title_keywords_summary(
+ task.table_html, max_keywords=3
+ )
+ return task.row_index, "table", (title, keywords, summary)
+
+ if isinstance(task, TextDeferredSummaryTask):
+ _, keywords, summary = extract_title_keywords_summary(
+ task.content,
+ max_keywords=3,
+ summary_len=deferred_input.summary_len,
+ )
+ return task.row_index, "text", (keywords, summary)
+ except Exception as exc:
+ logger.warning(
+ f"Deferred summary LLM call failed for idx={task.row_index}: {exc}"
+ )
+ return None
+
+ logger.warning(f"Unknown deferred markdown summary task type: {type(task).__name__}")
+ return None
+
+
+def _apply_deferred_summary_results(
+ deferred_input: MarkdownDeferredSummaryInput,
+ results: list[DeferredResult | None],
+) -> None:
+ deferred_by_index = {task.row_index: task for task in deferred_input.tasks}
+
+ for result in results:
+ if result is None:
+ continue
+
+ row_index, task_type, task_result = result
+ if task_type == "image":
+ if not _is_image_summary_result(task_result):
+ logger.warning(f"Invalid image deferred result for idx={row_index}")
+ continue
+ _apply_image_summary_result(
+ deferred_input.rows,
+ _get_image_task(deferred_by_index[row_index]),
+ row_index,
+ task_result,
+ )
+ elif task_type == "table":
+ if not _is_table_summary_result(task_result):
+ logger.warning(f"Invalid table deferred result for idx={row_index}")
+ continue
+ _apply_table_summary_result(
+ deferred_input.rows,
+ _get_table_task(deferred_by_index[row_index]),
+ row_index,
+ task_result,
+ )
+ elif task_type == "text":
+ if not _is_text_summary_result(task_result):
+ logger.warning(f"Invalid text deferred result for idx={row_index}")
+ continue
+ _apply_text_summary_result(deferred_input.rows, row_index, task_result)
+
+
+def _is_image_summary_result(result: object) -> TypeGuard[ImageSummaryResult]:
+ return (
+ isinstance(result, tuple)
+ and len(result) == 2
+ and all(isinstance(value, (str, type(None))) for value in result)
+ )
+
+
+def _is_table_summary_result(result: object) -> TypeGuard[TableSummaryResult]:
+ return (
+ isinstance(result, tuple)
+ and len(result) == 3
+ and all(isinstance(value, str) for value in result)
+ )
+
+
+def _is_text_summary_result(result: object) -> TypeGuard[TextSummaryResult]:
+ return (
+ isinstance(result, tuple)
+ and len(result) == 2
+ and all(isinstance(value, str) for value in result)
+ )
+
+
+def _get_image_task(task: MarkdownDeferredSummaryTask) -> ImageDeferredSummaryTask:
+ if isinstance(task, ImageDeferredSummaryTask):
+ return task
+ raise TypeError(f"Expected image deferred task, got {type(task).__name__}")
+
+
+def _get_table_task(task: MarkdownDeferredSummaryTask) -> TableDeferredSummaryTask:
+ if isinstance(task, TableDeferredSummaryTask):
+ return task
+ raise TypeError(f"Expected table deferred task, got {type(task).__name__}")
+
+
+def _apply_image_summary_result(
+ rows: list[list[str | int]],
+ original_task: ImageDeferredSummaryTask,
+ row_index: int,
+ result: ImageSummaryResult,
+) -> None:
+ img_title, img_summary = result
+ row = rows[row_index]
+ if img_summary:
+ image_index = str(row[5]).split("\n")[0] if row[5] else "image"
+ row[5] = f"{image_index}\n{img_summary}"
+
+ if not img_title:
+ return
+
+ image_dir = original_task.image_dir
+ old_img_name = original_task.image_name
+ image_suffix = original_task.image_suffix
+ safe_title = path_handle(str(img_title), mode="clean_single")
+ img_num_match = re.match(r"image-(\d+)", str(old_img_name))
+ img_num = (
+ img_num_match.group(1)
+ if img_num_match
+ else str(old_img_name).split("-")[1]
+ if "-" in str(old_img_name)
+ else "0"
+ )
+ new_img_name = path_handle(f"image-{img_num}-{safe_title}", mode="clean_single")
+ old_path = os.path.join(image_dir, f"{old_img_name}{image_suffix}")
+ new_path = os.path.join(image_dir, f"{new_img_name}{image_suffix}")
+ if old_path == new_path or not os.path.exists(old_path):
+ return
+
+ os.rename(old_path, new_path)
+ new_relative_path = f"images/{new_img_name}{image_suffix}"
+ replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path)
+ row[1] = new_relative_path
+
+
+def _apply_table_summary_result(
+ rows: list[list[str | int]],
+ original_task: TableDeferredSummaryTask,
+ row_index: int,
+ result: TableSummaryResult,
+) -> None:
+ title, keywords, summary = result
+ row = rows[row_index]
+ row[4] = keywords if isinstance(keywords, str) else ""
+ if summary:
+ table_index = str(row[5]) if "\n" not in str(row[5]) else str(row[5]).split("\n")[0]
+ row[5] = f"{table_index}\n{summary}"
+
+ if not title:
+ return
+
+ table_dir = original_task.table_dir
+ old_table_name = original_task.table_name
+ table_count = original_task.table_count
+ safe_title = sanitize_table_name_from_header(str(title))
+ new_table_name = path_handle(
+ f"table-{table_count} {safe_title}", mode="clean_single"
+ )
+ old_path = os.path.join(table_dir, f"{old_table_name}.html")
+ new_path = os.path.join(table_dir, f"{new_table_name}.html")
+ if old_path == new_path or not os.path.exists(old_path):
+ return
+
+ os.rename(old_path, new_path)
+ new_relative_path = f"tables/{new_table_name}.html"
+ replace_chunk_ref_in_rows(rows, str(row[1]), new_relative_path)
+ row[1] = new_relative_path
+
+
+def _apply_text_summary_result(
+ rows: list[list[str | int]], row_index: int, result: TextSummaryResult
+) -> None:
+ keywords, summary = result
+ rows[row_index][4] = keywords if isinstance(keywords, str) else ""
+ rows[row_index][5] = summary if isinstance(summary, str) else ""
diff --git a/apps/worker/app/services/document_parser/markdown_deferred_task.py b/apps/worker/app/services/document_parser/markdown_deferred_task.py
new file mode 100644
index 000000000..da935abaa
--- /dev/null
+++ b/apps/worker/app/services/document_parser/markdown_deferred_task.py
@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TypeAlias
+
+
+@dataclass(frozen=True)
+class ImageDeferredSummaryTask:
+ row_index: int
+ relative_path: str
+ image_dir: str
+ image_name: str
+ image_suffix: str
+
+
+@dataclass(frozen=True)
+class TableDeferredSummaryTask:
+ row_index: int
+ table_html: str
+ table_dir: str
+ table_name: str
+ table_count: int
+
+
+@dataclass(frozen=True)
+class TextDeferredSummaryTask:
+ row_index: int
+ content: str
+
+
+MarkdownDeferredSummaryTask: TypeAlias = (
+ ImageDeferredSummaryTask | TableDeferredSummaryTask | TextDeferredSummaryTask
+)
diff --git a/apps/worker/app/services/document_parser/markdown_image_asset.py b/apps/worker/app/services/document_parser/markdown_image_asset.py
new file mode 100644
index 000000000..10c93b5dc
--- /dev/null
+++ b/apps/worker/app/services/document_parser/markdown_image_asset.py
@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import cast
+
+from app.services.document_parser.identifiers import gen_str_codes
+from app.services.document_parser.image_parser import perceptual_hash
+from app.services.document_parser.inline_asset import build_image_asset_row
+from app.services.document_parser.markdown_deferred_task import (
+ ImageDeferredSummaryTask,
+ MarkdownDeferredSummaryTask,
+)
+from app.services.document_parser.markdown_parse_state import ParserRowValues
+from loguru import logger
+
+from shared.utils.chunk_refs import build_chunk_ref
+from shared.utils.file_utils import path_handle
+
+
+@dataclass(frozen=True)
+class MarkdownImageAsset:
+ content_item: str | None
+ row_values: ParserRowValues | None
+ cache_key: str | None
+ cache_entry: dict[str, str] | None
+ deferred_task: MarkdownDeferredSummaryTask | None
+ should_advance_image_count: bool
+
+
+@dataclass(frozen=True)
+class MarkdownImageAssetRequest:
+ output_dir: str
+ image_dir: str
+ image_path: str
+ image_name: str
+ image_count: int
+ last_context: str
+ image_summary: str | None
+ timestamp: str
+ current_page_number: int
+ seen_images: dict[str, dict[str, str]]
+ summary_image: bool
+ row_index: int
+
+
+def build_markdown_image_asset(
+ request: MarkdownImageAssetRequest,
+) -> MarkdownImageAsset:
+ image_suffix = os.path.splitext(request.image_path)[-1]
+ source_path = resolve_markdown_image_source_path(
+ request.output_dir,
+ request.image_path,
+ )
+ if source_path is None or not source_path.exists():
+ logger.warning(f"Image file not found, skipping rename: {request.image_path}")
+ return _empty_asset(should_advance_image_count=True)
+
+ with open(source_path, "rb") as image_file:
+ image_binary_hash = perceptual_hash(image_file.read())
+
+ if image_binary_hash in request.seen_images:
+ return _build_duplicate_image_asset(
+ source_path=source_path,
+ cache_entry=request.seen_images[image_binary_hash],
+ timestamp=request.timestamp,
+ current_page_number=request.current_page_number,
+ )
+
+ relative_image_path = f"images/{request.image_name}{image_suffix}"
+ target_image_path = os.path.join(
+ request.image_dir,
+ f"{request.image_name}{image_suffix}",
+ )
+ os.rename(source_path, target_image_path)
+
+ image_index = f"image-{request.image_count}"
+ effective_summary = request.image_summary or request.last_context or None
+ image_summary_field = (
+ f"{image_index}\n{effective_summary}" if effective_summary else image_index
+ )
+ image_content = _build_image_content(
+ relative_image_path=relative_image_path,
+ summary=effective_summary,
+ )
+ image_know_id = gen_str_codes(image_binary_hash)
+ row_values = _build_image_row_values(
+ content=image_content,
+ relative_path=relative_image_path,
+ summary=image_summary_field,
+ know_id=image_know_id,
+ timestamp=request.timestamp,
+ current_page_number=request.current_page_number,
+ )
+ cache_entry = {
+ "relative_img_path": relative_image_path,
+ "img_content": image_content,
+ "img_summary_field": image_summary_field,
+ "temp_uid": image_know_id,
+ }
+
+ deferred_task = None
+ if request.summary_image:
+ deferred_task = ImageDeferredSummaryTask(
+ row_index=request.row_index,
+ relative_path=relative_image_path,
+ image_dir=request.image_dir,
+ image_name=request.image_name,
+ image_suffix=image_suffix,
+ )
+
+ return MarkdownImageAsset(
+ content_item=image_content,
+ row_values=row_values,
+ cache_key=image_binary_hash,
+ cache_entry=cache_entry,
+ deferred_task=deferred_task,
+ should_advance_image_count=True,
+ )
+
+
+def build_markdown_image_name(*, image_count: int, last_context: str) -> str:
+ image_name_context = path_handle(last_context[:10], mode="clean_single")
+ return f"image-{str(image_count)}-{image_name_context}"
+
+
+def resolve_workspace_image_path(
+ candidate_path: Path, workspace_path: Path,
+) -> Path | None:
+ """Return the candidate only when it exists inside the current job workspace."""
+ resolved_path = candidate_path.resolve(strict=False)
+ try:
+ resolved_path.relative_to(workspace_path)
+ except ValueError:
+ return None
+ return resolved_path if resolved_path.exists() else None
+
+
+def resolve_markdown_image_source_path(output_dir: str, image_path: str) -> Path | None:
+ """Handle local absolute refs and container cwd-relative refs safely."""
+ if not image_path:
+ return None
+
+ workspace_path = Path(output_dir).resolve()
+ raw_path = Path(image_path).expanduser()
+ candidate_paths = (
+ [raw_path]
+ if raw_path.is_absolute()
+ else [
+ workspace_path / raw_path,
+ Path.cwd() / raw_path,
+ ]
+ )
+
+ for candidate_path in candidate_paths:
+ resolved_path = resolve_workspace_image_path(candidate_path, workspace_path)
+ if resolved_path is not None:
+ return resolved_path
+
+ return None
+
+
+def _build_duplicate_image_asset(
+ *,
+ source_path: Path,
+ cache_entry: dict[str, str],
+ timestamp: str,
+ current_page_number: int,
+) -> MarkdownImageAsset:
+ row_values = _build_image_row_values(
+ content=cache_entry["img_content"],
+ relative_path=cache_entry["relative_img_path"],
+ summary=cache_entry["img_summary_field"],
+ know_id=cache_entry["temp_uid"],
+ timestamp=timestamp,
+ current_page_number=current_page_number,
+ )
+ try:
+ source_path.unlink()
+ except OSError:
+ pass
+ logger.debug("Skipped duplicate image")
+ return MarkdownImageAsset(
+ content_item=cache_entry["img_content"],
+ row_values=row_values,
+ cache_key=None,
+ cache_entry=None,
+ deferred_task=None,
+ should_advance_image_count=False,
+ )
+
+
+def _build_image_content(*, relative_image_path: str, summary: str | None) -> str:
+ image_reference = build_chunk_ref(relative_image_path)
+ if summary:
+ return f"\n{summary}\n{image_reference}\n"
+ return f"\n{image_reference}\n"
+
+
+def _build_image_row_values(
+ *,
+ content: str,
+ relative_path: str,
+ summary: str,
+ know_id: str,
+ timestamp: str,
+ current_page_number: int,
+) -> ParserRowValues:
+ image_row = build_image_asset_row(
+ content=content,
+ relative_path=relative_path,
+ summary=summary,
+ know_id=know_id,
+ addtime=timestamp,
+ page_nums=str(current_page_number) if current_page_number > 0 else "",
+ )
+ return cast(ParserRowValues, image_row.to_list())
+
+
+def _empty_asset(*, should_advance_image_count: bool) -> MarkdownImageAsset:
+ return MarkdownImageAsset(
+ content_item=None,
+ row_values=None,
+ cache_key=None,
+ cache_entry=None,
+ deferred_task=None,
+ should_advance_image_count=should_advance_image_count,
+ )
diff --git a/apps/worker/app/services/document_parser/markdown_parse_state.py b/apps/worker/app/services/document_parser/markdown_parse_state.py
new file mode 100644
index 000000000..51f97b954
--- /dev/null
+++ b/apps/worker/app/services/document_parser/markdown_parse_state.py
@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import pandas as pd
+
+from app.services.document_parser.dataframe_helpers import process_dup_paths_df
+from app.services.document_parser.markdown_deferred_task import (
+ MarkdownDeferredSummaryTask,
+ TextDeferredSummaryTask,
+)
+from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder
+
+ParserRowValues = list[str | int]
+
+RowUpdater = Callable[
+ [list[ParserRowValues], list[str], str, dict[str, Any], str, str, int, bool],
+ list[ParserRowValues],
+]
+
+
+@dataclass
+class MarkdownParseState:
+ relative_root: str
+ split_char: str
+ llm_parameters: dict[str, Any]
+ timestamp: str
+ row_updater: RowUpdater
+ rows: list[ParserRowValues] = field(default_factory=list)
+ content_items: list[str] = field(default_factory=list)
+ path_stack: list[tuple[str, int]] = field(default_factory=list)
+ inner_paths: list[str] = field(default_factory=list)
+ error_line_numbers: list[int] = field(default_factory=list)
+ table_lines: list[str] = field(default_factory=list)
+ current_page_number: int = 0
+ chunk_pages: set[int] = field(default_factory=set)
+ base_level: int | None = None
+ path: str = ""
+ path_counter: dict[str, int] = field(default_factory=dict)
+ deferred_llm_tasks: list[MarkdownDeferredSummaryTask] = field(default_factory=list)
+ seen_images: dict[str, dict[str, str]] = field(default_factory=dict)
+ image_count: int = 1
+ table_count: int = 1
+
+ def __post_init__(self) -> None:
+ if not self.path:
+ self.path = self.relative_root
+
+ def record_page_marker(self, line: str) -> bool:
+ if "" not in line:
+ return False
+ if "page" not in line and "Slide number" not in line:
+ return False
+
+ page_match = re.search(r"page\s+(\d+)", line)
+ if page_match:
+ self.current_page_number = int(page_match.group(1))
+ else:
+ self.current_page_number += 1
+ self.chunk_pages.add(self.current_page_number)
+ return True
+
+ def flush_current_content(self) -> None:
+ page_numbers = self._format_chunk_pages()
+ self.rows = self.row_updater(
+ self.rows,
+ self.content_items,
+ self.path,
+ self.llm_parameters,
+ self.timestamp,
+ page_numbers,
+ 1500,
+ True,
+ )
+ self.content_items = []
+ self.chunk_pages = set()
+ if self.current_page_number > 0:
+ self.chunk_pages.add(self.current_page_number)
+
+ def flush_placeholder_chunk(self) -> None:
+ page_numbers = self._format_chunk_pages()
+ self.rows = self.row_updater(
+ self.rows,
+ [],
+ self.path,
+ self.llm_parameters,
+ self.timestamp,
+ page_numbers,
+ 1500,
+ True,
+ )
+
+ def enter_heading(self, heading: str, level: int) -> None:
+ if self.base_level is None:
+ self.base_level = level
+ elif level < self.base_level:
+ self.base_level = level
+
+ adjusted_level = level - self.base_level + 1
+ self.path_stack = [
+ (item_heading, item_level)
+ for item_heading, item_level in self.path_stack
+ if item_level < adjusted_level
+ ]
+
+ current_heading = (
+ heading.replace(self.split_char, "∕")
+ if self.split_char in heading
+ else heading
+ )
+ tentative_names = [item_heading for item_heading, _ in self.path_stack]
+ tentative_names.append(current_heading)
+ tentative_path_parts = [self.relative_root] if self.relative_root else []
+ tentative_path_parts.extend(tentative_names)
+ tentative_path = self.split_char.join(tentative_path_parts)
+
+ if tentative_path in self.path_counter:
+ self.path_counter[tentative_path] += 1
+ current_heading = f"{current_heading}_{self.path_counter[tentative_path]}"
+ else:
+ self.path_counter[tentative_path] = 1
+
+ self.path_stack.append((current_heading, adjusted_level))
+ heading_names = [item_heading for item_heading, _ in self.path_stack]
+ path_parts = [self.relative_root] if self.relative_root else []
+ path_parts.extend(heading_names)
+ self.inner_paths.append(self.split_char.join(heading_names))
+ self.path = self.split_char.join(path_parts)
+
+ def append_content_item(self, item: str) -> None:
+ self.content_items.append(item)
+
+ def append_plain_text(self, text: str) -> None:
+ self.content_items.append(text.strip())
+ if self.current_page_number > 0:
+ self.chunk_pages.add(self.current_page_number)
+
+ def append_row(self, row: ParserRowValues) -> None:
+ self.rows.append(row)
+
+ def schedule_deferred_task(self, task: MarkdownDeferredSummaryTask) -> None:
+ self.deferred_llm_tasks.append(task)
+
+ def collect_text_summary_tasks(self, summary_len: int) -> None:
+ if not self.llm_parameters.get("summary_txt"):
+ return
+
+ for index, entry in enumerate(self.rows):
+ marker = entry[2]
+ if isinstance(marker, str) and marker.strip().split("\n", 1)[0].lower() in {
+ "image",
+ "table",
+ }:
+ continue
+ content = str(entry[0])
+ if len(content) > summary_len and not entry[4] and not entry[5]:
+ self.deferred_llm_tasks.append(
+ TextDeferredSummaryTask(row_index=index, content=content)
+ )
+
+ def to_dataframe(self) -> pd.DataFrame:
+ rows_builder = ParsedRowsBuilder()
+ for row_values in self.rows:
+ rows_builder.append(
+ ParsedRow(
+ content=str(row_values[0]),
+ path=str(row_values[1]),
+ type=str(row_values[2]),
+ length=int(row_values[3]),
+ keywords=str(row_values[4]),
+ summary=str(row_values[5]),
+ know_id=str(row_values[6]),
+ tokens=str(row_values[7]),
+ connectto=str(row_values[8]),
+ addtime=str(row_values[9]),
+ page_nums=str(row_values[10]),
+ )
+ )
+ return process_dup_paths_df(rows_builder.to_dataframe())
+
+ def _format_chunk_pages(self) -> str:
+ return ",".join(str(page) for page in sorted(self.chunk_pages))
diff --git a/apps/worker/app/services/document_parser/markdown_table_asset.py b/apps/worker/app/services/document_parser/markdown_table_asset.py
new file mode 100644
index 000000000..a01e8ce97
--- /dev/null
+++ b/apps/worker/app/services/document_parser/markdown_table_asset.py
@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import cast
+
+from app.services.document_parser.html_parser import first_cols_rows_html
+from app.services.document_parser.identifiers import gen_str_codes
+from app.services.document_parser.inline_asset import build_table_asset_row
+from app.services.document_parser.markdown_deferred_task import (
+ MarkdownDeferredSummaryTask,
+ TableDeferredSummaryTask,
+)
+from app.services.document_parser.markdown_parse_state import ParserRowValues
+from app.services.document_parser.table_text_parser import sanitize_table_name_from_header
+
+from shared.utils.chunk_refs import build_chunk_ref
+from shared.utils.file_utils import path_handle
+
+
+@dataclass(frozen=True)
+class MarkdownTableAsset:
+ content_item: str
+ row_values: ParserRowValues
+ deferred_task: MarkdownDeferredSummaryTask | None
+ relative_path: str
+
+
+@dataclass(frozen=True)
+class MarkdownTableAssetRequest:
+ table_html: str
+ table_dir: str
+ table_count: int
+ timestamp: str
+ current_page_number: int
+ summary_table: bool
+ row_index: int
+
+
+def build_markdown_table_asset(
+ request: MarkdownTableAssetRequest,
+) -> MarkdownTableAsset:
+ first_row_text, _first_col_text = first_cols_rows_html(request.table_html)
+ table_index = f"table-{request.table_count}"
+
+ raw_table_name = (
+ sanitize_table_name_from_header(first_row_text) if first_row_text else ""
+ )
+ table_name = _sanitize_table_file_stem(
+ f"table-{str(request.table_count)} {raw_table_name}"
+ )
+ relative_table_path = f"tables/{table_name}.html"
+ table_ref = build_chunk_ref(relative_table_path)
+ table_content_item = f"\n{table_ref}\n"
+ table_path = os.path.join(request.table_dir, f"{table_name}.html")
+ _write_table_html(table_path=table_path, table_html=request.table_html)
+
+ table_row = build_table_asset_row(
+ content=request.table_html,
+ relative_path=relative_table_path,
+ summary=table_index,
+ keywords="",
+ know_id=gen_str_codes((request.table_html + str(request.table_count))),
+ addtime=request.timestamp,
+ page_nums=str(request.current_page_number)
+ if request.current_page_number > 0
+ else "",
+ )
+
+ deferred_task = None
+ if request.summary_table:
+ deferred_task = TableDeferredSummaryTask(
+ row_index=request.row_index,
+ table_html=request.table_html,
+ table_dir=request.table_dir,
+ table_name=table_name,
+ table_count=request.table_count - 1,
+ )
+
+ return MarkdownTableAsset(
+ content_item=table_content_item,
+ row_values=cast(ParserRowValues, table_row.to_list()),
+ deferred_task=deferred_task,
+ relative_path=relative_table_path,
+ )
+
+
+def _sanitize_table_file_stem(raw_name: str) -> str:
+ table_name = path_handle(raw_name, mode="clean_single")
+ if not isinstance(table_name, str) or not table_name:
+ raise ValueError(f"Failed to sanitize Markdown table name: {raw_name}")
+ return table_name
+
+
+def _write_table_html(*, table_path: str, table_html: str) -> None:
+ table_html_with_border = table_html.replace("", "").replace(
+ " Path | None:
- """Return the candidate only when it exists inside the current job workspace."""
- resolved_path = candidate_path.resolve(strict=False)
- try:
- resolved_path.relative_to(workspace_path)
- except ValueError:
- return None
- return resolved_path if resolved_path.exists() else None
-
-
-def resolve_markdown_image_source_path(output_dir: str, img_path: str) -> Path | None:
- """Handle local absolute refs and container cwd-relative refs safely."""
- if not img_path:
- return None
-
- workspace_path = Path(output_dir).resolve()
- raw_path = Path(img_path).expanduser()
- candidate_paths = (
- [raw_path]
- if raw_path.is_absolute()
- else [
- workspace_path / raw_path,
- Path.cwd() / raw_path,
- ]
- )
-
- for candidate_path in candidate_paths:
- resolved_path = resolve_workspace_image_path(candidate_path, workspace_path)
- if resolved_path is not None:
- return resolved_path
-
- return None
-
-
def find_surround_context(md_lines, lid):
def is_skip(line):
s = line.strip()
@@ -139,15 +105,17 @@ def eval_md_headings(
layout_json_path=None,
):
"""Evaluate markdown headings with optional TOC hierarchies context"""
- heading_preds = pred_titles(
- md_lines,
- source_type,
- toc_hierarchies=toc_hierarchies,
- enable_regx=True,
- smart_parse=smart_parse,
- model_name=model_name,
- output_dir=output_dir,
- layout_json_path=layout_json_path,
+ heading_preds = predict_heading_hierarchy(
+ HeadingHierarchyInput(
+ infos=md_lines,
+ doc_type=source_type,
+ toc_hierarchies=toc_hierarchies,
+ enable_regex=True,
+ smart_parse=smart_parse,
+ model_name=model_name,
+ output_dir=output_dir,
+ layout_json_path=layout_json_path,
+ )
)
if len(heading_preds) == 0:
@@ -184,24 +152,6 @@ def clean_md_table_lines(table_lines, start_line_num):
return cleaned_lines, error_lines
-def replace_chunk_ref_in_rows(df_list, old_path: str, new_path: str) -> None:
- """Rewrite readable chunk refs after deferred image/table renames."""
- old_ref = build_chunk_ref(old_path)
- new_ref = build_chunk_ref(new_path)
- if not old_ref or old_ref == new_ref:
- return
-
- for row in df_list:
- if len(row) > 0 and isinstance(row[0], str):
- row[0] = row[0].replace(old_ref, new_ref)
- if len(row) > 1 and row[1] == old_path:
- row[1] = new_path
- if len(row) > 2 and isinstance(row[2], str):
- row[2] = row[2].replace(old_ref, new_ref)
- if len(row) > 8 and isinstance(row[8], str):
- row[8] = row[8].replace(old_ref, new_ref)
-
-
def update_df_list(
df_list,
content_items,
@@ -241,19 +191,17 @@ def update_df_list(
)
df_list.append(
- [
- bottom_content,
- path,
- match_type,
- len(bottom_content),
- keywords,
- summary,
- know_id,
- bottom_tokens,
- "",
- time_stamp,
- page_nums,
- ]
+ ParsedRow(
+ content=bottom_content,
+ path=path,
+ type=match_type,
+ keywords=keywords,
+ summary=summary,
+ know_id=know_id,
+ tokens=bottom_tokens,
+ addtime=time_stamp,
+ page_nums=page_nums,
+ ).to_list()
)
return df_list
@@ -267,7 +215,7 @@ def parse_md(
relative_root=None,
):
if md_lines is None and file_path is not None:
- from shared.utils.CommonHelperSync import is_remote, load_file_bytes
+ from shared.utils.file_loading import is_remote, load_file_bytes
if is_remote(file_path):
file_bytes = load_file_bytes(file_path)
@@ -326,22 +274,13 @@ def parse_md(
# initialize vars
split_char = settings.SPLIT_CHAR or "/"
- df_list = []
- path_stack = [] # sxjg: uses (heading, level) tuples
- inner_paths = []
- error_line_numbers = []
- table_lines = []
- current_pg_num = 0
- chunk_pages = set() # collect all page numbers seen during current chunk
- base_level = None
- content_items = []
- # Use relative_root as initial path (not absolute output_dir)
- path = relative_root if relative_root else ""
- table_count = 1
- img_count = 1
- path_counter = {} # Track path occurrences for deduplication
- deferred_llm_tasks = [] # Collected during loop, executed in parallel after
- _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup
+ parser_state = MarkdownParseState(
+ relative_root=relative_root or "",
+ split_char=split_char,
+ llm_parameters=base_llm_paras,
+ timestamp=get_str_time(),
+ row_updater=update_df_list,
+ )
# Find layout.json path
layout_json_path = os.path.join(output_dir, "layout.json")
@@ -366,19 +305,10 @@ def parse_md(
layout_json_path=layout_json_path,
)
- time_stamp = get_str_time()
logger.debug("Parsing md data... total_lines={}", len(lines_with_heading))
for i, line in enumerate(lines_with_heading):
- if "" in line:
- if "page" in line or "Slide number" in line:
- # Parse actual page number from marker:
- pg_match = re.search(r"page\s+(\d+)", line)
- if pg_match:
- current_pg_num = int(pg_match.group(1))
- else:
- current_pg_num += 1 # fallback for Slide number or old format
- chunk_pages.add(current_pg_num)
- continue
+ if parser_state.record_page_marker(line):
+ continue
last_context = find_surround_context(
lines_with_heading, i
@@ -388,198 +318,58 @@ def parse_md(
if (
not current_heading_level == -1
): # indicate a new path should be evaluated or added
- if content_items: # record contents of the last path and reset content
- # Build page_nums from collected pages during this chunk
- chunk_page_str = (
- ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else ""
- )
- df_list = update_df_list(
- df_list,
- content_items,
- path,
- base_llm_paras,
- time_stamp,
- page_nums=chunk_page_str,
- skip_llm=True,
- )
- content_items = []
- chunk_pages = set() # reset for next chunk
- if current_pg_num > 0:
- chunk_pages.add(
- current_pg_num
- ) # carry current page into next chunk
- elif path and path != (relative_root or ""):
+ if parser_state.content_items:
+ parser_state.flush_current_content()
+ elif parser_state.path and parser_state.path != (relative_root or ""):
# Consecutive headings with no body text between them:
# Create a placeholder chunk so the previous heading's path
- chunk_page_str = (
- ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else ""
- )
- df_list = update_df_list(
- df_list,
- [],
- path,
- base_llm_paras,
- time_stamp,
- page_nums=chunk_page_str,
- skip_llm=True,
- )
+ parser_state.flush_placeholder_chunk()
- # update path based on path name and level
- if base_level is None:
- base_level = current_heading_level
- elif current_heading_level < base_level:
- base_level = current_heading_level
-
- adjusted_level = current_heading_level - base_level + 1
- path_stack = [(h, lvl) for h, lvl in path_stack if lvl < adjusted_level]
-
- # Build tentative path to check for duplicates
- # Sanitize heading: replace split_char in heading text to prevent path corruption
- current_heading = (
- current_heading.replace(split_char, "∕")
- if split_char in current_heading
- else current_heading
- )
- tentative_heading = current_heading
- tentative_names = [h for h, lvl in path_stack] + [tentative_heading]
- tentative_path_parts = [relative_root] if relative_root else []
- tentative_path_parts.extend(tentative_names)
- tentative_path = split_char.join(tentative_path_parts)
-
- # Deduplicate: if path already exists, add suffix
- if tentative_path in path_counter:
- path_counter[tentative_path] += 1
- suffix = path_counter[tentative_path]
- current_heading = (
- f"{current_heading}_{suffix}" # Modify heading with suffix
- )
- else:
- path_counter[tentative_path] = 1
-
- path_stack.append((current_heading, adjusted_level))
-
- # Extract pure heading names for path construction
- heading_names = [h for h, lvl in path_stack]
- # Use relative_root as prefix
- path_parts = [relative_root] if relative_root else []
- path_parts.extend(heading_names)
- inner_paths.append(split_char.join(heading_names))
- path = split_char.join(path_parts) # path with relative root
+ parser_state.enter_heading(current_heading, current_heading_level)
else: # no path change, remain in the same hierarchy
# a. handle lines containing images (LLM deferred to post-loop parallel batch)
- img_name_context = path_handle(last_context[:10], mode="clean_single")
- img_name = f"image-{str(img_count)}-{img_name_context}"
# Always skip inline LLM — vision calls are deferred to parallel batch
imgs = detect_summary_img_md(line, last_context, output_dir, mode=False)
+ image_name = build_markdown_image_name(
+ image_count=parser_state.image_count,
+ last_context=last_context,
+ )
- for img_path, img_title, img_summary in imgs:
- img_suffix = os.path.splitext(img_path)[-1]
- update_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}")
-
- # Check if source image file exists before renaming
- source_path = resolve_markdown_image_source_path(output_dir, img_path)
- if source_path is None or not source_path.exists():
- logger.warning(f"Image file not found, skipping rename: {img_path}")
- img_count += 1
- continue
-
- # Document-level dedup: perceptual hash for visual duplicates
- with open(source_path, "rb") as f:
- img_binary_hash = perceptual_hash(f.read())
-
- if img_binary_hash in _seen_images:
- cached = _seen_images[img_binary_hash]
- content_items.append(cached["img_content"])
- df_list.append(
- [
- cached["img_content"],
- cached["relative_img_path"],
- "image",
- len(cached["img_content"]),
- "",
- cached["img_summary_field"],
- cached["temp_uid"],
- "",
- "",
- time_stamp,
- str(current_pg_num) if current_pg_num > 0 else "",
- ]
- )
- logger.debug(
- f"Skipped duplicate image (hash={img_binary_hash[:12]}...)"
+ for img_path, _img_title, img_summary in imgs:
+ image_asset = build_markdown_image_asset(
+ MarkdownImageAssetRequest(
+ output_dir=output_dir,
+ image_dir=img_dir,
+ image_path=img_path,
+ image_name=image_name,
+ image_count=parser_state.image_count,
+ last_context=last_context,
+ image_summary=img_summary,
+ timestamp=parser_state.timestamp,
+ current_page_number=parser_state.current_page_number,
+ seen_images=parser_state.seen_images,
+ summary_image=bool(base_llm_paras["summary_image"]),
+ row_index=len(parser_state.rows),
)
- # Remove unused source file since we reuse the cached image
- try:
- source_path.unlink()
- except OSError:
- pass
- continue
-
- os.rename(source_path, update_img_path)
-
- # Image index (always present)
- image_index = f"image-{img_count}"
-
- # Fallback: LLM summary -> last_context -> None
- effective_summary = img_summary or last_context or None
-
- # Deterministic know_id: use image binary hash
- temp_uid = gen_str_codes(img_binary_hash)
- relative_img_path = f"images/{img_name}{img_suffix}"
- img_ref = build_chunk_ref(relative_img_path)
-
- # Build img_summary_field for df_list: image-n + optional summary
- if effective_summary:
- img_summary_field = f"{image_index}\n{effective_summary}"
- else:
- img_summary_field = image_index
-
- # Build image_ref for content: optional summary + image path ref
- if effective_summary:
- img_content = f"\n{effective_summary}\n{img_ref}\n"
- else:
- img_content = f"\n{img_ref}\n"
-
- content_items.append(img_content)
-
- df_list.append(
- [
- img_content,
- relative_img_path,
- "image",
- len(img_content),
- "",
- img_summary_field,
- temp_uid,
- "",
- "",
- time_stamp,
- str(current_pg_num) if current_pg_num > 0 else "",
- ]
)
-
- # Cache result for document-level dedup
- _seen_images[img_binary_hash] = {
- "relative_img_path": relative_img_path,
- "img_content": img_content,
- "img_summary_field": img_summary_field,
- "temp_uid": temp_uid,
- }
-
- if base_llm_paras["summary_image"]:
- # Store img_dir, img_name, img_suffix for post-loop rename (mirrors table deferred task)
- deferred_llm_tasks.append(
- (
- "image",
- len(df_list) - 1,
- relative_img_path,
- img_dir,
- img_name,
- img_suffix,
- )
+ if (
+ image_asset.content_item is not None
+ and image_asset.row_values is not None
+ ):
+ parser_state.append_content_item(image_asset.content_item)
+ parser_state.append_row(image_asset.row_values)
+ if (
+ image_asset.cache_key is not None
+ and image_asset.cache_entry is not None
+ ):
+ parser_state.seen_images[image_asset.cache_key] = (
+ image_asset.cache_entry
)
- img_count += 1
+ if image_asset.deferred_task is not None:
+ parser_state.schedule_deferred_task(image_asset.deferred_task)
+ if image_asset.should_advance_image_count:
+ parser_state.image_count += 1
# TODO for large and dense tables, such as "Epstein flight logs",
# integrate tabula-py as an independent extraction path to solve VLM hallucinations and misplacement
@@ -591,7 +381,7 @@ def parse_md(
tb_str = line
elif form == "md":
# For MD tables, accumulate lines until table ends
- table_lines.append(line)
+ parser_state.table_lines.append(line)
if i + 1 >= len(lines_with_heading):
tb_bool_next = False
else:
@@ -601,270 +391,54 @@ def parse_md(
if not tb_bool_next or i == len(lines_with_heading) - 1:
cleaned_table_lines, error_lines = clean_md_table_lines(
- table_lines, start_line_num=i
+ parser_state.table_lines, start_line_num=i
)
tb_str = "\n".join(cleaned_table_lines)
- error_line_numbers.extend(error_lines)
+ parser_state.error_line_numbers.extend(error_lines)
tb_str = extract_tables_by_forms(tb_str, form="md")
else:
continue # Keep accumulating MD table lines
else:
continue # Unknown form, skip
- # Extract first row and first column for fallback file naming only
- first_row_text, first_col_text = first_cols_rows_html(tb_str)
-
- # Table index (always present)
- table_index = f"table-{table_count}"
-
- # LLM title + keywords + summary deferred to post-loop parallel batch
- llm_title = None
- llm_summary = None
- tb_keywords = ""
-
- # Build tb_summary for df_list: table-n + optional LLM summary
- if llm_summary:
- tb_summary = f"{table_index}\n{llm_summary}"
- else:
- tb_summary = table_index
-
- raw_tb_name = (
- sanitize_table_name_from_header(first_row_text)
- if first_row_text
- else ""
- )
- # Use LLM title for filename when available, fallback to sanitized header
- effective_name = llm_title if llm_title else raw_tb_name
- tb_name = path_handle(
- f"table-{str(table_count)} {effective_name}", mode="clean_single"
- )
- temp_uid = gen_str_codes((tb_str + str(table_count)))
-
- relative_tb_path = f"tables/{tb_name}.html"
- tb_ref = build_chunk_ref(relative_tb_path)
-
- # Build table_ref for content: optional LLM summary + table path ref
- if llm_summary:
- content_items.append(f"\n{llm_summary}\n{tb_ref}\n")
- else:
- content_items.append(f"\n{tb_ref}\n")
- tb_path = os.path.join(tb_dir, f"{tb_name}.html")
- # Add border to HTML tables for consistent display
- tb_str_with_border = tb_str.replace(
- "", ""
- ).replace(" 0 else "",
- ]
- )
- if base_llm_paras["summary_table"]:
- deferred_llm_tasks.append(
- (
- "table",
- len(df_list) - 1,
- tb_str,
- tb_dir,
- tb_name,
- table_count - 1,
- )
+ table_asset = build_markdown_table_asset(
+ MarkdownTableAssetRequest(
+ table_html=tb_str,
+ table_dir=tb_dir,
+ table_count=parser_state.table_count,
+ timestamp=parser_state.timestamp,
+ current_page_number=parser_state.current_page_number,
+ summary_table=bool(base_llm_paras["summary_table"]),
+ row_index=len(parser_state.rows),
)
- table_lines = [] # Reset table_lines after storing the DataFrame
- table_count += 1
+ )
+ parser_state.append_content_item(table_asset.content_item)
+ parser_state.append_row(table_asset.row_values)
+ if table_asset.deferred_task is not None:
+ parser_state.schedule_deferred_task(table_asset.deferred_task)
+ parser_state.table_lines = []
+ parser_state.table_count += 1
# c. handle plain texts
if len(imgs) == 0 and not tb_bool:
- content_items.append(line.strip())
- if current_pg_num > 0:
- chunk_pages.add(current_pg_num) # track page for this content line
+ parser_state.append_plain_text(line)
- if content_items: # handle the remaining contents, append them to the last section
- chunk_page_str = (
- ",".join(str(p) for p in sorted(chunk_pages)) if chunk_pages else ""
- )
- df_list = update_df_list(
- df_list,
- content_items,
- path,
- base_llm_paras,
- time_stamp,
- page_nums=chunk_page_str,
- skip_llm=True,
- )
+ if parser_state.content_items:
+ parser_state.flush_current_content()
# Collect text chunk deferred tasks (entries needing summary/keywords)
summary_len = 1500
- if base_llm_paras.get("summary_txt"):
- for idx, entry in enumerate(df_list):
- marker = entry[2] # col 2: match_type / img_id / table_id
- if isinstance(marker, str) and marker.strip().split("\n", 1)[0].lower() in {
- "image",
- "table",
- }:
- continue
- if len(entry[0]) > summary_len and not entry[4] and not entry[5]:
- deferred_llm_tasks.append(("text", idx, entry[0]))
-
- # ── Post-loop: execute all deferred LLM calls in parallel via gevent ──
- if deferred_llm_tasks:
- image_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "image")
- table_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "table")
- text_task_count = sum(1 for task in deferred_llm_tasks if task[0] == "text")
- logger.info(
- f"Running {len(deferred_llm_tasks)} deferred summary LLM calls in parallel"
+ parser_state.collect_text_summary_tasks(summary_len)
+ apply_markdown_deferred_summaries(
+ MarkdownDeferredSummaryInput(
+ rows=parser_state.rows,
+ tasks=parser_state.deferred_llm_tasks,
+ output_dir=output_dir,
+ summary_len=summary_len,
)
- max_concurrent = getattr(settings, "SUMMARY_LLM_MAX_CONCURRENT", 8)
-
- with stage_timer(
- "md.deferred_summaries",
- total_tasks=len(deferred_llm_tasks),
- image_tasks=image_task_count,
- table_tasks=table_task_count,
- text_tasks=text_task_count,
- max_concurrent=min(max_concurrent, len(deferred_llm_tasks)),
- ):
-
- def _run_deferred(task):
- task_type, idx = task[0], task[1]
- try:
- if task_type == "image":
- relative_path = task[2]
- client = _get_vision_client()
- # TODO: Risk of missing text content if MinerU outputted a pure text image.
- # Consider adding judge-image-type and OCR fallback as done in image_parser.parse_image.
- llm_resp = ask_image(client, output_dir, paths_=[relative_path])
- if llm_resp:
- img_title, img_summary = split_title_summary(llm_resp)
- else:
- img_title, img_summary = None, None
- return idx, task_type, (img_title, img_summary)
- elif task_type == "table":
- tb_html = task[2]
- title, kw, summary = extract_title_keywords_summary(
- tb_html, max_keywords=3
- )
- return idx, task_type, (title, kw, summary)
- elif task_type == "text":
- text_content = task[2]
- _, kw, summary = extract_title_keywords_summary(
- text_content, max_keywords=3, summary_len=summary_len
- )
- return idx, task_type, (kw, summary)
- except Exception as e:
- logger.warning(
- f"Deferred {task_type} LLM call failed for idx={idx}: {e}"
- )
- return idx, task_type, None
-
- pool = GeventPool(size=min(max_concurrent, len(deferred_llm_tasks)))
- greenlets = [pool.spawn(_run_deferred, task) for task in deferred_llm_tasks]
- gevent.joinall(greenlets)
-
- # Build a lookup from deferred task list: idx -> original task tuple
- deferred_by_idx = {task[1]: task for task in deferred_llm_tasks}
-
- for g in greenlets:
- if g.value is None:
- continue
- idx, task_type, result = g.value
- if result is None:
- continue
- if task_type == "image":
- img_title, img_summary = result
- entry = df_list[idx]
- if img_summary:
- image_index = entry[5].split("\n")[0] if entry[5] else "image"
- entry[5] = f"{image_index}\n{img_summary}"
- # Rename image file if LLM provided a better title (mirrors table rename logic)
- if img_title:
- orig_task = deferred_by_idx[idx]
- i_dir, old_img_name, i_suffix = (
- orig_task[3],
- orig_task[4],
- orig_task[5],
- )
- safe_title = path_handle(img_title, mode="clean_single")
- # Derive image index number from old_img_name (e.g. "image-3-xxx" -> "3")
- img_num_match = re.match(r"image-(\d+)", old_img_name)
- img_num = (
- img_num_match.group(1)
- if img_num_match
- else (
- old_img_name.split("-")[1]
- if "-" in old_img_name
- else "0"
- )
- )
- new_img_name = path_handle(
- f"image-{img_num}-{safe_title}", mode="clean_single"
- )
- old_path = os.path.join(i_dir, f"{old_img_name}{i_suffix}")
- new_path = os.path.join(i_dir, f"{new_img_name}{i_suffix}")
- if old_path != new_path and os.path.exists(old_path):
- os.rename(old_path, new_path)
- new_relative_path = f"images/{new_img_name}{i_suffix}"
- replace_chunk_ref_in_rows(
- df_list, entry[1], new_relative_path
- )
- entry[1] = new_relative_path
- elif task_type == "table":
- title, kw, summary = result
- entry = df_list[idx]
- entry[4] = kw if isinstance(kw, str) else ""
- if summary:
- table_index = (
- entry[5]
- if "\n" not in entry[5]
- else entry[5].split("\n")[0]
- )
- entry[5] = f"{table_index}\n{summary}"
- # Rename table file if LLM provided a better title
- if title:
- orig_task = deferred_by_idx[idx]
- t_dir, old_tb_name, t_count = (
- orig_task[3],
- orig_task[4],
- orig_task[5],
- )
- safe_title = (
- sanitize_table_name_from_header(title) if title else ""
- )
- new_tb_name = path_handle(
- f"table-{t_count} {safe_title}", mode="clean_single"
- )
- old_path = os.path.join(t_dir, f"{old_tb_name}.html")
- new_path = os.path.join(t_dir, f"{new_tb_name}.html")
- if old_path != new_path and os.path.exists(old_path):
- os.rename(old_path, new_path)
- new_relative_path = f"tables/{new_tb_name}.html"
- replace_chunk_ref_in_rows(
- df_list, entry[1], new_relative_path
- )
- entry[1] = new_relative_path
- elif task_type == "text":
- kw, summary = result
- df_list[idx][4] = kw if isinstance(kw, str) else ""
- df_list[idx][5] = summary if isinstance(summary, str) else ""
-
- logger.info(
- f"Completed {len(deferred_llm_tasks)} deferred summary LLM calls"
- )
+ )
- with stage_timer("md.build_dataframe", row_count=len(df_list)):
- doc_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(","))
- doc_df = process_dup_paths_df(doc_df)
+ with stage_timer("md.build_dataframe", row_count=len(parser_state.rows)):
+ doc_df = parser_state.to_dataframe()
return doc_df
diff --git a/apps/worker/app/services/document_parser/mineru_client.py b/apps/worker/app/services/document_parser/mineru_client.py
new file mode 100644
index 000000000..1cf5e3b97
--- /dev/null
+++ b/apps/worker/app/services/document_parser/mineru_client.py
@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import requests
+from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager
+from loguru import logger
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from shared.core.config import settings
+from shared.core.exceptions.domain_exceptions import UnavailableException
+
+
+def build_mineru_session() -> requests.Session:
+ session = requests.Session()
+ retry_strategy = Retry(
+ total=settings.MINERU_UPLOAD_RETRY_TOTAL,
+ backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR,
+ status_forcelist=[429, 502, 503, 504],
+ allowed_methods=["GET", "POST", "PUT"],
+ raise_on_status=False,
+ )
+ adapter = HTTPAdapter(
+ max_retries=retry_strategy,
+ pool_connections=1,
+ pool_maxsize=settings.MINERU_POOL_MAXSIZE,
+ )
+ session.mount("https://", adapter)
+ session.mount("http://", adapter)
+ return session
+
+
+_mineru_session: Optional[requests.Session] = None
+
+
+def get_mineru_session() -> requests.Session:
+ global _mineru_session
+ if _mineru_session is None:
+ _mineru_session = build_mineru_session()
+ return _mineru_session
+
+
+def get_mineru_headers(api_key: str) -> dict[str, str]:
+ return {
+ "Content-Type": "application/json",
+ "Authorization": f"Bearer {api_key}",
+ }
+
+
+def mineru_logger(step: str, **fields: Any):
+ return logger.bind(service="mineru", step=step, **fields)
+
+
+def get_retry_after_seconds(
+ response: requests.Response, default_retry_after: int
+) -> int:
+ retry_after_header = response.headers.get("Retry-After")
+ if retry_after_header:
+ try:
+ return max(
+ 1,
+ min(
+ int(retry_after_header), settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER
+ ),
+ )
+ except ValueError:
+ logger.debug(f"Invalid MinerU Retry-After header: {retry_after_header}")
+
+ return max(1, min(default_retry_after, settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER))
+
+
+def raise_mineru_unavailable(
+ token_id: str, response: requests.Response, operation: str
+) -> None:
+ retry_after = get_retry_after_seconds(
+ response, settings.MINERU_TOKEN_COOLDOWN_SECONDS
+ )
+ quota_manager = get_mineru_quota_manager()
+ quota_manager.mark_rate_limited(token_id, retry_after)
+ mineru_logger(
+ "rate_limited",
+ operation=operation,
+ token_id=token_id,
+ status_code=response.status_code,
+ retry_after=retry_after,
+ ).warning("MinerU request rate-limited")
+ raise UnavailableException(
+ internal_message=f"MinerU rate limited during {operation}",
+ retry_after=retry_after,
+ limit=settings.MINERU_TOKEN_RPM_LIMIT,
+ period="minute",
+ user_message="Document processing is busy right now. Please retry shortly.",
+ )
diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/mineru_pdf_service.py
index 24a2145cc..bd4acdbc4 100644
--- a/apps/worker/app/services/document_parser/mineru_pdf_service.py
+++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py
@@ -1,27 +1,29 @@
-# pyright: reportUnusedExpression=false
import os
-import time
-from typing import Any, Callable, Optional
+from typing import Optional
import requests
+from app.services.document_parser.mineru_client import (
+ get_mineru_headers,
+ get_mineru_session,
+ mineru_logger,
+ raise_mineru_unavailable,
+)
from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager
+from app.services.document_parser.mineru_task_polling import (
+ get_batch_status,
+ poll_mineru_task,
+)
from app.services.document_parser.parser_log_utils import truncate_log_value
-from loguru import logger
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
from shared.core.config import settings
from shared.core.constants import APIConstants
from shared.core.exceptions.domain_exceptions import (
MinerUServiceException,
- PDFParsingException,
StorageServiceException,
- TimeoutException,
UnavailableException,
)
-from shared.core.exceptions.knowhere_exception import KnowhereException
-from shared.utils.CommonHelperSync import is_remote
-from shared.utils.FileDownUpUtils import s3_download_extract_zip
+from shared.services.storage.job_file_storage import JobFileStorage
+from shared.utils.file_loading import is_remote
MINERU_UPLOAD_TIMEOUT = (
settings.MINERU_UPLOAD_CONNECT_TIMEOUT,
@@ -29,49 +31,6 @@
)
-def _build_mineru_session() -> requests.Session:
- session = requests.Session()
- # Hybrid rate-limit control: urllib3 handles transient 429s with backoff
- # (respects Retry-After header); application-level handling in callers
- # covers persistent rate limits with Redis token marking and pool rotation.
- retry_strategy = Retry(
- total=settings.MINERU_UPLOAD_RETRY_TOTAL,
- backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR,
- status_forcelist=[429, 502, 503, 504],
- allowed_methods=["GET", "POST", "PUT"],
- raise_on_status=False,
- )
- adapter = HTTPAdapter(
- max_retries=retry_strategy,
- pool_connections=1,
- pool_maxsize=settings.MINERU_POOL_MAXSIZE,
- )
- session.mount("https://", adapter)
- session.mount("http://", adapter)
- return session
-
-
-_mineru_session: Optional[requests.Session] = None
-
-
-def get_mineru_session() -> requests.Session:
- global _mineru_session
- if _mineru_session is None:
- _mineru_session = _build_mineru_session()
- return _mineru_session
-
-
-def get_mineru_headers(api_key: str) -> dict[str, str]:
- return {
- "Content-Type": "application/json",
- "Authorization": f"Bearer {api_key}",
- }
-
-
-def _mineru_logger(step: str, **fields: Any):
- return logger.bind(service="mineru", step=step, **fields)
-
-
def _should_use_mineru_s3_url_mode(s3_key: Optional[str]) -> bool:
if settings.FORCE_MINERU_UPLOAD_ENABLED:
return False
@@ -85,7 +44,7 @@ def _log_mineru_url_mode_storage_fallback(
local_file_path: Optional[str],
exc: Exception,
) -> None:
- _mineru_logger(
+ mineru_logger(
"url_mode_storage_fallback",
operation=operation,
source_s3_key=s3_key,
@@ -103,7 +62,7 @@ def _log_mineru_url_mode_ingestion_fallback(
pdf_url: str,
exc: Exception,
) -> None:
- _mineru_logger(
+ mineru_logger(
"url_mode_ingestion_fallback",
operation=operation,
source_s3_key=s3_key,
@@ -120,10 +79,8 @@ def _inspect_mineru_source_s3_key(s3_key: Optional[str]) -> tuple[Optional[str],
return None, False
assert s3_key is not None
- from app.services.storage.sync_storage_service import verify_s3_file_exists
-
try:
- existing_file = verify_s3_file_exists(s3_key, settings.S3_BUCKET_NAME)
+ existing_file = JobFileStorage().verify_upload_exists(s3_key)
except Exception as exc:
_log_mineru_url_mode_storage_fallback(
operation="verify_source_object",
@@ -134,7 +91,7 @@ def _inspect_mineru_source_s3_key(s3_key: Optional[str]) -> tuple[Optional[str],
return None, False
if existing_file.get("exists"):
- _mineru_logger(
+ mineru_logger(
"url_mode_source_reused",
source_s3_key=s3_key,
).info("Reusing existing S3 source for MinerU URL mode")
@@ -165,10 +122,8 @@ def resolve_mineru_source_s3_key(
return None
assert s3_key is not None
- from app.services.storage.sync_storage_service import upload_to_s3
-
try:
- upload_to_s3(local_file_path, s3_key, settings.S3_BUCKET_NAME)
+ JobFileStorage().upload_source_file(local_file_path, s3_key)
except Exception as exc:
_log_mineru_url_mode_storage_fallback(
operation="upload_source_object",
@@ -178,7 +133,7 @@ def resolve_mineru_source_s3_key(
)
return None
- _mineru_logger(
+ mineru_logger(
"url_mode_source_uploaded",
source_s3_key=s3_key,
local_file_path=local_file_path,
@@ -186,293 +141,11 @@ def resolve_mineru_source_s3_key(
return s3_key
-def _get_retry_after_seconds(
- response: requests.Response, default_retry_after: int
-) -> int:
- """Parse Retry-After header with sane bounds for worker backoff."""
- retry_after_header = response.headers.get("Retry-After")
- if retry_after_header:
- try:
- return max(
- 1,
- min(
- int(retry_after_header), settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER
- ),
- )
- except ValueError:
- logger.debug(f"Invalid MinerU Retry-After header: {retry_after_header}")
-
- return max(1, min(default_retry_after, settings.MINERU_RATE_LIMIT_MAX_RETRY_AFTER))
-
-
-def _raise_mineru_unavailable(
- token_id: str, response: requests.Response, operation: str
-) -> None:
- retry_after = _get_retry_after_seconds(
- response, settings.MINERU_TOKEN_COOLDOWN_SECONDS
- )
- quota_manager = get_mineru_quota_manager()
- quota_manager.mark_rate_limited(token_id, retry_after)
- _mineru_logger(
- "rate_limited",
- operation=operation,
- token_id=token_id,
- status_code=response.status_code,
- retry_after=retry_after,
- ).warning("MinerU request rate-limited")
- raise UnavailableException(
- internal_message=f"MinerU rate limited during {operation}",
- retry_after=retry_after,
- limit=settings.MINERU_TOKEN_RPM_LIMIT,
- period="minute",
- user_message="Document processing is busy right now. Please retry shortly.",
- )
-
-
-def _polling_interval_for_state(state: str, attempt: int) -> float:
- """Return seconds to sleep before the next poll.
-
- Tuned so that ``WORKER_CONCURRENCY`` scales safely with MinerU limits.
- With a 4-token pool (300 RPM each), total budget is 1200 req/min:
-
- 150 tasks × (60 s / 15 s) = 600 req/min total
- 600 / 4 tokens = 150 req/min per token ← leaves headroom
-
- Observed data (Logfire, 2026-03-08 dev batch):
- - 99 % of tasks never enter ``running``; lifecycle is
- ``waiting-file`` → ``done`` in 2-4 s on MinerU's side.
- - Longest observed task: 21 s (5 poll attempts).
- - Peak burst: 351 concurrent tasks, 308 req/min → rate-limited.
- """
- if state == "pending":
- return min(20.0, 5.0 + attempt * 1.5)
- if state == "running":
- return 10.0
- # waiting-file, converting, unknown, etc.
- return 15.0
-
-
-def _get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]:
- extract_result = data.get("data", {}).get("extract_result")
- if isinstance(extract_result, list):
- return extract_result[0] if extract_result else None
- return extract_result
-
-
-def poll_mineru_task(
- status_url: str,
- task_id: str,
- output_dir: str,
- get_status: Callable[[dict[str, Any]], Optional[dict[str, Any]]],
- preferred_token_id: Optional[str] = None,
-) -> None:
- quota_manager = get_mineru_quota_manager()
- polling_logger = _mineru_logger(
- "poll_status",
- operation="poll_status",
- task_id=task_id,
- preferred_token_id=preferred_token_id,
- )
-
- max_polling_attempts = 120
- polling_interval = 5.0
- max_wait_time = 6000
-
- start_time = time.time()
- attempt = 0
- last_token_id: Optional[str] = None
- last_state: Optional[str] = None
-
- polling_logger.info("Starting MinerU polling")
-
- while attempt < max_polling_attempts:
- if time.time() - start_time > max_wait_time:
- polling_logger.bind(
- attempt=attempt + 1,
- max_polling_attempts=max_polling_attempts,
- max_wait_time=max_wait_time,
- ).warning("MinerU polling timed out")
- raise TimeoutException(
- internal_message=f"PDF parsing timed out, exceeded {max_wait_time} seconds",
- retry_after=60,
- user_message="PDF parsing timed out. Please try again.",
- )
-
- try:
- logger.debug(
- f"parse_pdfs status_url: {status_url} "
- f"(attempt {attempt + 1}/{max_polling_attempts})"
- )
- lease = quota_manager.acquire_request(
- operation="poll_status",
- preferred_token_id=preferred_token_id,
- )
- if lease.token_id != last_token_id:
- polling_logger.bind(
- token_id=lease.token_id,
- attempt=attempt + 1,
- ).info("Acquired MinerU token for polling")
- last_token_id = lease.token_id
-
- response = get_mineru_session().get(
- status_url,
- headers=get_mineru_headers(lease.api_key),
- timeout=settings.MINERU_API_TIMEOUT,
- )
-
- if response.status_code == 429:
- # urllib3 already retried with backoff — if we still see 429,
- # mark the token and let Celery handle task-level retry.
- _raise_mineru_unavailable(
- lease.token_id, response, operation="poll_status"
- )
-
- if response.status_code == 200:
- response_json = response.json()
- if response_json.get("code") != 0:
- response_message = str(response_json.get("msg") or "Unknown error")
- if "rate limit" in response_message.lower():
- quota_manager.mark_rate_limited(
- lease.token_id,
- settings.MINERU_TOKEN_COOLDOWN_SECONDS,
- )
- raise UnavailableException(
- internal_message=(
- f"MinerU rate limited during poll_status: {response_message}"
- ),
- retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS,
- limit=lease.rpm_limit,
- period="minute",
- user_message="Document processing is busy right now. Please retry shortly.",
- )
- raise MinerUServiceException(
- internal_message=f"MinerU API Error: {response_message}"
- )
-
- status = get_status(response_json)
- if not status:
- polling_logger.bind(
- token_id=lease.token_id,
- attempt=attempt + 1,
- ).warning("Received empty MinerU status payload")
- time.sleep(polling_interval)
- attempt += 1
- continue
-
- state = status.get("state", "unknown")
- if state != last_state:
- # polling_logger.bind(
- # token_id=lease.token_id,
- # attempt=attempt + 1,
- # state=state,
- # ).info("MinerU status changed")
- last_state = state
-
- if state == "done":
- s3_download_extract_zip(
- status["full_zip_url"],
- dest_dir=output_dir,
- keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"),
- exclude_patterns=("content_list", "middle.json", "model.json"),
- )
- polling_logger.bind(token_id=lease.token_id).info(
- "MinerU parsing completed"
- )
- break
-
- if state == "running":
- if "extract_progress" in status:
- try:
- (
- status["extract_progress"]["extracted_pages"]
- / status["extract_progress"]["total_pages"]
- )
- # polling_logger.bind(
- # token_id=lease.token_id,
- # progress=progress,
- # ).info("MinerU parsing progress updated")
- except (KeyError, ZeroDivisionError):
- polling_logger.bind(token_id=lease.token_id).info(
- "MinerU parsing in progress"
- )
- else:
- polling_logger.bind(token_id=lease.token_id).info(
- "MinerU parsing in progress"
- )
- elif state == "failed":
- error_message = status.get("err_msg", "Unknown error")
- polling_logger.bind(
- token_id=lease.token_id,
- error_message=error_message,
- ).error("MinerU parsing reported failed state")
- raise PDFParsingException(
- user_message="Failed to parse the PDF file",
- internal_message=f"MinerU failed with state 'failed': {error_message}",
- )
- elif state == "pending":
- polling_logger.bind(token_id=lease.token_id).debug(
- "MinerU parsing pending"
- )
- elif state == "waiting-file":
- polling_logger.bind(token_id=lease.token_id).debug(
- "MinerU waiting for file queueing"
- )
- elif state == "converting":
- polling_logger.bind(token_id=lease.token_id).debug(
- "MinerU converting file"
- )
- else:
- polling_logger.bind(
- token_id=lease.token_id,
- state=state,
- ).warning("MinerU returned unknown state")
-
- time.sleep(_polling_interval_for_state(state, attempt))
- attempt += 1
- else:
- polling_logger.bind(
- token_id=lease.token_id,
- attempt=attempt + 1,
- status_code=response.status_code,
- ).warning("MinerU status query failed")
- time.sleep(polling_interval * 2)
- attempt += 1
-
- except requests.RequestException as exc:
- polling_logger.bind(
- attempt=attempt + 1,
- error_message=str(exc),
- ).warning("MinerU polling network request failed")
- time.sleep(polling_interval * 2)
- attempt += 1
- except KnowhereException:
- raise
- except Exception as exc:
- polling_logger.bind(
- attempt=attempt + 1,
- error_message=str(exc),
- ).error("Unexpected error during MinerU polling")
- raise PDFParsingException(
- user_message="An unexpected error occurred while parsing the PDF",
- internal_message=str(exc),
- original_exception=exc,
- )
-
- if attempt >= max_polling_attempts:
- raise TimeoutException(
- internal_message=(
- f"minerU PDF parsing timed out after {max_polling_attempts} attempts, "
- f"Task ID: {task_id}"
- ),
- retry_after=60,
- user_message="PDF parsing timed out. Please try again.",
- )
-
def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]:
base_url = settings.MINERU_URL
quota_manager = get_mineru_quota_manager()
- upload_logger = _mineru_logger(
+ upload_logger = mineru_logger(
"upload_url",
operation="upload_url",
filename=filename,
@@ -504,7 +177,7 @@ def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]:
timeout=settings.MINERU_API_TIMEOUT,
)
if response.status_code == 429:
- _raise_mineru_unavailable(lease.token_id, response, operation="upload_url")
+ raise_mineru_unavailable(lease.token_id, response, operation="upload_url")
if response.status_code != 200:
upload_logger.bind(
token_id=lease.token_id,
@@ -554,7 +227,7 @@ def _request_upload_target(pdf_url: str, filename: str) -> tuple[str, str, str]:
def _upload_file_to_mineru(
pdf_url: str, filename: str, upload_url: str, token_id: str
) -> None:
- upload_logger = _mineru_logger(
+ upload_logger = mineru_logger(
"file_upload",
operation="file_upload",
filename=filename,
@@ -648,7 +321,7 @@ def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]:
"""
base_url = settings.MINERU_URL
quota_manager = get_mineru_quota_manager()
- submit_logger = _mineru_logger(
+ submit_logger = mineru_logger(
"submit_url_task",
operation="submit_url_task",
filename=filename,
@@ -678,7 +351,7 @@ def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]:
)
if response.status_code == 429:
- _raise_mineru_unavailable(lease.token_id, response, operation="submit_url_task")
+ raise_mineru_unavailable(lease.token_id, response, operation="submit_url_task")
if response.status_code != 200:
submit_logger.bind(
@@ -722,6 +395,8 @@ def parse_via_full(
output_dir: str,
s3_key: Optional[str] = None,
) -> None:
+ batch_id: str | None = None
+ token_id: str | None = None
resolved_s3_key = resolve_mineru_source_s3_key(
s3_key=s3_key,
local_file_path=None if is_remote(pdf_url) else pdf_url,
@@ -729,13 +404,11 @@ def parse_via_full(
if resolved_s3_key is not None:
try:
- from app.services.storage.sync_storage_service import generate_download_url
-
- presigned = generate_download_url(
+ presigned = JobFileStorage().generate_upload_download_url(
resolved_s3_key, expires_in=settings.MINERU_URL_MODE_PRESIGN_EXPIRY
)
presigned_url = presigned["download_url"]
- _mineru_logger("ingestion_mode", mode="s3_url").info(
+ mineru_logger("ingestion_mode", mode="s3_url").info(
"Using S3 URL mode for MinerU ingestion"
)
batch_id, token_id = _submit_url_task(presigned_url, filename)
@@ -749,16 +422,21 @@ def parse_via_full(
resolved_s3_key = None
if resolved_s3_key is None:
- _mineru_logger("ingestion_mode", mode="direct_upload").info(
+ mineru_logger("ingestion_mode", mode="direct_upload").info(
"Using direct upload mode for MinerU ingestion"
)
batch_id, upload_url, token_id = _request_upload_target(pdf_url, filename)
_upload_file_to_mineru(pdf_url, filename, upload_url, token_id)
+ if batch_id is None or token_id is None:
+ raise MinerUServiceException(
+ internal_message="MinerU task setup completed without a batch id or token"
+ )
+
poll_mineru_task(
status_url=f"{settings.MINERU_URL}/extract-results/batch/{batch_id}",
task_id=batch_id,
output_dir=output_dir,
- get_status=_get_batch_status,
+ get_status=get_batch_status,
preferred_token_id=token_id,
)
diff --git a/apps/worker/app/services/document_parser/mineru_task_polling.py b/apps/worker/app/services/document_parser/mineru_task_polling.py
new file mode 100644
index 000000000..c09c1ea86
--- /dev/null
+++ b/apps/worker/app/services/document_parser/mineru_task_polling.py
@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import Any, Optional
+
+import requests
+from app.services.document_parser.mineru_client import (
+ get_mineru_headers,
+ get_mineru_session,
+ mineru_logger,
+ raise_mineru_unavailable,
+)
+from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager
+from loguru import logger
+
+from shared.core.config import settings
+from shared.core.exceptions.domain_exceptions import (
+ MinerUServiceException,
+ PDFParsingException,
+ TimeoutException,
+ UnavailableException,
+)
+from shared.core.exceptions.knowhere_exception import KnowhereException
+from shared.utils.zip_download import download_and_extract_zip
+
+
+def get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]:
+ extract_result = data.get("data", {}).get("extract_result")
+ if isinstance(extract_result, list):
+ return extract_result[0] if extract_result else None
+ return extract_result
+
+
+def get_polling_interval_for_state(state: str, attempt: int) -> float:
+ """Return seconds to sleep before the next MinerU status poll."""
+ if state == "pending":
+ return min(20.0, 5.0 + attempt * 1.5)
+ if state == "running":
+ return 10.0
+ return 15.0
+
+
+def poll_mineru_task(
+ status_url: str,
+ task_id: str,
+ output_dir: str,
+ get_status: Callable[[dict[str, Any]], Optional[dict[str, Any]]],
+ preferred_token_id: Optional[str] = None,
+) -> None:
+ quota_manager = get_mineru_quota_manager()
+ polling_logger = mineru_logger(
+ "poll_status",
+ operation="poll_status",
+ task_id=task_id,
+ preferred_token_id=preferred_token_id,
+ )
+
+ max_polling_attempts = 120
+ polling_interval = 5.0
+ max_wait_time = 6000
+
+ start_time = time.time()
+ attempt = 0
+ last_token_id: Optional[str] = None
+ last_state: Optional[str] = None
+
+ polling_logger.info("Starting MinerU polling")
+
+ while attempt < max_polling_attempts:
+ if time.time() - start_time > max_wait_time:
+ polling_logger.bind(
+ attempt=attempt + 1,
+ max_polling_attempts=max_polling_attempts,
+ max_wait_time=max_wait_time,
+ ).warning("MinerU polling timed out")
+ raise TimeoutException(
+ internal_message=f"PDF parsing timed out, exceeded {max_wait_time} seconds",
+ retry_after=60,
+ user_message="PDF parsing timed out. Please try again.",
+ )
+
+ try:
+ logger.debug(
+ f"parse_pdfs status_url: {status_url} "
+ f"(attempt {attempt + 1}/{max_polling_attempts})"
+ )
+ lease = quota_manager.acquire_request(
+ operation="poll_status",
+ preferred_token_id=preferred_token_id,
+ )
+ if lease.token_id != last_token_id:
+ polling_logger.bind(
+ token_id=lease.token_id,
+ attempt=attempt + 1,
+ ).info("Acquired MinerU token for polling")
+ last_token_id = lease.token_id
+
+ response = get_mineru_session().get(
+ status_url,
+ headers=get_mineru_headers(lease.api_key),
+ timeout=settings.MINERU_API_TIMEOUT,
+ )
+
+ if response.status_code == 429:
+ raise_mineru_unavailable(
+ lease.token_id, response, operation="poll_status"
+ )
+
+ if response.status_code == 200:
+ response_json = response.json()
+ if response_json.get("code") != 0:
+ response_message = str(response_json.get("msg") or "Unknown error")
+ if "rate limit" in response_message.lower():
+ quota_manager.mark_rate_limited(
+ lease.token_id,
+ settings.MINERU_TOKEN_COOLDOWN_SECONDS,
+ )
+ raise UnavailableException(
+ internal_message=(
+ f"MinerU rate limited during poll_status: {response_message}"
+ ),
+ retry_after=settings.MINERU_TOKEN_COOLDOWN_SECONDS,
+ limit=lease.rpm_limit,
+ period="minute",
+ user_message="Document processing is busy right now. Please retry shortly.",
+ )
+ raise MinerUServiceException(
+ internal_message=f"MinerU API Error: {response_message}"
+ )
+
+ status = get_status(response_json)
+ if not status:
+ polling_logger.bind(
+ token_id=lease.token_id,
+ attempt=attempt + 1,
+ ).warning("Received empty MinerU status payload")
+ time.sleep(polling_interval)
+ attempt += 1
+ continue
+
+ state = status.get("state", "unknown")
+ if state != last_state:
+ last_state = state
+
+ if state == "done":
+ download_and_extract_zip(
+ status["full_zip_url"],
+ dest_dir=output_dir,
+ keep_exts=(".md", ".jpg", ".jpeg", ".png", ".gif", ".json"),
+ exclude_patterns=("content_list", "middle.json", "model.json"),
+ )
+ polling_logger.bind(token_id=lease.token_id).info(
+ "MinerU parsing completed"
+ )
+ break
+
+ if state == "running":
+ if "extract_progress" in status:
+ try:
+ _progress = (
+ status["extract_progress"]["extracted_pages"]
+ / status["extract_progress"]["total_pages"]
+ )
+ except (KeyError, ZeroDivisionError):
+ polling_logger.bind(token_id=lease.token_id).info(
+ "MinerU parsing in progress"
+ )
+ else:
+ polling_logger.bind(token_id=lease.token_id).info(
+ "MinerU parsing in progress"
+ )
+ elif state == "failed":
+ error_message = status.get("err_msg", "Unknown error")
+ polling_logger.bind(
+ token_id=lease.token_id,
+ error_message=error_message,
+ ).error("MinerU parsing reported failed state")
+ raise PDFParsingException(
+ user_message="Failed to parse the PDF file",
+ internal_message=f"MinerU failed with state 'failed': {error_message}",
+ )
+ elif state == "pending":
+ polling_logger.bind(token_id=lease.token_id).debug(
+ "MinerU parsing pending"
+ )
+ elif state == "waiting-file":
+ polling_logger.bind(token_id=lease.token_id).debug(
+ "MinerU waiting for file queueing"
+ )
+ elif state == "converting":
+ polling_logger.bind(token_id=lease.token_id).debug(
+ "MinerU converting file"
+ )
+ else:
+ polling_logger.bind(
+ token_id=lease.token_id,
+ state=state,
+ ).warning("MinerU returned unknown state")
+
+ time.sleep(get_polling_interval_for_state(state, attempt))
+ attempt += 1
+ else:
+ polling_logger.bind(
+ token_id=lease.token_id,
+ attempt=attempt + 1,
+ status_code=response.status_code,
+ ).warning("MinerU status query failed")
+ time.sleep(polling_interval * 2)
+ attempt += 1
+
+ except requests.RequestException as exc:
+ polling_logger.bind(
+ attempt=attempt + 1,
+ error_message=str(exc),
+ ).warning("MinerU polling network request failed")
+ time.sleep(polling_interval * 2)
+ attempt += 1
+ except KnowhereException:
+ raise
+ except Exception as exc:
+ polling_logger.bind(
+ attempt=attempt + 1,
+ error_message=str(exc),
+ ).error("Unexpected error during MinerU polling")
+ raise PDFParsingException(
+ user_message="An unexpected error occurred while parsing the PDF",
+ internal_message=str(exc),
+ original_exception=exc,
+ )
+
+ if attempt >= max_polling_attempts:
+ raise TimeoutException(
+ internal_message=(
+ f"minerU PDF parsing timed out after {max_polling_attempts} attempts, "
+ f"Task ID: {task_id}"
+ ),
+ retry_after=60,
+ user_message="PDF parsing timed out. Please try again.",
+ )
diff --git a/apps/worker/app/services/document_parser/orchestration/__init__.py b/apps/worker/app/services/document_parser/orchestration/__init__.py
new file mode 100644
index 000000000..6f12d8497
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/__init__.py
@@ -0,0 +1 @@
+"""Document parser orchestration modules."""
diff --git a/apps/worker/app/services/document_parser/orchestration/format_adapters.py b/apps/worker/app/services/document_parser/orchestration/format_adapters.py
new file mode 100644
index 000000000..f10f9edd1
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py
@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+import pandas as pd
+
+from app.services.document_parser.orchestration.parse_session import ParseSession
+
+
+class DocumentParseAdapter(Protocol):
+ @property
+ def document_format(self) -> object:
+ """Document format handled by this adapter."""
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ """Parse a document session into the parser output directory and DataFrame."""
+ ...
+
+
+@dataclass(frozen=True)
+class FragmentParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.fragment_parser import parse_fragment
+
+ full_output_dir, _relative_root, parsed_df = parse_fragment(
+ session.fragment_content,
+ filename=session.filename,
+ output_dir=session.output_dir,
+ kb_dir=session.kb_dir,
+ base_llm_paras=session.base_llm_paras,
+ )
+ return full_output_dir, parsed_df
+
+
+@dataclass(frozen=True)
+class TextParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.md_parser import parse_md
+ from app.services.document_parser.txt_parser import parse_texts
+
+ text_lines = parse_texts(file_path=session.file_full_path, baseurl=session.base_url)
+ parsed_df = parse_md(
+ session.full_output_dir,
+ source_type="md",
+ md_lines=text_lines,
+ base_llm_paras=session.base_llm_paras,
+ relative_root=session.relative_root,
+ )
+ return session.full_output_dir, parsed_df
+
+
+@dataclass(frozen=True)
+class ImageParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.image_parser import parse_image
+
+ parsed_df = parse_image(
+ session.file_full_path,
+ filename=session.filename,
+ output_dir=session.full_output_dir,
+ baseurl=session.base_url,
+ base_llm_paras=session.base_llm_paras,
+ relative_root=session.relative_root,
+ )
+ return session.full_output_dir, parsed_df
+
+
+@dataclass(frozen=True)
+class PdfParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.pdf_parser import parse_pdfs
+
+ parsed_df = parse_pdfs(
+ session.file_full_path,
+ filename=session.filename,
+ output_dir=session.full_output_dir,
+ base_llm_paras=session.base_llm_paras,
+ profile=session.profile,
+ relative_root=session.relative_root,
+ s3_key=session.s3_key,
+ )
+ return session.full_output_dir, parsed_df
+
+
+@dataclass(frozen=True)
+class DocParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.legacy_converter import doc_to_docx
+
+ converted_docx_path, _ = doc_to_docx(
+ session.file_full_path,
+ outdir=session.full_output_dir,
+ )
+ return _parse_docx_path(converted_docx_path, session)
+
+
+@dataclass(frozen=True)
+class DocxParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ return _parse_docx_path(session.file_full_path, session)
+
+
+@dataclass(frozen=True)
+class XlsParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.legacy_converter import xls_to_xlsx
+
+ converted_xlsx_path, _ = xls_to_xlsx(
+ session.file_full_path,
+ outdir=session.full_output_dir,
+ )
+ return _parse_xlsx_path(converted_xlsx_path, session)
+
+
+@dataclass(frozen=True)
+class XlsxParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ return _parse_xlsx_path(session.file_full_path, session)
+
+
+@dataclass(frozen=True)
+class PptxParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.pptx_parser import parse_pptx
+
+ parsed_df = parse_pptx(
+ session.file_full_path,
+ filename=session.filename,
+ output_dir=session.full_output_dir,
+ base_llm_paras=session.base_llm_paras,
+ strategy="to_pdf_api",
+ job_id=session.job_id,
+ relative_root=session.relative_root,
+ baseurl=session.base_url,
+ )
+ return session.full_output_dir, parsed_df
+
+
+@dataclass(frozen=True)
+class MarkdownParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.md_parser import parse_md
+
+ parsed_df = parse_md(
+ session.full_output_dir,
+ source_type="md",
+ file_path=session.file_full_path,
+ base_llm_paras=session.base_llm_paras,
+ relative_root=session.relative_root,
+ )
+ return session.full_output_dir, parsed_df
+
+
+@dataclass(frozen=True)
+class JsonParseAdapter:
+ document_format: object
+
+ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ return session.full_output_dir, None
+
+
+def _parse_docx_path(
+ docx_path: str,
+ session: ParseSession,
+) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx
+
+ parsed_structure, dataframe_list = parse_docx(
+ docx_path,
+ session.base_llm_paras,
+ session.full_output_dir,
+ session.filename,
+ session.base_url,
+ relative_root=session.relative_root,
+ )
+ parsed_df = convert_doc2dics(
+ parsed_structure,
+ dataframe_list,
+ session.full_output_dir,
+ base_llm_paras=session.base_llm_paras,
+ relative_root=session.relative_root,
+ )
+ return session.full_output_dir, parsed_df
+
+
+def _parse_xlsx_path(
+ xlsx_path: str,
+ session: ParseSession,
+) -> tuple[str, pd.DataFrame | None]:
+ from app.services.document_parser.excel_table_parser import parse_xlsx
+
+ parsed_df = parse_xlsx(
+ xlsx_path,
+ session.filename,
+ session.full_output_dir,
+ session.base_url,
+ base_llm_paras=session.base_llm_paras,
+ relative_root=session.relative_root,
+ )
+ return session.full_output_dir, parsed_df
diff --git a/apps/worker/app/services/document_parser/orchestration/format_router.py b/apps/worker/app/services/document_parser/orchestration/format_router.py
new file mode 100644
index 000000000..b94cf2787
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/format_router.py
@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import os
+from enum import Enum
+
+from app.services.document_parser.orchestration import format_adapters
+from app.services.document_parser.orchestration.format_adapters import (
+ DocumentParseAdapter,
+)
+from shared.core.exceptions.domain_exceptions import ValidationException
+
+
+class DocumentFormat(str, Enum):
+ TEXT = "text"
+ FRAGMENT = "fragment"
+ IMAGE = "image"
+ PDF = "pdf"
+ DOC = "doc"
+ DOCX = "docx"
+ XLS = "xls"
+ XLSX = "xlsx"
+ PPTX = "pptx"
+ MARKDOWN = "markdown"
+ JSON = "json"
+
+
+SUPPORTED_FILE_TYPES: tuple[str, ...] = (
+ ".txt",
+ ".fragment",
+ ".png",
+ ".jpg",
+ ".jpeg",
+ ".pdf",
+ ".doc",
+ ".docx",
+ ".xls",
+ ".xlsx",
+ ".pptx",
+ ".md",
+ ".json",
+)
+
+
+def resolve_document_format(file_path: str) -> DocumentFormat:
+ extension = os.path.splitext(file_path)[1].lower()
+ if extension == ".fragment":
+ return DocumentFormat.FRAGMENT
+ if extension == ".txt":
+ return DocumentFormat.TEXT
+ if extension in (".png", ".jpg", ".jpeg"):
+ return DocumentFormat.IMAGE
+ if extension == ".pdf":
+ return DocumentFormat.PDF
+ if extension == ".doc":
+ return DocumentFormat.DOC
+ if extension == ".docx":
+ return DocumentFormat.DOCX
+ if extension == ".xls":
+ return DocumentFormat.XLS
+ if extension == ".xlsx":
+ return DocumentFormat.XLSX
+ if extension == ".pptx":
+ return DocumentFormat.PPTX
+ if extension == ".md":
+ return DocumentFormat.MARKDOWN
+ if extension == ".json":
+ return DocumentFormat.JSON
+
+ raise ValidationException(
+ user_message=f"Unsupported file type: {extension}",
+ violations=[
+ {
+ "field": "file_type",
+ "description": f"Must be one of: {', '.join(SUPPORTED_FILE_TYPES)}",
+ }
+ ],
+ )
+
+
+def get_document_parse_adapter(document_format: DocumentFormat) -> DocumentParseAdapter:
+ adapter_by_format: dict[DocumentFormat, DocumentParseAdapter] = {
+ DocumentFormat.FRAGMENT: format_adapters.FragmentParseAdapter(document_format),
+ DocumentFormat.TEXT: format_adapters.TextParseAdapter(document_format),
+ DocumentFormat.IMAGE: format_adapters.ImageParseAdapter(document_format),
+ DocumentFormat.PDF: format_adapters.PdfParseAdapter(document_format),
+ DocumentFormat.DOC: format_adapters.DocParseAdapter(document_format),
+ DocumentFormat.DOCX: format_adapters.DocxParseAdapter(document_format),
+ DocumentFormat.XLS: format_adapters.XlsParseAdapter(document_format),
+ DocumentFormat.XLSX: format_adapters.XlsxParseAdapter(document_format),
+ DocumentFormat.PPTX: format_adapters.PptxParseAdapter(document_format),
+ DocumentFormat.MARKDOWN: format_adapters.MarkdownParseAdapter(document_format),
+ DocumentFormat.JSON: format_adapters.JsonParseAdapter(document_format),
+ }
+ return adapter_by_format[document_format]
diff --git a/apps/worker/app/services/document_parser/orchestration/parse_input.py b/apps/worker/app/services/document_parser/orchestration/parse_input.py
new file mode 100644
index 000000000..89249b3d0
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/parse_input.py
@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class ParseOptions:
+ llm_histories: int = 5
+ smart_title_parse: bool = True
+ summary_image: bool = True
+ summary_table: bool = True
+ summary_txt: bool = True
+ stopwords: list[str] | None = None
+ doc_type: str = "auto"
+ add_frag_desc: str = ""
+
+
+@dataclass(frozen=True)
+class ParseInput:
+ file_full_path: str
+ filename: str
+ output_dir: str
+ internal_output_filename: str
+ job_id: str | None = None
+ kb_dir: str = "Default_Root"
+ options: ParseOptions = field(default_factory=ParseOptions)
+ base_url: str = ""
+ fragment_content: str = ""
+ s3_key: str | None = None
diff --git a/apps/worker/app/services/document_parser/orchestration/parse_session.py b/apps/worker/app/services/document_parser/orchestration/parse_session.py
new file mode 100644
index 000000000..832e13969
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py
@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Any
+
+from app.services.document_parser.atlas_classifier import classify_atlas_with_vlm
+from app.services.document_parser.doc_profiler import profile_document
+from app.services.document_parser.orchestration.parse_input import ParseInput
+from app.services.document_parser.stage_profiler import stage_timer
+from loguru import logger
+
+from shared.core.config import settings
+from shared.core.exceptions.domain_exceptions import ValidationException
+from shared.utils.file_utils import path_handle
+
+PDF_PAGE_LIMIT = 600
+
+
+@dataclass(frozen=True)
+class ParseSession:
+ base_llm_paras: dict[str, object]
+ base_url: str
+ file_full_path: str
+ filename: str
+ fragment_content: str
+ full_output_dir: str
+ internal_output_filename: str
+ job_id: str | None
+ kb_dir: str
+ output_dir: str
+ profile: Any
+ relative_root: str
+ s3_key: str | None
+
+ @classmethod
+ def from_input(
+ cls,
+ *,
+ parse_input: ParseInput,
+ base_llm_paras: dict[str, object],
+ full_output_dir: str,
+ profile: Any,
+ relative_root: str,
+ ) -> "ParseSession":
+ return cls(
+ base_llm_paras=base_llm_paras,
+ base_url=parse_input.base_url,
+ file_full_path=parse_input.file_full_path,
+ filename=parse_input.filename,
+ fragment_content=parse_input.fragment_content,
+ full_output_dir=full_output_dir,
+ internal_output_filename=parse_input.internal_output_filename,
+ job_id=parse_input.job_id,
+ kb_dir=parse_input.kb_dir,
+ output_dir=parse_input.output_dir,
+ profile=profile,
+ relative_root=relative_root,
+ s3_key=parse_input.s3_key,
+ )
+
+
+def build_parse_session(parse_input: ParseInput) -> ParseSession:
+ """Build the parser routing session from explicit parse inputs."""
+ parse_options = parse_input.options
+ base_llm_paras = {
+ "llm_histories": parse_options.llm_histories,
+ "smart_title_parse": parse_options.smart_title_parse,
+ "summary_image": parse_options.summary_image,
+ "summary_table": parse_options.summary_table,
+ "summary_txt": parse_options.summary_txt,
+ "stopwords": parse_options.stopwords,
+ "doc_type": parse_options.doc_type,
+ "frag_desc": parse_options.add_frag_desc,
+ "model_name": settings.NORMOL_MODEL,
+ "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL,
+ }
+
+ logger.debug(f"baseurl: {parse_input.base_url}")
+ logger.debug(f"file_full_path: {parse_input.file_full_path}")
+
+ relative_root, full_output_dir = _resolve_output_paths(
+ filename=parse_input.filename,
+ internal_output_filename=parse_input.internal_output_filename,
+ kb_dir=parse_input.kb_dir,
+ output_dir=parse_input.output_dir,
+ )
+ logger.debug(f"relative_root: {relative_root}")
+ logger.debug(f"full_output_dir: {full_output_dir}")
+
+ with stage_timer("document.profile", filename=parse_input.filename):
+ profile = profile_document(
+ parse_input.file_full_path,
+ parse_input.internal_output_filename,
+ )
+ logger.info(f"📋 DocProfile: {profile.summary()}")
+ logger.debug(f"📋 Reasoning: {profile.reasoning}")
+
+ if profile.atlas_candidate and profile.doc_category not in ("atlas", "ppt_converted"):
+ logger.info(
+ f"🔍 Atlas candidate detected, running VLM visual check for {parse_input.filename}"
+ )
+ with stage_timer("document.atlas_vlm_check", filename=parse_input.filename):
+ vlm_is_atlas = classify_atlas_with_vlm(parse_input.file_full_path)
+ if vlm_is_atlas:
+ profile.doc_category = "atlas"
+ profile.reasoning += " | vlm_confirmed_atlas=True"
+ logger.info(f"✅ VLM confirmed atlas for {parse_input.filename}")
+ else:
+ profile.reasoning += " | vlm_confirmed_atlas=False"
+ logger.info(
+ f"ℹ️ VLM rejected atlas for {parse_input.filename}, routing as generic"
+ )
+
+ if profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT:
+ raise ValidationException(
+ user_message=(
+ f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. "
+ "Please split the document and upload in smaller batches."
+ ),
+ violations=[
+ {
+ "field": "page_count",
+ "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}",
+ }
+ ],
+ )
+
+ if profile.doc_category == "atlas":
+ filename, internal_output_filename, relative_root, full_output_dir = (
+ _rename_atlas_output(
+ filename=parse_input.filename,
+ internal_output_filename=parse_input.internal_output_filename,
+ kb_dir=parse_input.kb_dir,
+ output_dir=parse_input.output_dir,
+ )
+ )
+ logger.info(f"📐 Atlas output renamed: {filename}")
+ parse_input = ParseInput(
+ file_full_path=parse_input.file_full_path,
+ filename=filename,
+ output_dir=parse_input.output_dir,
+ internal_output_filename=internal_output_filename,
+ job_id=parse_input.job_id,
+ kb_dir=parse_input.kb_dir,
+ options=parse_input.options,
+ base_url=parse_input.base_url,
+ fragment_content=parse_input.fragment_content,
+ s3_key=parse_input.s3_key,
+ )
+
+ return ParseSession.from_input(
+ parse_input=parse_input,
+ base_llm_paras=base_llm_paras,
+ full_output_dir=full_output_dir,
+ profile=profile,
+ relative_root=relative_root,
+ )
+
+
+def _rename_atlas_output(
+ *,
+ filename: str,
+ internal_output_filename: str,
+ kb_dir: str,
+ output_dir: str,
+) -> tuple[str, str, str, str]:
+ name_base, _ = os.path.splitext(filename)
+ internal_name_base, _ = os.path.splitext(internal_output_filename)
+ atlas_filename = name_base + ".atlas"
+ atlas_internal_filename = internal_name_base + ".atlas"
+ relative_root, full_output_dir = _resolve_output_paths(
+ filename=atlas_filename,
+ internal_output_filename=atlas_internal_filename,
+ kb_dir=kb_dir,
+ output_dir=output_dir,
+ )
+ return atlas_filename, atlas_internal_filename, relative_root, full_output_dir
+
+
+def _resolve_output_paths(
+ *,
+ filename: str,
+ internal_output_filename: str,
+ kb_dir: str,
+ output_dir: str,
+) -> tuple[str, str]:
+ split_char = settings.SPLIT_CHAR or "/"
+ kb_dir_parts = kb_dir.split(split_char)
+
+ if filename and "images" not in kb_dir_parts:
+ relative_root = "/".join(kb_dir_parts + [filename])
+ else:
+ relative_root = "/".join(kb_dir_parts)
+
+ if internal_output_filename and "images" not in kb_dir_parts:
+ internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename])
+ else:
+ internal_relative_root = "/".join(kb_dir_parts)
+
+ full_output_dir = os.path.join(
+ output_dir,
+ internal_relative_root.replace("/", os.sep),
+ )
+ sanitized_output_dir = path_handle(full_output_dir, mode="sanitize")
+ if not isinstance(sanitized_output_dir, str) or not sanitized_output_dir:
+ raise ValueError(f"Failed to sanitize parser output directory: {full_output_dir}")
+ os.makedirs(sanitized_output_dir, exist_ok=True)
+
+ logger.debug(f"internal_relative_root: {internal_relative_root}")
+ return relative_root, sanitized_output_dir
diff --git a/apps/worker/app/services/document_parser/orchestration/postprocess.py b/apps/worker/app/services/document_parser/orchestration/postprocess.py
new file mode 100644
index 000000000..9cebc7c12
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/postprocess.py
@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import os
+import re
+
+import pandas as pd
+from app.services.document_parser.image_compressor import (
+ apply_rename_map_to_dataframe,
+ compress_output_images,
+)
+from app.services.document_parser.stage_profiler import stage_timer
+from loguru import logger
+
+
+def apply_parse_postprocess(
+ output_dir: str,
+ parsed_df: pd.DataFrame | None,
+) -> pd.DataFrame | None:
+ """Apply output cleanup and image compression after parsing."""
+ logger.debug(f"full_output_dir: {output_dir}")
+
+ with stage_timer("document.cleanup_unreferenced_images", output_dir=output_dir):
+ cleanup_unreferenced_images(output_dir)
+
+ with stage_timer("document.compress_images", output_dir=output_dir):
+ compress_stats = compress_output_images(output_dir)
+ if compress_stats.processed > 0:
+ logger.info(
+ f"📦 Image compression: {compress_stats.processed} processed "
+ f"({compress_stats.converted_png_to_jpg} PNG→JPG, "
+ f"{compress_stats.resized} resized), "
+ f"{compress_stats.bytes_before / 1024 / 1024:.1f}MB → "
+ f"{compress_stats.bytes_after / 1024 / 1024:.1f}MB"
+ )
+ if compress_stats.rename_map and parsed_df is not None:
+ return apply_rename_map_to_dataframe(parsed_df, compress_stats.rename_map)
+
+ return parsed_df
+
+
+def cleanup_unreferenced_images(output_dir: str) -> int:
+ """Remove UUID-named images that are not referenced by final parsed output."""
+ image_dir = os.path.join(output_dir, "images")
+ if not os.path.isdir(image_dir):
+ return 0
+
+ uuid_pattern = re.compile(
+ r"^[a-f0-9]{64}\.(?:jpg|jpeg|png|gif|webp)$",
+ re.IGNORECASE,
+ )
+ removed_count = 0
+
+ for filename in os.listdir(image_dir):
+ if not uuid_pattern.match(filename):
+ continue
+
+ file_path = os.path.join(image_dir, filename)
+ try:
+ os.remove(file_path)
+ removed_count += 1
+ logger.debug(f"Removed unreferenced image: {filename}")
+ except OSError as exc:
+ logger.warning(f"Failed to remove {filename}: {exc}")
+
+ if removed_count > 0:
+ logger.info(
+ f"Cleaned up {removed_count} unreferenced UUID-named images from {image_dir}"
+ )
+
+ return removed_count
diff --git a/apps/worker/app/services/document_parser/orchestration/route_parse.py b/apps/worker/app/services/document_parser/orchestration/route_parse.py
new file mode 100644
index 000000000..85402339e
--- /dev/null
+++ b/apps/worker/app/services/document_parser/orchestration/route_parse.py
@@ -0,0 +1,14 @@
+import pandas as pd
+
+from app.services.document_parser.orchestration.format_router import (
+ get_document_parse_adapter,
+ resolve_document_format,
+)
+from app.services.document_parser.orchestration.parse_session import ParseSession
+
+
+def route_document_parse(session: ParseSession) -> tuple[str, pd.DataFrame | None]:
+ """Route a parser session to the correct adapter and return its output."""
+ document_format = resolve_document_format(session.file_full_path)
+ adapter = get_document_parse_adapter(document_format)
+ return adapter.parse(session)
diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py
index ad5058e47..000cef7bf 100644
--- a/apps/worker/app/services/document_parser/parse_service.py
+++ b/apps/worker/app/services/document_parser/parse_service.py
@@ -1,68 +1,11 @@
-# pyright: reportArgumentType=false, reportReturnType=false
-"""
-main parsing service
-"""
-
-import os
-import re
+"""Stable parser seam backed by dedicated orchestration modules."""
import pandas as pd
-from app.services.document_parser.atlas_classifier import classify_atlas_with_vlm
-
-# document_parser imports
-from app.services.document_parser.doc_profiler import profile_document
-from app.services.document_parser.stage_profiler import stage_timer
-from loguru import logger
-
-from shared.core.config import settings
-from shared.core.exceptions.domain_exceptions import (
- ValidationException,
-)
-from shared.utils.file_utils import path_handle
-
-
-def cleanup_unreferenced_images(output_dir: str) -> int:
- """
- Clean up unreferenced UUID-named images from the images directory.
-
- After document parsing (PDF, DOCX, PPTX, etc.), the images/ directory may contain:
- 1. Processed images: renamed with semantic names like 'image-0-xxx.jpg'
- 2. Unreferenced images: UUID-named (64-char hex) that were parsed as tables/formulas
-
- This function removes the unreferenced UUID-named images to reduce final package size.
- Args:
- output_dir: The full output directory path
-
- Returns:
- Number of files removed
- """
- img_dir = os.path.join(output_dir, "images")
- if not os.path.isdir(img_dir):
- return 0
-
- # UUID pattern: 64 hex characters followed by image extension
- uuid_pattern = re.compile(
- r"^[a-f0-9]{64}\.(?:jpg|jpeg|png|gif|webp)$", re.IGNORECASE
- )
-
- removed_count = 0
- for filename in os.listdir(img_dir):
- if uuid_pattern.match(filename):
- file_path = os.path.join(img_dir, filename)
- try:
- os.remove(file_path)
- removed_count += 1
- logger.debug(f"Removed unreferenced image: {filename}")
- except OSError as e:
- logger.warning(f"Failed to remove {filename}: {e}")
-
- if removed_count > 0:
- logger.info(
- f"Cleaned up {removed_count} unreferenced UUID-named images from {img_dir}"
- )
-
- return removed_count
+from app.services.document_parser.orchestration.parse_input import ParseInput, ParseOptions
+from app.services.document_parser.orchestration.parse_session import build_parse_session
+from app.services.document_parser.orchestration.postprocess import apply_parse_postprocess
+from app.services.document_parser.orchestration.route_parse import route_document_parse
def checkerboard_inject_parse(
@@ -84,374 +27,29 @@ def checkerboard_inject_parse(
fragment_content: str = "",
s3_key: str | None = None,
) -> tuple[str, pd.DataFrame | None]:
- """
- main parsing function
-
- Args:
- file_full_path: source file path (local or URL)
- filename: file name
- output_dir: output directory (absolute path, caller provides)
- kb_dir: sub-directory name
- llm_histories: retained for downstream LLM settings
- smart_title_parse: enable smart heading parsing
- summary_image: enable image summaries
- summary_table: enable table summaries
- summary_txt: enable text summaries
- stopwords: optional stopword list
- doc_type: parser document type hint
- add_frag_desc: extra fragment description
- base_url: optional source base URL
- fragment_content: raw fragment content
- job_id: optional job identifier used for parser artifacts
- internal_output_filename: normalized internal folder name for on-disk output
- s3_key: optional S3 key for downstream parsers
-
- Returns:
- tuple: (output_dir, parsed_df)
- - output_dir: directory path after parsing
- - parsed_df: parsed content DataFrame
- """
- # Build base_llm_paras from explicit parameters
- base_llm_paras = {
- "llm_histories": llm_histories,
- "smart_title_parse": smart_title_parse,
- "summary_image": summary_image,
- "summary_table": summary_table,
- "summary_txt": summary_txt,
- "stopwords": stopwords,
- "doc_type": doc_type,
- "frag_desc": add_frag_desc,
- "model_name": settings.NORMOL_MODEL,
- "hierarchy_model_name": settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL,
- }
-
- baseurl = base_url
-
- logger.debug(f"baseurl: {baseurl}")
- logger.debug(f"file_full_path: {file_full_path}")
-
- # ========== Path handling ==========
- split_char = settings.SPLIT_CHAR or "/"
-
- # Develop relative root path for chunk path field
- kb_dir_parts = kb_dir.split(split_char)
- if filename and "images" not in kb_dir_parts:
- relative_root = "/".join(kb_dir_parts + [filename])
- else:
- relative_root = "/".join(kb_dir_parts)
-
- if internal_output_filename and "images" not in kb_dir_parts:
- internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename])
- else:
- internal_relative_root = "/".join(kb_dir_parts)
-
- # Develop full output directory (output_dir + relative_root)
- full_output_dir = os.path.join(
- output_dir, internal_relative_root.replace("/", os.sep)
- )
- full_output_dir = path_handle(full_output_dir, mode="sanitize")
- os.makedirs(full_output_dir, exist_ok=True)
-
- logger.debug(f"relative_root: {relative_root}")
- logger.debug(f"internal_relative_root: {internal_relative_root}")
- logger.debug(f"full_output_dir: {full_output_dir}")
-
- file_path_lower = file_full_path.lower()
- parsed_df = None
-
- # ── Agentic Profiler: classify document before routing ──
- with stage_timer("document.profile", filename=filename):
- profile = profile_document(file_full_path, internal_output_filename)
- logger.info(f"📋 DocProfile: {profile.summary()}")
- logger.debug(f"📋 Reasoning: {profile.reasoning}")
-
- # ── VLM second-pass: confirm atlas_candidate with visual check ──
- # Heuristics can miss atlases that have a rich OCR text layer on top of
- # scanned drawing pages (avg_text_density too high). VLM sees the actual
- # page layout and makes the final call.
- if profile.atlas_candidate and profile.doc_category not in (
- "atlas",
- "ppt_converted",
- ):
- logger.info(
- f"🔍 Atlas candidate detected, running VLM visual check for {filename}"
- )
- with stage_timer("document.atlas_vlm_check", filename=filename):
- vlm_is_atlas = classify_atlas_with_vlm(file_full_path)
- if vlm_is_atlas:
- profile.doc_category = "atlas"
- profile.reasoning += " | vlm_confirmed_atlas=True"
- logger.info(f"✅ VLM confirmed atlas for {filename}")
- else:
- profile.reasoning += " | vlm_confirmed_atlas=False"
- logger.info(f"ℹ️ VLM rejected atlas for {filename}, routing as generic")
-
- # ── Page count guard: reject oversized PDFs before routing ──
- PDF_PAGE_LIMIT = 600
- if profile and profile.file_type == "pdf" and profile.page_count > PDF_PAGE_LIMIT:
- raise ValidationException(
- user_message=(
- f"Document too large: {profile.page_count} pages exceeds the {PDF_PAGE_LIMIT}-page limit. "
- f"Please split the document and upload in smaller batches."
- ),
- violations=[
- {
- "field": "page_count",
- "description": f"PDF has {profile.page_count} pages, limit is {PDF_PAGE_LIMIT}",
- }
- ],
- )
-
- # Atlas routing: rename output folder from .pdf → .atlas for easy filtering
- if profile and profile.doc_category == "atlas":
- name_base, _ = os.path.splitext(filename)
- internal_name_base, _ = os.path.splitext(internal_output_filename)
- filename = name_base + ".atlas"
- internal_output_filename = internal_name_base + ".atlas"
- relative_root = "/".join(kb_dir_parts + [filename])
- internal_relative_root = "/".join(kb_dir_parts + [internal_output_filename])
- full_output_dir = os.path.join(
- output_dir, internal_relative_root.replace("/", os.sep)
- )
- full_output_dir = path_handle(full_output_dir, mode="sanitize")
- os.makedirs(full_output_dir, exist_ok=True)
- logger.info(f"📐 Atlas output renamed: {filename}")
-
- if ".fragment" in file_path_lower:
- logger.debug("file type is fragment")
- from app.services.document_parser.fragment_parser import parse_fragment
-
- full_output_dir, relative_root, parsed_df = parse_fragment(
- fragment_content,
- filename=filename,
- output_dir=output_dir,
- kb_dir=kb_dir,
- base_llm_paras=base_llm_paras,
- )
-
- elif ".txt" in file_path_lower:
- logger.debug("file type is txt")
- from app.services.document_parser.md_parser import parse_md
- from app.services.document_parser.txt_parser import parse_texts
-
- txt_lines = parse_texts(file_path=file_full_path, baseurl=baseurl)
- parsed_df = parse_md(
- full_output_dir,
- source_type="md",
- md_lines=txt_lines,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif (
- ".png" in file_path_lower
- or ".jpg" in file_path_lower
- or ".jpeg" in file_path_lower
- ):
- logger.debug("file type is image")
- from app.services.document_parser.image_parser import parse_image
-
- parsed_df = parse_image(
- file_full_path,
- filename=filename,
- output_dir=full_output_dir,
- baseurl=baseurl,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif ".pdf" in file_path_lower:
- logger.debug("file type is pdf")
- from app.services.document_parser.pdf_parser import parse_pdfs
-
- if filename and file_full_path:
- parsed_df = parse_pdfs(
- file_full_path,
- filename=filename,
- output_dir=full_output_dir,
- base_llm_paras=base_llm_paras,
- profile=profile,
- relative_root=relative_root,
- s3_key=s3_key,
- )
-
- elif ".doc" in file_path_lower and ".docx" not in file_path_lower:
- logger.debug("file type is doc")
- from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx
- from app.services.document_parser.legacy_converter import doc_to_docx
-
- if filename and file_full_path:
- converted_docx_path, _ = doc_to_docx(file_full_path, outdir=full_output_dir)
- parsed_structure, df_list = parse_docx(
- converted_docx_path,
- base_llm_paras,
- full_output_dir,
- filename,
- baseurl,
- relative_root=relative_root,
- )
- parsed_df = convert_doc2dics(
- parsed_structure,
- df_list,
- full_output_dir,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif ".docx" in file_path_lower:
- logger.debug("file type is docx")
- from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx
-
- if filename and file_full_path:
- parsed_structure, df_list = parse_docx(
- file_full_path,
- base_llm_paras,
- full_output_dir,
- filename,
- baseurl,
- relative_root=relative_root,
- )
- parsed_df = convert_doc2dics(
- parsed_structure,
- df_list,
- full_output_dir,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif ".xls" in file_path_lower and ".xlsx" not in file_path_lower:
- logger.debug("file type is xls")
- from app.services.document_parser.legacy_converter import xls_to_xlsx
- from app.services.document_parser.table_parser import parse_xlsx
-
- if filename and file_full_path:
- converted_xlsx_path, _ = xls_to_xlsx(file_full_path, outdir=full_output_dir)
- parsed_df = parse_xlsx(
- converted_xlsx_path,
- filename,
- full_output_dir,
- baseurl,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif ".xlsx" in file_path_lower:
- logger.debug("file type is xlsx")
- from app.services.document_parser.table_parser import parse_xlsx
-
- if filename and file_full_path:
- parsed_df = parse_xlsx(
- file_full_path,
- filename,
- full_output_dir,
- baseurl,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif ".pptx" in file_path_lower:
- logger.debug("file type is pptx")
- from app.services.document_parser.pptx_parser import parse_pptx
-
- if filename and file_full_path:
- # ====== iLoveAPI PPTX → PDF → MinerU (default production route) ======
- parsed_df = parse_pptx(
- file_full_path,
- filename=filename,
- output_dir=full_output_dir,
- base_llm_paras=base_llm_paras,
- strategy="to_pdf_api",
- job_id=job_id,
- relative_root=relative_root,
- baseurl=baseurl,
- )
-
- # ====== [EXPERIMENTAL] Directly send PPTX to MinerU via parse_pdfs ======
- # Uncomment the block below (and comment out parse_pptx above) to bypass iLoveAPI
- # from app.services.document_parser.pdf_parser import parse_pdfs
- # parsed_df = parse_pdfs(
- # file_full_path,
- # filename=filename,
- # output_dir=full_output_dir,
- # base_llm_paras=base_llm_paras,
- # profile=profile,
- # relative_root=relative_root,
- # s3_key=s3_key
- # )
-
- elif ".md" in file_path_lower:
- logger.debug("file type is md")
- from app.services.document_parser.md_parser import parse_md
-
- if filename and file_full_path:
- parsed_df = parse_md(
- full_output_dir,
- source_type="md",
- file_path=file_full_path,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- )
-
- elif ".json" in file_path_lower:
- logger.debug("file type is json")
- # JSON parsing not yet implemented
-
- else:
- # Unsupported file type
- file_ext = os.path.splitext(file_full_path)[1].lower()
- supported_types = [
- ".txt",
- ".fragment",
- ".png",
- ".jpg",
- ".jpeg",
- ".pdf",
- ".doc",
- ".docx",
- ".xls",
- ".xlsx",
- ".pptx",
- ".md",
- ".json",
- ]
- raise ValidationException(
- user_message=f"Unsupported file type: {file_ext}",
- violations=[
- {
- "field": "file_type",
- "description": f"Must be one of: {', '.join(supported_types)}",
- }
- ],
- )
-
- logger.debug(f"full_output_dir: {full_output_dir}")
-
- # Post-processing: clean up unreferenced UUID-named images
- with stage_timer(
- "document.cleanup_unreferenced_images", output_dir=full_output_dir
- ):
- cleanup_unreferenced_images(full_output_dir)
-
- # Post-processing: compress output images (PNG→JPEG, resize oversized)
- from app.services.document_parser.image_compressor import (
- apply_rename_map_to_dataframe,
- compress_output_images,
+ """Run the stable parser seam using dedicated orchestration modules."""
+ parse_input = ParseInput(
+ file_full_path=file_full_path,
+ filename=filename,
+ internal_output_filename=internal_output_filename,
+ job_id=job_id,
+ kb_dir=kb_dir,
+ output_dir=output_dir,
+ options=ParseOptions(
+ add_frag_desc=add_frag_desc,
+ doc_type=doc_type,
+ llm_histories=llm_histories,
+ smart_title_parse=smart_title_parse,
+ stopwords=stopwords,
+ summary_image=summary_image,
+ summary_table=summary_table,
+ summary_txt=summary_txt,
+ ),
+ base_url=base_url,
+ fragment_content=fragment_content,
+ s3_key=s3_key,
)
-
- with stage_timer("document.compress_images", output_dir=full_output_dir):
- compress_stats = compress_output_images(full_output_dir)
- if compress_stats.processed > 0:
- logger.info(
- f"📦 Image compression: {compress_stats.processed} processed "
- f"({compress_stats.converted_png_to_jpg} PNG→JPG, "
- f"{compress_stats.resized} resized), "
- f"{compress_stats.bytes_before / 1024 / 1024:.1f}MB → "
- f"{compress_stats.bytes_after / 1024 / 1024:.1f}MB"
- )
- # Update DataFrame references when PNG→JPG conversions occurred
- if compress_stats.rename_map and parsed_df is not None:
- parsed_df = apply_rename_map_to_dataframe(
- parsed_df, compress_stats.rename_map
- )
-
+ session = build_parse_session(parse_input)
+ full_output_dir, parsed_df = route_document_parse(session)
+ parsed_df = apply_parse_postprocess(full_output_dir, parsed_df)
return full_output_dir, parsed_df
diff --git a/apps/worker/app/services/document_parser/parser_rows.py b/apps/worker/app/services/document_parser/parser_rows.py
new file mode 100644
index 000000000..839ddd9f6
--- /dev/null
+++ b/apps/worker/app/services/document_parser/parser_rows.py
@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pandas as pd
+from pandas import Index
+
+from shared.core.config import settings
+
+PARSER_ROW_COLUMNS: tuple[str, ...] = tuple(settings.ALL_DF_COLS.split(","))
+
+
+@dataclass(frozen=True)
+class ParsedRow:
+ content: str
+ path: str
+ type: str
+ know_id: str
+ addtime: str
+ keywords: str = ""
+ summary: str = ""
+ tokens: str = ""
+ connectto: str = ""
+ page_nums: str = ""
+ length: int | None = None
+
+ def to_list(self) -> list[object]:
+ content_length = self.length if self.length is not None else len(self.content)
+ return [
+ self.content,
+ self.path,
+ self.type,
+ content_length,
+ self.keywords,
+ self.summary,
+ self.know_id,
+ self.tokens,
+ self.connectto,
+ self.addtime,
+ self.page_nums,
+ ]
+
+ def to_dict(self) -> dict[str, object]:
+ return dict(zip(PARSER_ROW_COLUMNS, self.to_list()))
+
+
+class ParsedRowsBuilder:
+ def __init__(self) -> None:
+ self._rows: list[ParsedRow] = []
+
+ def append(self, row: ParsedRow) -> None:
+ self._rows.append(row)
+
+ def extend(self, rows: list[ParsedRow]) -> None:
+ self._rows.extend(rows)
+
+ def to_dataframe(self) -> pd.DataFrame:
+ return pd.DataFrame(
+ [row.to_list() for row in self._rows],
+ columns=Index(PARSER_ROW_COLUMNS),
+ )
diff --git a/apps/worker/app/services/document_parser/path_helpers.py b/apps/worker/app/services/document_parser/path_helpers.py
new file mode 100644
index 000000000..41965e125
--- /dev/null
+++ b/apps/worker/app/services/document_parser/path_helpers.py
@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import os
+import re
+from typing import Any
+
+from bs4 import BeautifulSoup
+from shared.utils.chunk_refs import extract_chunk_refs
+from shared.utils.file_utils import path_handle
+
+SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3")
+
+
+def find_images(folder_path: str) -> list[str]:
+ """Find image files inside a folder tree."""
+ image_extensions = {".png", ".jpg", ".jpeg"}
+ image_files: list[str] = []
+
+ for _, _, files in os.walk(folder_path):
+ files.sort()
+ for file in files:
+ if os.path.splitext(file)[1].lower() in image_extensions:
+ image_files.append(file)
+ return image_files
+
+
+def find_matches_parsing(content: str, path: str) -> str:
+ """Parse table and image markers from content."""
+ matches = extract_chunk_refs(content)
+ match_type = "PTXT" if len(matches) == 0 else "\n".join((["PTXT"] + matches))
+
+ split_char = os.getenv("SPLIT_CHAR", "/")
+ if any(
+ f"{split_char}{summary_marker}" in path
+ for summary_marker in SUMMARY_PATH_MARKERS
+ ):
+ parent_path = path.split(split_char)[-2]
+ match_type = "SUMMARY_" + parent_path + "_SUMMARY"
+ return match_type
+
+
+def flatten_dic2paths(
+ d: dict[str, Any],
+ current_path: list[str] | None = None,
+ result: list[str] | None = None,
+) -> list[str]:
+ """Flatten a nested dict into path strings."""
+ if result is None:
+ result = []
+ if current_path is None:
+ current_path = []
+
+ for key, value in d.items():
+ if not isinstance(key, str):
+ continue
+ new_path = current_path + [key]
+ if isinstance(value, dict) and value:
+ flatten_dic2paths(value, new_path, result)
+ else:
+ split_char = os.getenv("SPLIT_CHAR", "/")
+ result.append(split_char.join(new_path))
+ return result
+
+
+def process_path_texts(path_: str, last: int = 50) -> str:
+ """Normalize path text for downstream use."""
+ temp_path = path_handle(path_, mode="sanitize")
+ if not isinstance(temp_path, str) or temp_path == "":
+ return ""
+ return "_".join(temp_path.split(os.sep))[:last]
+
+
+def remove_spaces(text: str, handle_punctuation: bool = False) -> str:
+ """Remove spaces between Chinese chars while keeping English word spacing."""
+ if handle_punctuation:
+ punctuation = (
+ r"""!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~，。、【】《》？；：''""（）…—-！"""
+ )
+ res_text = re.sub(f"[{re.escape(punctuation)}]", "", text)
+ else:
+ pattern = re.compile(r"([\u4e00-\u9fff])\s+|(?<=\s)([\u4e00-\u9fff])")
+
+ def replacer(match: re.Match[str]) -> str:
+ return match.group(1) or match.group(2)
+
+ res_text = pattern.sub(replacer, text)
+
+ res_text = re.sub(r"\s+", " ", res_text)
+ return res_text.strip()
+
+
+def traverse_dict(d: dict[str, Any], parent: str | None = None) -> list[str]:
+ """Traverse a dictionary and generate description text."""
+ dic_texts: list[str] = []
+ for key, value in d.items():
+ if value:
+ child_keys = ", ".join(value.keys())
+ text = f"'{key}' includes {child_keys}"
+ dic_texts.append(text)
+ dic_texts.extend(traverse_dict(value, key))
+ return dic_texts
+
+
+def restore_graph_by_paths(paths: list[str]) -> tuple[dict[str, Any], list[str]]:
+ """Rebuild a graph structure from path strings."""
+ root_dict: dict[str, Any] = {}
+ split_char = os.getenv("SPLIT_CHAR", "/")
+ for path in paths:
+ nodes = path.split(split_char)
+ current_dict = root_dict
+ for node in nodes:
+ if node not in current_dict:
+ current_dict[node] = {}
+ current_dict = current_dict[node]
+ dic_texts = traverse_dict(root_dict)
+ return root_dict, dic_texts
+
+
+def html2txt(html_text: str) -> str:
+ """Convert HTML into plain text."""
+ soup = BeautifulSoup(html_text, "html.parser")
+ return soup.get_text()
diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py
index dde943110..19063937a 100755
--- a/apps/worker/app/services/document_parser/pptx_parser.py
+++ b/apps/worker/app/services/document_parser/pptx_parser.py
@@ -6,20 +6,17 @@
import jwt
import requests
-from app.services.common.kb_utils import find_images
+from app.services.document_parser.path_helpers import find_images
from app.services.document_parser.legacy_converter import (
_convert_with_libreoffice,
)
from app.services.document_parser.md_parser import parse_md
-from app.services.document_parser.mineru_pdf_service import (
- get_existing_mineru_source_s3_key,
-)
from app.services.document_parser.parser_log_utils import truncate_log_value
-from app.services.document_parser.pdf_parser import parse_pdfs
-from app.services.document_parser.pptx_pdf_rendering import (
- render_pdf_to_image_pdf as _render_pdf_to_image_pdf,
+from app.services.document_parser.rendered_pdf_transform import (
+ build_rendered_pdf_s3_key,
+ parse_cached_rendered_pdf,
+ parse_rendered_pdf_bytes,
)
-from app.services.storage.sync_storage_service import download_s3_object_to_temp
from loguru import logger
from markitdown import MarkItDown
from pptx2md import ConversionConfig, convert
@@ -29,7 +26,7 @@
FileSystemException,
)
from shared.core.logging import LogEvent
-from shared.utils.CommonHelperSync import load_file_bytes
+from shared.utils.file_loading import load_file_bytes
from shared.utils.file_utils import path_handle
# ==================== LibreOffice conversion ====================
@@ -306,50 +303,6 @@ class _ILoveApiConcurrencyExceeded(Exception):
pass
-def _build_rendered_pdf_s3_key(job_id: str | None) -> str | None:
- """Store rendered parser artifacts under a stable transform/ prefix."""
- if settings.ENVIRONMENT == "development" or not job_id:
- return None
- return f"transform/{job_id}.rendered.pdf"
-
-
-def _parse_cached_rendered_pdf(
- rendered_pdf_s3_key: str | None,
- filename: str,
- output_dir: str,
- base_llm_paras,
- relative_root,
-):
- """Parse a previously rendered PPTX PDF from S3 without re-reading the source deck."""
- if rendered_pdf_s3_key is None:
- return None
-
- cached_rendered_pdf_s3_key = get_existing_mineru_source_s3_key(rendered_pdf_s3_key)
- if cached_rendered_pdf_s3_key is None:
- return None
-
- logger.info(
- f"[parse_pptx] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}"
- )
- cached_rendered_pdf_path = download_s3_object_to_temp(
- cached_rendered_pdf_s3_key,
- suffix=".pdf",
- temp_dir=output_dir,
- )
- try:
- return parse_pdfs(
- cached_rendered_pdf_path,
- filename,
- output_dir,
- base_llm_paras,
- relative_root=relative_root,
- s3_key=cached_rendered_pdf_s3_key,
- )
- finally:
- if os.path.exists(cached_rendered_pdf_path):
- os.remove(cached_rendered_pdf_path)
-
-
# ==================== main parsing entrance ====================
@@ -372,12 +325,12 @@ def parse_pptx(
- "to_pdf_api": use iLoveAPI to convert to PDF, then parse via MinerU (recommended)
"""
rendered_pdf_s3_key = (
- _build_rendered_pdf_s3_key(job_id)
+ build_rendered_pdf_s3_key(job_id)
if strategy in {"to_pdf_api", "to_pdf"}
else None
)
if strategy in {"to_pdf_api", "to_pdf"}:
- cached_result = _parse_cached_rendered_pdf(
+ cached_result = parse_cached_rendered_pdf(
rendered_pdf_s3_key=rendered_pdf_s3_key,
filename=filename,
output_dir=output_dir,
@@ -491,27 +444,14 @@ def _parse_pptx_via_api(
# Step 1: PPTX → PDF (in memory)
pdf_bytes = _pptx_bytes_to_pdf_bytes(pptx_data, filename)
- # Step 2: PDF → image-only PDF (in memory)
- img_pdf_bytes = _render_pdf_to_image_pdf(pdf_bytes)
-
- # Step 3: Write to output_dir for MinerU upload, then clean up
- tmp_path = os.path.join(output_dir, "_pptx_tmp.pdf")
- with open(tmp_path, "wb") as f:
- f.write(img_pdf_bytes)
-
- try:
- parsed_df = parse_pdfs(
- tmp_path,
- filename=filename,
- output_dir=output_dir,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- s3_key=rendered_pdf_s3_key,
- )
- return parsed_df
- finally:
- if os.path.exists(tmp_path):
- os.remove(tmp_path)
+ return parse_rendered_pdf_bytes(
+ pdf_bytes=pdf_bytes,
+ filename=filename,
+ output_dir=output_dir,
+ base_llm_paras=base_llm_paras,
+ relative_root=relative_root,
+ rendered_pdf_s3_key=rendered_pdf_s3_key,
+ )
def _parse_pptx_via_libreoffice(
@@ -552,25 +492,14 @@ def _parse_pptx_via_libreoffice(
finally:
shutil.rmtree(tmp_dir, ignore_errors=True)
- img_pdf_bytes = _render_pdf_to_image_pdf(pdf_bytes)
-
- tmp_path = os.path.join(output_dir, "_pptx_tmp.pdf")
- with open(tmp_path, "wb") as f:
- f.write(img_pdf_bytes)
-
- try:
- parsed_df = parse_pdfs(
- tmp_path,
- filename=filename,
- output_dir=output_dir,
- base_llm_paras=base_llm_paras,
- relative_root=relative_root,
- s3_key=rendered_pdf_s3_key,
- )
- return parsed_df
- finally:
- if os.path.exists(tmp_path):
- os.remove(tmp_path)
+ return parse_rendered_pdf_bytes(
+ pdf_bytes=pdf_bytes,
+ filename=filename,
+ output_dir=output_dir,
+ base_llm_paras=base_llm_paras,
+ relative_root=relative_root,
+ rendered_pdf_s3_key=rendered_pdf_s3_key,
+ )
def _parse_pptx_to_md(pptx_data, filename, output_dir, base_llm_paras, relative_root):
diff --git a/apps/worker/app/services/document_parser/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/pymupdf_subprocess.py
index c98b4678c..88bec26af 100644
--- a/apps/worker/app/services/document_parser/pymupdf_subprocess.py
+++ b/apps/worker/app/services/document_parser/pymupdf_subprocess.py
@@ -24,6 +24,7 @@
from multiprocessing.process import BaseProcess
from multiprocessing.queues import Queue as MultiprocessingQueue
from threading import RLock
+from typing import TYPE_CHECKING
from app.core.runtime_limits import read_pymupdf_max_concurrent
from loguru import logger
@@ -33,7 +34,8 @@
TimeoutException,
)
-from gevent.threadpool import ThreadPool as GeventThreadPool
+if TYPE_CHECKING:
+ from gevent.threadpool import ThreadPool as GeventThreadPool
# Default timeout for child processes (seconds)
DEFAULT_TIMEOUT = 3000
diff --git a/apps/worker/app/services/document_parser/rendered_pdf_transform.py b/apps/worker/app/services/document_parser/rendered_pdf_transform.py
new file mode 100644
index 000000000..cab7ca822
--- /dev/null
+++ b/apps/worker/app/services/document_parser/rendered_pdf_transform.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import os
+
+import pandas as pd
+
+from app.services.document_parser.mineru_pdf_service import (
+ get_existing_mineru_source_s3_key,
+)
+from app.services.document_parser.pdf_parser import parse_pdfs
+from app.services.document_parser.pptx_pdf_rendering import render_pdf_to_image_pdf
+from loguru import logger
+
+from shared.core.config import settings
+from shared.services.storage.job_file_storage import JobFileStorage
+
+RENDERED_PDF_TEMP_FILENAME = "_pptx_tmp.pdf"
+
+
+def build_rendered_pdf_s3_key(job_id: str | None) -> str | None:
+ """Store rendered parser artifacts under a stable transform/ prefix."""
+ if settings.ENVIRONMENT == "development" or not job_id:
+ return None
+ return f"transform/{job_id}.rendered.pdf"
+
+
+def parse_cached_rendered_pdf(
+ *,
+ rendered_pdf_s3_key: str | None,
+ filename: str,
+ output_dir: str,
+ base_llm_paras: dict[str, object],
+ relative_root: str | None,
+) -> pd.DataFrame | None:
+ """Parse a previously rendered PDF from S3 without re-reading the source deck."""
+ if rendered_pdf_s3_key is None:
+ return None
+
+ cached_rendered_pdf_s3_key = get_existing_mineru_source_s3_key(rendered_pdf_s3_key)
+ if cached_rendered_pdf_s3_key is None:
+ return None
+
+ logger.info(
+ f"[rendered_pdf_transform] Reusing rendered PDF for MinerU URL mode: {rendered_pdf_s3_key}"
+ )
+ cached_rendered_pdf_path = JobFileStorage().download_upload_to_temp(
+ cached_rendered_pdf_s3_key,
+ suffix=".pdf",
+ temp_dir=output_dir,
+ )
+ try:
+ return parse_pdfs(
+ cached_rendered_pdf_path,
+ filename,
+ output_dir,
+ base_llm_paras,
+ relative_root=relative_root,
+ s3_key=cached_rendered_pdf_s3_key,
+ )
+ finally:
+ if os.path.exists(cached_rendered_pdf_path):
+ os.remove(cached_rendered_pdf_path)
+
+
+def parse_rendered_pdf_bytes(
+ *,
+ pdf_bytes: bytes,
+ filename: str,
+ output_dir: str,
+ base_llm_paras: dict[str, object],
+ relative_root: str | None,
+ rendered_pdf_s3_key: str | None = None,
+) -> pd.DataFrame:
+ image_only_pdf_bytes = render_pdf_to_image_pdf(pdf_bytes)
+ temporary_pdf_path = os.path.join(output_dir, RENDERED_PDF_TEMP_FILENAME)
+ with open(temporary_pdf_path, "wb") as temporary_pdf_file:
+ temporary_pdf_file.write(image_only_pdf_bytes)
+
+ try:
+ return parse_pdfs(
+ temporary_pdf_path,
+ filename=filename,
+ output_dir=output_dir,
+ base_llm_paras=base_llm_paras,
+ relative_root=relative_root,
+ s3_key=rendered_pdf_s3_key,
+ )
+ finally:
+ if os.path.exists(temporary_pdf_path):
+ os.remove(temporary_pdf_path)
diff --git a/apps/worker/app/services/document_parser/table_asset_writer.py b/apps/worker/app/services/document_parser/table_asset_writer.py
new file mode 100644
index 000000000..91599f14d
--- /dev/null
+++ b/apps/worker/app/services/document_parser/table_asset_writer.py
@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from app.services.document_parser.parser_rows import ParsedRow
+
+
+@dataclass(frozen=True)
+class TableAssetInput:
+ html: str
+ output_dir: str
+ table_name: str
+ summary: str
+ keywords: str
+ know_id: str
+ addtime: str
+ page_nums: str = ""
+ content: str | None = None
+ tokens: str = ""
+ length: int | None = None
+
+
+def write_table_asset(table_input: TableAssetInput) -> ParsedRow:
+ table_dir = os.path.join(table_input.output_dir, "tables")
+ os.makedirs(table_dir, exist_ok=True)
+ table_filename = _ensure_html_extension(table_input.table_name)
+ table_path = os.path.join(table_dir, table_filename)
+ with open(table_path, "w", encoding="utf-8") as table_file:
+ table_file.write(table_input.html)
+ row_content = table_input.content if table_input.content is not None else table_input.html
+ return ParsedRow(
+ content=row_content,
+ path=f"tables/{table_filename}",
+ type="table",
+ keywords=table_input.keywords,
+ summary=table_input.summary,
+ know_id=table_input.know_id,
+ tokens=table_input.tokens,
+ connectto="",
+ addtime=table_input.addtime,
+ page_nums=table_input.page_nums,
+ length=table_input.length,
+ )
+
+
+def _ensure_html_extension(table_name: str) -> str:
+ return table_name if table_name.endswith(".html") else f"{table_name}.html"
diff --git a/apps/worker/app/services/document_parser/table_frame_parser.py b/apps/worker/app/services/document_parser/table_frame_parser.py
new file mode 100644
index 000000000..63f39feb8
--- /dev/null
+++ b/apps/worker/app/services/document_parser/table_frame_parser.py
@@ -0,0 +1,422 @@
+# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false, reportOperatorIssue=false, reportIndexIssue=false, reportAssignmentType=false, reportGeneralTypeIssues=false
+from __future__ import annotations
+
+import datetime
+import os
+import uuid
+from collections import OrderedDict
+
+import numpy as np
+import pandas as pd
+from app.services.document_parser.dataframe_html_renderer import df2html
+from app.services.document_parser.identifiers import gen_str_codes
+from app.services.document_parser.path_helpers import flatten_dic2paths
+from loguru import logger
+
+from shared.services.ai.prompt_service import build_prompt
+from shared.services.ai.response_process_service import eval_response
+from shared.utils.OpenAICompatibleClientSync import get_openai_client
+from shared.utils.text_utils import remove_duplicates_orderkept
+
+
+def parse_headers(
+ table_frame: pd.DataFrame,
+ paras: dict[str, object] | None = None,
+ header_window: int = 5,
+ smart_headers: bool = True,
+) -> pd.DataFrame:
+ llm_parameters = paras or {"summary_table": False}
+
+ def parse_headers_nonsmart(candidate_frame: pd.DataFrame) -> list[int]:
+ non_na_row = candidate_frame[candidate_frame.notna().any(axis=1)].head(1)
+ header_id = non_na_row.index[-1] if not non_na_row.empty else None
+ return list(range(header_id + 1))
+
+ if not pd.isna(table_frame.columns).all():
+ table_frame.loc[-1] = table_frame.columns
+ table_frame.index = table_frame.index + 1
+ table_frame = table_frame.sort_index()
+ table_frame.columns = [np.nan] * table_frame.shape[1]
+
+ if llm_parameters["summary_table"] and smart_headers:
+ try:
+ table_sample = table_frame.head(header_window)
+ table_sample_html = df2html(table_sample)
+ prompt, _temperature, _top_p, _max_tokens = build_prompt(
+ task="detect-table-headers",
+ texts=table_sample_html,
+ query="",
+ paras=llm_parameters,
+ )
+
+ messages = [
+ {"role": "system", "content": "You are a helpful assistant"},
+ {"role": "user", "content": prompt},
+ ]
+
+ context_task_id = gen_str_codes((str(uuid.uuid4()) + table_sample_html))
+
+ if os.getenv("LOCAL_DEBUG", "0") != "1":
+ from shared.services.redis.redis_sync_service import (
+ SyncRedisServiceFactory,
+ )
+
+ redis_service = SyncRedisServiceFactory.get_service()
+ redis_service.set(
+ f"task:{context_task_id}:status",
+ "processing",
+ ttl=7200,
+ )
+
+ header_response = get_openai_client().chat_completion(
+ messages=messages,
+ timeout=60,
+ )
+ parsed_response = eval_response(header_response)
+ if isinstance(parsed_response, dict):
+ answer = parsed_response.get("answer", [])
+ else:
+ answer = parsed_response if isinstance(parsed_response, list) else []
+
+ if not answer or len(answer) == 0:
+ logger.warning(
+ "AI returned empty list, cannot identify headers, falling back to traditional mode..."
+ )
+ header_rows = parse_headers_nonsmart(table_frame)
+ else:
+ try:
+ header_id = answer[-1]
+ header_rows = list(range(header_id + 1))
+ except Exception as exc:
+ logger.warning(
+ f"Failed to parse header row number: {exc}, falling back to traditional mode..."
+ )
+ header_rows = parse_headers_nonsmart(table_frame)
+
+ except Exception as exc:
+ logger.warning(
+ f"Smart header parsing failed: {exc}, falling back to traditional mode..."
+ )
+ header_rows = parse_headers_nonsmart(table_frame)
+ else:
+ header_rows = parse_headers_nonsmart(table_frame)
+
+ if len(header_rows) == 0 or (all(header is None for header in header_rows)):
+ logger.warning("No valid headers detected, fallback to using row 0 as header")
+ new_header = table_frame.iloc[0].ffill().bfill().tolist()
+ table_frame.columns = new_header
+ return table_frame.iloc[1:].reset_index(drop=True)
+
+ if len(header_rows) > 1:
+ header_levels = []
+ for header_index in range(0, len(header_rows)):
+ header_level = table_frame.iloc[header_index].ffill().bfill().tolist()
+ header_levels.append(header_level)
+ new_header = pd.MultiIndex.from_arrays(np.array(header_levels))
+ else:
+ new_header = table_frame.iloc[header_rows[-1]].ffill().bfill().tolist()
+
+ table_frame.columns = new_header
+ table_frame = table_frame.iloc[(header_rows[-1]) + 1 :]
+ return table_frame.reset_index(drop=True)
+
+
+def parse_tb_keywords(table_frame: pd.DataFrame, kw_spit: str = ">>>") -> str:
+ def parse_single_level(columns: list[object], keywords: list[str]) -> list[str]:
+ column_texts = [str(column) for column in columns]
+ for column_text in column_texts:
+ if kw_spit in column_text:
+ keyword = column_text.split(">>>")[0]
+ else:
+ keyword = column_text
+ if keyword not in keywords:
+ keywords.append(column_text)
+ return list({keyword.strip() for keyword in keywords})
+
+ table_keywords: list[str] = []
+ if isinstance(table_frame.columns, pd.MultiIndex):
+ multi_columns = table_frame.columns
+ columns_frame = pd.DataFrame(
+ multi_columns.tolist(),
+ columns=[f"level_{i}" for i in range(multi_columns.nlevels)],
+ )
+ for level_index in range(multi_columns.nlevels):
+ level_keywords: list[str] = []
+ level_keywords = parse_single_level(
+ columns_frame[f"level_{level_index}"].tolist(),
+ level_keywords,
+ )
+ table_keywords.extend(level_keywords)
+ else:
+ table_keywords = parse_single_level(table_frame.columns.tolist(), table_keywords)
+
+ table_keywords = remove_duplicates_orderkept(table_keywords)
+ table_keywords = [
+ keyword
+ for keyword in table_keywords
+ if isinstance(keyword, str)
+ and keyword.strip()
+ and keyword.strip() != "nan"
+ and "Unnamed" not in keyword
+ ]
+ return ";".join(table_keywords)
+
+
+def parse_tb_contents(
+ table_frame: pd.DataFrame,
+ parent_dic: dict[str, object] | None = None,
+ file_name: str = "",
+ sheet_name: str = "",
+ row_header_cols: int = 0,
+) -> tuple[list[str], str]:
+ if parent_dic is None:
+ parent_dic = {}
+
+ rendered_frame = table_frame.fillna("").infer_objects(copy=False)
+ table_html = df2html(rendered_frame, row_header_cols=row_header_cols)
+
+ table_tree = tb_columns_to_tree(table_frame, parent_dic, file_name, sheet_name)
+ table_paths = flatten_dic2paths(table_tree)
+ return table_paths, table_html
+
+
+def tb_columns_to_tree(
+ table_frame: pd.DataFrame,
+ parent_dic: dict[str, object],
+ file_name: str,
+ sheet_name: str,
+) -> dict[str, object]:
+ if isinstance(table_frame.columns, pd.MultiIndex):
+ columns = pd.DataFrame(table_frame.columns.tolist())
+ for level in range(columns.shape[1]):
+ columns[level] = process_duplicate_cols(columns[level])
+
+ new_columns = pd.MultiIndex.from_frame(columns)
+ tree_structure = multiindex_to_tree(new_columns)
+ else:
+ new_columns = process_duplicate_cols(table_frame.columns)
+ tree_structure = {column: {} for column in new_columns}
+
+ table_frame.columns = new_columns
+ if (not file_name == "") and (not sheet_name == ""):
+ parent_dic[file_name][sheet_name] = tree_structure
+ elif not sheet_name == "":
+ parent_dic[sheet_name] = tree_structure
+ elif not file_name == "":
+ parent_dic[file_name] = tree_structure
+ else:
+ parent_dic = tree_structure
+ return parent_dic
+
+
+def multiindex_to_tree(multiindex: pd.MultiIndex) -> dict[object, object]:
+ def tree() -> OrderedDict[object, object]:
+ return OrderedDict()
+
+ root = tree()
+ for keys in multiindex:
+ current_level = root
+ for key in keys:
+ if key not in current_level:
+ current_level[key] = tree()
+ current_level = current_level[key]
+
+ def convert_to_dict(value: object) -> object:
+ if isinstance(value, OrderedDict):
+ return {key: convert_to_dict(child) for key, child in value.items()}
+ return value
+
+ return dict(convert_to_dict(root))
+
+
+def postprocess_tb(table_frame: pd.DataFrame, drop: bool = False) -> pd.DataFrame:
+ if drop:
+ was_range_index = isinstance(table_frame.index, pd.RangeIndex)
+
+ source_row_columns = [
+ column
+ for column in table_frame.columns
+ if (isinstance(column, tuple) and column[0] == "_src_row")
+ or column == "_src_row"
+ ]
+ if source_row_columns:
+ data_columns = [
+ column for column in table_frame.columns if column not in source_row_columns
+ ]
+ mask = table_frame[data_columns].isna().all(axis=1)
+ table_frame = table_frame[~mask]
+ else:
+ table_frame = table_frame.dropna(how="all")
+
+ if was_range_index:
+ table_frame = table_frame.reset_index(drop=True)
+
+ cols_to_drop: list[int] = []
+ for column_index, column in enumerate(table_frame.columns):
+ if table_frame.iloc[:, column_index].isna().all():
+ has_meaningful_header = False
+ if isinstance(column, tuple):
+ for level in column:
+ if (
+ level
+ and str(level).strip()
+ and str(level).strip() not in ["None", "nan", "NaN"]
+ ):
+ has_meaningful_header = True
+ break
+ elif (
+ column
+ and str(column).strip()
+ and str(column).strip() not in ["None", "nan", "NaN"]
+ ):
+ has_meaningful_header = True
+
+ if not has_meaningful_header:
+ cols_to_drop.append(column_index)
+
+ if cols_to_drop:
+ cols_to_keep = [
+ index
+ for index in range(len(table_frame.columns))
+ if index not in cols_to_drop
+ ]
+ table_frame = table_frame.iloc[:, cols_to_keep]
+
+ logger.debug(f"Dropped {len(cols_to_drop)} empty columns")
+
+ if not isinstance(table_frame.index, pd.RangeIndex):
+ was_multiindex = isinstance(table_frame.columns, pd.MultiIndex)
+ column_level_count = table_frame.columns.nlevels if was_multiindex else 1
+ existing_column_set = set(table_frame.columns)
+
+ def make_padded(name: object) -> object:
+ if was_multiindex:
+ return (name,) + ("",) * (column_level_count - 1)
+ return name
+
+ if isinstance(table_frame.index, pd.MultiIndex):
+ seen_counts: dict[object, int] = {}
+ deduped_names: list[object | None] = []
+ for name in table_frame.index.names:
+ if name is None:
+ deduped_names.append(None)
+ continue
+ padded = make_padded(name)
+ if padded in existing_column_set or name in seen_counts:
+ deduped_names.append(None)
+ else:
+ deduped_names.append(name)
+ seen_counts[name] = seen_counts.get(name, 0) + 1
+ table_frame.index.names = deduped_names
+ elif hasattr(table_frame.index, "name") and table_frame.index.name is not None:
+ padded = make_padded(table_frame.index.name)
+ if padded in existing_column_set:
+ table_frame.index.name = None
+
+ table_frame = table_frame.reset_index()
+
+ if was_multiindex:
+ new_columns = []
+ for column in table_frame.columns:
+ if isinstance(column, str) and (
+ column.startswith("level_") or column == "index"
+ ):
+ new_columns.append(tuple([""] * column_level_count))
+ else:
+ new_columns.append(column)
+ table_frame.columns = pd.MultiIndex.from_tuples(new_columns)
+ else:
+ new_columns = []
+ for column in table_frame.columns:
+ if isinstance(column, str) and (
+ column.startswith("level_") or column == "index"
+ ):
+ new_columns.append("")
+ else:
+ new_columns.append(column)
+ table_frame.columns = new_columns
+ else:
+ table_frame.reset_index(drop=True, inplace=True)
+
+ if isinstance(table_frame.columns, pd.MultiIndex):
+ new_levels = []
+ for level_index in range(table_frame.columns.nlevels):
+ level_values = table_frame.columns.get_level_values(level_index)
+ cleaned = [
+ str(value).replace("\n", "") if value is not None else ""
+ for value in level_values
+ ]
+ new_levels.append(cleaned)
+ table_frame.columns = pd.MultiIndex.from_arrays(
+ new_levels,
+ names=table_frame.columns.names,
+ )
+
+ new_levels = []
+ for level_index in range(table_frame.columns.nlevels):
+ level_values = table_frame.columns.get_level_values(level_index)
+ cleaned = [np.nan if "Unnamed" in str(value) else value for value in level_values]
+ new_levels.append(cleaned)
+ table_frame.columns = pd.MultiIndex.from_arrays(
+ new_levels,
+ names=table_frame.columns.names,
+ )
+ else:
+ table_frame.columns = [
+ str(column).replace("\n", "") for column in table_frame.columns
+ ]
+ table_frame.columns = [
+ np.nan if "Unnamed" in str(column) else column
+ for column in table_frame.columns
+ ]
+
+ table_frame = table_frame.map(
+ lambda value: value.replace("\n", "") if isinstance(value, str) else value
+ )
+ return process_datetime_cells(table_frame)
+
+
+def process_datetime_cells(table_frame: pd.DataFrame) -> pd.DataFrame:
+ table_frame = table_frame.copy()
+
+ def convert(value: object) -> object:
+ if isinstance(value, (pd.Timestamp, datetime.datetime)):
+ return value.strftime("%Y-%m-%d %H:%M:%S")
+ return value
+
+ return table_frame.apply(lambda column: column.map(convert))
+
+
+def process_duplicate_cols(columns: object) -> list[object]:
+ column_counts: dict[object, int] = {}
+ new_columns: list[object] = []
+ for column in columns:
+ if column in column_counts:
+ new_columns.append(f"{column}>>>{column_counts[column]}")
+ column_counts[column] += 1
+ else:
+ new_columns.append(column)
+ column_counts[column] = 1
+ return new_columns
+
+
+def format_tb_scope(table_frame: pd.DataFrame, num: int) -> str:
+ if len(table_frame) > int(num * 3 + 1):
+ head_frame = table_frame.head(num)
+ tail_frame = table_frame.tail(num)
+ middle_frame = table_frame.iloc[num : len(table_frame) - num]
+
+ if len(middle_frame) >= num:
+ mid_sample_frame = middle_frame.sample(n=num, random_state=42)
+ else:
+ mid_sample_frame = middle_frame
+ scope_frame = pd.concat(
+ objs=[head_frame, mid_sample_frame, tail_frame],
+ ignore_index=True,
+ )
+ else:
+ scope_frame = table_frame
+ scope_frame = scope_frame.map(
+ lambda value: str(value).strip() if pd.notnull(value) else value
+ )
+ return df2html(scope_frame)
diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/table_parser.py
index 74860b867..2c189c0c7 100755
--- a/apps/worker/app/services/document_parser/table_parser.py
+++ b/apps/worker/app/services/document_parser/table_parser.py
@@ -1,1633 +1,49 @@
-# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportOptionalOperand=false, reportOptionalSubscript=false, reportReturnType=false
-import datetime
-import io
-import os
-import re
-import threading
-import uuid
-from collections import OrderedDict
-from typing import Dict, List, Optional, Tuple, Union
+from __future__ import annotations
-import numpy as np
-import openpyxl
import pandas as pd
-from app.services.common.kb_utils import (
- flatten_dic2paths,
- gen_str_codes,
- get_str_time,
- process_dup_paths_df,
- remove_spaces,
+from app.services.document_parser.table_frame_parser import (
+ format_tb_scope as format_tb_scope,
+ multiindex_to_tree as multiindex_to_tree,
+ parse_headers as parse_headers,
+ parse_tb_contents as parse_tb_contents,
+ parse_tb_keywords as parse_tb_keywords,
+ postprocess_tb as postprocess_tb,
+ process_datetime_cells as process_datetime_cells,
+ process_duplicate_cols as process_duplicate_cols,
+ tb_columns_to_tree as tb_columns_to_tree,
+)
+from app.services.document_parser.table_text_parser import (
+ clean_html_tb as clean_html_tb,
+ df2md as df2md,
+ extract_tables_by_forms as extract_tables_by_forms,
+ identify_tables as identify_tables,
+ sanitize_table_name_from_header as sanitize_table_name_from_header,
)
-from app.services.document_parser.html_parser import df2html
-from bs4 import BeautifulSoup
-from loguru import logger
-
-from shared.core.config import settings
-from shared.core.exceptions.domain_exceptions import TableParsingException
-from shared.core.exceptions.knowhere_exception import KnowhereException
-from shared.services.ai.prompt_service import build_prompt
-from shared.services.ai.response_process_service import eval_response
-from shared.utils.chunk_refs import build_chunk_ref
-from shared.utils.CommonHelperSync import load_file_bytes
-from shared.utils.file_utils import path_handle
-from shared.utils.OpenAICompatibleClientSync import get_openai_client
-from shared.utils.text_utils import remove_duplicates_orderkept, tokenize2stw_remove
-
-# ── Table filename sanitizer ────────────────────────────
-# Max byte-safe filename length. Most filesystems cap at 255 bytes; we leave
-# room for the "table-N " prefix (~10 chars) and ".html" suffix (5 chars).
-_MAX_TABLE_NAME_CHARS = 80
-
-
-def sanitize_table_name_from_header(raw_header_text: str) -> str:
- """Build a concise, filesystem-safe table name from raw first-row header text.
-
- Pipeline:
- 1. Split by common delimiters (' | ', '_br_'/'__br_', '\\n')
- 2. Strip whitespace, deduplicate (preserve order)
- 3. Drop trivial single-character tokens (single CJK char, single digit,
- single letter) — reuses ``_is_meaningful_token`` from shared text_utils
- 4. Rejoin with spaces and cap at ``_MAX_TABLE_NAME_CHARS``
-
- Args:
- raw_header_text: The raw first-row text, often pipe-separated.
-
- Returns:
- A cleaned string suitable for use in a filename (may be empty if all
- fields were trivial).
- """
- from shared.utils.text_utils import _is_meaningful_token
-
- if not raw_header_text:
- return ""
-
- # 1. Split on common header delimiters
- parts = re.split(r"\s*\|\s*|_+br_|\n", raw_header_text)
-
- # 2. Strip + deduplicate (order-preserved)
- seen: set[str] = set()
- unique: list[str] = []
- for p in parts:
- p = p.strip()
- if not p or p in seen:
- continue
- seen.add(p)
- unique.append(p)
-
- # 3. Keep only meaningful fields (drop single-char noise)
- meaningful = [f for f in unique if _is_meaningful_token(f)]
-
- # 4. Join and enforce length cap
- result = " ".join(meaningful)
- if len(result) > _MAX_TABLE_NAME_CHARS:
- result = result[:_MAX_TABLE_NAME_CHARS].rstrip()
- return result
-
-
-g_tbl_lock = threading.Lock()
-
-# ============================================================================
-# PRECISION MODE: Excel Header Detection with Merge Cell Metadata
-# ============================================================================
-
-
-def _get_merged_cell_value(ws, row: int, col: int, merged_ranges: list):
- """
- Get the value of a cell, accounting for merged cell regions.
- For merged cells, returns the value from the top-left corner of the merge range.
-
- Args:
- ws: openpyxl worksheet
- row: 1-indexed row number
- col: 1-indexed column number
- merged_ranges: list of merged cell ranges from ws.merged_cells.ranges
-
- Returns:
- The cell value (from merge origin if applicable)
- """
- for mr in merged_ranges:
- if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col:
- # This cell is part of a merged region, get value from top-left
- return ws.cell(mr.min_row, mr.min_col).value
- # Not a merged cell, return direct value
- return ws.cell(row, col).value
-
-
-# ============================================================================
-# NEW: Enhanced Header Detection with Row/Column MultiIndex Support
-# ============================================================================
-
-# Data types that indicate a cell is data, not header (parameterized for future extension)
-DATA_TYPES_TO_EXCLUDE = (int, float, datetime.datetime)
-
-
-def _get_unique_cells_in_row(
- ws, row: int, col_range: Tuple[int, int], merged_ranges: list
-) -> List[dict]:
- """Get all unique cells in a row, treating merged cells as single cells.
-
- Returns: List of {col_start, col_end, value, is_merged}
- """
- c_start, c_end = col_range
- cells = []
- visited_cols = set()
-
- for col in range(c_start, c_end + 1):
- if col in visited_cols:
- continue
-
- in_merge = False
- for mr in merged_ranges:
- if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col:
- val = ws.cell(mr.min_row, mr.min_col).value
- merge_col_end = min(mr.max_col, c_end)
-
- for mc in range(mr.min_col, merge_col_end + 1):
- visited_cols.add(mc)
-
- cells.append(
- {
- "col_start": mr.min_col,
- "col_end": merge_col_end,
- "value": val,
- "is_merged": True,
- }
- )
- in_merge = True
- break
-
- if not in_merge:
- val = ws.cell(row, col).value
- cells.append(
- {"col_start": col, "col_end": col, "value": val, "is_merged": False}
- )
- visited_cols.add(col)
-
- return cells
-
-
-def _get_unique_cells_in_col(
- ws, col: int, row_range: Tuple[int, int], merged_ranges: list
-) -> List[dict]:
- """Get all unique cells in a column, treating merged cells as single cells."""
- r_start, r_end = row_range
- cells = []
- visited_rows = set()
-
- for row in range(r_start, r_end + 1):
- if row in visited_rows:
- continue
-
- in_merge = False
- for mr in merged_ranges:
- if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col:
- val = ws.cell(mr.min_row, mr.min_col).value
- merge_row_end = min(mr.max_row, r_end)
-
- for mr_row in range(mr.min_row, merge_row_end + 1):
- visited_rows.add(mr_row)
-
- cells.append(
- {
- "row_start": mr.min_row,
- "row_end": merge_row_end,
- "value": val,
- "is_merged": True,
- }
- )
- in_merge = True
- break
-
- if not in_merge:
- val = ws.cell(row, col).value
- cells.append(
- {"row_start": row, "row_end": row, "value": val, "is_merged": False}
- )
- visited_rows.add(row)
-
- return cells
-
-
-def _is_candidate_header_row(
- ws,
- row: int,
- col_range: Tuple[int, int],
- merged_ranges: list,
- exclude_types: tuple = DATA_TYPES_TO_EXCLUDE,
-) -> bool:
- """Check if a row is a candidate header row.
-
- Logic: Row is a candidate if all cells are text (no numbers/dates).
- Merged cells are treated as single cells.
- """
- cells = _get_unique_cells_in_row(ws, row, col_range, merged_ranges)
-
- has_any_value = False
- for cell in cells:
- val = cell["value"]
- if val is None:
- continue
- has_any_value = True
-
- if isinstance(val, bool):
- continue
- if isinstance(val, exclude_types):
- return False
-
- return has_any_value
-
-
-def _is_candidate_header_col(
- ws,
- col: int,
- row_range: Tuple[int, int],
- merged_ranges: list,
- exclude_types: tuple = DATA_TYPES_TO_EXCLUDE,
-) -> bool:
- """Check if a column is a candidate header column (for row index)."""
- cells = _get_unique_cells_in_col(ws, col, row_range, merged_ranges)
-
- has_any_value = False
- for cell in cells:
- val = cell["value"]
- if val is None:
- continue
- has_any_value = True
-
- if isinstance(val, bool):
- continue
- if isinstance(val, exclude_types):
- return False
-
- return has_any_value
-
-
-def _detect_header_regions(
- ws, row_range: Tuple[int, int], col_range: Tuple[int, int], merged_ranges: list
-) -> Tuple[List[int], List[int]]:
- """Detect header rows and columns.
-
- Scans rows first, then scans columns only in the data region (excluding header rows).
- This prevents header row content from influencing column header detection.
-
- Returns:
- header_rows: List of candidate header row numbers (1-indexed)
- header_cols: List of candidate header column numbers (1-indexed)
- """
- r_start, r_end = row_range
- c_start, c_end = col_range
-
- # Scan for candidate header rows (top to bottom)
- header_rows = []
- for row in range(r_start, r_end + 1):
- if _is_candidate_header_row(ws, row, col_range, merged_ranges):
- header_rows.append(row)
- else:
- break
-
- # Determine data region (excluding header rows)
- data_row_start = header_rows[-1] + 1 if header_rows else r_start
-
- # Skip column scanning if no data rows remain
- if data_row_start > r_end:
- return header_rows, []
-
- # Scan for candidate header columns (left to right) - only in data region
- header_cols = []
- data_row_range = (data_row_start, r_end)
- for col in range(c_start, c_end + 1):
- if _is_candidate_header_col(ws, col, data_row_range, merged_ranges):
- header_cols.append(col)
- else:
- break
-
- return header_rows, header_cols
-
-
-def _build_column_multiindex(
- ws, header_rows: List[int], col_range: Tuple[int, int], merged_ranges: list
-) -> Union[pd.Index, pd.MultiIndex]:
- """Build column MultiIndex from header rows."""
- c_start, c_end = col_range
- levels = []
-
- for row in header_rows:
- row_values = []
- for col in range(c_start, c_end + 1):
- val = _get_merged_cell_value(ws, row, col, merged_ranges)
- row_values.append(str(val).strip() if val else "")
- levels.append(row_values)
-
- # Forward fill for merged cells
- for idx, level in enumerate(levels):
- filled = []
- last = ""
- for val in level:
- if val:
- last = val
- filled.append(last if last else val)
- levels[idx] = filled
-
- if len(levels) == 1:
- return pd.Index(levels[0])
- return pd.MultiIndex.from_arrays(levels)
-
-
-def _build_row_multiindex(
- ws,
- header_cols: List[int],
- row_range: Tuple[int, int],
- merged_ranges: list,
- header_rows: List[int] = None,
-) -> Union[pd.Index, pd.MultiIndex]:
- """Build row MultiIndex from header columns.
-
- Args:
- header_rows: If provided, use the last header row's values as index names
- """
- r_start, r_end = row_range
- levels = []
- names = []
-
- for col in header_cols:
- col_values = []
- for row in range(r_start, r_end + 1):
- val = _get_merged_cell_value(ws, row, col, merged_ranges)
- col_values.append(str(val).strip() if val else "")
- levels.append(col_values)
-
- # Get the column name from the last header row
- if header_rows:
- name_row = header_rows[-1]
- name_val = _get_merged_cell_value(ws, name_row, col, merged_ranges)
- names.append(str(name_val).strip() if name_val else None)
- else:
- names.append(None)
-
- # Forward fill for merged cells
- for idx, level in enumerate(levels):
- filled = []
- last = ""
- for val in level:
- if val:
- last = val
- filled.append(last if last else val)
- levels[idx] = filled
-
- if len(levels) == 1:
- idx = pd.Index(levels[0])
- idx.name = names[0] if names else None
- return idx
- return pd.MultiIndex.from_arrays(levels, names=names)
-
-
-def _parse_subtable(
- ws, row_range: Tuple[int, int], col_range: Tuple[int, int], merged_ranges: list
-) -> dict:
- """Parse a subtable with new header detection logic.
-
- Returns:
- dict with keys: df, header_rows, header_cols, fallback_col_header, fallback_row_header
- """
- r_start, r_end = row_range
- c_start, c_end = col_range
-
- header_rows, header_cols = _detect_header_regions(
- ws, row_range, col_range, merged_ranges
- )
-
- total_rows = r_end - r_start + 1
- total_cols = c_end - c_start + 1
-
- # Fall-back check: if all rows/cols are headers, treat as no-header
- fallback_col_header = len(header_rows) == total_rows
- fallback_row_header = len(header_cols) == total_cols
-
- # Determine data region
- if fallback_col_header:
- data_row_start = r_start
- columns = None
- else:
- data_row_start = header_rows[-1] + 1 if header_rows else r_start
- columns = (
- _build_column_multiindex(ws, header_rows, col_range, merged_ranges)
- if header_rows
- else None
- )
-
- if fallback_row_header:
- data_col_start = c_start
- row_index = None
- else:
- data_col_start = header_cols[-1] + 1 if header_cols else c_start
- row_index = (
- _build_row_multiindex(
- ws, header_cols, (data_row_start, r_end), merged_ranges, header_rows
- )
- if header_cols
- else None
- )
-
- # Read data
- data = []
- for row in range(data_row_start, r_end + 1):
- row_data = []
- for col in range(data_col_start, c_end + 1):
- val = _get_merged_cell_value(ws, row, col, merged_ranges)
- row_data.append(val)
- data.append(row_data)
-
- # Adjust column index if there are row index columns
- if columns is not None and header_cols and not fallback_row_header:
- if isinstance(columns, pd.MultiIndex):
- columns = columns[len(header_cols) :]
- else:
- columns = columns[len(header_cols) :]
-
- df = pd.DataFrame(data, columns=columns, index=row_index)
-
- # Append original Excel row numbers as the last column for cross-referencing
- excel_row_numbers = list(range(data_row_start, r_end + 1))
- if isinstance(df.columns, pd.MultiIndex):
- n_levels = df.columns.nlevels
- src_row_key = tuple(["_src_row"] + [""] * (n_levels - 1))
- df[src_row_key] = excel_row_numbers
- else:
- df["_src_row"] = excel_row_numbers
-
- return {
- "df": df,
- "header_rows": header_rows if not fallback_col_header else [],
- "header_cols": header_cols if not fallback_row_header else [],
- "fallback_col_header": fallback_col_header,
- "fallback_row_header": fallback_row_header,
- }
-
-
-# ============================================================================
-# Sheet Splitting: Detect true separators and split into subtables
-# ============================================================================
-
-
-def _find_effective_range(
- ws, row_range: Tuple[int, int], col_range: Tuple[int, int]
-) -> Tuple[Tuple[int, int], Tuple[int, int]]:
- """Find the effective (non-empty) row and column ranges within a region."""
- r_start, r_end = row_range
- c_start, c_end = col_range
-
- eff_r_start, eff_r_end = None, None
- eff_c_start, eff_c_end = None, None
-
- for row in range(r_start, r_end + 1):
- for col in range(c_start, c_end + 1):
- if ws.cell(row, col).value is not None:
- if eff_r_start is None:
- eff_r_start = row
- eff_r_end = row
- if eff_c_start is None or col < eff_c_start:
- eff_c_start = col
- if eff_c_end is None or col > eff_c_end:
- eff_c_end = col
-
- if eff_r_start is None:
- return ((r_start, r_start), (c_start, c_start))
-
- return ((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))
-
-
-def _is_true_separator_row(
- ws, row: int, effective_col_range: Tuple[int, int], merged_ranges: list = None
-) -> bool:
- """Check if a row is a true separator (all empty within effective column range).
-
- Considers merged cells - a cell is not empty if it's part of any merged range.
- """
- c_start, c_end = effective_col_range
- merged_ranges = merged_ranges or []
-
- for col in range(c_start, c_end + 1):
- # Check if cell has a value
- if ws.cell(row, col).value is not None:
- return False
- # Check if cell is part of a merged range
- for mr in merged_ranges:
- if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col:
- return False # Part of a merge, not truly empty
- return True
-
-
-def _is_true_separator_col(
- ws, col: int, effective_row_range: Tuple[int, int], merged_ranges: list = None
-) -> bool:
- """Check if a column is a true separator (all empty within effective row range).
-
- Considers merged cells - a cell is not empty if it's part of any merged range.
- """
- r_start, r_end = effective_row_range
- merged_ranges = merged_ranges or []
-
- for row in range(r_start, r_end + 1):
- # Check if cell has a value
- if ws.cell(row, col).value is not None:
- return False
- # Check if cell is part of a merged range
- for mr in merged_ranges:
- if mr.min_row <= row <= mr.max_row and mr.min_col <= col <= mr.max_col:
- return False # Part of a merge, not truly empty
- return True
-
-
-def _find_separator_groups(items: List[int]) -> List[List[int]]:
- """Group consecutive separator items together."""
- if not items:
- return []
-
- groups = []
- current_group = [items[0]]
-
- for i in range(1, len(items)):
- if items[i] == items[i - 1] + 1:
- current_group.append(items[i])
- else:
- groups.append(current_group)
- current_group = [items[i]]
-
- groups.append(current_group)
- return groups
-
-
-def _split_sheet_recursive(
- ws,
- row_range: Tuple[int, int],
- col_range: Tuple[int, int],
- merged_ranges: list = None,
- min_rows: int = 2,
- min_cols: int = 2,
-) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]:
- """
- Recursively split a sheet region into subtables based on true separators.
-
- Args:
- merged_ranges: List of merged cell ranges to consider when detecting separators
- Returns list of (row_range, col_range) tuples for each subtable.
- """
- r_start, r_end = row_range
- c_start, c_end = col_range
- merged_ranges = merged_ranges or []
- # Find effective range (trim empty edges)
- (eff_r_start, eff_r_end), (eff_c_start, eff_c_end) = _find_effective_range(
- ws, row_range, col_range
+def parse_xlsx(
+ file_path: str,
+ file_name: str,
+ output_dir: str,
+ baseurl: str,
+ base_llm_paras: dict[str, object] | None = None,
+ window_h: int = 10,
+ relative_root: str | None = None,
+ use_precision_mode: bool = True,
+ include_hidden_sheets: bool = False,
+) -> pd.DataFrame:
+ from app.services.document_parser.excel_table_parser import (
+ parse_xlsx as parse_excel_xlsx,
)
- # If region is too small or empty, return as-is or empty
- if eff_r_end - eff_r_start + 1 < min_rows or eff_c_end - eff_c_start + 1 < min_cols:
- if eff_r_start is not None:
- return [((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))]
- return []
-
- # Find true separator rows (considering merged cells)
- separator_rows = []
- for row in range(eff_r_start + 1, eff_r_end):
- if _is_true_separator_row(ws, row, (eff_c_start, eff_c_end), merged_ranges):
- separator_rows.append(row)
-
- # Find true separator columns (considering merged cells)
- separator_cols = []
- for col in range(eff_c_start + 1, eff_c_end):
- if _is_true_separator_col(ws, col, (eff_r_start, eff_r_end), merged_ranges):
- separator_cols.append(col)
-
- # Group consecutive separators
- row_groups = _find_separator_groups(separator_rows)
- col_groups = _find_separator_groups(separator_cols)
-
- # Choose split direction
- do_row_split = len(row_groups) > 0 and (
- len(col_groups) == 0 or len(row_groups) <= len(col_groups)
+ return parse_excel_xlsx(
+ file_path=file_path,
+ file_name=file_name,
+ output_dir=output_dir,
+ baseurl=baseurl,
+ base_llm_paras=base_llm_paras,
+ window_h=window_h,
+ relative_root=relative_root,
+ use_precision_mode=use_precision_mode,
+ include_hidden_sheets=include_hidden_sheets,
)
- do_col_split = len(col_groups) > 0 and not do_row_split
-
- if do_row_split:
- subtables = []
- prev_end = eff_r_start
- for group in row_groups:
- if group[0] > prev_end:
- sub_result = _split_sheet_recursive(
- ws,
- (prev_end, group[0] - 1),
- (eff_c_start, eff_c_end),
- merged_ranges,
- min_rows,
- min_cols,
- )
- subtables.extend(sub_result)
- prev_end = group[-1] + 1
- if prev_end <= eff_r_end:
- sub_result = _split_sheet_recursive(
- ws,
- (prev_end, eff_r_end),
- (eff_c_start, eff_c_end),
- merged_ranges,
- min_rows,
- min_cols,
- )
- subtables.extend(sub_result)
- return subtables
-
- elif do_col_split:
- subtables = []
- prev_end = eff_c_start
- for group in col_groups:
- if group[0] > prev_end:
- sub_result = _split_sheet_recursive(
- ws,
- (eff_r_start, eff_r_end),
- (prev_end, group[0] - 1),
- merged_ranges,
- min_rows,
- min_cols,
- )
- subtables.extend(sub_result)
- prev_end = group[-1] + 1
- if prev_end <= eff_c_end:
- sub_result = _split_sheet_recursive(
- ws,
- (eff_r_start, eff_r_end),
- (prev_end, eff_c_end),
- merged_ranges,
- min_rows,
- min_cols,
- )
- subtables.extend(sub_result)
- return subtables
-
- else:
- return [((eff_r_start, eff_r_end), (eff_c_start, eff_c_end))]
-
-
-# ============================================================================
-# Post-split Merge: Absorb small fragments into nearest neighbor
-# ============================================================================
-
-
-def _count_non_empty_cells(
- ws, row_range: Tuple[int, int], col_range: Tuple[int, int]
-) -> int:
- """Count non-empty cells in a region."""
- count = 0
- for r in range(row_range[0], row_range[1] + 1):
- for c in range(col_range[0], col_range[1] + 1):
- if ws.cell(r, c).value is not None:
- count += 1
- return count
-
-
-def _merge_small_subtables(
- ws, subtables: List[Tuple[Tuple[int, int], Tuple[int, int]]], min_cells: int = 4
-) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]:
- """Merge subtables that have too few non-empty cells into their nearest neighbor.
-
- This is a post-processing step after _split_sheet_recursive to prevent
- over-fragmentation. Fragments with fewer than min_cells non-empty cells
- are iteratively absorbed into the nearest neighbor subtable (by bounding-box
- distance), expanding the neighbor's bounding box to encompass both regions.
-
- Args:
- ws: openpyxl worksheet
- subtables: list of (row_range, col_range) tuples from _split_sheet_recursive
- min_cells: minimum non-empty cells for a subtable to be kept standalone
-
- Returns:
- Merged list of (row_range, col_range) tuples
- """
- if len(subtables) <= 1:
- return subtables
-
- # Build working list with cell counts
- items = []
- for rr, cr in subtables:
- count = _count_non_empty_cells(ws, rr, cr)
- items.append({"rr": rr, "cr": cr, "cells": count})
-
- # Iteratively merge the smallest sub-threshold fragment
- changed = True
- while changed and len(items) > 1:
- changed = False
-
- # Find the smallest fragment below threshold
- min_idx = None
- for i, item in enumerate(items):
- if item["cells"] < min_cells:
- if min_idx is None or item["cells"] < items[min_idx]["cells"]:
- min_idx = i
-
- if min_idx is None:
- break # All subtables are above threshold
-
- # Find nearest neighbor by bounding-box gap distance
- src = items[min_idx]
- best_j = None
- best_dist = float("inf")
- for j, tgt in enumerate(items):
- if j == min_idx:
- continue
- row_gap = max(
- 0, tgt["rr"][0] - src["rr"][1] - 1, src["rr"][0] - tgt["rr"][1] - 1
- )
- col_gap = max(
- 0, tgt["cr"][0] - src["cr"][1] - 1, src["cr"][0] - tgt["cr"][1] - 1
- )
- dist = row_gap + col_gap
- if dist < best_dist or (
- dist == best_dist and tgt["cells"] > items[best_j]["cells"]
- ):
- best_dist = dist
- best_j = j
-
- if best_j is None:
- break # Should not happen when len(items) > 1
-
- # Merge: expand neighbor's bounding box to encompass both
- tgt = items[best_j]
- merged_rr = (min(src["rr"][0], tgt["rr"][0]), max(src["rr"][1], tgt["rr"][1]))
- merged_cr = (min(src["cr"][0], tgt["cr"][0]), max(src["cr"][1], tgt["cr"][1]))
- items[best_j] = {
- "rr": merged_rr,
- "cr": merged_cr,
- "cells": src["cells"] + tgt["cells"],
- }
-
- logger.debug(
- f"Merged small fragment (rows={src['rr']}, cols={src['cr']}, "
- f"cells={src['cells']}) into neighbor (rows={tgt['rr']}, cols={tgt['cr']})"
- )
-
- del items[min_idx]
- changed = True
-
- return [(item["rr"], item["cr"]) for item in items]
-
-
-def parse_headers_from_excel(
- file_source: Union[str, io.BytesIO],
- sheet_name: Optional[str] = None,
- split_subtables: bool = True,
- include_hidden_sheets: bool = False,
-) -> Dict[str, pd.DataFrame]:
- """
- Parse Excel file using openpyxl to accurately detect headers via merged cell metadata.
-
- This is the PRECISION MODE for Excel parsing - it uses the actual merge cell
- information from the Excel file to build correct MultiIndex headers without
- relying on LLM or heuristics.
-
- Args:
- file_source: Path to Excel file or BytesIO stream
- sheet_name: Specific sheet to parse (None = all sheets)
- split_subtables: If True, split sheets into subtables based on empty row/column separators (default: True)
- include_hidden_sheets: If True, parse hidden/very-hidden sheets. Default False (skip them).
-
- Returns:
- Dictionary mapping sheet/subtable names to DataFrames with correctly set headers
- When split_subtables=True, keys are like 'SheetName', 'SheetName_2', 'SheetName_3' etc.
- """
- try:
- # Load workbook with data_only=True to get calculated values
- if isinstance(file_source, str):
- wb = openpyxl.load_workbook(file_source, data_only=True)
- else:
- # BytesIO stream
- file_source.seek(0) # Ensure we're at the start
- wb = openpyxl.load_workbook(file_source, data_only=True)
-
- results = {}
- sheets_to_parse = [sheet_name] if sheet_name else wb.sheetnames
-
- for sn in sheets_to_parse:
- if sn not in wb.sheetnames:
- logger.warning(f"Sheet '{sn}' not found in workbook, skipping")
- continue
-
- ws = wb[sn]
-
- # Skip hidden sheets unless explicitly included
- if not include_hidden_sheets and ws.sheet_state != "visible":
- logger.info(
- f"Sheet '{sn}' is hidden (state={ws.sheet_state}), skipping"
- )
- continue
-
- # Skip empty sheets
- if ws.max_row is None or ws.max_row == 0:
- logger.debug(f"Sheet '{sn}' is empty, skipping")
- continue
-
- # Get merged cell ranges
- merged_ranges = list(ws.merged_cells.ranges)
- logger.debug(f"Sheet '{sn}': found {len(merged_ranges)} merged cell ranges")
-
- if split_subtables:
- # Split sheet into subtables (considers merged cells)
- subtable_regions = _split_sheet_recursive(
- ws, (1, ws.max_row), (1, ws.max_column or 1), merged_ranges
- )
- # Merge back small fragments to prevent over-fragmentation
- before_count = len(subtable_regions)
- subtable_regions = _merge_small_subtables(ws, subtable_regions)
- if len(subtable_regions) != before_count:
- logger.info(
- f"Sheet '{sn}': merged {before_count} subtables → {len(subtable_regions)} "
- f"(absorbed {before_count - len(subtable_regions)} small fragments)"
- )
- logger.debug(
- f"Sheet '{sn}': {len(subtable_regions)} subtables after merge"
- )
-
- for idx, (row_range, col_range) in enumerate(subtable_regions):
- result = _parse_subtable(ws, row_range, col_range, merged_ranges)
- df = result["df"]
-
- # Store header_cols count in DataFrame attrs for later use in HTML rendering
- df.attrs["row_header_cols"] = len(result["header_cols"])
-
- # Generate unique key for each subtable
- if idx == 0:
- key = sn
- else:
- key = f"{sn}_{idx + 1}"
-
- logger.debug(
- f"Subtable '{key}': rows={row_range}, cols={col_range}, "
- f"header_rows={result['header_rows']}, header_cols={result['header_cols']}"
- )
-
- results[key] = df
- else:
- # Treat entire sheet as one subtable
- row_range = (1, ws.max_row)
- col_range = (1, ws.max_column or 1)
-
- result = _parse_subtable(ws, row_range, col_range, merged_ranges)
- df = result["df"]
-
- # Store header_cols count in DataFrame attrs for later use in HTML rendering
- df.attrs["row_header_cols"] = len(result["header_cols"])
-
- logger.debug(
- f"Sheet '{sn}': header_rows={result['header_rows']}, "
- f"header_cols={result['header_cols']}, "
- f"fallback_col={result['fallback_col_header']}, "
- f"fallback_row={result['fallback_row_header']}"
- )
-
- results[sn] = df
-
- wb.close()
- return results
-
- except Exception as e:
- logger.error(f"Error parsing Excel with precision mode: {e}")
- raise TableParsingException(
- user_message="Failed to parse Excel file headers",
- reason="EXCEL_PRECISION_PARSE_FAILED",
- internal_message=str(e),
- original_exception=e,
- )
-
-
-def identify_tables(line):
- """Identify if a line contains a table.
-
- Note: For HTML tables, use merge_html_tables() from html_parser.py
- to preprocess multi-line tables before calling this function.
- """
- # HTML table: complete in one line
- html_tb_pattern = r".*? "
- tables = re.findall(html_tb_pattern, line, re.DOTALL)
- if bool(tables):
- return True, "html", tables
-
- # MD table: lines starting and ending with |
- if line.startswith("|") and line.endswith("|"):
- return True, "md", []
-
- return False, None, None
-
-
-def df2md(tb_df: pd.DataFrame, *, index: bool = False, na_rep: str = "—") -> str:
- """Convert DataFrame to Markdown table format with dynamic column widths.
-
- Note: Truncation should be done externally using truncate_text before calling this function.
-
- Args:
- tb_df: Input DataFrame
- index: Whether to include index column
- na_rep: String to represent NA values
-
- Returns:
- Markdown table string
- """
- import unicodedata
-
- def get_display_width(text: str) -> int:
- """eval width for both ASCII and Chinese"""
- width = 0
- for char in text:
- if unicodedata.east_asian_width(char) in ("F", "W"):
- width += 2
- else:
- width += 1
- return width
-
- def pad_to_width(text: str, target_width: int) -> str:
- current_width = get_display_width(text)
- padding = target_width - current_width
- return text + " " * max(0, padding)
-
- df = tb_df.copy()
-
- # Handle index
- if index:
- df = df.reset_index()
-
- # Replace NA values
- df = df.fillna(na_rep)
-
- # Convert all values to string
- df = df.astype(str)
-
- # Calculate column widths based on actual display width (no truncation)
- col_widths = {}
- for col in df.columns:
- header_width = get_display_width(str(col))
- max_content_width = max(df[col].apply(get_display_width)) if len(df) > 0 else 0
- col_widths[col] = max(header_width, max_content_width)
-
- # Build header row
- header_cells = [pad_to_width(str(col), col_widths[col]) for col in df.columns]
- header_line = "| " + " | ".join(header_cells) + " |"
-
- # Build separator row
- separator_cells = ["-" * col_widths[col] for col in df.columns]
- separator_line = "|-" + "-|-".join(separator_cells) + "-|"
-
- # Build data rows
- data_lines = []
- for _, row in df.iterrows():
- cells = [pad_to_width(str(row[col]), col_widths[col]) for col in df.columns]
- data_lines.append("| " + " | ".join(cells) + " |")
-
- # Combine all parts
- lines = [header_line, separator_line] + data_lines
- return "\n".join(lines)
-
-
-def clean_html_tb(html: str) -> str:
- soup = BeautifulSoup(html, "html.parser")
- for row in soup.find_all("tr"):
- seen = set()
- unique_cells = []
- for cell in row.find_all("td", recursive=False):
- content = cell.encode_contents()
- if content not in seen:
- seen.add(content)
- unique_cells.append(cell)
- row.clear()
- for cell in unique_cells:
- row.append(cell)
- return soup.prettify()
-
-
-def extract_tables_by_forms(tb_txt, form):
- if form == "html":
- return tb_txt
- elif form == "md":
- tb_df = pd.read_table(
- pd.io.common.StringIO(tb_txt), sep="|", engine="python", on_bad_lines="skip"
- )
- tb_df = tb_df.drop(columns=tb_df.columns[0]) # Drop extra leading column
- tb_df = tb_df.drop(columns=tb_df.columns[-1]) # Drop extra trailing column
- tb_df.columns = tb_df.columns.str.strip() # Clean up headers
- # Filter out MD separator lines (e.g. "---", ":---:", "---:")
- separator_pattern = r"^[\s\-:]+$"
- tb_df = tb_df[
- ~tb_df.apply(
- lambda row: row.astype(str).str.match(separator_pattern).all(), axis=1
- )
- ]
- tb_strs = tb_df.to_html(index=False)
- else:
- tb_strs = None # UNDER DEVELOPMENT other forms of tables...
- return tb_strs
-
-
-def parse_headers(df_temp, paras=None, header_window=5, smart_headers=True):
- def parse_headers_nonsmart(df_):
- non_na_row = df_[df_.notna().any(axis=1)].head(1)
- header_id = non_na_row.index[-1] if not non_na_row.empty else None
- header_rows = list(range(header_id + 1))
- return header_rows
-
- if not pd.isna(
- df_temp.columns
- ).all(): # If columns are not all NaN, no need to add extra row
- df_temp.loc[-1] = df_temp.columns
- df_temp.index = df_temp.index + 1
- df_temp = df_temp.sort_index()
- df_temp.columns = [np.nan] * df_temp.shape[1]
-
- if paras["summary_table"] and smart_headers:
- try:
- tb_small = df_temp.head(header_window)
- tb_small_str = df2html(tb_small)
- prompt, temperature, top_p, max_tokens = build_prompt(
- task="detect-table-headers", texts=tb_small_str, query="", paras=paras
- )
-
- messages = [
- {"role": "system", "content": "You are a helpful assistant"},
- {"role": "user", "content": prompt},
- ]
-
- ctx_task_id = gen_str_codes((str(uuid.uuid4()) + tb_small_str))
-
- # Track task status via Redis (skip in LOCAL_DEBUG mode)
- import os
-
- if os.getenv("LOCAL_DEBUG", "0") != "1":
- from shared.services.redis.redis_sync_service import (
- SyncRedisServiceFactory,
- )
-
- redis_service = SyncRedisServiceFactory.get_service()
- redis_service.set(f"task:{ctx_task_id}:status", "processing", ttl=7200)
-
- # Use unified AI service
- header_res = get_openai_client().chat_completion(
- messages=messages, timeout=60
- )
- header_res = eval_response(header_res)
- # Extract answer field
- if isinstance(header_res, dict):
- answer = header_res.get("answer", [])
- else:
- answer = header_res if isinstance(header_res, list) else []
-
- # Check if answer is empty list
- if not answer or len(answer) == 0:
- logger.warning(
- "AI returned empty list, cannot identify headers, falling back to traditional mode..."
- )
- header_rows = parse_headers_nonsmart(df_temp)
- else:
- try:
- header_id = answer[-1]
- header_rows = list(range(header_id + 1))
- except Exception as e:
- logger.warning(
- f"Failed to parse header row number: {e}, falling back to traditional mode..."
- )
- header_rows = parse_headers_nonsmart(df_temp)
-
- except Exception as e:
- logger.warning(
- f"Smart header parsing failed: {e}, falling back to traditional mode..."
- )
- header_rows = parse_headers_nonsmart(df_temp)
- else:
- header_rows = parse_headers_nonsmart(df_temp)
-
- # improve table structure based on header rows
- if len(header_rows) == 0 or (all(h is None for h in header_rows)):
- logger.warning("No valid headers detected, fallback to using row 0 as header")
- new_header = df_temp.iloc[0].ffill().bfill().tolist()
- df_temp.columns = new_header
- df_temp = df_temp.iloc[1:].reset_index(drop=True)
- return df_temp
- elif len(header_rows) > 1:
- head_lst = []
- for i in range(0, len(header_rows)):
- temp_lst = df_temp.iloc[i].ffill().bfill().tolist()
- head_lst.append(temp_lst)
- new_header = pd.MultiIndex.from_arrays(np.array(head_lst))
- else:
- new_header = df_temp.iloc[header_rows[-1]].ffill().bfill().tolist()
-
- df_temp.columns = new_header
- df_temp = df_temp.iloc[(header_rows[-1]) + 1 :]
- df_temp = df_temp.reset_index(drop=True)
- return df_temp
-
-
-def parse_tb_keywords(
- tb_df, kw_spit=">>>"
-): # Extract keywords from headers (can also add LLM extraction)
- def parse_single_level_(cols, keywords):
- cols = [str(c) for c in cols]
- for col in cols:
- if kw_spit in col:
- tmp_kw = col.split(">>>")[0]
- else: # May be first occurrence
- tmp_kw = col
- if tmp_kw not in keywords:
- keywords.append(col)
- keywords_a_level = list(set([k.strip() for k in keywords]))
- return keywords_a_level
-
- tb_keywords = []
- if isinstance(tb_df.columns, pd.MultiIndex):
- multi_cols = tb_df.columns
- cols_df = pd.DataFrame(
- multi_cols.tolist(),
- columns=[f"level_{i}" for i in range(multi_cols.nlevels)],
- )
- for i in range(multi_cols.nlevels): # Extract each level as list
- level_kws = []
- level_kws = parse_single_level_(cols_df[f"level_{i}"].tolist(), level_kws)
- tb_keywords.extend(level_kws)
- else:
- tb_keywords = parse_single_level_(tb_df.columns, tb_keywords)
-
- # Remove duplicates while preserving column order
- tb_keywords = remove_duplicates_orderkept(tb_keywords)
- tb_keywords = [
- k
- for k in tb_keywords
- if isinstance(k, str)
- and k.strip()
- and k.strip() != "nan"
- and "Unnamed" not in k
- ]
- return ";".join(tb_keywords)
-
-
-def parse_tb_contents(
- df_temp, parent_dic=None, file_name="", sheet_name="", row_header_cols=0
-):
- """Parse table contents and generate HTML.
-
- Args:
- row_header_cols: Number of leftmost columns that are row headers (will be rendered as )
- """
- if parent_dic is None:
- parent_dic = {}
-
- tb_res = df_temp.fillna("").infer_objects(copy=False)
- tb_strs = df2html(tb_res, row_header_cols=row_header_cols)
-
- tb_tree = tb_columns_to_tree(df_temp, parent_dic, file_name, sheet_name)
- tb_paths = flatten_dic2paths(tb_tree)
- return tb_paths, tb_strs
-
-
-def tb_columns_to_tree(df, parent_dic, file_name, sheet_name):
- if isinstance(df.columns, pd.MultiIndex):
- # Convert MultiIndex columns to a nested dictionary (tree-like structure)
- columns = pd.DataFrame(df.columns.tolist())
- for level in range(columns.shape[1]):
- columns[level] = process_duplicate_cols(columns[level])
-
- new_columns = pd.MultiIndex.from_frame(columns)
- tree_structure = multiindex_to_tree(new_columns)
- else:
- # If columns are not MultiIndex, convert them to a dictionary with empty dictionaries as values
- new_columns = process_duplicate_cols(df.columns)
- tree_structure = {col: {} for col in new_columns}
-
- df.columns = new_columns
- if (not file_name == "") and (not sheet_name == ""):
- parent_dic[file_name][sheet_name] = tree_structure
- elif not sheet_name == "":
- parent_dic[sheet_name] = tree_structure
- elif not file_name == "":
- parent_dic[file_name] = tree_structure
- else:
- parent_dic = tree_structure
- return parent_dic
-
-
-def multiindex_to_tree(multiindex):
- """Convert a MultiIndex to a tree-like nested dictionary structure."""
-
- def tree():
- return OrderedDict()
-
- root = tree()
- for keys in multiindex:
- current_level = root
- for key in keys:
- if key not in current_level:
- current_level[key] = tree()
- current_level = current_level[key]
-
- def convert_to_dict(d):
- if isinstance(d, OrderedDict):
- d = {k: convert_to_dict(v) for k, v in d.items()}
- return d
-
- return convert_to_dict(root)
-
-
-def postprocess_tb(df, drop=False):
- if drop:
- # Track if index was originally a simple RangeIndex (no semantic meaning)
- # dropna(how='all') can turn RangeIndex into Int64Index by introducing gaps,
- # which would incorrectly trigger the "preserve row index" logic below.
- was_range_index = isinstance(df.index, pd.RangeIndex)
-
- # Drop rows where all data columns are empty (exclude _src_row from the check)
- # _src_row is always non-null, so including it would prevent any row from being dropped.
- src_row_cols = [
- c
- for c in df.columns
- if (isinstance(c, tuple) and c[0] == "_src_row") or c == "_src_row"
- ]
- if src_row_cols:
- data_cols = [c for c in df.columns if c not in src_row_cols]
- mask = df[data_cols].isna().all(axis=1)
- df = df[~mask]
- else:
- df = df.dropna(how="all")
-
- # If index was originally RangeIndex, re-number it to avoid gaps
- if was_range_index:
- df = df.reset_index(drop=True)
-
- # Drop columns that are all empty AND have no meaningful header
- # A column with a valid header should be preserved even if data is empty
- cols_to_drop = []
- for col_idx, col in enumerate(df.columns):
- # Check if all data values are NaN
- # Use iloc to avoid ambiguity when MultiIndex has duplicate tuple keys
- if df.iloc[:, col_idx].isna().all():
- # Check if the column header is meaningful
- # For MultiIndex: check if any level has a non-empty meaningful value
- # For simple index: check if the header is not None/empty
- has_meaningful_header = False
- if isinstance(col, tuple):
- # MultiIndex column - check if any level has meaningful content
- for level in col:
- if (
- level
- and str(level).strip()
- and str(level).strip() not in ["None", "nan", "NaN"]
- ):
- has_meaningful_header = True
- break
- else:
- # Simple column name
- if (
- col
- and str(col).strip()
- and str(col).strip() not in ["None", "nan", "NaN"]
- ):
- has_meaningful_header = True
-
- # Only drop if header is not meaningful
- if not has_meaningful_header:
- cols_to_drop.append(col_idx)
-
- if cols_to_drop:
- # Use positional indices to drop columns safely (avoids duplicate MultiIndex key issues)
- cols_to_keep = [i for i in range(len(df.columns)) if i not in cols_to_drop]
- df = df.iloc[:, cols_to_keep]
-
- logger.debug(f"Dropped {len(cols_to_drop)} empty columns")
-
- # Preserve meaningful row index (header columns) as regular columns
- # Only drop=True if it's a simple RangeIndex (no semantic meaning)
- if not isinstance(df.index, pd.RangeIndex):
- # Remember if columns were MultiIndex before reset
- was_multiindex = isinstance(df.columns, pd.MultiIndex)
- n_levels = df.columns.nlevels if was_multiindex else 1
-
- # Avoid name collision before reset_index().
- # Two collision sources:
- # A) An index level name, when padded into a tuple by pandas,
- # matches an existing column.
- # B) Multiple index levels share the same name → pandas tries
- # to insert duplicate columns (e.g. five levels all named
- # one merged header repeated across five padded columns.
- # Strategy: de-duplicate index.names so every level gets a unique
- # column name during reset_index, then clean up afterwards.
- existing_col_set = set(df.columns)
-
- def _make_padded(name):
- """Simulate the column name pandas would create for this index level."""
- if was_multiindex:
- return (name,) + ("",) * (n_levels - 1)
- return name
-
- if isinstance(df.index, pd.MultiIndex):
- seen_counts = {} # name → how many times seen so far
- deduped = []
- for n in df.index.names:
- if n is None:
- deduped.append(None)
- continue
- padded = _make_padded(n)
- # Collision with existing column OR with a previously-seen index name
- if padded in existing_col_set or n in seen_counts:
- deduped.append(
- None
- ) # let pandas auto-name it (level_0, level_1 …)
- else:
- deduped.append(n)
- seen_counts[n] = seen_counts.get(n, 0) + 1
- df.index.names = deduped
- elif hasattr(df.index, "name") and df.index.name is not None:
- padded = _make_padded(df.index.name)
- if padded in existing_col_set:
- df.index.name = None
-
- df = df.reset_index() # Converts index to columns
-
- # Clean up auto-generated column names like 'index', 'level_0', 'level_1'
- # For MultiIndex columns, we need to preserve the structure
- if was_multiindex:
- # Build new column tuples for the index columns
- new_cols = []
- for col in df.columns:
- if isinstance(col, str) and (
- col.startswith("level_") or col == "index"
- ):
- # Create a tuple with empty strings to match MultiIndex levels
- new_cols.append(tuple([""] * n_levels))
- else:
- new_cols.append(col)
- df.columns = pd.MultiIndex.from_tuples(new_cols)
- else:
- # For simple columns
- new_cols = []
- for col in df.columns:
- if isinstance(col, str) and (
- col.startswith("level_") or col == "index"
- ):
- new_cols.append("")
- else:
- new_cols.append(col)
- df.columns = new_cols
- else:
- df.reset_index(drop=True, inplace=True)
-
- # Clean column names - preserve MultiIndex structure if present
- if isinstance(df.columns, pd.MultiIndex):
- # For MultiIndex, clean each level's values while preserving structure
- new_levels = []
- for level_idx in range(df.columns.nlevels):
- level_vals = df.columns.get_level_values(level_idx)
- cleaned = [
- str(v).replace("\n", "") if v is not None else "" for v in level_vals
- ]
- new_levels.append(cleaned)
- df.columns = pd.MultiIndex.from_arrays(new_levels, names=df.columns.names)
- # Also handle 'Unnamed' in MultiIndex
- new_levels = []
- for level_idx in range(df.columns.nlevels):
- level_vals = df.columns.get_level_values(level_idx)
- cleaned = [np.nan if "Unnamed" in str(v) else v for v in level_vals]
- new_levels.append(cleaned)
- df.columns = pd.MultiIndex.from_arrays(new_levels, names=df.columns.names)
- else:
- df.columns = [
- str(col).replace("\n", "") for col in df.columns
- ] # Replace '\n' in column headers
- df.columns = [
- np.nan if "Unnamed" in str(col) else col for col in df.columns
- ] # Replace Unnamed with nan
- df = df.map(
- lambda x: x.replace("\n", "") if isinstance(x, str) else x
- ) # Replace '\n' in each cell
- df = process_datetime_cells(df)
- return df
-
-
-def process_datetime_cells(df):
- df = df.copy()
-
- def convert(x):
- if isinstance(x, (pd.Timestamp, datetime.datetime)):
- return x.strftime("%Y-%m-%d %H:%M:%S")
- return x
-
- return df.apply(lambda col: col.map(convert))
-
-
-def process_duplicate_cols(columns):
- col_count = {}
- new_columns = []
- for col in columns:
- if col in col_count:
- new_columns.append(f"{col}>>>{col_count[col]}")
- col_count[col] += 1
- else:
- new_columns.append(col)
- col_count[col] = 1
- return new_columns
-
-
-def format_tb_scope(df, num):
- if len(df) > int(num * 3 + 1):
- # Get head and tail rows
- head_df = df.head(num)
- tail_df = df.tail(num)
- # Middle portion excluding head and tail
- middle_df = df.iloc[num : len(df) - num]
-
- if len(middle_df) >= num:
- mid_sample_df = middle_df.sample(n=num, random_state=42)
- else: # If middle has less than num rows, take all
- mid_sample_df = middle_df
- scope_df = pd.concat(objs=[head_df, mid_sample_df, tail_df], ignore_index=True)
- else:
- scope_df = df
- scope_df = scope_df.applymap(lambda x: str(x).strip() if pd.notnull(x) else x)
- scope_str = df2html(scope_df)
- return scope_str
-
-
-def parse_xlsx(
- file_path,
- file_name,
- output_dir,
- baseurl,
- base_llm_paras=None,
- window_h=10,
- relative_root=None,
- use_precision_mode=True,
- include_hidden_sheets=False,
-):
- """
- Parse Excel file and extract table content.
-
- Args:
- file_path: Path or URL to the Excel file
- file_name: Display name for the file
- output_dir: Directory to save extracted tables
- baseurl: Base URL for file loading
- base_llm_paras: LLM parameters for summarization
- window_h: Window size for table scope
- relative_root: Root path for relative paths
- use_precision_mode: If True, use openpyxl merged cell metadata for accurate
- header detection. If False, use LLM/heuristic mode.
- Default is True for better accuracy.
- include_hidden_sheets: If True, parse hidden/very-hidden sheets. Default False.
-
- Returns:
- DataFrame with parsed table information
- """
- time_stamp = get_str_time()
- df_list = []
-
- table_data = load_file_bytes(file_path, file_url=baseurl)
- table_stream = io.BytesIO(table_data)
-
- tb_dir = os.path.join(output_dir, "tables")
- os.makedirs(tb_dir, exist_ok=True)
- all_tb_paths = []
- exist_sheets = []
-
- if use_precision_mode:
- # PRECISION MODE: Use openpyxl metadata for accurate header detection
- logger.info("Using precision mode for Excel header detection")
- try:
- sheets_dict = parse_headers_from_excel(
- table_stream, include_hidden_sheets=include_hidden_sheets
- )
- precision_mode_active = True
- except Exception as e:
- logger.warning(f"Precision mode failed, falling back to legacy mode: {e}")
- table_stream.seek(0) # Reset stream position
- sheets_dict = pd.read_excel(table_stream, sheet_name=None)
- precision_mode_active = False
- else:
- # LEGACY MODE: Use pandas read_excel + LLM/heuristic header detection
- sheets_dict = pd.read_excel(table_stream, sheet_name=None)
- precision_mode_active = False
-
- all_sheets = sheets_dict.items()
-
- for sheet_name, sheet_content in all_sheets:
- sheet_name = sheet_name.strip()
- if sheet_name in exist_sheets:
- sheet_name = sheet_name + str(len(exist_sheets))
- else:
- exist_sheets.append(sheet_name)
-
- sheet_tbs = [sheet_content]
- for tb in sheet_tbs:
- try:
- tb = postprocess_tb(tb, drop=True)
- if len(tb) == 0 or tb.empty or tb.isna().all().all():
- continue
-
- # In precision mode, headers are already correctly set by parse_headers_from_excel
- # In legacy mode, use LLM/heuristic header parsing
- if not precision_mode_active:
- tb = parse_headers(tb, paras=base_llm_paras)
-
- # Drop _src_row column before converting to HTML/keywords
- # (_src_row is a debug column added by _parse_subtable for cross-referencing)
- src_row_cols = [
- c
- for c in tb.columns
- if (isinstance(c, tuple) and c[0] == "_src_row") or c == "_src_row"
- ]
- if src_row_cols:
- tb = tb.drop(columns=src_row_cols)
-
- # Get row header column count from DataFrame attrs (set in parse_headers_from_excel)
- row_header_cols = tb.attrs.get("row_header_cols", 0)
-
- tb_paths, tb_strs = parse_tb_contents(
- tb,
- parent_dic={file_name: {sheet_name: {}}},
- file_name=file_name,
- sheet_name=sheet_name,
- row_header_cols=row_header_cols,
- )
-
- # Unified LLM extraction: title + keywords + summary in one call
- # (consistent with doc_parser.py and md_parser.py)
- llm_title = None
- llm_summary = None
- tb_keywords = ""
- if base_llm_paras["summary_table"]:
- from app.services.document_parser.txt_parser import (
- extract_title_keywords_summary,
- )
-
- llm_title, tb_keywords, llm_summary = (
- extract_title_keywords_summary(tb_strs, max_keywords=3)
- )
-
- # Build tb_summary: table index + optional LLM summary
- table_index = f"table-{sheet_name}"
- if llm_summary:
- tb_summary = f"{table_index}\n{llm_summary}"
- else:
- # Fallback: use mechanical column keywords when LLM is off
- tb_keywords_fallback = parse_tb_keywords(tb)
- tb_summary = table_index
- tb_keywords = tb_keywords if tb_keywords else tb_keywords_fallback
-
- # Use a filesystem-safe filename so LLM titles like "A/B" do not
- # accidentally create nested paths under tables/.
- effective_name = llm_title if llm_title else sheet_name
- tb_name = (
- path_handle(
- remove_spaces("table-" + effective_name), mode="clean_single"
- )
- + ".html"
- )
- tb_path = os.path.join(tb_dir, tb_name)
- soup = BeautifulSoup(tb_strs, features="html.parser")
- tb_html_str = soup.prettify()
- with open(tb_path, "w", encoding="utf-8") as f:
- f.write(tb_html_str)
-
- # Use same temp_uid for both marker and know_id (aligned with doc_parser/md_parser)
- temp_uid = gen_str_codes(tb_strs + str(sheet_name))
- relative_tb_path = f"tables/{tb_name}"
- tb_ref = build_chunk_ref(relative_tb_path)
- tb_bottom_content = f"{tb_ref}\nTable summary:\n{tb_summary}\nMain columns:\n{tb_keywords}"
-
- bottom_tokens = tokenize2stw_remove(
- [tb_bottom_content], base_llm_paras["stopwords"]
- )
-
- all_tb_paths.extend(tb_paths)
- # Use relative path for tables: "tables/xxx.html"
- df_list.append(
- [
- tb_bottom_content,
- relative_tb_path,
- "table",
- len(tb_strs),
- tb_keywords,
- tb_summary,
- temp_uid,
- bottom_tokens,
- "",
- time_stamp,
- "",
- ]
- )
-
- except KnowhereException:
- raise
- except Exception as e:
- logger.error(f"Table parsing failed: {e}")
- raise TableParsingException(
- user_message="Failed to parse Excel table content",
- reason="TABLE_PROCESSING_FAILED",
- internal_message=str(e),
- original_exception=e,
- )
-
- table_df = pd.DataFrame(df_list, columns=settings.ALL_DF_COLS.split(","))
- table_df = process_dup_paths_df(table_df)
- return table_df
diff --git a/apps/worker/app/services/document_parser/table_text_parser.py b/apps/worker/app/services/document_parser/table_text_parser.py
new file mode 100644
index 000000000..261bbe653
--- /dev/null
+++ b/apps/worker/app/services/document_parser/table_text_parser.py
@@ -0,0 +1,149 @@
+# pyright: reportArgumentType=false
+from __future__ import annotations
+
+import io
+import re
+import unicodedata
+
+import pandas as pd
+from bs4 import BeautifulSoup, Tag
+
+_MAX_TABLE_NAME_CHARS = 80
+
+
+def sanitize_table_name_from_header(raw_header_text: str) -> str:
+ """Build a concise, filesystem-safe table name from raw first-row header text."""
+ from shared.utils.text_utils import _is_meaningful_token
+
+ if not raw_header_text:
+ return ""
+
+ parts = re.split(r"\s*\|\s*|_+br_|\n", raw_header_text)
+
+ seen: set[str] = set()
+ unique: list[str] = []
+ for part in parts:
+ part = part.strip()
+ if not part or part in seen:
+ continue
+ seen.add(part)
+ unique.append(part)
+
+ meaningful = [field for field in unique if _is_meaningful_token(field)]
+ result = " ".join(meaningful)
+ if len(result) > _MAX_TABLE_NAME_CHARS:
+ result = result[:_MAX_TABLE_NAME_CHARS].rstrip()
+ return result
+
+
+def identify_tables(line: str) -> tuple[bool, str | None, list[str] | None]:
+ """Identify whether one logical Markdown line contains a table."""
+ html_table_pattern = r".*? | "
+ tables = re.findall(html_table_pattern, line, re.DOTALL)
+ if bool(tables):
+ return True, "html", tables
+
+ if line.startswith("|") and line.endswith("|"):
+ return True, "md", []
+
+ return False, None, None
+
+
+def df2md(table_frame: pd.DataFrame, *, index: bool = False, na_rep: str = "—") -> str:
+ """Convert a DataFrame to a Markdown table while preserving display width."""
+
+ def get_display_width(text: str) -> int:
+ width = 0
+ for character in text:
+ if unicodedata.east_asian_width(character) in ("F", "W"):
+ width += 2
+ else:
+ width += 1
+ return width
+
+ def pad_to_width(text: str, target_width: int) -> str:
+ current_width = get_display_width(text)
+ padding = target_width - current_width
+ return text + " " * max(0, padding)
+
+ table_frame = table_frame.copy()
+
+ if index:
+ table_frame = table_frame.reset_index()
+
+ table_frame = table_frame.fillna(na_rep).astype(str)
+
+ column_widths: dict[object, int] = {}
+ for column in table_frame.columns:
+ header_width = get_display_width(str(column))
+ max_content_width = (
+ max(table_frame[column].apply(get_display_width))
+ if len(table_frame) > 0
+ else 0
+ )
+ column_widths[column] = max(header_width, max_content_width)
+
+ header_cells = [
+ pad_to_width(str(column), column_widths[column])
+ for column in table_frame.columns
+ ]
+ header_line = "| " + " | ".join(header_cells) + " |"
+
+ separator_cells = ["-" * column_widths[column] for column in table_frame.columns]
+ separator_line = "|-" + "-|-".join(separator_cells) + "-|"
+
+ data_lines: list[str] = []
+ for _, row in table_frame.iterrows():
+ cells = [
+ pad_to_width(str(row[column]), column_widths[column])
+ for column in table_frame.columns
+ ]
+ data_lines.append("| " + " | ".join(cells) + " |")
+
+ return "\n".join([header_line, separator_line, *data_lines])
+
+
+def clean_html_tb(html: str) -> str:
+ soup = BeautifulSoup(html, "html.parser")
+ for row in soup.find_all("tr"):
+ if not isinstance(row, Tag):
+ continue
+ seen: set[bytes] = set()
+ unique_cells: list[Tag] = []
+ for cell in row.find_all("td", recursive=False):
+ if not isinstance(cell, Tag):
+ continue
+ content = cell.encode_contents()
+ if content not in seen:
+ seen.add(content)
+ unique_cells.append(cell)
+ row.clear()
+ for cell in unique_cells:
+ row.append(cell)
+ return str(soup.prettify())
+
+
+def extract_tables_by_forms(table_text: str, form: str) -> str | None:
+ if form == "html":
+ return table_text
+
+ if form != "md":
+ return None
+
+ table_frame = pd.read_table(
+ io.StringIO(table_text),
+ sep="|",
+ engine="python",
+ on_bad_lines="skip",
+ )
+ table_frame = table_frame.iloc[:, 1:-1]
+ table_frame.columns = table_frame.columns.astype(str).str.strip()
+
+ separator_pattern = r"^[\s\-:]+$"
+ table_frame = table_frame[
+ ~table_frame.apply(
+ lambda row: row.astype(str).str.match(separator_pattern).all(),
+ axis=1,
+ )
+ ]
+ return table_frame.to_html(index=False)
diff --git a/apps/worker/app/services/document_parser/text_helpers.py b/apps/worker/app/services/document_parser/text_helpers.py
new file mode 100644
index 000000000..b11fdc772
--- /dev/null
+++ b/apps/worker/app/services/document_parser/text_helpers.py
@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import re
+
+from shared.utils.text_utils import _CN_EN_NUM_RE
+
+_CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]")
+EN_START_LIMIT = 15
+CN_RATIO_THRESHOLD = 0.3
+
+
+def normalize_md(text: str) -> str:
+ """Normalize markdown string for comparison."""
+ text = re.sub(r"^\s*#+\s*", "", text)
+ text = re.sub(r"\s+", "", text)
+ return text.lower()
+
+
+def truncate_text(text: str, start_limit: int, end_limit: int) -> str:
+ """Truncate text by raw character count, keeping start and end parts."""
+ text = str(text)
+ total_limit = start_limit + end_limit
+ if len(text) <= total_limit:
+ return text
+ start_part = text[:start_limit]
+ end_part = text[-end_limit:] if end_limit > 0 else ""
+ return f"{start_part}...{end_part}"
+
+
+def detect_primary_lang(text: str) -> str:
+ """Detect whether text is primarily Chinese or English/other."""
+ if not text:
+ return "en"
+ tokens = _CN_EN_NUM_RE.findall(text)
+ if not tokens:
+ return "en"
+ cn_count = sum(1 for token in tokens if _CN_CHAR_RE.fullmatch(token))
+ return "zh" if (cn_count / len(tokens)) >= CN_RATIO_THRESHOLD else "en"
+
+
+def count_cn_en(text: str) -> int:
+ """Count semantic Chinese/English/number tokens in a string."""
+ return len(_CN_EN_NUM_RE.findall(str(text)))
+
+
+def truncate_text_by_tokens(
+ text: str,
+ start_limit: int,
+ end_limit: int,
+ lang_aware: bool = True,
+) -> str:
+ """Truncate text by semantic token count, preserving whole words."""
+ text = str(text)
+ matches = list(_CN_EN_NUM_RE.finditer(text))
+ total = len(matches)
+
+ if lang_aware and total > 0 and detect_primary_lang(text) == "en":
+ start_limit = min(start_limit, EN_START_LIMIT)
+
+ if total <= start_limit + end_limit:
+ return text
+
+ cut_start = matches[start_limit - 1].end() if start_limit > 0 else 0
+ cut_end = matches[total - end_limit].start() if end_limit > 0 else len(text)
+ if cut_start >= cut_end:
+ return text
+ return text[:cut_start] + "..." + text[cut_end:]
diff --git a/apps/worker/app/services/document_parser/toc_docx.py b/apps/worker/app/services/document_parser/toc_docx.py
new file mode 100644
index 000000000..86d612add
--- /dev/null
+++ b/apps/worker/app/services/document_parser/toc_docx.py
@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+import re
+
+from app.services.document_parser.heading_candidates import (
+ judge_by_conditions,
+ remove_by_conditions,
+)
+from app.services.document_parser.toc_hierarchy import build_toc_hierarchy_payload
+import lxml.etree as etree
+
+TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"}
+
+
+def parse_w_int_attr(elem, ns, attr_names):
+ if elem is None:
+ return None
+
+ for attr_name in attr_names:
+ raw_val = elem.get("{%s}%s" % (ns["w"], attr_name))
+ if raw_val is None:
+ continue
+ try:
+ return int(raw_val)
+ except (TypeError, ValueError):
+ continue
+ return None
+
+
+def get_docx_toc_layout_hints(elem, ns):
+ ppr = elem.find("./w:pPr", namespaces=ns)
+ if ppr is None:
+ ppr = elem.find(".//w:pPr", namespaces=ns)
+
+ if ppr is None:
+ return {
+ "outline_level": None,
+ "left_indent": None,
+ }
+
+ outline_elem = ppr.find("./w:outlineLvl", namespaces=ns)
+ outline_level = None
+ if outline_elem is not None:
+ outline_val = parse_w_int_attr(outline_elem, ns, ["val"])
+ if outline_val is not None:
+ outline_level = outline_val + 1
+
+ indent_elem = ppr.find("./w:ind", namespaces=ns)
+ left_indent = parse_w_int_attr(
+ indent_elem, ns, ["left", "start", "leftChars", "startChars"]
+ )
+
+ return {
+ "outline_level": outline_level,
+ "left_indent": left_indent,
+ }
+
+
+def infer_toc_level_from_text(text: str):
+ text_clean = str(text).strip()
+ if not text_clean:
+ return None
+
+ normalized = re.sub(r"\s+", " ", text_clean).lower()
+ if normalized in TOC_TITLE_KEYWORDS:
+ return None
+
+ raw_pos_code = judge_by_conditions(text_clean)
+ if not isinstance(raw_pos_code, list):
+ return None
+ positive_codes = [
+ int(value)
+ for value in raw_pos_code
+ if isinstance(value, int) and value > 0
+ ]
+ neg_code = remove_by_conditions(text_clean)
+ if any(value > 0 for value in neg_code) or not positive_codes:
+ return None
+
+ return max(positive_codes)
+
+
+def is_toc_title_text(text: str) -> bool:
+ normalized = re.sub(r"\s+", " ", str(text).strip()).lower()
+ return normalized in TOC_TITLE_KEYWORDS
+
+
+def infer_toc_levels_from_indentation(entries: list) -> None:
+ indent_values = sorted(
+ {
+ entry["left_indent"]
+ for entry in entries
+ if entry.get("level") is None
+ and entry.get("left_indent") is not None
+ and not is_toc_title_text(entry.get("heading", ""))
+ }
+ )
+
+ if not indent_values:
+ return
+
+ indent_to_level = {indent: idx + 1 for idx, indent in enumerate(indent_values)}
+ for entry in entries:
+ if entry.get("level") is not None:
+ continue
+ if is_toc_title_text(entry.get("heading", "")):
+ continue
+ left_indent = entry.get("left_indent")
+ if left_indent is None:
+ continue
+ entry["level"] = indent_to_level.get(left_indent)
+
+
+def get_docx_toc_style_info(elem, ns):
+ style = elem.find(".//w:pPr/w:pStyle", namespaces=ns)
+ if style is None:
+ return {
+ "is_toc_style": False,
+ "toc_level": None,
+ "style_name": None,
+ }
+
+ val = style.get("{%s}val" % ns["w"])
+ if not val:
+ return {
+ "is_toc_style": False,
+ "toc_level": None,
+ "style_name": None,
+ }
+
+ val_lower = val.lower().strip()
+ if "toc" not in val_lower and "目录" not in val:
+ return {
+ "is_toc_style": False,
+ "toc_level": None,
+ "style_name": val,
+ }
+
+ level = None
+ match = re.search(r"(?:toc|目录)\s*[_-]?(\d+)$", val_lower)
+ if match:
+ level = int(match.group(1))
+
+ layout_hints = get_docx_toc_layout_hints(elem, ns)
+ if level is None:
+ level = layout_hints["outline_level"]
+
+ return {
+ "is_toc_style": True,
+ "toc_level": level,
+ "style_name": val,
+ "outline_level": layout_hints["outline_level"],
+ "left_indent": layout_hints["left_indent"],
+ }
+
+
+def get_toc_level(elem, ns):
+ style_info = get_docx_toc_style_info(elem, ns)
+ if not style_info["is_toc_style"]:
+ return False
+
+ if style_info["toc_level"] is not None:
+ return style_info["toc_level"]
+ return True
+
+
+def detect_sdt_toc(elem, ns):
+ tag = etree.QName(elem.tag).localname if isinstance(elem.tag, str) else None
+
+ if tag != "sdt":
+ return {"is_toc_sdt": False, "gallery_type": None}
+
+ is_toc_sdt = False
+ gallery_type = None
+
+ sdt_pr = elem.find(".//w:sdtPr", namespaces=ns)
+ if sdt_pr is not None:
+ doc_part_obj = sdt_pr.find(".//w:docPartObj", namespaces=ns)
+ if doc_part_obj is not None:
+ doc_part_gallery = doc_part_obj.find(".//w:docPartGallery", namespaces=ns)
+ if doc_part_gallery is not None:
+ gallery_type = doc_part_gallery.get("{%s}val" % ns["w"])
+ if gallery_type and "table of contents" in gallery_type.lower():
+ is_toc_sdt = True
+
+ return {"is_toc_sdt": is_toc_sdt, "gallery_type": gallery_type}
+
+
+def detect_doc_tocs(elem, ns):
+ style_info = get_docx_toc_style_info(elem, ns)
+ is_style = style_info["is_toc_style"]
+ is_field_start = False
+
+ instrs = elem.findall(".//w:instrText", namespaces=ns)
+ for instr in instrs:
+ if instr.text:
+ instr_text_stripped = instr.text.strip()
+ instr_text_lower = instr_text_stripped.lower()
+ if (
+ instr_text_lower.startswith("toc")
+ or "table of contents" in instr_text_lower
+ or "目录" in instr_text_stripped
+ ):
+ is_field_start = True
+ break
+
+ is_field_end = False
+ fldchars = elem.findall(".//w:fldChar", namespaces=ns)
+ for fld in fldchars:
+ if fld.get("{%s}fldCharType" % ns["w"]) == "end":
+ is_field_end = True
+ break
+
+ return {
+ "is_style": is_style,
+ "toc_level": style_info["toc_level"],
+ "style_name": style_info["style_name"],
+ "outline_level": style_info.get("outline_level"),
+ "left_indent": style_info.get("left_indent"),
+ "is_field_start": is_field_start,
+ "is_field_end": is_field_end,
+ }
+
+
+def build_docx_toc_hierarchies(block_tuples: list) -> list:
+ toc_areas = []
+ current_area = []
+
+ for ele_num, block, label, meta in block_tuples:
+ if "TOC" in label:
+ current_area.append((ele_num, block, meta or {}))
+ continue
+
+ if current_area:
+ toc_areas.append(current_area)
+ current_area = []
+
+ if current_area:
+ toc_areas.append(current_area)
+
+ toc_hierarchies = []
+ for area in toc_areas:
+ toc_entries = []
+ for ele_num, block, meta in area:
+ toc_level = meta.get("toc_level")
+ try:
+ toc_level = int(toc_level) if toc_level is not None else None
+ except (TypeError, ValueError):
+ toc_level = None
+
+ text = getattr(block, "text", str(block)).strip()
+ if not text:
+ continue
+
+ if toc_level is None:
+ outline_level = meta.get("toc_outline_level")
+ try:
+ toc_level = (
+ int(outline_level) if outline_level is not None else None
+ )
+ except (TypeError, ValueError):
+ toc_level = None
+
+ if toc_level is None:
+ toc_level = infer_toc_level_from_text(text)
+
+ left_indent = meta.get("toc_left_indent")
+ try:
+ left_indent = int(left_indent) if left_indent is not None else None
+ except (TypeError, ValueError):
+ left_indent = None
+
+ toc_entries.append(
+ {
+ "id": ele_num,
+ "heading": text,
+ "level": toc_level if toc_level and toc_level > 0 else None,
+ "left_indent": left_indent,
+ }
+ )
+
+ infer_toc_levels_from_indentation(toc_entries)
+ payload = build_toc_hierarchy_payload(
+ toc_entries,
+ toc_range=(area[0][0], area[-1][0]),
+ scan_range=(area[0][0], area[-1][0]),
+ )
+ if payload:
+ toc_hierarchies.append(payload)
+
+ return toc_hierarchies
diff --git a/apps/worker/app/services/document_parser/toc_hierarchy.py b/apps/worker/app/services/document_parser/toc_hierarchy.py
new file mode 100644
index 000000000..b88420e32
--- /dev/null
+++ b/apps/worker/app/services/document_parser/toc_hierarchy.py
@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import pandas as pd
+from app.services.document_parser.layout_parser import hiearchy_llm
+from app.services.document_parser.stage_profiler import stage_timer
+from app.services.document_parser.table_text_parser import df2md
+from app.services.document_parser.text_helpers import normalize_md
+from loguru import logger
+from pandas import Index
+
+from shared.core.config import settings
+
+
+def resolve_hierarchy_model_name(model_name: str | None = None) -> str:
+ return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL
+
+
+def parse_toc_hierarchy(
+ toc_df: pd.DataFrame, max_depth: int = 6, model_name: str | None = None
+) -> list[dict]:
+ resolved_model_name = resolve_hierarchy_model_name(model_name)
+ try:
+ with stage_timer(
+ "toc.parse_hierarchy_llm",
+ model_name=resolved_model_name,
+ heading_count=len(toc_df),
+ max_depth=max_depth,
+ ):
+ toc_hierarchy = hiearchy_llm(
+ toc_df,
+ model_name=resolved_model_name,
+ max_depth=max_depth,
+ task="eval-toc-headings",
+ )
+ id_to_level = {item["id"]: item["level"] for item in toc_hierarchy}
+
+ toc_with_level = []
+ for _, row in toc_df.iterrows():
+ line_id = row["id"]
+ heading = row["heading"]
+ level = id_to_level.get(line_id, 1)
+ toc_with_level.append({"id": line_id, "heading": heading, "level": level})
+ return toc_with_level
+
+ except Exception as exc:
+ logger.error(f"LLM hierarchy analysis failed: {exc}")
+ return []
+
+
+def build_tree_tocs(toc_with_level: list[dict]) -> dict:
+ if not toc_with_level:
+ return {}
+
+ positive_levels = [item["level"] for item in toc_with_level if item["level"] > 0]
+ level_for_minus_one = max(positive_levels) + 1 if positive_levels else 1
+
+ root = {}
+ stack = [(root, 0)]
+
+ for item in toc_with_level:
+ heading = item["heading"]
+ original_level = item["level"]
+ normalized_level = (
+ level_for_minus_one if original_level == -1 else original_level
+ )
+ while len(stack) > 1 and stack[-1][1] >= normalized_level:
+ stack.pop()
+
+ parent_dict = stack[-1][0]
+ parent_dict[heading] = {}
+ stack.append((parent_dict[heading], normalized_level))
+ return root
+
+
+def build_toc_hierarchy_payload(
+ toc_entries: list[dict],
+ toc_range: tuple | None = None,
+ scan_range: tuple | None = None,
+) -> dict | None:
+ valid_entries = []
+ for entry in toc_entries:
+ heading = str(entry.get("heading", "")).strip()
+ level = entry.get("level")
+ if not heading or not isinstance(level, int) or level <= 0:
+ continue
+
+ valid_entries.append(
+ {
+ "id": entry.get("id"),
+ "heading": heading,
+ "level": level,
+ }
+ )
+
+ if not valid_entries:
+ return None
+
+ toc_df = pd.DataFrame(valid_entries, columns=Index(["id", "heading", "level"]))
+ payload = {
+ "toc_range": toc_range or (valid_entries[0]["id"], valid_entries[-1]["id"]),
+ "toc_with_level": df2md(toc_df, index=False),
+ "toc_tree": build_tree_tocs(valid_entries),
+ }
+ if scan_range is not None:
+ payload["scan_range"] = scan_range
+ return payload
+
+
+def eval_toc_levels(
+ toc_lines: list[str], model_name: str | None = None, max_depth: int = 6
+) -> tuple[str, dict]:
+ toc_title_keywords = {"目录", "目次", "tableofcontents", "contents"}
+ valid_data = []
+
+ for index, line in enumerate(toc_lines):
+ heading = line.strip()
+ if not heading:
+ continue
+ if normalize_md(heading) in toc_title_keywords:
+ logger.debug(
+ f"eval_toc_levels: skipping TOC keyword title line id={index}: {heading[:60]}"
+ )
+ continue
+
+ valid_data.append({"id": index, "heading": heading, "level": "Not Sure"})
+
+ toc_df = pd.DataFrame(valid_data)
+
+ if toc_df.empty:
+ logger.info("No valid TOC content, skip hierarchy analysis")
+ return "", {}
+
+ llm_result = parse_toc_hierarchy(toc_df, max_depth, model_name)
+ id_to_level = {item["id"]: item["level"] for item in llm_result}
+
+ valid_items_for_tree = []
+ for data in valid_data:
+ line_id = data["id"]
+ heading = data["heading"]
+ level = id_to_level.get(line_id, -1)
+ if level > 0:
+ valid_items_for_tree.append(
+ {"id": line_id, "heading": heading, "level": level}
+ )
+
+ payload = build_toc_hierarchy_payload(valid_items_for_tree)
+ if not payload:
+ return "", {}
+ return payload["toc_with_level"], payload["toc_tree"]
diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py
index e192bea40..2efb1a4cd 100644
--- a/apps/worker/app/services/document_parser/toc_parser.py
+++ b/apps/worker/app/services/document_parser/toc_parser.py
@@ -5,311 +5,25 @@
Provides functionality for:
- Detecting TOC (Table of Contents) candidates in markdown documents
-- Detecting TOC in DOCX documents (SDT containers, styles, field codes)
- Using LLM to determine precise TOC boundaries
-- Analyzing TOC hierarchy structure
-- Building nested tree structures from TOC
"""
import re
import gevent
import pandas as pd
-from app.services.common.kb_utils import (
- normalize_md,
- truncate_text_by_tokens,
-)
-from app.services.document_parser.layout_parser import (
- hiearchy_llm,
- judge_by_conditions,
- remove_by_conditions,
-)
+from app.services.document_parser.toc_hierarchy import eval_toc_levels
+from app.services.document_parser.text_helpers import normalize_md, truncate_text_by_tokens
from app.services.document_parser.stage_profiler import stage_timer
-from app.services.document_parser.table_parser import df2md
+from app.services.document_parser.table_text_parser import df2md
from gevent.pool import Pool as GeventPool
from loguru import logger
-from lxml import etree
-from shared.core.config import settings
from shared.services.ai.prompt_service import build_prompt
from shared.services.ai.response_process_service import eval_response
from shared.utils.OpenAICompatibleClientSync import get_openai_client
-def _resolve_hierarchy_model_name(model_name: str | None = None) -> str:
- """Resolve dedicated hierarchy model, falling back to the normal model."""
- return model_name or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL
-
-
-# ==================== DOCX TOC Detection Functions ====================
-
-TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"}
-
-
-def _parse_w_int_attr(elem, ns, attr_names):
- """Parse the first integer-valued OOXML attribute from a list of names."""
- if elem is None:
- return None
-
- for attr_name in attr_names:
- raw_val = elem.get("{%s}%s" % (ns["w"], attr_name))
- if raw_val is None:
- continue
- try:
- return int(raw_val)
- except (TypeError, ValueError):
- continue
- return None
-
-
-def get_docx_toc_layout_hints(elem, ns):
- """Extract outline/indent hints for TOC paragraphs without numeric TOC styles."""
- ppr = elem.find("./w:pPr", namespaces=ns)
- if ppr is None:
- ppr = elem.find(".//w:pPr", namespaces=ns)
-
- if ppr is None:
- return {
- "outline_level": None,
- "left_indent": None,
- }
-
- outline_elem = ppr.find("./w:outlineLvl", namespaces=ns)
- outline_level = None
- if outline_elem is not None:
- outline_val = _parse_w_int_attr(outline_elem, ns, ["val"])
- if outline_val is not None:
- outline_level = outline_val + 1
-
- indent_elem = ppr.find("./w:ind", namespaces=ns)
- left_indent = _parse_w_int_attr(
- indent_elem, ns, ["left", "start", "leftChars", "startChars"]
- )
-
- return {
- "outline_level": outline_level,
- "left_indent": left_indent,
- }
-
-
-def infer_toc_level_from_text(text: str):
- """Fallback TOC level inference from numbering patterns in TOC text."""
- text_clean = str(text).strip()
- if not text_clean:
- return None
-
- normalized = re.sub(r"\s+", " ", text_clean).lower()
- if normalized in TOC_TITLE_KEYWORDS:
- return None
-
- pos_code = judge_by_conditions(text_clean)
- neg_code = remove_by_conditions(text_clean)
- if any(x > 0 for x in neg_code) or not any(x > 0 for x in pos_code):
- return None
-
- return max(int(x) for x in pos_code)
-
-
-def is_toc_title_text(text: str) -> bool:
- """Return True when the line is likely the standalone TOC heading itself."""
- normalized = re.sub(r"\s+", " ", str(text).strip()).lower()
- return normalized in TOC_TITLE_KEYWORDS
-
-
-def infer_toc_levels_from_indentation(entries: list) -> None:
- """Populate missing TOC levels by ranking paragraph indentation within one TOC area."""
- indent_values = sorted(
- {
- entry["left_indent"]
- for entry in entries
- if entry.get("level") is None
- and entry.get("left_indent") is not None
- and not is_toc_title_text(entry.get("heading", ""))
- }
- )
-
- if not indent_values:
- return
-
- indent_to_level = {indent: idx + 1 for idx, indent in enumerate(indent_values)}
- for entry in entries:
- if entry.get("level") is not None:
- continue
- if is_toc_title_text(entry.get("heading", "")):
- continue
- left_indent = entry.get("left_indent")
- if left_indent is None:
- continue
- entry["level"] = indent_to_level.get(left_indent)
-
-
-def get_docx_toc_style_info(elem, ns):
- """
- Parse TOC style metadata from a DOCX paragraph element.
-
- Returns:
- dict: {
- 'is_toc_style': bool,
- 'toc_level': Optional[int],
- 'style_name': Optional[str]
- }
- """
- style = elem.find(".//w:pPr/w:pStyle", namespaces=ns)
- if style is None:
- return {
- "is_toc_style": False,
- "toc_level": None,
- "style_name": None,
- }
-
- val = style.get("{%s}val" % ns["w"])
- if not val:
- return {
- "is_toc_style": False,
- "toc_level": None,
- "style_name": None,
- }
-
- val_lower = val.lower().strip()
- if "toc" not in val_lower and "目录" not in val:
- return {
- "is_toc_style": False,
- "toc_level": None,
- "style_name": val,
- }
-
- level = None
- match = re.search(r"(?:toc|目录)\s*[_-]?(\d+)$", val_lower)
- if match:
- level = int(match.group(1))
-
- layout_hints = get_docx_toc_layout_hints(elem, ns)
- if level is None:
- level = layout_hints["outline_level"]
-
- return {
- "is_toc_style": True,
- "toc_level": level,
- "style_name": val,
- "outline_level": layout_hints["outline_level"],
- "left_indent": layout_hints["left_indent"],
- }
-
-
-def get_toc_level(elem, ns):
- """
- Detect whether a paragraph uses a TOC style.
-
- Args:
- elem: XML paragraph element.
- ns: XML namespace map.
-
- Returns:
- bool: True when the paragraph uses a TOC style.
- """
- style_info = get_docx_toc_style_info(elem, ns)
- if not style_info["is_toc_style"]:
- return False
-
- if style_info["toc_level"] is not None:
- return style_info["toc_level"]
- return True
-
-
-def detect_sdt_toc(elem, ns):
- """
- Detect an SDT (Structured Document Tag) TOC container.
- Word-generated TOCs are often wrapped in ``sdt`` elements.
-
- Args:
- elem: SDT element.
- ns: XML namespace map.
-
- Returns:
- dict: {
- 'is_toc_sdt': bool - whether the element is a TOC SDT,
- 'gallery_type': str - docPartGallery type
- }
- """
- tag = etree.QName(elem.tag).localname if isinstance(elem.tag, str) else None
-
- if tag != "sdt":
- return {"is_toc_sdt": False, "gallery_type": None}
-
- is_toc_sdt = False
- gallery_type = None
-
- sdt_pr = elem.find(".//w:sdtPr", namespaces=ns)
- if sdt_pr is not None:
- doc_part_obj = sdt_pr.find(".//w:docPartObj", namespaces=ns)
- if doc_part_obj is not None:
- doc_part_gallery = doc_part_obj.find(".//w:docPartGallery", namespaces=ns)
- if doc_part_gallery is not None:
- gallery_type = doc_part_gallery.get("{%s}val" % ns["w"])
- if gallery_type and "table of contents" in gallery_type.lower():
- is_toc_sdt = True
-
- return {"is_toc_sdt": is_toc_sdt, "gallery_type": gallery_type}
-
-
-def detect_doc_tocs(elem, ns):
- """
- Detect TOC regions using two strategies:
- 1. paragraph style detection (TOC styles)
- 2. field code detection (instrText)
-
- Note: SDT container detection is handled by ``detect_sdt_toc``.
-
- Args:
- elem: XML paragraph element.
- ns: XML namespace map.
-
- Returns:
- dict: {
- 'is_style': bool - whether the paragraph uses a TOC style,
- 'is_field_start': bool - whether this starts a TOC field,
- 'is_field_end': bool - whether this ends a field
- }
- """
- style_info = get_docx_toc_style_info(elem, ns)
- is_style = style_info["is_toc_style"]
- is_field_start = False
-
- instrs = elem.findall(".//w:instrText", namespaces=ns)
- for instr in instrs:
- if instr.text:
- instr_text_stripped = instr.text.strip()
- instr_text_lower = instr_text_stripped.lower()
- # Match standalone TOC field commands, NOT "PAGEREF _TocXXXX"
- # TOC fields start with "TOC" as the command word
- if (
- instr_text_lower.startswith("toc")
- or "table of contents" in instr_text_lower
- or "目录" in instr_text_stripped
- ):
- is_field_start = True
- break
-
- is_field_end = False
- # Always check for fldChar end, even on TOC-styled paragraphs,
- # so that the outer TOC field boundary can be properly closed.
- fldchars = elem.findall(".//w:fldChar", namespaces=ns)
- for fld in fldchars:
- if fld.get("{%s}fldCharType" % ns["w"]) == "end":
- is_field_end = True
- break
-
- return {
- "is_style": is_style,
- "toc_level": style_info["toc_level"],
- "style_name": style_info["style_name"],
- "outline_level": style_info.get("outline_level"),
- "left_indent": style_info.get("left_indent"),
- "is_field_start": is_field_start,
- "is_field_end": is_field_end,
- }
-
-
# ==================== Markdown TOC Detection Functions ====================
@@ -658,275 +372,6 @@ def _judge_single_area(idx, lines_, invalid_ids, area_start, area_end):
return toc_ranges
-def parse_toc_hierarchy(toc_df, max_depth: int = 6, model_name: str = None) -> list:
- """
- Parse TOC hierarchy using LLM
-
- Args:
- toc_df: DataFrame with id, heading columns
- max_depth: max depth of hierarchy
- model_name: model name (optional)
-
- Returns:
- List of dicts with id, heading, level
- """
- resolved_model_name = _resolve_hierarchy_model_name(model_name)
- try:
- with stage_timer(
- "toc.parse_hierarchy_llm",
- model_name=resolved_model_name,
- heading_count=len(toc_df),
- max_depth=max_depth,
- ):
- toc_hierarchy = hiearchy_llm(
- toc_df,
- model_name=resolved_model_name,
- max_depth=max_depth,
- task="eval-toc-headings",
- )
- id_to_level = {item["id"]: item["level"] for item in toc_hierarchy}
-
- toc_with_level = []
- for _, row in toc_df.iterrows():
- line_id = row["id"]
- heading = row["heading"]
- level = id_to_level.get(line_id, 1)
- toc_with_level.append({"id": line_id, "heading": heading, "level": level})
- return toc_with_level
-
- except Exception as e:
- logger.error(f"LLM hierarchy analysis failed: {e}")
- return []
-
-
-def build_tree_tocs(toc_with_level: list) -> dict:
- """
- Build nested JSON from TOC with level
-
- Args:
- toc_with_level: [{"id": line index, "heading": content, "level": level, "reason": ...}, ...]
- level: 1 for h1, 2 for h2..., -1 will be treated as the lowest level title
-
- Returns:
- nested JSON structure
-
- Notes:
- in the TOC scenario, all lines are treated as titles:
- - normal levels (1,2,3...) are treated as is
- - -1 is treated as a level deeper than all normal levels
- """
- if not toc_with_level:
- return {}
-
- # Step 1: collect all levels (exclude -1)
- positive_levels = [item["level"] for item in toc_with_level if item["level"] > 0]
-
- # Step 2: determine the level that -1 should be mapped to
- if positive_levels:
- # if there are normal levels, -1 is mapped to max + 1
- max_positive_level = max(positive_levels)
- level_for_minus_one = max_positive_level + 1
- else:
- level_for_minus_one = 1
-
- # Step 3: build nested structure
- root = {}
- stack = [(root, 0)]
-
- for item in toc_with_level:
- heading = item["heading"]
- original_level = item["level"]
-
- # normalize level: -1 -> level_for_minus_one
- normalized_level = (
- level_for_minus_one if original_level == -1 else original_level
- )
- while len(stack) > 1 and stack[-1][1] >= normalized_level:
- stack.pop()
-
- parent_dict = stack[-1][0]
- parent_dict[heading] = {}
- stack.append((parent_dict[heading], normalized_level))
- return root
-
-
-def build_toc_hierarchy_payload(
- toc_entries: list,
- toc_range: tuple | None = None,
- scan_range: tuple | None = None,
-) -> dict | None:
- """
- Build a toc_hierarchies-compatible payload from structured TOC entries.
- """
- valid_entries = []
- for entry in toc_entries:
- heading = str(entry.get("heading", "")).strip()
- level = entry.get("level")
- if not heading or not isinstance(level, int) or level <= 0:
- continue
-
- normalized_entry = {
- "id": entry.get("id"),
- "heading": heading,
- "level": level,
- }
- valid_entries.append(normalized_entry)
-
- if not valid_entries:
- return None
-
- result_df = pd.DataFrame(valid_entries)
- payload = {
- "toc_range": toc_range or (valid_entries[0]["id"], valid_entries[-1]["id"]),
- "toc_with_level": df2md(result_df[["id", "heading", "level"]], index=False),
- "toc_tree": build_tree_tocs(valid_entries),
- }
- if scan_range is not None:
- payload["scan_range"] = scan_range
- return payload
-
-
-def build_docx_toc_hierarchies(block_tuples: list) -> list:
- """
- Convert DOCX TOC blocks into the same toc_hierarchies structure used by MD/PDF.
- """
- toc_areas = []
- current_area = []
-
- for ele_num, block, label, meta in block_tuples:
- if "TOC" in label:
- current_area.append((ele_num, block, meta or {}))
- continue
-
- if current_area:
- toc_areas.append(current_area)
- current_area = []
-
- if current_area:
- toc_areas.append(current_area)
-
- toc_hierarchies = []
- for area in toc_areas:
- toc_entries = []
- for ele_num, block, meta in area:
- toc_level = meta.get("toc_level")
- try:
- toc_level = int(toc_level) if toc_level is not None else None
- except (TypeError, ValueError):
- toc_level = None
-
- text = getattr(block, "text", str(block)).strip()
- if not text:
- continue
-
- if toc_level is None:
- outline_level = meta.get("toc_outline_level")
- try:
- toc_level = (
- int(outline_level) if outline_level is not None else None
- )
- except (TypeError, ValueError):
- toc_level = None
-
- if toc_level is None:
- toc_level = infer_toc_level_from_text(text)
-
- left_indent = meta.get("toc_left_indent")
- try:
- left_indent = int(left_indent) if left_indent is not None else None
- except (TypeError, ValueError):
- left_indent = None
-
- toc_entries.append(
- {
- "id": ele_num,
- "heading": text,
- "level": toc_level if toc_level and toc_level > 0 else None,
- "left_indent": left_indent,
- }
- )
-
- infer_toc_levels_from_indentation(toc_entries)
- payload = build_toc_hierarchy_payload(
- toc_entries,
- toc_range=(area[0][0], area[-1][0]),
- scan_range=(area[0][0], area[-1][0]),
- )
- if payload:
- toc_hierarchies.append(payload)
-
- return toc_hierarchies
-
-
-def eval_toc_levels(
- toc_lines: list, model_name: str = None, max_depth: int = 6
-) -> tuple:
- """
- Analyze TOC hierarchy and generate nested JSON
-
- Args:
- toc_lines: list of pre-filtered valid TOC lines (invalid content already removed)
- model_name: model name (optional)
- max_depth: max depth of hierarchy
-
- Returns:
- (toc_with_level, toc_tree)
- - toc_with_level: list with level information
- Format: [{"id": int, "heading": str, "level": int, "reason": str}, ...]
- - toc_tree: nested JSON structure
- """
- # Build data for LLM judgment (all lines are valid, pre-filtered)
- # TOC title trigger lines are excluded from
- # the LLM input: they stay within toc_range so they are stripped from md_lines,
- # but they must not be sent to the hierarchy LLM to avoid a spurious Level=1 entry.
- _toc_title_keywords = {"目录", "目次", "tableofcontents", "contents"}
- valid_data = []
-
- for i, line in enumerate(toc_lines):
- heading = line.strip()
- if not heading:
- continue
- if normalize_md(heading) in _toc_title_keywords:
- logger.debug(
- f"eval_toc_levels: skipping TOC keyword title line id={i}: {heading[:60]}"
- )
- continue
-
- valid_data.append({"id": i, "heading": heading, "level": "Not Sure"})
-
- toc_df = pd.DataFrame(valid_data)
-
- if toc_df.empty:
- logger.info("No valid TOC content, skip hierarchy analysis")
- return "", {}
-
- # Evaluate TOC hierarchy with LLM
- llm_result = parse_toc_hierarchy(toc_df, max_depth, model_name)
-
- # Build id -> level mapping from LLM result
- id_to_level = {item["id"]: item["level"] for item in llm_result}
-
- # Build final result
- result_data = []
- valid_items_for_tree = []
-
- for data in valid_data:
- line_id = data["id"]
- heading = data["heading"]
- level = id_to_level.get(line_id, -1)
-
- if level > 0:
- result_data.append({"id": line_id, "heading": heading, "level": level})
- valid_items_for_tree.append(
- {"id": line_id, "heading": heading, "level": level}
- )
-
- payload = build_toc_hierarchy_payload(valid_items_for_tree)
- if not payload:
- return "", {}
- return payload["toc_with_level"], payload["toc_tree"]
-
-
def detect_tocs_in_texts(
md_lines: list,
model_name: str = None,
diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/txt_parser.py
index 3851817b5..e1d41956b 100755
--- a/apps/worker/app/services/document_parser/txt_parser.py
+++ b/apps/worker/app/services/document_parser/txt_parser.py
@@ -12,7 +12,7 @@
from shared.services.ai.prompt_service import build_prompt
from shared.services.ai.response_process_service import eval_response
from shared.utils.chunk_refs import CHUNK_REF_PATTERN
-from shared.utils.CommonHelperSync import load_file_bytes
+from shared.utils.file_loading import load_file_bytes
from shared.utils.OpenAICompatibleClientSync import get_openai_client
diff --git a/apps/worker/app/services/storage/sync_storage_service.py b/apps/worker/app/services/storage/sync_storage_service.py
deleted file mode 100644
index 268c68b90..000000000
--- a/apps/worker/app/services/storage/sync_storage_service.py
+++ /dev/null
@@ -1,141 +0,0 @@
-"""
-Sync storage operations for worker tasks.
-Provides S3 file operations and HTTP file downloads using sync adapters
-that yield cooperatively under gevent.
-"""
-
-import os
-import tempfile
-from typing import Any, Dict, Optional
-
-from loguru import logger
-
-from shared.core.config import settings
-from shared.core.config.storage import get_cached_storage_adapter
-from shared.core.exceptions.domain_exceptions import StorageServiceException
-from shared.utils.pinned_outbound_http import download_pinned_outbound_file
-from shared.utils.url_security import validate_http_url_and_resolve_ip
-
-
-def get_storage_adapter():
- """Get the storage adapter for direct sync S3 operations."""
- return get_cached_storage_adapter()
-
-
-def verify_s3_file_exists(s3_key: str, bucket: Optional[str] = None) -> Dict[str, Any]:
- """Verify S3 file exists using sync adapter calls."""
- adapter = get_storage_adapter()
- bucket_name = bucket or settings.S3_BUCKET_NAME
- try:
- if not adapter.exists(s3_key, bucket_name):
- return {"exists": False}
- size = adapter.get_object_size(s3_key, bucket_name)
- return {"exists": True, "size": size}
- except Exception as e:
- if "404" in str(e) or "not found" in str(e).lower():
- return {"exists": False}
- raise StorageServiceException(
- internal_message=f"S3 file verification failed: {e}",
- operation="verify_s3_file_exists",
- original_exception=e,
- )
-
-
-def generate_download_url(
- s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600
-) -> Dict[str, Any]:
- """Generate presigned download URL using sync adapter."""
- adapter = get_storage_adapter()
- bucket_name = bucket or settings.S3_BUCKET_NAME
- download_url = adapter.generate_presigned_url(
- s3_key, expiration=expires_in, bucket=bucket_name, method="GET"
- )
- return {"download_url": download_url, "expires_in": expires_in}
-
-
-def upload_to_s3(local_file_path: str, s3_key: str, bucket: str):
- """Upload file to S3 using sync adapter."""
- adapter = get_storage_adapter()
- adapter.upload_file(local_file_path, s3_key, bucket)
-
-
-def download_s3_object_to_temp(
- s3_key: str,
- suffix: str,
- temp_dir: str,
- bucket: Optional[str] = None,
-) -> str:
- """Download an object-storage file into a task-local temp file."""
- adapter = get_storage_adapter()
- bucket_name = bucket or settings.S3_BUCKET_NAME
- local_temp_path: str | None = None
-
- try:
- os.makedirs(temp_dir, exist_ok=True)
- with tempfile.NamedTemporaryFile(
- delete=False,
- suffix=suffix,
- dir=temp_dir,
- ) as temp_file:
- local_temp_path = temp_file.name
- adapter.download_file(s3_key, local_temp_path, bucket_name)
- return local_temp_path
- except Exception as e:
- if local_temp_path and os.path.exists(local_temp_path):
- os.remove(local_temp_path)
- raise StorageServiceException(
- internal_message=(
- f"Failed to download object-storage file to temp path: "
- f"s3_key={s3_key}, temp_dir={temp_dir}, error={e}"
- ),
- operation="download_s3_object_to_temp",
- original_exception=e,
- ) from e
-
-
-def upload_zip_result(job_id: str, zip_file_path: str) -> str:
- """Upload ZIP result file to S3 and cleanup temp file."""
- results_bucket = getattr(settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME)
- s3_key = f"results/{job_id}.zip"
- upload_to_s3(zip_file_path, s3_key, results_bucket)
- logger.info(f"Result ZIP uploaded: job_id={job_id}, key={s3_key}")
- try:
- if os.path.exists(zip_file_path):
- os.remove(zip_file_path)
- except Exception as e:
- logger.warning(f"Failed to cleanup temp ZIP: {e}")
- return s3_key
-
-
-def download_file_from_url(file_url: str) -> str:
- """Download a URL file through SSRF validation and IP pinning."""
- temp_file_path = ""
- try:
- validation = validate_http_url_and_resolve_ip(file_url)
- if not validation.is_valid or not validation.validated_ip:
- raise StorageServiceException(
- internal_message=f"Invalid URL: {validation.error_message}",
- operation="download_from_url",
- )
-
- temp_dir = getattr(settings, "TMP_PATH", "/tmp")
- os.makedirs(temp_dir, exist_ok=True)
- download_result = download_pinned_outbound_file(
- url=validation.url,
- pinned_ip=validation.validated_ip,
- timeout_seconds=300,
- user_agent="Knowhere-FileDownloader/1.0",
- temp_dir=temp_dir,
- )
- temp_file_path = download_result.temp_file_path
- return temp_file_path
- except StorageServiceException:
- raise
- except Exception as e:
- if os.path.exists(temp_file_path):
- os.remove(temp_file_path)
- raise StorageServiceException(
- internal_message=f"Failed to download file: {e}",
- operation="download_from_url",
- original_exception=e,
- )
diff --git a/apps/worker/app/services/workload/__init__.py b/apps/worker/app/services/workload/__init__.py
index b7df5ee6b..3ad3226c0 100644
--- a/apps/worker/app/services/workload/__init__.py
+++ b/apps/worker/app/services/workload/__init__.py
@@ -1,7 +1 @@
-"""
-Billing services for the worker.
-"""
-
-from .page_estimator import PageEstimator
-
-__all__ = ["PageEstimator"]
+"""Worker workload adapters."""
diff --git a/apps/worker/app/services/workload/url_upload_context.py b/apps/worker/app/services/workload/url_upload_context.py
new file mode 100644
index 000000000..af7fb807b
--- /dev/null
+++ b/apps/worker/app/services/workload/url_upload_context.py
@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from shared.core.exceptions.domain_exceptions import NotFoundException
+from shared.services.redis.redis_sync_service import (
+ SyncJobInfoRedisService,
+ SyncJobMetadataService,
+)
+
+
+@dataclass(frozen=True)
+class UrlUploadContext:
+ s3_key: str
+
+
+def load_url_upload_context(job_id: str, redis_service: Any) -> UrlUploadContext:
+ job_info_service = SyncJobInfoRedisService(redis_service)
+ job_info = job_info_service.get_job_info(job_id)
+
+ if job_info:
+ raw_s3_key = job_info.get("s3_key")
+ else:
+ metadata_service = SyncJobMetadataService(redis_service)
+ job_metadata = metadata_service.get_metadata(job_id)
+ if not job_metadata:
+ raise NotFoundException(
+ resource="JobInfo",
+ resource_id=job_id,
+ internal_message="Job info not found in Redis or Metadata",
+ )
+ raw_s3_key = job_metadata.get("s3_key")
+
+ if not raw_s3_key:
+ raise NotFoundException(
+ resource="JobInfo",
+ resource_id="s3_key",
+ internal_message=f"Missing s3_key in Redis job info for job_id={job_id}",
+ )
+
+ return UrlUploadContext(s3_key=str(raw_s3_key))
diff --git a/apps/worker/app/services/workload/url_upload_service.py b/apps/worker/app/services/workload/url_upload_service.py
new file mode 100644
index 000000000..cad1e2744
--- /dev/null
+++ b/apps/worker/app/services/workload/url_upload_service.py
@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+
+from app.services.workload.url_upload_context import load_url_upload_context
+from app.services.workload.url_upload_transfer import (
+ assert_temp_file_within_size_limit,
+ cleanup_temp_file,
+ download_source_url_to_temp,
+ resolve_supported_url_extension,
+ upload_temp_file_to_source_storage,
+ verify_source_upload,
+)
+from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service
+from shared.services.redis.redis_sync_service import SyncRedisServiceFactory
+
+
+def upload_url_file(
+ job_id: str,
+ source_url: str,
+ user_id: str | None,
+ job_type: str | None = None,
+) -> dict[str, Any]:
+ del user_id, job_type
+
+ lifecycle_service = get_sync_job_lifecycle_service()
+ redis_service = SyncRedisServiceFactory.get_service()
+ upload_context = load_url_upload_context(job_id, redis_service)
+
+ lifecycle_service.update_progress(
+ job_id, progress=3, message="Validating URL file type..."
+ )
+ file_extension = resolve_supported_url_extension(source_url)
+
+ lifecycle_service.update_progress(
+ job_id, progress=10, message="Downloading file from URL..."
+ )
+ temp_file_path = download_source_url_to_temp(source_url)
+
+ try:
+ lifecycle_service.update_progress(
+ job_id, progress=30, message="Validating file size..."
+ )
+ assert_temp_file_within_size_limit(
+ temp_file_path=temp_file_path,
+ file_extension=file_extension,
+ )
+
+ lifecycle_service.update_progress(
+ job_id, progress=50, message="Uploading file to S3..."
+ )
+ upload_temp_file_to_source_storage(
+ temp_file_path=temp_file_path,
+ s3_key=upload_context.s3_key,
+ )
+
+ finally:
+ cleanup_temp_file(temp_file_path)
+
+ lifecycle_service.update_progress(
+ job_id, progress=80, message="Verifying upload result..."
+ )
+ file_info = verify_source_upload(upload_context.s3_key)
+
+ lifecycle_service.update_progress(
+ job_id,
+ progress=100,
+ message="URL file upload complete, waiting for processing...",
+ )
+ logger.info(
+ "URL file upload complete, waiting for S3 webhook: "
+ f"{job_id} -> {upload_context.s3_key}"
+ )
+
+ return {
+ "status": "success",
+ "job_id": job_id,
+ "s3_key": upload_context.s3_key,
+ "file_size": file_info.get("size"),
+ }
diff --git a/apps/worker/app/services/workload/url_upload_transfer.py b/apps/worker/app/services/workload/url_upload_transfer.py
new file mode 100644
index 000000000..25f4faf13
--- /dev/null
+++ b/apps/worker/app/services/workload/url_upload_transfer.py
@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import os
+
+from loguru import logger
+
+from shared.core.config import settings
+from shared.core.exceptions.domain_exceptions import (
+ StorageServiceException,
+ ValidationException,
+)
+from shared.services.storage.job_file_storage import JobFileStorage
+from shared.utils.url_file_type import resolve_file_extension_sync
+
+
+def resolve_supported_url_extension(source_url: str) -> str:
+ file_extension = resolve_file_extension_sync(source_url)
+ if file_extension:
+ return file_extension
+
+ supported_formats = ", ".join(sorted(settings.get_supported_extensions()))
+ raise ValidationException(
+ user_message="Unsupported file type",
+ violations=[
+ {
+ "field": "file_extension",
+ "description": f"Must be one of: {supported_formats}",
+ }
+ ],
+ )
+
+
+def download_source_url_to_temp(source_url: str) -> str:
+ storage = JobFileStorage()
+ try:
+ return storage.download_file_from_url(
+ source_url,
+ temp_dir=getattr(settings, "TMP_PATH", "/tmp"),
+ )
+ except Exception as exc:
+ raise ValidationException(
+ user_message="Failed to download file from URL",
+ violations=[
+ {
+ "field": "source_url",
+ "description": "Could not download file from the provided URL",
+ }
+ ],
+ internal_message=(
+ f"Failed to download file from URL: {source_url}, error: {exc}"
+ ),
+ )
+
+
+def assert_temp_file_within_size_limit(
+ *,
+ temp_file_path: str,
+ file_extension: str,
+) -> int:
+ file_size = os.path.getsize(temp_file_path)
+ if file_size <= settings.MAX_FILE_SIZE:
+ return file_size
+
+ limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024)
+ raise ValidationException(
+ user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})",
+ violations=[
+ {
+ "field": "file_size",
+ "description": (
+ f"Size {file_size} bytes exceeds limit of "
+ f"{settings.MAX_FILE_SIZE} bytes"
+ ),
+ }
+ ],
+ )
+
+
+def upload_temp_file_to_source_storage(
+ *,
+ temp_file_path: str,
+ s3_key: str,
+) -> None:
+ storage = JobFileStorage()
+ storage.upload_source_file(temp_file_path, s3_key)
+ logger.info(f"File uploaded to S3: {s3_key}")
+
+
+def verify_source_upload(s3_key: str) -> dict[str, object]:
+ storage = JobFileStorage()
+ file_info = storage.verify_upload_exists(s3_key)
+ if file_info.get("exists"):
+ return dict(file_info)
+
+ raise StorageServiceException(
+ user_message="We failed to verify your file upload",
+ internal_message=f"S3 file verification failed for {s3_key}",
+ )
+
+
+def cleanup_temp_file(temp_file_path: str | None) -> None:
+ if not temp_file_path:
+ return
+ if os.path.exists(temp_file_path):
+ os.remove(temp_file_path)
+ logger.debug(f"Temp file cleaned up: {temp_file_path}")
diff --git a/apps/worker/tests/contract/test_document_parser_architecture_contract.py b/apps/worker/tests/contract/test_document_parser_architecture_contract.py
new file mode 100644
index 000000000..f0d5d381a
--- /dev/null
+++ b/apps/worker/tests/contract/test_document_parser_architecture_contract.py
@@ -0,0 +1,955 @@
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import pandas as pd
+
+
+def test_parse_input_builds_typed_llm_parameters(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.orchestration.parse_input import (
+ ParseInput,
+ ParseOptions,
+ )
+ from app.services.document_parser.orchestration.parse_session import (
+ build_parse_session,
+ )
+
+ monkeypatch.setattr(
+ "app.services.document_parser.orchestration.parse_session.profile_document",
+ lambda *_args, **_kwargs: SimpleNamespace(
+ file_type="pdf",
+ page_count=3,
+ atlas_candidate=False,
+ doc_category="generic",
+ summary=lambda: "profile",
+ reasoning="test",
+ ),
+ )
+
+ parse_input = ParseInput(
+ file_full_path=str(tmp_path / "sample.pdf"),
+ filename="sample.pdf",
+ output_dir=str(tmp_path),
+ internal_output_filename="internal.pdf",
+ job_id="job-1",
+ kb_dir="Default_Root",
+ options=ParseOptions(
+ doc_type="auto",
+ llm_histories=7,
+ smart_title_parse=False,
+ summary_image=False,
+ summary_table=True,
+ summary_txt=False,
+ stopwords=["the"],
+ add_frag_desc="fragment",
+ ),
+ s3_key="uploads/sample.pdf",
+ )
+
+ session = build_parse_session(parse_input)
+
+ assert session.base_llm_paras == {
+ "llm_histories": 7,
+ "smart_title_parse": False,
+ "summary_image": False,
+ "summary_table": True,
+ "summary_txt": False,
+ "stopwords": ["the"],
+ "doc_type": "auto",
+ "frag_desc": "fragment",
+ "model_name": session.base_llm_paras["model_name"],
+ "hierarchy_model_name": session.base_llm_paras["hierarchy_model_name"],
+ }
+ assert session.relative_root == "Default_Root/sample.pdf"
+
+
+def test_document_format_router_uses_adapters(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.orchestration.format_router import (
+ DocumentFormat,
+ get_document_parse_adapter,
+ resolve_document_format,
+ )
+ from app.services.document_parser.orchestration.parse_input import ParseInput
+ from app.services.document_parser.orchestration.parse_session import ParseSession
+ from app.services.document_parser.orchestration.route_parse import route_document_parse
+
+ assert resolve_document_format("/tmp/report.PDF") == DocumentFormat.PDF
+ assert resolve_document_format("/tmp/report.docx") == DocumentFormat.DOCX
+ assert get_document_parse_adapter(DocumentFormat.PDF).document_format == DocumentFormat.PDF
+
+ parsed_df = pd.DataFrame([{"content": "ok"}])
+
+ monkeypatch.setattr(
+ "app.services.document_parser.pdf_parser.parse_pdfs",
+ lambda *_args, **_kwargs: parsed_df,
+ )
+
+ profile = SimpleNamespace(route="standard", doc_category="generic")
+ parse_input = ParseInput(
+ file_full_path=str(tmp_path / "report.pdf"),
+ filename="report.pdf",
+ output_dir=str(tmp_path),
+ internal_output_filename="report.pdf",
+ )
+ session = ParseSession.from_input(
+ parse_input=parse_input,
+ base_llm_paras={},
+ full_output_dir=str(tmp_path),
+ profile=profile,
+ relative_root="Default_Root/report.pdf",
+ )
+
+ output_dir, actual_df = route_document_parse(session)
+
+ assert output_dir == str(tmp_path)
+ assert actual_df is parsed_df
+
+
+def test_rendered_pdf_transform_centralizes_temporary_pdf_cleanup(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.rendered_pdf_transform import (
+ parse_rendered_pdf_bytes,
+ )
+
+ parsed_df = pd.DataFrame([{"content": "pptx"}])
+ seen_pdf_bytes: list[bytes] = []
+
+ def fake_parse_pdfs(pdf_path: str, **_kwargs: Any) -> pd.DataFrame:
+ seen_pdf_bytes.append(Path(pdf_path).read_bytes())
+ assert Path(pdf_path).exists()
+ return parsed_df
+
+ monkeypatch.setattr(
+ "app.services.document_parser.rendered_pdf_transform.parse_pdfs",
+ fake_parse_pdfs,
+ )
+ monkeypatch.setattr(
+ "app.services.document_parser.rendered_pdf_transform.render_pdf_to_image_pdf",
+ lambda pdf_bytes: pdf_bytes,
+ )
+
+ actual_df = parse_rendered_pdf_bytes(
+ pdf_bytes=b"rendered",
+ filename="slides.pptx",
+ output_dir=str(tmp_path),
+ base_llm_paras={},
+ relative_root="Default_Root/slides.pptx",
+ rendered_pdf_s3_key="transforms/job.pdf",
+ )
+
+ assert actual_df is parsed_df
+ assert seen_pdf_bytes == [b"rendered"]
+ assert not (tmp_path / "_pptx_tmp.pdf").exists()
+
+
+def test_heading_hierarchy_module_wraps_prediction(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.heading_hierarchy import (
+ HeadingHierarchyInput,
+ predict_heading_hierarchy,
+ )
+
+ expected_df = pd.DataFrame(
+ [{"id": 1, "heading": "Intro", "level": 1, "reason": "test"}]
+ )
+ captured: dict[str, Any] = {}
+
+ def fake_pred_titles(*args: Any, **kwargs: Any) -> pd.DataFrame:
+ captured["args"] = args
+ captured["kwargs"] = kwargs
+ return expected_df
+
+ monkeypatch.setattr(
+ "app.services.document_parser.heading_hierarchy.pred_titles",
+ fake_pred_titles,
+ )
+
+ actual_df = predict_heading_hierarchy(
+ HeadingHierarchyInput(
+ infos=[(1, "Intro")],
+ doc_type="md",
+ smart_parse=True,
+ model_name="hierarchy-model",
+ output_dir=str(tmp_path),
+ layout_json_path=str(tmp_path / "layout.json"),
+ )
+ )
+
+ assert actual_df is expected_df
+ assert captured["kwargs"]["doc_type"] == "md"
+ assert captured["kwargs"]["smart_parse"] is True
+ assert captured["kwargs"]["model_name"] == "hierarchy-model"
+
+
+def test_parser_row_builder_owns_dataframe_column_order(
+ worker_contract_environment: None,
+) -> None:
+ from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder
+
+ builder = ParsedRowsBuilder()
+ builder.append(
+ ParsedRow(
+ content="chunk text",
+ path="Default_Root/doc/Section",
+ type="text",
+ keywords="alpha;beta",
+ summary="summary",
+ know_id="chunk-1",
+ tokens="alpha->beta",
+ connectto="",
+ page_nums="1,2",
+ addtime="now",
+ )
+ )
+
+ parsed_df = builder.to_dataframe()
+
+ assert list(parsed_df.columns) == [
+ "content",
+ "path",
+ "type",
+ "length",
+ "keywords",
+ "summary",
+ "know_id",
+ "tokens",
+ "connectto",
+ "addtime",
+ "page_nums",
+ ]
+ assert parsed_df.iloc[0].to_dict() == {
+ "content": "chunk text",
+ "path": "Default_Root/doc/Section",
+ "type": "text",
+ "length": len("chunk text"),
+ "keywords": "alpha;beta",
+ "summary": "summary",
+ "know_id": "chunk-1",
+ "tokens": "alpha->beta",
+ "connectto": "",
+ "addtime": "now",
+ "page_nums": "1,2",
+ }
+
+
+def test_inline_asset_module_builds_image_and_table_rows(
+ worker_contract_environment: None,
+) -> None:
+ from app.services.document_parser.inline_asset import (
+ build_image_asset_row,
+ build_table_asset_row,
+ )
+
+ image_row = build_image_asset_row(
+ content="\nImage summary\n[images/image-1.png]\n",
+ relative_path="images/image-1.png",
+ summary="image-1\nImage summary",
+ know_id="image-1",
+ addtime="now",
+ page_nums="3",
+ )
+ table_row = build_table_asset_row(
+ content="",
+ relative_path="tables/table-1.html",
+ summary="table-1\nTable summary",
+ keywords="column",
+ know_id="table-1",
+ addtime="now",
+ page_nums="4",
+ )
+
+ assert image_row.type == "image"
+ assert image_row.path == "images/image-1.png"
+ assert image_row.summary == "image-1\nImage summary"
+ assert table_row.type == "table"
+ assert table_row.path == "tables/table-1.html"
+ assert table_row.keywords == "column"
+
+
+def test_table_asset_writer_creates_table_row_and_html_file(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.table_asset_writer import (
+ TableAssetInput,
+ write_table_asset,
+ )
+
+ row = write_table_asset(
+ TableAssetInput(
+ html="",
+ output_dir=str(tmp_path),
+ table_name="table-1",
+ summary="table-1",
+ keywords="A",
+ know_id="table-1",
+ addtime="now",
+ )
+ )
+
+ assert (tmp_path / "tables" / "table-1.html").read_text(encoding="utf-8")
+ assert row.type == "table"
+ assert row.path == "tables/table-1.html"
+ assert row.content == ""
+
+
+def test_docx_asset_store_owns_asset_filesystem_lifecycle(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.docx_asset_store import DocxAssetStore
+
+ store = DocxAssetStore(str(tmp_path))
+ (tmp_path / "images").mkdir()
+ (tmp_path / "tables").mkdir()
+ (tmp_path / "images" / "stale.png").write_bytes(b"stale")
+ (tmp_path / "tables" / "stale.html").write_text("stale", encoding="utf-8")
+
+ store.reset()
+ image_asset = store.write_image("image-1 raw", ".png", b"image")
+ renamed_asset = store.rename_image(image_asset, "image-1 final")
+ table_asset = store.write_table("table-1 final", "")
+
+ assert not (tmp_path / "images" / "stale.png").exists()
+ assert not (tmp_path / "tables" / "stale.html").exists()
+ assert renamed_asset.relative_path == "images/image-1 final.png"
+ assert (tmp_path / "images" / "image-1 final.png").read_bytes() == b"image"
+ assert table_asset.relative_path == "tables/table-1 final.html"
+ assert (tmp_path / "tables" / "table-1 final.html").read_text(
+ encoding="utf-8"
+ ) == ""
+
+
+def test_docx_block_stream_emits_document_ordered_blocks(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.docx_block_stream import iter_block_items
+ from docx import Document
+ from docx.table import Table
+ from docx.text.paragraph import Paragraph
+
+ docx_path = tmp_path / "sample.docx"
+ document = Document()
+ document.add_paragraph("Intro")
+ table = document.add_table(rows=1, cols=1)
+ table.cell(0, 0).text = "Cell"
+ document.save(docx_path)
+
+ block_events = list(iter_block_items(docx_path.read_bytes()))
+
+ assert block_events[0][0] == 1
+ assert isinstance(block_events[0][1], Paragraph)
+ assert block_events[0][1].text == "Intro"
+ assert block_events[0][2] == "PTXT"
+ assert isinstance(block_events[1][1], Table)
+ assert block_events[1][2] == "TABLE"
+
+
+def test_html_table_modules_separate_docx_and_dataframe_rendering(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.dataframe_html_renderer import df2html
+ from app.services.document_parser.docx_table_html import table2html
+ from docx import Document
+
+ docx_path = tmp_path / "table.docx"
+ document = Document()
+ table = document.add_table(rows=1, cols=2)
+ table.cell(0, 0).text = "A"
+ table.cell(0, 1).text = "B"
+ document.save(docx_path)
+
+ loaded_table = Document(str(docx_path)).tables[0]
+ docx_html = table2html(loaded_table, cell_image_map={(0, 1): "image summary"})
+
+ dataframe_html = df2html(
+ pd.DataFrame([["North", "North", 3]], columns=["Region", "Group", "Value"]),
+ row_header_cols=2,
+ )
+
+ assert docx_html == (
+ ""
+ )
+ assert 'North | ' in dataframe_html
+ assert "3 | " in dataframe_html
+
+
+def test_table_text_parser_owns_markdown_table_text_contract(
+ worker_contract_environment: None,
+) -> None:
+ from app.services.document_parser.table_text_parser import (
+ df2md,
+ extract_tables_by_forms,
+ identify_tables,
+ sanitize_table_name_from_header,
+ )
+
+ markdown_table = "\n".join(
+ [
+ "| Product | Revenue |",
+ "| --- | ---: |",
+ "| Notebook | 42 |",
+ ]
+ )
+
+ is_table, table_form, _tables = identify_tables("| Product | Revenue |")
+ table_html = extract_tables_by_forms(markdown_table, form="md")
+ markdown_output = df2md(pd.DataFrame([{"City": "北京", "Value": 7}]))
+
+ assert is_table is True
+ assert table_form == "md"
+ assert table_html is not None
+ assert "Product | " in table_html
+ assert "Notebook | " in table_html
+ assert "---" not in table_html
+ assert sanitize_table_name_from_header("A | Revenue | Revenue | 市场") == (
+ "Revenue 市场"
+ )
+ assert "| City | Value |" in markdown_output
+ assert "| 北京 | 7 |" in markdown_output
+
+
+def test_table_frame_parser_owns_dataframe_table_contract(
+ worker_contract_environment: None,
+) -> None:
+ from app.services.document_parser.table_frame_parser import (
+ parse_tb_contents,
+ parse_tb_keywords,
+ postprocess_tb,
+ )
+
+ raw_frame = pd.DataFrame(
+ [["North\nAmerica", 42, None]],
+ columns=["Region\nName", "Revenue", None],
+ )
+
+ normalized_frame = postprocess_tb(raw_frame, drop=True)
+ paths, table_html = parse_tb_contents(
+ normalized_frame,
+ parent_dic={"budget.xlsx": {"Visible": {}}},
+ file_name="budget.xlsx",
+ sheet_name="Visible",
+ )
+ keywords = parse_tb_keywords(normalized_frame)
+
+ assert normalized_frame.columns.tolist() == ["RegionName", "Revenue"]
+ assert "NorthAmerica" in table_html
+ assert "RegionName" in keywords
+ assert "Revenue" in keywords
+ assert "budget.xlsx/Visible/RegionName" in paths
+ assert "budget.xlsx/Visible/Revenue" in paths
+
+
+def test_markdown_table_asset_module_owns_table_asset_contract(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.markdown_table_asset import (
+ MarkdownTableAssetRequest,
+ build_markdown_table_asset,
+ )
+ from app.services.document_parser.markdown_deferred_task import (
+ TableDeferredSummaryTask,
+ )
+
+ table_dir = tmp_path / "tables"
+ table_dir.mkdir()
+ table_html = (
+ "| Product | Revenue |
|---|
"
+ "| Notebook | 42 |
"
+ )
+
+ asset = build_markdown_table_asset(
+ MarkdownTableAssetRequest(
+ table_html=table_html,
+ table_dir=str(table_dir),
+ table_count=3,
+ timestamp="now",
+ current_page_number=9,
+ summary_table=True,
+ row_index=7,
+ )
+ )
+
+ assert asset.content_item == f"\n[{asset.relative_path}]\n"
+ assert asset.row_values[1] == asset.relative_path
+ assert asset.row_values[2] == "table"
+ assert asset.row_values[5] == "table-3"
+ assert asset.row_values[10] == "9"
+ assert asset.deferred_task == TableDeferredSummaryTask(
+ row_index=7,
+ table_html=table_html,
+ table_dir=str(table_dir),
+ table_name=Path(asset.relative_path).stem,
+ table_count=2,
+ )
+ assert "border='1'" in (tmp_path / asset.relative_path).read_text(
+ encoding="utf-8"
+ )
+
+
+def test_markdown_image_asset_module_owns_image_materialization_contract(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.markdown_image_asset import (
+ MarkdownImageAssetRequest,
+ build_markdown_image_name,
+ build_markdown_image_asset,
+ )
+ from app.services.document_parser.markdown_deferred_task import (
+ ImageDeferredSummaryTask,
+ )
+
+ image_dir = tmp_path / "images"
+ image_dir.mkdir()
+ source_image = tmp_path / "raw.png"
+ source_image.write_bytes(b"same pixels")
+ seen_images: dict[str, dict[str, str]] = {}
+
+ image_asset = build_markdown_image_asset(
+ MarkdownImageAssetRequest(
+ output_dir=str(tmp_path),
+ image_dir=str(image_dir),
+ image_path=str(source_image),
+ image_name=build_markdown_image_name(
+ image_count=2,
+ last_context="Revenue Chart",
+ ),
+ image_count=2,
+ last_context="Revenue Chart",
+ image_summary="Sales by region",
+ timestamp="now",
+ current_page_number=8,
+ seen_images=seen_images,
+ summary_image=True,
+ row_index=4,
+ )
+ )
+
+ assert image_asset.content_item is not None
+ assert image_asset.row_values is not None
+ assert image_asset.cache_key is not None
+ assert image_asset.cache_entry is not None
+ assert image_asset.should_advance_image_count is True
+ assert image_asset.row_values[1] == "images/image-2-Revenue Ch.png"
+ assert image_asset.row_values[2] == "image"
+ assert image_asset.row_values[5] == "image-2\nSales by region"
+ assert image_asset.row_values[10] == "8"
+ assert image_asset.deferred_task == ImageDeferredSummaryTask(
+ row_index=4,
+ relative_path="images/image-2-Revenue Ch.png",
+ image_dir=str(image_dir),
+ image_name="image-2-Revenue Ch",
+ image_suffix=".png",
+ )
+ assert (tmp_path / "images" / "image-2-Revenue Ch.png").read_bytes() == (
+ b"same pixels"
+ )
+ assert not source_image.exists()
+
+ seen_images[image_asset.cache_key] = image_asset.cache_entry
+ duplicate_source = tmp_path / "duplicate.png"
+ duplicate_source.write_bytes(b"same pixels")
+
+ duplicate_asset = build_markdown_image_asset(
+ MarkdownImageAssetRequest(
+ output_dir=str(tmp_path),
+ image_dir=str(image_dir),
+ image_path=str(duplicate_source),
+ image_name=build_markdown_image_name(
+ image_count=3,
+ last_context="Other Chart",
+ ),
+ image_count=3,
+ last_context="Other Chart",
+ image_summary=None,
+ timestamp="now",
+ current_page_number=9,
+ seen_images=seen_images,
+ summary_image=True,
+ row_index=5,
+ )
+ )
+
+ assert duplicate_asset.content_item == image_asset.content_item
+ assert duplicate_asset.row_values is not None
+ assert duplicate_asset.row_values[1] == "images/image-2-Revenue Ch.png"
+ assert duplicate_asset.deferred_task is None
+ assert duplicate_asset.should_advance_image_count is False
+ assert not duplicate_source.exists()
+
+
+def test_mineru_modules_separate_client_and_task_polling(
+ worker_contract_environment: None,
+) -> None:
+ from app.services.document_parser.mineru_client import get_mineru_headers
+ from app.services.document_parser.mineru_task_polling import (
+ get_batch_status,
+ get_polling_interval_for_state,
+ )
+
+ assert get_mineru_headers("token") == {
+ "Content-Type": "application/json",
+ "Authorization": "Bearer token",
+ }
+ assert get_batch_status({"data": {"extract_result": [{"state": "done"}]}}) == {
+ "state": "done"
+ }
+ assert get_batch_status({"data": {"extract_result": {"state": "failed"}}}) == {
+ "state": "failed"
+ }
+ assert get_polling_interval_for_state("pending", 2) == 8.0
+ assert get_polling_interval_for_state("running", 2) == 10.0
+ assert get_polling_interval_for_state("waiting-file", 2) == 15.0
+
+
+def test_doc_profile_model_owns_profile_contract_and_metadata(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ import json
+
+ from app.services.document_parser.doc_profile_model import (
+ DocProfile,
+ save_profile_metadata,
+ )
+
+ profile = DocProfile(
+ file_type="pdf",
+ route="fast",
+ decision_band="safe_fast",
+ page_count=3,
+ avg_text_density=123.4,
+ avg_image_coverage=0.05,
+ page_details=[{"page": 1}],
+ sample_text="hidden",
+ )
+
+ save_profile_metadata(profile, str(tmp_path))
+ saved_profile = json.loads((tmp_path / "profile.json").read_text(encoding="utf-8"))
+
+ assert "page_details" not in saved_profile
+ assert "sample_text" not in saved_profile
+ assert saved_profile["file_type"] == "pdf"
+ assert "route=fast" in profile.summary()
+
+
+def test_doc_profiler_dispatches_pdf_to_pdf_profile_module(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+) -> None:
+ from app.services.document_parser.doc_profile_model import DocProfile
+ from app.services.document_parser.doc_profiler import profile_document
+
+ called_paths: list[str] = []
+
+ def fake_profile_pdf(path: str) -> DocProfile:
+ called_paths.append(path)
+ return DocProfile(file_type="pdf", page_count=2)
+
+ monkeypatch.setattr(
+ "app.services.document_parser.doc_profiler.profile_pdf",
+ fake_profile_pdf,
+ )
+
+ pdf_profile = profile_document("/tmp/input.bin", filename="report.pdf")
+ docx_profile = profile_document("/tmp/input.bin", filename="report.docx")
+
+ assert called_paths == ["/tmp/input.bin"]
+ assert pdf_profile.file_type == "pdf"
+ assert docx_profile.file_type == "docx"
+ assert docx_profile.route == "standard"
+
+
+def test_excel_structure_parser_is_table_structure_seam(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ import openpyxl
+
+ from app.services.document_parser.excel_structure_parser import (
+ parse_excel_structure,
+ )
+
+ workbook = openpyxl.Workbook()
+ worksheet = workbook.active
+ worksheet.title = "Budget"
+ worksheet["A1"] = "Region"
+ worksheet["B1"] = "Value"
+ worksheet["A2"] = "North"
+ worksheet["B2"] = 10
+ workbook_path = tmp_path / "budget.xlsx"
+ workbook.save(workbook_path)
+
+ parsed_sheets = parse_excel_structure(str(workbook_path), split_subtables=False)
+
+ assert list(parsed_sheets.keys()) == ["Budget"]
+ assert parsed_sheets["Budget"].attrs["row_header_cols"] >= 0
+ assert "North" in parsed_sheets["Budget"].astype(str).to_string()
+
+
+def test_heading_hierarchy_exposes_candidate_and_tree_modules(
+ worker_contract_environment: None,
+) -> None:
+ from app.services.document_parser.heading_candidates import filter_markdown_headings
+ from app.services.document_parser.heading_tree import cleanup_heading_tree
+
+ candidates = filter_markdown_headings(["# Intro", "body", "## Detail"])
+ cleaned = cleanup_heading_tree(
+ pd.DataFrame(
+ [
+ {"id": 0, "heading": "Intro", "level": 1, "reason": ""},
+ {"id": 2, "heading": "Detail", "level": 2, "reason": ""},
+ ]
+ )
+ )
+
+ assert candidates[["id", "heading", "level"]].to_dict("records")[0] == {
+ "id": 0,
+ "heading": "Intro",
+ "level": 1,
+ }
+ assert cleaned["heading"].tolist() == ["Intro", "Detail"]
+
+
+def test_heading_llm_executor_owns_prompt_execution_and_fallback(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+) -> None:
+ from app.services.document_parser.heading_llm_executor import (
+ execute_llm_heading_hierarchy,
+ )
+
+ monkeypatch.setenv("KB_LAYOUT_LLM_COMPACT_INPUT", "true")
+ raw_preds = pd.DataFrame(
+ [
+ {"id": 0, "heading": "body", "level": -1, "reason": ""},
+ {"id": 1, "heading": "Intro", "level": -2, "reason": "POS [1] NEG [0]"},
+ {"id": 2, "heading": "body", "level": -1, "reason": ""},
+ ]
+ )
+ judged_prompts: list[pd.DataFrame] = []
+ saved_files: list[str] = []
+
+ def fake_hierarchy_judge(
+ df: pd.DataFrame,
+ *_args: Any,
+ **_kwargs: Any,
+ ) -> list[dict[str, int]]:
+ judged_prompts.append(df.copy())
+ return [{"id": 1, "level": 1}]
+
+ def unexpected_fallback(_df: pd.DataFrame) -> pd.DataFrame:
+ raise AssertionError("fallback should not run")
+
+ actual_df = execute_llm_heading_hierarchy(
+ raw_preds=raw_preds,
+ prompt_limt=4000,
+ hierarchy_judge=fake_hierarchy_judge,
+ fallback_hierarchy=unexpected_fallback,
+ save_intermediate_csv=lambda _df, _output_dir, filename: saved_files.append(
+ filename
+ ),
+ model_name="hierarchy-model",
+ )
+
+ assert actual_df["level"].tolist() == [-1, 1, -1]
+ assert "Intro" in judged_prompts[0]["heading"].tolist()
+ assert judged_prompts[0]["heading"].tolist().count("[1 BODY LINES]") == 2
+ assert saved_files == ["preds_3_llm_base", "preds_4_llm_final"]
+
+ body_only = pd.DataFrame(
+ [{"id": 0, "heading": "body", "level": -1, "reason": ""}]
+ )
+ skipped_df = execute_llm_heading_hierarchy(
+ raw_preds=body_only,
+ prompt_limt=4000,
+ hierarchy_judge=lambda *_args, **_kwargs: (_ for _ in ()).throw(
+ AssertionError("LLM should not run without heading candidates")
+ ),
+ fallback_hierarchy=unexpected_fallback,
+ save_intermediate_csv=lambda *_args: None,
+ )
+
+ assert skipped_df["level"].tolist() == [-1]
+
+
+def test_markdown_deferred_summary_module_updates_rows_and_refs(
+ worker_contract_environment: None,
+ monkeypatch: Any,
+ tmp_path: Path,
+) -> None:
+ import app.services.document_parser.markdown_deferred_summary as deferred_summary
+ from app.services.document_parser.markdown_deferred_summary import (
+ MarkdownDeferredSummaryInput,
+ apply_markdown_deferred_summaries,
+ )
+ from app.services.document_parser.markdown_deferred_task import (
+ ImageDeferredSummaryTask,
+ TableDeferredSummaryTask,
+ TextDeferredSummaryTask,
+ )
+
+ image_dir = tmp_path / "images"
+ table_dir = tmp_path / "tables"
+ image_dir.mkdir()
+ table_dir.mkdir()
+ (image_dir / "image-3-old.png").write_bytes(b"image")
+ (table_dir / "table-0 old.html").write_text("", encoding="utf-8")
+
+ rows: list[list[str | int]] = [
+ [
+ "[images/image-3-old.png]",
+ "images/image-3-old.png",
+ "image",
+ 24,
+ "",
+ "image-3",
+ "image-id",
+ "",
+ "",
+ "now",
+ "",
+ ],
+ [
+ "[tables/table-0 old.html]",
+ "tables/table-0 old.html",
+ "table",
+ 25,
+ "",
+ "table-0",
+ "table-id",
+ "",
+ "",
+ "now",
+ "",
+ ],
+ [
+ "long text",
+ "Root/Text",
+ "text",
+ 9,
+ "",
+ "",
+ "text-id",
+ "",
+ "",
+ "now",
+ "",
+ ],
+ ]
+
+ monkeypatch.setattr(deferred_summary, "_get_vision_client", lambda: object())
+ monkeypatch.setattr(
+ deferred_summary,
+ "ask_image",
+ lambda *_args, **_kwargs: "Better Image\nImage summary",
+ )
+
+ def fake_extract(text: str, **_kwargs: Any) -> tuple[str, str, str]:
+ if " None:
+ from app.services.document_parser.toc_docx import infer_toc_level_from_text
+ from app.services.document_parser.toc_hierarchy import build_toc_hierarchy_payload
+
+ assert infer_toc_level_from_text("1.2 Scope") == 2
+ payload = build_toc_hierarchy_payload(
+ [
+ {"id": 3, "heading": "1 Overview", "level": 1},
+ {"id": 4, "heading": "1.1 Detail", "level": 2},
+ ],
+ toc_range=(3, 4),
+ scan_range=(3, 5),
+ )
+
+ assert payload is not None
+ assert payload["toc_range"] == (3, 4)
+ assert payload["scan_range"] == (3, 5)
+ assert payload["toc_tree"] == {"1 Overview": {"1.1 Detail": {}}}
+
+
+def test_format_adapters_do_not_expose_lazy_any_wrappers(
+ worker_contract_environment: None,
+) -> None:
+ import app.services.document_parser.orchestration.format_adapters as format_adapters
+
+ wrapper_names = [
+ "parse_fragment",
+ "parse_texts",
+ "parse_md",
+ "parse_image",
+ "parse_pdfs",
+ "parse_docx",
+ "convert_doc2dics",
+ "doc_to_docx",
+ "xls_to_xlsx",
+ "parse_xlsx",
+ "parse_pptx",
+ ]
+
+ assert not any(hasattr(format_adapters, wrapper_name) for wrapper_name in wrapper_names)
diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py
new file mode 100644
index 000000000..0f86e6e96
--- /dev/null
+++ b/apps/worker/tests/contract/test_excel_parser_contract.py
@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from pytest import MonkeyPatch
+
+
+def _write_contract_workbook(workbook_path: Path) -> None:
+ import openpyxl
+
+ workbook = openpyxl.Workbook()
+ visible_sheet = workbook.active
+ visible_sheet.title = "Visible"
+ visible_sheet["A1"] = "Region"
+ visible_sheet["B1"] = "Value"
+ visible_sheet["A2"] = "North"
+ visible_sheet["B2"] = 10
+
+ hidden_sheet = workbook.create_sheet("Hidden")
+ hidden_sheet.sheet_state = "hidden"
+ hidden_sheet["A1"] = "Secret"
+ hidden_sheet["B1"] = "Value"
+ hidden_sheet["A2"] = "Hidden"
+ hidden_sheet["B2"] = 99
+
+ workbook.save(workbook_path)
+
+
+def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.parse_service import checkerboard_inject_parse
+
+ workbook_path = tmp_path / "budget.xlsx"
+ _write_contract_workbook(workbook_path)
+
+ full_output_dir, parsed_df = checkerboard_inject_parse(
+ file_full_path=str(workbook_path),
+ filename="budget.xlsx",
+ output_dir=str(tmp_path),
+ internal_output_filename="budget.xlsx",
+ summary_image=False,
+ summary_table=False,
+ summary_txt=False,
+ smart_title_parse=False,
+ stopwords=[],
+ )
+
+ assert full_output_dir.endswith("Default_Root/budget.xlsx")
+ assert parsed_df is not None
+ assert parsed_df["type"].tolist() == ["table"]
+ assert parsed_df["path"].tolist() == ["tables/table-Visible.html"]
+ assert parsed_df["summary"].tolist() == ["table-Visible"]
+ assert "Region" in parsed_df["keywords"].iloc[0]
+ assert "Value" in parsed_df["keywords"].iloc[0]
+
+ table_html = Path(full_output_dir) / "tables" / "table-Visible.html"
+ table_html_text = table_html.read_text(encoding="utf-8")
+
+ assert "North" in table_html_text
+ assert "10" in table_html_text
+ assert "Secret" not in table_html_text
+ assert "Hidden" not in table_html_text
+
+
+def test_xlsx_parser_contract_accepts_missing_llm_parameters(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from app.services.document_parser.excel_table_parser import parse_xlsx
+
+ workbook_path = tmp_path / "default-parameters.xlsx"
+ output_dir = tmp_path / "output"
+ _write_contract_workbook(workbook_path)
+
+ parsed_df = parse_xlsx(
+ file_path=str(workbook_path),
+ file_name="default-parameters.xlsx",
+ output_dir=str(output_dir),
+ baseurl="",
+ base_llm_paras=None,
+ )
+
+ assert parsed_df["type"].tolist() == ["table"]
+ assert parsed_df["path"].tolist() == ["tables/table-Visible.html"]
+ assert (output_dir / "tables" / "table-Visible.html").exists()
+
+
+def test_xlsx_parser_contract_falls_back_to_column_keywords_when_llm_summary_is_empty(
+ worker_contract_environment: None,
+ monkeypatch: MonkeyPatch,
+ tmp_path: Path,
+) -> None:
+ import app.services.document_parser.txt_parser as txt_parser
+ from app.services.document_parser.excel_table_parser import parse_xlsx
+
+ workbook_path = tmp_path / "empty-summary.xlsx"
+ output_dir = tmp_path / "output"
+ _write_contract_workbook(workbook_path)
+
+ monkeypatch.setattr(
+ txt_parser,
+ "extract_title_keywords_summary",
+ lambda *_args, **_kwargs: (None, "", ""),
+ )
+
+ parsed_df = parse_xlsx(
+ file_path=str(workbook_path),
+ file_name="empty-summary.xlsx",
+ output_dir=str(output_dir),
+ baseurl="",
+ base_llm_paras={"summary_table": True, "stopwords": []},
+ )
+
+ keywords = str(parsed_df["keywords"].iloc[0])
+ content = str(parsed_df["content"].iloc[0])
+
+ assert "Region" in keywords
+ assert "Value" in keywords
+ assert "Main columns:" in content
+ assert "Region" in content
+ assert "Value" in content
diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py
index 12f672267..b4c3aae49 100644
--- a/apps/worker/tests/contract/test_parse_task_contract.py
+++ b/apps/worker/tests/contract/test_parse_task_contract.py
@@ -35,8 +35,8 @@ def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]:
def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]:
import app.core.tasks.kb_tasks as kb_tasks
+ import app.services.document_ingestion.processing_run as parse_job_service
import app.services.document_parser.parse_service as parse_service
- import app.services.storage.sync_storage_service as sync_storage_service
from shared.core.database_sync import get_sync_engine
from shared.services.redis.redis_sync_service import (
SyncJobInfoRedisService,
@@ -47,7 +47,7 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]:
return (
kb_tasks,
parse_service,
- sync_storage_service,
+ parse_job_service,
get_sync_engine(),
SyncJobInfoRedisService,
SyncJobMetadataService,
@@ -55,6 +55,12 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]:
)
+def _load_worker_settings() -> Any:
+ from shared.core.config import settings
+
+ return settings
+
+
def _save_worker_task_cache(
*,
job_id: str,
@@ -84,6 +90,19 @@ def _save_worker_task_cache(
return redis_service
+def _patch_verify_upload_exists(
+ monkeypatch: MonkeyPatch,
+ file_info_for_storage_key: Any,
+) -> None:
+ from shared.services.storage.job_file_storage import JobFileStorage
+
+ monkeypatch.setattr(
+ JobFileStorage,
+ "verify_upload_exists",
+ lambda self, storage_key: file_info_for_storage_key(storage_key),
+ )
+
+
def _find_task_workspaces(root: Path, job_id: str) -> list[Path]:
return sorted(
path
@@ -124,12 +143,13 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state(
(
kb_tasks,
parse_service,
- sync_storage_service,
+ parse_job_service,
engine,
sync_job_info_service_cls,
sync_job_metadata_service_cls,
sync_redis_service_factory,
) = _load_parse_task_modules()
+ settings = _load_worker_settings()
user_id: str = f"worker-user-{uuid4().hex[:12]}"
job_id: str = f"job_parse_success_{uuid4().hex[:12]}"
@@ -166,8 +186,8 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state(
)
_bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks)
- monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path))
- monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", billing_enabled)
+ monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path))
+ monkeypatch.setattr(settings, "BILLING_ENABLED", billing_enabled)
def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
return {
@@ -175,25 +195,12 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
"size": _SAMPLE_PDF_PATH.stat().st_size,
}
- def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]:
- return {"download_url": f"https://example.test/{storage_key}"}
-
- monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists)
- monkeypatch.setattr(
- sync_storage_service,
- "verify_s3_file_exists",
- fake_verify_s3_file_exists,
- )
- monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url)
- monkeypatch.setattr(
- sync_storage_service,
- "generate_download_url",
- fake_generate_download_url,
- )
+ _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists)
def fake_download_s3_file_to_temp(
- file_url: str, file_ext: str, temp_dir: str
+ storage_key: str, file_ext: str, temp_dir: str
) -> str:
+ assert storage_key == s3_key
assert file_ext == ".pdf"
downloaded_path = Path(temp_dir) / f"downloaded{file_ext}"
shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path)
@@ -309,9 +316,9 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
raw_files={},
)
- monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
+ monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse)
- monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage())
+ monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage())
result = kb_tasks.parse_task.run(job_id, user_id, "kb_management")
@@ -338,8 +345,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
},
},
]
- expected_credits_charged = 3 * int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE)
- expected_initial_balance = int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
+ expected_credits_charged = 3 * int(settings.MICRO_DOLLARS_PER_PAGE)
+ expected_initial_balance = int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
assert result == {
"status": "success",
@@ -515,7 +522,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
SELECT transition_reason, to_state
FROM job_state_audit_logs
WHERE job_id = :job_id
- ORDER BY created_at ASC
+ ORDER BY id ASC
"""
),
{"job_id": job_id},
@@ -577,12 +584,13 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks
(
kb_tasks,
parse_service,
- sync_storage_service,
+ parse_job_service,
engine,
sync_job_info_service_cls,
sync_job_metadata_service_cls,
sync_redis_service_factory,
) = _load_parse_task_modules()
+ settings = _load_worker_settings()
user_id: str = f"worker-user-{uuid4().hex[:12]}"
existing_job_id: str = f"job_existing_{uuid4().hex[:12]}"
@@ -765,8 +773,8 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks
)
_bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks)
- monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path))
- monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", False)
+ monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path))
+ monkeypatch.setattr(settings, "BILLING_ENABLED", False)
def fake_cleanup_task_workspace(workspace_dir: str | None) -> bool:
captured_artifacts["workspace_dir"] = workspace_dir
@@ -778,12 +786,10 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
"size": _SAMPLE_PDF_PATH.stat().st_size,
}
- def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]:
- return {"download_url": f"https://example.test/{storage_key}"}
-
def fake_download_s3_file_to_temp(
- file_url: str, file_ext: str, temp_dir: str
+ storage_key: str, file_ext: str, temp_dir: str
) -> str:
+ assert storage_key == s3_key
downloaded_path = Path(temp_dir) / f"downloaded{file_ext}"
shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path)
return str(downloaded_path)
@@ -867,22 +873,11 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
raw_files={},
)
- monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists)
- monkeypatch.setattr(
- sync_storage_service,
- "verify_s3_file_exists",
- fake_verify_s3_file_exists,
- )
- monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url)
- monkeypatch.setattr(
- sync_storage_service,
- "generate_download_url",
- fake_generate_download_url,
- )
- monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
+ _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists)
+ monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse)
- monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage())
- monkeypatch.setattr(kb_tasks, "cleanup_task_workspace", fake_cleanup_task_workspace)
+ monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage())
+ monkeypatch.setattr(parse_job_service, "cleanup_task_workspace", fake_cleanup_task_workspace)
result = kb_tasks.parse_task.run(job_id, user_id, "kb_management")
@@ -965,12 +960,13 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks(
(
kb_tasks,
parse_service,
- sync_storage_service,
+ parse_job_service,
engine,
sync_job_info_service_cls,
sync_job_metadata_service_cls,
sync_redis_service_factory,
) = _load_parse_task_modules()
+ settings = _load_worker_settings()
user_id: str = f"worker-concurrent-user-{uuid4().hex[:12]}"
job_ids: list[str] = [f"job_cb_{index}_{uuid4().hex[:12]}" for index in range(2)]
@@ -1011,8 +1007,8 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks(
)
_bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks)
- monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path))
- monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", True)
+ monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path))
+ monkeypatch.setattr(settings, "BILLING_ENABLED", True)
def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
return {
@@ -1020,12 +1016,10 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
"size": _SAMPLE_PDF_PATH.stat().st_size,
}
- def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]:
- return {"download_url": f"https://example.test/{storage_key}"}
-
def fake_download_s3_file_to_temp(
- file_url: str, file_ext: str, temp_dir: str
+ storage_key: str, file_ext: str, temp_dir: str
) -> str:
+ assert storage_key in s3_keys.values()
assert file_ext == ".pdf"
downloaded_path = Path(temp_dir) / f"downloaded{file_ext}"
shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path)
@@ -1072,22 +1066,11 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
raw_files={},
)
- monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists)
- monkeypatch.setattr(
- sync_storage_service,
- "verify_s3_file_exists",
- fake_verify_s3_file_exists,
- )
- monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url)
- monkeypatch.setattr(
- sync_storage_service,
- "generate_download_url",
- fake_generate_download_url,
- )
- monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
- monkeypatch.setattr(kb_tasks.PageEstimator, "estimate", fake_estimate_page_count)
+ _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists)
+ monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
+ monkeypatch.setattr(parse_job_service.PageEstimator, "estimate", fake_estimate_page_count)
monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse)
- monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage())
+ monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage())
def run_parse_task(job_id: str) -> dict[str, Any]:
return dict(kb_tasks.parse_task.run(job_id, user_id, "kb_management"))
@@ -1095,9 +1078,9 @@ def run_parse_task(job_id: str) -> dict[str, Any]:
with ThreadPoolExecutor(max_workers=len(job_ids)) as executor:
results = list(executor.map(run_parse_task, job_ids))
- expected_credits_charged = int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE)
+ expected_credits_charged = int(settings.MICRO_DOLLARS_PER_PAGE)
expected_initial_balance = (
- int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
+ int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
)
with engine.begin() as connection:
@@ -1215,12 +1198,13 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal(
(
kb_tasks,
parse_service,
- sync_storage_service,
+ parse_job_service,
engine,
sync_job_info_service_cls,
sync_job_metadata_service_cls,
sync_redis_service_factory,
) = _load_parse_task_modules()
+ settings = _load_worker_settings()
user_id: str = f"worker-user-{uuid4().hex[:12]}"
job_id: str = f"job_parse_skipped_{uuid4().hex[:12]}"
@@ -1253,23 +1237,11 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal(
)
_bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks)
- monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path))
+ monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path))
def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
return {"exists": storage_key == s3_key, "size": 1024}
- monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists)
- monkeypatch.setattr(
- sync_storage_service,
- "verify_s3_file_exists",
- fake_verify_s3_file_exists,
- )
- monkeypatch.setattr(
- kb_tasks,
- "generate_download_url",
- lambda *_args, **_kwargs: (_ for _ in ()).throw(
- AssertionError("terminal parse task should not request a download URL")
- ),
- )
+ _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists)
monkeypatch.setattr(
parse_service,
"checkerboard_inject_parse",
@@ -1316,12 +1288,13 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi
(
kb_tasks,
parse_service,
- sync_storage_service,
+ parse_job_service,
engine,
sync_job_info_service_cls,
sync_job_metadata_service_cls,
sync_redis_service_factory,
) = _load_parse_task_modules()
+ settings = _load_worker_settings()
user_id: str = f"worker-user-{uuid4().hex[:12]}"
job_id: str = f"job_parse_failure_{uuid4().hex[:12]}"
@@ -1354,44 +1327,31 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi
)
_bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks)
- monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path))
+ monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path))
def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
return {
"exists": storage_key == s3_key,
"size": _SAMPLE_PDF_PATH.stat().st_size,
}
- def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]:
- return {"download_url": f"https://example.test/{storage_key}"}
-
- monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists)
- monkeypatch.setattr(
- sync_storage_service,
- "verify_s3_file_exists",
- fake_verify_s3_file_exists,
- )
- monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url)
- monkeypatch.setattr(
- sync_storage_service,
- "generate_download_url",
- fake_generate_download_url,
- )
+ _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists)
def fake_download_s3_file_to_temp(
- file_url: str, file_ext: str, temp_dir: str
+ storage_key: str, file_ext: str, temp_dir: str
) -> str:
+ assert storage_key == s3_key
downloaded_path = Path(temp_dir) / f"downloaded{file_ext}"
shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path)
return str(downloaded_path)
- monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
+ monkeypatch.setattr(parse_job_service, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
monkeypatch.setattr(
parse_service,
"checkerboard_inject_parse",
lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("parse failed")),
)
monkeypatch.setattr(
- kb_tasks,
+ parse_job_service,
"get_result_storage",
lambda: (_ for _ in ()).throw(
AssertionError("result storage should not run after parser failure")
@@ -1406,8 +1366,8 @@ def fake_download_s3_file_to_temp(
assert result.status == "FAILURE"
assert _find_task_workspaces(tmp_path, job_id) == []
- expected_credits_charged = 3 * int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE)
- expected_initial_balance = int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
+ expected_credits_charged = 3 * int(settings.MICRO_DOLLARS_PER_PAGE)
+ expected_initial_balance = int(settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
with engine.begin() as connection:
job_row = (
@@ -1460,7 +1420,7 @@ def fake_download_s3_file_to_temp(
SELECT transition_reason, to_state
FROM job_state_audit_logs
WHERE job_id = :job_id
- ORDER BY created_at ASC
+ ORDER BY id ASC
"""
),
{"job_id": job_id},
diff --git a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py
index 48cd82705..c5987ad8d 100644
--- a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py
+++ b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py
@@ -138,3 +138,89 @@ def test_should_skip_duplicate_beat_firing_with_the_real_periodic_redis_lock(
"status": "skipped",
"reason": "duplicate Beat firing",
}
+
+
+def test_should_record_retry_transition_through_sync_state_machine(
+ worker_contract_environment: None,
+) -> None:
+ from shared.core.state_machine.service_sync import SyncStateMachineService
+ from shared.core.database_sync import get_sync_db_context
+ from shared.services.redis.redis_sync_service import SyncRedisServiceFactory
+
+ job_id = f"job_retry_{uuid4().hex[:12]}"
+ user_id = f"worker-user-{uuid4().hex[:12]}"
+ _, engine = _load_worker_modules()
+
+ with engine.begin() as connection:
+ insert_contract_user(connection, user_id=user_id)
+ insert_contract_job(
+ connection,
+ job_id=job_id,
+ user_id=user_id,
+ status="failed",
+ source_type="file",
+ webhook_enabled=False,
+ job_metadata=_build_file_job_metadata(),
+ error_code="TRANSIENT",
+ error_message="temporary failure",
+ )
+
+ redis_service = SyncRedisServiceFactory.get_service()
+ state_machine = SyncStateMachineService(redis_service=redis_service)
+
+ with get_sync_db_context() as db:
+ did_retry = state_machine.handle_retry(
+ db,
+ job_id,
+ retry_metadata={"worker": "contract"},
+ )
+
+ assert did_retry is True
+
+ with engine.begin() as connection:
+ job_row = (
+ connection.execute(
+ text(
+ """
+ SELECT status
+ FROM jobs
+ WHERE job_id = :job_id
+ """
+ ),
+ {"job_id": job_id},
+ )
+ .mappings()
+ .one()
+ )
+ audit_log_row = (
+ connection.execute(
+ text(
+ """
+ SELECT from_state, to_state, transition_reason, operator_type, transition_metadata
+ FROM job_state_audit_logs
+ WHERE job_id = :job_id
+ ORDER BY id DESC
+ LIMIT 1
+ """
+ ),
+ {"job_id": job_id},
+ )
+ .mappings()
+ .one()
+ )
+
+ audit_metadata = dict(audit_log_row["transition_metadata"])
+ progress = redis_service.hgetall(f"task:{job_id}:progress")
+
+ assert job_row["status"] == "pending"
+ assert audit_log_row["from_state"] == "failed"
+ assert audit_log_row["to_state"] == "pending"
+ assert audit_log_row["transition_reason"] == "retry_transition"
+ assert audit_log_row["operator_type"] == "retry"
+ assert audit_metadata["worker"] == "contract"
+ assert audit_metadata["retry_reason"] == "task_retry"
+ assert audit_metadata["retry_count"] == 1
+ assert audit_metadata["retry_timestamp"]
+ assert progress["status"] == "pending"
+ assert progress["worker"] == "contract"
+ assert progress["retry_count"] == 1
diff --git a/apps/worker/tests/contract/test_url_upload_contract.py b/apps/worker/tests/contract/test_url_upload_contract.py
index 8c3ade508..cdf4a3de4 100644
--- a/apps/worker/tests/contract/test_url_upload_contract.py
+++ b/apps/worker/tests/contract/test_url_upload_contract.py
@@ -13,15 +13,22 @@
from support.contract_database import insert_contract_job, insert_contract_user
-def _load_upload_task_modules() -> tuple[Any, Engine, Any, Any]:
+def _load_upload_task_modules() -> tuple[Any, Any, Engine, Any, Any]:
import app.core.tasks.kb_tasks as kb_tasks
+ import app.services.workload.url_upload_service as url_upload_service
from shared.core.database_sync import get_sync_engine
from shared.services.redis.redis_sync_service import (
SyncJobInfoRedisService,
SyncRedisServiceFactory,
)
- return kb_tasks, get_sync_engine(), SyncJobInfoRedisService, SyncRedisServiceFactory
+ return (
+ kb_tasks,
+ url_upload_service,
+ get_sync_engine(),
+ SyncJobInfoRedisService,
+ SyncRedisServiceFactory,
+ )
def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progress(
@@ -29,9 +36,15 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres
monkeypatch: MonkeyPatch,
tmp_path: Path,
) -> None:
- kb_tasks, engine, sync_job_info_service_cls, sync_redis_service_factory = (
- _load_upload_task_modules()
- )
+ (
+ kb_tasks,
+ _url_upload_service,
+ engine,
+ sync_job_info_service_cls,
+ sync_redis_service_factory,
+ ) = _load_upload_task_modules()
+ from shared.core.config import settings
+ from shared.services.storage.job_file_storage import JobFileStorage
user_id = f"worker-user-{uuid4().hex[:12]}"
job_id = f"job_url_upload_{uuid4().hex[:12]}"
@@ -47,21 +60,21 @@ def resolve_public_address(
return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))]
monkeypatch.setattr(
- kb_tasks,
+ JobFileStorage,
"download_file_from_url",
- lambda _source_url: str(downloaded_path),
+ lambda self, _source_url, *, temp_dir=None: str(downloaded_path),
)
monkeypatch.setattr(
- kb_tasks,
- "upload_to_s3",
- lambda local_path, storage_key, bucket: uploaded_calls.append(
- (local_path, storage_key, bucket)
+ JobFileStorage,
+ "upload_source_file",
+ lambda self, local_path, storage_key: uploaded_calls.append(
+ (local_path, storage_key, self.uploads_bucket)
),
)
monkeypatch.setattr(
- kb_tasks,
- "verify_s3_file_exists",
- lambda storage_key: {"exists": storage_key == s3_key, "size": 3},
+ JobFileStorage,
+ "verify_upload_exists",
+ lambda self, storage_key: {"exists": storage_key == s3_key, "size": 3},
)
monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address)
@@ -112,7 +125,7 @@ def resolve_public_address(
"file_size": 3,
}
assert uploaded_calls == [
- (str(downloaded_path), s3_key, kb_tasks.settings.S3_BUCKET_NAME),
+ (str(downloaded_path), s3_key, settings.S3_BUCKET_NAME),
]
assert os.path.exists(downloaded_path) is False
@@ -140,4 +153,3 @@ def resolve_public_address(
assert job_row["status"] == "waiting-file"
assert job_row["source_type"] == "url"
assert job_row["s3_key"] == s3_key
-
diff --git a/apps/worker/tests/contract/test_webhook_recovery_contract.py b/apps/worker/tests/contract/test_webhook_recovery_contract.py
index b30633b9f..606e25818 100644
--- a/apps/worker/tests/contract/test_webhook_recovery_contract.py
+++ b/apps/worker/tests/contract/test_webhook_recovery_contract.py
@@ -36,7 +36,9 @@ def _insert_webhook_event(
created_at: datetime,
updated_at: datetime | None = None,
qstash_message_id: str | None = None,
+ payload: dict[str, Any] | None = None,
) -> None:
+ event_payload = payload or {"event": "job.failed", "job_id": job_id}
connection.execute(
text(
"""
@@ -69,7 +71,7 @@ def _insert_webhook_event(
"id": event_id,
"job_id": job_id,
"target_url": target_url,
- "payload": json.dumps({"event": "job.failed", "job_id": job_id}),
+ "payload": json.dumps(event_payload),
"status": status,
"attempts": attempts,
"next_retry_at": None,
@@ -80,6 +82,51 @@ def _insert_webhook_event(
)
+def _insert_job_result(
+ connection: Connection,
+ *,
+ job_result_id: str,
+ job_id: str,
+ result_s3_key: str,
+ inline_payload: dict[str, Any],
+) -> None:
+ timestamp = _utc_now()
+ connection.execute(
+ text(
+ """
+ INSERT INTO job_results (
+ id,
+ job_id,
+ delivery_mode,
+ inline_payload,
+ result_s3_key,
+ result_size,
+ created_at,
+ updated_at
+ ) VALUES (
+ :id,
+ :job_id,
+ 'url',
+ CAST(:inline_payload AS JSON),
+ :result_s3_key,
+ :result_size,
+ :created_at,
+ :updated_at
+ )
+ """
+ ),
+ {
+ "id": job_result_id,
+ "job_id": job_id,
+ "inline_payload": json.dumps(inline_payload),
+ "result_s3_key": result_s3_key,
+ "result_size": 123,
+ "created_at": timestamp,
+ "updated_at": timestamp,
+ },
+ )
+
+
def _load_worker_modules() -> tuple[Any, Any, Engine]:
import app.core.tasks.webhook_tasks as webhook_tasks
from shared.core.database_sync import get_sync_engine
@@ -132,8 +179,8 @@ def publish(self, **kwargs: Any) -> SimpleNamespace:
),
)
monkeypatch.setattr(
- publisher,
- "_get_client",
+ publisher._client_adapter,
+ "get_client",
lambda: SimpleNamespace(message=FakeMessageClient()),
)
@@ -271,6 +318,130 @@ def publish(self, **kwargs: Any) -> SimpleNamespace:
assert secrets_count_row["secrets_count"] == 1
+def test_should_publish_completed_webhook_with_result_delivery_payload(
+ worker_contract_environment: None,
+ monkeypatch: MonkeyPatch,
+) -> None:
+ _, qstash_publisher, engine = _load_worker_modules()
+ from shared.core.config import app_config
+ from shared.services.jobs.result_delivery import JobResultDeliveryResolver
+ from shared.services.storage.job_file_storage import JobFileStorage
+
+ user_id = f"worker-user-{uuid4().hex[:12]}"
+ target_url = "https://hooks.contract.test/worker"
+ job_id = f"job_completed_{uuid4().hex[:12]}"
+ event_id = str(uuid4())
+ result_s3_key = f"results/{job_id}.zip"
+ published_calls: list[dict[str, Any]] = []
+ signed_url_calls: list[dict[str, Any]] = []
+
+ class FakeMessageClient:
+ def publish(self, **kwargs: Any) -> SimpleNamespace:
+ published_calls.append(kwargs)
+ return SimpleNamespace(message_id=f"msg_{event_id}")
+
+ class FakeStorageAdapter:
+ def generate_presigned_url(
+ self,
+ key: str,
+ expiration: int = 3600,
+ bucket: str | None = None,
+ method: str = "GET",
+ headers: dict[str, str] | None = None,
+ ) -> str:
+ signed_url_calls.append(
+ {
+ "key": key,
+ "expiration": expiration,
+ "bucket": bucket,
+ "method": method,
+ "headers": headers,
+ }
+ )
+ return f"signed://{bucket}/{key}?expires={expiration}"
+
+ monkeypatch.setattr(
+ qstash_publisher,
+ "validate_http_url_and_resolve_ip",
+ lambda *args, **kwargs: SimpleNamespace(
+ is_valid=True,
+ error_message=None,
+ validated_ip="93.184.216.34",
+ hostname="hooks.contract.test",
+ ),
+ )
+ monkeypatch.setattr(
+ JobResultDeliveryResolver,
+ "__init__",
+ lambda self: setattr(
+ self,
+ "_storage",
+ JobFileStorage(storage_adapter=FakeStorageAdapter()),
+ ),
+ )
+ publisher = qstash_publisher.QStashWebhookPublisher()
+ monkeypatch.setattr(
+ publisher._client_adapter,
+ "get_client",
+ lambda: SimpleNamespace(message=FakeMessageClient()),
+ )
+
+ now = _utc_now()
+ with engine.begin() as connection:
+ insert_contract_user(connection, user_id=user_id)
+ insert_contract_job(
+ connection,
+ job_id=job_id,
+ user_id=user_id,
+ status="done",
+ source_type="file",
+ webhook_url=target_url,
+ webhook_enabled=True,
+ job_metadata=_build_file_job_metadata(),
+ billing_status="charged",
+ )
+ _insert_job_result(
+ connection,
+ job_result_id=str(uuid4()),
+ job_id=job_id,
+ result_s3_key=result_s3_key,
+ inline_payload={"checksum": "contract-checksum"},
+ )
+ _insert_webhook_event(
+ connection,
+ event_id=event_id,
+ job_id=job_id,
+ target_url=target_url,
+ status="pending",
+ attempts=0,
+ created_at=now,
+ payload={"event": "job.completed", "job_id": job_id},
+ )
+
+ message_id = publisher.publish_event(event_id)
+
+ assert message_id == f"msg_{event_id}"
+ assert len(published_calls) == 1
+ assert signed_url_calls == [
+ {
+ "key": result_s3_key,
+ "expiration": 3600,
+ "bucket": app_config.S3_RESULTS_BUCKET,
+ "method": "GET",
+ "headers": None,
+ }
+ ]
+
+ published_payload = json.loads(published_calls[0]["body"])
+ assert published_payload["event"] == "job.completed"
+ assert published_payload["job_id"] == job_id
+ assert published_payload["result"] == {"checksum": "contract-checksum"}
+ assert published_payload["result_url"] == (
+ f"signed://{app_config.S3_RESULTS_BUCKET}/{result_s3_key}"
+ "?expires=3600"
+ )
+
+
def test_should_reconcile_stale_delivering_webhook_events_from_qstash_logs(
worker_contract_environment: None,
monkeypatch: MonkeyPatch,
diff --git a/apps/worker/tests/contract/test_worker_job_file_storage_contract.py b/apps/worker/tests/contract/test_worker_job_file_storage_contract.py
new file mode 100644
index 000000000..b41b0bd9d
--- /dev/null
+++ b/apps/worker/tests/contract/test_worker_job_file_storage_contract.py
@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, BinaryIO
+
+
+class FakeStorageAdapter:
+ def __init__(self) -> None:
+ self.existing_keys: set[tuple[str, str]] = set()
+ self.object_sizes: dict[tuple[str, str], int] = {}
+ self.upload_calls: list[tuple[str, str, str]] = []
+ self.download_calls: list[tuple[str, str, str]] = []
+ self.presigned_calls: list[tuple[str, str, int, str]] = []
+
+ def generate_presigned_url(
+ self,
+ s3_key: str,
+ expiration: int = 3600,
+ bucket: str | None = None,
+ method: str = "GET",
+ headers: dict[str, str] | None = None,
+ ) -> str:
+ del headers
+ assert bucket is not None
+ self.presigned_calls.append((s3_key, bucket, expiration, method))
+ return f"https://storage.example.test/{bucket}/{s3_key}"
+
+ def exists(self, s3_key: str, bucket: str) -> bool:
+ return (s3_key, bucket) in self.existing_keys
+
+ def get_object_size(self, s3_key: str, bucket: str) -> int:
+ return self.object_sizes[(s3_key, bucket)]
+
+ def upload_file(self, file_path: str, s3_key: str, bucket: str) -> dict[str, Any]:
+ self.upload_calls.append((file_path, s3_key, bucket))
+ self.existing_keys.add((s3_key, bucket))
+ self.object_sizes[(s3_key, bucket)] = Path(file_path).stat().st_size
+ return {"bucket": bucket, "key": s3_key}
+
+ def upload_fileobj(
+ self,
+ file_obj: BinaryIO,
+ s3_key: str,
+ bucket: str,
+ content_type: str | None = None,
+ ) -> dict[str, Any]:
+ del file_obj, content_type
+ return {"bucket": bucket, "key": s3_key}
+
+ def download_file(self, s3_key: str, local_path: str, bucket: str) -> str:
+ self.download_calls.append((s3_key, local_path, bucket))
+ Path(local_path).write_bytes(b"downloaded")
+ return local_path
+
+
+def test_job_file_storage_should_hide_upload_bucket_rules_for_worker_source_files(
+ worker_contract_environment: None,
+ tmp_path: Path,
+) -> None:
+ from shared.services.storage.job_file_storage import JobFileStorage
+
+ del worker_contract_environment
+
+ storage_adapter = FakeStorageAdapter()
+ storage = JobFileStorage(
+ storage_adapter=storage_adapter,
+ uploads_bucket="uploads-bucket",
+ results_bucket="results-bucket",
+ )
+ source_path = tmp_path / "source.pdf"
+ source_path.write_bytes(b"pdf")
+
+ storage.upload_source_file(str(source_path), "uploads/job_123.pdf")
+ file_info = storage.verify_upload_exists("uploads/job_123.pdf")
+ download_info = storage.generate_upload_download_url(
+ "uploads/job_123.pdf",
+ expires_in=60,
+ )
+ downloaded_path = storage.download_upload_to_temp(
+ "uploads/job_123.pdf",
+ suffix=".pdf",
+ temp_dir=str(tmp_path),
+ )
+
+ assert file_info == {
+ "exists": True,
+ "size": 3,
+ "content_type": None,
+ "last_modified": None,
+ "etag": None,
+ }
+ assert download_info == {
+ "download_url": "https://storage.example.test/uploads-bucket/uploads/job_123.pdf",
+ "expires_in": 60,
+ }
+ assert Path(downloaded_path).read_bytes() == b"downloaded"
+ assert storage_adapter.upload_calls == [
+ (str(source_path), "uploads/job_123.pdf", "uploads-bucket")
+ ]
+ assert storage_adapter.download_calls == [
+ ("uploads/job_123.pdf", downloaded_path, "uploads-bucket")
+ ]
+ assert storage_adapter.presigned_calls == [
+ ("uploads/job_123.pdf", "uploads-bucket", 60, "GET")
+ ]
diff --git a/packages/shared-python/shared/core/async_utils.py b/packages/shared-python/shared/core/async_utils.py
deleted file mode 100644
index 1975d8b7a..000000000
--- a/packages/shared-python/shared/core/async_utils.py
+++ /dev/null
@@ -1,32 +0,0 @@
-import asyncio
-from typing import Any, Coroutine, TypeVar
-
-T = TypeVar("T")
-
-
-def run_async_task(coro: Coroutine[Any, Any, T]) -> T:
- """
- Run an async task in a synchronous context, reusing the event loop if possible.
-
- This function attempts to get the current event loop. If it's closed or missing,
- it creates a new one but DOES NOT close it after execution (unlike asyncio.run).
- This allows long-lived async resources to persist across tasks.
-
- Args:
- coro: The coroutine to run.
-
- Returns:
- The return value of the coroutine.
- """
- try:
- loop = asyncio.get_event_loop()
- if loop.is_closed():
- # Loop exists but closed - create new one
- loop = asyncio.new_event_loop()
- asyncio.set_event_loop(loop)
- except RuntimeError:
- # No loop in this thread - create new one
- loop = asyncio.new_event_loop()
- asyncio.set_event_loop(loop)
-
- return loop.run_until_complete(coro)
diff --git a/packages/shared-python/shared/core/celery_router.py b/packages/shared-python/shared/core/celery_router.py
index fe2ca4502..7c8fd86fb 100644
--- a/packages/shared-python/shared/core/celery_router.py
+++ b/packages/shared-python/shared/core/celery_router.py
@@ -214,16 +214,8 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str:
"""
try:
# TODO: temporary simplified path to avoid async work here.
- priority_level = 1 # Default free-subscription level.
-
- # Choose the queue by task type and priority level.
if job_type in ["kb_management", "kb_encoding"]:
- if priority_level >= 9:
- return "kb_high"
- elif priority_level >= 5:
- return "kb_medium"
- else:
- return "kb_low"
+ return "kb_low"
elif job_type in ["ai_query", "user_auth", "urgent_document"]:
return "ai_high_priority"
elif job_type in ["document_processing"]:
diff --git a/packages/shared-python/shared/core/logging.py b/packages/shared-python/shared/core/logging.py
index c2633b806..bb246f18e 100644
--- a/packages/shared-python/shared/core/logging.py
+++ b/packages/shared-python/shared/core/logging.py
@@ -3,11 +3,12 @@
from contextlib import contextmanager
from contextvars import ContextVar
from enum import Enum
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict
from loguru import logger
-from logfire.types import ExceptionCallbackHelper
+if TYPE_CHECKING:
+ from logfire.types import ExceptionCallbackHelper
_log_context: ContextVar[Dict[str, Any]] = ContextVar("log_context", default={})
_DEFAULT_CONSOLE_FORMAT = (
diff --git a/packages/shared-python/shared/core/state_machine/config.py b/packages/shared-python/shared/core/state_machine/config.py
deleted file mode 100644
index 2128df763..000000000
--- a/packages/shared-python/shared/core/state_machine/config.py
+++ /dev/null
@@ -1,92 +0,0 @@
-"""State-machine configuration."""
-
-from dataclasses import dataclass, field
-from typing import Dict
-
-
-DEFAULT_STATE_TIMEOUTS: Dict[str, int] = {
- "pending": 300,
- "uploading": 600,
- "processing": 1800,
- "completed": 0,
- "failed": 0,
-}
-
-
-@dataclass
-class StateMachineConfig:
- """State-machine settings."""
-
- max_retries: int = 3 # Maximum retry count.
- base_retry_delay: float = 0.1 # Base retry delay in seconds.
-
- # Timeout settings used with Redis Keyspace Notifications.
- state_timeouts: Dict[str, int] = field(
- default_factory=lambda: DEFAULT_STATE_TIMEOUTS.copy()
- )
-
- # Synchronization settings.
- sync_batch_size: int = 100 # Batch size for sync work.
- sync_interval: int = 300 # Sync interval in seconds.
-
- # Maintenance settings.
- maintenance_interval: int = 3600 # Maintenance interval in seconds.
- cleanup_interval: int = 1800 # Cleanup interval in seconds.
-
- # Redis Keyspace Notifications support.
- enable_keyspace_notifications: bool = True # Enable Keyspace Notifications.
-
-# Default state-machine configuration.
-DEFAULT_CONFIG = StateMachineConfig()
-
-
-def get_state_machine_config() -> StateMachineConfig:
- """Return the active state-machine configuration."""
- return DEFAULT_CONFIG
-
-
-def update_state_machine_config(
- *,
- max_retries: int | None = None,
- base_retry_delay: float | None = None,
- state_timeouts: Dict[str, int] | None = None,
- sync_batch_size: int | None = None,
- sync_interval: int | None = None,
- maintenance_interval: int | None = None,
- cleanup_interval: int | None = None,
- enable_keyspace_notifications: bool | None = None,
-) -> StateMachineConfig:
- """Update and return the active state-machine configuration."""
- global DEFAULT_CONFIG
- current = DEFAULT_CONFIG
- DEFAULT_CONFIG = StateMachineConfig(
- max_retries=current.max_retries if max_retries is None else max_retries,
- base_retry_delay=(
- current.base_retry_delay
- if base_retry_delay is None
- else base_retry_delay
- ),
- state_timeouts=(
- current.state_timeouts.copy()
- if state_timeouts is None
- else state_timeouts
- ),
- sync_batch_size=(
- current.sync_batch_size if sync_batch_size is None else sync_batch_size
- ),
- sync_interval=current.sync_interval if sync_interval is None else sync_interval,
- maintenance_interval=(
- current.maintenance_interval
- if maintenance_interval is None
- else maintenance_interval
- ),
- cleanup_interval=(
- current.cleanup_interval if cleanup_interval is None else cleanup_interval
- ),
- enable_keyspace_notifications=(
- current.enable_keyspace_notifications
- if enable_keyspace_notifications is None
- else enable_keyspace_notifications
- ),
- )
- return DEFAULT_CONFIG
diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py
index ee5f421a4..366649acb 100644
--- a/packages/shared-python/shared/core/state_machine/service.py
+++ b/packages/shared-python/shared/core/state_machine/service.py
@@ -7,7 +7,6 @@
import asyncio
import time
-from datetime import datetime, timezone
from typing import Any, Dict, Optional
from loguru import logger
@@ -19,18 +18,19 @@
JobStatus,
is_valid_transition,
)
+from shared.core.state_machine.transition_payloads import (
+ build_failure_transition_metadata,
+ build_progress_cache_payload,
+ build_retry_transition,
+ serialize_transition_metadata,
+ utc_now_naive,
+)
from shared.models.database.job import Job
from shared.models.database.job_state_audit_log import JobStateAuditLog
from shared.services.redis import RedisServiceFactory
-from shared.utils.error_details import normalize_error_details
-from shared.utils.json_utils import make_json_safe
from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder
-def _utc_now_naive() -> datetime:
- return datetime.now(timezone.utc).replace(tzinfo=None)
-
-
class AsyncStateMachineService:
"""Async state machine service — used by the API (FastAPI + asyncpg)."""
@@ -136,7 +136,14 @@ async def mark_failed(
) -> bool:
"""Mark a job as failed with error information."""
try:
- normalized_details = normalize_error_details(error_details)
+ normalized_details, transition_metadata = (
+ build_failure_transition_metadata(
+ error_message=error_message,
+ error_code=error_code,
+ error_details=error_details,
+ metadata=metadata,
+ )
+ )
await self._update_job_error(
db,
job_id,
@@ -145,12 +152,6 @@ async def mark_failed(
normalized_details,
)
- transition_metadata = (metadata or {}).copy()
- transition_metadata["error_message"] = error_message
- transition_metadata["error_code"] = error_code
- if normalized_details:
- transition_metadata["error_details"] = normalized_details
-
return await self.transition(
db,
job_id,
@@ -208,17 +209,11 @@ async def handle_retry(
logger.error(f"Job {job_id} has no status")
return False
- retry_target = (
- JobStatus.PENDING.value
- if current_state == JobStatus.FAILED.value
- else current_state
+ retry_target, retry_metadata = build_retry_transition(
+ current_state=current_state,
+ retry_metadata=retry_metadata,
)
- retry_metadata = retry_metadata or {}
- retry_metadata["retry_reason"] = "task_retry"
- retry_metadata["retry_timestamp"] = str(int(time.time()))
- retry_metadata["retry_count"] = retry_metadata.get("retry_count", 0) + 1
-
# Always use full transition() for CAS protection — even same-state
return await self.transition(
db,
@@ -291,7 +286,7 @@ async def _cas_update_state(
.values(
status=to_state,
version=old_version + 1,
- updated_at=_utc_now_naive(),
+ updated_at=utc_now_naive(),
)
)
return result.rowcount > 0
@@ -307,13 +302,7 @@ async def _record_audit_log(
operator_type: str,
metadata: Optional[Dict[str, Any]],
) -> None:
- serialized = None
- if metadata:
- try:
- serialized = make_json_safe(metadata)
- except Exception as e:
- logger.warning(f"Metadata serialization failed: {e}")
- serialized = {"error": "metadata_serialization_failed"}
+ serialized = serialize_transition_metadata(metadata)
db.add(
JobStateAuditLog(
@@ -392,13 +381,20 @@ async def _update_redis_cache(
)
progress_key = redis_key_builder.task_progress(job_id)
- progress_data: Dict[str, Any] = {
- "status": status,
- "timestamp": str(int(time.time())),
- }
+ progress_data: Dict[str, Any] = build_progress_cache_payload(
+ status=status,
+ metadata=None,
+ timestamp=int(time.time()),
+ )
if metadata:
try:
- progress_data.update(make_json_safe(metadata))
+ progress_data.update(
+ build_progress_cache_payload(
+ status=status,
+ metadata=metadata,
+ timestamp=int(time.time()),
+ )
+ )
except Exception as e:
logger.warning(f"Metadata serialization skipped: {e}")
diff --git a/packages/shared-python/shared/core/state_machine/service_sync.py b/packages/shared-python/shared/core/state_machine/service_sync.py
index 88e9aa3c6..175802459 100644
--- a/packages/shared-python/shared/core/state_machine/service_sync.py
+++ b/packages/shared-python/shared/core/state_machine/service_sync.py
@@ -7,7 +7,6 @@
"""
import time
-from datetime import datetime, timezone
from typing import Any, Dict, Optional
from loguru import logger
@@ -18,18 +17,19 @@
JobStatus,
is_valid_transition,
)
+from shared.core.state_machine.transition_payloads import (
+ build_failure_transition_metadata,
+ build_progress_cache_payload,
+ build_retry_transition,
+ serialize_transition_metadata,
+ utc_now_naive,
+)
from shared.models.database.job import Job
from shared.models.database.job_state_audit_log import JobStateAuditLog
from shared.services.redis.redis_sync_service import SyncRedisServiceFactory
-from shared.utils.error_details import normalize_error_details
-from shared.utils.json_utils import make_json_safe
from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder
-def _utc_now_naive() -> datetime:
- return datetime.now(timezone.utc).replace(tzinfo=None)
-
-
class SyncStateMachineService:
"""Sync state machine service — used by the Worker (gevent + psycopg2)."""
@@ -122,7 +122,14 @@ def mark_failed(
) -> bool:
"""Mark a job as failed with error information."""
try:
- normalized_details = normalize_error_details(error_details)
+ normalized_details, transition_metadata = (
+ build_failure_transition_metadata(
+ error_message=error_message,
+ error_code=error_code,
+ error_details=error_details,
+ metadata=metadata,
+ )
+ )
self._update_job_error(
db,
job_id,
@@ -131,12 +138,6 @@ def mark_failed(
normalized_details,
)
- transition_metadata = (metadata or {}).copy()
- transition_metadata["error_message"] = error_message
- transition_metadata["error_code"] = error_code
- if normalized_details:
- transition_metadata["error_details"] = normalized_details
-
return self.transition(
db,
job_id,
@@ -172,6 +173,48 @@ def mark_completed(
logger.error(f"Failed to mark Job {job_id} as completed: {e}")
return False
+ def handle_retry(
+ self,
+ db: Session,
+ job_id: str,
+ retry_metadata: Optional[Dict[str, Any]] = None,
+ operator_id: Optional[str] = None,
+ ) -> bool:
+ """Handle task retry — always goes through CAS-protected transition."""
+ try:
+ job = self._get_job(db, job_id)
+ if not job:
+ logger.error(f"Job {job_id} does not exist")
+ return False
+
+ current_state = job.status
+ if not current_state:
+ logger.error(f"Job {job_id} has no status")
+ return False
+
+ retry_target, retry_metadata = build_retry_transition(
+ current_state=current_state,
+ retry_metadata=retry_metadata,
+ )
+
+ return self.transition(
+ db,
+ job_id,
+ retry_target,
+ "retry_transition",
+ operator_id,
+ "retry",
+ retry_metadata,
+ )
+ except Exception as e:
+ logger.error(f"Job {job_id} retry failed: {e}")
+ try:
+ if db.is_active:
+ db.rollback()
+ except Exception as rollback_err:
+ logger.warning(f"Job {job_id} rollback failed: {rollback_err}")
+ return False
+
# ── Private helpers ─────────────────────────────────────────────────
def _get_job(self, db: Session, job_id: str) -> Optional[Job]:
@@ -199,7 +242,7 @@ def _cas_update_state(
.values(
status=to_state,
version=old_version + 1,
- updated_at=_utc_now_naive(),
+ updated_at=utc_now_naive(),
)
)
return result.rowcount > 0
@@ -215,13 +258,7 @@ def _record_audit_log(
operator_type: str,
metadata: Optional[Dict[str, Any]],
) -> None:
- serialized = None
- if metadata:
- try:
- serialized = make_json_safe(metadata)
- except Exception as e:
- logger.warning(f"Metadata serialization failed: {e}")
- serialized = {"error": "metadata_serialization_failed"}
+ serialized = serialize_transition_metadata(metadata)
db.add(
JobStateAuditLog(
@@ -281,13 +318,20 @@ def _update_redis_cache(
)
progress_key = redis_key_builder.task_progress(job_id)
- progress_data: Dict[str, Any] = {
- "status": status,
- "timestamp": str(int(time.time())),
- }
+ progress_data: Dict[str, Any] = build_progress_cache_payload(
+ status=status,
+ metadata=None,
+ timestamp=int(time.time()),
+ )
if metadata:
try:
- progress_data.update(make_json_safe(metadata))
+ progress_data.update(
+ build_progress_cache_payload(
+ status=status,
+ metadata=metadata,
+ timestamp=int(time.time()),
+ )
+ )
except Exception as e:
logger.warning(f"Metadata serialization skipped: {e}")
diff --git a/packages/shared-python/shared/core/state_machine/transition_payloads.py b/packages/shared-python/shared/core/state_machine/transition_payloads.py
new file mode 100644
index 000000000..c2a2cf58b
--- /dev/null
+++ b/packages/shared-python/shared/core/state_machine/transition_payloads.py
@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any
+
+from shared.core.state_machine.states import JobStatus
+from shared.utils.error_details import normalize_error_details
+from shared.utils.json_utils import make_json_safe
+
+
+def utc_now_naive() -> datetime:
+ return datetime.now(timezone.utc).replace(tzinfo=None)
+
+
+def serialize_transition_metadata(
+ metadata: dict[str, Any] | None,
+) -> dict[str, Any] | None:
+ if not metadata:
+ return None
+
+ try:
+ return make_json_safe(metadata)
+ except Exception:
+ return {"error": "metadata_serialization_failed"}
+
+
+def build_failure_transition_metadata(
+ *,
+ error_message: str,
+ error_code: str,
+ error_details: dict[str, Any] | None,
+ metadata: dict[str, Any] | None,
+) -> tuple[dict[str, Any] | None, dict[str, Any]]:
+ normalized_details = normalize_error_details(error_details)
+ transition_metadata = (metadata or {}).copy()
+ transition_metadata["error_message"] = error_message
+ transition_metadata["error_code"] = error_code
+ if normalized_details:
+ transition_metadata["error_details"] = normalized_details
+ return normalized_details, transition_metadata
+
+
+def build_retry_transition(
+ *,
+ current_state: str,
+ retry_metadata: dict[str, Any] | None,
+) -> tuple[str, dict[str, Any]]:
+ retry_target = (
+ JobStatus.PENDING.value
+ if current_state == JobStatus.FAILED.value
+ else current_state
+ )
+
+ resolved_metadata = retry_metadata or {}
+ resolved_metadata["retry_reason"] = "task_retry"
+ resolved_metadata["retry_timestamp"] = str(int(datetime.now(timezone.utc).timestamp()))
+ resolved_metadata["retry_count"] = resolved_metadata.get("retry_count", 0) + 1
+ return retry_target, resolved_metadata
+
+
+def build_progress_cache_payload(
+ *,
+ status: str,
+ metadata: dict[str, Any] | None,
+ timestamp: int,
+) -> dict[str, Any]:
+ progress_data: dict[str, Any] = {
+ "status": status,
+ "timestamp": str(timestamp),
+ }
+ if metadata:
+ progress_data.update(make_json_safe(metadata))
+ return progress_data
diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py
index 79cfbff81..a4d2c7d96 100644
--- a/packages/shared-python/shared/models/schemas/job_metadata.py
+++ b/packages/shared-python/shared/models/schemas/job_metadata.py
@@ -52,6 +52,35 @@ def create_from_request(request, **kwargs) -> Dict[str, Any]:
metadata.update(kwargs)
return metadata
+ @staticmethod
+ def set_document_scope(
+ metadata: Dict[str, Any],
+ *,
+ document_id: str,
+ namespace: str,
+ ) -> None:
+ """Store the effective retrieval document scope."""
+ metadata["document_id"] = document_id
+ metadata["namespace"] = namespace
+
+ @staticmethod
+ def set_file_source(metadata: Dict[str, Any], *, source_file_name: str) -> None:
+ """Store source metadata for direct file uploads."""
+ metadata["source_file_name"] = source_file_name
+ metadata["source_type"] = "file"
+
+ @staticmethod
+ def set_url_source(
+ metadata: Dict[str, Any],
+ *,
+ source_file_name: str,
+ source_url: str,
+ ) -> None:
+ """Store source metadata for URL ingestion."""
+ metadata["source_file_name"] = source_file_name
+ metadata["source_url"] = source_url
+ metadata["source_type"] = "url"
+
@staticmethod
def get_field(
metadata: Optional[Dict[str, Any]], field: str, default: Any = None
@@ -61,6 +90,53 @@ def get_field(
return default
return metadata.get(field, default)
+ @staticmethod
+ def get_string_field(
+ metadata: Optional[Dict[str, Any]], field: str, default: str | None = None
+ ) -> str | None:
+ """Read a string field from metadata."""
+ value = JobMetadataHelper.get_field(metadata, field, default)
+ return value if isinstance(value, str) else default
+
+ @staticmethod
+ def get_original_request(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+ """Return the stored creation request payload."""
+ original_request = JobMetadataHelper.get_field(metadata, "original_request", {})
+ return original_request if isinstance(original_request, dict) else {}
+
+ @staticmethod
+ def get_parsing_params_dict(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+ """Return stored parsing parameters as a dictionary."""
+ parsing_params = JobMetadataHelper.get_field(metadata, "parsing_params", {})
+ return parsing_params if isinstance(parsing_params, dict) else {}
+
+ @staticmethod
+ def get_namespace(
+ metadata: Optional[Dict[str, Any]], default: str | None = None
+ ) -> str | None:
+ """Return the retrieval namespace stored in metadata."""
+ return JobMetadataHelper.get_string_field(metadata, "namespace", default)
+
+ @staticmethod
+ def get_document_id(metadata: Optional[Dict[str, Any]]) -> str | None:
+ """Return the retrieval document id stored in metadata."""
+ return JobMetadataHelper.get_string_field(metadata, "document_id")
+
+ @staticmethod
+ def get_data_id(metadata: Optional[Dict[str, Any]]) -> str | None:
+ """Return the user-defined data id stored in metadata."""
+ return JobMetadataHelper.get_string_field(metadata, "data_id")
+
+ @staticmethod
+ def get_source_file_name(metadata: Optional[Dict[str, Any]]) -> str | None:
+ """Return the source file name stored in metadata."""
+ return JobMetadataHelper.get_string_field(metadata, "source_file_name")
+
+ @staticmethod
+ def get_source_url(metadata: Optional[Dict[str, Any]]) -> str | None:
+ """Return the source URL stored in metadata."""
+ return JobMetadataHelper.get_string_field(metadata, "source_url")
+
@staticmethod
def get_parsing_param(
metadata: Optional[Dict[str, Any]], param: str, default: Any = None
diff --git a/packages/shared-python/shared/models/schemas/s3_file.py b/packages/shared-python/shared/models/schemas/s3_file.py
deleted file mode 100644
index 4b3a0cbbc..000000000
--- a/packages/shared-python/shared/models/schemas/s3_file.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from pydantic import BaseModel
-
-
-class FliesDownload(BaseModel):
- message: str
- file_key: str
- download_url: str
- expires_in_seconds: int
diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/job_failure_sync.py
new file mode 100644
index 000000000..0888b1c27
--- /dev/null
+++ b/packages/shared-python/shared/services/job_failure_sync.py
@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from shared.core.response import build_standard_error_response
+from shared.core.state_machine.service_sync import SyncStateMachineService
+from shared.models.database.job import Job
+from shared.models.database.webhook import WebhookEvent
+from shared.services.billing.credits_sync_service import SyncCreditsService
+from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox
+from shared.utils.error_details import normalize_error_details
+
+
+class SyncJobFailureFinalizer:
+ """Finalize failed Jobs inside the lifecycle transaction."""
+
+ def __init__(
+ self,
+ *,
+ state_machine: SyncStateMachineService | None = None,
+ webhook_outbox: SyncJobWebhookOutbox | None = None,
+ credits_service: SyncCreditsService | None = None,
+ ) -> None:
+ self._state_machine = state_machine or SyncStateMachineService()
+ self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox()
+ self._credits_service = credits_service or SyncCreditsService()
+
+ def finalize(
+ self,
+ db: Session,
+ *,
+ job_id: str,
+ error_message: str,
+ error_code: str,
+ error_details: dict[str, Any] | None,
+ should_refund: bool,
+ ) -> tuple[bool, WebhookEvent | None]:
+ transition_ok = self._state_machine.mark_failed(
+ db,
+ job_id,
+ error_message,
+ error_code=error_code,
+ error_details=error_details,
+ )
+ if not transition_ok:
+ logger.error(f"Job {job_id} mark_failed transition failed")
+ return False, None
+
+ if should_refund:
+ self._try_refund_credits(db, job_id)
+
+ normalized_error_details = normalize_error_details(error_details)
+ webhook_event = self._webhook_outbox.create_event(
+ db,
+ job_id=job_id,
+ event_type="job.failed",
+ extra_payload={
+ "error": build_standard_error_response(
+ code=error_code,
+ message=error_message,
+ request_id=job_id,
+ details=normalized_error_details,
+ ),
+ },
+ )
+ return True, webhook_event
+
+ def enqueue_webhook_after_commit(self, webhook_event: WebhookEvent | None) -> None:
+ self._webhook_outbox.enqueue_after_commit(webhook_event)
+
+ def _try_refund_credits(self, db: Session, job_id: str) -> None:
+ try:
+ result = db.execute(select(Job).where(Job.job_id == job_id))
+ job = result.scalar_one_or_none()
+ if not job:
+ return
+
+ amount = getattr(job, "credits_charged", 0) or 0
+ billing_status = getattr(job, "billing_status", "")
+ if amount <= 0 or billing_status != "charged":
+ return
+
+ self._credits_service.refund_job_credits(
+ session=db,
+ user_id=str(job.user_id),
+ amount=amount,
+ job_id=job_id,
+ )
+ job.billing_status = "refunded"
+ logger.info(f"Refunded {amount} credits for job {job_id}")
+ except Exception as exc:
+ logger.error(f"Credit refund failed for job {job_id}: {exc}")
diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py
index cd9dcc8ef..9d79b4ce0 100644
--- a/packages/shared-python/shared/services/job_lifecycle_sync.py
+++ b/packages/shared-python/shared/services/job_lifecycle_sync.py
@@ -10,34 +10,19 @@
from __future__ import annotations
import time
-from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
-from uuid import uuid4
from loguru import logger
-from sqlalchemy import delete, select, update
-from sqlalchemy.orm import Session
from shared.core.database_sync import get_sync_db_context
-from shared.core.response import build_standard_error_response
-from shared.core.state_machine.service_sync import SyncStateMachineService
-from shared.models.database.job import Job
-from shared.models.database.document import DocumentSection
-from shared.models.database.job_result import JobChunk, JobResult
-from shared.models.database.webhook import WebhookEvent, WebhookEventStatus
-from shared.services.billing.credits_sync_service import SyncCreditsService
+from shared.services.job_failure_sync import SyncJobFailureFinalizer
+from shared.services.job_success_sync import SyncJobSuccessFinalizer
from shared.services.redis.redis_sync_service import (
SyncRedisServiceFactory,
)
-from shared.services.retrieval.publication_service import RetrievalPublicationService
-from shared.utils.error_details import normalize_error_details
from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder
-def _utc_now_naive() -> datetime:
- return datetime.now(timezone.utc).replace(tzinfo=None)
-
-
class SyncJobLifecycleService:
"""Manages job lifecycle transitions in the worker process (sync/gevent).
@@ -45,8 +30,8 @@ class SyncJobLifecycleService:
"""
def __init__(self) -> None:
- self._state_machine = SyncStateMachineService()
- self._retrieval_publication = RetrievalPublicationService()
+ self._success_finalizer = SyncJobSuccessFinalizer()
+ self._failure_finalizer = SyncJobFailureFinalizer()
# ── Public API ──────────────────────────────────────────────────────
@@ -75,88 +60,27 @@ def finalize_job_success(
with get_sync_db_context() as db:
try:
- inline_payload = {"checksum": checksum}
- job_result = self._upsert_job_result(
- db,
- job_id,
- delivery_mode,
- inline_payload=inline_payload,
- result_s3_key=result_s3_key,
- result_size=zip_size,
- )
-
- normalized_chunks = chunks or []
- self._replace_chunks(db, job_result.id, normalized_chunks)
- previous_document_scope = (
- self._retrieval_publication.get_existing_document_scope(
- db,
- job_id=job_id,
- )
- )
- published_document_state = (
- self._retrieval_publication.publish_document_state(
- db,
- job_id=job_id,
- job_result_id=job_result.id,
- chunks=normalized_chunks,
- )
- )
- if published_document_state is not None and not published_document_state.get("skipped_all_duplicate"):
- # Backfill DocumentSection.summary from enriched doc_nav data
- if section_summaries:
- self._backfill_section_summaries(
- db,
- document_id=published_document_state.get("document_id", ""),
- job_result_id=job_result.id,
- section_summaries=section_summaries,
- )
- self._retrieval_publication.publish_document_graph(
- db,
- job_id=job_id,
- job_result_id=job_result.id,
- )
- cache_invalidation = self._build_retrieval_cache_invalidation(
+ finalization = self._success_finalizer.finalize(
db,
job_id=job_id,
- published_document_state=published_document_state,
- previous_document_scope=previous_document_scope,
- )
-
- transition_ok = self._state_machine.mark_completed(
- db,
- job_id,
- result_metadata={
- "storage_completed": True,
- "stored_count": stored_count,
- "delivery_mode": delivery_mode,
- },
+ result_s3_key=result_s3_key,
+ checksum=checksum,
+ zip_size=zip_size,
+ chunks=chunks or [],
+ stored_count=stored_count,
+ delivery_mode=delivery_mode,
+ section_summaries=section_summaries,
)
- if not transition_ok:
- logger.error(f"Job {job_id} mark_completed transition failed")
+ if finalization.response.get("status") != "success":
db.rollback()
- return {
- "status": "failed",
- "job_id": job_id,
- "reason": "state_transition_failed",
- }
-
- webhook_event = self._maybe_create_webhook_event(
- db,
- job_id,
- event_type="job.completed",
- )
+ return finalization.response
db.commit()
logger.info(f"Job {job_id} success transaction committed")
- self._post_commit_invalidate_retrieval_cache(cache_invalidation)
- self._post_commit_enqueue_webhook(webhook_event)
+ self._success_finalizer.run_post_commit_actions(finalization)
- return {
- "status": "success",
- "job_id": job_id,
- "stored_count": stored_count,
- }
+ return finalization.response
except Exception as exc:
logger.error(f"Failed to finalize job success {job_id}: {exc}")
@@ -184,40 +108,22 @@ def finalize_job_failure(
with get_sync_db_context() as db:
try:
- transition_ok = self._state_machine.mark_failed(
+ transition_ok, webhook_event = self._failure_finalizer.finalize(
db,
- job_id,
- error_message,
+ job_id=job_id,
+ error_message=error_message,
error_code=error_code,
error_details=error_details,
+ should_refund=should_refund,
)
if not transition_ok:
- logger.error(f"Job {job_id} mark_failed transition failed")
db.rollback()
return False
- if should_refund:
- self._try_refund_credits(db, job_id)
-
- normalized_error_details = normalize_error_details(error_details)
- webhook_event = self._maybe_create_webhook_event(
- db,
- job_id,
- event_type="job.failed",
- extra_payload={
- "error": build_standard_error_response(
- code=error_code,
- message=error_message,
- request_id=job_id,
- details=normalized_error_details,
- ),
- },
- )
-
db.commit()
logger.info(f"Job {job_id} failure transaction committed")
- self._post_commit_enqueue_webhook(webhook_event)
+ self._failure_finalizer.enqueue_webhook_after_commit(webhook_event)
return True
@@ -259,248 +165,6 @@ def update_progress(
logger.warning(f"Failed to update progress for job {job_id}: {exc}")
return False
- # ── Private helpers ─────────────────────────────────────────────────
-
- def _backfill_section_summaries(
- self,
- db: Session,
- *,
- document_id: str,
- job_result_id: str,
- section_summaries: Dict[str, str],
- ) -> None:
- """Populate DocumentSection.summary from enriched doc_nav data.
-
- Runs UPDATE statements within the existing transaction so no extra
- commit is needed. Overwrites any existing summary value since the
- enriched doc_nav data is the authoritative source.
- """
- if not document_id or not section_summaries:
- return
- try:
- for path, summary in section_summaries.items():
- if not path or not summary:
- continue
- db.execute(
- update(DocumentSection)
- .where(DocumentSection.document_id == document_id)
- .where(DocumentSection.job_result_id == job_result_id)
- .where(DocumentSection.section_path == path)
- .values(summary=summary)
- )
- db.flush()
- logger.debug(
- f"Backfilled section summaries: document_id={document_id}, "
- f"count={len(section_summaries)}"
- )
- except Exception as exc:
- logger.warning(f"Section summary backfill failed (non-fatal): {exc}")
-
- def _upsert_job_result(
- self,
- db: Session,
- job_id: str,
- delivery_mode: str,
- *,
- inline_payload: Optional[Dict[str, Any]] = None,
- result_s3_key: Optional[str] = None,
- result_size: Optional[int] = None,
- ) -> JobResult:
- """Create or update JobResult row."""
- result = db.execute(select(JobResult).where(JobResult.job_id == job_id))
- existing = result.scalar_one_or_none()
-
- if existing:
- existing.delivery_mode = delivery_mode
- existing.inline_payload = inline_payload
- existing.result_s3_key = result_s3_key
- existing.result_size = result_size
- db.flush()
- return existing
-
- job_result = JobResult(
- job_id=job_id,
- delivery_mode=delivery_mode,
- document_metadata={},
- inline_payload=inline_payload,
- result_s3_key=result_s3_key,
- result_size=result_size,
- )
- db.add(job_result)
- db.flush()
- return job_result
-
- def _replace_chunks(
- self,
- db: Session,
- job_result_id: str,
- chunks: List[Dict[str, Any]],
- ) -> Optional[Dict[str, Any]]:
- """Delete existing chunks and insert new ones."""
- db.execute(delete(JobChunk).where(JobChunk.job_result_id == job_result_id))
-
- if not chunks:
- db.flush()
- return
-
- chunk_models = []
- for index, chunk in enumerate(chunks):
- chunk_identifier = chunk.get("chunk_id") or str(uuid4())
- chunk_models.append(
- JobChunk(
- job_result_id=job_result_id,
- chunk_id=chunk_identifier,
- chunk_type=chunk.get("type", "paragraph"),
- text=chunk.get("text"),
- path=chunk.get("metadata", {}).get("path"),
- chunk_metadata=chunk.get("metadata"),
- sort_order=chunk.get("order", index),
- )
- )
- db.add_all(chunk_models)
- db.flush()
-
- def _build_retrieval_cache_invalidation(
- self,
- db: Session,
- *,
- job_id: str,
- published_document_state: Optional[Dict[str, str]],
- previous_document_scope: Optional[Dict[str, str]],
- ) -> Optional[Dict[str, Any]]:
- job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none()
- if not job:
- return None
-
- namespaces: list[str] = []
- metadata = job.job_metadata or {}
- new_namespace = metadata.get("namespace") or "default"
- namespaces.append(new_namespace)
-
- if previous_document_scope and previous_document_scope.get("namespace"):
- namespaces.append(previous_document_scope["namespace"])
- if published_document_state and published_document_state.get("namespace"):
- namespaces.append(published_document_state["namespace"])
-
- return {"user_id": str(job.user_id), "namespaces": namespaces, "job_id": job_id}
-
- def _post_commit_invalidate_retrieval_cache(
- self, cache_invalidation: Optional[Dict[str, Any]]
- ) -> None:
- if not cache_invalidation:
- return
- try:
- redis_service = SyncRedisServiceFactory.get_service()
- user_id = cache_invalidation["user_id"]
- seen: set[str] = set()
- for namespace in cache_invalidation["namespaces"]:
- if not namespace or namespace in seen:
- continue
- seen.add(namespace)
- redis_service.incr(f"retrieval:version:{user_id}:{namespace}")
- except Exception as exc:
- logger.warning(
- f"Failed to invalidate retrieval cache after publication (ignored): job_id={cache_invalidation.get('job_id')}, error={exc}"
- )
-
- def _maybe_create_webhook_event(
- self,
- db: Session,
- job_id: str,
- event_type: str,
- extra_payload: Optional[Dict[str, Any]] = None,
- ) -> Optional[WebhookEvent]:
- """Create a WebhookEvent if the job has webhooks enabled."""
- result = db.execute(select(Job).where(Job.job_id == job_id))
- job = result.scalar_one_or_none()
-
- if not job:
- logger.warning(f"Job not found for webhook check: {job_id}")
- return None
-
- webhook_url = getattr(job, "webhook_url", None)
- if not job.webhook_enabled or not webhook_url:
- return None
-
- status = "completed" if event_type == "job.completed" else "failed"
- timestamp_key = f"{status}_at"
- payload: Dict[str, Any] = {
- "event": event_type,
- "job_id": job_id,
- "status": status,
- timestamp_key: _utc_now_naive().isoformat(),
- }
- if extra_payload:
- payload.update(extra_payload)
-
- event = WebhookEvent(
- job_id=job_id,
- target_url=webhook_url,
- payload=payload,
- status=WebhookEventStatus.PENDING,
- attempts=0,
- )
- db.add(event)
- db.flush()
- logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}")
- return event
-
- def _try_refund_credits(self, db: Session, job_id: str) -> None:
- """Attempt to refund credits for a failed job."""
- try:
- result = db.execute(select(Job).where(Job.job_id == job_id))
- job = result.scalar_one_or_none()
- if not job:
- return
-
- amount = getattr(job, "credits_charged", 0) or 0
- billing_status = getattr(job, "billing_status", "")
- if amount <= 0 or billing_status != "charged":
- return
-
- credits_service = SyncCreditsService()
- credits_service.refund_job_credits(
- session=db,
- user_id=str(job.user_id),
- amount=amount,
- job_id=job_id,
- )
- job.billing_status = "refunded"
- logger.info(f"Refunded {amount} credits for job {job_id}")
- except Exception as exc:
- logger.error(f"Credit refund failed for job {job_id}: {exc}")
-
- def _post_commit_enqueue_webhook(
- self,
- webhook_event: Optional[WebhookEvent],
- ) -> None:
- """Publish a persisted webhook via QStash after commit (best-effort)."""
- if not webhook_event:
- return
-
- try:
- from shared.services.webhook.qstash_publisher import (
- get_qstash_webhook_publisher,
- )
-
- publisher = get_qstash_webhook_publisher()
- message_id = publisher.publish_event(webhook_event.id)
- if not message_id:
- logger.warning(
- f"Webhook publish failed after commit: event_id={webhook_event.id}"
- )
- return
- logger.info(
- f"Webhook published after commit: event_id={webhook_event.id}, "
- f"message_id={message_id}"
- )
- except Exception as exc:
- logger.error(
- f"Failed to publish webhook after commit (event persisted): "
- f"event_id={webhook_event.id}, error={exc}"
- )
-
-
# Module-level singleton
_lifecycle_service: Optional[SyncJobLifecycleService] = None
diff --git a/packages/shared-python/shared/services/job_publication_sync.py b/packages/shared-python/shared/services/job_publication_sync.py
new file mode 100644
index 000000000..b6b2d69a1
--- /dev/null
+++ b/packages/shared-python/shared/services/job_publication_sync.py
@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select, update
+from sqlalchemy.orm import Session
+
+from shared.models.database.document import DocumentSection
+from shared.models.database.job import Job
+from shared.models.schemas.job_metadata import JobMetadataHelper
+from shared.services.redis.redis_sync_service import SyncRedisServiceFactory
+from shared.services.retrieval.publication_service import RetrievalPublicationService
+
+
+@dataclass(frozen=True)
+class JobPublicationOutcome:
+ published_document_state: dict[str, str] | None
+ cache_invalidation: dict[str, Any] | None
+
+
+class SyncJobPublicationFinalizer:
+ """Publish terminal parse results and invalidate retrieval cache after commit."""
+
+ def __init__(
+ self,
+ *,
+ retrieval_publication: RetrievalPublicationService | None = None,
+ ) -> None:
+ self._retrieval_publication = (
+ retrieval_publication or RetrievalPublicationService()
+ )
+
+ def publish_result(
+ self,
+ db: Session,
+ *,
+ job_id: str,
+ job_result_id: str,
+ chunks: list[dict[str, Any]],
+ section_summaries: dict[str, str] | None,
+ ) -> JobPublicationOutcome:
+ previous_document_scope = self._retrieval_publication.get_existing_document_scope(
+ db,
+ job_id=job_id,
+ )
+ published_document_state = self._retrieval_publication.publish_document_state(
+ db,
+ job_id=job_id,
+ job_result_id=job_result_id,
+ chunks=chunks,
+ )
+ if _should_publish_document_graph(published_document_state):
+ assert published_document_state is not None
+ if section_summaries:
+ self._backfill_section_summaries(
+ db,
+ document_id=published_document_state.get("document_id", ""),
+ job_result_id=job_result_id,
+ section_summaries=section_summaries,
+ )
+ self._retrieval_publication.publish_document_graph(
+ db,
+ job_id=job_id,
+ job_result_id=job_result_id,
+ )
+
+ cache_invalidation = self._build_cache_invalidation(
+ db,
+ job_id=job_id,
+ published_document_state=published_document_state,
+ previous_document_scope=previous_document_scope,
+ )
+ return JobPublicationOutcome(
+ published_document_state=published_document_state,
+ cache_invalidation=cache_invalidation,
+ )
+
+ def invalidate_cache_after_commit(
+ self,
+ cache_invalidation: dict[str, Any] | None,
+ ) -> None:
+ if not cache_invalidation:
+ return
+
+ try:
+ redis_service = SyncRedisServiceFactory.get_service()
+ user_id = cache_invalidation["user_id"]
+ seen: set[str] = set()
+ for namespace in cache_invalidation["namespaces"]:
+ if not namespace or namespace in seen:
+ continue
+ seen.add(namespace)
+ redis_service.incr(f"retrieval:version:{user_id}:{namespace}")
+ except Exception as exc:
+ logger.warning(
+ "Failed to invalidate retrieval cache after publication "
+ f"(ignored): job_id={cache_invalidation.get('job_id')}, error={exc}"
+ )
+
+ def _backfill_section_summaries(
+ self,
+ db: Session,
+ *,
+ document_id: str,
+ job_result_id: str,
+ section_summaries: dict[str, str],
+ ) -> None:
+ if not document_id or not section_summaries:
+ return
+
+ try:
+ for path, summary in section_summaries.items():
+ if not path or not summary:
+ continue
+ db.execute(
+ update(DocumentSection)
+ .where(DocumentSection.document_id == document_id)
+ .where(DocumentSection.job_result_id == job_result_id)
+ .where(DocumentSection.section_path == path)
+ .values(summary=summary)
+ )
+ db.flush()
+ logger.debug(
+ f"Backfilled section summaries: document_id={document_id}, "
+ f"count={len(section_summaries)}"
+ )
+ except Exception as exc:
+ logger.warning(f"Section summary backfill failed (non-fatal): {exc}")
+
+ def _build_cache_invalidation(
+ self,
+ db: Session,
+ *,
+ job_id: str,
+ published_document_state: dict[str, str] | None,
+ previous_document_scope: dict[str, str] | None,
+ ) -> dict[str, Any] | None:
+ job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none()
+ if not job:
+ return None
+
+ metadata = job.job_metadata or {}
+ namespaces = [
+ JobMetadataHelper.get_namespace(metadata, "default") or "default",
+ ]
+ if previous_document_scope and previous_document_scope.get("namespace"):
+ namespaces.append(previous_document_scope["namespace"])
+ if published_document_state and published_document_state.get("namespace"):
+ namespaces.append(published_document_state["namespace"])
+
+ return {"user_id": str(job.user_id), "namespaces": namespaces, "job_id": job_id}
+
+
+def _should_publish_document_graph(
+ published_document_state: dict[str, str] | None,
+) -> bool:
+ return published_document_state is not None and not published_document_state.get(
+ "skipped_all_duplicate"
+ )
diff --git a/packages/shared-python/shared/services/job_result_sync.py b/packages/shared-python/shared/services/job_result_sync.py
new file mode 100644
index 000000000..2c7b2f5aa
--- /dev/null
+++ b/packages/shared-python/shared/services/job_result_sync.py
@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import Any
+from uuid import uuid4
+
+from sqlalchemy import delete, select
+from sqlalchemy.orm import Session
+
+from shared.models.database.job_result import JobChunk, JobResult
+
+
+class SyncJobResultWriter:
+ """Persist terminal Job Result artifacts inside an existing transaction."""
+
+ def upsert_job_result(
+ self,
+ db: Session,
+ job_id: str,
+ delivery_mode: str,
+ *,
+ inline_payload: dict[str, Any] | None = None,
+ result_s3_key: str | None = None,
+ result_size: int | None = None,
+ ) -> JobResult:
+ result = db.execute(select(JobResult).where(JobResult.job_id == job_id))
+ existing = result.scalar_one_or_none()
+
+ if existing:
+ existing.delivery_mode = delivery_mode
+ existing.inline_payload = inline_payload
+ existing.result_s3_key = result_s3_key
+ existing.result_size = result_size
+ db.flush()
+ return existing
+
+ job_result = JobResult(
+ job_id=job_id,
+ delivery_mode=delivery_mode,
+ document_metadata={},
+ inline_payload=inline_payload,
+ result_s3_key=result_s3_key,
+ result_size=result_size,
+ )
+ db.add(job_result)
+ db.flush()
+ return job_result
+
+ def replace_chunks(
+ self,
+ db: Session,
+ job_result_id: str,
+ chunks: list[dict[str, Any]],
+ ) -> None:
+ db.execute(delete(JobChunk).where(JobChunk.job_result_id == job_result_id))
+
+ if not chunks:
+ db.flush()
+ return
+
+ chunk_models = []
+ for index, chunk in enumerate(chunks):
+ chunk_identifier = chunk.get("chunk_id") or str(uuid4())
+ metadata = chunk.get("metadata")
+ chunk_models.append(
+ JobChunk(
+ job_result_id=job_result_id,
+ chunk_id=chunk_identifier,
+ chunk_type=chunk.get("type", "paragraph"),
+ text=chunk.get("text"),
+ path=metadata.get("path") if isinstance(metadata, dict) else None,
+ chunk_metadata=metadata,
+ sort_order=chunk.get("order", index),
+ )
+ )
+ db.add_all(chunk_models)
+ db.flush()
diff --git a/packages/shared-python/shared/services/job_success_sync.py b/packages/shared-python/shared/services/job_success_sync.py
new file mode 100644
index 000000000..110e46df3
--- /dev/null
+++ b/packages/shared-python/shared/services/job_success_sync.py
@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from loguru import logger
+from sqlalchemy.orm import Session
+
+from shared.core.state_machine.service_sync import SyncStateMachineService
+from shared.models.database.webhook import WebhookEvent
+from shared.services.job_publication_sync import SyncJobPublicationFinalizer
+from shared.services.job_result_sync import SyncJobResultWriter
+from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox
+
+
+@dataclass(frozen=True)
+class JobSuccessFinalization:
+ response: dict[str, Any]
+ cache_invalidation: dict[str, Any] | None
+ webhook_event: WebhookEvent | None
+
+
+class SyncJobSuccessFinalizer:
+ """Finalize successful Jobs inside the lifecycle transaction."""
+
+ def __init__(
+ self,
+ *,
+ state_machine: SyncStateMachineService | None = None,
+ result_writer: SyncJobResultWriter | None = None,
+ publication_finalizer: SyncJobPublicationFinalizer | None = None,
+ webhook_outbox: SyncJobWebhookOutbox | None = None,
+ ) -> None:
+ self._state_machine = state_machine or SyncStateMachineService()
+ self._result_writer = result_writer or SyncJobResultWriter()
+ self._publication_finalizer = (
+ publication_finalizer or SyncJobPublicationFinalizer()
+ )
+ self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox()
+
+ def finalize(
+ self,
+ db: Session,
+ *,
+ job_id: str,
+ result_s3_key: str,
+ checksum: str,
+ zip_size: int,
+ chunks: list[dict[str, Any]],
+ stored_count: int,
+ delivery_mode: str,
+ section_summaries: dict[str, str] | None,
+ ) -> JobSuccessFinalization:
+ job_result = self._result_writer.upsert_job_result(
+ db,
+ job_id,
+ delivery_mode,
+ inline_payload={"checksum": checksum},
+ result_s3_key=result_s3_key,
+ result_size=zip_size,
+ )
+ self._result_writer.replace_chunks(db, job_result.id, chunks)
+ publication_outcome = self._publication_finalizer.publish_result(
+ db,
+ job_id=job_id,
+ job_result_id=job_result.id,
+ chunks=chunks,
+ section_summaries=section_summaries,
+ )
+
+ transition_ok = self._state_machine.mark_completed(
+ db,
+ job_id,
+ result_metadata={
+ "storage_completed": True,
+ "stored_count": stored_count,
+ "delivery_mode": delivery_mode,
+ },
+ )
+ if not transition_ok:
+ logger.error(f"Job {job_id} mark_completed transition failed")
+ return JobSuccessFinalization(
+ response={
+ "status": "failed",
+ "job_id": job_id,
+ "reason": "state_transition_failed",
+ },
+ cache_invalidation=None,
+ webhook_event=None,
+ )
+
+ webhook_event = self._webhook_outbox.create_event(
+ db,
+ job_id=job_id,
+ event_type="job.completed",
+ )
+ return JobSuccessFinalization(
+ response={
+ "status": "success",
+ "job_id": job_id,
+ "stored_count": stored_count,
+ },
+ cache_invalidation=publication_outcome.cache_invalidation,
+ webhook_event=webhook_event,
+ )
+
+ def run_post_commit_actions(
+ self,
+ finalization: JobSuccessFinalization,
+ ) -> None:
+ self._publication_finalizer.invalidate_cache_after_commit(
+ finalization.cache_invalidation,
+ )
+ self._webhook_outbox.enqueue_after_commit(finalization.webhook_event)
diff --git a/packages/shared-python/shared/services/job_webhook_outbox_sync.py b/packages/shared-python/shared/services/job_webhook_outbox_sync.py
new file mode 100644
index 000000000..d699a26d5
--- /dev/null
+++ b/packages/shared-python/shared/services/job_webhook_outbox_sync.py
@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from shared.models.database.job import Job
+from shared.models.database.webhook import WebhookEvent, WebhookEventStatus
+
+
+class SyncJobWebhookOutbox:
+ """Create webhook events in-transaction and publish them after commit."""
+
+ def create_event(
+ self,
+ db: Session,
+ *,
+ job_id: str,
+ event_type: str,
+ extra_payload: dict[str, Any] | None = None,
+ ) -> WebhookEvent | None:
+ result = db.execute(select(Job).where(Job.job_id == job_id))
+ job = result.scalar_one_or_none()
+
+ if not job:
+ logger.warning(f"Job not found for webhook check: {job_id}")
+ return None
+
+ webhook_url = getattr(job, "webhook_url", None)
+ if not job.webhook_enabled or not webhook_url:
+ return None
+
+ status = "completed" if event_type == "job.completed" else "failed"
+ timestamp_key = f"{status}_at"
+ payload: dict[str, Any] = {
+ "event": event_type,
+ "job_id": job_id,
+ "status": status,
+ timestamp_key: _utc_now_naive().isoformat(),
+ }
+ if extra_payload:
+ payload.update(extra_payload)
+
+ event = WebhookEvent(
+ job_id=job_id,
+ target_url=webhook_url,
+ payload=payload,
+ status=WebhookEventStatus.PENDING,
+ attempts=0,
+ )
+ db.add(event)
+ db.flush()
+ logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}")
+ return event
+
+ def enqueue_after_commit(self, webhook_event: WebhookEvent | None) -> None:
+ if not webhook_event:
+ return
+
+ try:
+ from shared.services.webhook.qstash_publisher import (
+ get_qstash_webhook_publisher,
+ )
+
+ publisher = get_qstash_webhook_publisher()
+ message_id = publisher.publish_event(webhook_event.id)
+ if not message_id:
+ logger.warning(
+ f"Webhook publish failed after commit: event_id={webhook_event.id}"
+ )
+ return
+ logger.info(
+ f"Webhook published after commit: event_id={webhook_event.id}, "
+ f"message_id={message_id}"
+ )
+ except Exception as exc:
+ logger.error(
+ "Failed to publish webhook after commit (event persisted): "
+ f"event_id={webhook_event.id}, error={exc}"
+ )
+
+
+def _utc_now_naive() -> datetime:
+ return datetime.now(timezone.utc).replace(tzinfo=None)
diff --git a/packages/shared-python/shared/services/jobs/result_delivery.py b/packages/shared-python/shared/services/jobs/result_delivery.py
new file mode 100644
index 000000000..0a5c9b0d6
--- /dev/null
+++ b/packages/shared-python/shared/services/jobs/result_delivery.py
@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Any
+
+from shared.services.storage.job_file_storage import JobFileStorage
+from shared.utils.utc_now import utc_now_naive
+
+
+@dataclass(frozen=True)
+class JobResultDelivery:
+ result: dict[str, Any] | None
+ result_url: str | None
+ result_url_expires_at: datetime | None
+
+
+class JobResultDeliveryResolver:
+ """Resolve the public delivery fields exposed for a terminal Job Result."""
+
+ def __init__(self, *, storage: JobFileStorage | None = None) -> None:
+ self._storage = storage or JobFileStorage()
+
+ def resolve(
+ self,
+ job_result: Any | None,
+ *,
+ default_expires_at: datetime | None = None,
+ ) -> JobResultDelivery:
+ result = None
+ result_url = None
+ result_url_expires_at = default_expires_at
+
+ if not job_result:
+ return JobResultDelivery(
+ result=result,
+ result_url=result_url,
+ result_url_expires_at=result_url_expires_at,
+ )
+
+ inline_payload = getattr(job_result, "inline_payload", None)
+ if inline_payload:
+ result = inline_payload
+
+ result_s3_key = getattr(job_result, "result_s3_key", None)
+ if result_s3_key:
+ url_info = self._storage.generate_download_url(
+ result_s3_key,
+ bucket=self._storage.results_bucket,
+ )
+ result_url = url_info["download_url"]
+ expires_in = int(url_info.get("expires_in", 3600))
+ result_url_expires_at = utc_now_naive() + timedelta(seconds=expires_in)
+
+ return JobResultDelivery(
+ result=result,
+ result_url=result_url,
+ result_url_expires_at=result_url_expires_at,
+ )
+
+ def enrich_payload(
+ self,
+ payload: dict[str, Any],
+ *,
+ job_result: Any | None,
+ ) -> dict[str, Any]:
+ if payload.get("event") != "job.completed":
+ return payload
+
+ delivery = self.resolve(job_result)
+ enriched = dict(payload)
+ if delivery.result_url:
+ enriched["result_url"] = delivery.result_url
+ if delivery.result:
+ enriched["result"] = delivery.result
+ return enriched
diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py
index 832e81eb6..4c1a56a95 100644
--- a/packages/shared-python/shared/services/retrieval/__init__.py
+++ b/packages/shared-python/shared/services/retrieval/__init__.py
@@ -6,7 +6,8 @@
invalidate_retrieval_cache_namespaces,
set_cached_retrieval_query_result,
)
-from .graph_service import DocumentGraphService, GraphQueryService, GraphScope
+from .graph_query_service import GraphQueryService
+from .graph_service import DocumentGraphService, GraphScope
from .hit_stats_service import record_retrieval_hits
from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn
diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py
deleted file mode 100644
index d5b1bed78..000000000
--- a/packages/shared-python/shared/services/retrieval/agent_navigate.py
+++ /dev/null
@@ -1,1137 +0,0 @@
-"""Shared helpers for agentic KG document routing and scope navigation."""
-from __future__ import annotations
-
-import json
-import re
-from typing import Any, Sequence, TYPE_CHECKING
-
-if TYPE_CHECKING:
- from shared.services.retrieval.agentic.types import DocTreeNode
-
-from loguru import logger
-from sqlalchemy import func, select, or_
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphNode, GraphEdge
-from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path
-from shared.utils.text_utils import tokenize_for_retrieval
-
-_MAX_OVERVIEW_FILES = 50
-
-_FILE_SELECT_PROMPT = """\
-You are a document routing assistant.
-
-{budget_block}
-Below is a knowledge base overview showing all available documents,
-their navigation summaries, chunk counts, and media counts.
-
-=== Knowledge Base Overview ===
-{overview}
-=== End Overview ===
-
-User query: {query}
-{revision_context}
-Based on the query, select documents that may contain relevant information.
-If NO document in the knowledge base is relevant to the query, return an EMPTY array [].
-Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"]
-Do not include any explanation.
-"""
-
-
-
-_DISCOVERY_SELECT_PROMPT = """\
-You are a document navigation assistant.
-
-Document: "{doc_name}"
-
-{budget_block}
-After navigating the document's section tree, the following section paths
-were additionally discovered via keyword and semantic search.
-They may contain relevant evidence not found through hierarchical navigation.
-
-=== Discovery Candidates ===
-{items}
-=== End Discovery Candidates ===
-
-User query: {query}
-{revision_context}
-Select section paths whose content is needed to answer the query.
-If none are relevant, return an EMPTY list [].
-
-Return ONLY a JSON object:
-{{"selections": [{{"path": "...", "confidence": }}, ...]}}
-Do not include any explanation.
-"""
-
-
-_ACTION_PROMPT = """\
-You are a document navigation agent.
-
-Document: "{doc_name}" (id: {doc_id})
-
-{budget_block}
-{scope_header}
-Below is the document's section tree.
-Sections tagged [SELECT] are within the current scope and may be selected.
-Other sections are shown as structural context only (not selectable).
-Nodes marked [Leaf] have no further sub-sections.
-
-=== Section Tree ===
-{items_overview}
-=== End Section Tree ===
-
-User query: {query}
-
-=== Available Actions ===
-
-Choose ONE action:
-
-NAVIGATE — Drill into selected sections for detailed content.
- Consider this when the query targets specific topics and you need deeper text evidence.
- Select one or more [SELECT] sections.
-
-STOP — Current scope evidence is sufficient. No further drill-down.
- Consider this when:
- - The query asks for an outline, overview, or summary
- - The query is broad/global, the tree section can fulfill it without drilling into individual sections.
- - You have already collected enough evidence at this level.
-
-{tools_block}
-
-When action is NAVIGATE, provide selections:
-- You may ONLY select sections marked with [SELECT].
-
-When action is STOP, selections must be empty.
-
-Return ONLY a JSON object:
-{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}}
-or
-{{"action": "STOP", "tools": [...], "selections": []}}
-Do not include any explanation.
-"""
-
-
-def _parse_action_response(text: str) -> dict:
- """Parse the unified action response from LLM.
-
- Returns dict with keys:
- action: 'NAVIGATE' | 'STOP'
- tools: list[str] (subset of FIND_IMAGES, FIND_TABLES)
- selections: list[dict] (each has 'path' and optional 'confidence')
-
- When action is STOP, selections are forced to empty.
- """
- import json as _json
- import re as _re
-
- text = text.strip()
- _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'}
- default = {'action': 'NAVIGATE', 'tools': [], 'selections': []}
-
- def _extract(data: dict) -> dict:
- action = str(data.get('action', 'NAVIGATE')).strip().upper()
- if action not in ('NAVIGATE', 'STOP'):
- action = 'NAVIGATE'
-
- tools_val = data.get('tools') or []
- if isinstance(tools_val, list):
- tools = [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS]
- else:
- tools = []
-
- # STOP → no selections allowed
- if action == 'STOP':
- return {'action': action, 'tools': tools, 'selections': []}
-
- selections_val = data.get('selections') or []
- selections = []
- if isinstance(selections_val, list):
- for s in selections_val:
- if isinstance(s, dict) and s.get('path'):
- conf = _normalize_confidence(s.get('confidence', 0.7))
- selections.append({'path': str(s['path']), 'confidence': conf or 0.7})
-
- return {'action': action, 'tools': tools, 'selections': selections}
-
- # Try JSON parse
- try:
- data = _json.loads(text)
- if isinstance(data, dict):
- return _extract(data)
- except (ValueError, _json.JSONDecodeError):
- pass
-
- # Try extracting JSON from markdown fences
- fence_match = _re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, _re.DOTALL)
- if fence_match:
- try:
- data = _json.loads(fence_match.group(1).strip())
- if isinstance(data, dict):
- return _extract(data)
- except (ValueError, _json.JSONDecodeError):
- pass
-
- # Try finding a JSON object anywhere
- brace_match = _re.search(r'\{.*\}', text, _re.DOTALL)
- if brace_match:
- try:
- data = _json.loads(brace_match.group())
- if isinstance(data, dict):
- return _extract(data)
- except (ValueError, _json.JSONDecodeError):
- pass
-
- return default
-
-
-def _format_budget_block(snapshot: dict | None) -> str:
- if not snapshot:
- return ""
- planning = snapshot.get("planning") or {}
- context = snapshot.get("context") or {}
- return (
- "=== Resource Status ===\n"
- f"Planning Budget: {planning.get('status', 'HEALTHY')} "
- f"({planning.get('used_pct', 0)}% used)\n"
- f"Context Budget: {context.get('status', 'HEALTHY')} "
- f"({context.get('used_pct', 0)}% used)\n"
- f"KG Coverage: {snapshot.get('explored_chunks', 0)}/"
- f"{snapshot.get('total_chunks', 0)} chunks explored\n"
- f"Docs Explored: {snapshot.get('explored_docs', 0)}/"
- f"{snapshot.get('total_docs', 0)}\n"
- "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. "
- "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n"
- "=== End Resource Status ===\n"
- )
-
-
-def _extract_json_array_payload(text: str) -> list[Any]:
- """Best-effort extraction of a JSON array payload from LLM response text."""
- text = text.strip()
- try:
- result = json.loads(text)
- if isinstance(result, list):
- return result
- except (json.JSONDecodeError, ValueError):
- pass
- match = re.search(r'\[.*?\]', text, re.DOTALL)
- if match:
- try:
- result = json.loads(match.group())
- if isinstance(result, list):
- return result
- except (json.JSONDecodeError, ValueError):
- pass
- return []
-
-
-def _parse_json_array(text: str) -> list[str]:
- """Best-effort extraction of a JSON array of strings from LLM response text."""
- result = _extract_json_array_payload(text)
- return [str(x) for x in result]
-
-
-def _normalize_confidence(value: Any) -> float | None:
- if value is None:
- return None
- if isinstance(value, str):
- value = value.strip().rstrip('%')
- try:
- parsed = float(value)
- except (TypeError, ValueError):
- return None
- if parsed > 1.0:
- parsed = parsed / 100.0
- return max(0.0, min(parsed, 1.0))
-
-
-async def _build_knowledge_map_overview(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
-) -> tuple[str, dict[str, str]]:
- """Build a file-level knowledge map overview for LLM file selection.
-
- Returns (overview_text, doc_id_to_name) where doc_id_to_name maps
- document_id -> source_file_name for validation after LLM response.
- """
- doc_stmt = (
- select(Document)
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(Document.current_job_result_id.is_not(None))
- .order_by(Document.updated_at.desc())
- .limit(_MAX_OVERVIEW_FILES)
- )
- doc_result = await db.execute(doc_stmt)
- documents = list(doc_result.scalars())
-
- if not documents:
- return '(empty)', {}
-
- doc_ids = [d.document_id for d in documents]
- doc_id_to_name: dict[str, str] = {
- d.document_id: (d.source_file_name or d.document_id)
- for d in documents
- }
-
- chunk_stats_stmt = (
- select(
- DocumentChunk.document_id,
- func.count(DocumentChunk.id).label('chunk_count'),
- func.count(func.nullif(DocumentChunk.chunk_type, 'text')).label('media_count'),
- )
- .join(Document, (Document.document_id == DocumentChunk.document_id) & (Document.current_job_result_id == DocumentChunk.job_result_id))
- .where(DocumentChunk.document_id.in_(doc_ids))
- .group_by(DocumentChunk.document_id)
- )
- chunk_stats_result = await db.execute(chunk_stats_stmt)
- chunk_stats: dict[str, dict[str, int]] = {}
- for row in chunk_stats_result.all():
- chunk_stats[row[0]] = {'total': row[1], 'media': row[2]}
-
- graph_summary_stmt = (
- select(GraphNode.owner_document_id, GraphNode.properties)
- .where(GraphNode.owner_document_id.in_(doc_ids))
- .where(GraphNode.node_kind == 'document')
- )
- graph_summary_result = await db.execute(graph_summary_stmt)
- doc_top_summaries: dict[str, str] = {}
- for did, properties in graph_summary_result.all():
- if not isinstance(properties, dict):
- continue
- top_summary = str(properties.get('top_summary') or '').strip()
- if top_summary:
- doc_top_summaries[did] = top_summary
-
- lines: list[str] = []
- for doc in documents:
- did = doc.document_id
- name = doc_id_to_name[did]
- stats = chunk_stats.get(did, {'total': 0, 'media': 0})
- top_summary = doc_top_summaries.get(did, '')
-
- line = f'- [{did}] {name} chunks={stats["total"]}'
- if stats['media'] > 0:
- line += f' media={stats["media"]}'
- if top_summary:
- line += f'\n top_summary:\n{_indent_block(top_summary, 4)}'
- lines.append(line)
-
- return '\n'.join(lines), doc_id_to_name
-
-
-def _indent_block(text: str, spaces: int) -> str:
- prefix = ' ' * spaces
- return '\n'.join(f'{prefix}{line}' for line in str(text or '').splitlines())
-
-
-def _format_items_for_llm(
- items: list[dict],
- max_chars: int = 20000,
-) -> tuple[str, bool]:
- """Format items with ▸ └ [Leaf] hierarchy for scope navigation.
-
- Supports arbitrary depth levels via absolute ``level`` field.
- Items with ``show_summary=False`` render title only (structural context).
- ``[LN]`` tags indicate the absolute document depth of each section.
- ``[Leaf]`` tags indicate bottom-level sections with no further children.
- Summaries are included when within budget, dropped on overflow.
-
- Returns (text, overflowed).
- """
- from shared.utils.text_utils import truncate_content_preview
-
- if not items:
- return '(no items available)', False
-
- SUMMARY_HEAD_TOKENS = 80
-
- def _render_item(item: dict, include_summary: bool) -> str:
- level = item.get('level', 1)
- show = item.get('show_summary', True)
- is_leaf = item.get('is_leaf', False)
- leaf_tag = ' [Leaf]' if is_leaf else ''
- path = item.get('path', '')
- summary = item.get('summary') or ''
-
- # Build chunk count tags (only for current-scope items)
- counts_str = ''
- if show:
- count_parts: list[str] = []
- chunk_count = item.get('chunk_count', 0)
- if chunk_count > 0:
- count_parts.append(f'text={chunk_count}')
- image_count = item.get('image_count', 0)
- if image_count > 0:
- count_parts.append(f'image={image_count}')
- table_count = item.get('table_count', 0)
- if table_count > 0:
- count_parts.append(f'table={table_count}')
- counts_str = f' [{" ".join(count_parts)}]' if count_parts else ''
-
- indent = " " * (level - 1)
- prefix = '▸' if level == 1 else '└'
- level_tag = f'[L{level}]'
- select_tag = '[SELECT] ' if item.get('selectable', False) else ''
-
- lines: list[str] = []
- lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}')
-
- if include_summary and show and summary:
- sub_indent = " " * level
- clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0)
- lines.append(f'{sub_indent}{clipped}')
-
- return '\n'.join(lines)
-
- # Try full render (with summaries for show_summary=True items)
- full_lines = [_render_item(item, include_summary=True) for item in items]
- full_text = '\n'.join(full_lines)
- if len(full_text) <= max_chars:
- return full_text, False
-
- # Overflow: render without summaries
- slim_lines = [_render_item(item, include_summary=False) for item in items]
- slim_text = '\n'.join(slim_lines)
- return slim_text[:max_chars], True
-
-
-# ------------------------------------------------------------------
-# GREP document discovery (aligned with KB do_discover_files)
-# ------------------------------------------------------------------
-
-async def _grep_discover_document_ids(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- query: str,
- exclude_document_ids: Sequence[str] = (),
- limit: int = 10,
-) -> list[str]:
- """GREP discovery: search term_search_text for query terms, return parent document_ids.
-
- Aligned with KB's do_discover_files(): if a chunk's term_search_text
- contains query terms, its parent document is included in the KG scope.
- """
- units = tokenize_for_retrieval(query, dedupe=True)
- logger.info(f' GREP tokenized units (cap 8): {units[:8]} (total={len(units)})')
- if not units:
- return []
-
- # Build OR conditions for ILIKE matching
- conditions = []
- params: dict[str, str] = {
- 'user_id': user_id,
- 'namespace': namespace,
- }
- for i, unit in enumerate(units[:8]): # cap at 8 terms to avoid huge queries
- param_name = f'unit_{i}'
- params[param_name] = f'%{unit}%'
- conditions.append(DocumentChunk.term_search_text.ilike(f'%{unit}%'))
-
- if not conditions:
- return []
-
- stmt = (
- select(Document.document_id)
- .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id)
- & (DocumentChunk.job_result_id == Document.current_job_result_id))
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(DocumentChunk.term_search_text.is_not(None))
- .where(or_(*conditions))
- .distinct()
- .limit(limit)
- )
- if exclude_document_ids:
- stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
-
- result = await db.execute(stmt)
- return [row[0] for row in result.all()]
-
-
-# ------------------------------------------------------------------
-# Edge expansion (aligned with KB KGIndex.neighbors)
-# ------------------------------------------------------------------
-
-async def _expand_by_edges(
- db: AsyncSession,
- *,
- document_ids: list[str],
- user_id: str,
- namespace: str,
- hops: int = 1,
-) -> list[str]:
- """Expand document set by following GraphEdge relationships.
-
- Aligned with KB's KGIndex.neighbors(): traverse edges to include
- related documents. Only queries document-level nodes (no section nodes).
- No weight filtering — edges already passed threshold during publication.
- """
- if not document_ids:
- return document_ids
-
- current = set(document_ids)
-
- for hop_idx in range(hops):
- # Find document-level graph nodes for current document set
- doc_node_ids = [f"doc:{did}" for did in current]
- node_stmt = (
- select(GraphNode.node_id, GraphNode.owner_document_id)
- .where(GraphNode.user_id == user_id)
- .where(GraphNode.namespace == namespace)
- .where(GraphNode.node_kind == 'document')
- .where(GraphNode.node_id.in_(doc_node_ids))
- )
- node_result = await db.execute(node_stmt)
- node_rows = node_result.all()
- logger.info(f' edge_expand hop={hop_idx}: doc_nodes_found={len(node_rows)} (of {len(doc_node_ids)} requested)')
-
- if not node_rows:
- break
-
- node_ids = {row[0] for row in node_rows}
-
- # Follow edges from/to these document nodes
- edge_stmt = (
- select(GraphEdge.source_node_id, GraphEdge.target_node_id)
- .where(GraphEdge.user_id == user_id)
- .where(GraphEdge.namespace == namespace)
- .where(or_(
- GraphEdge.source_node_id.in_(list(node_ids)),
- GraphEdge.target_node_id.in_(list(node_ids)),
- ))
- )
- edge_result = await db.execute(edge_stmt)
- edge_rows = edge_result.all()
-
- neighbor_node_ids: set[str] = set()
- for src, tgt in edge_rows:
- if src in node_ids:
- neighbor_node_ids.add(tgt)
- if tgt in node_ids:
- neighbor_node_ids.add(src)
- logger.info(f' edge_expand hop={hop_idx}: edges_traversed={len(edge_rows)} neighbor_nodes={len(neighbor_node_ids)}')
-
- if not neighbor_node_ids:
- break
-
- # Resolve neighbor nodes to document_ids
- neighbor_doc_stmt = (
- select(GraphNode.owner_document_id)
- .where(GraphNode.node_id.in_(list(neighbor_node_ids)))
- .where(GraphNode.node_kind == 'document')
- )
- neighbor_doc_result = await db.execute(neighbor_doc_stmt)
- for (doc_id,) in neighbor_doc_result.all():
- current.add(doc_id)
-
- # Preserve original order, append new ones at end
- ordered = list(document_ids)
- for doc_id in current:
- if doc_id not in document_ids:
- ordered.append(doc_id)
- return ordered
-
-
-# ---------------------------------------------------------------------------
-# Unified scope navigation: load child sections (2-level)
-# ---------------------------------------------------------------------------
-
-async def _load_child_sections(
- db: AsyncSession,
- document_id: str,
- job_result_id: str,
- scope_path: str | list[str] | None = None,
- exclude_paths: set[str] | None = None,
-) -> list[dict]:
- """Load the Continuous Context Tree for *scope_path*.
-
- Returns a flat list sorted by document order, each item:
- {path, title, summary, chunk_count, image_count, table_count,
- level, show_summary, is_leaf}
-
- scope_path can be:
- - None: root scope, all items are selectable (2 depth bands).
- - str: single scope, descendants are selectable.
- - list[str]: multi-scope, descendants of ALL paths are selectable
- simultaneously — used when the LLM selected multiple drill-down
- paths in the previous step.
-
- - level: absolute depth in the document (1-based)
- - show_summary: controls whether _format_items_for_llm renders summary
- - exclude_paths: paths already hydrated; skipped from selectable items
- """
- # ── Fetch all sections for this document revision ────────────────────
- stmt = (
- select(
- DocumentSection.section_id,
- DocumentSection.section_title,
- DocumentSection.section_path,
- DocumentSection.summary,
- DocumentSection.sort_order,
- )
- .where(DocumentSection.document_id == document_id)
- .where(DocumentSection.job_result_id == job_result_id)
- .order_by(DocumentSection.sort_order)
- )
- section_rows = (await db.execute(stmt)).all()
- if not section_rows:
- return []
-
- # ── Normalize scope(s) ───────────────────────────────────────────────
- # Multi-scope: list of paths to expand simultaneously
- if isinstance(scope_path, list):
- scope_list = [normalize_section_path(p) for p in scope_path]
- elif scope_path:
- scope_list = [normalize_section_path(scope_path)]
- else:
- scope_list = [] # root
-
- # For logging, derive representative scope info
- scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0
-
- logger.debug(
- f' _load_child_sections: scopes={scope_list or ["root"]} '
- f'scope_depth={scope_depth} exclude_paths={_excl if (_excl := exclude_paths or set()) else "none"} '
- f'total_sections={len(section_rows)}'
- )
-
- # Build full section metadata index
- all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id, parts, depth}
- for section_id, title, path, summary, sort_order in section_rows:
- if not path:
- continue
- path = normalize_section_path(path)
- parts = split_section_path(path)
- all_sections[path] = {
- 'title': title or parts[-1] if parts else path,
- 'summary': summary or '',
- 'sort_order': int(sort_order or 0),
- 'section_id': section_id,
- 'parts': parts,
- 'depth': len(parts),
- }
-
- # ── Build the set of ancestor prefixes for pruning ────────────────────
- # For multi-scope, union all ancestor prefixes from all scope paths
- ancestor_prefixes: set[str] = set()
- for sp in scope_list:
- sp_parts = split_section_path(sp)
- for i in range(1, len(sp_parts) + 1):
- ancestor_prefixes.add(' / '.join(sp_parts[:i]))
-
- # ── Classify each section ────────────────────────────────────────────
- _excl = exclude_paths or set()
- items_by_path: dict[str, dict] = {}
- # Per-scope depth bands: track child depths separately per scope
- per_scope_child_depths: dict[str, set[int]] = {sp: set() for sp in scope_list} if scope_list else {}
- root_child_depths: set[int] = set() # used when scope_list is empty (root)
-
- def _make_item(path: str, meta: dict, show_summary: bool) -> dict:
- return {
- 'path': path,
- 'title': meta['title'],
- 'summary': meta['summary'],
- 'level': meta['depth'],
- 'sort_order': meta['sort_order'],
- 'chunk_count': 0,
- 'image_count': 0,
- 'table_count': 0,
- 'section_id': meta['section_id'],
- 'show_summary': show_summary,
- }
-
- def _is_excluded(path: str) -> bool:
- return bool(_excl and any(
- path == ep or path.startswith(ep + ' / ') for ep in _excl
- ))
-
- for path, meta in all_sections.items():
- parts = meta['parts']
- depth = meta['depth']
-
- if not scope_list:
- # Root scope: everything is a potential child
- if depth < 1 or _is_excluded(path):
- continue
- root_child_depths.add(depth)
- items_by_path[path] = _make_item(path, meta, show_summary=True)
- continue
-
- # --- Non-root scope(s) ---
- # Check if this path is a descendant of ANY scope in scope_list
- matched_scope: str | None = None
- for sp in scope_list:
- sp_parts = split_section_path(sp)
- sp_depth = len(sp_parts)
- if depth > sp_depth and parts[:sp_depth] == sp_parts:
- matched_scope = sp
- break
-
- if matched_scope:
- # Category 2: descendant of a scope path → selectable
- if _is_excluded(path):
- continue
- per_scope_child_depths[matched_scope].add(depth)
- items_by_path[path] = _make_item(path, meta, show_summary=True)
- continue
-
- # Category 1: structural context (ancestors of scope paths only)
- # Only show nodes that are on the ancestor chain of a scope path.
- # Non-scope siblings (e.g. 法律声明, 前言 when navigating into
- # chapters 一~六) are pruned to reduce token waste and prevent
- # summary overflow in _format_items_for_llm.
- max_scope_depth = max(len(split_section_path(sp)) for sp in scope_list)
- if depth <= max_scope_depth:
- if depth == 1:
- if path in ancestor_prefixes:
- items_by_path.setdefault(path, _make_item(path, meta, show_summary=False))
- else:
- parent_prefix = ' / '.join(parts[:-1])
- if parent_prefix in ancestor_prefixes:
- items_by_path.setdefault(path, _make_item(path, meta, show_summary=False))
- continue
-
- # Category 3: pruned
-
- if not items_by_path:
- return []
-
- # ── Limit children to 2 depth bands (relative to each scope) ────────
- allowed_set: set[int] = set()
- if scope_list:
- for sp, depths in per_scope_child_depths.items():
- if depths:
- allowed_set.update(sorted(depths)[:2])
- else:
- if root_child_depths:
- allowed_set.update(sorted(root_child_depths)[:2])
-
- if allowed_set:
- to_remove = [
- path for path, item in items_by_path.items()
- if item['show_summary'] and item['level'] not in allowed_set
- ]
- for path in to_remove:
- del items_by_path[path]
-
- if not items_by_path:
- return []
-
- # ── Count chunks per section (text / image / table) ──────────────────
- # Only count for show_summary=True items (current scope children)
- scope_item_sids = {item['section_id'] for item in items_by_path.values() if item['show_summary']}
- # Also need all section_ids for upward aggregation
- all_section_ids = [meta['section_id'] for meta in all_sections.values()]
- if all_section_ids and scope_item_sids:
- from sqlalchemy import case, literal_column
- chunk_stmt = (
- select(
- DocumentChunk.section_id,
- func.count(
- case(
- (DocumentChunk.chunk_type.notin_(['image', 'table']), literal_column('1')),
- )
- ).label('text_count'),
- func.count(
- case(
- (DocumentChunk.chunk_type == 'image', literal_column('1')),
- )
- ).label('image_count'),
- func.count(
- case(
- (DocumentChunk.chunk_type == 'table', literal_column('1')),
- )
- ).label('table_count'),
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.section_id.in_(all_section_ids))
- .group_by(DocumentChunk.section_id)
- )
- chunk_rows = (await db.execute(chunk_stmt)).all()
- section_id_counts: dict[str, tuple[int, int, int]] = {
- sid: (int(tc), int(ic), int(tbc)) for sid, tc, ic, tbc in chunk_rows
- }
- else:
- section_id_counts = {}
-
- # Build section_id → path mapping for aggregation
- sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()}
-
- # Aggregate chunk counts upward: each show_summary item gets counts from itself + descendants
- # Phase 1: Direct section assignment — counts from chunks directly under each section
- for sid, (text_c, img_c, tbl_c) in section_id_counts.items():
- chunk_path = sid_to_path.get(sid, '')
- if not chunk_path:
- continue
-
- for item_path, item in items_by_path.items():
- if not item['show_summary']:
- continue
- if chunk_path == item_path or chunk_path.startswith(item_path + ' / '):
- item['chunk_count'] += text_c
- item['image_count'] += img_c
- item['table_count'] += tbl_c
-
- # Phase 2: connect_to reference tracing — Root-level standalone assets
- # Images/tables often live in the Root section but are referenced via connect_to
- # from text chunks in deeper sections. Trace these references to attribute
- # assets to the sections that actually use them.
- #
- # Algorithm: for each show_summary item, find all text chunks under its subtree,
- # collect their connect_to targets, and count how many are image/table chunks.
- scope_items_with_zero_assets = [
- item for item in items_by_path.values()
- if item['show_summary'] and item['image_count'] == 0 and item['table_count'] == 0
- ]
- if scope_items_with_zero_assets:
- # Load connect_to metadata for text chunks under all scope sections
- scope_section_ids = {item['section_id'] for item in items_by_path.values() if item.get('section_id')}
- if scope_section_ids:
- from sqlalchemy import literal_column
- connect_stmt = (
- select(
- DocumentChunk.section_id,
- DocumentChunk.chunk_metadata,
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.section_id.in_(list(scope_section_ids)))
- .where(DocumentChunk.chunk_type == 'text')
- )
- connect_result = (await db.execute(connect_stmt)).all()
-
- # Map section_id → set of connected target chunk_ids
- section_target_ids: dict[str, set[str]] = {}
- for sec_id, metadata in connect_result:
- if not isinstance(metadata, dict):
- continue
- for conn in metadata.get('connect_to') or []:
- target_id = conn.get('target', '')
- if target_id:
- section_target_ids.setdefault(sec_id, set()).add(target_id)
-
- if section_target_ids:
- # Collect all target chunk_ids and look up their types
- all_target_ids = set()
- for tids in section_target_ids.values():
- all_target_ids.update(tids)
-
- target_type_stmt = (
- select(
- DocumentChunk.chunk_id,
- DocumentChunk.chunk_type,
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.chunk_id.in_(list(all_target_ids)))
- .where(DocumentChunk.chunk_type.in_(['image', 'table']))
- )
- target_type_result = (await db.execute(target_type_stmt)).all()
- target_types: dict[str, str] = {cid: ctype for cid, ctype in target_type_result}
-
- # Aggregate connected asset counts per section path → upward to items
- for sec_id, target_ids in section_target_ids.items():
- ref_path = sid_to_path.get(sec_id, '')
- if not ref_path:
- continue
- ref_img = sum(1 for tid in target_ids if target_types.get(tid) == 'image')
- ref_tbl = sum(1 for tid in target_ids if target_types.get(tid) == 'table')
- if ref_img == 0 and ref_tbl == 0:
- continue
- for item_path, item in items_by_path.items():
- if not item['show_summary']:
- continue
- if ref_path == item_path or ref_path.startswith(item_path + ' / '):
- item['image_count'] += ref_img
- item['table_count'] += ref_tbl
-
- # ── Sort by native document order ─────────────────────────────────────
- sorted_items = sorted(items_by_path.values(), key=lambda x: x['sort_order'])
- # Clean up internal fields
- for item in sorted_items:
- item.pop('sort_order', None)
- item.pop('section_id', None)
-
- # ── Detect leaf status ────────────────────────────────────────────────
- # A section is a leaf if no other section in the database for this
- # document has a path that descends from it.
- all_section_paths = set(all_sections.keys())
- for item in sorted_items:
- item_path = item['path']
- has_descendants = any(
- p != item_path and p.startswith(item_path + ' / ')
- for p in all_section_paths
- )
- item['is_leaf'] = not has_descendants
-
- # ── Assign selectability ──────────────────────────────────────────────
- # Rule: in the 2-band window, only the DEEPER band is selectable.
- # Leaf nodes at the shallower band are still selectable (no children
- # to drill into). Structural context (show_summary=False) is never
- # selectable.
- if allowed_set:
- shallowest_band = min(allowed_set)
- for item in sorted_items:
- if not item.get('show_summary', True):
- # Structural context → never selectable
- item['selectable'] = False
- elif item['level'] == shallowest_band and not item.get('is_leaf', False):
- # Shallowest band, non-leaf → grouping header, not selectable
- item['selectable'] = False
- else:
- item['selectable'] = True
- else:
- for item in sorted_items:
- item['selectable'] = item.get('show_summary', True)
-
- return sorted_items
-
-
-# ---------------------------------------------------------------------------
-
-
-# ---------------------------------------------------------------------------
-# Unified document tree rendering (DocTreeNode → single coherent hierarchy)
-# ---------------------------------------------------------------------------
-
-def _render_leaf_chunks(
- parts: list[str],
- chunks: list[dict[str, Any]],
- indent: str,
- asset_lookup: dict[str, str] | None = None,
-) -> None:
- """Render hydrated leaf chunks inline with table/image inlining and dedup.
-
- Uses ``connect_to`` metadata to resolve asset references — the same
- pattern as ``assemble_retrieval_results``:
- - **Tables**: inline HTML content at the ``ref`` placeholder
- - **Images**: inline the ``file_path`` (S3-compatible URL) at the
- placeholder for multimodal LLMs
-
- Connected target chunks (images/tables) are expected to already be
- present in ``chunks`` via ``hydrate_connected_target_rows``.
-
- Phase 2: After rendering all text chunks, standalone image/table
- chunks that were NOT inlined via connect_to are rendered separately.
- This handles cases where assets exist at root/section level without
- a parent text chunk referencing them.
- """
- chunk_by_id: dict[str, dict] = {
- c.get('chunk_id', ''): c for c in chunks if c.get('chunk_id')
- }
- rendered_ids: set[str] = set()
-
- # Phase 1: Render text chunks with inline asset resolution
- for chunk in chunks:
- cid = chunk.get('chunk_id', '')
- if cid and cid in rendered_ids:
- continue
-
- chunk_type = (chunk.get('chunk_type') or chunk.get('type') or 'text').strip().lower()
-
- # Skip standalone image/table chunks — they'll be rendered in Phase 2
- # if not inlined via connect_to from a parent text chunk.
- # NOTE: do NOT add to rendered_ids here — Phase 2 needs to see them.
- if chunk_type in ('image', 'table'):
- continue
-
- if cid:
- rendered_ids.add(cid)
-
- content = str(chunk.get('content', '')).strip()
-
- # Resolve connected assets via connect_to metadata
- for conn in (chunk.get('chunk_metadata') or {}).get('connect_to') or []:
- target = chunk_by_id.get(conn.get('target', ''))
- if not target:
- continue
- target_cid = target.get('chunk_id', '')
- target_type = (target.get('chunk_type') or target.get('type') or '').strip().lower()
- ref_str = conn.get('ref', '')
- if not ref_str or ref_str not in content:
- continue
-
- if target_cid:
- rendered_ids.add(target_cid)
-
- if target_type == 'table':
- table_html = str(target.get('content', '')).strip()
- content = content.replace(ref_str, f'\n[表格内容]\n{table_html}\n')
- elif target_type == 'image':
- file_path = target.get('file_path') or ''
- img_desc = str(target.get('content', '')).strip()
- # Strip self-reference from image description
- if ref_str in img_desc:
- img_desc = img_desc.replace(ref_str, '').strip()
- # Use pre-generated asset URL if available, fall back to file_path
- asset_url = (asset_lookup or {}).get(target_cid, '') if target_cid else ''
- display_ref = asset_url or file_path
- if display_ref:
- content = content.replace(ref_str, f'\n[图片: {display_ref}]\n{img_desc}\n')
- elif img_desc:
- content = content.replace(ref_str, f'\n[图片描述]\n{img_desc}\n')
-
- for line in content.split('\n'):
- if line.strip():
- parts.append(f'{indent}┈ {line}')
-
- # Phase 2: Render standalone image/table chunks not inlined via connect_to
- for chunk in chunks:
- cid = chunk.get('chunk_id', '')
- if cid and cid in rendered_ids:
- continue
- if cid:
- rendered_ids.add(cid)
-
- chunk_type = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower()
- if chunk_type == 'image':
- file_path = chunk.get('file_path') or ''
- img_desc = str(chunk.get('content', '')).strip()
- asset_url = (asset_lookup or {}).get(cid, '') if cid else ''
- display_ref = asset_url or file_path
- if display_ref:
- parts.append(f'{indent}┈ [图片: {display_ref}]')
- if img_desc:
- for line in img_desc.split('\n'):
- if line.strip():
- parts.append(f'{indent}┈ {line}')
- elif chunk_type == 'table':
- table_html = str(chunk.get('content', '')).strip()
- parts.append(f'{indent}┈ [表格内容]')
- if table_html:
- for line in table_html.split('\n'):
- if line.strip():
- parts.append(f'{indent}┈ {line}')
-
-
-def render_unified_doc_tree(
- node: DocTreeNode,
- doc_name: str,
- depth: int = 0,
- asset_lookup: dict[str, str] | None = None,
-) -> str:
- """Render a DocTreeNode as a single coherent hierarchy.
-
- Summaries are navigation-only aids and NEVER appear in evidence.
- The rendered output contains:
- 1. Structural titles for ALL sections (positioning context)
- 2. Hydrated chunk content (┈ lines) ONLY for selected leaf paths
-
- Asset references (tables/images) are resolved via ``connect_to``
- metadata in hydrated chunks — no separate lookup needed.
- """
-
- parts: list[str] = []
- indent = ' ' * depth
-
- if depth == 0:
- parts.append(f'【文档】{doc_name}\n')
-
- # Collect children keys for path-hierarchy dedup:
- child_prefixes = set(node.children.keys())
-
- # Helper: min sort_order of a leaf_content entry
- def _min_sort(path: str) -> float:
- chunks = node.leaf_content.get(path, [])
- return min((c.get('sort_order') or float('inf') for c in chunks), default=float('inf'))
-
- # ── Build a unified render queue ──
- # Each entry: (sort_key, render_type, data)
- # render_type: 'outline' | 'orphan_leaf' | 'orphan_child'
- render_queue: list[tuple[float, str, dict | str]] = []
-
- outline_paths: set[str] = set()
- # Position counter for outline-only items (no leaf content) to preserve
- # their relative ordering among themselves.
- outline_position = 0.0
-
- for item in node.outline_items:
- path = item.get('path', '')
- # Skip items belonging to a drilled-into child's subtree
- if any(path.startswith(cp + ' / ') for cp in child_prefixes):
- continue
- outline_paths.add(path)
-
- # Determine sort_key: use chunk sort_order if content exists,
- # else use a synthetic position to maintain outline ordering.
- if path in node.leaf_content or path in node.children:
- sort_key = _min_sort(path) if path in node.leaf_content else outline_position
- else:
- sort_key = outline_position
- outline_position = max(outline_position, sort_key) + 0.001
-
- render_queue.append((sort_key, 'outline', item))
-
- # Add orphan leaf_content paths (not covered by outline_items)
- for path in node.leaf_content:
- if path not in outline_paths:
- render_queue.append((_min_sort(path), 'orphan_leaf', path))
-
- # Add orphan children (not covered by outline_items)
- for path in node.children:
- if path not in outline_paths:
- render_queue.append((float('inf'), 'orphan_child', path))
-
- # Sort by sort_key (stable sort preserves insertion order for ties)
- render_queue.sort(key=lambda x: x[0])
-
- from typing import cast
-
- # ── Render the unified queue ──
- for _sort_key, rtype, data in render_queue:
- if rtype == 'outline':
- item = cast(dict, data)
- path = item.get('path', '')
- title = item.get('title', '')
- is_leaf = item.get('is_leaf', False)
- level = item.get('level', 1)
- leaf_tag = ' [Leaf]' if is_leaf else ''
-
- level_tag = f'[L{level}] ' if level else ''
- if level <= 1:
- parts.append(f'{indent}▸ {level_tag}{title}{leaf_tag}')
- else:
- parts.append(f'{indent}└ {level_tag}{title}{leaf_tag}')
-
- sub_indent = indent + ' '
-
- # Case 1: drilled-into child → render child tree
- if path in node.children:
- child = node.children[path]
- if path in node.leaf_content:
- _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
- child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup)
- if child_text.strip():
- parts.append(child_text)
-
- # Case 2: hydrated leaf → show chunk content
- elif path in node.leaf_content:
- _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
-
- # Case 3: unselected → title only (already rendered above)
-
- elif rtype == 'orphan_leaf':
- path = cast(str, data)
- title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path
- parts.append(f'{indent}▸ [Leaf] {title}')
- sub_indent = indent + ' '
- _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
-
- elif rtype == 'orphan_child':
- path = cast(str, data)
- title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path
- parts.append(f'{indent}▸ {title} [DrillDown]')
- child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup)
- if child_text.strip():
- parts.append(child_text)
-
- return '\n'.join(parts)
-
diff --git a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py
new file mode 100644
index 000000000..48403e3ca
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py
@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import func as sa_func
+from sqlalchemy import or_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, DocumentSection
+from shared.models.database.job_result import JobResult
+
+
+def build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]:
+ owner_map: dict[str, str] = {}
+ for chunk in text_chunks:
+ if (chunk.get("chunk_type") or "text") != "text":
+ continue
+ section_path = chunk.get("section_path") or ""
+ if not section_path:
+ continue
+ metadata = chunk.get("chunk_metadata") or {}
+ if not isinstance(metadata, dict):
+ continue
+ for conn in metadata.get("connect_to") or []:
+ if not isinstance(conn, dict):
+ continue
+ target_id = str(conn.get("target") or "").strip()
+ if target_id and target_id not in owner_map:
+ owner_map[target_id] = section_path
+ return owner_map
+
+
+async def _load_scope_sections(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ scope_paths: list[str],
+) -> list[tuple[str, str]]:
+ section_stmt = (
+ select(DocumentSection.section_id, DocumentSection.section_path)
+ .where(DocumentSection.document_id == document_id)
+ .where(DocumentSection.job_result_id == job_result_id)
+ )
+ if scope_paths:
+ scope_filters = []
+ for scope in scope_paths:
+ scope_filters.append(DocumentSection.section_path == scope)
+ scope_filters.append(DocumentSection.section_path.like(f"{scope} / %"))
+ section_stmt = section_stmt.where(or_(*scope_filters))
+ rows = (await db.execute(section_stmt)).all()
+ return [(section_id, section_path or "") for section_id, section_path in rows]
+
+
+async def count_assets_under_scope(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ scope_paths: list[str],
+) -> tuple[int, int]:
+ section_rows = await _load_scope_sections(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ scope_paths=scope_paths,
+ )
+ all_section_ids = [section_id for section_id, _section_path in section_rows]
+
+ if not all_section_ids:
+ return 0, 0
+
+ count_stmt = (
+ select(
+ DocumentChunk.chunk_type,
+ sa_func.count(DocumentChunk.id),
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.section_id.in_(all_section_ids))
+ .where(DocumentChunk.chunk_type.in_(["image", "table"]))
+ .group_by(DocumentChunk.chunk_type)
+ )
+ count_result = await db.execute(count_stmt)
+
+ total_images = 0
+ total_tables = 0
+ for chunk_type, count in count_result.all():
+ if chunk_type == "image":
+ total_images = count
+ elif chunk_type == "table":
+ total_tables = count
+ return total_images, total_tables
+
+
+def build_asset_tools_block(total_images: int, total_tables: int) -> str:
+ if total_images <= 0 and total_tables <= 0:
+ return ""
+
+ tools_lines = ["\nOptional asset tools (usable with NAVIGATE or STOP):\n"]
+ if total_images > 0:
+ tools_lines.append(
+ f" FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n"
+ )
+ if total_tables > 0:
+ tools_lines.append(
+ f" FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n"
+ )
+ tools_lines.append(
+ " Note: with NAVIGATE selections, asset tools are limited to the selected sections; "
+ "with STOP or no selections, they use the current scope.\n"
+ )
+ return "".join(tools_lines)
+
+
+async def resolve_root_asset_owners(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ chunks: list[dict[str, Any]],
+) -> dict[str, str]:
+ root_asset_ids = [
+ str(chunk.get("chunk_id") or "")
+ for chunk in chunks
+ if not chunk.get("owner_section_path")
+ and (chunk.get("section_path") or "") == "Root"
+ and (chunk.get("chunk_type") or "").lower() in ("image", "table")
+ and chunk.get("chunk_id")
+ ]
+ if not root_asset_ids:
+ return {}
+
+ root_asset_set = set(root_asset_ids)
+ text_stmt = (
+ select(
+ DocumentChunk.chunk_metadata,
+ DocumentSection.section_path,
+ )
+ .outerjoin(
+ DocumentSection,
+ DocumentSection.section_id == DocumentChunk.section_id,
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.chunk_type == "text")
+ )
+ result = await db.execute(text_stmt)
+
+ owner_map: dict[str, str] = {}
+ for metadata, section_path in result.all():
+ if not isinstance(metadata, dict) or not section_path:
+ continue
+ for conn in metadata.get("connect_to") or []:
+ if not isinstance(conn, dict):
+ continue
+ target_id = str(conn.get("target") or "").strip()
+ if target_id in root_asset_set and target_id not in owner_map:
+ owner_map[target_id] = section_path
+
+ if owner_map:
+ logger.info(
+ f" resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} "
+ f"Root assets to their owner sections"
+ )
+ return owner_map
+
+
+async def asset_filter_step(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ scope_path: str | list[str] | None,
+ asset_type: str,
+) -> list[dict[str, Any]]:
+ t0 = time.monotonic()
+ try:
+ scope_list = (
+ scope_path
+ if isinstance(scope_path, list)
+ else [scope_path]
+ if scope_path
+ else []
+ )
+
+ section_rows = await _load_scope_sections(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ scope_paths=scope_list,
+ )
+ section_ids = {row[0] for row in section_rows}
+
+ if not section_ids:
+ logger.info(f" asset_filter_step: no sections found under scope={scope_path}")
+ return []
+
+ section_path_by_id = {
+ section_id: section_path for section_id, section_path in section_rows
+ }
+ asset_rows = (
+ await db.execute(
+ select(
+ DocumentChunk.chunk_id,
+ DocumentChunk.chunk_type,
+ DocumentChunk.content,
+ DocumentChunk.file_path,
+ DocumentChunk.section_id,
+ DocumentChunk.source_chunk_path,
+ DocumentChunk.chunk_metadata,
+ DocumentChunk.sort_order,
+ DocumentChunk.job_result_id,
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.section_id.in_(list(section_ids)))
+ .where(DocumentChunk.chunk_type == asset_type)
+ .order_by(DocumentChunk.sort_order)
+ )
+ ).all()
+
+ text_rows = (
+ await db.execute(
+ select(
+ DocumentChunk.section_id,
+ DocumentChunk.chunk_type,
+ DocumentChunk.chunk_metadata,
+ DocumentChunk.source_chunk_path,
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.section_id.in_(list(section_ids)))
+ .where(DocumentChunk.chunk_type == "text")
+ )
+ ).all()
+ text_row_dicts = [
+ {
+ "chunk_type": chunk_type,
+ "chunk_metadata": metadata or {},
+ "section_id": section_id,
+ "section_path": section_path_by_id.get(section_id, ""),
+ "source_chunk_path": source_chunk_path,
+ }
+ for section_id, chunk_type, metadata, source_chunk_path in text_rows
+ ]
+ owner_by_target_id = build_connected_owner_map(text_row_dicts)
+
+ if any(value == "Root" for value in owner_by_target_id.values()):
+ doc_stmt = select(Document.source_file_name).where(
+ Document.document_id == document_id
+ )
+ doc_file_name = (await db.execute(doc_stmt)).scalar() or ""
+ if doc_file_name:
+ for target_id in list(owner_by_target_id):
+ if owner_by_target_id[target_id] == "Root":
+ owner_by_target_id[target_id] = doc_file_name
+
+ connected_target_ids: set[str] = set(owner_by_target_id.keys())
+ if connected_target_ids:
+ connected_rows = (
+ await db.execute(
+ select(
+ DocumentChunk.chunk_id,
+ DocumentChunk.chunk_type,
+ DocumentChunk.content,
+ DocumentChunk.file_path,
+ DocumentChunk.section_id,
+ DocumentChunk.source_chunk_path,
+ DocumentChunk.chunk_metadata,
+ DocumentChunk.sort_order,
+ DocumentChunk.job_result_id,
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.chunk_id.in_(list(connected_target_ids)))
+ .where(DocumentChunk.chunk_type == asset_type)
+ .order_by(DocumentChunk.sort_order)
+ )
+ ).all()
+ else:
+ connected_rows = []
+
+ job_id = (
+ await db.execute(select(JobResult.job_id).where(JobResult.id == job_result_id))
+ ).scalar() or ""
+ seen_ids: set[str] = set()
+ chunks: list[dict[str, Any]] = []
+ for row in list(asset_rows) + list(connected_rows):
+ chunk_id = row[0]
+ if chunk_id in seen_ids:
+ continue
+ seen_ids.add(chunk_id)
+
+ owner_section_path = owner_by_target_id.get(chunk_id)
+ if not owner_section_path:
+ own_section_path = section_path_by_id.get(row[4])
+ if own_section_path and own_section_path == "Root":
+ logger.warning(
+ " asset_filter_step: rejecting root-level owner fallback "
+ f"chunk_id={chunk_id} section_path={own_section_path}"
+ )
+ own_section_path = None
+ owner_section_path = own_section_path
+
+ if not owner_section_path:
+ logger.warning(
+ f" asset_filter_step unresolved owner: chunk_id={chunk_id} "
+ f"file_path={row[3]} scope={scope_path or 'root'}"
+ )
+ continue
+
+ chunks.append(
+ {
+ "document_id": document_id,
+ "chunk_id": chunk_id,
+ "chunk_type": row[1],
+ "content": row[2],
+ "file_path": row[3],
+ "section_id": row[4],
+ "section_path": owner_section_path,
+ "owner_section_path": owner_section_path,
+ "source_chunk_path": row[5],
+ "chunk_metadata": row[6] or {},
+ "sort_order": row[7],
+ "job_result_id": job_result_id,
+ "job_id": job_id,
+ }
+ )
+
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.info(
+ f" asset_filter_step scope={scope_path or 'root'} "
+ f"type={asset_type}: {len(chunks)} chunks found, {latency}ms"
+ )
+ return chunks
+
+ except Exception as exc:
+ logger.error(f" asset_filter_step failed: {exc}")
+ return []
diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py
new file mode 100644
index 000000000..2af9ce690
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py
@@ -0,0 +1,252 @@
+"""Discovery and document selection phase for agentic retrieval."""
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document
+from shared.services.retrieval.agentic import tools
+from shared.services.retrieval.agentic.budget import BudgetExceeded
+from shared.services.retrieval.agentic.trace import TraceRecorder
+from shared.services.retrieval.agentic.types import AgentState, CandidateDoc, ToolResult
+from shared.services.retrieval.llm_adapter import LLMFn
+
+
+async def run_initial_discovery(
+ db: AsyncSession,
+ *,
+ state: AgentState,
+ trace: TraceRecorder,
+ trace_enabled: bool,
+ user_id: str,
+ namespace: str,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int,
+ signal_paths: list[str] | None,
+ filter_mode: str,
+ channels: list[str] | None,
+ channel_weights: dict[str, float] | None,
+ bootstrap_llm_fn: LLMFn | None,
+) -> list[dict[str, Any]]:
+ discovery_kwargs: dict[str, Any] = {
+ "user_id": user_id,
+ "namespace": namespace,
+ "query": query,
+ "top_k": top_k,
+ "exclude_document_ids": exclude_document_ids,
+ "exclude_sections": exclude_sections,
+ "data_type": data_type,
+ "signal_paths": signal_paths,
+ "filter_mode": filter_mode,
+ "channels": channels,
+ "channel_weights": channel_weights,
+ }
+
+ logger.info(" agentic: Phase 1 — discovery + document selection")
+ discovery_result = await tools.bottom_discovery(db, **discovery_kwargs)
+ state.step_count += 1
+ discovery_rows = (
+ discovery_result.payload.get("fused_rows", [])
+ if discovery_result.status != "error"
+ else []
+ )
+ state.discovery_top_doc_ids = (
+ discovery_result.payload.get("top_doc_ids", [])
+ if discovery_result.status != "error"
+ else []
+ )
+
+ if trace_enabled:
+ trace.record_step(
+ "bottom_discovery",
+ discovery_result,
+ decision_reason="phase_1_mandatory",
+ )
+
+ logger.info(
+ f" agentic step {state.step_count}: bottom_discovery "
+ f"status={discovery_result.status} latency={discovery_result.latency_ms}ms"
+ )
+
+ if bootstrap_llm_fn is not None:
+ await _select_documents(
+ db,
+ state=state,
+ trace=trace,
+ trace_enabled=trace_enabled,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ exclude_document_ids=exclude_document_ids,
+ bootstrap_llm_fn=bootstrap_llm_fn,
+ )
+
+ return discovery_rows
+
+
+async def register_discovery_documents(
+ db: AsyncSession,
+ *,
+ state: AgentState,
+ discovery_by_doc: dict[str, list[dict[str, Any]]],
+) -> None:
+ selected_doc_ids = {doc.document_id for doc in state.selected_docs}
+ for doc_id in discovery_by_doc:
+ if doc_id in selected_doc_ids or doc_id in state.ever_explored_doc_ids:
+ continue
+ doc_stmt = (
+ select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
+ .where(Document.document_id == doc_id)
+ )
+ doc_result = await db.execute(doc_stmt)
+ row_data = doc_result.first()
+ if row_data is None:
+ continue
+ did, fname, job_result_id = row_data
+ state.selected_docs.append(
+ CandidateDoc(
+ document_id=did,
+ source_file_name=fname or did,
+ confidence=0.4,
+ reason="discovery_auto (not in KG selection)",
+ source="discovery_auto",
+ )
+ )
+ state.doc_id_to_name[did] = fname or did
+ if job_result_id:
+ state.doc_job_map[did] = job_result_id
+
+
+async def select_revision_documents(
+ db: AsyncSession,
+ *,
+ state: AgentState,
+ trace: TraceRecorder,
+ trace_enabled: bool,
+ user_id: str,
+ namespace: str,
+ query: str,
+ exclude_document_ids: list[str],
+ bootstrap_llm_fn: LLMFn,
+ revision_hint: str,
+) -> str | None:
+ try:
+ kg_result = await tools.kg_document_select(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ llm_fn=bootstrap_llm_fn,
+ exclude_document_ids=list(set(exclude_document_ids)),
+ revision_hint=revision_hint,
+ budget_snapshot=state.ledger.snapshot() if state.ledger else None,
+ )
+ except BudgetExceeded:
+ logger.info(" agentic: bootstrap budget exhausted during revision doc selection")
+ if trace_enabled:
+ trace.record_budget_stop("bootstrap_exhausted")
+ return "bootstrap_budget"
+ state.step_count += 1
+ _append_selected_docs(state, kg_result)
+ return None
+
+
+async def _select_documents(
+ db: AsyncSession,
+ *,
+ state: AgentState,
+ trace: TraceRecorder,
+ trace_enabled: bool,
+ user_id: str,
+ namespace: str,
+ query: str,
+ exclude_document_ids: list[str],
+ bootstrap_llm_fn: LLMFn,
+) -> None:
+ try:
+ kg_result = await tools.kg_document_select(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ llm_fn=bootstrap_llm_fn,
+ exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)),
+ budget_snapshot=state.ledger.snapshot() if state.ledger else None,
+ )
+ except BudgetExceeded:
+ logger.info(" agentic: bootstrap budget exhausted during document selection")
+ if trace_enabled:
+ trace.record_budget_stop("bootstrap_exhausted")
+ kg_result = ToolResult(
+ status="no_confident_doc",
+ payload={"reason": "bootstrap budget exhausted"},
+ )
+ state.step_count += 1
+
+ if trace_enabled:
+ trace.record_step(
+ "kg_document_select",
+ kg_result,
+ decision_reason="phase_1_doc_selection",
+ )
+
+ _append_selected_docs(state, kg_result)
+ if not state.selected_docs and state.discovery_top_doc_ids:
+ await _append_discovery_hints(db, state=state)
+
+ logger.info(
+ f" agentic step {state.step_count}: kg_document_select "
+ f"status={kg_result.status} docs={len(state.selected_docs)} "
+ f"latency={kg_result.latency_ms}ms"
+ )
+
+
+async def _append_discovery_hints(db: AsyncSession, *, state: AgentState) -> None:
+ hint_ids = [
+ doc_id
+ for doc_id in state.discovery_top_doc_ids
+ if doc_id not in state.ever_explored_doc_ids
+ ]
+ if not hint_ids:
+ return
+ doc_stmt = (
+ select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
+ .where(Document.document_id.in_(hint_ids))
+ )
+ doc_result = await db.execute(doc_stmt)
+ for doc_id, source_file_name, job_result_id in doc_result.all():
+ state.selected_docs.append(
+ CandidateDoc(
+ document_id=doc_id,
+ source_file_name=source_file_name or doc_id,
+ confidence=0.5,
+ reason="discovery_hint (KG returned 0)",
+ source="discovery_hint",
+ )
+ )
+ state.doc_id_to_name[doc_id] = source_file_name or doc_id
+ if job_result_id:
+ state.doc_job_map[doc_id] = job_result_id
+
+
+def _append_selected_docs(state: AgentState, kg_result: ToolResult) -> None:
+ if kg_result.status != "selected_docs":
+ return
+ for doc_data in kg_result.payload.get("candidate_docs", []):
+ state.selected_docs.append(
+ CandidateDoc(
+ document_id=doc_data.get("document_id", ""),
+ source_file_name=doc_data.get("source_file_name", ""),
+ confidence=doc_data.get("confidence", 0.0),
+ reason=doc_data.get("reason", ""),
+ source=doc_data.get("source", ""),
+ )
+ )
+ state.doc_id_to_name.update(kg_result.payload.get("doc_id_to_name", {}))
+ state.doc_job_map.update(kg_result.payload.get("doc_job_map", {}))
diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py
new file mode 100644
index 000000000..d54df4e5f
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py
@@ -0,0 +1,206 @@
+"""Post-navigation discovery selection for agentic retrieval."""
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.agentic.budget import BudgetExceeded
+from shared.services.retrieval.agentic.prompts import (
+ DISCOVERY_SELECT_PROMPT,
+ format_budget_block,
+ parse_action_response,
+)
+from shared.services.retrieval.agentic.selection_hydration import (
+ hydrate_path_selections_into_node,
+)
+from shared.services.retrieval.agentic.types import DocTreeNode
+from shared.services.retrieval.lexical_text import normalize_section_path
+from shared.services.retrieval.llm_adapter import LLMFn
+
+
+_MAX_DISCOVERY_PER_DOC = 3
+
+
+async def discovery_select_step(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ query: str,
+ llm_fn: LLMFn,
+ user_id: str,
+ namespace: str,
+ doc_name: str = "",
+ discovery_hints: list[dict[str, Any]],
+ exclude_paths: set[str] | None = None,
+ revision_hint: str | None = None,
+ budget_snapshot: dict | None = None,
+) -> DocTreeNode:
+ """Select and hydrate discovery-found sections after BFS navigation."""
+ node = DocTreeNode(scope_path=None)
+ if not discovery_hints:
+ return node
+
+ hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC]
+
+ t0 = time.monotonic()
+ try:
+ hint_lines, hint_by_path, root_path_selections = _project_discovery_hints(
+ hints,
+ exclude_paths=exclude_paths,
+ )
+ if not hint_lines and not root_path_selections:
+ return node
+
+ selections: list[dict[str, Any]] = []
+ if hint_lines:
+ prompt = _build_discovery_selection_prompt(
+ document_id=document_id,
+ doc_name=doc_name,
+ query=query,
+ hint_lines=hint_lines,
+ revision_hint=revision_hint,
+ budget_snapshot=budget_snapshot,
+ )
+ response = await llm_fn(prompt)
+ parsed = parse_action_response(response)
+ selections = parsed.get("selections", [])
+
+ logger.info(
+ f' discovery_select_step doc="{doc_name}": '
+ f"hints={len(hints)} selections={len(selections)} "
+ f"root_selections={len(root_path_selections)}"
+ )
+
+ path_selections = _build_discovery_path_selections(
+ selections=selections,
+ hint_by_path=hint_by_path,
+ root_path_selections=root_path_selections,
+ node=node,
+ )
+ await hydrate_path_selections_into_node(
+ db,
+ node=node,
+ path_selections=path_selections,
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
+
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.info(
+ f" discovery_select_step done: hydrated={len(node.leaf_content)} "
+ f"latency={latency}ms"
+ )
+ return node
+
+ except BudgetExceeded:
+ raise
+ except Exception as exc:
+ logger.error(f" discovery_select_step failed for doc={document_id}: {exc}")
+ return node
+
+
+def _project_discovery_hints(
+ hints: list[dict[str, Any]],
+ *,
+ exclude_paths: set[str] | None,
+) -> tuple[list[str], dict[str, dict], list[dict[str, Any]]]:
+ exclude_set = {
+ normalize_section_path(path)
+ for path in (exclude_paths or set())
+ if path
+ }
+ hint_lines: list[str] = []
+ hint_by_path: dict[str, dict] = {}
+ root_path_selections: list[dict[str, Any]] = []
+ for hint in hints:
+ section_path = normalize_section_path(hint.get("section_path", ""))
+ if not section_path:
+ continue
+ if section_path in exclude_set:
+ continue
+ if section_path in hint_by_path:
+ continue
+
+ hint_by_path[section_path] = hint
+ if section_path == "Root":
+ root_path_selections.append({
+ "path": section_path,
+ "confidence": float(
+ hint.get("discovery_score") or hint.get("score") or 0.7
+ ),
+ "hydrate_mode": "self_only",
+ })
+ continue
+
+ summary = hint.get("summary", "") or ""
+ hint_lines.append(f'▸ path="{section_path}"')
+ if summary:
+ hint_lines.append(f" {summary[:300]}")
+
+ return hint_lines, hint_by_path, root_path_selections
+
+
+def _build_discovery_selection_prompt(
+ *,
+ document_id: str,
+ doc_name: str,
+ query: str,
+ hint_lines: list[str],
+ revision_hint: str | None,
+ budget_snapshot: dict | None,
+) -> str:
+ revision_context = ""
+ if revision_hint:
+ revision_context = (
+ "\nIMPORTANT: This is a REVISION round. "
+ "The previous search attempt failed because:\n"
+ f'"{revision_hint}"\n'
+ "Adjust your selection accordingly. "
+ "If no candidate is relevant, return an EMPTY list [].\n"
+ )
+
+ return DISCOVERY_SELECT_PROMPT.format(
+ doc_name=doc_name or document_id,
+ budget_block=format_budget_block(budget_snapshot),
+ items="\n".join(hint_lines),
+ query=query,
+ revision_context=revision_context,
+ )
+
+
+def _build_discovery_path_selections(
+ *,
+ selections: list[dict[str, Any]],
+ hint_by_path: dict[str, dict],
+ root_path_selections: list[dict[str, Any]],
+ node: DocTreeNode,
+) -> list[dict[str, Any]]:
+ valid_selections = [
+ selection for selection in selections if selection["path"] in hint_by_path
+ ]
+ path_selections = list(root_path_selections)
+ for selection in valid_selections:
+ path = selection["path"]
+ confidence = selection.get("confidence", 0.7)
+ node.confidence[path] = confidence
+ path_selections.append({"path": path, "confidence": confidence})
+
+ if not path_selections and hint_by_path:
+ fallback_path, fallback_hint = next(iter(hint_by_path.items()))
+ fallback_confidence = float(
+ fallback_hint.get("discovery_score")
+ or fallback_hint.get("score")
+ or 0.5
+ )
+ node.confidence[fallback_path] = fallback_confidence
+ path_selections.append({
+ "path": fallback_path,
+ "confidence": fallback_confidence,
+ "hydrate_mode": "self_only",
+ })
+
+ return path_selections
diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py
new file mode 100644
index 000000000..5f32e36cb
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py
@@ -0,0 +1,294 @@
+"""Agentic retrieval discovery tools.
+
+This Module owns phase-1 retrieval: lexical bottom discovery and document
+selection from the document-level knowledge map. The public tool adapter stays
+in ``tools.py`` so orchestrator call sites keep a stable interface.
+"""
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document
+from shared.services.retrieval.agentic.budget import BudgetExceeded
+from shared.services.retrieval.agentic.knowledge_map import build_knowledge_map_overview
+from shared.services.retrieval.agentic.prompts import (
+ FILE_SELECT_PROMPT,
+ format_budget_block,
+ parse_json_array,
+)
+from shared.services.retrieval.agentic.types import ToolResult
+from shared.services.retrieval.channels import content_channel, path_channel, term_channel
+from shared.services.retrieval.llm_adapter import LLMFn
+from shared.services.retrieval.scoring import (
+ merge_channels_rrf,
+ merge_same_section_rows,
+ normalize_row_scores,
+)
+from shared.services.retrieval.settings import (
+ CHANNEL_WEIGHT_CONTENT,
+ CHANNEL_WEIGHT_PATH,
+ CHANNEL_WEIGHT_TERM,
+ INTERNAL_RECALL_K_MULTIPLIER,
+ resolve_allowed_chunk_types,
+)
+
+
+async def bottom_discovery(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int = 1,
+ signal_paths: list[str] | None = None,
+ filter_mode: str = "delete",
+ channels: list[str] | None = None,
+ channel_weights: dict[str, float] | None = None,
+ internal_recall_k: int | None = None,
+ **_kwargs: Any,
+) -> ToolResult:
+ """Run 3-channel BM25 discovery plus RRF fusion."""
+ t0 = time.monotonic()
+ try:
+ allowed_chunk_types = resolve_allowed_chunk_types(data_type)
+ effective_recall_k = (
+ internal_recall_k
+ if internal_recall_k is not None
+ else top_k * INTERNAL_RECALL_K_MULTIPLIER
+ )
+ active_channels = set(channels) if channels else {"path", "content", "term"}
+
+ path_rows: list[dict[str, Any]] = []
+ content_rows: list[dict[str, Any]] = []
+ term_rows: list[dict[str, Any]] = []
+
+ if "path" in active_channels:
+ path_rows = await path_channel(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=effective_recall_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ allowed_chunk_types=allowed_chunk_types,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ )
+
+ if "content" in active_channels:
+ content_rows = await content_channel(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=effective_recall_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ allowed_chunk_types=allowed_chunk_types,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ )
+
+ if "term" in active_channels:
+ term_rows = await term_channel(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=effective_recall_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ allowed_chunk_types=allowed_chunk_types,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ )
+
+ default_weights = {
+ "path": CHANNEL_WEIGHT_PATH,
+ "content": CHANNEL_WEIGHT_CONTENT,
+ "term": CHANNEL_WEIGHT_TERM,
+ }
+ effective_weights = {**default_weights, **(channel_weights or {})}
+
+ channel_lists: list[list[dict[str, Any]]] = []
+ weight_list: list[float] = []
+ if path_rows:
+ channel_lists.append(path_rows)
+ weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH))
+ if content_rows:
+ channel_lists.append(content_rows)
+ weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT))
+ if term_rows:
+ channel_lists.append(term_rows)
+ weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM))
+
+ fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else []
+ fused_rows = merge_same_section_rows(fused_rows)
+
+ if fused_rows:
+ normalize_row_scores(
+ fused_rows,
+ source_field="score",
+ target_field="discovery_score",
+ default=0.5,
+ )
+
+ doc_id_counts: dict[str, int] = {}
+ for row in fused_rows:
+ did = row.get("document_id", "")
+ if did:
+ doc_id_counts[did] = doc_id_counts.get(did, 0) + 1
+ top_doc_ids = sorted(
+ doc_id_counts,
+ key=lambda document_id: doc_id_counts[document_id],
+ reverse=True,
+ )[:5]
+
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.info(
+ f" agentic.bottom_discovery: {len(fused_rows)} fused rows, "
+ f"top_doc_ids={top_doc_ids}, {latency}ms"
+ )
+ return ToolResult(
+ status="discovery_done",
+ payload={
+ "fused_rows": fused_rows,
+ "top_doc_ids": top_doc_ids,
+ "channel_counts": {
+ "path": len(path_rows),
+ "content": len(content_rows),
+ "term": len(term_rows),
+ },
+ },
+ latency_ms=latency,
+ )
+ except Exception as exc:
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.error(f" agentic.bottom_discovery failed: {exc}")
+ return ToolResult(status="error", error=str(exc), latency_ms=latency)
+
+
+async def kg_document_select(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ llm_fn: LLMFn | None,
+ exclude_document_ids: list[str],
+ revision_hint: str | None = None,
+ **_kwargs: Any,
+) -> ToolResult:
+ """Select candidate documents from document-level KG."""
+ t0 = time.monotonic()
+ try:
+ overview_text, doc_id_to_name = await build_knowledge_map_overview(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ )
+ if overview_text == "(empty)":
+ latency = int((time.monotonic() - t0) * 1000)
+ return ToolResult(
+ status="no_confident_doc",
+ payload={"reason": "no active documents in namespace"},
+ latency_ms=latency,
+ )
+
+ if llm_fn is None:
+ latency = int((time.monotonic() - t0) * 1000)
+ return ToolResult(
+ status="no_confident_doc",
+ payload={"reason": "LLM not available"},
+ latency_ms=latency,
+ )
+
+ revision_context = ""
+ if revision_hint:
+ revision_context = (
+ "\nIMPORTANT: This is a REVISION round. "
+ "The previous search attempt failed because:\n"
+ f'"{revision_hint}"\n'
+ "Adjust your document selection accordingly. "
+ "If no document can address this, return an EMPTY array [].\n"
+ )
+
+ file_prompt = FILE_SELECT_PROMPT.format(
+ overview=overview_text,
+ query=query,
+ revision_context=revision_context,
+ budget_block=format_budget_block(_kwargs.get("budget_snapshot")),
+ )
+ file_response = await llm_fn(file_prompt)
+ selected_ids = parse_json_array(file_response)
+
+ exclude_set = set(exclude_document_ids)
+ valid_ids = [
+ document_id
+ for document_id in selected_ids
+ if document_id in doc_id_to_name and document_id not in exclude_set
+ ]
+
+ if not valid_ids:
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.info(
+ f" agentic.kg_document_select: LLM returned no valid docs, {latency}ms"
+ )
+ return ToolResult(
+ status="no_confident_doc",
+ payload={
+ "reason": "LLM returned no valid document IDs",
+ "raw_ids": selected_ids,
+ },
+ latency_ms=latency,
+ )
+
+ doc_job_map: dict[str, str] = {}
+ doc_stmt = (
+ select(Document.document_id, Document.current_job_result_id)
+ .where(Document.document_id.in_(valid_ids))
+ )
+ doc_result = await db.execute(doc_stmt)
+ for document_id, job_result_id in doc_result.all():
+ if job_result_id:
+ doc_job_map[document_id] = job_result_id
+
+ candidate_docs = [
+ {
+ "document_id": document_id,
+ "source_file_name": doc_id_to_name.get(document_id, ""),
+ "confidence": 1.0,
+ "reason": "LLM selected from KG overview",
+ "source": "kg_llm_select",
+ }
+ for document_id in valid_ids
+ ]
+
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.info(
+ f" agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms"
+ )
+ return ToolResult(
+ status="selected_docs",
+ payload={
+ "candidate_docs": candidate_docs,
+ "doc_id_to_name": doc_id_to_name,
+ "doc_job_map": doc_job_map,
+ },
+ latency_ms=latency,
+ )
+ except BudgetExceeded:
+ raise
+ except Exception as exc:
+ latency = int((time.monotonic() - t0) * 1000)
+ logger.error(f" agentic.kg_document_select failed: {exc}")
+ return ToolResult(status="error", error=str(exc), latency_ms=latency)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py b/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py
new file mode 100644
index 000000000..14b3d5f63
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py
@@ -0,0 +1,414 @@
+"""Per-document navigation for agentic retrieval."""
+from __future__ import annotations
+
+from typing import Any, cast
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.agentic import tools
+from shared.services.retrieval.agentic.budget import BudgetExceeded
+from shared.services.retrieval.agentic.evidence import reconcile_deferred_assets
+from shared.services.retrieval.agentic.runtime import AgentLlmBudget
+from shared.services.retrieval.agentic.trace import TraceRecorder
+from shared.services.retrieval.agentic.types import (
+ AgentRunConfig,
+ AgentState,
+ CandidateDoc,
+ DocTreeNode,
+ ToolResult,
+)
+from shared.services.retrieval.llm_adapter import LLMFn
+
+
+class DocumentNavigationRunner:
+ def __init__(
+ self,
+ *,
+ db: AsyncSession,
+ state: AgentState,
+ trace: TraceRecorder,
+ trace_enabled: bool,
+ user_id: str,
+ namespace: str,
+ query: str,
+ config: AgentRunConfig,
+ discovery_by_doc: dict[str, list[dict[str, Any]]],
+ llm_fn: LLMFn | None,
+ llm_budget: AgentLlmBudget,
+ ) -> None:
+ self._db = db
+ self._state = state
+ self._trace = trace
+ self._trace_enabled = trace_enabled
+ self._user_id = user_id
+ self._namespace = namespace
+ self._query = query
+ self._config = config
+ self._discovery_by_doc = discovery_by_doc
+ self._llm_fn = llm_fn
+ self._llm_budget = llm_budget
+
+ async def navigate_selected_documents(self, *, revision_hint: str | None) -> None:
+ logger.info(
+ f" agentic: Phase 2 — navigating {len(self._state.selected_docs)} documents"
+ )
+ for doc in self._state.selected_docs:
+ if self._state.elapsed_ms >= self._config.latency_budget_ms:
+ logger.info(" agentic: latency budget hit during Phase 2, stopping")
+ break
+ await self._navigate_document(doc, revision_hint=revision_hint)
+
+ async def _navigate_document(
+ self,
+ doc: CandidateDoc,
+ *,
+ revision_hint: str | None,
+ ) -> None:
+ job_result_id = self._state.doc_job_map.get(doc.document_id, "")
+ if not job_result_id:
+ logger.info(f" agentic: skipping doc {doc.document_id} — no job_result_id")
+ self._state.ever_explored_doc_ids.add(doc.document_id)
+ return
+
+ doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "")
+ is_discovery_only_doc = doc.source == "discovery_auto"
+ root = DocTreeNode(scope_path=None)
+ doc_pending_assets: list[dict[str, Any]] = []
+
+ if not is_discovery_only_doc:
+ doc_pending_assets = await self._navigate_bfs(
+ doc=doc,
+ root=root,
+ doc_name=doc_name,
+ job_result_id=job_result_id,
+ revision_hint=revision_hint,
+ )
+
+ await self._hydrate_discovery_hints(
+ doc=doc,
+ root=root,
+ doc_name=doc_name,
+ revision_hint=revision_hint,
+ )
+
+ if not is_discovery_only_doc and doc_pending_assets:
+ self._reconcile_pending_assets(
+ doc=doc,
+ root=root,
+ doc_name=doc_name,
+ doc_pending_assets=doc_pending_assets,
+ )
+
+ if doc.document_id in self._state.doc_trees:
+ self._state.doc_trees[doc.document_id].merge(root)
+ else:
+ self._state.doc_trees[doc.document_id] = root
+ self._state.ever_explored_doc_ids.add(doc.document_id)
+ if self._state.ledger is not None:
+ self._state.ledger.mark_explored(docs=1)
+
+ async def _navigate_bfs(
+ self,
+ *,
+ doc: CandidateDoc,
+ root: DocTreeNode,
+ doc_name: str,
+ job_result_id: str,
+ revision_hint: str | None,
+ ) -> list[dict[str, Any]]:
+ doc_exclude: set[str] = {
+ key.split("::", 1)[1]
+ for key in self._state.seen_section_keys
+ if key.startswith(f"{doc.document_id}::")
+ } if self._state.seen_section_keys else set()
+ pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)]
+ doc_pending_assets: list[dict[str, Any]] = []
+
+ while pending:
+ if self._state.elapsed_ms >= self._config.latency_budget_ms:
+ break
+
+ scope, parent_node, depth = pending.pop(0)
+ if depth >= self._config.max_nav_depth:
+ continue
+ if self._llm_fn is None:
+ break
+ if self._state.ledger and self._state.ledger.status("planning") in ("CRITICAL", "EXHAUSTED"):
+ logger.info(" agentic: planning budget critical, ending BFS for current doc")
+ break
+
+ doc_llm_fn = self._llm_budget.for_document(
+ cast(LLMFn, self._llm_fn),
+ doc_id=doc.document_id,
+ depth=depth,
+ )
+ try:
+ action, asset_tools, step_node, drill_paths = await tools.navigate_step(
+ self._db,
+ document_id=doc.document_id,
+ job_result_id=job_result_id,
+ query=self._query,
+ llm_fn=doc_llm_fn,
+ user_id=self._user_id,
+ namespace=self._namespace,
+ doc_name=doc_name,
+ scope_path=scope,
+ exclude_paths=doc_exclude,
+ revision_hint=revision_hint if depth == 0 else None,
+ budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None,
+ )
+ except BudgetExceeded:
+ logger.info(" agentic: planning budget exhausted during navigation")
+ if self._trace_enabled:
+ self._trace.record_budget_stop("planning_exhausted")
+ break
+ self._state.step_count += 1
+
+ await self._collect_assets(
+ doc=doc,
+ scope=scope,
+ step_node=step_node,
+ asset_tools=asset_tools,
+ pending_assets=doc_pending_assets,
+ round_scope="nav",
+ )
+ _merge_step_node(parent_node, step_node)
+ _update_excluded_leaf_paths(doc_exclude, step_node, drill_paths)
+ _queue_drill_paths(pending, parent_node, drill_paths, depth)
+ parent_node.reparent_leaf_content()
+ self._record_navigation_step(
+ doc=doc,
+ scope=scope,
+ depth=depth,
+ action=action,
+ asset_tools=asset_tools,
+ step_node=step_node,
+ drill_paths=drill_paths,
+ )
+ if self._state.ledger is not None:
+ self._state.ledger.mark_explored(
+ chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()),
+ )
+
+ return doc_pending_assets
+
+ async def _collect_assets(
+ self,
+ *,
+ doc: CandidateDoc,
+ scope: str | list[str] | None,
+ step_node: DocTreeNode,
+ asset_tools: list[str],
+ pending_assets: list[dict[str, Any]],
+ round_scope: str,
+ ) -> None:
+ selected_asset_scopes = list(step_node.confidence.keys())
+ asset_scope = selected_asset_scopes or scope
+ for asset_tool in asset_tools:
+ if asset_tool not in ("FIND_IMAGES", "FIND_TABLES"):
+ continue
+ asset_type = "image" if asset_tool == "FIND_IMAGES" else "table"
+ asset_chunks = await tools.asset_filter_step(
+ self._db,
+ document_id=doc.document_id,
+ job_result_id=self._state.doc_job_map.get(doc.document_id, ""),
+ scope_path=asset_scope,
+ asset_type=asset_type,
+ )
+ if asset_chunks:
+ pending_assets.extend(asset_chunks)
+
+ scope_display = (
+ asset_scope
+ if isinstance(asset_scope, list)
+ else (asset_scope or "root")
+ )
+ if self._trace_enabled:
+ self._trace.record_step(
+ "asset_filter_step",
+ ToolResult(
+ status="filtered" if asset_chunks else "empty",
+ payload={
+ "document_id": doc.document_id,
+ "scope": scope_display,
+ "navigation_scope": scope if isinstance(scope, str) else (scope or "root"),
+ "asset_type": asset_type,
+ "chunks_found": len(asset_chunks) if asset_chunks else 0,
+ },
+ ),
+ decision_reason=f"asset_{round_scope}_{doc.source_file_name}",
+ )
+ logger.info(
+ f" agentic step {self._state.step_count}: asset_filter_step "
+ f'doc="{doc.source_file_name}" scope={scope_display} '
+ f"type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}"
+ )
+
+ async def _hydrate_discovery_hints(
+ self,
+ *,
+ doc: CandidateDoc,
+ root: DocTreeNode,
+ doc_name: str,
+ revision_hint: str | None,
+ ) -> None:
+ doc_hints = self._discovery_by_doc.get(doc.document_id, [])
+ if not doc_hints or self._llm_fn is None:
+ return
+ if self._state.elapsed_ms >= self._config.latency_budget_ms:
+ return
+
+ discovery_exclude_paths = {
+ key.split("::", 1)[1]
+ for key in root.collect_all_paths(doc.document_id)
+ }
+ doc_discovery_llm_fn = self._llm_budget.for_discovery(
+ cast(LLMFn, self._llm_fn),
+ doc_id=doc.document_id,
+ low_priority=root.has_content(),
+ )
+ try:
+ discovery_node = await tools.discovery_select_step(
+ self._db,
+ document_id=doc.document_id,
+ query=self._query,
+ llm_fn=doc_discovery_llm_fn,
+ user_id=self._user_id,
+ namespace=self._namespace,
+ doc_name=doc_name,
+ discovery_hints=doc_hints,
+ exclude_paths=discovery_exclude_paths,
+ revision_hint=revision_hint,
+ budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None,
+ )
+ except BudgetExceeded:
+ logger.info(" agentic: planning budget exhausted during discovery selection")
+ if self._trace_enabled:
+ self._trace.record_budget_stop("planning_exhausted")
+ discovery_node = DocTreeNode(scope_path=None)
+ self._state.step_count += 1
+
+ if self._trace_enabled:
+ self._trace.record_step(
+ "discovery_select_step",
+ ToolResult(
+ status="selected" if discovery_node.has_content() else "empty",
+ payload={
+ "document_id": doc.document_id,
+ "hints_count": len(doc_hints),
+ "hydrated_count": len(discovery_node.leaf_content),
+ },
+ ),
+ decision_reason=f"discovery_{doc.source_file_name}",
+ )
+ root.merge(discovery_node)
+ if self._state.ledger is not None:
+ self._state.ledger.mark_explored(
+ chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()),
+ )
+
+ def _reconcile_pending_assets(
+ self,
+ *,
+ doc: CandidateDoc,
+ root: DocTreeNode,
+ doc_name: str,
+ doc_pending_assets: list[dict[str, Any]],
+ ) -> None:
+ if doc_name and not root.children and not any(
+ item.get("path") == doc_name for item in root.outline_items
+ ):
+ root.outline_items.insert(0, {"path": doc_name, "level": 0})
+ reconcile_deferred_assets(root, doc_pending_assets)
+ if self._trace_enabled:
+ self._trace.record_step(
+ "deferred_asset_reconcile",
+ ToolResult(
+ status="reconciled",
+ payload={
+ "document_id": doc.document_id,
+ "pending_count": len(doc_pending_assets),
+ "placed_count": sum(
+ 1 for asset in doc_pending_assets
+ if str(asset.get("chunk_id") or "") in {
+ str(row.get("chunk_id") or "")
+ for row in root.flatten_chunk_rows()
+ }
+ ),
+ },
+ ),
+ decision_reason=f"deferred_reconcile_{doc.source_file_name}",
+ )
+
+ def _record_navigation_step(
+ self,
+ *,
+ doc: CandidateDoc,
+ scope: str | list[str] | None,
+ depth: int,
+ action: str,
+ asset_tools: list[str],
+ step_node: DocTreeNode,
+ drill_paths: list[dict[str, Any]],
+ ) -> None:
+ if self._trace_enabled:
+ self._trace.record_step(
+ "navigate_step",
+ ToolResult(
+ status=f"{action.lower()}" + (" (content)" if step_node.has_content() else ""),
+ payload={
+ "document_id": doc.document_id,
+ "scope": scope if isinstance(scope, str) else (scope or "root"),
+ "depth": depth,
+ "action": action,
+ "asset_tools": asset_tools,
+ "outline_count": len(step_node.outline_items),
+ "leaf_count": len(step_node.leaf_content),
+ "pending_drills": len(drill_paths),
+ },
+ ),
+ decision_reason=f"nav_d{depth}_{doc.source_file_name}",
+ )
+ scope_log = scope if isinstance(scope, str) else (", ".join(scope) if scope else "root")
+ logger.info(
+ f" agentic step {self._state.step_count}: navigate_step "
+ f'doc="{doc.source_file_name}" scope={scope_log} '
+ f"depth={depth} action={action} tools={asset_tools} "
+ f"outline={len(step_node.outline_items)} "
+ f"leaves={len(step_node.leaf_content)} "
+ f"drills={len(drill_paths)}"
+ )
+
+
+def _merge_step_node(parent_node: DocTreeNode, step_node: DocTreeNode) -> None:
+ parent_node.outline_items = step_node.outline_items
+ for leaf_path, chunks in step_node.leaf_content.items():
+ parent_node.add_leaf_chunks(leaf_path, chunks)
+ parent_node.confidence = step_node.confidence
+
+
+def _update_excluded_leaf_paths(
+ doc_exclude: set[str],
+ step_node: DocTreeNode,
+ drill_paths: list[dict[str, Any]],
+) -> None:
+ drill_path_set = {str(selection["path"]) for selection in drill_paths}
+ for leaf_path in step_node.leaf_content:
+ if leaf_path not in drill_path_set:
+ doc_exclude.add(leaf_path)
+
+
+def _queue_drill_paths(
+ pending: list[tuple[str | list[str] | None, DocTreeNode, int]],
+ parent_node: DocTreeNode,
+ drill_paths: list[dict[str, Any]],
+ depth: int,
+) -> None:
+ if not drill_paths:
+ return
+ for selection in drill_paths:
+ child = DocTreeNode(scope_path=selection["path"])
+ parent_node.children[selection["path"]] = child
+ batch_scope = [selection["path"] for selection in drill_paths]
+ pending.append((batch_scope, parent_node, depth + 1))
diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence.py
new file mode 100644
index 000000000..3d54fb47b
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/evidence.py
@@ -0,0 +1,325 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import RetrievalHitStat
+from shared.services.retrieval.agentic.budget import BudgetLedger
+from shared.services.retrieval.agentic.types import DocTreeNode
+from shared.services.retrieval.assets import build_retrieval_asset_url_map
+from shared.services.retrieval.hit_stats_service import compute_importance_score
+from shared.utils.token_estimate import estimate_tokens
+
+
+def with_context_prompt_projection(
+ snapshot: dict[str, object],
+ *,
+ prompt_tokens: int,
+) -> dict[str, object]:
+ projected: dict[str, object] = dict(snapshot)
+ context_raw = projected.get("context") or {}
+ if not isinstance(context_raw, dict):
+ return projected
+
+ context = dict(context_raw)
+ used = int(context.get("used", 0) or 0)
+ reserved = int(context.get("reserved", 0) or 0)
+ capacity = int(context.get("capacity", 0) or 0)
+ projected_used = min(capacity, used + max(int(prompt_tokens), 0))
+ projected_remaining = max(capacity - projected_used - reserved, 0)
+ context.update(
+ {
+ "used_projected_before_answer": projected_used,
+ "answer_prompt_estimate": max(int(prompt_tokens), 0),
+ "remaining": projected_remaining,
+ "used_pct": 100
+ if capacity <= 0
+ else min(100, int(round((projected_used + reserved) * 100 / capacity))),
+ }
+ )
+ if projected_remaining <= 0:
+ context["status"] = "EXHAUSTED"
+ elif context["used_pct"] >= 80:
+ context["status"] = "CRITICAL"
+ elif context["used_pct"] >= 50:
+ context["status"] = "TIGHT"
+ else:
+ context["status"] = "HEALTHY"
+ projected["context"] = context
+ return projected
+
+
+def collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]:
+ media: list[dict[str, Any]] = []
+ for chunks in node.leaf_content.values():
+ for chunk in chunks:
+ chunk_type = (
+ chunk.get("chunk_type") or chunk.get("type") or ""
+ ).strip().lower()
+ if chunk_type in ("image", "table"):
+ media.append(chunk)
+ for child in node.children.values():
+ media.extend(collect_media_chunks(child))
+ return media
+
+
+def collect_media_chunks_all(
+ doc_trees: dict[str, DocTreeNode],
+) -> list[dict[str, Any]]:
+ media: list[dict[str, Any]] = []
+ for tree in doc_trees.values():
+ media.extend(collect_media_chunks(tree))
+ return media
+
+
+async def build_asset_url_map(
+ media_chunks: list[dict[str, Any]],
+) -> dict[str, str]:
+ return await build_retrieval_asset_url_map(
+ media_chunks,
+ log_context="agentic evidence",
+ )
+
+
+def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]:
+ paths = set(node.leaf_content.keys())
+ for child in node.children.values():
+ paths.update(_collect_all_leaf_paths(child))
+ return paths
+
+
+def _collect_visible_paths(node: DocTreeNode) -> set[str]:
+ paths = {item["path"] for item in node.outline_items if item.get("path")}
+ for child in node.children.values():
+ paths.update(_collect_visible_paths(child))
+ return paths
+
+
+def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None:
+ parts = path.split(" / ")
+ for i in range(len(parts) - 1, 0, -1):
+ ancestor = " / ".join(parts[:i])
+ if ancestor in target_paths:
+ return ancestor
+ return None
+
+
+def reconcile_deferred_assets(
+ tree: DocTreeNode,
+ pending_assets: list[dict],
+) -> None:
+ final_paths = _collect_all_leaf_paths(tree)
+ visible_paths = _collect_visible_paths(tree)
+ all_target_paths = final_paths | visible_paths
+
+ if not all_target_paths:
+ return
+
+ existing_ids = {
+ str(row.get("chunk_id") or "")
+ for row in tree.flatten_chunk_rows()
+ if row.get("chunk_id")
+ }
+
+ placed = 0
+ ancestor_placed = 0
+ for asset in pending_assets:
+ chunk_id = str(asset.get("chunk_id") or "")
+ if chunk_id and chunk_id in existing_ids:
+ continue
+
+ owner_path = asset.get("owner_section_path") or asset.get("section_path")
+ if not owner_path:
+ continue
+
+ target_path = owner_path if owner_path in all_target_paths else None
+ if target_path is None:
+ target_path = _find_closest_ancestor(owner_path, all_target_paths)
+ if target_path:
+ ancestor_placed += 1
+
+ if target_path is None:
+ continue
+
+ tree.add_leaf_chunks(target_path, [asset])
+ if chunk_id:
+ existing_ids.add(chunk_id)
+ placed += 1
+
+ if placed:
+ tree.reparent_leaf_content()
+ logger.info(
+ f" deferred asset reconcile: {placed}/{len(pending_assets)} "
+ f"assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths "
+ f"(ancestor_fallback={ancestor_placed})"
+ )
+
+
+async def render_evidence(
+ db: AsyncSession,
+ doc_trees: dict[str, DocTreeNode],
+ doc_id_to_name: dict[str, str],
+) -> str:
+ del db
+
+ from shared.services.retrieval.agentic.evidence_renderer import render_unified_doc_tree
+
+ asset_url_map = await build_asset_url_map(collect_media_chunks_all(doc_trees))
+
+ evidence_parts: list[str] = []
+ for doc_id, doc_tree in doc_trees.items():
+ if doc_tree.has_content():
+ doc_name = doc_id_to_name.get(doc_id, doc_id)
+ rendered = render_unified_doc_tree(
+ doc_tree,
+ doc_name,
+ asset_lookup=asset_url_map,
+ )
+ if rendered.strip():
+ evidence_parts.append(rendered)
+
+ return "\n\n".join(evidence_parts) if evidence_parts else "(no evidence collected)"
+
+
+def _iter_leaf_content(node: DocTreeNode):
+ for path, chunks in node.leaf_content.items():
+ yield path, chunks
+ for child in node.children.values():
+ yield from _iter_leaf_content(child)
+
+
+def _collect_confidences(node: DocTreeNode) -> dict[str, float]:
+ values = dict(node.confidence)
+ for child in node.children.values():
+ for path, score in _collect_confidences(child).items():
+ values[path] = max(values.get(path, 0.0), score)
+ return values
+
+
+def _pop_leaf_path(node: DocTreeNode, path: str) -> bool:
+ if path in node.leaf_content:
+ node.leaf_content.pop(path)
+ return True
+ for child in node.children.values():
+ if _pop_leaf_path(child, path):
+ return True
+ return False
+
+
+def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int:
+ text = "\n".join(str(chunk.get("content") or "") for chunk in chunks)
+ return estimate_tokens(text)
+
+
+async def _fetch_importance_norm_scores(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ chunk_ids: list[str],
+) -> dict[str, float]:
+ if not chunk_ids:
+ return {}
+ stmt = (
+ select(
+ RetrievalHitStat.chunk_id,
+ RetrievalHitStat.hit_count,
+ RetrievalHitStat.last_hit_at,
+ RetrievalHitStat.created_at,
+ )
+ .where(RetrievalHitStat.user_id == user_id)
+ .where(RetrievalHitStat.namespace == namespace)
+ .where(RetrievalHitStat.hit_kind == "chunk")
+ .where(RetrievalHitStat.chunk_id.in_(chunk_ids))
+ )
+ result = await db.execute(stmt)
+ scores: dict[str, float] = {}
+ for chunk_id, hit_count, last_hit_at, created_at in result.all():
+ if chunk_id and last_hit_at and created_at:
+ scores[str(chunk_id)] = compute_importance_score(
+ hit_count,
+ last_hit_at,
+ created_at,
+ )
+ return scores
+
+
+async def trim_evidence_to_budget(
+ db: AsyncSession,
+ *,
+ doc_trees: dict[str, DocTreeNode],
+ doc_id_to_name: dict[str, str],
+ context_remaining: int,
+ user_id: str,
+ namespace: str,
+ ledger: BudgetLedger | None,
+ safety_margin: float = 0.9,
+) -> str:
+ full_text = await render_evidence(db, doc_trees, doc_id_to_name)
+ target = int(max(context_remaining, 0) * safety_margin)
+ if estimate_tokens(full_text) <= target:
+ return full_text
+
+ candidates: list[tuple[str, str, tuple[float, float, float], int]] = []
+ for doc_id, tree in doc_trees.items():
+ confidence = _collect_confidences(tree)
+ for path, chunks in _iter_leaf_content(tree):
+ chunk_ids = [
+ str(chunk.get("chunk_id"))
+ for chunk in chunks
+ if chunk.get("chunk_id")
+ ]
+ importance = 0.0
+ importance_scores = await _fetch_importance_norm_scores(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ chunk_ids=chunk_ids,
+ )
+ if importance_scores:
+ importance = max(importance_scores.values())
+ discovery_score = (
+ float(chunks[0].get("discovery_score", 0.0) or 0.0)
+ if chunks
+ else 0.0
+ )
+ score = (
+ float(confidence.get(path, 0.0) or 0.0),
+ discovery_score,
+ importance,
+ )
+ candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks)))
+
+ current_estimate = estimate_tokens(full_text)
+ removed: list[dict[str, Any]] = []
+ for doc_id, path, score, token_estimate in sorted(
+ candidates,
+ key=lambda item: (item[2], -item[3]),
+ ):
+ if current_estimate <= target:
+ break
+ if _pop_leaf_path(doc_trees[doc_id], path):
+ confidence_score, discovery_score, importance_score = score
+ removed.append(
+ {
+ "document_id": doc_id,
+ "document_name": doc_id_to_name.get(doc_id, doc_id),
+ "path": path,
+ "confidence_score": round(confidence_score, 4),
+ "discovery_score": round(discovery_score, 4),
+ "importance_score": round(importance_score, 4),
+ "token_estimate": token_estimate,
+ }
+ )
+ current_estimate = max(current_estimate - token_estimate, 0)
+
+ if ledger is not None:
+ ledger.trimmed_paths.extend(removed)
+ logger.info(
+ f" agentic.trim_evidence: removed={len(removed)} "
+ f"est_tokens={current_estimate} target={target}"
+ )
+ return await render_evidence(db, doc_trees, doc_id_to_name)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py
new file mode 100644
index 000000000..22326d484
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py
@@ -0,0 +1,182 @@
+"""Render agentic document trees into evidence text."""
+from __future__ import annotations
+
+from typing import Any, cast
+
+from shared.services.retrieval.agentic.types import DocTreeNode
+
+
+def render_unified_doc_tree(
+ node: DocTreeNode,
+ doc_name: str,
+ depth: int = 0,
+ asset_lookup: dict[str, str] | None = None,
+) -> str:
+ """Render a DocTreeNode as one coherent hierarchy."""
+ parts: list[str] = []
+ indent = " " * depth
+
+ if depth == 0:
+ parts.append(f"【文档】{doc_name}\n")
+
+ child_prefixes = set(node.children.keys())
+
+ def min_sort(path: str) -> float:
+ chunks = node.leaf_content.get(path, [])
+ return min((chunk.get("sort_order") or float("inf") for chunk in chunks), default=float("inf"))
+
+ render_queue: list[tuple[float, str, dict | str]] = []
+ outline_paths: set[str] = set()
+ outline_position = 0.0
+
+ for item in node.outline_items:
+ path = item.get("path", "")
+ if any(path.startswith(child_prefix + " / ") for child_prefix in child_prefixes):
+ continue
+ outline_paths.add(path)
+
+ if path in node.leaf_content or path in node.children:
+ sort_key = min_sort(path) if path in node.leaf_content else outline_position
+ else:
+ sort_key = outline_position
+ outline_position = max(outline_position, sort_key) + 0.001
+
+ render_queue.append((sort_key, "outline", item))
+
+ for path in node.leaf_content:
+ if path not in outline_paths:
+ render_queue.append((min_sort(path), "orphan_leaf", path))
+
+ for path in node.children:
+ if path not in outline_paths:
+ render_queue.append((float("inf"), "orphan_child", path))
+
+ render_queue.sort(key=lambda item: item[0])
+
+ for _sort_key, render_type, data in render_queue:
+ if render_type == "outline":
+ item = cast(dict, data)
+ path = item.get("path", "")
+ title = item.get("title", "")
+ is_leaf = item.get("is_leaf", False)
+ level = item.get("level", 1)
+ leaf_tag = " [Leaf]" if is_leaf else ""
+
+ level_tag = f"[L{level}] " if level else ""
+ if level <= 1:
+ parts.append(f"{indent}▸ {level_tag}{title}{leaf_tag}")
+ else:
+ parts.append(f"{indent}└ {level_tag}{title}{leaf_tag}")
+
+ sub_indent = indent + " "
+ if path in node.children:
+ child = node.children[path]
+ if path in node.leaf_content:
+ render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
+ child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup)
+ if child_text.strip():
+ parts.append(child_text)
+ elif path in node.leaf_content:
+ render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
+
+ elif render_type == "orphan_leaf":
+ path = cast(str, data)
+ title = path.rsplit(" / ", 1)[-1] if " / " in path else path
+ parts.append(f"{indent}▸ [Leaf] {title}")
+ render_leaf_chunks(parts, node.leaf_content[path], indent + " ", asset_lookup=asset_lookup)
+
+ elif render_type == "orphan_child":
+ path = cast(str, data)
+ title = path.rsplit(" / ", 1)[-1] if " / " in path else path
+ parts.append(f"{indent}▸ {title} [DrillDown]")
+ child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup)
+ if child_text.strip():
+ parts.append(child_text)
+
+ return "\n".join(parts)
+
+
+def render_leaf_chunks(
+ parts: list[str],
+ chunks: list[dict[str, Any]],
+ indent: str,
+ asset_lookup: dict[str, str] | None = None,
+) -> None:
+ chunk_by_id = {
+ chunk.get("chunk_id", ""): chunk
+ for chunk in chunks
+ if chunk.get("chunk_id")
+ }
+ rendered_ids: set[str] = set()
+
+ for chunk in chunks:
+ chunk_id = chunk.get("chunk_id", "")
+ if chunk_id and chunk_id in rendered_ids:
+ continue
+
+ chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "text").strip().lower()
+ if chunk_type in ("image", "table"):
+ continue
+
+ if chunk_id:
+ rendered_ids.add(chunk_id)
+
+ content = str(chunk.get("content", "")).strip()
+ for connection in (chunk.get("chunk_metadata") or {}).get("connect_to") or []:
+ target = chunk_by_id.get(connection.get("target", ""))
+ if not target:
+ continue
+ target_id = target.get("chunk_id", "")
+ target_type = (target.get("chunk_type") or target.get("type") or "").strip().lower()
+ ref_str = connection.get("ref", "")
+ if not ref_str or ref_str not in content:
+ continue
+
+ if target_id:
+ rendered_ids.add(target_id)
+
+ if target_type == "table":
+ table_html = str(target.get("content", "")).strip()
+ content = content.replace(ref_str, f"\n[表格内容]\n{table_html}\n")
+ elif target_type == "image":
+ file_path = target.get("file_path") or ""
+ image_description = str(target.get("content", "")).strip()
+ if ref_str in image_description:
+ image_description = image_description.replace(ref_str, "").strip()
+ asset_url = (asset_lookup or {}).get(target_id, "") if target_id else ""
+ display_ref = asset_url or file_path
+ if display_ref:
+ content = content.replace(ref_str, f"\n[图片: {display_ref}]\n{image_description}\n")
+ elif image_description:
+ content = content.replace(ref_str, f"\n[图片描述]\n{image_description}\n")
+
+ for line in content.split("\n"):
+ if line.strip():
+ parts.append(f"{indent}┈ {line}")
+
+ for chunk in chunks:
+ chunk_id = chunk.get("chunk_id", "")
+ if chunk_id and chunk_id in rendered_ids:
+ continue
+ if chunk_id:
+ rendered_ids.add(chunk_id)
+
+ chunk_type = (chunk.get("chunk_type") or chunk.get("type") or "").strip().lower()
+ if chunk_type == "image":
+ file_path = chunk.get("file_path") or ""
+ image_description = str(chunk.get("content", "")).strip()
+ asset_url = (asset_lookup or {}).get(chunk_id, "") if chunk_id else ""
+ display_ref = asset_url or file_path
+ if display_ref:
+ parts.append(f"{indent}┈ [图片: {display_ref}]")
+ if image_description:
+ for line in image_description.split("\n"):
+ if line.strip():
+ parts.append(f"{indent}┈ {line}")
+ elif chunk_type == "table":
+ table_html = str(chunk.get("content", "")).strip()
+ parts.append(f"{indent}┈ [表格内容]")
+ if table_html:
+ for line in table_html.split("\n"):
+ if line.strip():
+ parts.append(f"{indent}┈ {line}")
diff --git a/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py b/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py
new file mode 100644
index 000000000..c0196be66
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py
@@ -0,0 +1,93 @@
+"""Knowledge-map overview for agentic document selection."""
+from __future__ import annotations
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, GraphNode
+
+
+_MAX_OVERVIEW_FILES = 50
+
+
+async def build_knowledge_map_overview(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+) -> tuple[str, dict[str, str]]:
+ """Build a file-level knowledge map overview for LLM file selection."""
+ doc_stmt = (
+ select(Document)
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == "active")
+ .where(Document.current_job_result_id.is_not(None))
+ .order_by(Document.updated_at.desc())
+ .limit(_MAX_OVERVIEW_FILES)
+ )
+ doc_result = await db.execute(doc_stmt)
+ documents = list(doc_result.scalars())
+
+ if not documents:
+ return "(empty)", {}
+
+ doc_ids = [document.document_id for document in documents]
+ doc_id_to_name = {
+ document.document_id: (document.source_file_name or document.document_id)
+ for document in documents
+ }
+
+ chunk_stats_stmt = (
+ select(
+ DocumentChunk.document_id,
+ func.count(DocumentChunk.id).label("chunk_count"),
+ func.count(func.nullif(DocumentChunk.chunk_type, "text")).label("media_count"),
+ )
+ .join(
+ Document,
+ (Document.document_id == DocumentChunk.document_id)
+ & (Document.current_job_result_id == DocumentChunk.job_result_id),
+ )
+ .where(DocumentChunk.document_id.in_(doc_ids))
+ .group_by(DocumentChunk.document_id)
+ )
+ chunk_stats_result = await db.execute(chunk_stats_stmt)
+ chunk_stats: dict[str, dict[str, int]] = {}
+ for document_id, chunk_count, media_count in chunk_stats_result.all():
+ chunk_stats[document_id] = {"total": chunk_count, "media": media_count}
+
+ graph_summary_stmt = (
+ select(GraphNode.owner_document_id, GraphNode.properties)
+ .where(GraphNode.owner_document_id.in_(doc_ids))
+ .where(GraphNode.node_kind == "document")
+ )
+ graph_summary_result = await db.execute(graph_summary_stmt)
+ doc_top_summaries: dict[str, str] = {}
+ for document_id, properties in graph_summary_result.all():
+ if not isinstance(properties, dict):
+ continue
+ top_summary = str(properties.get("top_summary") or "").strip()
+ if top_summary:
+ doc_top_summaries[document_id] = top_summary
+
+ lines: list[str] = []
+ for document in documents:
+ document_id = document.document_id
+ name = doc_id_to_name[document_id]
+ stats = chunk_stats.get(document_id, {"total": 0, "media": 0})
+ top_summary = doc_top_summaries.get(document_id, "")
+
+ line = f'- [{document_id}] {name} chunks={stats["total"]}'
+ if stats["media"] > 0:
+ line += f' media={stats["media"]}'
+ if top_summary:
+ line += f"\n top_summary:\n{indent_block(top_summary, 4)}"
+ lines.append(line)
+
+ return "\n".join(lines), doc_id_to_name
+
+
+def indent_block(text: str, spaces: int) -> str:
+ prefix = " " * spaces
+ return "\n".join(f"{prefix}{line}" for line in str(text or "").splitlines())
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py
new file mode 100644
index 000000000..98f8fd424
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py
@@ -0,0 +1,186 @@
+"""Agentic retrieval navigation tools.
+
+This Module owns document-scope navigation and post-navigation discovery
+selection. It keeps the LLM prompt, section traversal, hydration, and asset
+owner reconciliation local to the navigation seam.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.agentic.asset_tools import (
+ build_asset_tools_block,
+ count_assets_under_scope,
+)
+from shared.services.retrieval.agentic.budget import BudgetExceeded
+from shared.services.retrieval.agentic.prompts import (
+ ACTION_PROMPT,
+ format_budget_block,
+ parse_action_response,
+)
+from shared.services.retrieval.agentic.section_prompt_projection import format_items_for_llm
+from shared.services.retrieval.agentic.section_tree import load_child_sections
+from shared.services.retrieval.agentic.selection_hydration import (
+ hydrate_path_selections_into_node,
+)
+from shared.services.retrieval.agentic.types import DocTreeNode
+from shared.services.retrieval.llm_adapter import LLMFn
+
+
+async def navigate_step(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ query: str,
+ llm_fn: LLMFn,
+ user_id: str,
+ namespace: str,
+ doc_name: str = "",
+ scope_path: str | list[str] | None = None,
+ exclude_paths: set[str] | None = None,
+ revision_hint: str | None = None,
+ budget_snapshot: dict | None = None,
+) -> tuple[str, list[str], DocTreeNode, list[dict]]:
+ """Navigate one document scope and hydrate selected sections."""
+ scope_paths = (
+ scope_path if isinstance(scope_path, list)
+ else [scope_path] if scope_path
+ else []
+ )
+ scope_path_set = set(scope_paths)
+
+ empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None)
+
+ try:
+ items = await load_child_sections(
+ db,
+ document_id,
+ job_result_id,
+ scope_path,
+ exclude_paths=exclude_paths,
+ )
+ if not items:
+ return "STOP", [], empty, []
+
+ selectable = {
+ item["path"]: item for item in items if item.get("selectable", False)
+ }
+ total_images, total_tables = await count_assets_under_scope(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ scope_paths=scope_paths,
+ )
+ tools_block = build_asset_tools_block(total_images, total_tables)
+
+ items_text, overflowed = format_items_for_llm(items)
+ prompt = _build_navigation_prompt(
+ document_id=document_id,
+ doc_name=doc_name,
+ query=query,
+ scope_paths=scope_paths,
+ budget_snapshot=budget_snapshot,
+ items_text=items_text,
+ tools_block=tools_block,
+ revision_hint=revision_hint,
+ )
+
+ response = await llm_fn(prompt)
+ parsed = parse_action_response(response)
+ action = parsed["action"]
+ selected_tools = parsed["tools"]
+ selections = parsed["selections"]
+
+ scope_label = ", ".join(scope_paths) if scope_paths else "root"
+ logger.info(
+ f" navigate_step scope={scope_label}: "
+ f"action={action} tools={selected_tools} "
+ f"selections={len(selections)} selectable={len(selectable)} "
+ f"overflowed={overflowed}"
+ )
+
+ node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None)
+ node.outline_items = [item for item in items if item.get("show_summary", True)]
+
+ valid_selections = [
+ selection
+ for selection in selections
+ if selection["path"] in selectable and selection["path"] not in scope_path_set
+ ]
+
+ pending: list[dict] = []
+ path_selections: list[dict[str, Any]] = []
+ for selection in valid_selections:
+ path = selection["path"]
+ confidence = selection.get("confidence", 0.7)
+ item = selectable[path]
+ node.confidence[path] = confidence
+
+ if item.get("is_leaf"):
+ path_selections.append({
+ "path": path,
+ "confidence": confidence,
+ "hydrate_mode": "chunks",
+ })
+ else:
+ pending.append({"path": path, "confidence": confidence})
+ path_selections.append({
+ "path": path,
+ "confidence": confidence,
+ "hydrate_mode": "self_only",
+ })
+
+ await hydrate_path_selections_into_node(
+ db,
+ node=node,
+ path_selections=path_selections,
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ )
+
+ return action, selected_tools, node, pending
+
+ except BudgetExceeded:
+ raise
+ except Exception as exc:
+ logger.error(f" navigate_step failed for doc={document_id}: {exc}")
+ return "STOP", [], empty, []
+def _build_navigation_prompt(
+ *,
+ document_id: str,
+ doc_name: str,
+ query: str,
+ scope_paths: list[str],
+ budget_snapshot: dict | None,
+ items_text: str,
+ tools_block: str,
+ revision_hint: str | None,
+) -> str:
+ if not scope_paths:
+ scope_header = "Current scope: root (document top level)"
+ elif len(scope_paths) == 1:
+ scope_header = f'Current scope: navigating into "{scope_paths[0]}"'
+ else:
+ scope_header = f"Current scope: navigating into {len(scope_paths)} sections"
+
+ prompt = ACTION_PROMPT.format(
+ doc_name=doc_name or document_id,
+ doc_id=document_id,
+ scope_header=scope_header,
+ budget_block=format_budget_block(budget_snapshot),
+ items_overview=items_text,
+ query=query,
+ tools_block=tools_block,
+ )
+ if revision_hint:
+ prompt += (
+ "\n\nIMPORTANT: Previous round feedback: "
+ f'"{revision_hint}". Adjust your selections accordingly.'
+ )
+ return prompt
diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
index 8aa58964d..c8fdfa19a 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
@@ -18,446 +18,38 @@
from __future__ import annotations
import os
-import json
from typing import Any, cast
from loguru import logger
-from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.models.database.document import Document, DocumentChunk, RetrievalHitStat
-
-from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger, BudgetPoolName
+from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger
+from shared.services.retrieval.agentic.discovery_phase import (
+ register_discovery_documents,
+ run_initial_discovery,
+ select_revision_documents,
+)
+from shared.services.retrieval.agentic.document_navigation import DocumentNavigationRunner
+from shared.services.retrieval.agentic.evidence import (
+ build_asset_url_map as _build_asset_url_map,
+ collect_media_chunks_all as _collect_media_chunks_all,
+ render_evidence as _render_evidence,
+ trim_evidence_to_budget as _trim_evidence_to_budget,
+ with_context_prompt_projection as _with_context_prompt_projection,
+)
+from shared.services.retrieval.agentic.runtime import (
+ AgentLlmBudget,
+ build_config_from_env as _build_config_from_env,
+ load_budget_inventory as _load_budget_inventory,
+)
from shared.services.retrieval.agentic.trace import TraceRecorder
from shared.services.retrieval.agentic.types import (
AgentRunConfig,
AgentState,
AgenticResult,
- CandidateDoc,
- DocTreeNode,
ToolResult,
)
-from shared.services.retrieval.app_service import (
- generate_retrieval_asset_url,
- _is_client_result_artifact_ref,
-)
from shared.services.retrieval.llm_adapter import LLMFn
-from shared.services.retrieval.llm_adapter import current_llm_usage
-from shared.services.retrieval.hit_stats_service import compute_importance_score
-from shared.utils.token_estimate import estimate_tokens
-
-
-
-def _with_context_prompt_projection(
- snapshot: dict[str, object],
- *,
- prompt_tokens: int,
-) -> dict[str, object]:
- """Return a display snapshot that includes the upcoming answer prompt."""
- projected: dict[str, object] = dict(snapshot)
- context_raw = projected.get('context') or {}
- if not isinstance(context_raw, dict):
- return projected
- context = dict(context_raw)
- used = int(context.get('used', 0) or 0)
- reserved = int(context.get('reserved', 0) or 0)
- capacity = int(context.get('capacity', 0) or 0)
- projected_used = min(capacity, used + max(int(prompt_tokens), 0))
- projected_remaining = max(capacity - projected_used - reserved, 0)
- context.update({
- 'used_projected_before_answer': projected_used,
- 'answer_prompt_estimate': max(int(prompt_tokens), 0),
- 'remaining': projected_remaining,
- 'used_pct': 100 if capacity <= 0 else min(
- 100,
- int(round((projected_used + reserved) * 100 / capacity)),
- ),
- })
- if projected_remaining <= 0:
- context['status'] = 'EXHAUSTED'
- elif context['used_pct'] >= 80:
- context['status'] = 'CRITICAL'
- elif context['used_pct'] >= 50:
- context['status'] = 'TIGHT'
- else:
- context['status'] = 'HEALTHY'
- projected['context'] = context
- return projected
-
-
-
-
-
-def _collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]:
- """Recursively collect image/table chunks from a doc tree's leaf_content."""
- media: list[dict[str, Any]] = []
- for chunks in node.leaf_content.values():
- for chunk in chunks:
- ct = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower()
- if ct in ('image', 'table'):
- media.append(chunk)
- for child in node.children.values():
- media.extend(_collect_media_chunks(child))
- return media
-
-
-def _collect_media_chunks_all(doc_trees: dict[str, DocTreeNode]) -> list[dict[str, Any]]:
- """Collect media chunks from all doc trees."""
- media: list[dict[str, Any]] = []
- for tree in doc_trees.values():
- media.extend(_collect_media_chunks(tree))
- return media
-
-
-async def _build_asset_url_map(
- media_chunks: list[dict[str, Any]],
-) -> dict[str, str]:
- """Generate presigned asset URLs for media chunks.
-
- Uses the same ``generate_retrieval_asset_url`` as ``_to_public_response``
- in ``app_service.py`` — no separate logic.
- """
- url_map: dict[str, str] = {}
- for chunk in media_chunks:
- chunk_id = str(chunk.get('chunk_id') or '').strip()
- file_path = chunk.get('file_path') or ''
- job_id = chunk.get('job_id') or ''
- if not chunk_id or not file_path or not job_id:
- continue
- if not _is_client_result_artifact_ref(file_path):
- continue
- try:
- url = await generate_retrieval_asset_url(
- job_id=str(job_id),
- artifact_ref=str(file_path),
- )
- if url:
- url_map[chunk_id] = url
- except Exception as e:
- logger.warning(f'Failed to generate asset URL for {chunk_id} (ignored): {e}')
- return url_map
-
-
-def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]:
- """Recursively collect all leaf_content keys across the entire tree."""
- paths = set(node.leaf_content.keys())
- for child in node.children.values():
- paths.update(_collect_all_leaf_paths(child))
- return paths
-
-
-def _collect_visible_paths(node: DocTreeNode) -> set[str]:
- """Collect all outline_items paths across the entire tree.
-
- These are sections that are "visible" in the rendered tree (shown to the
- LLM during navigation) even if no chunks have been hydrated into them yet.
- Used as fallback targets for asset reconciliation.
- """
- paths = {item['path'] for item in node.outline_items if item.get('path')}
- for child in node.children.values():
- paths.update(_collect_visible_paths(child))
- return paths
-
-
-def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None:
- """Walk up a section path to find the closest ancestor in target_paths.
-
- Example: path="kb/file/Ch1/S1.1/S1.1.1", target_paths={"kb/file/Ch1/S1.1"}
- → returns "kb/file/Ch1/S1.1"
-
- Uses the ' / ' separator convention from the section path format.
- """
- parts = path.split(' / ')
- # Walk from most specific to least specific (skip the full path itself)
- for i in range(len(parts) - 1, 0, -1):
- ancestor = ' / '.join(parts[:i])
- if ancestor in target_paths:
- return ancestor
- return None
-
-
-def _reconcile_deferred_assets(
- tree: DocTreeNode,
- pending_assets: list[dict],
-) -> None:
- """Place collected assets into the tree based on final navigated paths.
-
- Called ONCE after the entire BFS + discovery merge completes for a
- document. Asset placement uses a two-tier strategy:
-
- 1. **Exact match**: If the asset's ``owner_section_path`` matches a
- leaf_content key, place directly (existing behavior).
- 2. **Closest visible ancestor**: If exact match fails, walk up the
- owner_section_path hierarchy to find the nearest ancestor that
- appears in either leaf_content or outline_items. This handles
- the case where the LLM stopped navigation early (e.g. at root)
- but still requested images/tables — assets at L3 get attributed
- to the visible L2 section on their path.
- """
- final_paths = _collect_all_leaf_paths(tree)
- visible_paths = _collect_visible_paths(tree)
- all_target_paths = final_paths | visible_paths
-
- if not all_target_paths:
- return
-
- # Collect existing chunk_ids to avoid duplicates
- existing_ids = {
- str(row.get('chunk_id') or '')
- for row in tree.flatten_chunk_rows()
- if row.get('chunk_id')
- }
-
- placed = 0
- ancestor_placed = 0
- for asset in pending_assets:
- chunk_id = str(asset.get('chunk_id') or '')
- if chunk_id and chunk_id in existing_ids:
- continue # already in tree via hydrate_connected_target_rows
-
- owner_path = (
- asset.get('owner_section_path')
- or asset.get('section_path')
- )
- if not owner_path:
- continue
-
- # Tier 1: exact match in leaf_content or visible outline
- target_path = owner_path if owner_path in all_target_paths else None
-
- # Tier 2: closest visible ancestor fallback
- if target_path is None:
- target_path = _find_closest_ancestor(owner_path, all_target_paths)
- if target_path:
- ancestor_placed += 1
-
- if target_path is None:
- continue # no visible ancestor → discard
-
- # Place into root; reparent_leaf_content will move to correct child
- tree.add_leaf_chunks(target_path, [asset])
- if chunk_id:
- existing_ids.add(chunk_id)
- placed += 1
-
- if placed:
- tree.reparent_leaf_content()
- logger.info(
- f' deferred asset reconcile: {placed}/{len(pending_assets)} '
- f'assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths '
- f'(ancestor_fallback={ancestor_placed})'
- )
-
-
-def _build_config_from_env() -> AgentRunConfig:
- """Read agent config from environment, with sensible defaults."""
- return AgentRunConfig(
- max_revisions=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_REVISIONS', '2')),
- max_nav_depth=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_NAV_DEPTH', '3')),
- latency_budget_ms=int(os.environ.get('RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS', '12000')),
- token_budget_total=int(os.environ.get('RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL', '40000')),
- planning_ratio=float(os.environ.get('RETRIEVAL_AGENTIC_PLANNING_RATIO', '0.5')),
- bootstrap_budget=int(os.environ.get('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', '2000')),
- per_doc_min_share=int(os.environ.get('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', '1500')),
- inventory_aware=os.environ.get('RETRIEVAL_AGENTIC_INVENTORY_AWARE', 'true') == 'true',
- )
-
-
-def _stringify_llm_input(prompt: Any) -> str:
- if isinstance(prompt, str):
- return prompt
- try:
- return json.dumps(prompt, ensure_ascii=False, default=str)
- except Exception:
- return str(prompt)
-
-
-async def _render_evidence(
- db: AsyncSession,
- doc_trees: dict[str, DocTreeNode],
- doc_id_to_name: dict[str, str],
-) -> str:
- """Render unified evidence text from doc trees.
-
- Discovery paths are now handled by ``discovery_select_step`` in Phase 2
- and merged into doc_trees — no separate fallback needed.
- """
- from shared.services.retrieval.agent_navigate import render_unified_doc_tree
-
- # Build asset URL map for all media chunks (images/tables)
- # — same pattern as _to_public_response in app_service.py
- all_media_chunks: list[dict[str, Any]] = []
- for doc_tree in doc_trees.values():
- all_media_chunks.extend(_collect_media_chunks(doc_tree))
- asset_url_map = await _build_asset_url_map(all_media_chunks)
-
- # Render unified evidence from doc trees
- evidence_parts: list[str] = []
- for doc_id, doc_tree in doc_trees.items():
- if doc_tree.has_content():
- doc_name = doc_id_to_name.get(doc_id, doc_id)
- rendered = render_unified_doc_tree(doc_tree, doc_name, asset_lookup=asset_url_map)
- if rendered.strip():
- evidence_parts.append(rendered)
-
- return '\n\n'.join(evidence_parts) if evidence_parts else '(no evidence collected)'
-
-
-async def _load_budget_inventory(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- exclude_document_ids: list[str],
-) -> tuple[int, int, dict[str, int]]:
- stmt = (
- select(Document.document_id, func.count(DocumentChunk.id))
- .join(
- DocumentChunk,
- (DocumentChunk.document_id == Document.document_id)
- & (DocumentChunk.job_result_id == Document.current_job_result_id),
- )
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .group_by(Document.document_id)
- )
- if exclude_document_ids:
- stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
-
- result = await db.execute(stmt)
- doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()}
- return sum(doc_chunks.values()), len(doc_chunks), doc_chunks
-
-
-def _iter_leaf_content(node: DocTreeNode):
- for path, chunks in node.leaf_content.items():
- yield path, chunks
- for child in node.children.values():
- yield from _iter_leaf_content(child)
-
-
-def _collect_confidences(node: DocTreeNode) -> dict[str, float]:
- values = dict(node.confidence)
- for child in node.children.values():
- for path, score in _collect_confidences(child).items():
- values[path] = max(values.get(path, 0.0), score)
- return values
-
-
-def _pop_leaf_path(node: DocTreeNode, path: str) -> bool:
- if path in node.leaf_content:
- node.leaf_content.pop(path)
- return True
- for child in node.children.values():
- if _pop_leaf_path(child, path):
- return True
- return False
-
-
-def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int:
- text = '\n'.join(str(chunk.get('content') or '') for chunk in chunks)
- return estimate_tokens(text)
-
-
-async def _fetch_importance_norm_scores(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- chunk_ids: list[str],
-) -> dict[str, float]:
- if not chunk_ids:
- return {}
- stmt = (
- select(
- RetrievalHitStat.chunk_id,
- RetrievalHitStat.hit_count,
- RetrievalHitStat.last_hit_at,
- RetrievalHitStat.created_at,
- )
- .where(RetrievalHitStat.user_id == user_id)
- .where(RetrievalHitStat.namespace == namespace)
- .where(RetrievalHitStat.hit_kind == 'chunk')
- .where(RetrievalHitStat.chunk_id.in_(chunk_ids))
- )
- result = await db.execute(stmt)
- scores: dict[str, float] = {}
- for chunk_id, hit_count, last_hit_at, created_at in result.all():
- if chunk_id and last_hit_at and created_at:
- scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at)
- return scores
-
-
-async def _trim_evidence_to_budget(
- db: AsyncSession,
- *,
- doc_trees: dict[str, DocTreeNode],
- doc_id_to_name: dict[str, str],
- context_remaining: int,
- user_id: str,
- namespace: str,
- ledger: BudgetLedger | None,
- safety_margin: float = 0.9,
-) -> str:
- full_text = await _render_evidence(db, doc_trees, doc_id_to_name)
- target = int(max(context_remaining, 0) * safety_margin)
- if estimate_tokens(full_text) <= target:
- return full_text
-
- candidates: list[tuple[str, str, tuple[float, float, float], int]] = []
- for doc_id, tree in doc_trees.items():
- confidence = _collect_confidences(tree)
- for path, chunks in _iter_leaf_content(tree):
- chunk_ids = [
- str(chunk.get('chunk_id'))
- for chunk in chunks
- if chunk.get('chunk_id')
- ]
- importance = 0.0
- importance_scores = await _fetch_importance_norm_scores(
- db,
- user_id=user_id,
- namespace=namespace,
- chunk_ids=chunk_ids,
- )
- if importance_scores:
- importance = max(importance_scores.values())
- discovery_score = (
- float(chunks[0].get('discovery_score', 0.0) or 0.0)
- if chunks else 0.0
- )
- score = (float(confidence.get(path, 0.0) or 0.0), discovery_score, importance)
- candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks)))
-
- current_estimate = estimate_tokens(full_text)
- removed: list[dict[str, Any]] = []
- for doc_id, path, _score, token_estimate in sorted(
- candidates,
- key=lambda item: (item[2], -item[3]),
- ):
- if current_estimate <= target:
- break
- if _pop_leaf_path(doc_trees[doc_id], path):
- confidence_score, discovery_score, importance_score = _score
- removed.append({
- 'document_id': doc_id,
- 'document_name': doc_id_to_name.get(doc_id, doc_id),
- 'path': path,
- 'confidence_score': round(confidence_score, 4),
- 'discovery_score': round(discovery_score, 4),
- 'importance_score': round(importance_score, 4),
- 'token_estimate': token_estimate,
- })
- current_estimate = max(current_estimate - token_estimate, 0)
-
- if ledger is not None:
- ledger.trimmed_paths.extend(removed)
- logger.info(
- f' agentic.trim_evidence: removed={len(removed)} '
- f'est_tokens={current_estimate} target={target}'
- )
- return await _render_evidence(db, doc_trees, doc_id_to_name)
class RetrievalAgent:
@@ -477,82 +69,6 @@ class RetrievalAgent:
If ``llm_fn`` is None, the run returns discovery-only results.
"""
- async def _call_llm_with_budget(
- self,
- state: AgentState,
- llm_fn: LLMFn,
- prompt: Any,
- *,
- pool: BudgetPoolName,
- doc_id: str | None = None,
- priority: str = 'normal',
- ) -> str:
- ledger = state.ledger
- if ledger is None:
- return await llm_fn(prompt)
-
- prompt_text = _stringify_llm_input(prompt)
- est = estimate_tokens(prompt_text)
- reserved = await ledger.try_reserve(
- pool,
- est,
- doc_id=doc_id,
- priority='low' if priority == 'low' else 'normal',
- )
- if not reserved:
- raise BudgetExceeded(f'{pool} budget exhausted')
-
- try:
- response = await llm_fn(prompt)
- except Exception:
- await ledger.refund(pool, est=est, doc_id=doc_id)
- raise
-
- usage = current_llm_usage.get() or {}
- actual = int(usage.get('prompt_tokens') or est)
- await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id)
- return response
-
- def _budgeted_doc_llm_fn(
- self,
- state: AgentState,
- llm_fn: LLMFn,
- *,
- doc_id: str,
- depth: int,
- ) -> LLMFn:
- async def _call(prompt):
- return await self._call_llm_with_budget(
- state,
- llm_fn,
- prompt,
- pool='planning',
- doc_id=doc_id,
- priority='low' if depth >= 2 else 'normal',
- )
-
- return _call
-
- def _budgeted_discovery_llm_fn(
- self,
- state: AgentState,
- llm_fn: LLMFn,
- *,
- doc_id: str,
- low_priority: bool,
- ) -> LLMFn:
- async def _call(prompt):
- return await self._call_llm_with_budget(
- state,
- llm_fn,
- prompt,
- pool='planning',
- doc_id=doc_id,
- priority='low' if low_priority else 'normal',
- )
-
- return _call
-
async def run(
self,
db: AsyncSession,
@@ -581,7 +97,6 @@ async def run(
errors are captured in trace and the best available result
is returned.
"""
- from shared.services.retrieval.agentic import tools
from shared.services.retrieval.agentic.policy import (
attempt_answer,
estimate_attempt_answer_prompt_tokens,
@@ -636,135 +151,32 @@ async def run(
if llm_fn is None:
logger.warning('agentic: no llm_fn provided — running discovery-only mode')
- planning_llm_fn: LLMFn | None = None
bootstrap_llm_fn: LLMFn | None = None
context_llm_fn: LLMFn | None = None
+ llm_budget = AgentLlmBudget(state)
if llm_fn is not None:
- base_llm_fn = llm_fn
-
- async def _planning_llm_call(prompt):
- return await self._call_llm_with_budget(
- state, base_llm_fn, prompt, pool='planning'
- )
-
- async def _bootstrap_llm_call(prompt):
- return await self._call_llm_with_budget(
- state, base_llm_fn, prompt, pool='bootstrap'
- )
-
- async def _context_llm_call(prompt):
- return await self._call_llm_with_budget(
- state, base_llm_fn, prompt, pool='context'
- )
-
- planning_llm_fn = _planning_llm_call
- bootstrap_llm_fn = _bootstrap_llm_call
- context_llm_fn = _context_llm_call
-
- # Shared kwargs for bottom_discovery
- discovery_kwargs: dict[str, Any] = {
- 'user_id': user_id,
- 'namespace': namespace,
- 'query': query,
- 'top_k': top_k,
- 'exclude_document_ids': exclude_document_ids,
- 'exclude_sections': exclude_sections,
- 'data_type': data_type,
- 'signal_paths': signal_paths,
- 'filter_mode': filter_mode,
- 'channels': channels,
- 'channel_weights': channel_weights,
- }
-
- # ══════════════════════════════════════════════════════════════════
- # Phase 1: Discovery + Document Selection
- # ══════════════════════════════════════════════════════════════════
- logger.info(' agentic: Phase 1 — discovery + document selection')
+ bootstrap_llm_fn = llm_budget.for_pool(llm_fn, pool='bootstrap')
+ context_llm_fn = llm_budget.for_pool(llm_fn, pool='context')
- # 1a. Bottom discovery (always runs)
- discovery_result = await tools.bottom_discovery(db, **discovery_kwargs)
- state.step_count += 1
- discovery_rows = discovery_result.payload.get('fused_rows', []) if discovery_result.status != 'error' else []
- state.discovery_top_doc_ids = discovery_result.payload.get('top_doc_ids', []) if discovery_result.status != 'error' else []
-
- if trace_enabled:
- trace.record_step(
- 'bottom_discovery', discovery_result,
- decision_reason='phase_1_mandatory',
- )
-
- logger.info(
- f' agentic step {state.step_count}: bottom_discovery '
- f'status={discovery_result.status} latency={discovery_result.latency_ms}ms'
+ discovery_rows = await run_initial_discovery(
+ db,
+ state=state,
+ trace=trace,
+ trace_enabled=trace_enabled,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ bootstrap_llm_fn=bootstrap_llm_fn,
)
- # 1b. KG document selection (requires LLM)
- if bootstrap_llm_fn is not None:
- try:
- kg_result = await tools.kg_document_select(
- db,
- user_id=user_id,
- namespace=namespace,
- query=query,
- llm_fn=bootstrap_llm_fn,
- exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)),
- budget_snapshot=state.ledger.snapshot() if state.ledger else None,
- )
- except BudgetExceeded:
- logger.info(' agentic: bootstrap budget exhausted during document selection')
- if trace_enabled:
- trace.record_budget_stop('bootstrap_exhausted')
- kg_result = ToolResult(
- status='no_confident_doc',
- payload={'reason': 'bootstrap budget exhausted'},
- )
- state.step_count += 1
-
- if trace_enabled:
- trace.record_step(
- 'kg_document_select', kg_result,
- decision_reason='phase_1_doc_selection',
- )
-
- if kg_result.status == 'selected_docs':
- for doc_data in kg_result.payload.get('candidate_docs', []):
- state.selected_docs.append(CandidateDoc(
- document_id=doc_data.get('document_id', ''),
- source_file_name=doc_data.get('source_file_name', ''),
- confidence=doc_data.get('confidence', 0.0),
- reason=doc_data.get('reason', ''),
- source=doc_data.get('source', ''),
- ))
- state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {}))
- state.doc_job_map.update(kg_result.payload.get('doc_job_map', {}))
-
- # If KG returned nothing, use discovery hints
- if not state.selected_docs and state.discovery_top_doc_ids:
- hint_ids = [d for d in state.discovery_top_doc_ids if d not in state.ever_explored_doc_ids]
- if hint_ids:
- doc_stmt = (
- select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
- .where(Document.document_id.in_(hint_ids))
- )
- doc_result = await db.execute(doc_stmt)
- for did, fname, jrid in doc_result.all():
- state.selected_docs.append(CandidateDoc(
- document_id=did,
- source_file_name=fname or did,
- confidence=0.5,
- reason='discovery_hint (KG returned 0)',
- source='discovery_hint',
- ))
- state.doc_id_to_name[did] = fname or did
- if jrid:
- state.doc_job_map[did] = jrid
-
- logger.info(
- f' agentic step {state.step_count}: kg_document_select '
- f'status={kg_result.status} docs={len(state.selected_docs)} '
- f'latency={kg_result.latency_ms}ms'
- )
-
# If no LLM or no docs selected, return discovery rows directly
if not state.selected_docs:
logger.info('agentic: no documents selected — returning discovery results')
@@ -796,38 +208,17 @@ async def _context_llm_call(prompt):
router_used='agentic_discovery_only',
)
- # ══════════════════════════════════════════════════════════════════
- # Discovery → Navigation integration
- # Group discovery_rows by document for post-BFS discovery selection
- # ══════════════════════════════════════════════════════════════════
discovery_by_doc: dict[str, list[dict[str, Any]]] = {}
for row in discovery_rows:
doc_id = row.get('document_id', '')
if doc_id:
discovery_by_doc.setdefault(doc_id, []).append(row)
- # Auto-register B-class docs (discovery-only, not selected by KG)
- selected_doc_ids = {d.document_id for d in state.selected_docs}
- for doc_id in discovery_by_doc:
- if doc_id not in selected_doc_ids and doc_id not in state.ever_explored_doc_ids:
- doc_stmt = (
- select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
- .where(Document.document_id == doc_id)
- )
- doc_result = await db.execute(doc_stmt)
- row_data = doc_result.first()
- if row_data:
- did, fname, jrid = row_data
- state.selected_docs.append(CandidateDoc(
- document_id=did,
- source_file_name=fname or did,
- confidence=0.4,
- reason='discovery_auto (not in KG selection)',
- source='discovery_auto',
- ))
- state.doc_id_to_name[did] = fname or did
- if jrid:
- state.doc_job_map[did] = jrid
+ await register_discovery_documents(
+ db,
+ state=state,
+ discovery_by_doc=discovery_by_doc,
+ )
if state.ledger is not None:
await state.ledger.allocate_doc_caps({
@@ -849,286 +240,20 @@ async def _context_llm_call(prompt):
stop_reason = 'latency_budget'
break
- # ── Phase 2: Per-Document Navigation ────────────────────────
- logger.info(
- f' agentic: Phase 2 (round {round_idx}) — '
- f'navigating {len(state.selected_docs)} documents'
+ navigation_runner = DocumentNavigationRunner(
+ db=db,
+ state=state,
+ trace=trace,
+ trace_enabled=trace_enabled,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ config=config,
+ discovery_by_doc=discovery_by_doc,
+ llm_fn=llm_fn,
+ llm_budget=llm_budget,
)
-
- for doc in state.selected_docs:
- if state.elapsed_ms >= config.latency_budget_ms:
- logger.info(' agentic: latency budget hit during Phase 2, stopping')
- break
-
- job_result_id = state.doc_job_map.get(doc.document_id, '')
- if not job_result_id:
- logger.info(f' agentic: skipping doc {doc.document_id} — no job_result_id')
- state.ever_explored_doc_ids.add(doc.document_id)
- continue
-
- doc_name = doc.source_file_name or state.doc_id_to_name.get(doc.document_id, '')
-
- # B-class docs (discovery_auto) skip BFS, go to discovery_select
- is_b_class = doc.source == 'discovery_auto'
-
- if not is_b_class:
- # Build exclude_paths for this doc from seen_section_keys
- # Starts with revision-carried paths, then accumulates
- # leaf paths hydrated during THIS BFS round to prevent
- # re-selection in deeper drill-downs.
- doc_exclude: set[str] = {
- key.split('::', 1)[1]
- for key in state.seen_section_keys
- if key.startswith(f'{doc.document_id}::')
- } if state.seen_section_keys else set()
-
- # BFS queue: (scope_path(s), parent_node, depth)
- # scope can be: None (root), str, or list[str] (multi-scope)
- root = DocTreeNode(scope_path=None)
- pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)]
- doc_pending_assets: list[dict] = [] # deferred asset reconcile
-
- while pending:
- if state.elapsed_ms >= config.latency_budget_ms:
- break
-
- scope, parent_node, depth = pending.pop(0)
- if depth >= config.max_nav_depth:
- continue
-
- if planning_llm_fn is None:
- break
- if state.ledger and state.ledger.status('planning') in ('CRITICAL', 'EXHAUSTED'):
- logger.info(' agentic: planning budget critical, ending BFS for current doc')
- break
-
- doc_llm_fn = self._budgeted_doc_llm_fn(
- state,
- cast(LLMFn, llm_fn),
- doc_id=doc.document_id,
- depth=depth,
- )
-
- # ★ Unified navigate step (supports multi-scope batching)
- try:
- action, asset_tools, step_node, drill_paths = await tools.navigate_step(
- db,
- document_id=doc.document_id,
- job_result_id=job_result_id,
- query=query,
- llm_fn=doc_llm_fn,
- user_id=user_id,
- namespace=namespace,
- doc_name=doc_name,
- scope_path=scope,
- exclude_paths=doc_exclude,
- revision_hint=revision_hint if depth == 0 else None,
- budget_snapshot=state.ledger.snapshot() if state.ledger else None,
- )
- except BudgetExceeded:
- logger.info(' agentic: planning budget exhausted during navigation')
- if trace_enabled:
- trace.record_budget_stop('planning_exhausted')
- break
- state.step_count += 1
-
- # ★ Asset collection (deferred reconcile) — runs if LLM selected tools.
- # If this navigation call selected sections, bind asset tools to
- # those selections; otherwise keep the current scope (STOP/root).
- selected_asset_scopes = list(step_node.confidence.keys())
- asset_scope = selected_asset_scopes or scope
- for asset_tool in asset_tools:
- if asset_tool not in ('FIND_IMAGES', 'FIND_TABLES'):
- continue
- asset_type = 'image' if asset_tool == 'FIND_IMAGES' else 'table'
- asset_chunks = await tools.asset_filter_step(
- db,
- document_id=doc.document_id,
- job_result_id=job_result_id,
- scope_path=asset_scope,
- asset_type=asset_type,
- )
- if asset_chunks:
- doc_pending_assets.extend(asset_chunks)
-
- scope_display = (
- asset_scope if isinstance(asset_scope, list)
- else (asset_scope or 'root')
- )
- if trace_enabled:
- trace.record_step(
- 'asset_filter_step', ToolResult(
- status='filtered' if asset_chunks else 'empty',
- payload={
- 'document_id': doc.document_id,
- 'scope': scope_display,
- 'navigation_scope': scope if isinstance(scope, str) else (scope or 'root'),
- 'asset_type': asset_type,
- 'chunks_found': len(asset_chunks) if asset_chunks else 0,
- },
- ),
- decision_reason=f'asset_r{round_idx}_d{depth}_{doc.source_file_name}',
- )
-
- logger.info(
- f' agentic step {state.step_count}: asset_filter_step '
- f'doc="{doc.source_file_name}" scope={scope_display} '
- f'type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}'
- )
-
- # Merge step result into parent node
- parent_node.outline_items = step_node.outline_items
- for leaf_path, chunks in step_node.leaf_content.items():
- parent_node.add_leaf_chunks(leaf_path, chunks)
- parent_node.confidence = step_node.confidence
-
- # Accumulate hydrated leaf paths into doc_exclude
- drill_path_set = {sel['path'] for sel in drill_paths}
- for leaf_path in step_node.leaf_content:
- if leaf_path not in drill_path_set:
- doc_exclude.add(leaf_path)
-
- # Queue non-leaf selections as a SINGLE batched item
- # (all drill paths expand simultaneously in the next call)
- if drill_paths:
- for sel in drill_paths:
- child = DocTreeNode(scope_path=sel['path'])
- parent_node.children[sel['path']] = child
- batch_scope = [sel['path'] for sel in drill_paths]
- pending.append((batch_scope, parent_node, depth + 1))
-
- # Re-parent leaf paths that belong to a child's subtree
- parent_node.reparent_leaf_content()
-
- if trace_enabled:
- trace.record_step(
- 'navigate_step', ToolResult(
- status=f'{action.lower()}' + (' (content)' if step_node.has_content() else ''),
- payload={
- 'document_id': doc.document_id,
- 'scope': scope if isinstance(scope, str) else (scope or 'root'),
- 'depth': depth,
- 'action': action,
- 'asset_tools': asset_tools,
- 'outline_count': len(step_node.outline_items),
- 'leaf_count': len(step_node.leaf_content),
- 'pending_drills': len(drill_paths),
- },
- ),
- decision_reason=f'nav_r{round_idx}_d{depth}_{doc.source_file_name}',
- )
-
- scope_log = scope if isinstance(scope, str) else (', '.join(scope) if scope else 'root')
- logger.info(
- f' agentic step {state.step_count}: navigate_step '
- f'doc="{doc.source_file_name}" scope={scope_log} '
- f'depth={depth} action={action} tools={asset_tools} '
- f'outline={len(step_node.outline_items)} '
- f'leaves={len(step_node.leaf_content)} '
- f'drills={len(drill_paths)}'
- )
- if state.ledger is not None:
- state.ledger.mark_explored(
- chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()),
- )
- else:
- # B-class: no BFS, create empty root
- root = DocTreeNode(scope_path=None)
-
- # ── Post-BFS: Discovery selection step ─────────────────────
- doc_hints = discovery_by_doc.get(doc.document_id, [])
- if doc_hints and planning_llm_fn is not None and state.elapsed_ms < config.latency_budget_ms:
- discovery_exclude_paths = {
- key.split('::', 1)[1]
- for key in root.collect_all_paths(doc.document_id)
- }
- doc_discovery_llm_fn = self._budgeted_discovery_llm_fn(
- state,
- cast(LLMFn, llm_fn),
- doc_id=doc.document_id,
- low_priority=root.has_content(),
- )
- try:
- discovery_node = await tools.discovery_select_step(
- db,
- document_id=doc.document_id,
- query=query,
- llm_fn=doc_discovery_llm_fn,
- user_id=user_id,
- namespace=namespace,
- doc_name=doc_name,
- discovery_hints=doc_hints,
- exclude_paths=discovery_exclude_paths,
- revision_hint=revision_hint,
- budget_snapshot=state.ledger.snapshot() if state.ledger else None,
- )
- except BudgetExceeded:
- logger.info(' agentic: planning budget exhausted during discovery selection')
- if trace_enabled:
- trace.record_budget_stop('planning_exhausted')
- discovery_node = DocTreeNode(scope_path=None)
- state.step_count += 1
-
- if trace_enabled:
- trace.record_step(
- 'discovery_select_step', ToolResult(
- status='selected' if discovery_node.has_content() else 'empty',
- payload={
- 'document_id': doc.document_id,
- 'hints_count': len(doc_hints),
- 'hydrated_count': len(discovery_node.leaf_content),
- },
- ),
- decision_reason=f'discovery_r{round_idx}_{doc.source_file_name}',
- )
-
- # Merge discovery results into BFS tree
- root.merge(discovery_node)
- if state.ledger is not None:
- state.ledger.mark_explored(
- chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()),
- )
-
- # ── Deferred asset reconcile ──────────────────────────────
- # Assets were collected across all BFS depths but NOT placed
- # into the tree yet. Now that the final navigated paths are
- # known (BFS + discovery), filter and place only those assets
- # whose owner path matches a navigated leaf.
- if not is_b_class and doc_pending_assets:
- # Inject doc file name as a visible root-level path, but
- # ONLY when BFS stopped at root (no children = STOP action).
- if doc_name and not root.children and not any(
- item.get('path') == doc_name for item in root.outline_items
- ):
- root.outline_items.insert(0, {'path': doc_name, 'level': 0})
- _reconcile_deferred_assets(root, doc_pending_assets)
- if trace_enabled:
- trace.record_step(
- 'deferred_asset_reconcile', ToolResult(
- status='reconciled',
- payload={
- 'document_id': doc.document_id,
- 'pending_count': len(doc_pending_assets),
- 'placed_count': sum(
- 1 for a in doc_pending_assets
- if str(a.get('chunk_id') or '') in {
- str(r.get('chunk_id') or '')
- for r in root.flatten_chunk_rows()
- }
- ),
- },
- ),
- decision_reason=f'deferred_reconcile_r{round_idx}_{doc.source_file_name}',
- )
-
- # Merge or store doc tree
- if doc.document_id in state.doc_trees:
- state.doc_trees[doc.document_id].merge(root)
- else:
- state.doc_trees[doc.document_id] = root
- state.ever_explored_doc_ids.add(doc.document_id)
- if state.ledger is not None:
- state.ledger.mark_explored(docs=1)
+ await navigation_runner.navigate_selected_documents(revision_hint=revision_hint)
# ── Phase 3: Render evidence + attempt_answer ────────────────
budget_snapshot_before_answer = state.ledger.snapshot() if state.ledger else None
@@ -1201,9 +326,7 @@ async def _context_llm_call(prompt):
]
async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
- return await self._call_llm_with_budget(
- state, cast(LLMFn, _vlm_fn), prompt, pool='context'
- )
+ return await llm_budget.call(cast(LLMFn, _vlm_fn), prompt, pool='context')
# Auto-trigger attempt_answer (VLM if images present)
try:
@@ -1270,36 +393,21 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
if bootstrap_llm_fn is None:
stop_reason = 'no_llm'
break
- try:
- kg_result = await tools.kg_document_select(
- db,
- user_id=user_id,
- namespace=namespace,
- query=query,
- llm_fn=bootstrap_llm_fn,
- exclude_document_ids=list(set(exclude_document_ids)),
- revision_hint=revision_hint,
- budget_snapshot=state.ledger.snapshot() if state.ledger else None,
- )
- except BudgetExceeded:
- logger.info(' agentic: bootstrap budget exhausted during revision doc selection')
- if trace_enabled:
- trace.record_budget_stop('bootstrap_exhausted')
- stop_reason = 'bootstrap_budget'
+ revision_stop_reason = await select_revision_documents(
+ db,
+ state=state,
+ trace=trace,
+ trace_enabled=trace_enabled,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ exclude_document_ids=exclude_document_ids,
+ bootstrap_llm_fn=bootstrap_llm_fn,
+ revision_hint=revision_hint,
+ )
+ if revision_stop_reason is not None:
+ stop_reason = revision_stop_reason
break
- state.step_count += 1
-
- if kg_result.status == 'selected_docs':
- for doc_data in kg_result.payload.get('candidate_docs', []):
- state.selected_docs.append(CandidateDoc(
- document_id=doc_data.get('document_id', ''),
- source_file_name=doc_data.get('source_file_name', ''),
- confidence=doc_data.get('confidence', 0.0),
- reason=doc_data.get('reason', ''),
- source=doc_data.get('source', ''),
- ))
- state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {}))
- state.doc_job_map.update(kg_result.payload.get('doc_job_map', {}))
if not state.selected_docs:
logger.info(' agentic: revision found no new docs — stopping')
diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py
new file mode 100644
index 000000000..7c6295913
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py
@@ -0,0 +1,223 @@
+"""Prompt templates and response parsers for agentic retrieval."""
+from __future__ import annotations
+
+import json
+import re
+from typing import Any
+
+
+FILE_SELECT_PROMPT = """\
+You are a document routing assistant.
+
+{budget_block}
+Below is a knowledge base overview showing all available documents,
+their navigation summaries, chunk counts, and media counts.
+
+=== Knowledge Base Overview ===
+{overview}
+=== End Overview ===
+
+User query: {query}
+{revision_context}
+Based on the query, select documents that may contain relevant information.
+If NO document in the knowledge base is relevant to the query, return an EMPTY array [].
+Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"]
+Do not include any explanation.
+"""
+
+
+DISCOVERY_SELECT_PROMPT = """\
+You are a document navigation assistant.
+
+Document: "{doc_name}"
+
+{budget_block}
+After navigating the document's section tree, the following section paths
+were additionally discovered via keyword and semantic search.
+They may contain relevant evidence not found through hierarchical navigation.
+
+=== Discovery Candidates ===
+{items}
+=== End Discovery Candidates ===
+
+User query: {query}
+{revision_context}
+Select section paths whose content is needed to answer the query.
+If none are relevant, return an EMPTY list [].
+
+Return ONLY a JSON object:
+{{"selections": [{{"path": "...", "confidence": }}, ...]}}
+Do not include any explanation.
+"""
+
+
+ACTION_PROMPT = """\
+You are a document navigation agent.
+
+Document: "{doc_name}" (id: {doc_id})
+
+{budget_block}
+{scope_header}
+Below is the document's section tree.
+Sections tagged [SELECT] are within the current scope and may be selected.
+Other sections are shown as structural context only (not selectable).
+Nodes marked [Leaf] have no further sub-sections.
+
+=== Section Tree ===
+{items_overview}
+=== End Section Tree ===
+
+User query: {query}
+
+=== Available Actions ===
+
+Choose ONE action:
+
+NAVIGATE — Drill into selected sections for detailed content.
+ Consider this when the query targets specific topics and you need deeper text evidence.
+ Select one or more [SELECT] sections.
+
+STOP — Current scope evidence is sufficient. No further drill-down.
+ Consider this when:
+ - The query asks for an outline, overview, or summary
+ - The query is broad/global, the tree section can fulfill it without drilling into individual sections.
+ - You have already collected enough evidence at this level.
+
+{tools_block}
+
+When action is NAVIGATE, provide selections:
+- You may ONLY select sections marked with [SELECT].
+
+When action is STOP, selections must be empty.
+
+Return ONLY a JSON object:
+{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}}
+or
+{{"action": "STOP", "tools": [...], "selections": []}}
+Do not include any explanation.
+"""
+
+
+def parse_action_response(text: str) -> dict:
+ """Parse the unified navigation response from an LLM."""
+ text = text.strip()
+ asset_tools = {"FIND_IMAGES", "FIND_TABLES"}
+ default = {"action": "NAVIGATE", "tools": [], "selections": []}
+
+ def extract(data: dict) -> dict:
+ action = str(data.get("action", "NAVIGATE")).strip().upper()
+ if action not in ("NAVIGATE", "STOP"):
+ action = "NAVIGATE"
+
+ tools_val = data.get("tools") or []
+ if isinstance(tools_val, list):
+ tools = [
+ str(tool).strip().upper()
+ for tool in tools_val
+ if str(tool).strip().upper() in asset_tools
+ ]
+ else:
+ tools = []
+
+ if action == "STOP":
+ return {"action": action, "tools": tools, "selections": []}
+
+ selections_val = data.get("selections") or []
+ selections = []
+ if isinstance(selections_val, list):
+ for selection in selections_val:
+ if isinstance(selection, dict) and selection.get("path"):
+ confidence = normalize_confidence(selection.get("confidence", 0.7))
+ selections.append({
+ "path": str(selection["path"]),
+ "confidence": confidence or 0.7,
+ })
+
+ return {"action": action, "tools": tools, "selections": selections}
+
+ try:
+ data = json.loads(text)
+ if isinstance(data, dict):
+ return extract(data)
+ except (ValueError, json.JSONDecodeError):
+ pass
+
+ fence_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL)
+ if fence_match:
+ try:
+ data = json.loads(fence_match.group(1).strip())
+ if isinstance(data, dict):
+ return extract(data)
+ except (ValueError, json.JSONDecodeError):
+ pass
+
+ brace_match = re.search(r"\{.*\}", text, re.DOTALL)
+ if brace_match:
+ try:
+ data = json.loads(brace_match.group())
+ if isinstance(data, dict):
+ return extract(data)
+ except (ValueError, json.JSONDecodeError):
+ pass
+
+ return default
+
+
+def format_budget_block(snapshot: dict | None) -> str:
+ if not snapshot:
+ return ""
+ planning = snapshot.get("planning") or {}
+ context = snapshot.get("context") or {}
+ return (
+ "=== Resource Status ===\n"
+ f"Planning Budget: {planning.get('status', 'HEALTHY')} "
+ f"({planning.get('used_pct', 0)}% used)\n"
+ f"Context Budget: {context.get('status', 'HEALTHY')} "
+ f"({context.get('used_pct', 0)}% used)\n"
+ f"KG Coverage: {snapshot.get('explored_chunks', 0)}/"
+ f"{snapshot.get('total_chunks', 0)} chunks explored\n"
+ f"Docs Explored: {snapshot.get('explored_docs', 0)}/"
+ f"{snapshot.get('total_docs', 0)}\n"
+ "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. "
+ "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n"
+ "=== End Resource Status ===\n"
+ )
+
+
+def parse_json_array(text: str) -> list[str]:
+ """Best-effort extraction of a JSON array of strings from LLM response text."""
+ result = extract_json_array_payload(text)
+ return [str(item) for item in result]
+
+
+def extract_json_array_payload(text: str) -> list[Any]:
+ text = text.strip()
+ try:
+ result = json.loads(text)
+ if isinstance(result, list):
+ return result
+ except (json.JSONDecodeError, ValueError):
+ pass
+ match = re.search(r"\[.*?\]", text, re.DOTALL)
+ if match:
+ try:
+ result = json.loads(match.group())
+ if isinstance(result, list):
+ return result
+ except (json.JSONDecodeError, ValueError):
+ pass
+ return []
+
+
+def normalize_confidence(value: Any) -> float | None:
+ if value is None:
+ return None
+ if isinstance(value, str):
+ value = value.strip().rstrip("%")
+ try:
+ parsed = float(value)
+ except (TypeError, ValueError):
+ return None
+ if parsed > 1.0:
+ parsed = parsed / 100.0
+ return max(0.0, min(parsed, 1.0))
diff --git a/packages/shared-python/shared/services/retrieval/agentic/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/runtime.py
new file mode 100644
index 000000000..d658cd562
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/runtime.py
@@ -0,0 +1,146 @@
+"""Runtime setup helpers for agentic retrieval."""
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk
+from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetPoolName
+from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState
+from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage
+from shared.utils.token_estimate import estimate_tokens
+
+
+def build_config_from_env() -> AgentRunConfig:
+ return AgentRunConfig(
+ max_revisions=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_REVISIONS", "2")),
+ max_nav_depth=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_DEPTH", "3")),
+ latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "12000")),
+ token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")),
+ planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")),
+ bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")),
+ per_doc_min_share=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE", "1500")),
+ inventory_aware=os.environ.get("RETRIEVAL_AGENTIC_INVENTORY_AWARE", "true") == "true",
+ )
+
+
+async def load_budget_inventory(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ exclude_document_ids: list[str],
+) -> tuple[int, int, dict[str, int]]:
+ stmt = (
+ select(Document.document_id, func.count(DocumentChunk.id))
+ .join(
+ DocumentChunk,
+ (DocumentChunk.document_id == Document.document_id)
+ & (DocumentChunk.job_result_id == Document.current_job_result_id),
+ )
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == "active")
+ .group_by(Document.document_id)
+ )
+ if exclude_document_ids:
+ stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
+
+ result = await db.execute(stmt)
+ doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()}
+ return sum(doc_chunks.values()), len(doc_chunks), doc_chunks
+
+
+class AgentLlmBudget:
+ def __init__(self, state: AgentState) -> None:
+ self._state = state
+
+ async def call(
+ self,
+ llm_fn: LLMFn,
+ prompt: Any,
+ *,
+ pool: BudgetPoolName,
+ doc_id: str | None = None,
+ priority: str = "normal",
+ ) -> str:
+ ledger = self._state.ledger
+ if ledger is None:
+ return await llm_fn(prompt)
+
+ prompt_text = _stringify_llm_input(prompt)
+ est = estimate_tokens(prompt_text)
+ reserved = await ledger.try_reserve(
+ pool,
+ est,
+ doc_id=doc_id,
+ priority="low" if priority == "low" else "normal",
+ )
+ if not reserved:
+ raise BudgetExceeded(f"{pool} budget exhausted")
+
+ try:
+ response = await llm_fn(prompt)
+ except Exception:
+ await ledger.refund(pool, est=est, doc_id=doc_id)
+ raise
+
+ usage = current_llm_usage.get() or {}
+ actual = int(usage.get("prompt_tokens") or est)
+ await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id)
+ return response
+
+ def for_pool(self, llm_fn: LLMFn, *, pool: BudgetPoolName) -> LLMFn:
+ async def _call(prompt: Any) -> str:
+ return await self.call(llm_fn, prompt, pool=pool)
+
+ return _call
+
+ def for_document(
+ self,
+ llm_fn: LLMFn,
+ *,
+ doc_id: str,
+ depth: int,
+ ) -> LLMFn:
+ async def _call(prompt: Any) -> str:
+ return await self.call(
+ llm_fn,
+ prompt,
+ pool="planning",
+ doc_id=doc_id,
+ priority="low" if depth >= 2 else "normal",
+ )
+
+ return _call
+
+ def for_discovery(
+ self,
+ llm_fn: LLMFn,
+ *,
+ doc_id: str,
+ low_priority: bool,
+ ) -> LLMFn:
+ async def _call(prompt: Any) -> str:
+ return await self.call(
+ llm_fn,
+ prompt,
+ pool="planning",
+ doc_id=doc_id,
+ priority="low" if low_priority else "normal",
+ )
+
+ return _call
+
+
+def _stringify_llm_input(prompt: Any) -> str:
+ if isinstance(prompt, str):
+ return prompt
+ try:
+ return json.dumps(prompt, ensure_ascii=False, default=str)
+ except Exception:
+ return str(prompt)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_counts.py b/packages/shared-python/shared/services/retrieval/agentic/section_counts.py
new file mode 100644
index 000000000..74ba465d3
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/section_counts.py
@@ -0,0 +1,174 @@
+"""Section count aggregation for agentic navigation."""
+from __future__ import annotations
+
+from sqlalchemy import case, func, literal_column, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import DocumentChunk
+
+
+async def attach_section_counts(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ all_sections: dict[str, dict],
+ items_by_path: dict[str, dict],
+) -> None:
+ """Attach direct chunk and connected asset counts to visible section items."""
+ scope_item_sids = {
+ item["section_id"]
+ for item in items_by_path.values()
+ if item["show_summary"]
+ }
+ all_section_ids = [meta["section_id"] for meta in all_sections.values()]
+ if not all_section_ids or not scope_item_sids:
+ return
+
+ section_id_counts = await _load_direct_chunk_counts(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ all_section_ids=all_section_ids,
+ )
+
+ sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()}
+ for section_id, (text_count, image_count, table_count) in section_id_counts.items():
+ chunk_path = sid_to_path.get(section_id, "")
+ if not chunk_path:
+ continue
+
+ for item_path, item in items_by_path.items():
+ if not item["show_summary"]:
+ continue
+ if chunk_path == item_path or chunk_path.startswith(item_path + " / "):
+ item["chunk_count"] += text_count
+ item["image_count"] += image_count
+ item["table_count"] += table_count
+
+ await _attach_connected_asset_counts(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ items_by_path=items_by_path,
+ sid_to_path=sid_to_path,
+ )
+
+
+async def _load_direct_chunk_counts(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ all_section_ids: list[str],
+) -> dict[str, tuple[int, int, int]]:
+ chunk_stmt = (
+ select(
+ DocumentChunk.section_id,
+ func.count(
+ case(
+ (DocumentChunk.chunk_type.notin_(["image", "table"]), literal_column("1")),
+ )
+ ).label("text_count"),
+ func.count(
+ case(
+ (DocumentChunk.chunk_type == "image", literal_column("1")),
+ )
+ ).label("image_count"),
+ func.count(
+ case(
+ (DocumentChunk.chunk_type == "table", literal_column("1")),
+ )
+ ).label("table_count"),
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.section_id.in_(all_section_ids))
+ .group_by(DocumentChunk.section_id)
+ )
+ chunk_rows = (await db.execute(chunk_stmt)).all()
+ return {
+ section_id: (int(text_count), int(image_count), int(table_count))
+ for section_id, text_count, image_count, table_count in chunk_rows
+ }
+
+
+async def _attach_connected_asset_counts(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ items_by_path: dict[str, dict],
+ sid_to_path: dict[str, str],
+) -> None:
+ scope_items_with_zero_assets = [
+ item
+ for item in items_by_path.values()
+ if item["show_summary"] and item["image_count"] == 0 and item["table_count"] == 0
+ ]
+ if not scope_items_with_zero_assets:
+ return
+
+ scope_section_ids = {
+ item["section_id"]
+ for item in items_by_path.values()
+ if item.get("section_id")
+ }
+ if not scope_section_ids:
+ return
+
+ connect_stmt = (
+ select(
+ DocumentChunk.section_id,
+ DocumentChunk.chunk_metadata,
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.section_id.in_(list(scope_section_ids)))
+ .where(DocumentChunk.chunk_type == "text")
+ )
+ connect_result = (await db.execute(connect_stmt)).all()
+
+ section_target_ids: dict[str, set[str]] = {}
+ for section_id, metadata in connect_result:
+ if not isinstance(metadata, dict):
+ continue
+ for connection in metadata.get("connect_to") or []:
+ target_id = connection.get("target", "")
+ if target_id:
+ section_target_ids.setdefault(section_id, set()).add(target_id)
+
+ if not section_target_ids:
+ return
+
+ all_target_ids: set[str] = set()
+ for target_ids in section_target_ids.values():
+ all_target_ids.update(target_ids)
+
+ target_type_stmt = (
+ select(
+ DocumentChunk.chunk_id,
+ DocumentChunk.chunk_type,
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.chunk_id.in_(list(all_target_ids)))
+ .where(DocumentChunk.chunk_type.in_(["image", "table"]))
+ )
+ target_type_result = (await db.execute(target_type_stmt)).all()
+ target_types = {chunk_id: chunk_type for chunk_id, chunk_type in target_type_result}
+
+ for section_id, target_ids in section_target_ids.items():
+ ref_path = sid_to_path.get(section_id, "")
+ if not ref_path:
+ continue
+ referenced_images = sum(1 for target_id in target_ids if target_types.get(target_id) == "image")
+ referenced_tables = sum(1 for target_id in target_ids if target_types.get(target_id) == "table")
+ if referenced_images == 0 and referenced_tables == 0:
+ continue
+ for item_path, item in items_by_path.items():
+ if not item["show_summary"]:
+ continue
+ if ref_path == item_path or ref_path.startswith(item_path + " / "):
+ item["image_count"] += referenced_images
+ item["table_count"] += referenced_tables
diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py b/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py
new file mode 100644
index 000000000..b7a251805
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py
@@ -0,0 +1,59 @@
+"""Prompt projection for agentic section navigation."""
+from __future__ import annotations
+
+from shared.utils.text_utils import truncate_content_preview
+
+
+def format_items_for_llm(
+ items: list[dict],
+ max_chars: int = 20000,
+) -> tuple[str, bool]:
+ """Format section items with hierarchy, selectability, counts, and summaries."""
+ if not items:
+ return "(no items available)", False
+
+ full_text = "\n".join(_render_item(item, include_summary=True) for item in items)
+ if len(full_text) <= max_chars:
+ return full_text, False
+
+ slim_text = "\n".join(_render_item(item, include_summary=False) for item in items)
+ return slim_text[:max_chars], True
+
+
+def _render_item(item: dict, include_summary: bool) -> str:
+ level = item.get("level", 1)
+ show_summary = item.get("show_summary", True)
+ is_leaf = item.get("is_leaf", False)
+ leaf_tag = " [Leaf]" if is_leaf else ""
+ path = item.get("path", "")
+ summary = item.get("summary") or ""
+
+ counts_str = ""
+ if show_summary:
+ count_parts: list[str] = []
+ chunk_count = item.get("chunk_count", 0)
+ if chunk_count > 0:
+ count_parts.append(f"text={chunk_count}")
+ image_count = item.get("image_count", 0)
+ if image_count > 0:
+ count_parts.append(f"image={image_count}")
+ table_count = item.get("table_count", 0)
+ if table_count > 0:
+ count_parts.append(f"table={table_count}")
+ counts_str = f' [{" ".join(count_parts)}]' if count_parts else ""
+
+ indent = " " * (level - 1)
+ prefix = "▸" if level == 1 else "└"
+ level_tag = f"[L{level}]"
+ select_tag = "[SELECT] " if item.get("selectable", False) else ""
+
+ lines = [
+ f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}'
+ ]
+
+ if include_summary and show_summary and summary:
+ sub_indent = " " * level
+ clipped = truncate_content_preview(summary, head=80, tail=0)
+ lines.append(f"{sub_indent}{clipped}")
+
+ return "\n".join(lines)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py
new file mode 100644
index 000000000..53507f35d
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py
@@ -0,0 +1,231 @@
+"""Section-tree loading and prompt projection for agentic navigation."""
+from __future__ import annotations
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import DocumentSection
+from shared.services.retrieval.agentic.section_counts import attach_section_counts
+from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path
+
+
+async def load_child_sections(
+ db: AsyncSession,
+ document_id: str,
+ job_result_id: str,
+ scope_path: str | list[str] | None = None,
+ exclude_paths: set[str] | None = None,
+) -> list[dict]:
+ """Load the continuous context tree for a navigation scope."""
+ stmt = (
+ select(
+ DocumentSection.section_id,
+ DocumentSection.section_title,
+ DocumentSection.section_path,
+ DocumentSection.summary,
+ DocumentSection.sort_order,
+ )
+ .where(DocumentSection.document_id == document_id)
+ .where(DocumentSection.job_result_id == job_result_id)
+ .order_by(DocumentSection.sort_order)
+ )
+ section_rows = (await db.execute(stmt)).all()
+ if not section_rows:
+ return []
+
+ if isinstance(scope_path, list):
+ scope_list = [normalize_section_path(path) for path in scope_path]
+ elif scope_path:
+ scope_list = [normalize_section_path(scope_path)]
+ else:
+ scope_list = []
+
+ scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0
+ excluded_paths = exclude_paths or set()
+
+ logger.debug(
+ f" load_child_sections: scopes={scope_list or ['root']} "
+ f"scope_depth={scope_depth} exclude_paths={excluded_paths if excluded_paths else 'none'} "
+ f"total_sections={len(section_rows)}"
+ )
+
+ all_sections: dict[str, dict] = {}
+ for section_id, title, path, summary, sort_order in section_rows:
+ if not path:
+ continue
+ normalized_path = normalize_section_path(path)
+ parts = split_section_path(normalized_path)
+ all_sections[normalized_path] = {
+ "title": title or parts[-1] if parts else normalized_path,
+ "summary": summary or "",
+ "sort_order": int(sort_order or 0),
+ "section_id": section_id,
+ "parts": parts,
+ "depth": len(parts),
+ }
+
+ ancestor_prefixes: set[str] = set()
+ for scope in scope_list:
+ scope_parts = split_section_path(scope)
+ for index in range(1, len(scope_parts) + 1):
+ ancestor_prefixes.add(" / ".join(scope_parts[:index]))
+
+ items_by_path = _select_scope_items(
+ all_sections,
+ scope_list=scope_list,
+ ancestor_prefixes=ancestor_prefixes,
+ exclude_paths=excluded_paths,
+ )
+ if not items_by_path:
+ return []
+
+ allowed_set = _resolve_allowed_depths(items_by_path, scope_list)
+ if allowed_set:
+ to_remove = [
+ path
+ for path, item in items_by_path.items()
+ if item["show_summary"] and item["level"] not in allowed_set
+ ]
+ for path in to_remove:
+ del items_by_path[path]
+
+ if not items_by_path:
+ return []
+
+ await attach_section_counts(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ all_sections=all_sections,
+ items_by_path=items_by_path,
+ )
+
+ sorted_items = sorted(items_by_path.values(), key=lambda item: item["sort_order"])
+ for item in sorted_items:
+ item.pop("sort_order", None)
+ item.pop("section_id", None)
+
+ _mark_leaf_and_selectable(sorted_items, all_section_paths=set(all_sections.keys()), allowed_set=allowed_set)
+ return sorted_items
+
+
+def _select_scope_items(
+ all_sections: dict[str, dict],
+ *,
+ scope_list: list[str],
+ ancestor_prefixes: set[str],
+ exclude_paths: set[str],
+) -> dict[str, dict]:
+ items_by_path: dict[str, dict] = {}
+
+ def is_excluded(path: str) -> bool:
+ return bool(
+ exclude_paths
+ and any(path == excluded or path.startswith(excluded + " / ") for excluded in exclude_paths)
+ )
+
+ for path, meta in all_sections.items():
+ parts = meta["parts"]
+ depth = meta["depth"]
+
+ if not scope_list:
+ if depth < 1 or is_excluded(path):
+ continue
+ items_by_path[path] = _make_item(path, meta, show_summary=True)
+ continue
+
+ matched_scope = _find_matched_scope(parts, depth=depth, scope_list=scope_list)
+ if matched_scope:
+ if is_excluded(path):
+ continue
+ items_by_path[path] = _make_item(path, meta, show_summary=True)
+ continue
+
+ max_scope_depth = max(len(split_section_path(scope)) for scope in scope_list)
+ if depth <= max_scope_depth:
+ if depth == 1 and path in ancestor_prefixes:
+ items_by_path.setdefault(path, _make_item(path, meta, show_summary=False))
+ elif depth > 1:
+ parent_prefix = " / ".join(parts[:-1])
+ if parent_prefix in ancestor_prefixes:
+ items_by_path.setdefault(path, _make_item(path, meta, show_summary=False))
+
+ return items_by_path
+
+
+def _make_item(path: str, meta: dict, show_summary: bool) -> dict:
+ return {
+ "path": path,
+ "title": meta["title"],
+ "summary": meta["summary"],
+ "level": meta["depth"],
+ "sort_order": meta["sort_order"],
+ "chunk_count": 0,
+ "image_count": 0,
+ "table_count": 0,
+ "section_id": meta["section_id"],
+ "show_summary": show_summary,
+ }
+
+
+def _find_matched_scope(parts: list[str], *, depth: int, scope_list: list[str]) -> str | None:
+ for scope in scope_list:
+ scope_parts = split_section_path(scope)
+ scope_depth = len(scope_parts)
+ if depth > scope_depth and parts[:scope_depth] == scope_parts:
+ return scope
+ return None
+
+
+def _resolve_allowed_depths(items_by_path: dict[str, dict], scope_list: list[str]) -> set[int]:
+ if not scope_list:
+ depths = {
+ item["level"]
+ for item in items_by_path.values()
+ if item.get("show_summary", True)
+ }
+ return set(sorted(depths)[:2])
+
+ allowed_set: set[int] = set()
+ for scope in scope_list:
+ scope_parts = split_section_path(scope)
+ scope_depth = len(scope_parts)
+ child_depths = {
+ item["level"]
+ for item in items_by_path.values()
+ if item.get("show_summary", True)
+ and item["level"] > scope_depth
+ and split_section_path(item["path"])[:scope_depth] == scope_parts
+ }
+ if child_depths:
+ allowed_set.update(sorted(child_depths)[:2])
+ return allowed_set
+
+
+def _mark_leaf_and_selectable(
+ sorted_items: list[dict],
+ *,
+ all_section_paths: set[str],
+ allowed_set: set[int],
+) -> None:
+ for item in sorted_items:
+ item_path = item["path"]
+ has_descendants = any(
+ path != item_path and path.startswith(item_path + " / ")
+ for path in all_section_paths
+ )
+ item["is_leaf"] = not has_descendants
+
+ if allowed_set:
+ shallowest_band = min(allowed_set)
+ for item in sorted_items:
+ if not item.get("show_summary", True):
+ item["selectable"] = False
+ elif item["level"] == shallowest_band and not item.get("is_leaf", False):
+ item["selectable"] = False
+ else:
+ item["selectable"] = True
+ else:
+ for item in sorted_items:
+ item["selectable"] = item.get("show_summary", True)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py
new file mode 100644
index 000000000..2fa86c690
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py
@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.agentic import asset_tools
+from shared.services.retrieval.agentic.types import DocTreeNode
+from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows
+from shared.services.retrieval.path_hydration import hydrate_paths_to_rows
+
+
+async def hydrate_path_selections_into_node(
+ db: AsyncSession,
+ *,
+ node: DocTreeNode,
+ path_selections: list[dict[str, Any]],
+ user_id: str,
+ namespace: str,
+ document_id: str,
+ job_result_id: str | None = None,
+) -> None:
+ chunks = await hydrate_paths_to_rows(
+ db,
+ path_selections=path_selections,
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
+ if not chunks:
+ return
+
+ chunks = await _append_connected_asset_targets(db, chunks)
+ resolved_job_result_id = job_result_id or _find_job_result_id(chunks)
+ if resolved_job_result_id:
+ await _attach_root_asset_owners(
+ db,
+ document_id=document_id,
+ job_result_id=resolved_job_result_id,
+ chunks=chunks,
+ )
+
+ add_chunks_to_node(node, chunks)
+
+
+async def _append_connected_asset_targets(
+ db: AsyncSession, chunks: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+ connected = await hydrate_connected_target_rows(
+ db=db,
+ rows=chunks,
+ exclude_document_ids=[],
+ exclude_sections=[],
+ )
+ if not connected:
+ return chunks
+
+ owner_map = asset_tools.build_connected_owner_map(chunks)
+ for chunk in connected:
+ if not chunk.get("owner_section_path"):
+ chunk["owner_section_path"] = owner_map.get(str(chunk.get("chunk_id") or ""))
+ return [*chunks, *connected]
+
+
+async def _attach_root_asset_owners(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ job_result_id: str,
+ chunks: list[dict[str, Any]],
+) -> None:
+ root_map = await asset_tools.resolve_root_asset_owners(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ chunks=chunks,
+ )
+ if not root_map:
+ return
+
+ for chunk in chunks:
+ if chunk.get("owner_section_path"):
+ continue
+ chunk_id = str(chunk.get("chunk_id") or "")
+ if chunk_id in root_map:
+ chunk["owner_section_path"] = root_map[chunk_id]
+
+
+def _find_job_result_id(chunks: list[dict[str, Any]]) -> str | None:
+ return next(
+ (str(chunk["job_result_id"]) for chunk in chunks if chunk.get("job_result_id")),
+ None,
+ )
+
+
+def add_chunks_to_node(node: DocTreeNode, chunks: list[dict[str, Any]]) -> None:
+ for chunk in chunks:
+ real_path = (
+ chunk.get("owner_section_path")
+ or chunk.get("section_path")
+ or chunk.get("source_chunk_path")
+ )
+ if real_path:
+ node.add_leaf_chunks(str(real_path), [chunk])
diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py
index 10ab54757..4fde17d1d 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py
@@ -1,151 +1,24 @@
-"""Agentic retrieval tools — thin wrappers around existing retrieval components.
+"""Agentic retrieval tool adapters.
-Each tool:
- 1. Calls existing functions from channels.py, agent_navigate.py, app_service.py
- 2. Returns a unified ToolResult
- 3. Never raises — errors are captured in ToolResult.error
+Concrete tool implementations live in focused Modules. This file is the stable
+adapter seam used by the workflow orchestrator and contract tests.
"""
from __future__ import annotations
-import time
from typing import Any
-from loguru import logger
-from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.models.database.document import Document
-from shared.services.retrieval.agentic.budget import BudgetExceeded
-from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult
-from shared.services.retrieval.agent_navigate import (
- _build_knowledge_map_overview,
- _format_items_for_llm,
- _load_child_sections,
- _parse_json_array,
- _parse_action_response,
- _ACTION_PROMPT,
- _DISCOVERY_SELECT_PROMPT,
- _FILE_SELECT_PROMPT,
- _format_budget_block,
-)
-from shared.services.retrieval.app_service import (
- _CHANNEL_WEIGHT_CONTENT,
- _CHANNEL_WEIGHT_PATH,
- _CHANNEL_WEIGHT_TERM,
- _INTERNAL_RECALL_K_MULTIPLIER,
- _merge_same_section_rows,
- _normalize_row_scores,
- _resolve_allowed_chunk_types,
- hydrate_connected_target_rows,
- merge_channels_rrf,
+from shared.services.retrieval.agentic import (
+ asset_tools,
+ discovery_selection,
+ discovery_tools,
+ navigation_tools,
)
-from shared.services.retrieval.channels import content_channel, path_channel, term_channel
-from shared.services.retrieval.lexical_text import normalize_section_path
+from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult
from shared.services.retrieval.llm_adapter import LLMFn
-# ---------------------------------------------------------------------------
-# Helper: resolve connected asset → owner text chunk section_path
-# ---------------------------------------------------------------------------
-
-def _build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]:
- """Build target_chunk_id → owner text chunk section_path mapping.
-
- When text chunks reference images/tables via connect_to metadata,
- the referenced assets live in Root section. This map lets us attribute
- those assets back to the text chunk's section for correct tree placement.
- """
- owner_map: dict[str, str] = {}
- for chunk in text_chunks:
- if (chunk.get('chunk_type') or 'text') != 'text':
- continue
- section_path = chunk.get('section_path') or ''
- if not section_path:
- continue
- metadata = chunk.get('chunk_metadata') or {}
- if not isinstance(metadata, dict):
- continue
- for conn in metadata.get('connect_to') or []:
- if not isinstance(conn, dict):
- continue
- target_id = str(conn.get('target') or '').strip()
- if target_id and target_id not in owner_map:
- owner_map[target_id] = section_path
- return owner_map
-
-
-async def _resolve_root_asset_owners(
- db: AsyncSession,
- *,
- document_id: str,
- job_result_id: str,
- chunks: list[dict[str, Any]],
-) -> dict[str, str]:
- """Resolve owner section_path for Root-stranded image/table chunks.
-
- When Root is hydrated directly (e.g. via discovery selection), the
- batch contains standalone image/table chunks whose section_path is
- 'Root'. ``_build_connected_owner_map`` cannot help because the
- referencing text chunks live in other sections outside the batch.
-
- This function queries the *entire document* for text chunks with
- connect_to metadata, using the same logic as
- ``_build_connected_owner_map``, to resolve the true owner.
-
- Returns target_chunk_id → owner_section_path for Root assets only.
- Returns empty dict when there are no Root assets (zero DB overhead).
- """
- from shared.models.database.document import DocumentChunk, DocumentSection
-
- root_asset_ids = [
- str(c.get('chunk_id') or '')
- for c in chunks
- if not c.get('owner_section_path') # skip if already resolved by batch-level owner map
- and (c.get('section_path') or '') == 'Root'
- and (c.get('chunk_type') or '').lower() in ('image', 'table')
- and c.get('chunk_id')
- ]
- if not root_asset_ids:
- return {}
-
- root_asset_set = set(root_asset_ids)
-
- # Query all text chunks in this document for connect_to metadata
- text_stmt = (
- select(
- DocumentChunk.chunk_metadata,
- DocumentSection.section_path,
- )
- .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.chunk_type == 'text')
- )
- result = await db.execute(text_stmt)
-
- owner_map: dict[str, str] = {}
- for metadata, section_path in result.all():
- if not isinstance(metadata, dict) or not section_path:
- continue
- for conn in metadata.get('connect_to') or []:
- if not isinstance(conn, dict):
- continue
- target_id = str(conn.get('target') or '').strip()
- if target_id in root_asset_set and target_id not in owner_map:
- owner_map[target_id] = section_path
-
- if owner_map:
- logger.info(
- f' _resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} '
- f'Root assets to their owner sections'
- )
- return owner_map
-
-
-# ---------------------------------------------------------------------------
-# Tool: bottom_discovery
-# ---------------------------------------------------------------------------
-
async def bottom_discovery(
db: AsyncSession,
*,
@@ -157,108 +30,29 @@ async def bottom_discovery(
exclude_sections: list[dict[str, str]],
data_type: int = 1,
signal_paths: list[str] | None = None,
- filter_mode: str = 'delete',
+ filter_mode: str = "delete",
channels: list[str] | None = None,
channel_weights: dict[str, float] | None = None,
internal_recall_k: int | None = None,
- **_kwargs: Any,
+ **kwargs: Any,
) -> ToolResult:
- """Run 3-channel BM25 discovery + RRF fusion."""
- t0 = time.monotonic()
- try:
- allowed_chunk_types = _resolve_allowed_chunk_types(data_type)
- effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER
- active_channels = set(channels) if channels else {'path', 'content', 'term'}
-
- path_rows: list[dict[str, Any]] = []
- content_rows: list[dict[str, Any]] = []
- term_rows: list[dict[str, Any]] = []
-
- if 'path' in active_channels:
- path_rows = await path_channel(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=effective_recall_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths, filter_mode=filter_mode,
- )
-
- if 'content' in active_channels:
- content_rows = await content_channel(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=effective_recall_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths, filter_mode=filter_mode,
- )
-
- if 'term' in active_channels:
- term_rows = await term_channel(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=effective_recall_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths, filter_mode=filter_mode,
- )
-
- # RRF fusion
- default_weights = {
- 'path': _CHANNEL_WEIGHT_PATH,
- 'content': _CHANNEL_WEIGHT_CONTENT,
- 'term': _CHANNEL_WEIGHT_TERM,
- }
- effective_weights = {**default_weights, **(channel_weights or {})}
-
- channel_lists: list[list[dict[str, Any]]] = []
- weight_list: list[float] = []
- if path_rows:
- channel_lists.append(path_rows)
- weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH))
- if content_rows:
- channel_lists.append(content_rows)
- weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT))
- if term_rows:
- channel_lists.append(term_rows)
- weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM))
-
- fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else []
- fused_rows = _merge_same_section_rows(fused_rows)
-
- if fused_rows:
- _normalize_row_scores(fused_rows, source_field='score', target_field='discovery_score', default=0.5)
-
- # Extract top document IDs as hints for KG selection
- doc_id_counts: dict[str, int] = {}
- for row in fused_rows:
- did = row.get('document_id', '')
- if did:
- doc_id_counts[did] = doc_id_counts.get(did, 0) + 1
- top_doc_ids = sorted(doc_id_counts, key=lambda d: doc_id_counts[d], reverse=True)[:5]
-
- latency = int((time.monotonic() - t0) * 1000)
- logger.info(
- f' agentic.bottom_discovery: {len(fused_rows)} fused rows, '
- f'top_doc_ids={top_doc_ids}, {latency}ms'
- )
- return ToolResult(
- status='discovery_done',
- payload={
- 'fused_rows': fused_rows,
- 'top_doc_ids': top_doc_ids,
- 'channel_counts': {
- 'path': len(path_rows),
- 'content': len(content_rows),
- 'term': len(term_rows),
- },
- },
- latency_ms=latency,
- )
- except Exception as e:
- latency = int((time.monotonic() - t0) * 1000)
- logger.error(f' agentic.bottom_discovery failed: {e}')
- return ToolResult(status='error', error=str(e), latency_ms=latency)
-
+ return await discovery_tools.bottom_discovery(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ internal_recall_k=internal_recall_k,
+ **kwargs,
+ )
-# ---------------------------------------------------------------------------
-# Tool: kg_document_select
-# ---------------------------------------------------------------------------
async def kg_document_select(
db: AsyncSession,
@@ -269,104 +63,19 @@ async def kg_document_select(
llm_fn: LLMFn | None,
exclude_document_ids: list[str],
revision_hint: str | None = None,
- **_kwargs: Any,
+ **kwargs: Any,
) -> ToolResult:
- """Select candidate documents from document-level KG."""
- t0 = time.monotonic()
- try:
- overview_text, doc_id_to_name = await _build_knowledge_map_overview(
- db, user_id=user_id, namespace=namespace,
- )
- if overview_text == '(empty)':
- latency = int((time.monotonic() - t0) * 1000)
- return ToolResult(
- status='no_confident_doc',
- payload={'reason': 'no active documents in namespace'},
- latency_ms=latency,
- )
-
- if llm_fn is None:
- latency = int((time.monotonic() - t0) * 1000)
- return ToolResult(
- status='no_confident_doc',
- payload={'reason': 'LLM not available'},
- latency_ms=latency,
- )
-
- revision_context = ''
- if revision_hint:
- revision_context = (
- f'\nIMPORTANT: This is a REVISION round. '
- f'The previous search attempt failed because:\n'
- f'"{revision_hint}"\n'
- f'Adjust your document selection accordingly. '
- f'If no document can address this, return an EMPTY array [].\n'
- )
-
- file_prompt = _FILE_SELECT_PROMPT.format(
- overview=overview_text, query=query,
- revision_context=revision_context,
- budget_block=_format_budget_block(_kwargs.get('budget_snapshot')),
- )
- file_response = await llm_fn(file_prompt)
- selected_ids = _parse_json_array(file_response)
-
- exclude_set = set(exclude_document_ids)
- valid_ids = [did for did in selected_ids if did in doc_id_to_name and did not in exclude_set]
-
- if not valid_ids:
- latency = int((time.monotonic() - t0) * 1000)
- logger.info(f' agentic.kg_document_select: LLM returned no valid docs, {latency}ms')
- return ToolResult(
- status='no_confident_doc',
- payload={'reason': 'LLM returned no valid document IDs', 'raw_ids': selected_ids},
- latency_ms=latency,
- )
-
- # Load job_result_ids for selected documents
- doc_job_map: dict[str, str] = {}
- doc_stmt = (
- select(Document.document_id, Document.current_job_result_id)
- .where(Document.document_id.in_(valid_ids))
- )
- doc_result = await db.execute(doc_stmt)
- for did, jrid in doc_result.all():
- if jrid:
- doc_job_map[did] = jrid
-
- candidate_docs = []
- for did in valid_ids:
- candidate_docs.append({
- 'document_id': did,
- 'source_file_name': doc_id_to_name.get(did, ''),
- 'confidence': 1.0,
- 'reason': 'LLM selected from KG overview',
- 'source': 'kg_llm_select',
- })
-
- latency = int((time.monotonic() - t0) * 1000)
- logger.info(f' agentic.kg_document_select: {len(candidate_docs)} docs selected, {latency}ms')
- return ToolResult(
- status='selected_docs',
- payload={
- 'candidate_docs': candidate_docs,
- 'doc_id_to_name': doc_id_to_name,
- 'doc_job_map': doc_job_map,
- },
- latency_ms=latency,
- )
- except BudgetExceeded:
- raise
- except Exception as e:
- latency = int((time.monotonic() - t0) * 1000)
- logger.error(f' agentic.kg_document_select failed: {e}')
- return ToolResult(status='error', error=str(e), latency_ms=latency)
-
-
+ return await discovery_tools.kg_document_select(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ llm_fn=llm_fn,
+ exclude_document_ids=exclude_document_ids,
+ revision_hint=revision_hint,
+ **kwargs,
+ )
-# ---------------------------------------------------------------------------
-# Tool: asset_filter_step (programmatic asset extraction)
-# ---------------------------------------------------------------------------
async def asset_filter_step(
db: AsyncSession,
@@ -374,218 +83,16 @@ async def asset_filter_step(
document_id: str,
job_result_id: str,
scope_path: str | list[str] | None,
- asset_type: str, # 'image' | 'table'
+ asset_type: str,
) -> list[dict[str, Any]]:
- """Extract assets from all descendants under scope_path.
-
- Terminal action — no LLM involved.
- Algorithm: load all text chunks under scope → parse connect_to metadata →
- batch-load target image/table chunks → return directly.
-
- Also collects standalone asset chunks (image/table) that exist directly
- under the scope but are not referenced via connect_to.
-
- scope_path can be:
- - None: root scope (entire document)
- - str: single scope path
- - list[str]: multiple scope paths (queried simultaneously)
- """
- from shared.models.database.document import DocumentChunk, DocumentSection
-
- t0 = time.monotonic()
- try:
- # 1. Find all section_ids under scope_path(s)
- # Normalize scope to list for uniform handling
- scope_list = (
- scope_path if isinstance(scope_path, list)
- else [scope_path] if scope_path
- else []
- )
-
- section_stmt = (
- select(DocumentSection.section_id, DocumentSection.section_path)
- .where(DocumentSection.document_id == document_id)
- .where(DocumentSection.job_result_id == job_result_id)
- )
- if scope_list:
- from sqlalchemy import or_
- scope_filters = []
- for sp in scope_list:
- scope_filters.append(DocumentSection.section_path == sp)
- scope_filters.append(DocumentSection.section_path.like(f'{sp} / %'))
- section_stmt = section_stmt.where(or_(*scope_filters))
- section_result = await db.execute(section_stmt)
- section_rows = section_result.all()
- section_ids = {row[0] for row in section_rows}
-
- if not section_ids:
- logger.info(f' asset_filter_step: no sections found under scope={scope_path}')
- return []
-
- # 2. Load target asset chunks directly (standalone assets in the scope)
- asset_stmt = (
- select(
- DocumentChunk.chunk_id,
- DocumentChunk.chunk_type,
- DocumentChunk.content,
- DocumentChunk.file_path,
- DocumentChunk.section_id,
- DocumentChunk.source_chunk_path,
- DocumentChunk.chunk_metadata,
- DocumentChunk.sort_order,
- DocumentChunk.job_result_id,
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.section_id.in_(list(section_ids)))
- .where(DocumentChunk.chunk_type == asset_type)
- .order_by(DocumentChunk.sort_order)
- )
- asset_result = await db.execute(asset_stmt)
- asset_rows = asset_result.all()
-
- section_path_by_id = {section_id: section_path for section_id, section_path in section_rows}
-
- # 3. Resolve media → owner text section via connect_to tracing
- text_stmt = (
- select(
- DocumentChunk.section_id,
- DocumentChunk.chunk_type,
- DocumentChunk.chunk_metadata,
- DocumentChunk.source_chunk_path,
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.section_id.in_(list(section_ids)))
- .where(DocumentChunk.chunk_type == 'text')
- )
- text_result = await db.execute(text_stmt)
- text_row_dicts = [
- {
- 'chunk_type': chunk_type,
- 'chunk_metadata': metadata or {},
- 'section_id': sid,
- 'section_path': section_path_by_id.get(sid, ''),
- 'source_chunk_path': scp,
- }
- for sid, chunk_type, metadata, scp in text_result.all()
- ]
- owner_by_target_id = _build_connected_owner_map(text_row_dicts)
-
- # Replace synthetic "Root" owner with the document's source_file_name.
- # Root is a hybrid node whose real path is the file name (e.g.
- # "32_安全大模型技术与市场研究报告_1.docx"); the DB stores the
- # synthetic label "Root" which cannot match any outline node.
- if any(v == 'Root' for v in owner_by_target_id.values()):
- doc_stmt = select(Document.source_file_name).where(
- Document.document_id == document_id
- )
- doc_file_name = (await db.execute(doc_stmt)).scalar() or ''
- if doc_file_name:
- for tid in list(owner_by_target_id):
- if owner_by_target_id[tid] == 'Root':
- owner_by_target_id[tid] = doc_file_name
-
- # Collect connected target IDs for batch-loading
- connected_target_ids: set[str] = set(owner_by_target_id.keys())
-
- # Load connected targets that match asset_type
- if connected_target_ids:
- connected_stmt = (
- select(
- DocumentChunk.chunk_id,
- DocumentChunk.chunk_type,
- DocumentChunk.content,
- DocumentChunk.file_path,
- DocumentChunk.section_id,
- DocumentChunk.source_chunk_path,
- DocumentChunk.chunk_metadata,
- DocumentChunk.sort_order,
- DocumentChunk.job_result_id,
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.chunk_id.in_(list(connected_target_ids)))
- .where(DocumentChunk.chunk_type == asset_type)
- .order_by(DocumentChunk.sort_order)
- )
- connected_result = await db.execute(connected_stmt)
- connected_rows = connected_result.all()
- else:
- connected_rows = []
-
- # 4. Merge and deduplicate
- seen_ids: set[str] = set()
- chunks: list[dict[str, Any]] = []
-
- # Helper to look up job_id from job_result
- from shared.models.database.job_result import JobResult
- job_stmt = (
- select(JobResult.job_id)
- .where(JobResult.id == job_result_id)
- )
- job_result_row = await db.execute(job_stmt)
- job_id = job_result_row.scalar() or ''
-
- for row in list(asset_rows) + list(connected_rows):
- chunk_id = row[0]
- if chunk_id in seen_ids:
- continue
- seen_ids.add(chunk_id)
-
- # Owner resolution: prefer connect_to-based owner
- owner_section_path = owner_by_target_id.get(chunk_id)
-
- # Fallback: media's own section_id path, but guard against
- # Root / top-level aggregation sections
- if not owner_section_path:
- own_section_path = section_path_by_id.get(row[4])
- if own_section_path and own_section_path == 'Root':
- # Reject only the synthetic Root aggregation label;
- # legitimate L1 sections (e.g. "前言") are valid owners.
- logger.warning(
- f' asset_filter_step: rejecting root-level owner fallback '
- f'chunk_id={chunk_id} section_path={own_section_path}'
- )
- own_section_path = None
- owner_section_path = own_section_path
-
- if not owner_section_path:
- logger.warning(
- f' asset_filter_step unresolved owner: chunk_id={chunk_id} '
- f'file_path={row[3]} scope={scope_path or "root"}'
- )
- continue
- chunks.append({
- 'document_id': document_id,
- 'chunk_id': chunk_id,
- 'chunk_type': row[1],
- 'content': row[2],
- 'file_path': row[3],
- 'section_id': row[4],
- 'section_path': owner_section_path,
- 'owner_section_path': owner_section_path,
- 'source_chunk_path': row[5],
- 'chunk_metadata': row[6] or {},
- 'sort_order': row[7],
- 'job_result_id': job_result_id,
- 'job_id': job_id,
- })
-
- latency = int((time.monotonic() - t0) * 1000)
- logger.info(
- f' asset_filter_step scope={scope_path or "root"} '
- f'type={asset_type}: {len(chunks)} chunks found, {latency}ms'
- )
- return chunks
-
- except Exception as e:
- logger.error(f' asset_filter_step failed: {e}')
- return []
+ return await asset_tools.asset_filter_step(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ scope_path=scope_path,
+ asset_type=asset_type,
+ )
-# ---------------------------------------------------------------------------
-# Tool: navigate_step (unified action — merges tool_select + scope_navigate)
-# ---------------------------------------------------------------------------
async def navigate_step(
db: AsyncSession,
@@ -596,226 +103,26 @@ async def navigate_step(
llm_fn: LLMFn,
user_id: str,
namespace: str,
- doc_name: str = '',
+ doc_name: str = "",
scope_path: str | list[str] | None = None,
exclude_paths: set[str] | None = None,
revision_hint: str | None = None,
budget_snapshot: dict | None = None,
) -> tuple[str, list[str], DocTreeNode, list[dict]]:
- """Unified navigation step — one LLM call for action + tools + selections.
-
- scope_path can be:
- - None: root scope
- - str: single scope to drill into
- - list[str]: multiple scopes to expand simultaneously
-
- Returns:
- - action: 'STOP' | 'NAVIGATE'
- - asset_tools: list of asset tools to run (FIND_IMAGES, FIND_TABLES)
- - node: DocTreeNode with outline_items and leaf_content
- - pending: list of {path, confidence} for non-leaf drill-downs (empty when STOP)
- """
- from shared.services.retrieval.app_service import _hydrate_paths_to_rows
-
- # Normalize scope for internal use
- scope_paths: list[str] = (
- scope_path if isinstance(scope_path, list)
- else [scope_path] if scope_path
- else []
+ return await navigation_tools.navigate_step(
+ db,
+ document_id=document_id,
+ job_result_id=job_result_id,
+ query=query,
+ llm_fn=llm_fn,
+ user_id=user_id,
+ namespace=namespace,
+ doc_name=doc_name,
+ scope_path=scope_path,
+ exclude_paths=exclude_paths,
+ revision_hint=revision_hint,
+ budget_snapshot=budget_snapshot,
)
- # Set of scope path strings (for filtering selections)
- scope_path_set = set(scope_paths)
-
- empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None)
-
- try:
- # 1. Load continuous context tree (supports multi-scope)
- items = await _load_child_sections(
- db, document_id, job_result_id, scope_path,
- exclude_paths=exclude_paths,
- )
- if not items:
- return 'STOP', [], empty, []
-
- # 2. Build selectable index
- selectable = {item['path']: item for item in items if item.get('selectable', False)}
-
- # 3. Count ALL image/table chunks under the scope subtree(s)
- from shared.models.database.document import DocumentChunk, DocumentSection
- from sqlalchemy import func as sa_func
-
- scope_section_stmt = (
- select(DocumentSection.section_id)
- .where(DocumentSection.document_id == document_id)
- .where(DocumentSection.job_result_id == job_result_id)
- )
- if scope_paths:
- from sqlalchemy import or_
- scope_filters = []
- for sp in scope_paths:
- scope_filters.append(DocumentSection.section_path == sp)
- scope_filters.append(DocumentSection.section_path.like(f'{sp} / %'))
- scope_section_stmt = scope_section_stmt.where(or_(*scope_filters))
- scope_section_ids = await db.execute(scope_section_stmt)
- all_section_ids = [r[0] for r in scope_section_ids.all()]
-
- total_images = 0
- total_tables = 0
- if all_section_ids:
- count_stmt = (
- select(
- DocumentChunk.chunk_type,
- sa_func.count(DocumentChunk.id),
- )
- .where(DocumentChunk.document_id == document_id)
- .where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.section_id.in_(all_section_ids))
- .where(DocumentChunk.chunk_type.in_(['image', 'table']))
- .group_by(DocumentChunk.chunk_type)
- )
- count_result = await db.execute(count_stmt)
- for chunk_type, cnt in count_result.all():
- if chunk_type == 'image':
- total_images = cnt
- elif chunk_type == 'table':
- total_tables = cnt
-
- tools_block = ''
- if total_images > 0 or total_tables > 0:
- tools_lines = ['\nOptional asset tools (usable with NAVIGATE or STOP):\n']
- if total_images > 0:
- tools_lines.append(
- f' FIND_IMAGES — Extract image/chart assets under the current scope ({total_images} available).\n'
- )
- if total_tables > 0:
- tools_lines.append(
- f' FIND_TABLES — Extract table/data assets under the current scope ({total_tables} available).\n'
- )
- tools_lines.append(
- ' Note: with NAVIGATE selections, asset tools are limited to the selected sections; '
- 'with STOP or no selections, they use the current scope.\n'
- )
- tools_block = ''.join(tools_lines)
-
- # 4. Format tree and build prompt
- text, overflowed = _format_items_for_llm(items)
- if not scope_paths:
- scope_header = 'Current scope: root (document top level)'
- elif len(scope_paths) == 1:
- scope_header = f'Current scope: navigating into "{scope_paths[0]}"'
- else:
- scope_header = f'Current scope: navigating into {len(scope_paths)} sections'
- prompt = _ACTION_PROMPT.format(
- doc_name=doc_name or document_id,
- doc_id=document_id,
- scope_header=scope_header,
- budget_block=_format_budget_block(budget_snapshot),
- items_overview=text,
- query=query,
- tools_block=tools_block,
- )
- if revision_hint:
- prompt += (
- f'\n\nIMPORTANT: Previous round feedback: '
- f'"{revision_hint}". Adjust your selections accordingly.'
- )
-
- # 5. Single LLM call
- response = await llm_fn(prompt)
- parsed = _parse_action_response(response)
- action = parsed['action']
- asset_tools = parsed['tools']
- selections = parsed['selections']
-
- scope_label = ', '.join(scope_paths) if scope_paths else 'root'
- logger.info(
- f' navigate_step scope={scope_label}: '
- f'action={action} tools={asset_tools} '
- f'selections={len(selections)} selectable={len(selectable)} '
- f'overflowed={overflowed}'
- )
-
- # 6. Build node with LOCAL items only
- node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None)
- local_items = [item for item in items if item.get('show_summary', True)]
- node.outline_items = local_items
-
- # 7. Dispatch selections (only present when action == NAVIGATE)
- valid_selections = [
- s for s in selections
- if s['path'] in selectable and s['path'] not in scope_path_set
- ]
-
- pending: list[dict] = []
- path_selections = []
- for sel in valid_selections:
- path = sel['path']
- conf = sel.get('confidence', 0.7)
- item = selectable[path]
- node.confidence[path] = conf
-
- if item.get('is_leaf'):
- path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'})
- else:
- # Non-leaf → will be batched into a single next call
- pending.append({'path': path, 'confidence': conf})
- path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'self_only'})
-
- if path_selections:
- chunks = await _hydrate_paths_to_rows(
- db,
- path_selections=path_selections,
- user_id=user_id,
- namespace=namespace,
- document_id=document_id,
- )
- if chunks:
- connected = await hydrate_connected_target_rows(
- db=db,
- rows=chunks,
- exclude_document_ids=[],
- exclude_sections=[],
- )
- if connected:
- _owner_map = _build_connected_owner_map(chunks)
- for c in connected:
- if not c.get('owner_section_path'):
- c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or ''))
- chunks = chunks + connected
-
- _root_map = await _resolve_root_asset_owners(
- db,
- document_id=document_id,
- job_result_id=job_result_id,
- chunks=chunks,
- )
- if _root_map:
- for c in chunks:
- if c.get('owner_section_path'):
- continue
- cid = str(c.get('chunk_id') or '')
- if cid in _root_map:
- c['owner_section_path'] = _root_map[cid]
-
- for chunk in chunks:
- real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path')
- if real_path:
- node.add_leaf_chunks(str(real_path), [chunk])
-
- return action, asset_tools, node, pending
-
- except BudgetExceeded:
- raise
- except Exception as e:
- logger.error(f' navigate_step failed for doc={document_id}: {e}')
- return 'STOP', [], empty, []
-
-
-# ---------------------------------------------------------------------------
-# Tool: discovery_select_step (post-navigation discovery selection)
-# ---------------------------------------------------------------------------
-
-_MAX_DISCOVERY_PER_DOC = 3
async def discovery_select_step(
@@ -826,152 +133,22 @@ async def discovery_select_step(
llm_fn: LLMFn,
user_id: str,
namespace: str,
- doc_name: str = '',
+ doc_name: str = "",
discovery_hints: list[dict[str, Any]],
exclude_paths: set[str] | None = None,
revision_hint: str | None = None,
budget_snapshot: dict | None = None,
) -> DocTreeNode:
- """Post-navigation discovery selection step.
-
- After BFS navigation exhausts for a document, present discovery-found
- section paths (from bottom_discovery BM25) to the LLM for selection.
- Selected paths are hydrated as leaf content.
-
- For B-class documents (discovery-only, not KG-selected), this is the
- only navigation step — no prior BFS.
- """
- from shared.services.retrieval.app_service import _hydrate_paths_to_rows
-
- node = DocTreeNode(scope_path=None)
- if not discovery_hints:
- return node
-
- # Limit hints per document
- hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC]
-
- t0 = time.monotonic()
- try:
- # 1. Format hints for LLM (deduplicate by section_path)
- exclude_set = {
- normalize_section_path(path)
- for path in (exclude_paths or set())
- if path
- }
- hint_lines: list[str] = []
- hint_by_path: dict[str, dict] = {}
- for h in hints:
- sp = normalize_section_path(h.get('section_path', ''))
- if not sp or sp == 'Root':
- continue
- if sp in exclude_set:
- continue
- if sp in hint_by_path:
- continue # skip duplicate section_path
- summary = h.get('summary', '') or ''
- hint_lines.append(f'▸ path="{sp}"')
- if summary:
- clipped = summary[:300]
- hint_lines.append(f' {clipped}')
- hint_by_path[sp] = h
-
- if not hint_lines:
- return node
-
- items_text = '\n'.join(hint_lines)
-
- revision_context = ''
- if revision_hint:
- revision_context = (
- f'\nIMPORTANT: This is a REVISION round. '
- f'The previous search attempt failed because:\n'
- f'"{revision_hint}"\n'
- f'Adjust your selection accordingly. '
- f'If no candidate is relevant, return an EMPTY list [].\n'
- )
-
- prompt = _DISCOVERY_SELECT_PROMPT.format(
- doc_name=doc_name or document_id,
- budget_block=_format_budget_block(budget_snapshot),
- items=items_text,
- query=query,
- revision_context=revision_context,
- )
- response = await llm_fn(prompt)
- # Parse {"selections": [...]} response — reuse action parser's extraction
- parsed = _parse_action_response(response)
- selections = parsed.get('selections', [])
-
- logger.info(
- f' discovery_select_step doc="{doc_name}": '
- f'hints={len(hints)} selections={len(selections)}'
- )
-
- # 2. Hydrate selected paths
- valid_selections = [s for s in selections if s['path'] in hint_by_path]
- path_selections = []
- for sel in valid_selections:
- path = sel['path']
- conf = sel.get('confidence', 0.7)
- node.confidence[path] = conf
- path_selections.append({'path': path, 'confidence': conf})
-
- if path_selections:
- chunks = await _hydrate_paths_to_rows(
- db,
- path_selections=path_selections,
- user_id=user_id,
- namespace=namespace,
- document_id=document_id,
- )
- if chunks:
- connected = await hydrate_connected_target_rows(
- db=db,
- rows=chunks,
- exclude_document_ids=[],
- exclude_sections=[],
- )
- if connected:
- _owner_map = _build_connected_owner_map(chunks)
- for c in connected:
- if not c.get('owner_section_path'):
- c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or ''))
- chunks = chunks + connected
-
- # Resolve Root-stranded assets to their true owner sections
- _disc_job_result_id = next(
- (str(c['job_result_id']) for c in chunks if c.get('job_result_id')),
- None,
- )
- _root_map = await _resolve_root_asset_owners(
- db,
- document_id=document_id,
- job_result_id=_disc_job_result_id,
- chunks=chunks,
- ) if _disc_job_result_id else {}
- if _root_map:
- for c in chunks:
- if c.get('owner_section_path'):
- continue # already resolved by batch-level owner map
- cid = str(c.get('chunk_id') or '')
- if cid in _root_map:
- c['owner_section_path'] = _root_map[cid]
-
- for chunk in chunks:
- # Distribute chunk to its real path or fallback to the selection path
- real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path')
- if real_path:
- node.add_leaf_chunks(str(real_path), [chunk])
-
- latency = int((time.monotonic() - t0) * 1000)
- logger.info(
- f' discovery_select_step done: hydrated={len(node.leaf_content)} '
- f'latency={latency}ms'
- )
- return node
-
- except BudgetExceeded:
- raise
- except Exception as e:
- logger.error(f' discovery_select_step failed for doc={document_id}: {e}')
- return node
+ return await discovery_selection.discovery_select_step(
+ db,
+ document_id=document_id,
+ query=query,
+ llm_fn=llm_fn,
+ user_id=user_id,
+ namespace=namespace,
+ doc_name=doc_name,
+ discovery_hints=discovery_hints,
+ exclude_paths=exclude_paths,
+ revision_hint=revision_hint,
+ budget_snapshot=budget_snapshot,
+ )
diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py
index 96c125e49..13c60625e 100644
--- a/packages/shared-python/shared/services/retrieval/app_service.py
+++ b/packages/shared-python/shared/services/retrieval/app_service.py
@@ -1,966 +1,13 @@
from __future__ import annotations
-import asyncio
-import os
-import re
-import time
from typing import Any
-from loguru import logger
-from sqlalchemy import and_, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.core.database import get_db_context
-from shared.models.database.document import Document, DocumentChunk, DocumentSection, RetrievalHitStat
-from shared.services.retrieval.graph_service import GraphQueryService, is_excluded_section
-from shared.services.retrieval.lexical_text import normalize_section_path
-from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result
-from shared.services.retrieval.hit_stats_service import compute_importance_score, record_retrieval_hits
-from shared.services.retrieval.channels import path_channel, content_channel, term_channel
-from shared.services.storage.result_storage import get_result_storage
-from shared.models.database.job_result import JobResult
+from shared.services.retrieval.execution_plan import RetrievalExecutionPlan
+from shared.services.retrieval.scoring import merge_channels_rrf
-
-_MEDIA_CHUNK_TYPES = {'image', 'table'}
-
-_RRF_K = 60
-_CHANNEL_WEIGHT_PATH = 1.0
-_CHANNEL_WEIGHT_CONTENT = 2.0
-_CHANNEL_WEIGHT_TERM = 1.5
-_INTERNAL_RECALL_K_MULTIPLIER = 2
-_pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set()
-
-_DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = {
- 1: None,
- 2: {'text'},
- 3: {'image'},
- 4: {'table'},
- 5: {'text', 'image'},
- 6: {'text', 'table'},
-}
-
-
-def _resolve_allowed_chunk_types(data_type: int) -> set[str] | None:
- return _DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type)
-
-
-_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]')
-
-
-def _clean_content(content: str) -> str:
- return _PATH_REF_RE.sub('', content).strip()
-
-_PUBLIC_RESULT_FIELDS = {
- 'chunk_type', 'content', 'score', 'asset_url',
-}
-
-_PUBLIC_SOURCE_FIELDS = {
- 'document_id', 'source_file_name', 'section_path',
-}
-
-
-def _normalize_chunk_type(raw: str | None) -> str:
- return str(raw or '').strip().split('\n', 1)[0].lower()
-
-
-def _filter_excluded_rows(
- rows: list[dict[str, Any]],
- *,
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
-) -> list[dict[str, Any]]:
- filtered: list[dict[str, Any]] = []
- excluded_documents = set(exclude_document_ids)
- for row in rows:
- document_id = row.get('document_id')
- if document_id in excluded_documents:
- continue
- if is_excluded_section(
- document_id=document_id,
- section_path=row.get('section_path'),
- exclude_sections=exclude_sections,
- ):
- continue
- filtered.append(row)
- return filtered
-
-
-def _iter_connected_target_ids(row: dict[str, Any]) -> list[str]:
- metadata = row.get('chunk_metadata') or {}
- if not isinstance(metadata, dict):
- return []
-
- target_ids: list[str] = []
- for item in metadata.get('connect_to') or []:
- if not isinstance(item, dict):
- continue
- target_id = str(item.get('target') or '').strip()
- if target_id:
- target_ids.append(target_id)
- return target_ids
-
-
-async def hydrate_connected_target_rows(
- *,
- db: AsyncSession | None,
- rows: list[dict[str, Any]],
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
-) -> list[dict[str, Any]]:
- if db is None:
- return []
-
- existing_chunk_ids = {
- str(row.get('chunk_id') or '').strip()
- for row in rows
- if row.get('chunk_id')
- }
- target_ids_by_revision: dict[tuple[str, str], set[str]] = {}
- for row in rows:
- if _normalize_chunk_type(row.get('chunk_type')) != 'text':
- continue
- document_id = str(row.get('document_id') or '').strip()
- job_result_id = str(row.get('job_result_id') or '').strip()
- if not document_id or not job_result_id:
- continue
- for target_id in _iter_connected_target_ids(row):
- if target_id in existing_chunk_ids:
- continue
- target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(target_id)
-
- if not target_ids_by_revision:
- return []
-
- revision_filters = [
- and_(
- DocumentChunk.document_id == document_id,
- DocumentChunk.job_result_id == job_result_id,
- DocumentChunk.chunk_id.in_(sorted(target_ids)),
- )
- for (document_id, job_result_id), target_ids in target_ids_by_revision.items()
- if target_ids
- ]
- if not revision_filters:
- return []
-
- stmt = (
- select(Document, DocumentChunk, DocumentSection, JobResult)
- .join(DocumentChunk, DocumentChunk.document_id == Document.document_id)
- .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
- .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
- .where(or_(*revision_filters))
- .order_by(DocumentChunk.sort_order)
- )
- result = await db.execute(stmt)
-
- hydrated_rows: list[dict[str, Any]] = []
- for document, chunk, section, job_result in result.all():
- section_path = section.section_path if section else None
- hydrated_rows.append(
- {
- 'document_id': document.document_id,
- 'chunk_id': chunk.chunk_id,
- 'section_id': chunk.section_id,
- 'section_path': section_path,
- 'source_file_name': document.source_file_name,
- 'chunk_type': chunk.chunk_type,
- 'content': chunk.content,
- 'score': 0.0,
- 'file_path': chunk.file_path,
- 'chunk_metadata': chunk.chunk_metadata or {},
- 'job_result_id': chunk.job_result_id,
- 'job_id': job_result.job_id if job_result else None,
- 'sort_order': chunk.sort_order,
- }
- )
-
- return _filter_excluded_rows(
- hydrated_rows,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- )
-
-
-async def assemble_retrieval_results(
- *,
- db: AsyncSession | None = None,
- rows: list[dict[str, Any]],
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
- allowed_chunk_types: set[str] | None = None,
-) -> list[dict[str, Any]]:
- filtered_rows = _filter_excluded_rows(
- rows,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- )
- if allowed_chunk_types is not None:
- filtered_rows = [
- row for row in filtered_rows
- if _normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types
- ]
- hydrated_rows = await hydrate_connected_target_rows(
- db=db,
- rows=filtered_rows,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- )
- rows_by_chunk_id = {
- str(row.get('chunk_id') or ''): row
- for row in [*filtered_rows, *hydrated_rows]
- if row.get('chunk_id')
- }
-
- embedded_targets: set[str] = set()
- for row in filtered_rows:
- for target_id in _iter_connected_target_ids(row):
- if target_id in rows_by_chunk_id:
- embedded_targets.add(target_id)
-
- assembled: list[dict[str, Any]] = []
- for row in filtered_rows:
- if row.get('chunk_id') in embedded_targets:
- continue
- metadata = row.get('chunk_metadata') or {}
- if not isinstance(metadata, dict):
- metadata = {}
- assembled_row = dict(row)
- base_content = str(row.get('content') or '')
- if _normalize_chunk_type(row.get('chunk_type')) == 'text':
- connected_targets: list[tuple[int, str]] = []
- for target_id in _iter_connected_target_ids(row):
- target_row = rows_by_chunk_id.get(target_id)
- if not target_row:
- continue
- if _normalize_chunk_type(target_row.get('chunk_type')) != 'table':
- continue
- target_content = str(target_row.get('content') or '').strip()
- if target_content:
- sort_key = int(target_row.get('sort_order', 0) or 0)
- connected_targets.append((sort_key, target_content))
- connected_targets.sort(key=lambda x: x[0])
- related_parts = [content for _, content in connected_targets]
- if base_content and related_parts:
- assembled_row['content'] = '\n\n'.join([base_content, *related_parts])
- else:
- assembled_row['content'] = base_content
- else:
- assembled_row['content'] = base_content
- assembled_row['content'] = _clean_content(assembled_row['content'])
- assembled.append(assembled_row)
- return assembled
-
-
-
-
-
-def _merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
- if not rows:
- return rows
- groups: dict[str, list[dict[str, Any]]] = {}
- order: list[str] = []
- for row in rows:
- sp = row.get('section_path')
- if sp:
- key = f"{row.get('document_id', '')}::{sp}"
- else:
- key = row.get('chunk_id', '')
- if key not in groups:
- groups[key] = []
- order.append(key)
- groups[key].append(row)
-
- merged: list[dict[str, Any]] = []
- for key in order:
- group = groups[key]
- if len(group) == 1:
- merged.append(group[0])
- continue
- base = dict(group[0])
- base['content'] = '\n'.join(str(r.get('content', '')) for r in group)
- base['score'] = max(r.get('score', 0.0) for r in group)
- merged.append(base)
- return merged
-
-
-def merge_channels_rrf(
- channels: list[list[dict[str, Any]]],
- weights: list[float],
- top_k: int,
- k: int = _RRF_K,
-) -> list[dict[str, Any]]:
- """Reciprocal Rank Fusion across multiple retrieval channels."""
- score_dict: dict[str, float] = {}
- row_by_chunk_id: dict[str, dict[str, Any]] = {}
-
- for channel_idx, channel_rows in enumerate(channels):
- w = weights[channel_idx] if channel_idx < len(weights) else 1.0
- for rank, row in enumerate(channel_rows):
- chunk_id = str(row.get('chunk_id') or '')
- if not chunk_id:
- continue
- rrf_score = w / (k + rank + 1)
- score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score
- if chunk_id not in row_by_chunk_id:
- row_by_chunk_id[chunk_id] = row
-
- ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True)
- results: list[dict[str, Any]] = []
- for chunk_id, fused_score in ranked[:top_k]:
- row = row_by_chunk_id[chunk_id]
- results.append(dict(row, score=round(fused_score, 6)))
- return results
-
-
-async def list_graph_routed_chunks(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- query: str,
- top_k: int,
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
-) -> list[dict[str, Any]]:
- service = GraphQueryService()
- entry_document_ids = await service.find_entry_documents(
- db,
- user_id=user_id,
- namespace=namespace,
- query=query,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- )
- return await service.collect_candidate_chunks(
- db,
- user_id=user_id,
- namespace=namespace,
- entry_document_ids=entry_document_ids,
- query=query,
- top_k=top_k * _INTERNAL_RECALL_K_MULTIPLIER,
- exclude_sections=exclude_sections,
- )
-
-
-def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None:
- _pending_retrieval_hit_stat_tasks.discard(task)
-
- try:
- task.result()
- except asyncio.CancelledError:
- pass
- except Exception as e:
- logger.warning(f'Failed to record retrieval hit stats (ignored): {e}')
-
-
-def schedule_retrieval_hit_stats_update(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None:
- try:
- task = asyncio.create_task(
- _record_retrieval_hit_stats_best_effort(
- user_id=user_id,
- namespace=namespace,
- results=results,
- ),
- name=f'retrieval_hit_stats:{user_id}:{namespace}',
- )
- _pending_retrieval_hit_stat_tasks.add(task)
- task.add_done_callback(_finalize_retrieval_hit_stats_task)
- except Exception as e:
- logger.warning(f'Failed to schedule retrieval hit stats update (ignored): {e}')
-
-
-async def drain_retrieval_hit_stats_updates(timeout_seconds: float = 2.0) -> None:
- if not _pending_retrieval_hit_stat_tasks:
- return
-
- pending_tasks = tuple(_pending_retrieval_hit_stat_tasks)
-
- try:
- await asyncio.wait_for(
- asyncio.gather(*pending_tasks, return_exceptions=True),
- timeout=timeout_seconds,
- )
- except asyncio.TimeoutError:
- for task in pending_tasks:
- if not task.done():
- task.cancel()
-
- await asyncio.gather(*pending_tasks, return_exceptions=True)
-
-
-async def _record_retrieval_hit_stats_best_effort(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None:
- try:
- async with get_db_context() as db:
- await record_retrieval_hits(db, user_id=user_id, namespace=namespace, results=results)
- await db.commit()
- except Exception as e:
- logger.warning(f'Failed to record retrieval hit stats (ignored): {e}')
-
-
-def _with_citation(row: dict[str, Any]) -> dict[str, Any]:
- citation = {
- 'document_id': row.get('document_id'),
- 'chunk_id': row.get('chunk_id'),
- 'source_file_name': row.get('source_file_name'),
- 'section_path': row.get('section_path'),
- }
- return {**row, 'citation': citation}
-
-
-def _to_public_source(row: dict[str, Any]) -> dict[str, Any]:
- return {field: row.get(field) for field in _PUBLIC_SOURCE_FIELDS}
-
-
-def _is_media_chunk(row: dict[str, Any]) -> bool:
- return _normalize_chunk_type(row.get('chunk_type')) in _MEDIA_CHUNK_TYPES
-
-
-async def generate_retrieval_asset_url(*, job_id: str, artifact_ref: str) -> str | None:
- return get_result_storage().generate_artifact_url(job_id=job_id, artifact_ref=artifact_ref)
-
-
-def _is_client_result_artifact_ref(asset_ref: str | None) -> bool:
- return get_result_storage().normalize_artifact_ref(asset_ref) is not None
-
-
-async def _to_public_response(response: dict[str, Any]) -> dict[str, Any]:
- public_response = {
- 'namespace': response.get('namespace'),
- 'query': response.get('query'),
- 'router_used': response.get('router_used'),
- 'results': [],
- }
-
- # Forward agentic evidence fields when present
- if response.get('answer_text') is not None:
- public_response['answer_text'] = response['answer_text']
- if response.get('referenced_chunks') is not None:
- public_response['referenced_chunks'] = response['referenced_chunks']
-
- public_results: list[dict[str, Any]] = []
- for row in response.get('results', []):
- artifact_ref = row.get('file_path')
- asset_url = None
- if _is_media_chunk(row) and _is_client_result_artifact_ref(artifact_ref) and row.get('job_id'):
- try:
- asset_url = await generate_retrieval_asset_url(
- job_id=str(row['job_id']),
- artifact_ref=str(artifact_ref),
- )
- except Exception as e:
- logger.warning(f'Failed to generate retrieval asset URL (ignored): {e}')
-
- public_row: dict[str, Any] = {}
- for field in _PUBLIC_RESULT_FIELDS:
- if field == 'asset_url':
- if asset_url:
- public_row['asset_url'] = asset_url
- elif field in row:
- public_row[field] = row[field]
- if 'source' in row:
- public_row['source'] = row['source']
- else:
- public_row['source'] = _to_public_source(row)
- public_results.append(public_row)
-
- public_response['results'] = public_results
- return public_response
-
-
-def _get_row_path(row: dict[str, Any]) -> str:
- """Extract the canonical path from a row for deduplication."""
- return str(row.get('section_path') or row.get('source_chunk_path') or '')
-
-
-def _get_candidate_key(row: dict[str, Any]) -> str:
- path = _get_row_path(row)
- if path:
- return f'path:{path}'
- chunk_id = str(row.get('chunk_id') or '').strip()
- return f'chunk:{chunk_id}' if chunk_id else ''
-
-
-def _normalize_row_scores(
- rows: list[dict[str, Any]],
- *,
- source_field: str,
- target_field: str,
- default: float,
-) -> None:
- if not rows:
- return
- values = [float(row.get(source_field, 0.0) or 0.0) for row in rows]
- min_score = min(values)
- max_score = max(values)
- if max_score <= 0.0 and min_score <= 0.0:
- for row in rows:
- row[target_field] = 0.0
- return
- if max_score == min_score:
- for row in rows:
- row[target_field] = default
- return
- denominator = max_score - min_score
- for row in rows:
- raw_score = float(row.get(source_field, 0.0) or 0.0)
- row[target_field] = round((raw_score - min_score) / denominator, 6)
-
-
-def _importance_multiplier(
- rows: list[dict[str, Any]],
- *,
- raw_field: str = 'importance_raw_score',
- low: float = 0.1,
- high: float = 2.0,
-) -> None:
- """Apply adaptive sigmoid-based importance boost to agent/discovery scores.
-
- Uses median of ``raw_field`` as center and IQR as spread so the curve
- adapts to any KB size without hard-coded thresholds. When all values
- are identical (IQR ≈ 0) the multiplier is 1.0 (neutral).
-
- Output range ``[low, high]`` — default [0.1, 2.0] — is the only
- configured constant: max 2× boost, min 10%. The function modifies
- ``agent_score`` and ``discovery_score`` **in place**.
- """
- import math
-
- if not rows:
- return
-
- values = sorted(float(r.get(raw_field, 0.0) or 0.0) for r in rows)
- n = len(values)
- median = values[n // 2] if n % 2 else (values[n // 2 - 1] + values[n // 2]) / 2
- q1 = values[n // 4] if n >= 4 else values[0]
- q3 = values[3 * n // 4] if n >= 4 else values[-1]
- iqr = q3 - q1
-
- for row in rows:
- raw = float(row.get(raw_field, 0.0) or 0.0)
- if iqr <= 1e-9:
- mult = 1.0
- else:
- z = (raw - median) / iqr
- s = 1.0 / (1.0 + math.exp(-z))
- mult = low + (high - low) * s
- row['importance_multiplier'] = round(mult, 4)
- row['agent_score'] = round(
- float(row.get('agent_score', 0.0) or 0.0) * mult, 6,
- )
- row['discovery_score'] = round(
- float(row.get('discovery_score', 0.0) or 0.0) * mult, 6,
- )
-
-
-async def _load_chunk_importance_scores(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- rows: list[dict[str, Any]],
-) -> dict[str, float]:
- chunk_ids = sorted({
- str(row.get('chunk_id') or '').strip()
- for row in rows
- if row.get('chunk_id')
- })
- if not chunk_ids:
- return {}
- stmt = (
- select(
- RetrievalHitStat.chunk_id,
- RetrievalHitStat.hit_count,
- RetrievalHitStat.last_hit_at,
- RetrievalHitStat.created_at,
- )
- .where(RetrievalHitStat.user_id == user_id)
- .where(RetrievalHitStat.namespace == namespace)
- .where(RetrievalHitStat.hit_kind == 'chunk')
- .where(RetrievalHitStat.chunk_id.in_(chunk_ids))
- )
- result = await db.execute(stmt)
- importance_scores: dict[str, float] = {}
- for chunk_id, hit_count, last_hit_at, created_at in result.all():
- if not chunk_id:
- continue
- importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at)
- return importance_scores
-
-
-def _rank_candidates_by_path(
- discovery_rows: list[dict[str, Any]],
- routed_rows: list[dict[str, Any]],
- top_k: int,
-) -> list[dict[str, Any]]:
- """Rank discovery and routed candidates in one comparable path space."""
- merged: dict[str, dict[str, Any]] = {}
- insertion_order: dict[str, int] = {}
- counter = 0
-
- for row in discovery_rows:
- key = _get_candidate_key(row)
- if not key:
- continue
- candidate = dict(row)
- candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0)
- candidate['agent_score'] = 0.0
- candidate.setdefault('hydrate_mode', 'chunks')
- merged[key] = candidate
- insertion_order[key] = counter
- counter += 1
-
- for row in routed_rows:
- key = _get_candidate_key(row)
- if not key:
- continue
- routed_agent_score = float(row.get('agent_score', 0.0) or 0.0)
- if key not in merged:
- candidate = dict(row)
- candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0)
- candidate['agent_score'] = routed_agent_score
- merged[key] = candidate
- insertion_order[key] = counter
- counter += 1
- continue
- candidate = merged[key]
- candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score)
- if not candidate.get('source_chunk_path') and row.get('source_chunk_path'):
- candidate['source_chunk_path'] = row.get('source_chunk_path')
- if not candidate.get('section_path') and row.get('section_path'):
- candidate['section_path'] = row.get('section_path')
-
- # ── Dual-priority ranking ────────────────────────────────────────────
- # When the agent produced results (routed_rows non-empty), rows with
- # agent_score=0 are demoted to a fallback pool. Primary sort is by
- # agent_score (includes importance boost from _importance_multiplier),
- # with discovery_score as tiebreaker.
- has_agent_results = len(routed_rows) > 0
-
- primary_rows: list[dict[str, Any]] = []
- fallback_rows: list[dict[str, Any]] = []
-
- for key, row in merged.items():
- agent_score = float(row.get('agent_score', 0.0) or 0.0)
- discovery_score = float(row.get('discovery_score', 0.0) or 0.0)
- row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6)
- row['score'] = row['evidence_score']
- row['_candidate_order'] = insertion_order[key]
-
- if has_agent_results and agent_score <= 0.0:
- fallback_rows.append(row)
- else:
- primary_rows.append(row)
-
- def _sort_key(row):
- return (
- float(row.get('agent_score', 0.0) or 0.0),
- float(row.get('discovery_score', 0.0) or 0.0),
- -int(row.get('_candidate_order', 0) or 0),
- )
-
- primary_rows.sort(key=_sort_key, reverse=True)
- ranked_rows = primary_rows[:top_k]
-
- # Back-fill from fallback if primary results are insufficient
- if len(ranked_rows) < top_k and fallback_rows:
- fallback_rows.sort(key=_sort_key, reverse=True)
- ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)])
-
- for row in ranked_rows:
- row.pop('_candidate_order', None)
- return ranked_rows
-
-
-async def _count_scoped_chunks(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- exclude_document_ids: list[str],
- allowed_chunk_types: set[str] | None,
-) -> int:
- stmt = (
- select(func.count(DocumentChunk.id))
- .join(Document, (Document.document_id == DocumentChunk.document_id) & (Document.current_job_result_id == DocumentChunk.job_result_id))
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- )
- if exclude_document_ids:
- stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
- if allowed_chunk_types is not None:
- stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types)))
- result = await db.execute(stmt)
- return result.scalar() or 0
-
-
-async def _load_all_scoped_chunks(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
- allowed_chunk_types: set[str] | None,
- signal_paths: list[str],
- filter_mode: str,
-) -> list[dict[str, Any]]:
- stmt = (
- select(Document, DocumentChunk, DocumentSection, JobResult)
- .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id))
- .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
- .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .order_by(DocumentChunk.sort_order)
- )
- if exclude_document_ids:
- stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
- if allowed_chunk_types is not None:
- stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types)))
-
- result = await db.execute(stmt)
- rows: list[dict[str, Any]] = []
- for document, chunk, section, job_result in result.all():
- section_path = section.section_path if section else None
- if is_excluded_section(document_id=document.document_id, section_path=section_path, exclude_sections=exclude_sections):
- continue
- if signal_paths and section_path:
- path_lower = section_path.lower()
- matches_any = any(kw.lower() in path_lower for kw in signal_paths)
- if filter_mode == 'keep' and not matches_any:
- continue
- if filter_mode == 'delete' and matches_any:
- continue
- rows.append({
- 'document_id': document.document_id,
- 'chunk_id': chunk.chunk_id,
- 'section_id': chunk.section_id,
- 'section_path': section_path,
- 'source_file_name': document.source_file_name,
- 'chunk_type': chunk.chunk_type,
- 'content': chunk.content,
- 'score': 1.0,
- 'file_path': chunk.file_path,
- 'chunk_metadata': chunk.chunk_metadata or {},
- 'job_result_id': chunk.job_result_id,
- 'job_id': job_result.job_id if job_result else None,
- 'sort_order': chunk.sort_order,
- })
- return rows
-
-
-async def _hydrate_paths_to_rows(
- db: AsyncSession,
- *,
- path_selections: list[dict[str, Any]],
- user_id: str,
- namespace: str,
- document_id: str | None = None,
-) -> list[dict[str, Any]]:
- """Load full chunk rows by section_path or source_chunk_path.
-
- When *document_id* is provided the query is scoped to that single
- document, preventing cross-document collisions on generic paths
- such as ``Root``.
-
- Supports hydrate_mode branching:
- - 'chunks' (default): all chunk types under the section subtree
- - 'outline': synthetic row from section metadata, no real chunks
- - 'assets_only': only image + table chunks
- - 'image_only': only image chunks
- - 'table_only': only table chunks
- """
- if not path_selections:
- return []
-
- # Group selections by hydrate_mode
- confidence_by_path: dict[str, float] = {}
- mode_by_path: dict[str, str] = {}
- ordered_paths: list[str] = []
- for item in path_selections:
- raw_path = str(item.get('path') or '').strip()
- path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path
- if not path:
- continue
- confidence = float(item.get('confidence', 0.0) or 0.0)
- hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower()
- if path not in confidence_by_path:
- ordered_paths.append(path)
- confidence_by_path[path] = confidence
- mode_by_path[path] = hydrate_mode
- else:
- confidence_by_path[path] = max(confidence_by_path[path], confidence)
- if not ordered_paths:
- return []
-
- # Separate outline paths from chunk-loading paths
- outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline']
- chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline']
-
- rows: list[dict[str, Any]] = []
-
- # ── Outline mode: synthesize rows from section metadata ──────────────
- if outline_paths:
- outline_section_filters = []
- for path in outline_paths:
- outline_section_filters.append(DocumentSection.section_path == path)
-
- outline_stmt = (
- select(Document, DocumentSection)
- .join(DocumentSection, (DocumentSection.document_id == Document.document_id)
- & (DocumentSection.job_result_id == Document.current_job_result_id))
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(or_(*outline_section_filters))
- )
- if document_id:
- outline_stmt = outline_stmt.where(Document.document_id == document_id)
- outline_result = await db.execute(outline_stmt)
- for document, section in outline_result.all():
- agent_score = confidence_by_path.get(section.section_path, 0.0)
- summary_text = (section.summary or '').strip()
- title_text = (section.section_title or '').strip()
- content = f'[Outline] {title_text}'
- if summary_text:
- content += f'\n{summary_text}'
- rows.append({
- 'document_id': document.document_id,
- 'chunk_id': f'outline_{section.section_id}',
- 'section_id': section.section_id,
- 'section_path': section.section_path,
- 'source_file_name': document.source_file_name,
- 'chunk_type': 'outline',
- 'content': content,
- 'score': agent_score,
- 'agent_score': agent_score,
- 'file_path': None,
- 'chunk_metadata': {},
- 'job_result_id': section.job_result_id,
- 'job_id': None,
- 'source_chunk_path': None,
- 'sort_order': section.sort_order,
- 'hydrate_mode': 'outline',
- })
-
- # ── Chunk modes: load real chunks with optional type filters ─────────
- if chunk_paths:
- section_path_filters = []
- # Separate self_only paths (exact match only, no descendant LIKE)
- # from regular chunk paths (exact + descendant subtree match)
- self_only_paths = {p for p in chunk_paths if mode_by_path.get(p) == 'self_only'}
- for path in chunk_paths:
- section_path_filters.append(DocumentSection.section_path == path)
- if path not in self_only_paths:
- section_path_filters.append(DocumentSection.section_path.like(f'{path} / %'))
-
- stmt = (
- select(Document, DocumentChunk, DocumentSection, JobResult)
- .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id)
- & (DocumentChunk.job_result_id == Document.current_job_result_id))
- .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
- .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(
- or_(
- *section_path_filters,
- DocumentChunk.source_chunk_path.in_(chunk_paths),
- )
- )
- )
- if document_id:
- stmt = stmt.where(Document.document_id == document_id)
- result = await db.execute(stmt)
-
- # Build a map of path → allowed chunk_types based on hydrate_mode
- _MODE_ALLOWED_TYPES: dict[str, set[str] | None] = {
- 'chunks': None, # all types
- 'self_only': None, # all types, but without descendant filtering
- 'assets_only': {'image', 'table'},
- 'image_only': {'image'},
- 'table_only': {'table'},
- }
-
- seen_paths: set[str] = set()
- for document, chunk, section, job_result in result.all():
- row_path = (section.section_path if section else None) or chunk.source_chunk_path or ''
- if row_path in seen_paths:
- continue
-
- # Find which ordered path this row belongs to
- matched_path = row_path
- if section and section.section_path not in confidence_by_path:
- matched_path = next(
- (
- path for path in chunk_paths
- if section.section_path == path or section.section_path.startswith(f'{path} / ')
- ),
- row_path,
- )
-
- # Check chunk_type filter based on hydrate_mode
- path_mode = mode_by_path.get(matched_path, 'chunks')
- allowed_types = _MODE_ALLOWED_TYPES.get(path_mode)
- if allowed_types is not None:
- chunk_type_lower = (chunk.chunk_type or '').strip().lower()
- if chunk_type_lower not in allowed_types:
- continue
-
- seen_paths.add(row_path)
- agent_score = confidence_by_path.get(matched_path, 0.0)
- rows.append({
- 'document_id': document.document_id,
- 'chunk_id': chunk.chunk_id,
- 'section_id': chunk.section_id,
- 'section_path': section.section_path if section else None,
- 'source_file_name': document.source_file_name,
- 'chunk_type': chunk.chunk_type,
- 'content': chunk.content,
- 'score': agent_score,
- 'agent_score': agent_score,
- 'file_path': chunk.file_path,
- 'chunk_metadata': chunk.chunk_metadata or {},
- 'job_result_id': chunk.job_result_id,
- 'job_id': job_result.job_id if job_result else None,
- 'source_chunk_path': chunk.source_chunk_path,
- 'sort_order': chunk.sort_order,
- 'hydrate_mode': path_mode,
- })
-
- # ── Sort by agent-selected order ─────────────────────────────────────
- path_order = {p: idx for idx, p in enumerate(ordered_paths)}
-
- def _row_sort_key(row: dict[str, Any]) -> int:
- row_path = _get_row_path(row)
- if row_path in path_order:
- return path_order[row_path]
- for path, idx in path_order.items():
- if row_path.startswith(f'{path} / '):
- return idx
- return 10**9
-
- rows.sort(key=_row_sort_key)
- hydrated_paths = {_get_row_path(r) for r in rows}
- resolved_inputs = {
- path for path in ordered_paths
- if path in hydrated_paths or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths)
- }
- # Outline paths are always resolved (synthesized)
- resolved_inputs |= set(outline_paths)
- missed = len(ordered_paths) - len(resolved_inputs)
- if missed > 0:
- missing_paths = [p for p in ordered_paths if p not in resolved_inputs]
- logger.warning(
- f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); '
- f'missing[:5]={missing_paths[:5]}'
- )
- else:
- logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved')
- return rows
+__all__ = ["merge_channels_rrf", "run_retrieval_query"]
async def run_retrieval_query(
@@ -974,7 +21,7 @@ async def run_retrieval_query(
exclude_sections: list[dict[str, str]],
data_type: int = 1,
signal_paths: list[str] | None = None,
- filter_mode: str = 'delete',
+ filter_mode: str = "delete",
channels: list[str] | None = None,
channel_weights: dict[str, float] | None = None,
rerank: bool = False,
@@ -982,415 +29,22 @@ async def run_retrieval_query(
internal_recall_k: int | None = None,
use_agentic: bool | None = None,
) -> dict[str, Any]:
- """Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly."""
- t_start = time.monotonic()
- query = query.strip()
- logger.info('\n' + '█' * 70)
- logger.info(' 🚀 RETRIEVAL PIPELINE START')
- logger.info(f' query="{query}"')
- logger.info(f' user={user_id} ns={namespace} top_k={top_k} data_type={data_type}')
- logger.info(f' exclude_docs={exclude_document_ids} exclude_secs={len(exclude_sections)}')
- logger.info('█' * 70)
-
- if not query:
- logger.info(' ⛔ Empty query filtered, skipping retrieval pipeline')
- return {
- "namespace": namespace,
- "query": query,
- "router_used": "empty_query_filtered",
- "results": [],
- }
-
- allowed_chunk_types = _resolve_allowed_chunk_types(data_type)
- effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER
- logger.info(f' allowed_chunk_types={allowed_chunk_types} effective_recall_k={effective_recall_k} signal_paths={signal_paths} filter_mode={filter_mode} rerank={rerank} threshold={threshold}')
-
- cache_extra = dict(
- data_type=data_type,
- signal_paths=signal_paths,
- filter_mode=filter_mode,
- channels=channels,
- channel_weights=channel_weights,
- rerank=rerank,
- threshold=threshold,
- internal_recall_k=internal_recall_k,
- # Always True: agentic mode now always routes through workflow
- decomposition_enabled=True,
- )
-
- cache_version: int | None = None
- try:
- cache_version, cached = await get_cached_retrieval_query_result(
- user_id=user_id,
- namespace=namespace,
- query=query,
- top_k=top_k,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- **cache_extra,
- )
- if cached:
- logger.info(f'retrieval: cache_hit=True version={cache_version}')
- try:
- schedule_retrieval_hit_stats_update(
- user_id=user_id,
- namespace=namespace,
- results=cached.get("results", []),
- )
- except Exception as e:
- logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}")
- return await _to_public_response(cached)
- except Exception as e:
- logger.warning(f"Failed to read retrieval cache (ignored): {e}")
-
- logger.debug(f' 📦 Cache miss (version={cache_version}), running full pipeline')
-
- # ── Small KB optimization ──
- try:
- total_chunk_count = await _count_scoped_chunks(
- db, user_id=user_id, namespace=namespace,
- exclude_document_ids=exclude_document_ids,
- allowed_chunk_types=allowed_chunk_types,
- )
- except Exception as e:
- logger.warning(f"Failed to count scoped chunks, skipping small KB optimization: {e}")
- total_chunk_count = top_k + 1
- logger.info(f'\n 📊 Total chunks in scope: {total_chunk_count}')
- if total_chunk_count <= top_k:
- logger.info(f' ⚡ Small KB optimization: {total_chunk_count} chunks <= top_k={top_k}, returning all')
- all_rows = await _load_all_scoped_chunks(
- db, user_id=user_id, namespace=namespace,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths or [],
- filter_mode=filter_mode,
- )
- logger.info(f' small_kb load: loaded={len(all_rows)} rows after signal/exclude filters')
- assembled_rows = await assemble_retrieval_results(
- db=db, rows=all_rows,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- allowed_chunk_types=allowed_chunk_types,
- )
- results = [_with_citation(row) for row in assembled_rows]
- response = {
- "namespace": namespace, "query": query,
- "router_used": "small_kb_all", "results": results,
- }
- if cache_version is not None:
- try:
- await set_cached_retrieval_query_result(
- user_id=user_id, namespace=namespace, version=cache_version,
- query=query, top_k=top_k,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- response=response, **cache_extra,
- )
- except Exception as e:
- logger.warning(f"Failed to write retrieval cache (ignored): {e}")
- try:
- schedule_retrieval_hit_stats_update(user_id=user_id, namespace=namespace, results=results)
- except Exception as e:
- logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}")
- elapsed_total = round((time.monotonic() - t_start) * 1000)
- logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms')
- return await _to_public_response(response)
-
- # ══ Route: agentic (unified workflow) vs legacy ══
- if use_agentic is not None:
- _agentic_enabled = use_agentic
- else:
- _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'true') == 'true'
- if _agentic_enabled:
- # ── Unified agentic path via WorkflowOrchestrator ──
- # Simple queries: planner returns a single-step plan (no decomposition).
- # Complex queries: planner returns a multi-step plan with synthesize.
- # Both go through the same code path.
- from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator
-
- workflow = WorkflowOrchestrator()
- workflow_result = await workflow.run(
- db,
- user_id=user_id,
- namespace=namespace,
- query=query,
- top_k=top_k,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- data_type=data_type,
- signal_paths=signal_paths,
- filter_mode=filter_mode,
- channels=channels,
- channel_weights=channel_weights,
- )
-
- # Enrich referenced_chunks with asset URLs (images/tables)
- enriched_refs: list[dict[str, Any]] = []
- for ref in workflow_result.referenced_chunks:
- enriched = dict(ref)
- chunk_type = _normalize_chunk_type(ref.get('chunk_type'))
- artifact_ref = ref.get('file_path', '')
- job_id = ref.get('job_id', '')
- if chunk_type in _MEDIA_CHUNK_TYPES and job_id and _is_client_result_artifact_ref(artifact_ref):
- try:
- asset_url = await generate_retrieval_asset_url(
- job_id=str(job_id), artifact_ref=str(artifact_ref),
- )
- if asset_url:
- enriched['asset_url'] = asset_url
- except Exception as e:
- logger.warning(f'Failed to generate agentic asset URL (ignored): {e}')
- enriched_refs.append(enriched)
-
- response = workflow_result.to_api_response()
- # Override referenced_chunks with enriched versions
- response['referenced_chunks'] = enriched_refs
-
- if cache_version is not None:
- try:
- await set_cached_retrieval_query_result(
- user_id=user_id, namespace=namespace, version=cache_version,
- query=query, top_k=top_k,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- response=response, **cache_extra,
- )
- except Exception as e:
- logger.warning(f"Failed to write retrieval cache (ignored): {e}")
-
- try:
- schedule_retrieval_hit_stats_update(
- user_id=user_id, namespace=namespace,
- results=enriched_refs,
- )
- except Exception as e:
- logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}")
-
- elapsed_total = round((time.monotonic() - t_start) * 1000)
- logger.info(
- f'\n{"█" * 70}\n'
- f' ✅ AGENTIC RETRIEVAL COMPLETE: '
- f'{len(enriched_refs)} chunks | '
- f'answer={len(workflow_result.answer_text)} chars | '
- f'router={workflow_result.router_used} | {elapsed_total}ms\n'
- f'{"█" * 70}'
- )
-
- return await _to_public_response(response)
-
- else:
-
- # ── LEGACY path (existing code, unchanged) ──
-
- # ── Channel execution ──
- active_channels = set(channels) if channels else {'path', 'content', 'term'}
- logger.info(f'\n 📡 PHASE 1: Bottom-Layer Discovery (channels={sorted(active_channels)})')
- logger.info(f' effective_recall_k={effective_recall_k}')
-
- path_rows: list[dict[str, Any]] = []
- content_rows: list[dict[str, Any]] = []
- term_rows: list[dict[str, Any]] = []
-
- if 'path' in active_channels:
- t_ch = time.monotonic()
- path_rows = await path_channel(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=effective_recall_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths, filter_mode=filter_mode,
- )
- elapsed_ch = round((time.monotonic() - t_ch) * 1000)
- logger.info(f'\n 📡 path_channel: {len(path_rows)} rows in {elapsed_ch}ms')
- for i, r in enumerate(path_rows[:5]):
- logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}')
- if len(path_rows) > 5:
- logger.info(f' ... and {len(path_rows) - 5} more')
-
- if 'content' in active_channels:
- t_ch = time.monotonic()
- content_rows = await content_channel(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=effective_recall_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths, filter_mode=filter_mode,
- )
- elapsed_ch = round((time.monotonic() - t_ch) * 1000)
- logger.info(f'\n 📡 content_channel: {len(content_rows)} rows in {elapsed_ch}ms')
- for i, r in enumerate(content_rows[:5]):
- logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} content={str(r.get("content",""))[:80]}')
- if len(content_rows) > 5:
- logger.info(f' ... and {len(content_rows) - 5} more')
-
- if 'term' in active_channels:
- t_ch = time.monotonic()
- term_rows = await term_channel(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=effective_recall_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections, allowed_chunk_types=allowed_chunk_types,
- signal_paths=signal_paths, filter_mode=filter_mode,
- )
- elapsed_ch = round((time.monotonic() - t_ch) * 1000)
- logger.info(f'\n 📡 term_channel: {len(term_rows)} rows in {elapsed_ch}ms')
- for i, r in enumerate(term_rows[:5]):
- logger.info(f' [{i}] score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")} type={r.get("chunk_type","?")}')
- if len(term_rows) > 5:
- logger.info(f' ... and {len(term_rows) - 5} more')
-
- # ── RRF fusion with configurable weights ──
- default_weights = {
- 'path': _CHANNEL_WEIGHT_PATH,
- 'content': _CHANNEL_WEIGHT_CONTENT,
- 'term': _CHANNEL_WEIGHT_TERM,
- }
- effective_weights = {**default_weights, **(channel_weights or {})}
-
- channel_lists: list[list[dict[str, Any]]] = []
- weight_list: list[float] = []
-
- if path_rows:
- channel_lists.append(path_rows)
- weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH))
- if content_rows:
- channel_lists.append(content_rows)
- weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT))
- if term_rows:
- channel_lists.append(term_rows)
- weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM))
-
- fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else []
- logger.info(f'\n 🔀 RRF Fusion: {len(fused_rows)} rows from {len(channel_lists)} channels (weights={dict(zip(["path","content","term"][:len(weight_list)], weight_list))})')
- for i, r in enumerate(fused_rows[:5]):
- logger.info(f' [{i}] rrf_score={r.get("score",0):.4f} path={r.get("section_path","") or r.get("source_chunk_path","")}')
- if len(fused_rows) > 5:
- logger.info(f' ... and {len(fused_rows) - 5} more')
-
- # ── Section merging ──
- pre_merge = len(fused_rows)
- fused_rows = _merge_same_section_rows(fused_rows)
- if len(fused_rows) != pre_merge:
- logger.info(f'retrieval: section_merge={pre_merge}->{len(fused_rows)}')
-
- # ── Threshold filtering ──
- if threshold > 0.0 and fused_rows:
- pre_count = len(fused_rows)
- fused_rows = [row for row in fused_rows if row.get('score', 0.0) >= threshold]
- logger.info(f'retrieval: threshold_filter={pre_count}->{len(fused_rows)} (threshold={threshold})')
-
- if fused_rows:
- _normalize_row_scores(
- fused_rows,
- source_field='score',
- target_field='discovery_score',
- default=0.5,
- )
-
- # ── Legacy graph routing ──
- logger.info('\n 🧭 PHASE 2: Legacy Graph Routing')
- router_used = 'discovery_only'
- agent_rows: list[dict[str, Any]] = []
-
- try:
- agent_rows = await list_graph_routed_chunks(
- db, user_id=user_id, namespace=namespace, query=query,
- top_k=top_k, exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- )
- if agent_rows:
- router_used = 'discovery+graph'
- logger.info(f' 📊 Graph routing: {len(agent_rows)} rows')
- except Exception as exc:
- logger.error(f' ❌ Graph routing failed (ignored): {exc}')
- agent_rows = []
-
- if agent_rows:
- _normalize_row_scores(
- agent_rows,
- source_field='score',
- target_field='agent_score',
- default=0.5,
- )
-
- combined_rows = [*fused_rows, *agent_rows]
- if combined_rows:
- try:
- chunk_importance_scores = await _load_chunk_importance_scores(
- db,
- user_id=user_id,
- namespace=namespace,
- rows=combined_rows,
- )
- except Exception as exc:
- logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}')
- chunk_importance_scores = {}
- for row in combined_rows:
- row['importance_raw_score'] = float(chunk_importance_scores.get(str(row.get('chunk_id') or ''), 0.0) or 0.0)
-
- ranked_rows = _rank_candidates_by_path(fused_rows, agent_rows, top_k)
- if ranked_rows:
- logger.info(f'\n 🧮 Unified candidate ranking: {len(ranked_rows)} rows')
- for i, row in enumerate(ranked_rows[:10]):
- logger.info(
- ' '
- f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} '
- f'discovery={row.get("discovery_score", 0.0):.4f} '
- f'agent={row.get("agent_score", 0.0):.4f} '
- f'path={_get_row_path(row)}'
- )
-
- assembled_rows = await assemble_retrieval_results(
- db=db,
- rows=ranked_rows,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- allowed_chunk_types=allowed_chunk_types,
- )
- results = [_with_citation(row) for row in assembled_rows]
-
- response = {
+ request: dict[str, Any] = {
+ "db": db,
+ "user_id": user_id,
"namespace": namespace,
"query": query,
- "router_used": router_used,
- "results": results,
+ "top_k": top_k,
+ "exclude_document_ids": exclude_document_ids,
+ "exclude_sections": exclude_sections,
+ "data_type": data_type,
+ "signal_paths": signal_paths,
+ "filter_mode": filter_mode,
+ "channels": channels,
+ "channel_weights": channel_weights,
+ "rerank": rerank,
+ "threshold": threshold,
+ "internal_recall_k": internal_recall_k,
+ "use_agentic": use_agentic,
}
-
- if cache_version is not None:
- try:
- await set_cached_retrieval_query_result(
- user_id=user_id,
- namespace=namespace,
- version=cache_version,
- query=query,
- top_k=top_k,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- response=response,
- **cache_extra,
- )
- except Exception as e:
- logger.warning(f"Failed to write retrieval cache (ignored): {e}")
-
- try:
- schedule_retrieval_hit_stats_update(
- user_id=user_id,
- namespace=namespace,
- results=results,
- )
- except Exception as e:
- logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}")
-
- elapsed_total = round((time.monotonic() - t_start) * 1000)
- logger.info(f'\n{"█" * 70}')
- logger.info(f' ✅ RETRIEVAL COMPLETE: {len(results)} results | router={router_used} | {elapsed_total}ms')
- for i, r in enumerate(results[:10]):
- src = r.get('source', {})
- logger.info(
- f' [{i+1}] type={r.get("chunk_type","?")} score={r.get("score",0):.4f}'
- f' path={src.get("section_path","")}'
- f' file={src.get("source_file_name","")}'
- )
- if len(results) > 10:
- logger.info(f' ... and {len(results) - 10} more')
- logger.info(f'{"█" * 70}')
-
- return await _to_public_response(response)
+ return await RetrievalExecutionPlan(request).execute()
diff --git a/packages/shared-python/shared/services/retrieval/assets.py b/packages/shared-python/shared/services/retrieval/assets.py
new file mode 100644
index 000000000..035825430
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/assets.py
@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+
+from shared.services.retrieval.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type
+from shared.services.storage.result_storage import get_result_storage
+
+
+def _normalize_artifact_ref(asset_ref: object) -> str | None:
+ return get_result_storage().normalize_artifact_ref(
+ None if asset_ref is None else str(asset_ref)
+ )
+
+
+def _is_retrieval_media_row(row: dict[str, Any]) -> bool:
+ raw_chunk_type = row.get("chunk_type") or row.get("type")
+ return normalize_chunk_type(raw_chunk_type) in MEDIA_CHUNK_TYPES
+
+
+def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None:
+ job_id = str(row.get("job_id") or "").strip()
+ if not job_id or not _is_retrieval_media_row(row):
+ return None
+
+ artifact_ref = _normalize_artifact_ref(row.get("file_path"))
+ if artifact_ref is None:
+ return None
+
+ return job_id, artifact_ref
+
+
+async def _generate_retrieval_asset_url(
+ *,
+ row: dict[str, Any],
+ log_context: str,
+) -> str | None:
+ request = _resolve_asset_request(row)
+ if request is None:
+ return None
+
+ job_id, artifact_ref = request
+ try:
+ return get_result_storage().generate_artifact_url(
+ job_id=job_id,
+ artifact_ref=artifact_ref,
+ )
+ except Exception as exc:
+ logger.warning(f"Failed to generate {log_context} asset URL (ignored): {exc}")
+ return None
+
+
+async def enrich_rows_with_retrieval_asset_urls(
+ rows: list[dict[str, Any]],
+ *,
+ log_context: str,
+) -> list[dict[str, Any]]:
+ enriched_rows: list[dict[str, Any]] = []
+ for row in rows:
+ enriched = dict(row)
+ asset_url = await _generate_retrieval_asset_url(
+ row=row,
+ log_context=log_context,
+ )
+ if asset_url:
+ enriched["asset_url"] = asset_url
+ enriched_rows.append(enriched)
+ return enriched_rows
+
+
+async def build_retrieval_asset_url_map(
+ rows: list[dict[str, Any]],
+ *,
+ log_context: str,
+) -> dict[str, str]:
+ url_map: dict[str, str] = {}
+ for row in rows:
+ chunk_id = str(row.get("chunk_id") or "").strip()
+ if not chunk_id:
+ continue
+
+ asset_url = await _generate_retrieval_asset_url(
+ row=row,
+ log_context=log_context,
+ )
+ if asset_url:
+ url_map[chunk_id] = asset_url
+ return url_map
+
+
+def is_client_result_artifact_ref(asset_ref: str | None) -> bool:
+ return _normalize_artifact_ref(asset_ref) is not None
diff --git a/packages/shared-python/shared/services/retrieval/channels.py b/packages/shared-python/shared/services/retrieval/channels.py
index a597307ca..35966ed23 100644
--- a/packages/shared-python/shared/services/retrieval/channels.py
+++ b/packages/shared-python/shared/services/retrieval/channels.py
@@ -8,12 +8,14 @@
from typing import Any
-from loguru import logger
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.services.retrieval.graph_service import is_excluded_section
-from shared.utils.text_utils import tokenize_for_retrieval
+from shared.services.retrieval.lexical_ranker import (
+ rank_rows_by_bm25,
+ tokenize_query_for_ranker,
+)
+from shared.services.retrieval.section_filters import is_excluded_section
_SCOPED_CORPUS_CTE = """
@@ -193,56 +195,6 @@ async def content_channel(
)
-def _tokenize_query(query: str) -> list[str]:
- return tokenize_for_retrieval(query, dedupe=True)
-
-
-def _bm25_rerank(
- rows: list[dict[str, Any]],
- query_tokens: list[str],
- *,
- search_field: str,
-) -> list[dict[str, Any]]:
- """Rank matching rows with BM25 over pre-tokenized search text."""
- try:
- from rank_bm25 import BM25Okapi
- except ImportError:
- logger.warning("rank_bm25 not installed, skipping BM25 re-rank")
- ranked_rows: list[dict[str, Any]] = []
- query_token_set = set(query_tokens)
- for row in rows:
- tokens = [token for token in str(row.get(search_field) or "").split() if token]
- overlap = len(query_token_set.intersection(tokens))
- if overlap <= 0:
- continue
- row["score"] = float(overlap)
- ranked_rows.append(row)
- ranked_rows.sort(key=lambda r: r["score"], reverse=True)
- return ranked_rows
-
- corpus: list[list[str]] = []
- ranked_rows: list[dict[str, Any]] = []
- query_token_set = set(query_tokens)
- for row in rows:
- tokens = [token for token in str(row.get(search_field) or "").split() if token]
- if not tokens or not query_token_set.intersection(tokens):
- continue
- corpus.append(tokens)
- ranked_rows.append(row)
-
- if not corpus or not query_tokens:
- return []
-
- bm25 = BM25Okapi(corpus)
- scores = bm25.get_scores(query_tokens)
-
- for i, row in enumerate(ranked_rows):
- row["score"] = float(scores[i])
-
- ranked_rows.sort(key=lambda r: r["score"], reverse=True)
- return ranked_rows
-
-
async def _bm25_channel(
db: AsyncSession,
*,
@@ -260,7 +212,7 @@ async def _bm25_channel(
if search_field not in {"content_search_text", "path_search_text"}:
raise ValueError(f"Unsupported search_field: {search_field}")
- query_tokens = _tokenize_query(query)
+ query_tokens = tokenize_query_for_ranker(query)
if not query_tokens:
return []
@@ -287,7 +239,7 @@ async def _bm25_channel(
rows = [_row_to_dict(r) for r in result.all()]
rows = _filter_excluded_sections(rows, exclude_sections)
- ranked_rows = _bm25_rerank(rows, query_tokens, search_field=search_field)
+ ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field)
return ranked_rows[:top_k]
@@ -312,7 +264,7 @@ async def term_channel(
Note: top_k is already effective_recall_k from app_service.
"""
query_lower = query.lower().strip()
- query_tokens = _tokenize_query(query)
+ query_tokens = tokenize_query_for_ranker(query)
if not query_lower or not query_tokens:
return []
diff --git a/packages/shared-python/shared/services/retrieval/connected_hydration.py b/packages/shared-python/shared/services/retrieval/connected_hydration.py
new file mode 100644
index 000000000..c5c2a181b
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/connected_hydration.py
@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import and_, or_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, DocumentSection
+from shared.models.database.job_result import JobResult
+from shared.services.retrieval.row_utils import (
+ filter_excluded_rows,
+ iter_connected_target_ids,
+ normalize_chunk_type,
+)
+
+
+async def hydrate_connected_target_rows(
+ *,
+ db: AsyncSession | None,
+ rows: list[dict[str, Any]],
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+) -> list[dict[str, Any]]:
+ if db is None:
+ return []
+
+ existing_chunk_ids = {
+ str(row.get('chunk_id') or '').strip()
+ for row in rows
+ if row.get('chunk_id')
+ }
+ target_ids_by_revision: dict[tuple[str, str], set[str]] = {}
+ for row in rows:
+ if normalize_chunk_type(row.get('chunk_type')) != 'text':
+ continue
+ document_id = str(row.get('document_id') or '').strip()
+ job_result_id = str(row.get('job_result_id') or '').strip()
+ if not document_id or not job_result_id:
+ continue
+ for target_id in iter_connected_target_ids(row):
+ if target_id in existing_chunk_ids:
+ continue
+ target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(
+ target_id
+ )
+
+ if not target_ids_by_revision:
+ return []
+
+ revision_filters = [
+ and_(
+ DocumentChunk.document_id == document_id,
+ DocumentChunk.job_result_id == job_result_id,
+ DocumentChunk.chunk_id.in_(sorted(target_ids)),
+ )
+ for (document_id, job_result_id), target_ids in target_ids_by_revision.items()
+ if target_ids
+ ]
+ if not revision_filters:
+ return []
+
+ stmt = (
+ select(Document, DocumentChunk, DocumentSection, JobResult)
+ .join(DocumentChunk, DocumentChunk.document_id == Document.document_id)
+ .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
+ .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
+ .where(or_(*revision_filters))
+ .order_by(DocumentChunk.sort_order)
+ )
+ result = await db.execute(stmt)
+
+ hydrated_rows: list[dict[str, Any]] = []
+ for document, chunk, section, job_result in result.all():
+ section_path = section.section_path if section else None
+ hydrated_rows.append(
+ {
+ 'document_id': document.document_id,
+ 'chunk_id': chunk.chunk_id,
+ 'section_id': chunk.section_id,
+ 'section_path': section_path,
+ 'source_file_name': document.source_file_name,
+ 'chunk_type': chunk.chunk_type,
+ 'content': chunk.content,
+ 'score': 0.0,
+ 'file_path': chunk.file_path,
+ 'chunk_metadata': chunk.chunk_metadata or {},
+ 'job_result_id': chunk.job_result_id,
+ 'job_id': job_result.job_id if job_result else None,
+ 'sort_order': chunk.sort_order,
+ }
+ )
+
+ return filter_excluded_rows(
+ hydrated_rows,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ )
diff --git a/packages/shared-python/shared/services/retrieval/execution_plan.py b/packages/shared-python/shared/services/retrieval/execution_plan.py
new file mode 100644
index 000000000..71669f3ed
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/execution_plan.py
@@ -0,0 +1,320 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.cache_service import (
+ get_cached_retrieval_query_result,
+ set_cached_retrieval_query_result,
+)
+from shared.services.retrieval.execution_routes import run_retrieval_route
+from shared.services.retrieval.route_types import RetrievalRouteContext
+from shared.services.retrieval.hit_stats_recorder import (
+ schedule_retrieval_hit_stats_update,
+)
+from shared.services.retrieval.response_projection import (
+ project_public_retrieval_response,
+)
+from shared.services.retrieval.settings import (
+ INTERNAL_RECALL_K_MULTIPLIER,
+ resolve_allowed_chunk_types,
+)
+
+
+async def run_retrieval_query(
+ *,
+ db: AsyncSession,
+ user_id: str,
+ namespace: str,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int = 1,
+ signal_paths: list[str] | None = None,
+ filter_mode: str = "delete",
+ channels: list[str] | None = None,
+ channel_weights: dict[str, float] | None = None,
+ rerank: bool = False,
+ threshold: float = 0.0,
+ internal_recall_k: int | None = None,
+ use_agentic: bool | None = None,
+) -> dict[str, Any]:
+ """Run retrieval through the plan module."""
+ return await RetrievalExecutionPlan(
+ {
+ "db": db,
+ "user_id": user_id,
+ "namespace": namespace,
+ "query": query,
+ "top_k": top_k,
+ "exclude_document_ids": exclude_document_ids,
+ "exclude_sections": exclude_sections,
+ "data_type": data_type,
+ "signal_paths": signal_paths,
+ "filter_mode": filter_mode,
+ "channels": channels,
+ "channel_weights": channel_weights,
+ "rerank": rerank,
+ "threshold": threshold,
+ "internal_recall_k": internal_recall_k,
+ "use_agentic": use_agentic,
+ }
+ ).execute()
+
+
+class RetrievalExecutionPlan:
+ def __init__(self, request: dict[str, Any]) -> None:
+ self.request = request
+
+ async def execute(self) -> dict[str, Any]:
+ db: AsyncSession = self.request["db"]
+ user_id: str = self.request["user_id"]
+ namespace: str = self.request["namespace"]
+ query: str = str(self.request["query"]).strip()
+ top_k: int = self.request["top_k"]
+ exclude_document_ids: list[str] = self.request["exclude_document_ids"]
+ exclude_sections: list[dict[str, str]] = self.request["exclude_sections"]
+ data_type: int = self.request["data_type"]
+ signal_paths: list[str] | None = self.request["signal_paths"]
+ filter_mode: str = self.request["filter_mode"]
+ channels: list[str] | None = self.request["channels"]
+ channel_weights: dict[str, float] | None = self.request["channel_weights"]
+ rerank: bool = self.request["rerank"]
+ threshold: float = self.request["threshold"]
+ internal_recall_k: int | None = self.request["internal_recall_k"]
+ use_agentic: bool | None = self.request["use_agentic"]
+
+ start_time = time.monotonic()
+ _log_retrieval_start(
+ query=query,
+ user_id=user_id,
+ namespace=namespace,
+ top_k=top_k,
+ data_type=data_type,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ )
+
+ if not query:
+ logger.info(" ⛔ Empty query filtered, skipping retrieval pipeline")
+ return {
+ "namespace": namespace,
+ "query": query,
+ "router_used": "empty_query_filtered",
+ "results": [],
+ }
+
+ allowed_chunk_types: set[str] | None = resolve_allowed_chunk_types(data_type)
+ effective_recall_k = (
+ internal_recall_k
+ if internal_recall_k is not None
+ else top_k * INTERNAL_RECALL_K_MULTIPLIER
+ )
+ logger.info(
+ f" allowed_chunk_types={allowed_chunk_types} "
+ f"effective_recall_k={effective_recall_k} "
+ f"signal_paths={signal_paths} filter_mode={filter_mode} "
+ f"rerank={rerank} threshold={threshold}"
+ )
+
+ cache_extra = {
+ "data_type": data_type,
+ "signal_paths": signal_paths,
+ "filter_mode": filter_mode,
+ "channels": channels,
+ "channel_weights": channel_weights,
+ "rerank": rerank,
+ "threshold": threshold,
+ "internal_recall_k": internal_recall_k,
+ "decomposition_enabled": True,
+ }
+ cache_version, cached_response = await _read_cached_response(
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ cache_extra=cache_extra,
+ )
+ if cached_response is not None:
+ return cached_response
+
+ logger.debug(f" 📦 Cache miss (version={cache_version}), running full pipeline")
+
+ route_context = RetrievalRouteContext(
+ db=db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ allowed_chunk_types=allowed_chunk_types,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ threshold=threshold,
+ effective_recall_k=effective_recall_k,
+ use_agentic=use_agentic,
+ )
+ outcome = await run_retrieval_route(route_context)
+
+ if cache_version is not None:
+ await _write_cached_response(
+ user_id=user_id,
+ namespace=namespace,
+ version=cache_version,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ response=outcome.response,
+ cache_extra=cache_extra,
+ )
+
+ _schedule_hit_stats_update(
+ user_id=user_id,
+ namespace=namespace,
+ results=outcome.hit_stats_results,
+ )
+ _log_retrieval_complete(
+ outcome=outcome.response,
+ label=outcome.completion_label,
+ count=outcome.completion_count,
+ detail=outcome.completion_detail,
+ elapsed_ms=round((time.monotonic() - start_time) * 1000),
+ )
+ return await project_public_retrieval_response(outcome.response)
+
+
+async def _read_cached_response(
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ cache_extra: dict[str, Any],
+) -> tuple[int | None, dict[str, Any] | None]:
+ cache_version: int | None = None
+ try:
+ cache_version, cached = await get_cached_retrieval_query_result(
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ **cache_extra,
+ )
+ if cached:
+ logger.info(f"retrieval: cache_hit=True version={cache_version}")
+ _schedule_hit_stats_update(
+ user_id=user_id,
+ namespace=namespace,
+ results=cached.get("results", []),
+ )
+ return cache_version, await project_public_retrieval_response(cached)
+ except Exception as exc:
+ logger.warning(f"Failed to read retrieval cache (ignored): {exc}")
+ return cache_version, None
+
+
+async def _write_cached_response(
+ *,
+ user_id: str,
+ namespace: str,
+ version: int,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ response: dict[str, Any],
+ cache_extra: dict[str, Any],
+) -> None:
+ try:
+ await set_cached_retrieval_query_result(
+ user_id=user_id,
+ namespace=namespace,
+ version=version,
+ query=query,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ response=response,
+ **cache_extra,
+ )
+ except Exception as exc:
+ logger.warning(f"Failed to write retrieval cache (ignored): {exc}")
+
+
+def _schedule_hit_stats_update(
+ *,
+ user_id: str,
+ namespace: str,
+ results: list[dict[str, Any]],
+) -> None:
+ try:
+ schedule_retrieval_hit_stats_update(
+ user_id=user_id,
+ namespace=namespace,
+ results=results,
+ )
+ except Exception as exc:
+ logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {exc}")
+
+
+def _log_retrieval_start(
+ *,
+ query: str,
+ user_id: str,
+ namespace: str,
+ top_k: int,
+ data_type: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+) -> None:
+ logger.info("\n" + "█" * 70)
+ logger.info(" 🚀 RETRIEVAL PIPELINE START")
+ logger.info(f' query="{query}"')
+ logger.info(
+ f" user={user_id} ns={namespace} top_k={top_k} data_type={data_type}"
+ )
+ logger.info(
+ f" exclude_docs={exclude_document_ids} "
+ f"exclude_secs={len(exclude_sections)}"
+ )
+ logger.info("█" * 70)
+
+
+def _log_retrieval_complete(
+ *,
+ outcome: dict[str, Any],
+ label: str,
+ count: int,
+ detail: str,
+ elapsed_ms: int,
+) -> None:
+ logger.info(f'\n{"█" * 70}')
+ logger.info(f" ✅ {label} COMPLETE: {count} {detail} | {elapsed_ms}ms")
+ results = outcome.get("results", [])
+ if isinstance(results, list):
+ for index, result in enumerate(results[:10]):
+ source = result.get("source", {})
+ logger.info(
+ f" [{index + 1}] type={result.get('chunk_type', '?')} "
+ f"score={result.get('score', 0):.4f}"
+ f" path={source.get('section_path', '')}"
+ f" file={source.get('source_file_name', '')}"
+ )
+ if len(results) > 10:
+ logger.info(f" ... and {len(results) - 10} more")
+ logger.info(f'{"█" * 70}')
diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py
new file mode 100644
index 000000000..ee5516159
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/execution_routes.py
@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import os
+
+from loguru import logger
+
+from shared.services.retrieval.legacy_route import run_legacy_retrieval_route
+from shared.services.retrieval.reference_hydration import hydrate_referenced_chunk_rows
+from shared.services.retrieval.result_assembly import assemble_retrieval_results
+from shared.services.retrieval.response_projection import (
+ attach_citation,
+ enrich_referenced_chunks_with_asset_urls,
+)
+from shared.services.retrieval.route_types import (
+ RetrievalRouteContext,
+ RetrievalRouteOutcome,
+)
+from shared.services.retrieval.scoped_corpus import (
+ count_scoped_chunks,
+ load_all_scoped_chunks,
+)
+
+
+async def run_retrieval_route(
+ context: RetrievalRouteContext,
+) -> RetrievalRouteOutcome:
+ small_kb_outcome = await _try_run_small_kb_route(context)
+ if small_kb_outcome is not None:
+ return small_kb_outcome
+
+ if _should_use_agentic_route(context.use_agentic):
+ return await _run_agentic_route(context)
+
+ return await run_legacy_retrieval_route(context)
+
+
+async def _try_run_small_kb_route(
+ context: RetrievalRouteContext,
+) -> RetrievalRouteOutcome | None:
+ try:
+ total_chunk_count = await count_scoped_chunks(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ exclude_document_ids=context.exclude_document_ids,
+ allowed_chunk_types=context.allowed_chunk_types,
+ )
+ except Exception as exc:
+ logger.warning(
+ f"Failed to count scoped chunks, skipping small KB optimization: {exc}"
+ )
+ total_chunk_count = context.top_k + 1
+
+ logger.info(f"\n Total chunks in scope: {total_chunk_count}")
+ if total_chunk_count > context.top_k:
+ return None
+
+ logger.info(
+ f" Small KB optimization: {total_chunk_count} chunks "
+ f"<= top_k={context.top_k}, returning all"
+ )
+ all_rows = await load_all_scoped_chunks(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ signal_paths=context.signal_paths or [],
+ filter_mode=context.filter_mode,
+ )
+ logger.info(
+ f" small_kb load: loaded={len(all_rows)} rows after signal/exclude filters"
+ )
+ assembled_rows = await assemble_retrieval_results(
+ db=context.db,
+ rows=all_rows,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ )
+ results = [attach_citation(row) for row in assembled_rows]
+ response = {
+ "namespace": context.namespace,
+ "query": context.query,
+ "router_used": "small_kb_all",
+ "results": results,
+ }
+ return RetrievalRouteOutcome(
+ response=response,
+ hit_stats_results=results,
+ completion_label="Small KB",
+ completion_count=len(results),
+ completion_detail="results",
+ )
+
+
+def _should_use_agentic_route(use_agentic: bool | None) -> bool:
+ if use_agentic is not None:
+ return use_agentic
+ return os.environ.get("RETRIEVAL_AGENTIC_ENABLED", "true") == "true"
+
+
+async def _run_agentic_route(
+ context: RetrievalRouteContext,
+) -> RetrievalRouteOutcome:
+ from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator
+
+ workflow = WorkflowOrchestrator()
+ workflow_result = await workflow.run(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ query=context.query,
+ top_k=context.top_k,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ data_type=context.data_type,
+ signal_paths=context.signal_paths,
+ filter_mode=context.filter_mode,
+ channels=context.channels,
+ channel_weights=context.channel_weights,
+ )
+
+ enriched_refs = await enrich_referenced_chunks_with_asset_urls(
+ workflow_result.referenced_chunks,
+ )
+
+ workflow_result_rows = await hydrate_referenced_chunk_rows(
+ db=context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ refs=enriched_refs,
+ )
+ scoped_reference_keys = {
+ (
+ str(row.get("document_id") or "").strip(),
+ str(row.get("chunk_id") or "").strip(),
+ )
+ for row in workflow_result_rows
+ }
+ enriched_refs = [
+ ref
+ for ref in enriched_refs
+ if (
+ str(ref.get("document_id") or "").strip(),
+ str(ref.get("chunk_id") or "").strip(),
+ )
+ in scoped_reference_keys
+ ]
+ assembled_workflow_rows = await assemble_retrieval_results(
+ db=context.db,
+ rows=workflow_result_rows,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ )
+ response = workflow_result.to_api_response()
+ response["referenced_chunks"] = enriched_refs
+ response["results"] = [attach_citation(row) for row in assembled_workflow_rows]
+
+ completion_detail = (
+ f"chunks | answer={len(workflow_result.answer_text)} chars | "
+ f"router={workflow_result.router_used}"
+ )
+ return RetrievalRouteOutcome(
+ response=response,
+ hit_stats_results=enriched_refs,
+ completion_label="AGENTIC RETRIEVAL",
+ completion_count=len(enriched_refs),
+ completion_detail=completion_detail,
+ )
diff --git a/packages/shared-python/shared/services/retrieval/graph_keywords.py b/packages/shared-python/shared/services/retrieval/graph_keywords.py
new file mode 100644
index 000000000..5769bd355
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/graph_keywords.py
@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import math
+import re
+from typing import Any
+
+MIN_KEYWORD_OVERLAP = 3
+KEYWORD_SCORE_WEIGHT = 1.0
+MIN_SCORE_THRESHOLD = 0.8
+
+
+def normalize_keyword(keyword: str) -> str:
+ """Normalize a keyword: lowercase, strip, collapse spaces."""
+ keyword = keyword.lower().strip()
+ return re.sub(r'\s+', ' ', keyword)
+
+
+def extract_keywords_from_chunk_metadata(meta: dict) -> list[str]:
+ """Extract keywords from chunk metadata."""
+ if not isinstance(meta, dict):
+ return []
+
+ keywords = meta.get('keywords', [])
+ if isinstance(keywords, list) and keywords:
+ return [str(keyword) for keyword in keywords if keyword]
+
+ tokens = meta.get('tokens', [])
+ if isinstance(tokens, list) and tokens:
+ return [str(token) for token in tokens if token and len(str(token)) > 1]
+
+ return []
+
+
+def compute_tfidf_keywords(
+ chunk_metadata_list: list[dict[str, Any]],
+ top_k: int = 10,
+) -> list[str]:
+ """Compute TF-IDF keywords from chunk metadata."""
+ df_count: dict[str, int] = {}
+ tf_count: dict[str, int] = {}
+ total = len(chunk_metadata_list) or 1
+ for meta in chunk_metadata_list:
+ keywords = extract_keywords_from_chunk_metadata(meta)
+ seen: set[str] = set()
+ for keyword in keywords:
+ if len(str(keyword)) <= 1 or re.match(r'^\d+[.,%]*$', str(keyword)):
+ continue
+ normalized = normalize_keyword(str(keyword))
+ if not normalized:
+ continue
+ tf_count[normalized] = tf_count.get(normalized, 0) + 1
+ if normalized not in seen:
+ df_count[normalized] = df_count.get(normalized, 0) + 1
+ seen.add(normalized)
+ scored = [
+ (term, freq * (math.log(total / (df_count.get(term, 1))) + 1))
+ for term, freq in tf_count.items()
+ ]
+ scored.sort(key=lambda item: item[1], reverse=True)
+ return [term for term, _ in scored[:top_k]]
+
+
+def compute_keyword_score(
+ shared_keywords: set[str],
+ keywords_a: set[str],
+ keywords_b: set[str],
+ weight: float = 1.0,
+) -> float:
+ """Character-length-weighted keyword overlap score."""
+ weighted_a = sum(len(keyword) for keyword in keywords_a)
+ weighted_b = sum(len(keyword) for keyword in keywords_b)
+ denominator = min(weighted_a, weighted_b)
+ if denominator == 0:
+ return 0.0
+ weighted_shared = sum(len(keyword) for keyword in shared_keywords)
+ return weight * weighted_shared / denominator
+
+
+def get_normalized_keyword_set(chunk_metadata_list: list[dict[str, Any]]) -> set[str]:
+ """Collect all normalized keywords from chunk metadata for a document."""
+ result: set[str] = set()
+ for meta in chunk_metadata_list:
+ for keyword in extract_keywords_from_chunk_metadata(meta):
+ normalized = normalize_keyword(str(keyword))
+ if normalized and len(normalized) > 1 and not re.match(
+ r'^\d+[.,%]*$', normalized
+ ):
+ result.add(normalized)
+ return result
+
+
+def extract_document_top_summary(chunk_metadata_list: list[dict[str, Any]]) -> str:
+ """Read the parser-injected top summary from chunk metadata."""
+ for meta in chunk_metadata_list:
+ if not isinstance(meta, dict):
+ continue
+ summary = str(meta.get('document_top_summary') or '').strip()
+ if summary:
+ return summary
+ return ''
diff --git a/packages/shared-python/shared/services/retrieval/graph_query_service.py b/packages/shared-python/shared/services/retrieval/graph_query_service.py
new file mode 100644
index 000000000..de912a0ca
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/graph_query_service.py
@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, DocumentSection
+from shared.models.database.job_result import JobResult
+from shared.services.retrieval.section_filters import is_excluded_section
+
+_SECTION_EXCLUSION_PAGE_MULTIPLIER = 2
+
+
+def _build_lexical_match_predicate(query: str):
+ like = f'%{query}%'
+ return (
+ DocumentChunk.content_lexical_text.ilike(like)
+ | DocumentChunk.path_lexical_text.ilike(like)
+ )
+
+
+class GraphQueryService:
+ """Read-side graph routing before canonical chunk hydration."""
+
+ async def find_entry_documents(
+ self,
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ exclude_document_ids: Iterable[str] = (),
+ exclude_sections: Iterable[dict[str, str]] = (),
+ ) -> list[str]:
+ query_lc = query.lower().strip()
+ excluded_document_ids = set(exclude_document_ids)
+
+ if query_lc:
+ section_matches = await self._find_documents_by_section(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query_lc,
+ exclude_document_ids=excluded_document_ids,
+ exclude_sections=exclude_sections,
+ )
+ if section_matches:
+ return section_matches
+
+ return await self._find_documents_by_content(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query_lc,
+ exclude_document_ids=excluded_document_ids,
+ )
+
+ async def _find_documents_by_section(
+ self,
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ exclude_document_ids: set[str],
+ exclude_sections: Iterable[dict[str, str]],
+ ) -> list[str]:
+ like = f'%{query}%'
+ stmt = (
+ select(DocumentSection.document_id)
+ .join(
+ Document,
+ (Document.document_id == DocumentSection.document_id)
+ & (Document.current_job_result_id == DocumentSection.job_result_id),
+ )
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .where(
+ DocumentSection.section_title.ilike(like)
+ | DocumentSection.section_path.ilike(like)
+ )
+ .distinct()
+ )
+ if exclude_document_ids:
+ stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
+ for item in exclude_sections or ():
+ if not isinstance(item, dict):
+ continue
+ excluded_document_id = str(item.get('document_id') or '').strip()
+ excluded_path = str(item.get('section_path') or '').strip()
+ if excluded_document_id and excluded_path:
+ stmt = stmt.where(
+ ~(
+ (DocumentSection.document_id == excluded_document_id)
+ & (
+ (DocumentSection.section_path == excluded_path)
+ | DocumentSection.section_path.like(f'{excluded_path} / %')
+ )
+ )
+ )
+
+ result = await db.execute(stmt)
+ return [document_id for (document_id,) in result.all()]
+
+ async def _find_documents_by_content(
+ self,
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ exclude_document_ids: set[str],
+ ) -> list[str]:
+ like = f'%{query}%'
+ stmt = (
+ select(Document.document_id)
+ .join(
+ DocumentChunk,
+ (DocumentChunk.document_id == Document.document_id)
+ & (DocumentChunk.job_result_id == Document.current_job_result_id),
+ )
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .where(DocumentChunk.content_lexical_text.ilike(like))
+ )
+ if exclude_document_ids:
+ stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
+ result = await db.execute(stmt)
+ seen: list[str] = []
+ for (document_id,) in result.all():
+ if document_id and document_id not in seen:
+ seen.append(document_id)
+ return seen
+
+ async def collect_candidate_chunks(
+ self,
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ entry_document_ids: Sequence[str],
+ query: str,
+ top_k: int,
+ exclude_sections: Iterable[dict[str, str]] = (),
+ ) -> list[dict[str, Any]]:
+ if not entry_document_ids:
+ return []
+ page_size = top_k
+ if exclude_sections:
+ page_size = max(top_k, top_k * _SECTION_EXCLUSION_PAGE_MULTIPLIER)
+ base_stmt = (
+ select(Document, DocumentChunk, DocumentSection, JobResult)
+ .join(
+ DocumentChunk,
+ (DocumentChunk.document_id == Document.document_id)
+ & (DocumentChunk.job_result_id == Document.current_job_result_id),
+ )
+ .outerjoin(
+ DocumentSection,
+ DocumentSection.section_id == DocumentChunk.section_id,
+ )
+ .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .where(Document.document_id.in_(list(entry_document_ids)))
+ .where(_build_lexical_match_predicate(query))
+ .order_by(DocumentChunk.sort_order)
+ )
+ rows: list[dict[str, Any]] = []
+ offset = 0
+ while len(rows) < top_k:
+ result = await db.execute(base_stmt.limit(page_size).offset(offset))
+ result_rows = result.all()
+ if not result_rows:
+ break
+ for document, chunk, section, job_result in result_rows:
+ section_path = section.section_path if section else None
+ if is_excluded_section(
+ document_id=document.document_id,
+ section_path=section_path,
+ exclude_sections=exclude_sections,
+ ):
+ continue
+ rows.append({
+ 'document_id': document.document_id,
+ 'chunk_id': chunk.chunk_id,
+ 'section_id': chunk.section_id,
+ 'section_path': section_path,
+ 'source_file_name': document.source_file_name,
+ 'chunk_type': chunk.chunk_type,
+ 'content': chunk.content,
+ 'score': 2.0,
+ 'file_path': chunk.file_path,
+ 'chunk_metadata': chunk.chunk_metadata or {},
+ 'job_result_id': chunk.job_result_id,
+ 'job_id': job_result.job_id if job_result else None,
+ })
+ if len(rows) >= top_k:
+ break
+ if len(result_rows) < page_size:
+ break
+ offset += page_size
+ return rows
diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py
index 8244a8a72..0d96a76e3 100644
--- a/packages/shared-python/shared/services/retrieval/graph_service.py
+++ b/packages/shared-python/shared/services/retrieval/graph_service.py
@@ -1,160 +1,31 @@
from __future__ import annotations
import logging
-import math
-import re
from collections import defaultdict
from dataclasses import dataclass
-from typing import Any, Iterable, Sequence
from sqlalchemy import delete, or_, select
-from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
-from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphEdge, GraphNode
-from shared.models.database.job_result import JobResult
+from shared.models.database.document import (
+ Document,
+ DocumentChunk,
+ GraphEdge,
+ GraphNode,
+)
+from shared.services.retrieval.graph_keywords import (
+ KEYWORD_SCORE_WEIGHT,
+ MIN_KEYWORD_OVERLAP,
+ MIN_SCORE_THRESHOLD,
+ compute_keyword_score,
+ compute_tfidf_keywords,
+ extract_document_top_summary,
+ get_normalized_keyword_set,
+ normalize_keyword,
+)
logger = logging.getLogger(__name__)
-_SECTION_EXCLUSION_PAGE_MULTIPLIER = 2
-
-# ── Keyword overlap config (aligned with connect_builder DEFAULT_CONFIG) ──
-_MIN_KEYWORD_OVERLAP = 3
-_KEYWORD_SCORE_WEIGHT = 1.0
-_MIN_SCORE_THRESHOLD = 0.8
-_CROSS_FILE_ONLY = True
-_MAX_CONTENT_OVERLAP = 0.8
-
-
-def _build_lexical_match_predicate(query: str):
- like = f'%{query}%'
- return (
- DocumentChunk.content_lexical_text.ilike(like)
- | DocumentChunk.path_lexical_text.ilike(like)
- )
-
-
-def is_excluded_section(
- *,
- document_id: str | None,
- section_path: str | None,
- exclude_sections: Iterable[dict[str, str]],
-) -> bool:
- document_id = str(document_id or '').strip()
- section_path = str(section_path or '').strip()
- if not document_id or not section_path:
- return False
- for item in exclude_sections:
- if not isinstance(item, dict):
- continue
- exc_doc = str(item.get('document_id') or '').strip()
- exc_path = str(item.get('section_path') or '').strip()
- if document_id == exc_doc and (section_path == exc_path or section_path.startswith(exc_path + ' / ')):
- return True
- return False
-
-
-# ── Keyword extraction & scoring (aligned with connect_builder/builder.py) ──
-
-def _normalize_keyword(keyword: str) -> str:
- """Normalize a keyword: lowercase, strip, collapse spaces."""
- kw = keyword.lower().strip()
- return re.sub(r'\s+', ' ', kw)
-
-
-def _extract_keywords_from_chunk_metadata(meta: dict) -> list[str]:
- """Extract keywords from chunk metadata, same logic as builder._get_keywords."""
- if not isinstance(meta, dict):
- return []
- # Try metadata.keywords
- kws = meta.get('keywords', [])
- if isinstance(kws, list) and kws:
- return [str(k) for k in kws if k]
- # Fallback: tokens
- tokens = meta.get('tokens', [])
- if isinstance(tokens, list) and tokens:
- return [str(t) for t in tokens if t and len(str(t)) > 1]
- return []
-
-
-def _compute_tfidf_keywords(
- chunk_metadata_list: list[dict[str, Any]],
- top_k: int = 10,
-) -> list[str]:
- """Compute TF-IDF keywords from chunk metadata, aligned with graph_builder."""
- df_count: dict[str, int] = {}
- tf_count: dict[str, int] = {}
- total = len(chunk_metadata_list) or 1
- for meta in chunk_metadata_list:
- kws = _extract_keywords_from_chunk_metadata(meta)
- seen: set[str] = set()
- for k in kws:
- if len(str(k)) <= 1 or re.match(r'^\d+[.,%]*$', str(k)):
- continue
- lower = _normalize_keyword(str(k))
- if not lower:
- continue
- tf_count[lower] = tf_count.get(lower, 0) + 1
- if lower not in seen:
- df_count[lower] = df_count.get(lower, 0) + 1
- seen.add(lower)
- scored = [
- (term, freq * (math.log(total / (df_count.get(term, 1))) + 1))
- for term, freq in tf_count.items()
- ]
- scored.sort(key=lambda x: x[1], reverse=True)
- return [s[0] for s in scored[:top_k]]
-
-
-def _compute_keyword_score(
- shared_kws: set[str],
- kws_a: set[str],
- kws_b: set[str],
- weight: float = 1.0,
-) -> float:
- """Character-length-weighted keyword overlap score (aligned with builder.py).
-
- Longer tokens contribute more: '施工现场'(4) has 2x weight of '交底'(2).
- Formula: score = weight * sum(len(kw) for shared) / min(sum(len) for A, sum(len) for B)
- """
- weighted_a = sum(len(k) for k in kws_a)
- weighted_b = sum(len(k) for k in kws_b)
- denominator = min(weighted_a, weighted_b)
- if denominator == 0:
- return 0.0
- weighted_shared = sum(len(k) for k in shared_kws)
- return weight * weighted_shared / denominator
-
-
-def _get_normalized_keyword_set(chunk_metadata_list: list[dict[str, Any]]) -> set[str]:
- """Collect all normalized keywords from chunk metadata for a document."""
- result: set[str] = set()
- for meta in chunk_metadata_list:
- for k in _extract_keywords_from_chunk_metadata(meta):
- normalized = _normalize_keyword(str(k))
- if normalized and len(normalized) > 1 and not re.match(r'^\d+[.,%]*$', normalized):
- result.add(normalized)
- return result
-
-
-def _extract_document_top_summary(
- chunk_metadata_list: list[dict[str, Any]],
- section_titles: Sequence[str],
-) -> str:
- """Extract document_top_summary from chunk metadata.
-
- The summary is injected by kb_tasks.py via load_nav_top_summary()
- at parse time, so it should always be present. If missing, return empty
- string rather than fabricating a low-quality fallback.
- """
- for meta in chunk_metadata_list:
- if not isinstance(meta, dict):
- continue
- summary = str(meta.get('document_top_summary') or '').strip()
- if summary:
- return summary
- return ''
-
@dataclass
class GraphScope:
@@ -169,10 +40,17 @@ class DocumentGraphService:
- Only document-level nodes (no section nodes)
- Document nodes carry rich metadata: top_keywords, chunks_count, types, top_summary
- Edges are keyword-overlap-based cross-document connections with meaningful scores
- - Edge scoring uses connect_builder DEFAULT_CONFIG thresholds
"""
- def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, document_id: str, job_result_id: str) -> None:
+ def publish_document_graph(
+ self,
+ db: Session,
+ *,
+ user_id: str,
+ namespace: str,
+ document_id: str,
+ job_result_id: str,
+ ) -> None:
document = db.execute(
select(Document).where(Document.document_id == document_id)
).scalar_one_or_none()
@@ -190,29 +68,22 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d
chunk_metadata_list = [row[1] or {} for row in chunk_meta_rows]
# Compute document-level metadata (aligned with KB knowledge_graph.json files dict)
- top_keywords = _compute_tfidf_keywords(chunk_metadata_list)
- new_doc_kws = _get_normalized_keyword_set(chunk_metadata_list)
+ top_keywords = compute_tfidf_keywords(chunk_metadata_list)
+ new_doc_kws = get_normalized_keyword_set(chunk_metadata_list)
types_breakdown: dict[str, int] = defaultdict(int)
for chunk_type, _ in chunk_meta_rows:
types_breakdown[chunk_type or 'text'] += 1
chunks_count = len(chunk_meta_rows)
- sections = [
- section_title
- for section_title in db.execute(
- select(DocumentSection.section_title)
- .where(DocumentSection.document_id == document_id)
- .where(DocumentSection.job_result_id == job_result_id)
- .where(DocumentSection.section_level <= 2)
- .order_by(DocumentSection.sort_order)
- ).scalars()
- if section_title is not None
- ]
- top_summary = _extract_document_top_summary(chunk_metadata_list, sections)
+ top_summary = extract_document_top_summary(chunk_metadata_list)
# ── Clean up old graph data for this document ──
- self.remove_document_graph(db, scope=GraphScope(user_id=user_id, namespace=namespace), document_id=document_id)
+ self.remove_document_graph(
+ db,
+ scope=GraphScope(user_id=user_id, namespace=namespace),
+ document_id=document_id,
+ )
# ── Create document-level node (no section nodes — aligned with KB KG) ──
document_node_id = f"doc:{document_id}"
@@ -237,9 +108,8 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d
)
db.flush()
- # ── Keyword-overlap-based cross-document edges (aligned with KB edges) ──
- # Only create edges where keyword overlap score >= threshold,
- # matching connect_builder DEFAULT_CONFIG parameters.
+ # ── Keyword-overlap-based cross-document edges ──
+ # Only create edges where keyword overlap score >= threshold.
other_doc_nodes = list(
db.execute(
select(GraphNode)
@@ -257,7 +127,7 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d
# Build normalized keyword sets for comparison
peer_kws: set[str] = set()
for k in peer_keywords:
- normalized = _normalize_keyword(str(k))
+ normalized = normalize_keyword(str(k))
if normalized:
peer_kws.add(normalized)
@@ -266,17 +136,17 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d
# Find shared keywords
shared_kws = new_doc_kws & peer_kws
- if len(shared_kws) < _MIN_KEYWORD_OVERLAP:
+ if len(shared_kws) < MIN_KEYWORD_OVERLAP:
continue
# Compute character-length-weighted score
- score = _compute_keyword_score(
- shared_kws=shared_kws,
- kws_a=new_doc_kws,
- kws_b=peer_kws,
- weight=_KEYWORD_SCORE_WEIGHT,
+ score = compute_keyword_score(
+ shared_keywords=shared_kws,
+ keywords_a=new_doc_kws,
+ keywords_b=peer_kws,
+ weight=KEYWORD_SCORE_WEIGHT,
)
- if score < _MIN_SCORE_THRESHOLD:
+ if score < MIN_SCORE_THRESHOLD:
continue
# Create edge with meaningful weight and metadata
@@ -307,7 +177,9 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d
f"keywords={len(top_keywords)} chunks={chunks_count}"
)
- def remove_document_graph(self, db: Session, *, scope: GraphScope | None, document_id: str) -> None:
+ def remove_document_graph(
+ self, db: Session, *, scope: GraphScope | None, document_id: str
+ ) -> None:
document_node_id = f"doc:{document_id}"
edge_delete = delete(GraphEdge).where(
or_(
@@ -329,152 +201,3 @@ def remove_document_graph(self, db: Session, *, scope: GraphScope | None, docume
db.execute(edge_delete)
db.execute(node_delete)
db.flush()
-
-
-class GraphQueryService:
- """Read-side graph service for document routing before canonical chunk hydration."""
-
- async def find_entry_documents(
- self,
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- query: str,
- exclude_document_ids: Iterable[str] = (),
- exclude_sections: Iterable[dict[str, str]] = (),
- ) -> list[str]:
- query_lc = query.lower().strip()
- exclude_document_ids = set(exclude_document_ids)
-
- if query_lc:
- like = f'%{query_lc}%'
- stmt = (
- select(DocumentSection.document_id)
- .join(Document, (Document.document_id == DocumentSection.document_id) & (Document.current_job_result_id == DocumentSection.job_result_id))
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(
- DocumentSection.section_title.ilike(like)
- | DocumentSection.section_path.ilike(like)
- )
- .distinct()
- )
- if exclude_document_ids:
- stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
- for exc in (exclude_sections or ()):
- if not isinstance(exc, dict):
- continue
- exc_doc = str(exc.get('document_id') or '').strip()
- exc_path = str(exc.get('section_path') or '').strip()
- if exc_doc and exc_path:
- stmt = stmt.where(
- ~((DocumentSection.document_id == exc_doc) & (
- (DocumentSection.section_path == exc_path) |
- DocumentSection.section_path.like(f'{exc_path} / %')
- ))
- )
- result = await db.execute(stmt)
- seen = [row[0] for row in result.all()]
- if seen:
- return seen
-
- return await self._find_documents_by_content(
- db,
- user_id=user_id,
- namespace=namespace,
- query=query_lc,
- exclude_document_ids=exclude_document_ids,
- )
-
- async def _find_documents_by_content(
- self,
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- query: str,
- exclude_document_ids: set[str],
- ) -> list[str]:
- like = f'%{query}%'
- stmt = (
- select(Document.document_id)
- .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id))
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(DocumentChunk.content_lexical_text.ilike(like))
- )
- if exclude_document_ids:
- stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
- result = await db.execute(stmt)
- seen: list[str] = []
- for (doc_id,) in result.all():
- if doc_id and doc_id not in seen:
- seen.append(doc_id)
- return seen
-
- async def collect_candidate_chunks(
- self,
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- entry_document_ids: Sequence[str],
- query: str,
- top_k: int,
- exclude_sections: Iterable[dict[str, str]] = (),
- ) -> list[dict[str, Any]]:
- if not entry_document_ids:
- return []
- page_size = top_k
- if exclude_sections:
- page_size = max(top_k, top_k * _SECTION_EXCLUSION_PAGE_MULTIPLIER)
- base_stmt = (
- select(Document, DocumentChunk, DocumentSection, JobResult)
- .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id))
- .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
- .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
- .where(Document.user_id == user_id)
- .where(Document.namespace == namespace)
- .where(Document.status == 'active')
- .where(Document.document_id.in_(list(entry_document_ids)))
- .where(_build_lexical_match_predicate(query))
- .order_by(DocumentChunk.sort_order)
- )
- rows = []
- offset = 0
- while len(rows) < top_k:
- result = await db.execute(base_stmt.limit(page_size).offset(offset))
- result_rows = result.all()
- if not result_rows:
- break
- for document, chunk, section, job_result in result_rows:
- section_path = section.section_path if section else None
- if is_excluded_section(
- document_id=document.document_id,
- section_path=section_path,
- exclude_sections=exclude_sections,
- ):
- continue
- rows.append({
- 'document_id': document.document_id,
- 'chunk_id': chunk.chunk_id,
- 'section_id': chunk.section_id,
- 'section_path': section_path,
- 'source_file_name': document.source_file_name,
- 'chunk_type': chunk.chunk_type,
- 'content': chunk.content,
- 'score': 2.0,
- 'file_path': chunk.file_path,
- 'chunk_metadata': chunk.chunk_metadata or {},
- 'job_result_id': chunk.job_result_id,
- 'job_id': job_result.job_id if job_result else None,
- })
- if len(rows) >= top_k:
- break
- if len(result_rows) < page_size:
- break
- offset += page_size
- return rows
diff --git a/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py b/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py
new file mode 100644
index 000000000..b101292fa
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py
@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from loguru import logger
+
+from shared.services.retrieval.hit_stats_service import record_retrieval_hits
+
+
+_pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set()
+
+
+def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None:
+ _pending_retrieval_hit_stat_tasks.discard(task)
+
+ try:
+ task.result()
+ except asyncio.CancelledError:
+ pass
+ except Exception as exc:
+ logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}')
+
+
+def schedule_retrieval_hit_stats_update(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None:
+ try:
+ task = asyncio.create_task(
+ _record_retrieval_hit_stats_best_effort(
+ user_id=user_id,
+ namespace=namespace,
+ results=results,
+ ),
+ name=f'retrieval_hit_stats:{user_id}:{namespace}',
+ )
+ _pending_retrieval_hit_stat_tasks.add(task)
+ task.add_done_callback(_finalize_retrieval_hit_stats_task)
+ except Exception as exc:
+ logger.warning(f'Failed to schedule retrieval hit stats update (ignored): {exc}')
+
+
+async def drain_retrieval_hit_stats_updates(timeout_seconds: float = 2.0) -> None:
+ if not _pending_retrieval_hit_stat_tasks:
+ return
+
+ pending_tasks = tuple(_pending_retrieval_hit_stat_tasks)
+
+ try:
+ await asyncio.wait_for(
+ asyncio.gather(*pending_tasks, return_exceptions=True),
+ timeout=timeout_seconds,
+ )
+ except asyncio.TimeoutError:
+ for task in pending_tasks:
+ if not task.done():
+ task.cancel()
+
+ await asyncio.gather(*pending_tasks, return_exceptions=True)
+
+
+async def _record_retrieval_hit_stats_best_effort(*, user_id: str, namespace: str, results: list[dict[str, Any]]) -> None:
+ try:
+ from shared.core.database import get_db_context
+
+ async with get_db_context() as db:
+ await record_retrieval_hits(db, user_id=user_id, namespace=namespace, results=results)
+ await db.commit()
+ except Exception as exc:
+ logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}')
diff --git a/packages/shared-python/shared/services/retrieval/legacy_route.py b/packages/shared-python/shared/services/retrieval/legacy_route.py
new file mode 100644
index 000000000..e8642840b
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/legacy_route.py
@@ -0,0 +1,347 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.channels import content_channel, path_channel, term_channel
+from shared.services.retrieval.graph_query_service import GraphQueryService
+from shared.services.retrieval.ranking import rank_retrieval_candidates
+from shared.services.retrieval.response_projection import attach_citation
+from shared.services.retrieval.result_assembly import assemble_retrieval_results
+from shared.services.retrieval.route_types import (
+ RetrievalRouteContext,
+ RetrievalRouteOutcome,
+)
+from shared.services.retrieval.scoring import (
+ get_row_path,
+ merge_channels_rrf,
+ merge_same_section_rows,
+ normalize_row_scores,
+)
+from shared.services.retrieval.settings import (
+ CHANNEL_WEIGHT_CONTENT,
+ CHANNEL_WEIGHT_PATH,
+ CHANNEL_WEIGHT_TERM,
+ INTERNAL_RECALL_K_MULTIPLIER,
+)
+
+
+async def run_legacy_retrieval_route(
+ context: RetrievalRouteContext,
+) -> RetrievalRouteOutcome:
+ active_channels = set(context.channels) if context.channels else {
+ "path",
+ "content",
+ "term",
+ }
+ logger.info(
+ f"\n PHASE 1: Bottom-Layer Discovery "
+ f"(channels={sorted(active_channels)})"
+ )
+ logger.info(f" effective_recall_k={context.effective_recall_k}")
+
+ path_rows = await _load_path_rows(context, active_channels)
+ content_rows = await _load_content_rows(context, active_channels)
+ term_rows = await _load_term_rows(context, active_channels)
+
+ fused_rows = _fuse_legacy_rows(
+ context=context,
+ path_rows=path_rows,
+ content_rows=content_rows,
+ term_rows=term_rows,
+ )
+ router_used, graph_rows = await _run_legacy_graph_routing(context)
+
+ ranked_rows = await rank_retrieval_candidates(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ discovery_rows=fused_rows,
+ routed_rows=graph_rows,
+ top_k=context.top_k,
+ )
+ _log_ranked_rows(ranked_rows)
+
+ assembled_rows = await assemble_retrieval_results(
+ db=context.db,
+ rows=ranked_rows,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ )
+ results = [attach_citation(row) for row in assembled_rows]
+ response = {
+ "namespace": context.namespace,
+ "query": context.query,
+ "router_used": router_used,
+ "results": results,
+ }
+ return RetrievalRouteOutcome(
+ response=response,
+ hit_stats_results=results,
+ completion_label="RETRIEVAL",
+ completion_count=len(results),
+ completion_detail=f"results | router={router_used}",
+ )
+
+
+async def list_graph_routed_chunks(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+) -> list[dict[str, Any]]:
+ service = GraphQueryService()
+ entry_document_ids = await service.find_entry_documents(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ )
+ return await service.collect_candidate_chunks(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ entry_document_ids=entry_document_ids,
+ query=query,
+ top_k=top_k * INTERNAL_RECALL_K_MULTIPLIER,
+ exclude_sections=exclude_sections,
+ )
+
+
+async def _run_legacy_graph_routing(
+ context: RetrievalRouteContext,
+) -> tuple[str, list[dict[str, Any]]]:
+ logger.info("\n PHASE 2: Legacy Graph Routing")
+ try:
+ graph_rows = await list_graph_routed_chunks(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ query=context.query,
+ top_k=context.top_k,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ )
+ if graph_rows:
+ logger.info(f" Graph routing: {len(graph_rows)} rows")
+ normalize_row_scores(
+ graph_rows,
+ source_field="score",
+ target_field="agent_score",
+ default=0.5,
+ )
+ return "discovery+graph", graph_rows
+ except Exception as exc:
+ logger.error(f" Graph routing failed (ignored): {exc}")
+
+ return "discovery_only", []
+
+
+async def _load_path_rows(
+ context: RetrievalRouteContext,
+ active_channels: set[str],
+) -> list[dict[str, Any]]:
+ if "path" not in active_channels:
+ return []
+
+ start_time = time.monotonic()
+ rows = await path_channel(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ query=context.query,
+ top_k=context.effective_recall_k,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ signal_paths=context.signal_paths,
+ filter_mode=context.filter_mode,
+ )
+ elapsed_ms = round((time.monotonic() - start_time) * 1000)
+ logger.info(f"\n path_channel: {len(rows)} rows in {elapsed_ms}ms")
+ for index, row in enumerate(rows[:5]):
+ logger.info(
+ f" [{index}] score={row.get('score', 0):.4f} "
+ f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} "
+ f"type={row.get('chunk_type', '?')}"
+ )
+ if len(rows) > 5:
+ logger.info(f" ... and {len(rows) - 5} more")
+ return rows
+
+
+async def _load_content_rows(
+ context: RetrievalRouteContext,
+ active_channels: set[str],
+) -> list[dict[str, Any]]:
+ if "content" not in active_channels:
+ return []
+
+ start_time = time.monotonic()
+ rows = await content_channel(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ query=context.query,
+ top_k=context.effective_recall_k,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ signal_paths=context.signal_paths,
+ filter_mode=context.filter_mode,
+ )
+ elapsed_ms = round((time.monotonic() - start_time) * 1000)
+ logger.info(f"\n content_channel: {len(rows)} rows in {elapsed_ms}ms")
+ for index, row in enumerate(rows[:5]):
+ logger.info(
+ f" [{index}] score={row.get('score', 0):.4f} "
+ f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} "
+ f"content={str(row.get('content', ''))[:80]}"
+ )
+ if len(rows) > 5:
+ logger.info(f" ... and {len(rows) - 5} more")
+ return rows
+
+
+async def _load_term_rows(
+ context: RetrievalRouteContext,
+ active_channels: set[str],
+) -> list[dict[str, Any]]:
+ if "term" not in active_channels:
+ return []
+
+ start_time = time.monotonic()
+ rows = await term_channel(
+ context.db,
+ user_id=context.user_id,
+ namespace=context.namespace,
+ query=context.query,
+ top_k=context.effective_recall_k,
+ exclude_document_ids=context.exclude_document_ids,
+ exclude_sections=context.exclude_sections,
+ allowed_chunk_types=context.allowed_chunk_types,
+ signal_paths=context.signal_paths,
+ filter_mode=context.filter_mode,
+ )
+ elapsed_ms = round((time.monotonic() - start_time) * 1000)
+ logger.info(f"\n term_channel: {len(rows)} rows in {elapsed_ms}ms")
+ for index, row in enumerate(rows[:5]):
+ logger.info(
+ f" [{index}] score={row.get('score', 0):.4f} "
+ f"path={row.get('section_path', '') or row.get('source_chunk_path', '')} "
+ f"type={row.get('chunk_type', '?')}"
+ )
+ if len(rows) > 5:
+ logger.info(f" ... and {len(rows) - 5} more")
+ return rows
+
+
+def _fuse_legacy_rows(
+ *,
+ context: RetrievalRouteContext,
+ path_rows: list[dict[str, Any]],
+ content_rows: list[dict[str, Any]],
+ term_rows: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ default_weights = {
+ "path": CHANNEL_WEIGHT_PATH,
+ "content": CHANNEL_WEIGHT_CONTENT,
+ "term": CHANNEL_WEIGHT_TERM,
+ }
+ effective_weights = {**default_weights, **(context.channel_weights or {})}
+
+ channel_lists: list[list[dict[str, Any]]] = []
+ weight_list: list[float] = []
+
+ if path_rows:
+ channel_lists.append(path_rows)
+ weight_list.append(effective_weights.get("path", CHANNEL_WEIGHT_PATH))
+ if content_rows:
+ channel_lists.append(content_rows)
+ weight_list.append(effective_weights.get("content", CHANNEL_WEIGHT_CONTENT))
+ if term_rows:
+ channel_lists.append(term_rows)
+ weight_list.append(effective_weights.get("term", CHANNEL_WEIGHT_TERM))
+
+ if channel_lists:
+ fused_rows = merge_channels_rrf(
+ channel_lists,
+ weight_list,
+ context.effective_recall_k,
+ )
+ else:
+ fused_rows = []
+ logger.info(
+ f"\n RRF Fusion: {len(fused_rows)} rows from "
+ f"{len(channel_lists)} channels "
+ f"(weights={dict(zip(['path', 'content', 'term'][:len(weight_list)], weight_list))})"
+ )
+ for index, row in enumerate(fused_rows[:5]):
+ logger.info(
+ f" [{index}] rrf_score={row.get('score', 0):.4f} "
+ f"path={row.get('section_path', '') or row.get('source_chunk_path', '')}"
+ )
+ if len(fused_rows) > 5:
+ logger.info(f" ... and {len(fused_rows) - 5} more")
+
+ pre_merge = len(fused_rows)
+ fused_rows = merge_same_section_rows(fused_rows)
+ if len(fused_rows) != pre_merge:
+ logger.info(f"retrieval: section_merge={pre_merge}->{len(fused_rows)}")
+
+ if context.channel_weights is not None:
+ logger.debug(f"retrieval: channel_weights={context.channel_weights}")
+
+ fused_rows = _filter_rows_by_threshold(fused_rows, context)
+ if fused_rows:
+ normalize_row_scores(
+ fused_rows,
+ source_field="score",
+ target_field="discovery_score",
+ default=0.5,
+ )
+
+ return fused_rows
+
+
+def _filter_rows_by_threshold(
+ rows: list[dict[str, Any]],
+ context: RetrievalRouteContext,
+) -> list[dict[str, Any]]:
+ if context.threshold <= 0.0 or not rows:
+ return rows
+
+ pre_count = len(rows)
+ filtered_rows = [
+ row for row in rows if row.get("score", 0.0) >= context.threshold
+ ]
+ logger.info(
+ f"retrieval: threshold_filter={pre_count}->{len(filtered_rows)} "
+ f"(threshold={context.threshold})"
+ )
+ return filtered_rows
+
+
+def _log_ranked_rows(ranked_rows: list[dict[str, Any]]) -> None:
+ if not ranked_rows:
+ return
+
+ logger.info(f"\n Unified candidate ranking: {len(ranked_rows)} rows")
+ for index, row in enumerate(ranked_rows[:10]):
+ logger.info(
+ " "
+ f"[{index}] evidence={row.get('evidence_score', 0.0):.4f} "
+ f"discovery={row.get('discovery_score', 0.0):.4f} "
+ f"agent={row.get('agent_score', 0.0):.4f} "
+ f"path={get_row_path(row)}"
+ )
diff --git a/packages/shared-python/shared/services/retrieval/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/lexical_ranker.py
new file mode 100644
index 000000000..8b50cd404
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/lexical_ranker.py
@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+
+from shared.utils.text_utils import tokenize_for_retrieval
+
+
+def tokenize_query_for_ranker(query: str) -> list[str]:
+ return tokenize_for_retrieval(query, dedupe=True)
+
+
+def rank_rows_by_bm25(
+ rows: list[dict[str, Any]],
+ query_tokens: list[str],
+ *,
+ search_field: str,
+) -> list[dict[str, Any]]:
+ """Rank matching rows with BM25 over pre-tokenized search text."""
+ try:
+ from rank_bm25 import BM25Okapi
+ except ImportError:
+ return _rank_rows_by_token_overlap(
+ rows,
+ query_tokens,
+ search_field=search_field,
+ )
+
+ corpus: list[list[str]] = []
+ ranked_rows: list[dict[str, Any]] = []
+ query_token_set = set(query_tokens)
+ for row in rows:
+ tokens = _get_search_tokens(row, search_field=search_field)
+ if not tokens or not query_token_set.intersection(tokens):
+ continue
+ corpus.append(tokens)
+ ranked_rows.append(row)
+
+ if not corpus or not query_tokens:
+ return []
+
+ bm25 = BM25Okapi(corpus)
+ scores = bm25.get_scores(query_tokens)
+
+ for index, row in enumerate(ranked_rows):
+ row["score"] = float(scores[index])
+
+ ranked_rows.sort(key=lambda row: row["score"], reverse=True)
+ return ranked_rows
+
+
+def _rank_rows_by_token_overlap(
+ rows: list[dict[str, Any]],
+ query_tokens: list[str],
+ *,
+ search_field: str,
+) -> list[dict[str, Any]]:
+ logger.warning("rank_bm25 not installed, skipping BM25 re-rank")
+ ranked_rows: list[dict[str, Any]] = []
+ query_token_set = set(query_tokens)
+ for row in rows:
+ tokens = _get_search_tokens(row, search_field=search_field)
+ overlap = len(query_token_set.intersection(tokens))
+ if overlap <= 0:
+ continue
+ row["score"] = float(overlap)
+ ranked_rows.append(row)
+ ranked_rows.sort(key=lambda row: row["score"], reverse=True)
+ return ranked_rows
+
+
+def _get_search_tokens(row: dict[str, Any], *, search_field: str) -> list[str]:
+ return [token for token in str(row.get(search_field) or "").split() if token]
diff --git a/packages/shared-python/shared/services/retrieval/path_hydration.py b/packages/shared-python/shared/services/retrieval/path_hydration.py
new file mode 100644
index 000000000..d973f8a8b
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/path_hydration.py
@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import or_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, DocumentSection
+from shared.models.database.job_result import JobResult
+from shared.services.retrieval.lexical_text import normalize_section_path
+from shared.services.retrieval.scoring import get_row_path
+
+
+async def hydrate_paths_to_rows(
+ db: AsyncSession,
+ *,
+ path_selections: list[dict[str, Any]],
+ user_id: str,
+ namespace: str,
+ document_id: str | None = None,
+) -> list[dict[str, Any]]:
+ """Load full chunk rows by section_path or source_chunk_path."""
+ if not path_selections:
+ return []
+
+ confidence_by_path: dict[str, float] = {}
+ mode_by_path: dict[str, str] = {}
+ ordered_paths: list[str] = []
+ for item in path_selections:
+ raw_path = str(item.get('path') or '').strip()
+ path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path
+ if not path:
+ continue
+ confidence = float(item.get('confidence', 0.0) or 0.0)
+ hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower()
+ if path not in confidence_by_path:
+ ordered_paths.append(path)
+ confidence_by_path[path] = confidence
+ mode_by_path[path] = hydrate_mode
+ else:
+ confidence_by_path[path] = max(confidence_by_path[path], confidence)
+ if not ordered_paths:
+ return []
+
+ outline_paths = [path for path in ordered_paths if mode_by_path.get(path) == 'outline']
+ chunk_paths = [path for path in ordered_paths if mode_by_path.get(path) != 'outline']
+
+ rows: list[dict[str, Any]] = []
+
+ if outline_paths:
+ rows.extend(
+ await _hydrate_outline_paths(
+ db,
+ outline_paths=outline_paths,
+ confidence_by_path=confidence_by_path,
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
+ )
+
+ if chunk_paths:
+ rows.extend(
+ await _hydrate_chunk_paths(
+ db,
+ chunk_paths=chunk_paths,
+ confidence_by_path=confidence_by_path,
+ mode_by_path=mode_by_path,
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
+ )
+
+ _sort_rows_by_selection_order(rows, ordered_paths)
+ _log_hydration_resolution(rows=rows, ordered_paths=ordered_paths, outline_paths=outline_paths)
+ return rows
+
+
+async def _hydrate_outline_paths(
+ db: AsyncSession,
+ *,
+ outline_paths: list[str],
+ confidence_by_path: dict[str, float],
+ user_id: str,
+ namespace: str,
+ document_id: str | None,
+) -> list[dict[str, Any]]:
+ outline_section_filters = [
+ DocumentSection.section_path == path
+ for path in outline_paths
+ ]
+ outline_stmt = (
+ select(Document, DocumentSection)
+ .join(
+ DocumentSection,
+ (DocumentSection.document_id == Document.document_id)
+ & (DocumentSection.job_result_id == Document.current_job_result_id),
+ )
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .where(or_(*outline_section_filters))
+ )
+ if document_id:
+ outline_stmt = outline_stmt.where(Document.document_id == document_id)
+
+ rows: list[dict[str, Any]] = []
+ outline_result = await db.execute(outline_stmt)
+ for document, section in outline_result.all():
+ agent_score = confidence_by_path.get(section.section_path, 0.0)
+ summary_text = (section.summary or '').strip()
+ title_text = (section.section_title or '').strip()
+ content = f'[Outline] {title_text}'
+ if summary_text:
+ content += f'\n{summary_text}'
+ rows.append({
+ 'document_id': document.document_id,
+ 'chunk_id': f'outline_{section.section_id}',
+ 'section_id': section.section_id,
+ 'section_path': section.section_path,
+ 'source_file_name': document.source_file_name,
+ 'chunk_type': 'outline',
+ 'content': content,
+ 'score': agent_score,
+ 'agent_score': agent_score,
+ 'file_path': None,
+ 'chunk_metadata': {},
+ 'job_result_id': section.job_result_id,
+ 'job_id': None,
+ 'source_chunk_path': None,
+ 'sort_order': section.sort_order,
+ 'hydrate_mode': 'outline',
+ })
+ return rows
+
+
+async def _hydrate_chunk_paths(
+ db: AsyncSession,
+ *,
+ chunk_paths: list[str],
+ confidence_by_path: dict[str, float],
+ mode_by_path: dict[str, str],
+ user_id: str,
+ namespace: str,
+ document_id: str | None,
+) -> list[dict[str, Any]]:
+ section_path_filters = []
+ self_only_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'self_only'}
+ for path in chunk_paths:
+ section_path_filters.append(DocumentSection.section_path == path)
+ if path not in self_only_paths:
+ section_path_filters.append(DocumentSection.section_path.like(f'{path} / %'))
+
+ stmt = (
+ select(Document, DocumentChunk, DocumentSection, JobResult)
+ .join(
+ DocumentChunk,
+ (DocumentChunk.document_id == Document.document_id)
+ & (DocumentChunk.job_result_id == Document.current_job_result_id),
+ )
+ .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
+ .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .where(
+ or_(
+ *section_path_filters,
+ DocumentChunk.source_chunk_path.in_(chunk_paths),
+ )
+ )
+ )
+ if document_id:
+ stmt = stmt.where(Document.document_id == document_id)
+ result = await db.execute(stmt)
+
+ rows: list[dict[str, Any]] = []
+ seen_paths: set[str] = set()
+ for document, chunk, section, job_result in result.all():
+ row_path = (section.section_path if section else None) or chunk.source_chunk_path or ''
+ if row_path in seen_paths:
+ continue
+
+ matched_path = row_path
+ if section and section.section_path not in confidence_by_path:
+ matched_path = next(
+ (
+ path for path in chunk_paths
+ if section.section_path == path
+ or section.section_path.startswith(f'{path} / ')
+ ),
+ row_path,
+ )
+
+ path_mode = mode_by_path.get(matched_path, 'chunks')
+ allowed_types = _get_allowed_types_for_mode(path_mode)
+ if allowed_types is not None:
+ chunk_type_lower = (chunk.chunk_type or '').strip().lower()
+ if chunk_type_lower not in allowed_types:
+ continue
+
+ seen_paths.add(row_path)
+ agent_score = confidence_by_path.get(matched_path, 0.0)
+ rows.append({
+ 'document_id': document.document_id,
+ 'chunk_id': chunk.chunk_id,
+ 'section_id': chunk.section_id,
+ 'section_path': section.section_path if section else None,
+ 'source_file_name': document.source_file_name,
+ 'chunk_type': chunk.chunk_type,
+ 'content': chunk.content,
+ 'score': agent_score,
+ 'agent_score': agent_score,
+ 'file_path': chunk.file_path,
+ 'chunk_metadata': chunk.chunk_metadata or {},
+ 'job_result_id': chunk.job_result_id,
+ 'job_id': job_result.job_id if job_result else None,
+ 'source_chunk_path': chunk.source_chunk_path,
+ 'sort_order': chunk.sort_order,
+ 'hydrate_mode': path_mode,
+ })
+ return rows
+
+
+def _get_allowed_types_for_mode(path_mode: str) -> set[str] | None:
+ mode_allowed_types: dict[str, set[str] | None] = {
+ 'chunks': None,
+ 'self_only': None,
+ 'assets_only': {'image', 'table'},
+ 'image_only': {'image'},
+ 'table_only': {'table'},
+ }
+ return mode_allowed_types.get(path_mode)
+
+
+def _sort_rows_by_selection_order(rows: list[dict[str, Any]], ordered_paths: list[str]) -> None:
+ path_order = {path: index for index, path in enumerate(ordered_paths)}
+
+ def row_sort_key(row: dict[str, Any]) -> int:
+ row_path = get_row_path(row)
+ if row_path in path_order:
+ return path_order[row_path]
+ for path, index in path_order.items():
+ if row_path.startswith(f'{path} / '):
+ return index
+ return 10**9
+
+ rows.sort(key=row_sort_key)
+
+
+def _log_hydration_resolution(
+ *, rows: list[dict[str, Any]], ordered_paths: list[str], outline_paths: list[str]
+) -> None:
+ hydrated_paths = {get_row_path(row) for row in rows}
+ resolved_inputs = {
+ path for path in ordered_paths
+ if path in hydrated_paths
+ or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths)
+ }
+ resolved_inputs |= set(outline_paths)
+ missed = len(ordered_paths) - len(resolved_inputs)
+ if missed > 0:
+ missing_paths = [path for path in ordered_paths if path not in resolved_inputs]
+ logger.warning(
+ f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); '
+ f'missing[:5]={missing_paths[:5]}'
+ )
+ else:
+ logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved')
diff --git a/packages/shared-python/shared/services/retrieval/ranking.py b/packages/shared-python/shared/services/retrieval/ranking.py
new file mode 100644
index 000000000..bd1cfd9fd
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/ranking.py
@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import RetrievalHitStat
+from shared.services.retrieval.hit_stats_service import compute_importance_score
+from shared.services.retrieval.scoring import get_row_path
+
+
+def get_candidate_key(row: dict[str, Any]) -> str:
+ path = get_row_path(row)
+ if path:
+ return f'path:{path}'
+ chunk_id = str(row.get('chunk_id') or '').strip()
+ return f'chunk:{chunk_id}' if chunk_id else ''
+
+
+async def load_chunk_importance_scores(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ rows: list[dict[str, Any]],
+) -> dict[str, float]:
+ chunk_ids = sorted({
+ str(row.get('chunk_id') or '').strip()
+ for row in rows
+ if row.get('chunk_id')
+ })
+ if not chunk_ids:
+ return {}
+ stmt = (
+ select(
+ RetrievalHitStat.chunk_id,
+ RetrievalHitStat.hit_count,
+ RetrievalHitStat.last_hit_at,
+ RetrievalHitStat.created_at,
+ )
+ .where(RetrievalHitStat.user_id == user_id)
+ .where(RetrievalHitStat.namespace == namespace)
+ .where(RetrievalHitStat.hit_kind == 'chunk')
+ .where(RetrievalHitStat.chunk_id.in_(chunk_ids))
+ )
+ result = await db.execute(stmt)
+ importance_scores: dict[str, float] = {}
+ for chunk_id, hit_count, last_hit_at, created_at in result.all():
+ if not chunk_id:
+ continue
+ importance_scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at)
+ return importance_scores
+
+
+def apply_importance_multiplier(
+ rows: list[dict[str, Any]],
+ *,
+ raw_field: str = 'importance_raw_score',
+ low: float = 0.1,
+ high: float = 2.0,
+) -> None:
+ if not rows:
+ return
+
+ values = sorted(float(row.get(raw_field, 0.0) or 0.0) for row in rows)
+ item_count = len(values)
+ median = values[item_count // 2] if item_count % 2 else (values[item_count // 2 - 1] + values[item_count // 2]) / 2
+ q1 = values[item_count // 4] if item_count >= 4 else values[0]
+ q3 = values[3 * item_count // 4] if item_count >= 4 else values[-1]
+ iqr = q3 - q1
+
+ for row in rows:
+ raw_score = float(row.get(raw_field, 0.0) or 0.0)
+ if iqr <= 1e-9:
+ multiplier = 1.0
+ else:
+ z_score = (raw_score - median) / iqr
+ sigmoid_score = 1.0 / (1.0 + math.exp(-z_score))
+ multiplier = low + (high - low) * sigmoid_score
+ row['importance_multiplier'] = round(multiplier, 4)
+ row['agent_score'] = round(
+ float(row.get('agent_score', 0.0) or 0.0) * multiplier,
+ 6,
+ )
+ row['discovery_score'] = round(
+ float(row.get('discovery_score', 0.0) or 0.0) * multiplier,
+ 6,
+ )
+
+
+def rank_candidates_by_path(
+ discovery_rows: list[dict[str, Any]],
+ routed_rows: list[dict[str, Any]],
+ top_k: int,
+ *,
+ importance_scores: dict[str, float] | None = None,
+) -> list[dict[str, Any]]:
+ merged: dict[str, dict[str, Any]] = {}
+ insertion_order: dict[str, int] = {}
+ counter = 0
+
+ for row in discovery_rows:
+ key = get_candidate_key(row)
+ if not key:
+ continue
+ candidate = dict(row)
+ candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0)
+ candidate['agent_score'] = 0.0
+ candidate.setdefault('hydrate_mode', 'chunks')
+ merged[key] = candidate
+ insertion_order[key] = counter
+ counter += 1
+
+ for row in routed_rows:
+ key = get_candidate_key(row)
+ if not key:
+ continue
+ routed_agent_score = float(row.get('agent_score', 0.0) or 0.0)
+ if key not in merged:
+ candidate = dict(row)
+ candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0)
+ candidate['agent_score'] = routed_agent_score
+ merged[key] = candidate
+ insertion_order[key] = counter
+ counter += 1
+ continue
+ candidate = merged[key]
+ candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score)
+ if not candidate.get('source_chunk_path') and row.get('source_chunk_path'):
+ candidate['source_chunk_path'] = row.get('source_chunk_path')
+ if not candidate.get('section_path') and row.get('section_path'):
+ candidate['section_path'] = row.get('section_path')
+
+ for row in merged.values():
+ row['importance_raw_score'] = float(
+ (importance_scores or {}).get(str(row.get('chunk_id') or ''), 0.0) or 0.0
+ )
+ apply_importance_multiplier(list(merged.values()))
+
+ has_agent_results = len(routed_rows) > 0
+ primary_rows: list[dict[str, Any]] = []
+ fallback_rows: list[dict[str, Any]] = []
+
+ for key, row in merged.items():
+ agent_score = float(row.get('agent_score', 0.0) or 0.0)
+ discovery_score = float(row.get('discovery_score', 0.0) or 0.0)
+ row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6)
+ row['score'] = row['evidence_score']
+ row['_candidate_order'] = insertion_order[key]
+
+ if has_agent_results and agent_score <= 0.0:
+ fallback_rows.append(row)
+ else:
+ primary_rows.append(row)
+
+ def get_sort_key(row: dict[str, Any]) -> tuple[float, float, int]:
+ return (
+ float(row.get('agent_score', 0.0) or 0.0),
+ float(row.get('discovery_score', 0.0) or 0.0),
+ -int(row.get('_candidate_order', 0) or 0),
+ )
+
+ primary_rows.sort(key=get_sort_key, reverse=True)
+ ranked_rows = primary_rows[:top_k]
+
+ if len(ranked_rows) < top_k and fallback_rows:
+ fallback_rows.sort(key=get_sort_key, reverse=True)
+ ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)])
+
+ for row in ranked_rows:
+ row.pop('_candidate_order', None)
+ return ranked_rows
+
+
+async def rank_retrieval_candidates(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ discovery_rows: list[dict[str, Any]],
+ routed_rows: list[dict[str, Any]],
+ top_k: int,
+) -> list[dict[str, Any]]:
+ try:
+ importance_scores = await load_chunk_importance_scores(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ rows=[*discovery_rows, *routed_rows],
+ )
+ except Exception as exc:
+ logger.warning(f'Failed to load chunk importance scores, continuing without importance: {exc}')
+ importance_scores = {}
+ return rank_candidates_by_path(
+ discovery_rows,
+ routed_rows,
+ top_k,
+ importance_scores=importance_scores,
+ )
diff --git a/packages/shared-python/shared/services/retrieval/reference_hydration.py b/packages/shared-python/shared/services/retrieval/reference_hydration.py
new file mode 100644
index 000000000..299024e2e
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/reference_hydration.py
@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, DocumentSection
+from shared.models.database.job_result import JobResult
+from shared.services.retrieval.row_utils import (
+ ReferenceLookupKey,
+ build_reference_lookup_key,
+)
+
+
+async def hydrate_referenced_chunk_rows(
+ *,
+ db: AsyncSession | None,
+ user_id: str,
+ namespace: str,
+ refs: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ if db is None or not refs:
+ return []
+
+ ref_keys = [
+ build_reference_lookup_key(
+ document_id=ref.get('document_id'),
+ chunk_id=ref.get('chunk_id'),
+ section_path=ref.get('section_path'),
+ file_path=ref.get('file_path'),
+ )
+ for ref in refs
+ ]
+ ref_keys = [key for key in ref_keys if key[0] and key[1]]
+ if not ref_keys:
+ return []
+
+ document_ids = sorted({document_id for document_id, _, _, _ in ref_keys})
+ chunk_ids = sorted({chunk_id for _, chunk_id, _, _ in ref_keys})
+ stmt = (
+ select(Document, DocumentChunk, DocumentSection, JobResult)
+ .join(
+ DocumentChunk,
+ (DocumentChunk.document_id == Document.document_id)
+ & (DocumentChunk.job_result_id == Document.current_job_result_id),
+ )
+ .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
+ .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .where(Document.document_id.in_(document_ids))
+ .where(DocumentChunk.chunk_id.in_(chunk_ids))
+ .order_by(DocumentChunk.sort_order)
+ )
+ result = await db.execute(stmt)
+
+ rows_by_key: dict[ReferenceLookupKey, dict[str, Any]] = {}
+ rows_by_base_key: dict[tuple[str, str], list[dict[str, Any]]] = {}
+ for document, chunk, section, job_result in result.all():
+ row = {
+ 'document_id': document.document_id,
+ 'chunk_id': chunk.chunk_id,
+ 'section_id': chunk.section_id,
+ 'section_path': section.section_path if section else None,
+ 'source_file_name': document.source_file_name,
+ 'chunk_type': chunk.chunk_type,
+ 'content': chunk.content,
+ 'score': 1.0,
+ 'file_path': chunk.file_path,
+ 'chunk_metadata': chunk.chunk_metadata or {},
+ 'job_result_id': chunk.job_result_id,
+ 'job_id': job_result.job_id if job_result else None,
+ 'source_chunk_path': chunk.source_chunk_path,
+ 'sort_order': chunk.sort_order,
+ }
+ key = build_reference_lookup_key(
+ document_id=row['document_id'],
+ chunk_id=row['chunk_id'],
+ section_path=row['section_path'],
+ file_path=row['file_path'],
+ )
+ rows_by_key[key] = row
+ rows_by_base_key.setdefault((key[0], key[1]), []).append(row)
+
+ rows: list[dict[str, Any]] = []
+ seen_keys: set[ReferenceLookupKey] = set()
+ for key in ref_keys:
+ row = rows_by_key.get(key)
+ if row is None:
+ candidates = rows_by_base_key.get((key[0], key[1]), [])
+ row = next(
+ (
+ candidate
+ for candidate in candidates
+ if key[2]
+ and str(candidate.get('section_path') or '').strip() == key[2]
+ ),
+ None,
+ )
+ if row is None:
+ row = next(
+ (
+ candidate
+ for candidate in candidates
+ if build_reference_lookup_key(
+ document_id=candidate.get('document_id'),
+ chunk_id=candidate.get('chunk_id'),
+ section_path=candidate.get('section_path'),
+ file_path=candidate.get('file_path'),
+ )
+ not in seen_keys
+ ),
+ None,
+ )
+ if row is not None:
+ row_key = build_reference_lookup_key(
+ document_id=row.get('document_id'),
+ chunk_id=row.get('chunk_id'),
+ section_path=row.get('section_path'),
+ file_path=row.get('file_path'),
+ )
+ if row_key in seen_keys:
+ continue
+ seen_keys.add(row_key)
+ rows.append(row)
+ return rows
diff --git a/packages/shared-python/shared/services/retrieval/response_projection.py b/packages/shared-python/shared/services/retrieval/response_projection.py
new file mode 100644
index 000000000..f12638a5c
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/response_projection.py
@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Any
+
+from shared.services.retrieval.assets import enrich_rows_with_retrieval_asset_urls
+from shared.services.retrieval.row_utils import (
+ PUBLIC_RESULT_FIELDS,
+ PUBLIC_SOURCE_FIELDS,
+)
+
+
+def attach_citation(row: dict[str, Any]) -> dict[str, Any]:
+ citation = {
+ 'document_id': row.get('document_id'),
+ 'chunk_id': row.get('chunk_id'),
+ 'source_file_name': row.get('source_file_name'),
+ 'section_path': row.get('section_path'),
+ }
+ return {**row, 'citation': citation}
+
+
+def to_public_source(row: dict[str, Any]) -> dict[str, Any]:
+ return {field: row.get(field) for field in PUBLIC_SOURCE_FIELDS}
+
+
+async def enrich_referenced_chunks_with_asset_urls(refs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ return await enrich_rows_with_retrieval_asset_urls(
+ refs,
+ log_context='agentic referenced chunk',
+ )
+
+
+async def project_public_retrieval_response(response: dict[str, Any]) -> dict[str, Any]:
+ public_response = {
+ 'namespace': response.get('namespace'),
+ 'query': response.get('query'),
+ 'router_used': response.get('router_used'),
+ 'results': [],
+ }
+
+ if response.get('answer_text') is not None:
+ public_response['answer_text'] = response['answer_text']
+ if response.get('referenced_chunks') is not None:
+ public_response['referenced_chunks'] = response['referenced_chunks']
+
+ projected_rows = await enrich_rows_with_retrieval_asset_urls(
+ response.get('results', []),
+ log_context='retrieval result',
+ )
+ public_results: list[dict[str, Any]] = []
+ for row in projected_rows:
+ public_row: dict[str, Any] = {}
+ for field in PUBLIC_RESULT_FIELDS:
+ if field in row:
+ public_row[field] = row[field]
+ if 'source' in row:
+ public_row['source'] = row['source']
+ else:
+ public_row['source'] = to_public_source(row)
+ public_results.append(public_row)
+
+ public_response['results'] = public_results
+ return public_response
diff --git a/packages/shared-python/shared/services/retrieval/result_assembly.py b/packages/shared-python/shared/services/retrieval/result_assembly.py
new file mode 100644
index 000000000..1ef948078
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/result_assembly.py
@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows
+from shared.services.retrieval.row_utils import (
+ clean_content,
+ filter_excluded_rows,
+ iter_connected_target_ids,
+ normalize_chunk_type,
+)
+
+
+async def assemble_retrieval_results(
+ *,
+ db: AsyncSession | None = None,
+ rows: list[dict[str, Any]],
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ allowed_chunk_types: set[str] | None = None,
+) -> list[dict[str, Any]]:
+ filtered_rows = filter_excluded_rows(
+ rows,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ )
+ if allowed_chunk_types is not None:
+ filtered_rows = [
+ row for row in filtered_rows
+ if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types
+ ]
+ hydrated_rows = await hydrate_connected_target_rows(
+ db=db,
+ rows=filtered_rows,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ )
+ rows_by_chunk_id = {
+ str(row.get('chunk_id') or ''): row
+ for row in [*filtered_rows, *hydrated_rows]
+ if row.get('chunk_id')
+ }
+
+ embedded_targets: set[str] = set()
+ for row in filtered_rows:
+ for target_id in iter_connected_target_ids(row):
+ if target_id in rows_by_chunk_id:
+ embedded_targets.add(target_id)
+
+ assembled: list[dict[str, Any]] = []
+ for row in filtered_rows:
+ if row.get('chunk_id') in embedded_targets:
+ continue
+ assembled_row = dict(row)
+ base_content = str(row.get('content') or '')
+ if normalize_chunk_type(row.get('chunk_type')) == 'text':
+ connected_targets: list[tuple[int, str]] = []
+ for target_id in iter_connected_target_ids(row):
+ target_row = rows_by_chunk_id.get(target_id)
+ if not target_row:
+ continue
+ if normalize_chunk_type(target_row.get('chunk_type')) != 'table':
+ continue
+ target_content = str(target_row.get('content') or '').strip()
+ if target_content:
+ sort_key = int(target_row.get('sort_order', 0) or 0)
+ connected_targets.append((sort_key, target_content))
+ connected_targets.sort(key=lambda item: item[0])
+ related_parts = [content for _, content in connected_targets]
+ if base_content and related_parts:
+ assembled_row['content'] = '\n\n'.join([base_content, *related_parts])
+ else:
+ assembled_row['content'] = base_content
+ else:
+ assembled_row['content'] = base_content
+ assembled_row['content'] = clean_content(assembled_row['content'])
+ assembled.append(assembled_row)
+ return assembled
diff --git a/packages/shared-python/shared/services/retrieval/route_types.py b/packages/shared-python/shared/services/retrieval/route_types.py
new file mode 100644
index 000000000..f1aa94160
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/route_types.py
@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+@dataclass(frozen=True)
+class RetrievalRouteContext:
+ db: AsyncSession
+ user_id: str
+ namespace: str
+ query: str
+ top_k: int
+ exclude_document_ids: list[str]
+ exclude_sections: list[dict[str, str]]
+ allowed_chunk_types: set[str] | None
+ data_type: int
+ signal_paths: list[str] | None
+ filter_mode: str
+ channels: list[str] | None
+ channel_weights: dict[str, float] | None
+ threshold: float
+ effective_recall_k: int
+ use_agentic: bool | None
+
+
+@dataclass(frozen=True)
+class RetrievalRouteOutcome:
+ response: dict[str, Any]
+ hit_stats_results: list[dict[str, Any]]
+ completion_label: str
+ completion_count: int
+ completion_detail: str
diff --git a/packages/shared-python/shared/services/retrieval/row_utils.py b/packages/shared-python/shared/services/retrieval/row_utils.py
new file mode 100644
index 000000000..0784830b0
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/row_utils.py
@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from shared.services.retrieval.section_filters import is_excluded_section
+
+MEDIA_CHUNK_TYPES = {'image', 'table'}
+PUBLIC_RESULT_FIELDS = {
+ 'chunk_type', 'content', 'score', 'asset_url',
+}
+PUBLIC_SOURCE_FIELDS = {
+ 'document_id', 'source_file_name', 'section_path',
+}
+
+ReferenceLookupKey = tuple[str, str, str, str]
+
+_PATH_REF_RE = re.compile(r'\[(?:images|tables)/[^\]\n]+\]')
+
+
+def clean_content(content: str) -> str:
+ return _PATH_REF_RE.sub('', content).strip()
+
+
+def normalize_chunk_type(raw: object) -> str:
+ return str(raw or '').strip().split('\n', 1)[0].lower()
+
+
+def is_media_chunk(row: dict[str, Any]) -> bool:
+ return normalize_chunk_type(row.get('chunk_type')) in MEDIA_CHUNK_TYPES
+
+
+def build_reference_lookup_key(
+ *,
+ document_id: object,
+ chunk_id: object,
+ section_path: object = '',
+ file_path: object = '',
+) -> ReferenceLookupKey:
+ return (
+ str(document_id or '').strip(),
+ str(chunk_id or '').strip(),
+ str(section_path or '').strip(),
+ str(file_path or '').strip(),
+ )
+
+
+def filter_excluded_rows(
+ rows: list[dict[str, Any]],
+ *,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+) -> list[dict[str, Any]]:
+ filtered: list[dict[str, Any]] = []
+ excluded_documents = set(exclude_document_ids)
+ for row in rows:
+ document_id = row.get('document_id')
+ if document_id in excluded_documents:
+ continue
+ if is_excluded_section(
+ document_id=document_id,
+ section_path=row.get('section_path'),
+ exclude_sections=exclude_sections,
+ ):
+ continue
+ filtered.append(row)
+ return filtered
+
+
+def iter_connected_target_ids(row: dict[str, Any]) -> list[str]:
+ metadata = row.get('chunk_metadata') or {}
+ if not isinstance(metadata, dict):
+ return []
+
+ target_ids: list[str] = []
+ for item in metadata.get('connect_to') or []:
+ if not isinstance(item, dict):
+ continue
+ target_id = str(item.get('target') or '').strip()
+ if target_id:
+ target_ids.append(target_id)
+ return target_ids
diff --git a/packages/shared-python/shared/services/retrieval/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/scoped_corpus.py
new file mode 100644
index 000000000..012ba7129
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/scoped_corpus.py
@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.models.database.document import Document, DocumentChunk, DocumentSection
+from shared.models.database.job_result import JobResult
+from shared.services.retrieval.section_filters import is_excluded_section
+
+
+async def count_scoped_chunks(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ exclude_document_ids: list[str],
+ allowed_chunk_types: set[str] | None,
+) -> int:
+ stmt = (
+ select(func.count(DocumentChunk.id))
+ .join(
+ Document,
+ (Document.document_id == DocumentChunk.document_id)
+ & (Document.current_job_result_id == DocumentChunk.job_result_id),
+ )
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ )
+ if exclude_document_ids:
+ stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
+ if allowed_chunk_types is not None:
+ stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types)))
+ result = await db.execute(stmt)
+ return result.scalar() or 0
+
+
+async def load_all_scoped_chunks(
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ allowed_chunk_types: set[str] | None,
+ signal_paths: list[str],
+ filter_mode: str,
+) -> list[dict[str, Any]]:
+ stmt = (
+ select(Document, DocumentChunk, DocumentSection, JobResult)
+ .join(
+ DocumentChunk,
+ (DocumentChunk.document_id == Document.document_id)
+ & (DocumentChunk.job_result_id == Document.current_job_result_id),
+ )
+ .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id)
+ .join(JobResult, JobResult.id == DocumentChunk.job_result_id)
+ .where(Document.user_id == user_id)
+ .where(Document.namespace == namespace)
+ .where(Document.status == 'active')
+ .order_by(DocumentChunk.sort_order)
+ )
+ if exclude_document_ids:
+ stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
+ if allowed_chunk_types is not None:
+ stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types)))
+
+ result = await db.execute(stmt)
+ rows: list[dict[str, Any]] = []
+ for document, chunk, section, job_result in result.all():
+ section_path = section.section_path if section else None
+ if is_excluded_section(
+ document_id=document.document_id,
+ section_path=section_path,
+ exclude_sections=exclude_sections,
+ ):
+ continue
+ if signal_paths and section_path:
+ path_lower = section_path.lower()
+ matches_any = any(keyword.lower() in path_lower for keyword in signal_paths)
+ if filter_mode == 'keep' and not matches_any:
+ continue
+ if filter_mode == 'delete' and matches_any:
+ continue
+ rows.append({
+ 'document_id': document.document_id,
+ 'chunk_id': chunk.chunk_id,
+ 'section_id': chunk.section_id,
+ 'section_path': section_path,
+ 'source_file_name': document.source_file_name,
+ 'chunk_type': chunk.chunk_type,
+ 'content': chunk.content,
+ 'score': 1.0,
+ 'file_path': chunk.file_path,
+ 'chunk_metadata': chunk.chunk_metadata or {},
+ 'job_result_id': chunk.job_result_id,
+ 'job_id': job_result.job_id if job_result else None,
+ 'sort_order': chunk.sort_order,
+ })
+ return rows
diff --git a/packages/shared-python/shared/services/retrieval/scoring.py b/packages/shared-python/shared/services/retrieval/scoring.py
new file mode 100644
index 000000000..848a3adac
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/scoring.py
@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import Any
+
+from shared.services.retrieval.settings import RRF_K
+
+
+def get_row_path(row: dict[str, Any]) -> str:
+ """Extract the canonical path from a row for deduplication."""
+ return str(row.get('section_path') or row.get('source_chunk_path') or '')
+
+
+def merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ if not rows:
+ return rows
+ groups: dict[str, list[dict[str, Any]]] = {}
+ order: list[str] = []
+ for row in rows:
+ section_path = row.get('section_path')
+ if section_path:
+ key = f"{row.get('document_id', '')}::{section_path}"
+ else:
+ key = row.get('chunk_id', '')
+ if key not in groups:
+ groups[key] = []
+ order.append(key)
+ groups[key].append(row)
+
+ merged: list[dict[str, Any]] = []
+ for key in order:
+ group = groups[key]
+ if len(group) == 1:
+ merged.append(group[0])
+ continue
+ base = dict(group[0])
+ base['content'] = '\n'.join(str(row.get('content', '')) for row in group)
+ base['score'] = max(row.get('score', 0.0) for row in group)
+ merged.append(base)
+ return merged
+
+
+def merge_channels_rrf(
+ channels: list[list[dict[str, Any]]],
+ weights: list[float],
+ top_k: int,
+ k: int = RRF_K,
+) -> list[dict[str, Any]]:
+ """Reciprocal Rank Fusion across multiple retrieval channels."""
+ score_dict: dict[str, float] = {}
+ row_by_chunk_id: dict[str, dict[str, Any]] = {}
+
+ for channel_idx, channel_rows in enumerate(channels):
+ weight = weights[channel_idx] if channel_idx < len(weights) else 1.0
+ for rank, row in enumerate(channel_rows):
+ chunk_id = str(row.get('chunk_id') or '')
+ if not chunk_id:
+ continue
+ rrf_score = weight / (k + rank + 1)
+ score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score
+ if chunk_id not in row_by_chunk_id:
+ row_by_chunk_id[chunk_id] = row
+
+ ranked = sorted(score_dict.items(), key=lambda x: x[1], reverse=True)
+ results: list[dict[str, Any]] = []
+ for chunk_id, fused_score in ranked[:top_k]:
+ row = row_by_chunk_id[chunk_id]
+ results.append(dict(row, score=round(fused_score, 6)))
+ return results
+
+
+def normalize_row_scores(
+ rows: list[dict[str, Any]],
+ *,
+ source_field: str,
+ target_field: str,
+ default: float,
+) -> None:
+ if not rows:
+ return
+ values = [float(row.get(source_field, 0.0) or 0.0) for row in rows]
+ min_score = min(values)
+ max_score = max(values)
+ if max_score <= 0.0 and min_score <= 0.0:
+ for row in rows:
+ row[target_field] = 0.0
+ return
+ if max_score == min_score:
+ for row in rows:
+ row[target_field] = default
+ return
+ denominator = max_score - min_score
+ for row in rows:
+ raw_score = float(row.get(source_field, 0.0) or 0.0)
+ row[target_field] = round((raw_score - min_score) / denominator, 6)
diff --git a/packages/shared-python/shared/services/retrieval/section_filters.py b/packages/shared-python/shared/services/retrieval/section_filters.py
new file mode 100644
index 000000000..f74b7f3c3
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/section_filters.py
@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+
+def is_excluded_section(
+ *,
+ document_id: str | None,
+ section_path: str | None,
+ exclude_sections: Iterable[dict[str, str]],
+) -> bool:
+ document_id = str(document_id or '').strip()
+ section_path = str(section_path or '').strip()
+ if not document_id or not section_path:
+ return False
+ for item in exclude_sections:
+ if not isinstance(item, dict):
+ continue
+ exc_doc = str(item.get('document_id') or '').strip()
+ exc_path = str(item.get('section_path') or '').strip()
+ if document_id == exc_doc and (
+ section_path == exc_path or section_path.startswith(exc_path + ' / ')
+ ):
+ return True
+ return False
diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py
new file mode 100644
index 000000000..49f2461b5
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/settings.py
@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+CHANNEL_WEIGHT_PATH = 1.0
+CHANNEL_WEIGHT_CONTENT = 2.0
+CHANNEL_WEIGHT_TERM = 1.5
+INTERNAL_RECALL_K_MULTIPLIER = 2
+RRF_K = 60
+
+DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = {
+ 1: None,
+ 2: {'text'},
+ 3: {'image'},
+ 4: {'table'},
+ 5: {'text', 'image'},
+ 6: {'text', 'table'},
+}
+
+
+def resolve_allowed_chunk_types(data_type: int) -> set[str] | None:
+ return DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type)
diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
index abce68b36..3fd425521 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
@@ -2,37 +2,45 @@
from __future__ import annotations
import asyncio
-import os
import time
-from typing import Any
+from collections.abc import Callable
+from contextlib import AbstractAsyncContextManager
from uuid import uuid4
from loguru import logger
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.core.database import get_db_context
from shared.services.retrieval.agentic.budget import BudgetLedger
-from shared.services.retrieval.agentic.orchestrator import RetrievalAgent, _load_budget_inventory
-from shared.services.retrieval.agentic.types import AgenticResult
-from shared.services.retrieval.cache_service import (
- get_cached_workflow_plan,
- set_cached_workflow_plan,
-)
+from shared.services.retrieval.agentic.orchestrator import _load_budget_inventory
from shared.services.retrieval.llm_adapter import (
create_retrieval_llm_fn,
create_retrieval_planner_fn,
)
-from shared.services.retrieval.workflow.planner import QueryPlanner
-from shared.services.retrieval.workflow.synthesizer import compose_final_answer, synthesize_step
-from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult, WorkflowResult
+from shared.services.retrieval.workflow.plan_service import WorkflowPlanService
+from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection
+from shared.services.retrieval.workflow.runtime_config import WorkflowRuntimeConfig
+from shared.services.retrieval.workflow.step_runner import WorkflowStepRunner
+from shared.services.retrieval.workflow.synthesizer import compose_final_answer
+from shared.services.retrieval.workflow.types import StepResult, WorkflowResult
from shared.services.retrieval.workflow.wallet import BudgetWallet
+DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]]
+
class WorkflowOrchestrator:
"""Plan and execute a query workflow DAG."""
- def __init__(self) -> None:
+ def __init__(self, db_factory: DbSessionFactory | None = None) -> None:
self.parent_run_id = f'wret_{uuid4().hex[:12]}'
+ self._db_factory = db_factory
+
+ def _get_db_factory(self) -> DbSessionFactory:
+ if self._db_factory is not None:
+ return self._db_factory
+
+ from shared.core.database import get_db_context
+
+ return get_db_context
async def run(
self,
@@ -52,18 +60,14 @@ async def run(
llm_fn=None,
) -> WorkflowResult:
t0 = time.monotonic()
+ config = WorkflowRuntimeConfig.from_env()
llm_fn = llm_fn or create_retrieval_llm_fn()
planner_llm = create_retrieval_planner_fn(thinking=True)
- planner_budget = _env_int('RETRIEVAL_PLANNER_THINKING_BUDGET', 4000)
- wallet_total = _env_int('RETRIEVAL_WALLET_TOTAL_BUDGET', 200000)
- per_retrieve = _env_int('RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET', 40000)
- per_synthesize = _env_int('RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET', 6000)
- max_steps = _env_int('RETRIEVAL_DECOMPOSITION_MAX_STEPS', 5)
planner_ledger = BudgetLedger(
- total=planner_budget,
+ total=config.planner_budget,
planning_ratio=0.0,
- bootstrap=planner_budget,
+ bootstrap=config.planner_budget,
per_doc_min_share=0,
)
total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory(
@@ -74,33 +78,36 @@ async def run(
)
planner_ledger.total_chunks = total_chunks
planner_ledger.total_docs = total_docs
- plan = await self._load_or_plan(
+ plan = await WorkflowPlanService().load_or_create(
user_id=user_id,
namespace=namespace,
query=query,
planner_llm=planner_llm,
planner_ledger=planner_ledger,
- max_steps=max_steps,
- wallet_total=wallet_total,
- per_retrieve=per_retrieve,
+ max_steps=config.max_steps,
+ wallet_total=config.wallet_total_budget,
+ per_retrieve=config.per_retrieve_step_budget,
kb_total_docs=total_docs,
kb_total_chunks=total_chunks,
)
wallet = BudgetWallet(
- total=wallet_total,
- per_retrieve_step_default=per_retrieve,
- per_synthesize_step_default=per_synthesize,
+ total=config.wallet_total_budget,
+ per_retrieve_step_default=config.per_retrieve_step_budget,
+ per_synthesize_step_default=config.per_synthesize_step_budget,
)
ledgers = await wallet.allocate(plan)
results_by_id: dict[str, StepResult] = {}
- sem = asyncio.Semaphore(_env_int('RETRIEVAL_WORKFLOW_PARALLEL_MAX', 3))
+ sem = asyncio.Semaphore(config.parallel_max)
+ step_runner = WorkflowStepRunner(
+ db_factory=self._get_db_factory(),
+ parent_run_id=self.parent_run_id,
+ )
for batch in plan.topological_batches():
await asyncio.gather(
*[
- self._run_step(
- db,
+ step_runner.run_step(
step=step,
ledger=ledgers[step.id],
results_by_id=results_by_id,
@@ -125,10 +132,11 @@ async def run(
answer_text = compose_final_answer(plan, results_by_id)
ordered_results = [results_by_id[step.id] for step in plan.steps if step.id in results_by_id]
- referenced_chunks = _dedupe_references(
+ reference_projection = WorkflowReferenceProjection()
+ referenced_chunks = reference_projection.dedupe(
ref for step_result in ordered_results for ref in step_result.referenced_chunks
)
- api_results = _references_to_results(referenced_chunks)
+ api_results = reference_projection.to_api_results(referenced_chunks)
elapsed_ms = int((time.monotonic() - t0) * 1000)
logger.info(
'workflow retrieval DONE: steps={} refs={} answer_chars={} elapsed={}ms',
@@ -151,256 +159,3 @@ async def run(
planner_snapshot=planner_ledger.snapshot(),
parent_run_id=self.parent_run_id,
)
-
- async def _load_or_plan(
- self,
- *,
- user_id: str,
- namespace: str,
- query: str,
- planner_llm,
- planner_ledger: BudgetLedger,
- max_steps: int,
- wallet_total: int,
- per_retrieve: int,
- kb_total_docs: int,
- kb_total_chunks: int,
- ) -> QueryPlan:
- try:
- cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query)
- if cached:
- return QueryPlan.from_dict(cached, original_query=query)
- except Exception as exc:
- logger.warning(f'workflow plan cache read failed (ignored): {exc}')
-
- planner = QueryPlanner(
- llm_fn=planner_llm,
- planner_ledger=planner_ledger,
- max_steps=max_steps,
- total_budget=wallet_total,
- per_step_budget=per_retrieve,
- )
- plan = await planner.plan(
- query=query,
- kb_total_docs=kb_total_docs,
- kb_total_chunks=kb_total_chunks,
- )
- try:
- await set_cached_workflow_plan(
- user_id=user_id,
- namespace=namespace,
- query=query,
- plan=plan.to_dict(),
- )
- except Exception as exc:
- logger.warning(f'workflow plan cache write failed (ignored): {exc}')
- return plan
-
- async def _run_step(
- self,
- db: AsyncSession,
- *,
- step: PlannedStep,
- ledger: BudgetLedger,
- results_by_id: dict[str, StepResult],
- semaphore: asyncio.Semaphore,
- user_id: str,
- namespace: str,
- top_k: int,
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
- data_type: int,
- signal_paths: list[str] | None,
- filter_mode: str,
- channels: list[str] | None,
- channel_weights: dict[str, float] | None,
- llm_fn,
- ) -> None:
- async with semaphore:
- if step.step_kind == 'synthesize':
- await self._run_synthesize_step(step, ledger, results_by_id, llm_fn)
- return
- await self._run_retrieve_step(
- db,
- step=step,
- ledger=ledger,
- results_by_id=results_by_id,
- user_id=user_id,
- namespace=namespace,
- top_k=top_k,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- data_type=data_type,
- signal_paths=signal_paths,
- filter_mode=filter_mode,
- channels=channels,
- channel_weights=channel_weights,
- llm_fn=llm_fn,
- )
-
- async def _run_retrieve_step(
- self,
- db: AsyncSession,
- *,
- step: PlannedStep,
- ledger: BudgetLedger,
- results_by_id: dict[str, StepResult],
- user_id: str,
- namespace: str,
- top_k: int,
- exclude_document_ids: list[str],
- exclude_sections: list[dict[str, str]],
- data_type: int,
- signal_paths: list[str] | None,
- filter_mode: str,
- channels: list[str] | None,
- channel_weights: dict[str, float] | None,
- llm_fn,
- ) -> None:
- try:
- # AsyncSession is not safe for concurrent use. Workflow steps may
- # run in the same topological batch, so each retrieve step opens an
- # isolated session and leaves the parent session untouched.
- async with get_db_context() as step_db:
- agentic_result = await RetrievalAgent().run(
- step_db,
- user_id=user_id,
- namespace=namespace,
- query=step.sub_query,
- top_k=top_k,
- llm_fn=llm_fn,
- exclude_document_ids=exclude_document_ids,
- exclude_sections=exclude_sections,
- data_type=data_type,
- signal_paths=signal_paths,
- filter_mode=filter_mode,
- channels=channels,
- channel_weights=channel_weights,
- ledger=ledger,
- parent_run_id=self.parent_run_id,
- workflow_step_id=step.id,
- )
- results_by_id[step.id] = _step_result_from_agentic(step, agentic_result)
- except Exception as exc:
- logger.exception(f'workflow retrieve step failed: step_id={step.id}')
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status='error',
- error=str(exc),
- budget_snapshot=ledger.snapshot(),
- )
-
- async def _run_synthesize_step(
- self,
- step: PlannedStep,
- ledger: BudgetLedger,
- results_by_id: dict[str, StepResult],
- llm_fn,
- ) -> None:
- if llm_fn is None:
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status='skipped',
- answer_text='',
- error='llm unavailable for synthesis',
- budget_snapshot=ledger.snapshot(),
- )
- return
- prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id}
- try:
- answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger)
- refs = _dedupe_references(
- ref for result in prior.values() for ref in result.referenced_chunks
- )
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status='done',
- answer_text=answer,
- referenced_chunks=refs,
- budget_snapshot=ledger.snapshot(),
- )
- except Exception as exc:
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status='budget_stop' if 'budget' in str(exc).lower() else 'error',
- answer_text='(budget exhausted)' if 'budget' in str(exc).lower() else '',
- error=str(exc),
- budget_snapshot=ledger.snapshot(),
- )
-
-
-def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult:
- if result.answer_text:
- status = 'done'
- elif result.failure_reason:
- status = 'not_found'
- elif 'budget' in (result.stop_reason or ''):
- status = 'budget_stop'
- else:
- status = 'done'
- return StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status=status, # type: ignore[arg-type]
- answer_text=result.answer_text,
- evidence_text=result.evidence_text,
- referenced_chunks=result.referenced_chunks,
- budget_snapshot=result.budget_snapshot,
- router_used=result.router_used,
- stop_reason=result.stop_reason,
- failure_reason=result.failure_reason,
- )
-
-
-def _dedupe_references(refs) -> list[dict[str, Any]]:
- seen: set[str] = set()
- out: list[dict[str, Any]] = []
- for ref in refs:
- chunk_id = str(ref.get('chunk_id') or '')
- key = chunk_id or str(ref)
- if key in seen:
- continue
- seen.add(key)
- out.append(dict(ref))
- return out
-
-
-def _references_to_results(refs: list[dict[str, Any]]) -> list[dict[str, Any]]:
- return [
- {
- 'chunk_id': ref.get('chunk_id'),
- 'document_id': ref.get('document_id'),
- 'chunk_type': ref.get('chunk_type'),
- 'source': {
- 'document_id': ref.get('document_id'),
- 'section_path': ref.get('section_path'),
- },
- }
- for ref in refs
- ]
-
-
-def _env_int(name: str, default: int) -> int:
- try:
- return int(os.environ.get(name, str(default)))
- except (TypeError, ValueError):
- return default
diff --git a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py
new file mode 100644
index 000000000..a5e3a0d43
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py
@@ -0,0 +1,59 @@
+"""Plan loading and creation for decomposed retrieval workflows."""
+from __future__ import annotations
+
+from loguru import logger
+
+from shared.services.retrieval.agentic.budget import BudgetLedger
+from shared.services.retrieval.cache_service import (
+ get_cached_workflow_plan,
+ set_cached_workflow_plan,
+)
+from shared.services.retrieval.llm_adapter import LLMFn
+from shared.services.retrieval.workflow.planner import QueryPlanner
+from shared.services.retrieval.workflow.types import QueryPlan
+
+
+class WorkflowPlanService:
+ async def load_or_create(
+ self,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ planner_llm: LLMFn | None,
+ planner_ledger: BudgetLedger,
+ max_steps: int,
+ wallet_total: int,
+ per_retrieve: int,
+ kb_total_docs: int,
+ kb_total_chunks: int,
+ ) -> QueryPlan:
+ try:
+ cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query)
+ if cached:
+ return QueryPlan.from_dict(cached, original_query=query)
+ except Exception as exc:
+ logger.warning(f"workflow plan cache read failed (ignored): {exc}")
+
+ planner = QueryPlanner(
+ llm_fn=planner_llm,
+ planner_ledger=planner_ledger,
+ max_steps=max_steps,
+ total_budget=wallet_total,
+ per_step_budget=per_retrieve,
+ )
+ plan = await planner.plan(
+ query=query,
+ kb_total_docs=kb_total_docs,
+ kb_total_chunks=kb_total_chunks,
+ )
+ try:
+ await set_cached_workflow_plan(
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ plan=plan.to_dict(),
+ )
+ except Exception as exc:
+ logger.warning(f"workflow plan cache write failed (ignored): {exc}")
+ return plan
diff --git a/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py b/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py
new file mode 100644
index 000000000..992085c00
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/reference_projection.py
@@ -0,0 +1,40 @@
+"""Reference projection for decomposed retrieval workflows."""
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+
+class WorkflowReferenceProjection:
+ def dedupe(self, refs: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
+ seen: set[str] = set()
+ out: list[dict[str, Any]] = []
+ for ref in refs:
+ document_id = str(ref.get("document_id") or "").strip()
+ chunk_id = str(ref.get("chunk_id") or "").strip()
+ section_path = str(ref.get("section_path") or "").strip()
+ file_path = str(ref.get("file_path") or "").strip()
+ key = (
+ f"{document_id}:{chunk_id}:{section_path}:{file_path}"
+ if document_id and chunk_id
+ else str(ref)
+ )
+ if key in seen:
+ continue
+ seen.add(key)
+ out.append(dict(ref))
+ return out
+
+ def to_api_results(self, refs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ return [
+ {
+ "chunk_id": ref.get("chunk_id"),
+ "document_id": ref.get("document_id"),
+ "chunk_type": ref.get("chunk_type"),
+ "source": {
+ "document_id": ref.get("document_id"),
+ "section_path": ref.get("section_path"),
+ },
+ }
+ for ref in refs
+ ]
diff --git a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py
new file mode 100644
index 000000000..bcb6ab262
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py
@@ -0,0 +1,33 @@
+"""Runtime configuration for decomposed retrieval workflows."""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class WorkflowRuntimeConfig:
+ planner_budget: int = 4000
+ wallet_total_budget: int = 200000
+ per_retrieve_step_budget: int = 40000
+ per_synthesize_step_budget: int = 6000
+ max_steps: int = 5
+ parallel_max: int = 3
+
+ @classmethod
+ def from_env(cls) -> "WorkflowRuntimeConfig":
+ return cls(
+ planner_budget=_env_int("RETRIEVAL_PLANNER_THINKING_BUDGET", 4000),
+ wallet_total_budget=_env_int("RETRIEVAL_WALLET_TOTAL_BUDGET", 200000),
+ per_retrieve_step_budget=_env_int("RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET", 40000),
+ per_synthesize_step_budget=_env_int("RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET", 6000),
+ max_steps=_env_int("RETRIEVAL_DECOMPOSITION_MAX_STEPS", 5),
+ parallel_max=_env_int("RETRIEVAL_WORKFLOW_PARALLEL_MAX", 3),
+ )
+
+
+def _env_int(name: str, default: int) -> int:
+ try:
+ return int(os.environ.get(name, str(default)))
+ except (TypeError, ValueError):
+ return default
diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
new file mode 100644
index 000000000..8a6be3d9d
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
@@ -0,0 +1,192 @@
+"""Step execution for decomposed retrieval workflows."""
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from contextlib import AbstractAsyncContextManager
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.services.retrieval.agentic.budget import BudgetLedger
+from shared.services.retrieval.agentic.orchestrator import RetrievalAgent
+from shared.services.retrieval.agentic.types import AgenticResult
+from shared.services.retrieval.llm_adapter import LLMFn
+from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection
+from shared.services.retrieval.workflow.synthesizer import synthesize_step
+from shared.services.retrieval.workflow.types import PlannedStep, StepResult, StepStatus
+
+DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]]
+
+
+class WorkflowStepRunner:
+ def __init__(self, *, db_factory: DbSessionFactory, parent_run_id: str) -> None:
+ self._db_factory = db_factory
+ self._parent_run_id = parent_run_id
+ self._references = WorkflowReferenceProjection()
+
+ async def run_step(
+ self,
+ *,
+ step: PlannedStep,
+ ledger: BudgetLedger,
+ results_by_id: dict[str, StepResult],
+ semaphore: asyncio.Semaphore,
+ user_id: str,
+ namespace: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int,
+ signal_paths: list[str] | None,
+ filter_mode: str,
+ channels: list[str] | None,
+ channel_weights: dict[str, float] | None,
+ llm_fn: LLMFn | None,
+ ) -> None:
+ async with semaphore:
+ if step.step_kind == "synthesize":
+ await self._run_synthesize_step(step, ledger, results_by_id, llm_fn)
+ return
+ await self._run_retrieve_step(
+ step=step,
+ ledger=ledger,
+ results_by_id=results_by_id,
+ user_id=user_id,
+ namespace=namespace,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ llm_fn=llm_fn,
+ )
+
+ async def _run_retrieve_step(
+ self,
+ *,
+ step: PlannedStep,
+ ledger: BudgetLedger,
+ results_by_id: dict[str, StepResult],
+ user_id: str,
+ namespace: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int,
+ signal_paths: list[str] | None,
+ filter_mode: str,
+ channels: list[str] | None,
+ channel_weights: dict[str, float] | None,
+ llm_fn: LLMFn | None,
+ ) -> None:
+ try:
+ async with self._db_factory() as step_db:
+ agentic_result = await RetrievalAgent().run(
+ step_db,
+ user_id=user_id,
+ namespace=namespace,
+ query=step.sub_query,
+ top_k=top_k,
+ llm_fn=llm_fn,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ ledger=ledger,
+ parent_run_id=self._parent_run_id,
+ workflow_step_id=step.id,
+ )
+ results_by_id[step.id] = _step_result_from_agentic(step, agentic_result)
+ except Exception as exc:
+ logger.exception(f"workflow retrieve step failed: step_id={step.id}")
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status="error",
+ error=str(exc),
+ budget_snapshot=ledger.snapshot(),
+ )
+
+ async def _run_synthesize_step(
+ self,
+ step: PlannedStep,
+ ledger: BudgetLedger,
+ results_by_id: dict[str, StepResult],
+ llm_fn: LLMFn | None,
+ ) -> None:
+ if llm_fn is None:
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status="skipped",
+ answer_text="",
+ error="llm unavailable for synthesis",
+ budget_snapshot=ledger.snapshot(),
+ )
+ return
+ prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id}
+ try:
+ answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger)
+ refs = self._references.dedupe(ref for result in prior.values() for ref in result.referenced_chunks)
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status="done",
+ answer_text=answer,
+ referenced_chunks=refs,
+ budget_snapshot=ledger.snapshot(),
+ )
+ except Exception as exc:
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status="budget_stop" if "budget" in str(exc).lower() else "error",
+ answer_text="(budget exhausted)" if "budget" in str(exc).lower() else "",
+ error=str(exc),
+ budget_snapshot=ledger.snapshot(),
+ )
+
+
+def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult:
+ if result.answer_text:
+ status: StepStatus = "done"
+ elif result.failure_reason:
+ status = "not_found"
+ elif "budget" in (result.stop_reason or ""):
+ status = "budget_stop"
+ else:
+ status = "done"
+ return StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status=status,
+ answer_text=result.answer_text,
+ evidence_text=result.evidence_text,
+ referenced_chunks=result.referenced_chunks,
+ budget_snapshot=result.budget_snapshot,
+ router_used=result.router_used,
+ stop_reason=result.stop_reason,
+ failure_reason=result.failure_reason,
+ )
diff --git a/packages/shared-python/shared/services/storage/file_encryptor_service.py b/packages/shared-python/shared/services/storage/file_encryptor_service.py
deleted file mode 100755
index e542a705a..000000000
--- a/packages/shared-python/shared/services/storage/file_encryptor_service.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import os
-import pickle
-from typing import Any
-
-from cryptography.fernet import Fernet
-
-
-class FernetPickleEncryptor:
- encrypt = False
-
- def __init__(self, key: bytes = b"nc1BPZSkNb7Oc82_Wo3QoZTmJCEnQtpKZ2n-Z5F4CwY="):
- self.cipher = Fernet(key)
-
- def save_to_file(self, data: Any, file_path: str) -> None:
- serialized_data = pickle.dumps(data) # Serialize the input payload.
- encrypted_data = self.cipher.encrypt(
- serialized_data
- ) # Encrypt the serialized bytes.
- with open(file_path, "wb") as f:
- f.write(encrypted_data)
-
- def load_from_file(self, file_path: str) -> Any:
- if not os.path.exists(file_path):
- raise FileNotFoundError(f"File {file_path} does not exist.")
- with open(file_path, "rb") as f:
- encrypted_data = f.read()
- decrypted_data = self.cipher.decrypt(encrypted_data) # Decrypt the file bytes.
- loaded_data = pickle.loads(decrypted_data) # Deserialize the decrypted payload.
- return loaded_data
-
-
-encryptor = FernetPickleEncryptor()
-
-if __name__ == "__main__":
- # 2. Encrypt the payload.
- data = {"key": "value"}
- # 3. Save the encrypted payload to a file.
- file_path = "data.pkl"
- encryptor.save_to_file(data, file_path)
- print(f"Encrypted data saved to file: {file_path}")
-
- # 4. Load the encrypted payload back from disk.
- decrypted_data = encryptor.load_from_file(file_path)
- print("Decrypted data loaded from file:", decrypted_data)
diff --git a/packages/shared-python/shared/services/storage/file_upload_service.py b/packages/shared-python/shared/services/storage/file_upload_service.py
index c60f3c557..8b24793a4 100644
--- a/packages/shared-python/shared/services/storage/file_upload_service.py
+++ b/packages/shared-python/shared/services/storage/file_upload_service.py
@@ -1,147 +1,30 @@
-"""Storage upload service."""
+"""Async adapter for shared Job file storage."""
import asyncio
-import json
-import os
-from typing import Any, Dict, Optional
+from typing import Any, Optional
from loguru import logger
-from shared.core.config import settings
-from shared.core.exceptions.domain_exceptions import (
- KnowhereException,
- StorageServiceException,
-)
-from shared.utils.pinned_outbound_http import (
- download_pinned_outbound_file_async,
-)
-from shared.utils.url_security import validate_http_url_and_resolve_ip_async
+from shared.core.exceptions.domain_exceptions import StorageServiceException
+from shared.services.storage.job_file_storage import JobFileStorage
class FileUploadService:
- """File upload service supporting S3, OSS, and MinIO."""
+ """Async adapter over the shared Job file storage module."""
- def __init__(self):
- self.adapter = settings.get_storage_adapter()
- self.uploads_bucket = settings.S3_BUCKET_NAME
- self.results_bucket = getattr(
- settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME
- )
-
- async def handle_direct_upload(self, file_path: str, job_id: str) -> str:
- """
- Handle a direct file upload.
-
- Args:
- file_path: Local file path.
- job_id: Job ID.
-
- Returns:
- str: Storage key.
- """
- try:
- # Build the storage key.
- file_extension = os.path.splitext(file_path)[1]
- s3_key = f"uploads/{job_id}{file_extension}"
-
- # Upload the file.
- await self._upload_to_s3(file_path, s3_key, self.uploads_bucket)
-
- logger.info(f"Direct file upload succeeded: {file_path} -> {s3_key}")
- return s3_key
-
- except KnowhereException:
- raise
- except Exception as e:
- logger.error(f"Direct file upload failed: {e}")
- raise StorageServiceException(
- internal_message=f"Direct file upload failed: {str(e)}",
- operation="direct_upload",
- original_exception=e,
- )
-
- async def handle_url_upload(self, file_url: str, job_id: str) -> str:
- """
- Handle a URL-based upload flow.
-
- Args:
- file_url: File URL.
- job_id: Job ID.
-
- Returns:
- str: Storage key.
- """
- try:
- # Download the file into a temporary location first.
- temp_file_path = await self._download_file_from_url(file_url)
-
- try:
- # Build the storage key.
- file_extension = os.path.splitext(file_url.split("?")[0])[1]
- s3_key = f"uploads/{job_id}{file_extension}"
-
- # Upload the downloaded file.
- await self._upload_to_s3(temp_file_path, s3_key, self.uploads_bucket)
-
- logger.info(
- f"URL file download and upload succeeded: {file_url} -> {s3_key}"
- )
- return s3_key
-
- finally:
- # Clean up the temporary file.
- if os.path.exists(temp_file_path):
- os.remove(temp_file_path)
-
- except KnowhereException:
- raise
- except Exception as e:
- logger.error(f"URL file handling failed: {e}")
- raise StorageServiceException(
- internal_message=f"URL file handling failed: {str(e)}",
- operation="url_upload",
- original_exception=e,
- )
+ def __init__(self, *, storage: JobFileStorage | None = None) -> None:
+ self._storage = storage or JobFileStorage()
async def generate_upload_url(
self, job_id: str, file_extension: str = ""
- ) -> Dict[str, Any]:
- """
- Generate a presigned upload URL.
-
- Args:
- job_id: Job ID.
- file_extension: File extension.
-
- Returns:
- Dict: Upload URL payload including the storage key.
- """
+ ) -> dict[str, Any]:
try:
- s3_key = f"uploads/{job_id}{file_extension}"
-
- # Infer a Content-Type from the file extension.
- content_type = self.get_content_type(file_extension)
-
- # Use the job waiting expiry as the upload URL TTL.
- upload_url = self.adapter.generate_presigned_url(
- s3_key,
- expiration=settings.JOB_WAITING_EXPIRE_SECONDS,
- bucket=self.uploads_bucket,
- method="PUT",
- headers={"Content-Type": content_type},
+ return await asyncio.to_thread(
+ self._storage.generate_upload_url,
+ job_id=job_id,
+ file_extension=file_extension,
)
- logger.info(f"Generated presigned upload URL: {upload_url}")
-
- return {
- "upload_url": upload_url,
- "s3_key": s3_key,
- "expires_in": settings.JOB_WAITING_EXPIRE_SECONDS,
- "upload_headers": {"Content-Type": content_type},
- }
-
- except KnowhereException:
- raise
except Exception as e:
logger.error(f"Failed to generate upload URL: {e}")
raise StorageServiceException(
@@ -152,29 +35,15 @@ async def generate_upload_url(
async def generate_download_url(
self, s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600
- ) -> Dict[str, Any]:
- """
- Generate a presigned download URL.
-
- Args:
- s3_key: Storage key.
- bucket: Optional bucket name.
-
- Returns:
- str: Download URL.
- """
+ ) -> dict[str, Any]:
try:
- bucket_name = bucket or self.results_bucket
-
- # Generate a one-hour presigned URL by default.
- download_url = self.adapter.generate_presigned_url(
- s3_key, expiration=expires_in, bucket=bucket_name, method="GET"
+ return await asyncio.to_thread(
+ self._storage.generate_download_url,
+ s3_key,
+ bucket=bucket or self._storage.results_bucket,
+ expires_in=expires_in,
)
- return {"download_url": download_url, "expires_in": expires_in}
-
- except KnowhereException:
- raise
except Exception as e:
logger.error(f"Failed to generate download URL: {e}")
raise StorageServiceException(
@@ -183,411 +52,19 @@ async def generate_download_url(
original_exception=e,
)
- async def get_file_info(
+ async def verify_s3_file_exists(
self, s3_key: str, bucket: Optional[str] = None
- ) -> Optional[Dict[str, Any]]:
- """
- Get file information.
-
- Args:
- s3_key: Storage key.
- bucket: Optional bucket name.
-
- Returns:
- Dict: File metadata.
- """
- try:
- bucket_name = bucket or self.results_bucket
-
- # Check existence and load the object size.
- if not self.adapter.exists(s3_key, bucket_name):
- return None
-
- size = self.adapter.get_object_size(s3_key, bucket_name)
- return {
- "size": size,
- "content_type": None, # The adapter interface does not expose content_type yet.
- "last_modified": None,
- "etag": None,
- }
-
- except Exception as e:
- # Treat not-found responses as a missing object.
- if "404" in str(e) or "not found" in str(e).lower():
- return None
- logger.error(f"Failed to get file info: {e}")
- raise StorageServiceException(
- internal_message=f"Failed to get file info: {str(e)}",
- operation="get_file_info",
- original_exception=e,
- )
-
- async def upload_result_file(
- self, local_file_path: str, job_id: str, file_extension: str = ""
- ) -> str:
- """
- Upload a result file.
-
- Args:
- local_file_path: Local file path.
- job_id: Job ID.
- file_extension: File extension.
-
- Returns:
- str: Storage key.
- """
- try:
- s3_key = f"results/{job_id}{file_extension}"
- await self._upload_to_s3(local_file_path, s3_key, self.results_bucket)
-
- logger.info(f"Result file upload succeeded: {local_file_path} -> {s3_key}")
- return s3_key
-
- except KnowhereException:
- raise
- except Exception as e:
- logger.error(f"Result file upload failed: {e}")
- raise StorageServiceException(
- internal_message=f"Result file upload failed: {str(e)}",
- operation="upload_result_file",
- original_exception=e,
- )
-
- async def upload_json_result(
- self,
- job_id: str,
- result_data: Dict[str, Any],
- *,
- content_type: str = "application/json",
- ) -> str:
- """Upload a JSON result file; deprecated but kept for compatibility."""
+ ) -> dict[str, Any]:
try:
- s3_key = f"results/{job_id}.json"
- from io import BytesIO
-
- body = json.dumps(result_data, ensure_ascii=False).encode("utf-8")
- self.adapter.upload_fileobj(
- BytesIO(body),
+ return await asyncio.to_thread(
+ self._storage.verify_exists,
s3_key,
- bucket=self.results_bucket,
- content_type=content_type,
+ bucket=bucket or self._storage.uploads_bucket,
)
- logger.info(f"Result JSON upload succeeded: job_id={job_id}, key={s3_key}")
- return s3_key
- except KnowhereException:
- raise
except Exception as e:
- logger.error(f"Failed to upload result JSON: {e}")
- raise StorageServiceException(
- internal_message=f"Failed to upload result JSON: {str(e)}",
- operation="upload_json_result",
- original_exception=e,
- )
-
- async def upload_zip_result(
- self,
- job_id: str,
- zip_file_path: str,
- ) -> str:
- """Upload a ZIP result file."""
- try:
- s3_key = f"results/{job_id}.zip"
- await self._upload_to_s3(zip_file_path, s3_key, self.results_bucket)
- logger.info(f"Result ZIP upload succeeded: job_id={job_id}, key={s3_key}")
-
- # Clean up the temporary ZIP after upload.
- try:
- if os.path.exists(zip_file_path):
- os.remove(zip_file_path)
- except Exception as e:
- logger.warning(f"Failed to clean up temporary ZIP file: {e}")
-
- return s3_key
- except KnowhereException:
- raise
- except Exception as e:
- logger.error(f"Failed to upload result ZIP: {e}")
- raise StorageServiceException(
- internal_message=f"Failed to upload result ZIP: {str(e)}",
- operation="upload_zip_result",
- original_exception=e,
- )
-
- def _ensure_bucket_exists(self, bucket_name: str) -> bool:
- """
- Ensure the bucket is accessible.
-
- Args:
- bucket_name: Bucket name.
-
- Returns:
- bool: Whether the bucket check succeeded.
- """
- try:
- # In adapter mode, probe accessibility by listing objects.
- adapter = settings.get_storage_adapter()
- list(adapter.list_objects(prefix="", bucket=bucket_name))
- logger.debug(f"Bucket {bucket_name} is accessible")
- return True
- except Exception as e:
- # The bucket is missing or inaccessible.
- # For OSS, buckets should already exist; only accessibility is checked here.
- logger.warning(
- f"Bucket {bucket_name} may not exist or may be inaccessible: {e}"
- )
- # In production, buckets should already be provisioned, so continue.
- # Return False here instead if strict enforcement is ever needed.
- return True
-
- async def _ensure_bucket_exists_async(self, bucket_name: str) -> bool:
- """
- Asynchronously ensure the bucket is accessible.
-
- Args:
- bucket_name: Bucket name.
-
- Returns:
- bool: Whether the bucket check succeeded.
- """
-
- def _check_and_create():
- try:
- # In adapter mode, probe accessibility by listing objects.
- adapter = settings.get_storage_adapter()
- list(adapter.list_objects(prefix="", bucket=bucket_name))
- logger.debug(f"Bucket {bucket_name} is accessible")
- return True
- except Exception as e:
- # The bucket is missing or inaccessible.
- logger.warning(
- f"Bucket {bucket_name} may not exist or may be inaccessible: {e}"
- )
- # In production, buckets should already be provisioned, so continue.
- return True
-
- # Run the synchronous probe in a thread pool.
- loop = asyncio.get_event_loop()
- return await loop.run_in_executor(None, _check_and_create)
-
- async def _upload_to_s3(self, local_file_path: str, s3_key: str, bucket: str):
- """Upload a file to storage."""
- # Ensure the bucket is accessible before uploading.
- if not await self._ensure_bucket_exists_async(bucket):
- raise StorageServiceException(
- internal_message=f"Could not ensure bucket {bucket} exists",
- operation="ensure_bucket",
- )
-
- def _upload():
- self.adapter.upload_file(local_file_path, s3_key, bucket)
-
- # Run the blocking upload in a thread pool.
- loop = asyncio.get_event_loop()
- await loop.run_in_executor(None, _upload)
-
- async def download_from_s3(self, s3_key: str, bucket: Optional[str] = None) -> str:
- """Download a file from storage into a local temporary directory."""
- import uuid
-
- if bucket is None:
- bucket = settings.S3_BUCKET_NAME
-
- # Create the temporary destination directory.
- temp_dir = getattr(settings, "TMP_PATH", "/tmp")
- os.makedirs(temp_dir, exist_ok=True)
-
- # Generate a temporary filename while preserving the original extension.
- file_extension = os.path.splitext(s3_key)[1]
- temp_filename = f"temp_{uuid.uuid4().hex}{file_extension}"
- temp_file_path = os.path.join(temp_dir, temp_filename)
-
- try:
- # Use the adapter to download the file.
- def _download():
- self.adapter.download_file(s3_key, temp_file_path, bucket)
-
- # Run the blocking download in the event loop executor.
- loop = asyncio.get_event_loop()
- await loop.run_in_executor(None, _download)
-
- return temp_file_path
-
- except KnowhereException:
- if os.path.exists(temp_file_path):
- os.remove(temp_file_path)
- raise
- except Exception as e:
- # Clean up the temporary file on failure.
- if os.path.exists(temp_file_path):
- os.remove(temp_file_path)
- raise StorageServiceException(
- internal_message=f"Failed to download file from S3: {str(e)}",
- operation="download_from_s3",
- original_exception=e,
- )
-
- async def _download_file_from_url(self, file_url: str) -> str:
- """Download a file from a URL into a temporary directory."""
- temp_file_path = ""
- try:
- validation = await validate_http_url_and_resolve_ip_async(file_url)
- if not validation.is_valid or not validation.validated_ip:
- raise StorageServiceException(
- internal_message=f"Invalid URL: {validation.error_message}",
- operation="download_from_url",
- )
-
- temp_dir = getattr(settings, "TMP_PATH", "/tmp")
- os.makedirs(temp_dir, exist_ok=True)
- download_result = await download_pinned_outbound_file_async(
- url=validation.url,
- pinned_ip=validation.validated_ip,
- timeout_seconds=300,
- user_agent="Knowhere-FileDownloader/1.0",
- temp_dir=temp_dir,
- )
- temp_file_path = download_result.temp_file_path
- return temp_file_path
-
- except KnowhereException:
- raise
- except Exception as e:
- if os.path.exists(temp_file_path):
- os.remove(temp_file_path)
- raise StorageServiceException(
- internal_message=f"Failed to download file: {str(e)}",
- operation="download_from_url",
- original_exception=e,
- )
-
- async def verify_s3_file_exists(
- self, s3_key: str, bucket: Optional[str] = None
- ) -> Dict[str, Any]:
- """
- Verify whether a file exists in storage.
-
- Args:
- s3_key: Storage key.
- bucket: Optional bucket name.
-
- Returns:
- Dict: File info payload, or `{"exists": False}` when missing.
- """
- try:
- bucket_name = bucket or self.uploads_bucket
-
- # Use the adapter to check object existence.
- exists = self.adapter.exists(s3_key, bucket_name)
- if not exists:
- return {"exists": False}
-
- size = self.adapter.get_object_size(s3_key, bucket_name)
- return {
- "exists": True,
- "size": size,
- "content_type": None,
- "last_modified": None,
- "etag": None,
- }
-
- except Exception as e:
- # Treat not-found responses as a missing object.
- if "404" in str(e) or "not found" in str(e).lower():
- return {"exists": False}
logger.error(f"Failed to verify file existence: {e}")
raise StorageServiceException(
internal_message=f"Failed to verify file existence: {str(e)}",
operation="verify_s3_file_exists",
original_exception=e,
)
-
- def get_content_type(self, file_extension: str) -> str:
- """
- Return a Content-Type for a file extension.
-
- Args:
- file_extension: File extension, such as `.pdf` or `.docx`.
-
- Returns:
- str: Content-Type
- """
- content_types = {
- ".pdf": "application/pdf",
- ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
- ".doc": "application/msword",
- ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
- ".xls": "application/vnd.ms-excel",
- ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
- ".ppt": "application/vnd.ms-powerpoint",
- ".csv": "text/csv",
- ".txt": "text/plain",
- ".md": "text/markdown",
- ".json": "application/json",
- ".xml": "application/xml",
- ".html": "text/html",
- ".htm": "text/html",
- ".png": "image/png",
- ".jpg": "image/jpeg",
- ".jpeg": "image/jpeg",
- ".gif": "image/gif",
- ".bmp": "image/bmp",
- ".tiff": "image/tiff",
- ".svg": "image/svg+xml",
- ".zip": "application/zip",
- ".rar": "application/x-rar-compressed",
- ".7z": "application/x-7z-compressed",
- ".tar": "application/x-tar",
- ".gz": "application/gzip",
- }
- return content_types.get(file_extension.lower(), "application/octet-stream")
-
- async def get_file_url(
- self, s3_key: str, bucket: Optional[str] = None, expires_in: int = 3600
- ) -> str:
- """
- Get a file URL from a storage key.
-
- Args:
- s3_key: Storage key.
- bucket: Optional bucket name.
- expires_in: URL TTL in seconds, defaulting to one hour.
-
- Returns:
- str: File URL.
- """
- try:
- bucket_name = bucket or self.uploads_bucket
-
- # Generate a presigned GET URL.
- file_url = self.adapter.generate_presigned_url(
- s3_key, expiration=expires_in, bucket=bucket_name, method="GET"
- )
-
- logger.info(f"Generated file URL successfully: {s3_key} -> {file_url}")
- return file_url
-
- except KnowhereException:
- raise
- except Exception as e:
- logger.error(f"Failed to get file URL: {e}")
- raise StorageServiceException(
- internal_message=f"Failed to get file URL: {str(e)}",
- operation="get_file_url",
- original_exception=e,
- )
-
- def generate_s3_key(
- self, job_id: str, file_extension: str = "", prefix: str = "uploads"
- ) -> str:
- """
- Generate a storage key.
-
- Args:
- job_id: Job ID.
- file_extension: File extension.
- prefix: Key prefix such as `uploads` or `results`.
-
- Returns:
- str: Storage key.
- """
- return f"{prefix}/{job_id}{file_extension}"
diff --git a/packages/shared-python/shared/services/storage/job_file_storage.py b/packages/shared-python/shared/services/storage/job_file_storage.py
new file mode 100644
index 000000000..cc46f6bd3
--- /dev/null
+++ b/packages/shared-python/shared/services/storage/job_file_storage.py
@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import os
+import tempfile
+from typing import Any, BinaryIO
+
+from shared.core.config import settings
+from shared.core.config.storage import get_cached_storage_adapter
+from shared.core.exceptions.domain_exceptions import StorageServiceException
+from shared.services.storage.storage_adapter import StorageAdapter
+from shared.utils.pinned_outbound_http import download_pinned_outbound_file
+from shared.utils.url_security import validate_http_url_and_resolve_ip
+
+
+class JobFileStorage:
+ """Own storage rules for Job source files and Job Result bundles."""
+
+ def __init__(
+ self,
+ *,
+ storage_adapter: StorageAdapter | None = None,
+ uploads_bucket: str | None = None,
+ results_bucket: str | None = None,
+ ) -> None:
+ self._storage_adapter = storage_adapter
+ self.uploads_bucket = uploads_bucket or settings.S3_BUCKET_NAME
+ self.results_bucket = results_bucket or getattr(
+ settings,
+ "S3_RESULTS_BUCKET",
+ settings.S3_BUCKET_NAME,
+ )
+
+ @property
+ def storage_adapter(self) -> StorageAdapter:
+ if self._storage_adapter is None:
+ self._storage_adapter = get_cached_storage_adapter()
+ return self._storage_adapter
+
+ def build_upload_key(self, *, job_id: str, file_extension: str = "") -> str:
+ return f"uploads/{job_id}{file_extension}"
+
+ def build_result_key(self, *, job_id: str, file_extension: str = "") -> str:
+ return f"results/{job_id}{file_extension}"
+
+ def build_result_zip_key(self, *, job_id: str) -> str:
+ return self.build_result_key(job_id=job_id, file_extension=".zip")
+
+ def build_result_raw_prefix(self, *, job_id: str) -> str:
+ return f"results/{job_id}/"
+
+ def generate_upload_url(
+ self,
+ *,
+ job_id: str,
+ file_extension: str = "",
+ ) -> dict[str, Any]:
+ storage_key = self.build_upload_key(
+ job_id=job_id,
+ file_extension=file_extension,
+ )
+ content_type = self.get_content_type(file_extension)
+ upload_url = self.storage_adapter.generate_presigned_url(
+ storage_key,
+ expiration=settings.JOB_WAITING_EXPIRE_SECONDS,
+ bucket=self.uploads_bucket,
+ method="PUT",
+ headers={"Content-Type": content_type},
+ )
+ return {
+ "upload_url": upload_url,
+ "s3_key": storage_key,
+ "expires_in": settings.JOB_WAITING_EXPIRE_SECONDS,
+ "upload_headers": {"Content-Type": content_type},
+ }
+
+ def generate_download_url(
+ self,
+ storage_key: str,
+ *,
+ bucket: str,
+ expires_in: int = 3600,
+ ) -> dict[str, Any]:
+ download_url = self.storage_adapter.generate_presigned_url(
+ storage_key,
+ expiration=expires_in,
+ bucket=bucket,
+ method="GET",
+ )
+ return {"download_url": download_url, "expires_in": expires_in}
+
+ def generate_upload_download_url(
+ self,
+ storage_key: str,
+ *,
+ expires_in: int = 3600,
+ ) -> dict[str, Any]:
+ return self.generate_download_url(
+ storage_key,
+ bucket=self.uploads_bucket,
+ expires_in=expires_in,
+ )
+
+ def verify_exists(
+ self,
+ storage_key: str,
+ *,
+ bucket: str,
+ ) -> dict[str, Any]:
+ try:
+ if not self.storage_adapter.exists(storage_key, bucket):
+ return {"exists": False}
+
+ size = self.storage_adapter.get_object_size(storage_key, bucket)
+ return {
+ "exists": True,
+ "size": size,
+ "content_type": None,
+ "last_modified": None,
+ "etag": None,
+ }
+ except Exception as exc:
+ if "404" in str(exc) or "not found" in str(exc).lower():
+ return {"exists": False}
+ raise StorageServiceException(
+ internal_message=f"Storage file verification failed: {exc}",
+ operation="verify_exists",
+ original_exception=exc,
+ ) from exc
+
+ def verify_upload_exists(self, storage_key: str) -> dict[str, Any]:
+ return self.verify_exists(storage_key, bucket=self.uploads_bucket)
+
+ def upload_local_file(
+ self,
+ local_file_path: str,
+ storage_key: str,
+ *,
+ bucket: str,
+ ) -> dict[str, Any]:
+ try:
+ return self.storage_adapter.upload_file(local_file_path, storage_key, bucket)
+ except Exception as exc:
+ raise StorageServiceException(
+ internal_message=f"Storage upload failed: {exc}",
+ operation="upload_local_file",
+ original_exception=exc,
+ ) from exc
+
+ def upload_source_file(
+ self,
+ local_file_path: str,
+ storage_key: str,
+ ) -> dict[str, Any]:
+ return self.upload_local_file(
+ local_file_path,
+ storage_key,
+ bucket=self.uploads_bucket,
+ )
+
+ def upload_fileobj(
+ self,
+ file_obj: BinaryIO,
+ storage_key: str,
+ *,
+ bucket: str,
+ content_type: str | None = None,
+ ) -> dict[str, Any]:
+ try:
+ return self.storage_adapter.upload_fileobj(
+ file_obj,
+ storage_key,
+ bucket=bucket,
+ content_type=content_type,
+ )
+ except Exception as exc:
+ raise StorageServiceException(
+ internal_message=f"Storage upload file object failed: {exc}",
+ operation="upload_fileobj",
+ original_exception=exc,
+ ) from exc
+
+ def download_to_path(
+ self,
+ storage_key: str,
+ local_path: str,
+ *,
+ bucket: str,
+ ) -> str:
+ try:
+ return self.storage_adapter.download_file(storage_key, local_path, bucket)
+ except Exception as exc:
+ raise StorageServiceException(
+ internal_message=f"Storage download failed: {exc}",
+ operation="download_to_path",
+ original_exception=exc,
+ ) from exc
+
+ def download_to_temp(
+ self,
+ storage_key: str,
+ *,
+ suffix: str,
+ temp_dir: str,
+ bucket: str,
+ ) -> str:
+ local_temp_path: str | None = None
+
+ try:
+ os.makedirs(temp_dir, exist_ok=True)
+ with tempfile.NamedTemporaryFile(
+ delete=False,
+ suffix=suffix,
+ dir=temp_dir,
+ ) as temp_file:
+ local_temp_path = temp_file.name
+
+ self.download_to_path(
+ storage_key,
+ local_temp_path,
+ bucket=bucket,
+ )
+ return local_temp_path
+ except Exception as exc:
+ if local_temp_path and os.path.exists(local_temp_path):
+ os.remove(local_temp_path)
+ raise StorageServiceException(
+ internal_message=(
+ "Failed to download object-storage file to temp path: "
+ f"storage_key={storage_key}, temp_dir={temp_dir}, error={exc}"
+ ),
+ operation="download_to_temp",
+ original_exception=exc,
+ ) from exc
+
+ def download_upload_to_temp(
+ self,
+ storage_key: str,
+ *,
+ suffix: str,
+ temp_dir: str,
+ ) -> str:
+ return self.download_to_temp(
+ storage_key,
+ suffix=suffix,
+ temp_dir=temp_dir,
+ bucket=self.uploads_bucket,
+ )
+
+ def download_file_from_url(
+ self,
+ file_url: str,
+ *,
+ temp_dir: str | None = None,
+ ) -> str:
+ temp_file_path = ""
+ try:
+ validation = validate_http_url_and_resolve_ip(file_url)
+ if not validation.is_valid or not validation.validated_ip:
+ raise StorageServiceException(
+ internal_message=f"Invalid URL: {validation.error_message}",
+ operation="download_from_url",
+ )
+
+ effective_temp_dir = temp_dir or getattr(settings, "TMP_PATH", "/tmp")
+ os.makedirs(effective_temp_dir, exist_ok=True)
+ download_result = download_pinned_outbound_file(
+ url=validation.url,
+ pinned_ip=validation.validated_ip,
+ timeout_seconds=300,
+ user_agent="Knowhere-FileDownloader/1.0",
+ temp_dir=effective_temp_dir,
+ )
+ temp_file_path = download_result.temp_file_path
+ return temp_file_path
+ except StorageServiceException:
+ raise
+ except Exception as exc:
+ if temp_file_path and os.path.exists(temp_file_path):
+ os.remove(temp_file_path)
+ raise StorageServiceException(
+ internal_message=f"Failed to download file: {exc}",
+ operation="download_from_url",
+ original_exception=exc,
+ ) from exc
+
+ @staticmethod
+ def get_content_type(file_extension: str) -> str:
+ content_types = {
+ ".pdf": "application/pdf",
+ ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+ ".doc": "application/msword",
+ ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+ ".xls": "application/vnd.ms-excel",
+ ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+ ".ppt": "application/vnd.ms-powerpoint",
+ ".csv": "text/csv",
+ ".txt": "text/plain",
+ ".md": "text/markdown",
+ ".json": "application/json",
+ ".xml": "application/xml",
+ ".html": "text/html",
+ ".htm": "text/html",
+ ".png": "image/png",
+ ".jpg": "image/jpeg",
+ ".jpeg": "image/jpeg",
+ ".gif": "image/gif",
+ ".bmp": "image/bmp",
+ ".tiff": "image/tiff",
+ ".svg": "image/svg+xml",
+ ".zip": "application/zip",
+ ".rar": "application/x-rar-compressed",
+ ".7z": "application/x-7z-compressed",
+ ".tar": "application/x-tar",
+ ".gz": "application/gzip",
+ }
+ return content_types.get(file_extension.lower(), "application/octet-stream")
diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py
index 52d26783f..650019385 100644
--- a/packages/shared-python/shared/services/storage/result_storage.py
+++ b/packages/shared-python/shared/services/storage/result_storage.py
@@ -2,9 +2,13 @@
import os
from dataclasses import dataclass
+from collections.abc import Iterator
from pathlib import Path
from typing import Protocol
+from shared.services.storage.job_file_storage import JobFileStorage
+from shared.services.storage.storage_adapter import StorageAdapter
+
_EXCLUDED_FILE_NAMES = {".DS_Store", "Thumbs.db"}
_EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"}
_CLIENT_ARTIFACT_DIRS = {"images", "tables"}
@@ -29,32 +33,24 @@ def generate_artifact_url(
def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: ...
-class ResultS3:
+class JobResultStorage:
def __init__(
- self, *, results_bucket: str | None = None, storage_adapter=None
+ self,
+ *,
+ results_bucket: str | None = None,
+ storage_adapter: StorageAdapter | None = None,
) -> None:
- if results_bucket is None:
- from shared.core.config import settings
-
- results_bucket = getattr(
- settings, "S3_RESULTS_BUCKET", settings.S3_BUCKET_NAME
- )
- self.results_bucket = results_bucket
- self._storage_adapter = storage_adapter
-
- @property
- def storage_adapter(self):
- if self._storage_adapter is None:
- from shared.core.config.storage import get_cached_storage_adapter
-
- self._storage_adapter = get_cached_storage_adapter()
- return self._storage_adapter
+ self._job_file_storage = JobFileStorage(
+ storage_adapter=storage_adapter,
+ results_bucket=results_bucket,
+ )
+ self.results_bucket = self._job_file_storage.results_bucket
def build_zip_key(self, *, job_id: str) -> str:
- return f"results/{job_id}.zip"
+ return self._job_file_storage.build_result_zip_key(job_id=job_id)
def build_raw_prefix(self, *, job_id: str) -> str:
- return f"results/{job_id}/"
+ return self._job_file_storage.build_result_raw_prefix(job_id=job_id)
def build_raw_key(self, *, job_id: str, relative_path: str) -> str:
normalized = self._normalize_raw_relative_path(relative_path)
@@ -82,15 +78,21 @@ def upload(
if not zip_path.is_file():
raise ValueError(f"Result ZIP file does not exist: {zip_file_path}")
zip_key = self.build_zip_key(job_id=job_id)
- self.storage_adapter.upload_file(str(zip_path), zip_key, self.results_bucket)
+ self._job_file_storage.upload_local_file(
+ str(zip_path),
+ zip_key,
+ bucket=self.results_bucket,
+ )
self._cleanup_file(zip_path)
raw_files: dict[str, str] = {}
for file_path in self._iter_raw_files(result_path):
relative_path = file_path.relative_to(result_path).as_posix()
raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path)
- self.storage_adapter.upload_file(
- str(file_path), raw_key, self.results_bucket
+ self._job_file_storage.upload_local_file(
+ str(file_path),
+ raw_key,
+ bucket=self.results_bucket,
)
raw_files[relative_path] = raw_key
@@ -101,12 +103,11 @@ def upload(
)
def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | None:
- return self.storage_adapter.generate_presigned_url(
+ return self._job_file_storage.generate_download_url(
storage_key,
- expiration=expires_in,
bucket=self.results_bucket,
- method="GET",
- )
+ expires_in=expires_in,
+ )["download_url"]
def generate_artifact_url(
self, *, job_id: str, artifact_ref: str, expires_in: int = 3600
@@ -119,7 +120,7 @@ def generate_artifact_url(
expires_in=expires_in,
)
- def _iter_raw_files(self, result_dir: Path):
+ def _iter_raw_files(self, result_dir: Path) -> Iterator[Path]:
for root, dir_names, file_names in os.walk(result_dir):
dir_names[:] = [
dir_name
@@ -160,4 +161,4 @@ def _cleanup_file(self, file_path: Path) -> None:
def get_result_storage() -> ResultStorage:
- return ResultS3()
+ return JobResultStorage()
diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py
new file mode 100644
index 000000000..410bfa111
--- /dev/null
+++ b/packages/shared-python/shared/services/storage/zip_result_schema.py
@@ -0,0 +1,345 @@
+"""Schema projection for Knowhere ZIP result packages."""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from shared.services.chunks.chunk_connections import (
+ build_resource_target_map,
+ convert_refs_to_embed_connections,
+ merge_connections,
+ normalize_connect_to_targets,
+ parse_relationship_refs,
+)
+from shared.utils.text_utils import truncate_content_preview
+from shared.utils.utc_now import utc_now_naive
+
+
+class ZipResultSchemaBuilder:
+ def calculate_statistics(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
+ total_chunks = len(chunks)
+ text_chunks = 0
+ image_chunks = 0
+ table_chunks = 0
+
+ for chunk in chunks:
+ chunk_type = chunk.get("type", "")
+ raw_type = str(chunk_type).strip()
+ normalized_type = raw_type.split("\n", 1)[0].lower()
+ if normalized_type == "image":
+ image_chunks += 1
+ elif normalized_type == "table":
+ table_chunks += 1
+ else:
+ text_chunks += 1
+
+ return {
+ "total_chunks": total_chunks,
+ "text_chunks": text_chunks,
+ "image_chunks": image_chunks,
+ "table_chunks": table_chunks,
+ "total_pages": None,
+ }
+
+ def format_chunks(
+ self,
+ chunks: List[Dict[str, Any]],
+ image_files_map: Dict[str, Dict[str, Any]],
+ table_files_map: Dict[str, Dict[str, Any]],
+ ) -> List[Dict[str, Any]]:
+ resource_target_map = build_resource_target_map(
+ chunks,
+ image_files_map=image_files_map,
+ table_files_map=table_files_map,
+ )
+
+ formatted = []
+ for chunk in chunks:
+ chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id"))
+ chunk_type_str = chunk.get("type", "")
+ raw_type = str(chunk_type_str).strip()
+ normalized_type = raw_type.split("\n", 1)[0].lower()
+ img_info = image_files_map.get(chunk_id)
+
+ if normalized_type == "image":
+ chunk_type = "image"
+ elif normalized_type == "table":
+ chunk_type = "table"
+ else:
+ chunk_type = "text"
+
+ content = chunk.get("text") or chunk.get("content", "")
+ path = chunk.get("path", "")
+ existing_metadata = chunk.get("metadata", {})
+ metadata = {
+ "length": existing_metadata.get("length") or len(content),
+ "summary": existing_metadata.get("summary") or chunk.get("summary", ""),
+ "page_nums": existing_metadata.get("page_nums", []),
+ }
+ document_top_summary = str(
+ existing_metadata.get("document_top_summary") or ""
+ ).strip()
+ if document_top_summary:
+ metadata["document_top_summary"] = document_top_summary
+
+ if chunk_type == "text":
+ metadata["tokens"] = existing_metadata.get("tokens") or chunk.get(
+ "tokens", 0
+ )
+ metadata["keywords"] = existing_metadata.get("keywords") or chunk.get(
+ "keywords", []
+ )
+ relationship_refs = parse_relationship_refs(
+ chunk.get("type_raw") or chunk_type_str,
+ str(content),
+ )
+ embed_connections = convert_refs_to_embed_connections(
+ relationship_refs, resource_target_map
+ )
+ related_connections = normalize_connect_to_targets(
+ existing_metadata.get("connect_to")
+ or chunk.get("connect_to")
+ or chunk.get("connectto"),
+ resource_target_map,
+ )
+ metadata["connect_to"] = merge_connections(
+ embed_connections, related_connections
+ )
+
+ elif chunk_type == "image":
+ if img_info:
+ metadata["file_path"] = img_info["file_path"]
+ metadata["keywords"] = existing_metadata.get("keywords") or chunk.get(
+ "keywords", []
+ )
+ metadata["tokens"] = []
+
+ elif chunk_type == "table":
+ file_path = existing_metadata.get("file_path")
+ if not file_path:
+ table_info = table_files_map.get(chunk_id)
+ if table_info:
+ file_path = table_info["file_path"]
+ else:
+ table_name = (
+ path.split("/")[-1]
+ if "/" in path
+ else f"table_{chunk_id}.html"
+ )
+ file_path = f"tables/{table_name}"
+
+ metadata["file_path"] = file_path
+ metadata["keywords"] = existing_metadata.get("keywords") or chunk.get(
+ "keywords", []
+ )
+ metadata["tokens"] = []
+
+ formatted.append(
+ {
+ "chunk_id": chunk_id,
+ "type": chunk_type,
+ "content": content,
+ "path": path,
+ "metadata": metadata,
+ }
+ )
+
+ return formatted
+
+ def generate_manifest(
+ self,
+ *,
+ job_id: str,
+ data_id: Optional[str],
+ source_file_name: str,
+ statistics: Dict[str, Any],
+ job_metadata: Dict[str, Any],
+ hierarchy: Optional[Dict[str, Any]] = None,
+ ) -> Dict[str, Any]:
+ return {
+ "version": "2.0",
+ "job_id": job_id,
+ "data_id": data_id,
+ "source_file_name": source_file_name,
+ "processing_date": utc_now_naive().isoformat() + "Z",
+ "processing": {
+ "page_count": job_metadata.get("page_count"),
+ "billing_status": job_metadata.get("billing_status"),
+ "cost": {
+ "micro_dollars": job_metadata.get("billing_amount_micro_dollars"),
+ "credits": job_metadata.get("billing_credits"),
+ },
+ "timing": {
+ "started_at": job_metadata.get("processing_started_at"),
+ "completed_at": job_metadata.get("processing_completed_at"),
+ "duration_ms": job_metadata.get("processing_duration_ms"),
+ },
+ },
+ "statistics": statistics,
+ "HIERARCHY": hierarchy or {},
+ }
+
+ def build_hierarchy_dict(
+ self,
+ sections: List[Dict[str, Any]],
+ ) -> Dict[str, Any]:
+ hierarchy: Dict[str, Any] = {}
+ title_counts: Dict[str, int] = {}
+
+ for section in sections:
+ raw_title = str(section.get("title") or "").strip()
+ if not raw_title:
+ continue
+
+ title_counts[raw_title] = title_counts.get(raw_title, 0) + 1
+ title = (
+ raw_title
+ if title_counts[raw_title] == 1
+ else f"{raw_title} ({title_counts[raw_title]})"
+ )
+ hierarchy[title] = self.build_hierarchy_dict(
+ section.get("children") or []
+ )
+
+ return hierarchy
+
+ def build_doc_nav(
+ self,
+ formatted_chunks: List[Dict[str, Any]],
+ source_file_name: str,
+ ) -> Dict[str, Any]:
+ text_chunks: List[Dict[str, Any]] = []
+ image_resources: List[Dict[str, Any]] = []
+ table_resources: List[Dict[str, Any]] = []
+
+ stats = {
+ "total_chunks": 0,
+ "text_chunks": 0,
+ "image_chunks": 0,
+ "table_chunks": 0,
+ "max_depth": 0,
+ }
+
+ for formatted_chunk in formatted_chunks:
+ chunk_type = formatted_chunk.get("type", "text")
+ path = formatted_chunk.get("path", "")
+ metadata = formatted_chunk.get("metadata") or {}
+ summary_raw = (metadata.get("summary") or "").strip()
+ content_raw = (formatted_chunk.get("content") or "").strip()
+ summary = " ".join(summary_raw.split()) if summary_raw else ""
+ content_preview = truncate_content_preview(content_raw) if content_raw else ""
+
+ stats["total_chunks"] += 1
+ if chunk_type == "image":
+ stats["image_chunks"] += 1
+ image_resources.append(
+ {
+ "path": path,
+ "summary": summary or content_preview,
+ }
+ )
+ elif chunk_type == "table":
+ stats["table_chunks"] += 1
+ table_resources.append(
+ {
+ "path": path,
+ "summary": summary or content_preview,
+ }
+ )
+ else:
+ stats["text_chunks"] += 1
+ text_chunks.append(
+ {
+ "path": path,
+ "summary": summary or content_preview,
+ }
+ )
+
+ sections = self._build_section_tree(text_chunks)
+ stats["max_depth"] = _max_depth(sections)
+ return {
+ "version": "1.0",
+ "file_name": source_file_name or "",
+ "stats": stats,
+ "sections": sections,
+ "resources": {
+ "images": image_resources,
+ "tables": table_resources,
+ },
+ }
+
+ def _build_section_tree(
+ self,
+ text_chunks: List[Dict[str, Any]],
+ ) -> List[Dict[str, Any]]:
+ root_children: Dict[str, dict] = {}
+
+ for chunk in text_chunks:
+ path = chunk.get("path", "")
+ parts = [part.strip() for part in path.split("/") if part.strip()]
+ section_parts = parts[2:] if len(parts) > 2 else []
+
+ if not section_parts:
+ key = "__root__"
+ if key not in root_children:
+ root_children[key] = {
+ "title": "Root",
+ "path": "/".join(parts[:2]) if len(parts) >= 2 else path,
+ "summary": chunk.get("summary", ""),
+ "chunk_count": 0,
+ "_children_map": {},
+ }
+ root_children[key]["chunk_count"] += 1
+ if not root_children[key]["summary"]:
+ root_children[key]["summary"] = chunk.get("summary", "")
+ continue
+
+ current_level = root_children
+ full_section_path_parts = parts[:2]
+ for index, part in enumerate(section_parts):
+ full_section_path_parts.append(part)
+ if part not in current_level:
+ current_level[part] = {
+ "title": part,
+ "path": "/".join(full_section_path_parts),
+ "summary": "",
+ "chunk_count": 0,
+ "_children_map": {},
+ }
+ node = current_level[part]
+ if index == len(section_parts) - 1:
+ node["chunk_count"] += 1
+ if not node["summary"]:
+ node["summary"] = chunk.get("summary", "")
+ current_level = node["_children_map"]
+
+ return _section_tree_to_output(root_children)
+
+
+def _max_depth(nodes: list, depth: int = 1) -> int:
+ max_depth = depth if nodes else 0
+ for node in nodes:
+ max_depth = max(max_depth, _max_depth(node.get("children", []), depth + 1))
+ return max_depth
+
+
+def _section_tree_to_output(
+ children_map: Dict[str, dict],
+ level: int = 1,
+) -> List[Dict[str, Any]]:
+ result = []
+ for node in children_map.values():
+ children = _section_tree_to_output(node["_children_map"], level + 1)
+ total_chunks = node["chunk_count"] + sum(
+ child.get("chunk_count", 0) for child in children
+ )
+ result.append(
+ {
+ "title": node["title"],
+ "path": node["path"],
+ "level": level,
+ "summary": node["summary"],
+ "chunk_count": total_chunks,
+ "children": children,
+ }
+ )
+ return result
diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py
index 68eb1b0fc..033c2db25 100644
--- a/packages/shared-python/shared/services/storage/zip_result_service.py
+++ b/packages/shared-python/shared/services/storage/zip_result_service.py
@@ -13,29 +13,20 @@
from loguru import logger
from PIL import Image
-from shared.services.chunks.chunk_connections import (
- build_resource_target_map,
- convert_refs_to_embed_connections,
- merge_connections,
- normalize_connect_to_targets,
- parse_relationship_refs,
-)
-from shared.utils.text_utils import truncate_content_preview
-
import pandas as pd
from shared.core.exceptions.domain_exceptions import (
KnowhereException,
StorageServiceException,
)
-from shared.utils.utc_now import utc_now_naive
+from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder
class ZipResultService:
"""ZIP Result Package Generation Service"""
def __init__(self):
- pass
+ self._schema = ZipResultSchemaBuilder()
def generate_zip_package(
self,
@@ -85,10 +76,10 @@ def generate_zip_package(
table_files_map = {tb["id"]: tb for tb in table_files_info}
# Convert chunks data format (using file info)
- formatted_chunks = self._format_chunks(
+ formatted_chunks = self._schema.format_chunks(
chunks, image_files_map, table_files_map
)
- statistics = self._calculate_statistics(formatted_chunks)
+ statistics = self._schema.calculate_statistics(formatted_chunks)
doc_nav: Dict[str, Any] = {}
hierarchy: Dict[str, Any] = {}
@@ -136,8 +127,8 @@ def generate_zip_package(
# 5. Generate doc_nav.json — unified navigation file
try:
- doc_nav = self._build_doc_nav(formatted_chunks, source_file_name)
- hierarchy = self._build_hierarchy_dict(doc_nav.get("sections", []))
+ doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name)
+ hierarchy = self._schema.build_hierarchy_dict(doc_nav.get("sections", []))
doc_nav_json = json.dumps(doc_nav, ensure_ascii=False, indent=2)
zip_file.writestr("doc_nav.json", doc_nav_json.encode("utf-8"))
logger.info("Added doc_nav.json")
@@ -145,7 +136,7 @@ def generate_zip_package(
logger.warning(f"generate doc_nav.json fail {e}")
# 6. Generate manifest.json (checksum not included, stored in database)
- manifest = self._generate_manifest(
+ manifest = self._schema.generate_manifest(
job_id=job_id,
data_id=data_id,
source_file_name=source_file_name,
@@ -179,184 +170,6 @@ def generate_zip_package(
original_exception=e,
)
- def _calculate_statistics(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
- """Calculate statistics"""
- total_chunks = len(chunks)
- text_chunks = 0
- image_chunks = 0
- table_chunks = 0
-
- for chunk in chunks:
- chunk_type = chunk.get("type", "")
- raw_type = str(chunk_type).strip()
- normalized_type = raw_type.split("\n", 1)[0].lower()
- if normalized_type == "image":
- image_chunks += 1
- elif normalized_type == "table":
- table_chunks += 1
- else:
- text_chunks += 1
-
- return {
- "total_chunks": total_chunks,
- "text_chunks": text_chunks,
- "image_chunks": image_chunks,
- "table_chunks": table_chunks,
- "total_pages": None, # Cannot determine page count at this point
- }
-
- def _format_chunks(
- self,
- chunks: List[Dict[str, Any]],
- image_files_map: Dict[str, Dict[str, Any]],
- table_files_map: Dict[str, Dict[str, Any]],
- ) -> List[Dict[str, Any]]:
- """Convert chunks data to ZIP specification format"""
- resource_target_map = build_resource_target_map(
- chunks,
- image_files_map=image_files_map,
- table_files_map=table_files_map,
- )
-
- formatted = []
- for chunk in chunks:
- chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id"))
- chunk_type_str = chunk.get("type", "")
- raw_type = str(chunk_type_str).strip()
- normalized_type = raw_type.split("\n", 1)[0].lower()
- img_info = image_files_map.get(chunk_id)
-
- # Determine chunk type
- if normalized_type == "image":
- chunk_type = "image"
- elif normalized_type == "table":
- chunk_type = "table"
- else:
- chunk_type = "text"
-
- # Get content
- content = chunk.get("text") or chunk.get("content", "")
-
- # Use original path directly to match kb.csv
- path = chunk.get("path", "")
-
- # Get or build base metadata
- existing_metadata = chunk.get("metadata", {})
- metadata = {
- "length": existing_metadata.get("length") or len(content),
- "summary": existing_metadata.get("summary") or chunk.get("summary", ""),
- "page_nums": existing_metadata.get("page_nums", []),
- }
- document_top_summary = str(
- existing_metadata.get("document_top_summary") or ""
- ).strip()
- if document_top_summary:
- metadata["document_top_summary"] = document_top_summary
-
- # Add type-specific fields
- if chunk_type == "text":
- metadata["tokens"] = existing_metadata.get("tokens") or chunk.get(
- "tokens", 0
- )
- metadata["keywords"] = existing_metadata.get("keywords") or chunk.get(
- "keywords", []
- )
-
- # Convert in-text resource refs into embeds edges.
- relationship_refs = parse_relationship_refs(
- chunk.get("type_raw") or chunk_type_str,
- str(content),
- )
- embed_connections = convert_refs_to_embed_connections(
- relationship_refs, resource_target_map
- )
- related_connections = normalize_connect_to_targets(
- existing_metadata.get("connect_to")
- or chunk.get("connect_to")
- or chunk.get("connectto"),
- resource_target_map,
- )
- metadata["connect_to"] = merge_connections(
- embed_connections, related_connections
- )
-
- elif chunk_type == "image":
- if img_info:
- metadata["file_path"] = img_info["file_path"]
- # Unified schema: include keywords and tokens for all chunk types
- metadata["keywords"] = existing_metadata.get("keywords") or chunk.get(
- "keywords", []
- )
- metadata["tokens"] = []
-
- elif chunk_type == "table":
- # Get table info from existing_metadata or table_files_map
- file_path = existing_metadata.get("file_path")
-
- if not file_path:
- # Get table info from table_files_map
- tb_info = table_files_map.get(chunk_id)
- if tb_info:
- file_path = tb_info["file_path"]
- else:
- # Extract from path or use default
- tbl_name = (
- path.split("/")[-1]
- if "/" in path
- else f"table_{chunk_id}.html"
- )
- file_path = f"tables/{tbl_name}"
-
- metadata["file_path"] = file_path
- # Unified schema: include keywords and tokens for all chunk types
- metadata["keywords"] = existing_metadata.get("keywords") or chunk.get(
- "keywords", []
- )
- metadata["tokens"] = []
-
- formatted_chunk = {
- "chunk_id": chunk_id,
- "type": chunk_type,
- "content": content,
- "path": path,
- "metadata": metadata,
- }
- formatted.append(formatted_chunk)
-
- return formatted
-
- def _clean_path(self, path: str) -> str:
- """Clean path, keep only logical path"""
- if not path:
- return "/"
-
- # Remove filesystem path prefix
- # Example: .-->users-->KB_DATA_xxx-->dir-->file.pdf-->chapter-->section
- # Should extract: chapter-->section
-
- # Find the last .pdf, .docx, etc. file extension
- import re
-
- # Match filename pattern (with extension)
- file_pattern = r"[^/]+\.(pdf|docx|doc|txt|md|xlsx|xls|pptx|ppt)"
- match = re.search(file_pattern, path, re.IGNORECASE)
-
- if match:
- # Extract the part after filename
- path_after_file = path[match.end() :]
- # Clean path separators
- path_after_file = path_after_file.replace("-->", "/").strip("/")
- if path_after_file:
- return path_after_file
-
- # If no file pattern found, try to clean common prefixes
- path = path.replace("-->", "/")
- # Remove leading path separators and empty segments
- path = "/".join(
- [p for p in path.split("/") if p and p not in ["", ".", "users"]]
- )
- return path if path else "/"
-
def _collect_image_files(
self, chunks: List[Dict[str, Any]], images_dir: str
) -> List[Dict[str, Any]]:
@@ -642,41 +455,6 @@ def resolve_source_path(
return table_files
- def _generate_manifest(
- self,
- job_id: str,
- data_id: Optional[str],
- source_file_name: str,
- statistics: Dict[str, Any],
- job_metadata: Dict[str, Any],
- hierarchy: Optional[Dict[str, Any]] = None,
- ) -> Dict[str, Any]:
- """Generate manifest.json"""
- manifest = {
- "version": "2.0",
- "job_id": job_id,
- "data_id": data_id,
- "source_file_name": source_file_name,
- "processing_date": utc_now_naive().isoformat() + "Z",
- "processing": {
- "page_count": job_metadata.get("page_count"),
- "billing_status": job_metadata.get("billing_status"),
- "cost": {
- "micro_dollars": job_metadata.get("billing_amount_micro_dollars"),
- "credits": job_metadata.get("billing_credits"),
- },
- "timing": {
- "started_at": job_metadata.get("processing_started_at"),
- "completed_at": job_metadata.get("processing_completed_at"),
- "duration_ms": job_metadata.get("processing_duration_ms"),
- },
- },
- "statistics": statistics,
- "HIERARCHY": hierarchy or {},
- }
-
- return manifest
-
def _calculate_zip_checksum(self, zip_file_path: str) -> str:
"""Calculate SHA-256 checksum of ZIP file"""
sha256_hash = hashlib.sha256()
@@ -684,193 +462,3 @@ def _calculate_zip_checksum(self, zip_file_path: str) -> str:
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest().lower()
-
- def _build_hierarchy_dict(
- self,
- sections: List[Dict[str, Any]],
- ) -> Dict[str, Any]:
- """Build a title-only nested hierarchy from doc_nav sections."""
- hierarchy: Dict[str, Any] = {}
- title_counts: Dict[str, int] = {}
-
- for section in sections:
- raw_title = str(section.get("title") or "").strip()
- if not raw_title:
- continue
-
- title_counts[raw_title] = title_counts.get(raw_title, 0) + 1
- title = (
- raw_title
- if title_counts[raw_title] == 1
- else f"{raw_title} ({title_counts[raw_title]})"
- )
- hierarchy[title] = self._build_hierarchy_dict(
- section.get("children") or []
- )
-
- return hierarchy
-
- def _build_doc_nav(
- self,
- formatted_chunks: List[Dict[str, Any]],
- source_file_name: str,
- ) -> Dict[str, Any]:
- """Build doc_nav.json — unified navigation file.
-
- Structured file serving both human demo and LLM navigation.
-
- The output contains:
- - ``sections``: tree of text sections with summaries and chunk counts.
- - ``resources``: flat lists of image/table chunks with summaries.
- - ``stats``: chunk counts by type.
-
- Each leaf section carries a ``summary`` derived from:
- 1. chunk.metadata.summary (LLM-generated, highest quality)
- 2. chunk.content[:300] (fallback truncation)
-
- Non-leaf section summaries are left empty at this stage and are
- populated later by ``summary_builder.enrich_doc_nav_summaries``.
- """
- # ── Separate text chunks from resource chunks ──
- text_chunks: List[Dict[str, Any]] = []
- image_resources: List[Dict[str, Any]] = []
- table_resources: List[Dict[str, Any]] = []
-
- stats = {"total_chunks": 0, "text_chunks": 0, "image_chunks": 0, "table_chunks": 0, "max_depth": 0}
-
- for fc in formatted_chunks:
- ctype = fc.get("type", "text")
- path = fc.get("path", "")
- meta = fc.get("metadata") or {}
- summary_raw = (meta.get("summary") or "").strip()
- content_raw = (fc.get("content") or "").strip()
- # Normalize whitespace
- summary = " ".join(summary_raw.split()) if summary_raw else ""
- content_preview = truncate_content_preview(content_raw) if content_raw else ""
-
- stats["total_chunks"] += 1
-
- if ctype == "image":
- stats["image_chunks"] += 1
- image_resources.append({
- "path": path,
- "summary": summary or content_preview,
- })
- elif ctype == "table":
- stats["table_chunks"] += 1
- table_resources.append({
- "path": path,
- "summary": summary or content_preview,
- })
- else:
- stats["text_chunks"] += 1
- text_chunks.append({
- "path": path,
- "summary": summary or content_preview,
- })
-
- # ── Build section tree from text chunk paths ──
- # Each text chunk path looks like: "kb_root/filename.pdf/Section/Subsection"
- # We strip the kb_root and filename prefix to get relative section paths.
- sections = self._build_section_tree(text_chunks)
-
- # Compute max depth
- def _max_depth(nodes: list, d: int = 1) -> int:
- m = d if nodes else 0
- for n in nodes:
- m = max(m, _max_depth(n.get("children", []), d + 1))
- return m
-
- stats["max_depth"] = _max_depth(sections)
-
- return {
- "version": "1.0",
- "file_name": source_file_name or "",
- "stats": stats,
- "sections": sections,
- "resources": {
- "images": image_resources,
- "tables": table_resources,
- },
- }
-
- def _build_section_tree(
- self,
- text_chunks: List[Dict[str, Any]],
- ) -> List[Dict[str, Any]]:
- """Build a tree of sections from flat text chunk paths.
-
- Each text chunk has a ``path`` like ``"kb/file.pdf/Sec1/Sub1"``.
- We extract section parts (after kb_root + filename) and build a
- tree using ``children`` arrays.
-
- Returns a list of top-level section nodes.
- """
- # Internal tree node: {title, summary, chunk_count, children: {title: node}}
- root_children: Dict[str, dict] = {} # ordered dict of top-level titles
-
- for chunk in text_chunks:
- path = chunk.get("path", "")
- parts = [p.strip() for p in path.split("/") if p.strip()]
- # Skip kb_root + filename → section parts start at index 2
- section_parts = parts[2:] if len(parts) > 2 else []
-
- if not section_parts:
- # Root-level chunk (no section hierarchy)
- key = "__root__"
- if key not in root_children:
- root_children[key] = {
- "title": "Root",
- "path": "/".join(parts[:2]) if len(parts) >= 2 else path,
- "summary": chunk.get("summary", ""),
- "chunk_count": 0,
- "_children_map": {},
- }
- root_children[key]["chunk_count"] += 1
- # Use the first chunk's summary for root
- if not root_children[key]["summary"]:
- root_children[key]["summary"] = chunk.get("summary", "")
- continue
-
- # Walk the tree, creating nodes as needed
- current_level = root_children
- full_section_path_parts = parts[:2] # start with kb_root/filename
- for i, part in enumerate(section_parts):
- full_section_path_parts.append(part)
- if part not in current_level:
- current_level[part] = {
- "title": part,
- "path": "/".join(full_section_path_parts),
- "summary": "",
- "chunk_count": 0,
- "_children_map": {},
- }
- node = current_level[part]
- if i == len(section_parts) - 1:
- # Leaf — this is the chunk's actual section
- node["chunk_count"] += 1
- if not node["summary"]:
- node["summary"] = chunk.get("summary", "")
- current_level = node["_children_map"]
-
- # Convert internal tree to output format
- def _to_output(children_map: Dict[str, dict], level: int = 1) -> List[Dict[str, Any]]:
- result = []
- for node in children_map.values():
- children = _to_output(node["_children_map"], level + 1)
- # Compute total chunk_count including descendants
- total_chunks = node["chunk_count"] + sum(
- c.get("chunk_count", 0) for c in children
- )
- out = {
- "title": node["title"],
- "path": node["path"],
- "level": level,
- "summary": node["summary"],
- "chunk_count": total_chunks,
- "children": children,
- }
- result.append(out)
- return result
-
- return _to_output(root_children)
diff --git a/packages/shared-python/shared/services/webhook/delivery_client.py b/packages/shared-python/shared/services/webhook/delivery_client.py
new file mode 100644
index 000000000..232f9a875
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/delivery_client.py
@@ -0,0 +1,151 @@
+"""Pinned HTTP delivery for outbound webhooks."""
+
+import asyncio
+import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from loguru import logger
+
+from shared.utils.pinned_outbound_http import send_pinned_outbound_request
+from shared.utils.url_security import validate_http_url_and_resolve_ip_async
+
+HTTP_TIMEOUT_SECONDS = 10
+
+
+@dataclass(frozen=True)
+class WebhookDeliveryTarget:
+ target_url: str
+ pinned_ip: str
+
+
+@dataclass(frozen=True)
+class WebhookDeliveryResult:
+ success: bool
+ status_code: int | None
+ duration_ms: int
+ error_message: str | None
+
+
+@dataclass(frozen=True)
+class WebhookTargetValidation:
+ target: WebhookDeliveryTarget | None
+ failure: WebhookDeliveryResult | None
+
+
+class WebhookDeliveryClient:
+ """Validate and send direct webhook HTTP requests with DNS pinning."""
+
+ async def validate_target(
+ self, *, event_id: str, target_url: str
+ ) -> WebhookTargetValidation:
+ validation = await validate_http_url_and_resolve_ip_async(target_url)
+
+ if not validation.is_valid:
+ logger.warning(
+ f"SSRF validation failed: event_id={event_id}, "
+ f"error={validation.error_message}"
+ )
+ return WebhookTargetValidation(
+ target=None,
+ failure=WebhookDeliveryResult(
+ success=False,
+ status_code=400,
+ duration_ms=0,
+ error_message=f"SSRF: {validation.error_message}",
+ ),
+ )
+
+ if not validation.validated_ip:
+ return WebhookTargetValidation(
+ target=None,
+ failure=WebhookDeliveryResult(
+ success=False,
+ status_code=400,
+ duration_ms=0,
+ error_message="SSRF validation did not return a pinned IP",
+ ),
+ )
+
+ return WebhookTargetValidation(
+ target=WebhookDeliveryTarget(
+ target_url=target_url,
+ pinned_ip=validation.validated_ip,
+ ),
+ failure=None,
+ )
+
+ async def post_json(
+ self,
+ *,
+ event_id: str,
+ target: WebhookDeliveryTarget,
+ payload: Mapping[str, Any],
+ headers: Mapping[str, str],
+ ) -> WebhookDeliveryResult:
+ start_time = time.time()
+
+ try:
+ response = await send_pinned_outbound_request(
+ method="POST",
+ url=target.target_url,
+ pinned_ip=target.pinned_ip,
+ timeout_seconds=HTTP_TIMEOUT_SECONDS,
+ headers=headers,
+ json_body=payload,
+ )
+ duration_ms = int((time.time() - start_time) * 1000)
+
+ if 200 <= response.status < 300:
+ logger.info(
+ f"Webhook delivered: event_id={event_id}, status={response.status}"
+ )
+ return WebhookDeliveryResult(
+ success=True,
+ status_code=response.status,
+ duration_ms=duration_ms,
+ error_message=None,
+ )
+
+ if 300 <= response.status < 400:
+ logger.warning(
+ f"Webhook redirect blocked (SSRF protection): "
+ f"event_id={event_id}, status={response.status}"
+ )
+ return WebhookDeliveryResult(
+ success=False,
+ status_code=response.status,
+ duration_ms=duration_ms,
+ error_message=f"Redirect blocked: HTTP {response.status}",
+ )
+
+ logger.warning(
+ f"Webhook failed: event_id={event_id}, status={response.status}"
+ )
+ return WebhookDeliveryResult(
+ success=False,
+ status_code=response.status,
+ duration_ms=duration_ms,
+ error_message=f"HTTP {response.status}",
+ )
+
+ except asyncio.TimeoutError:
+ duration_ms = int((time.time() - start_time) * 1000)
+ logger.error(f"Webhook timeout: event_id={event_id}")
+ return WebhookDeliveryResult(
+ success=False,
+ status_code=None,
+ duration_ms=duration_ms,
+ error_message="Connection timeout",
+ )
+
+ except Exception as error:
+ duration_ms = int((time.time() - start_time) * 1000)
+ logger.error(f"Webhook error: event_id={event_id}, error={error}")
+ return WebhookDeliveryResult(
+ success=False,
+ status_code=None,
+ duration_ms=duration_ms,
+ error_message=str(error),
+ )
diff --git a/packages/shared-python/shared/services/webhook/dispatcher.py b/packages/shared-python/shared/services/webhook/dispatcher.py
index e5ee4610d..0f3c4b953 100644
--- a/packages/shared-python/shared/services/webhook/dispatcher.py
+++ b/packages/shared-python/shared/services/webhook/dispatcher.py
@@ -1,44 +1,25 @@
"""
Webhook Dispatcher Service
-Dispatches webhook events via HTTP requests with HMAC signing and delivery logging.
-Called by Celery task for async processing.
+Dispatches webhook events with retry policy. Direct HTTP delivery details live
+behind WebhookEventDelivery.
"""
-import asyncio
-import hashlib
-import hmac
-import json
import threading
-import time
-import uuid
from datetime import datetime, timezone
-from typing import Any, Dict, Optional, Tuple
+from typing import Optional
from loguru import logger
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
-# Use standard db context - run_async_task handles the loop reuse
from shared.core.database import get_db_context
-from shared.core.exceptions.domain_exceptions import (
- SystemSettingInvalidException,
- SystemSettingMissingException,
-)
from shared.core.exceptions.webhook_exceptions import WebhookDeliveryException
-from shared.models.database.job import Job
from shared.models.database.webhook import WebhookEvent, WebhookEventStatus
-from shared.models.database.webhook_log import WebhookLog
-from shared.utils.pinned_outbound_http import (
- send_pinned_outbound_request,
-)
-from shared.utils.url_security import (
- HTTPURLValidationResult,
- validate_http_url_and_resolve_ip_async,
-)
+from shared.services.webhook.delivery_client import WebhookDeliveryResult
+from shared.services.webhook.event_delivery import WebhookEventDelivery
# Constants
-HTTP_TIMEOUT_SECONDS = 10
MAX_ATTEMPTS = 6
@@ -54,6 +35,9 @@ class WebhookDispatcher:
5. On failure, signals the caller to schedule retry
"""
+ def __init__(self, event_delivery: WebhookEventDelivery | None = None) -> None:
+ self._event_delivery = event_delivery or WebhookEventDelivery()
+
async def dispatch(self, event_id: str) -> bool:
"""
Dispatch a webhook event.
@@ -87,27 +71,24 @@ async def dispatch(self, event_id: str) -> bool:
await self._mark_failed(db, event)
return True # ACK
- # 4. Dispatch the webhook
- # Logging is now handled inside _send_webhook
- success, status_code, duration_ms, error_message = await self._send_webhook(
+ delivery_result = await self._event_delivery.send(
db=db, event=event, is_manual=False
)
- # 6. Handle result (Logging already done)
- if success:
+ if delivery_result.success:
await self._mark_delivered(db, event)
return True # Success
else:
# Determine if error is retryable
# Retryable: 5xx, timeout (None), 429 (rate limit)
# NOT retryable: 4xx (except 429) - client errors won't be fixed by retrying
- is_retryable = self._is_retryable_error(status_code)
+ is_retryable = self._is_retryable_error(delivery_result.status_code)
if not is_retryable:
# Permanent failure - don't retry
logger.warning(
f"WebhookEvent permanent failure (non-retryable): "
- f"event_id={event_id}, status={status_code}"
+ f"event_id={event_id}, status={delivery_result.status_code}"
)
await self._mark_failed(db, event)
return True # ACK - no point retrying
@@ -121,9 +102,11 @@ async def dispatch(self, event_id: str) -> bool:
# Raise exception so Celery task will retry
raise WebhookDeliveryException(
- internal_message=f"Webhook delivery failed: {error_message}",
+ internal_message=(
+ f"Webhook delivery failed: {delivery_result.error_message}"
+ ),
retryable=True,
- status_code=status_code,
+ status_code=delivery_result.status_code,
)
async def mark_event_failed(self, event_id: str) -> None:
@@ -147,226 +130,11 @@ async def _fetch_event(
)
return result.scalar_one_or_none()
- async def _send_webhook(
- self, db: AsyncSession, event: WebhookEvent, is_manual: bool = False
- ) -> Tuple[bool, Optional[int], int, Optional[str]]:
- """
- Send HTTP POST request to webhook target and log the attempt.
-
- Args:
- db: Database session
- event: WebhookEvent object
- is_manual: True if manually triggered (adds 'trigger': 'manual' to payload)
-
- Returns:
- Tuple of (success, status_code, duration_ms, error_message)
- """
-
- # Generate attempt ID
- attempt_id = str(uuid.uuid4())
-
- # SSRF Protection
- validation: HTTPURLValidationResult = await validate_http_url_and_resolve_ip_async(
- event.target_url,
- )
- if not validation.is_valid:
- logger.warning(
- f"SSRF validation failed: event_id={event.id}, error={validation.error_message}"
- )
- return False, 400, 0, f"SSRF: {validation.error_message}"
-
- # Enrich payload with job result data at delivery time
- enriched_payload = await self._enrich_payload(event)
-
- # Add manual mark if requested
- if is_manual:
- enriched_payload["trigger"] = "manual"
-
- # Helper to get user_id from job
- async def _get_job_owner(job_id: str) -> Optional[str]:
- result = await db.execute(select(Job.user_id).where(Job.job_id == job_id))
- return result.scalar_one_or_none()
-
- # Resolve secret (Lazy creation)
- secret = None
- try:
- user_id = await _get_job_owner(event.job_id)
- if user_id:
- secret = await self._resolve_secret(db, user_id, event.target_url)
- else:
- logger.warning(
- f"Could not resolve secret: Job {event.job_id} has no user_id"
- )
- except (SystemSettingMissingException, SystemSettingInvalidException) as e:
- logger.error(f"Configuration error during secret resolution: {e}")
- # Return 424 (Failed Dependency) to ensure it's treated as a non-retryable error
- return False, 424, 0, f"Configuration Error: {e}"
- except Exception as e:
- logger.error(f"Secret resolution failed: {e}")
-
- if not secret:
- logger.error(f"No secret found or created/resolved for event {event.id}")
- # Default to non-retryable error for any secret resolution failure
- return False, 424, 0, "Secret resolution failed"
-
- # Sign payload
- signature = self._sign_payload(enriched_payload, secret)
-
- # Build headers
- headers = {
- "Content-Type": "application/json",
- "X-Knowhere-Signature": signature,
- "X-Knowhere-Attempt-ID": attempt_id,
- "User-Agent": "Knowhere-Webhook/1.0",
- }
-
- start_time = time.time()
- status_code = None
- error_message = None
- success = False
-
- try:
- pinned_ip = validation.validated_ip
- if not pinned_ip:
- return False, 400, 0, "SSRF validation did not return a pinned IP"
-
- response = await send_pinned_outbound_request(
- method="POST",
- url=event.target_url,
- pinned_ip=pinned_ip,
- timeout_seconds=HTTP_TIMEOUT_SECONDS,
- headers=headers,
- json_body=enriched_payload,
- )
- duration_ms = int((time.time() - start_time) * 1000)
- status_code = response.status
-
- if 200 <= response.status < 300:
- logger.info(
- f"Webhook delivered: event_id={event.id}, status={response.status}"
- )
- success = True
- elif 300 <= response.status < 400:
- logger.warning(
- f"Webhook redirect blocked (SSRF protection): "
- f"event_id={event.id}, status={response.status}"
- )
- error_message = f"Redirect blocked: HTTP {response.status}"
- success = False
- else:
- logger.warning(
- f"Webhook failed: event_id={event.id}, status={response.status}"
- )
- error_message = f"HTTP {response.status}"
- success = False
-
- except asyncio.TimeoutError:
- duration_ms = int((time.time() - start_time) * 1000)
- logger.error(f"Webhook timeout: event_id={event.id}")
- error_message = "Connection timeout"
- success = False
-
- except Exception as e:
- duration_ms = int((time.time() - start_time) * 1000)
- logger.error(f"Webhook error: event_id={event.id}, error={e}")
- error_message = str(e)
- success = False
-
- # Log delivery attempt
- # If manual, event_id is None to avoid FK violation
- log_event_id = None if is_manual else event.id
-
- try:
- # Combine headers and payload
- combined_payload = {"header": headers, "payload": enriched_payload}
-
- log = WebhookLog(
- job_id=event.job_id,
- event_id=log_event_id,
- webhook_url=event.target_url,
- attempt_number=event.attempts + 1,
- request_payload=combined_payload,
- signature=signature,
- idempotency_key=str(uuid.uuid4()),
- response_status_code=status_code,
- error_message=error_message,
- duration_ms=duration_ms,
- )
- db.add(log)
- # If auto-commit is needed?
- # Dispatcher.dispatch uses passed 'db' session which is managed by 'async with get_db_context()'.
- # It commits inside _mark_delivered etc.
- # We should probably commit/flush here to persist log even if update fails?
- await db.commit()
-
- except Exception as e:
- logger.error(f"Failed to log webhook delivery: {e}")
-
- return success, status_code, duration_ms, error_message
-
- async def _enrich_payload(self, event: WebhookEvent) -> Dict[str, Any]:
- """
- Enrich webhook payload with job result data at delivery time.
-
- For job.completed events:
- - Adds result_url (fresh download URL for result zip)
- - Adds result (inline payload with checksum/statistics)
-
- This ensures download URLs are generated fresh (they expire)
- and data is current at delivery time.
- """
- payload = dict(event.payload) # Copy to avoid mutating stored payload
-
- # Only enrich completion events
- if payload.get("event") != "job.completed":
- return payload
-
- try:
- # Fetch job with result
- from sqlalchemy.orm import selectinload
-
- from shared.models.database.job import Job
-
- async with get_db_context() as db:
- result = await db.execute(
- select(Job)
- .options(selectinload(Job.job_result))
- .where(Job.job_id == event.job_id)
- )
- job = result.scalar_one_or_none()
-
- if not job or not job.job_result:
- logger.warning(
- f"Job or result not found for enrichment: job_id={event.job_id}"
- )
- return payload
-
- job_result = job.job_result
-
- # Add result_url (fresh download link)
- if job_result.result_s3_key:
- from shared.services.storage.file_upload_service import (
- FileUploadService,
- )
-
- upload_service = FileUploadService()
- url_info = await upload_service.generate_download_url(
- job_result.result_s3_key
- )
- payload["result_url"] = url_info["download_url"]
- logger.debug(
- f"Enriched payload with result_url for job {event.job_id}"
- )
-
- # Add result (inline payload)
- if job_result.inline_payload:
- payload["result"] = job_result.inline_payload
-
- except Exception as e:
- logger.error(f"Failed to enrich payload for event {event.id}: {e}")
- # Continue with original payload if enrichment fails
-
- return payload
+ async def send_manual_webhook(
+ self, db: AsyncSession, event: WebhookEvent
+ ) -> WebhookDeliveryResult:
+ """Send a webhook immediately for an operator-triggered retry."""
+ return await self._event_delivery.send(db=db, event=event, is_manual=True)
def _is_retryable_error(self, status_code: Optional[int]) -> bool:
"""
@@ -402,62 +170,6 @@ def _is_retryable_error(self, status_code: Optional[int]) -> bool:
# Examples: 400 Bad Request, 401 Unauthorized, 404 Not Found
return False
- async def _resolve_secret(
- self, db: AsyncSession, user_id: str, endpoint: str
- ) -> Optional[str]:
- """
- Resolve webhook secret using repository (Lazy creation).
-
- 1. Try to get existing active secret for user/endpoint.
- 2. If not found, create a new one.
- 3. Decrypt and return the raw secret string.
- """
- try:
- # Import here to avoid circular dependency with WebhookDispatcher
- from shared.repositories.webhook_secret_repository import (
- WebhookSecretRepository,
- )
-
- repo = WebhookSecretRepository()
- secret_obj = await repo.get_or_create_secret(db, user_id, endpoint=endpoint)
-
- # Update usage timestamp
- if secret_obj:
- secret_obj.last_used_at = datetime.now(timezone.utc).replace(
- tzinfo=None
- )
- db.add(secret_obj)
- # We don't commit here to avoid side effects if the caller aborts,
- # but the session will eventually be committed by the caller.
-
- # Decrypt
- return repo.decrypt_secret(secret_obj)
- except (SystemSettingMissingException, SystemSettingInvalidException):
- # Re-raise configuration errors so they can be handled as non-retryable
- raise
- except Exception as e:
- logger.error(f"Failed to resolve/create secret for user {user_id}: {e}")
- return None
-
- def _sign_payload(self, payload: Dict[str, Any], secret: str) -> str:
- """
- Generate timestamped HMAC-SHA256 signature.
-
- Format: t=,v1=
- Signed content: "{timestamp}.{json_payload}"
-
- This prevents replay attacks by binding the signature to the current time.
- """
- timestamp = int(time.time())
- payload_str = json.dumps(payload, separators=(",", ":"))
- signed_content = f"{timestamp}.{payload_str}"
-
- signature = hmac.new(
- secret.encode("utf-8"), signed_content.encode("utf-8"), hashlib.sha256
- ).hexdigest()
-
- return f"t={timestamp},v1={signature}"
-
async def _mark_delivered(self, db: AsyncSession, event: WebhookEvent) -> None:
"""Mark event as delivered."""
event.status = WebhookEventStatus.DELIVERED
diff --git a/packages/shared-python/shared/services/webhook/event_delivery.py b/packages/shared-python/shared/services/webhook/event_delivery.py
new file mode 100644
index 000000000..21b5509fa
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/event_delivery.py
@@ -0,0 +1,127 @@
+"""Direct WebhookEvent delivery attempt orchestration."""
+
+import uuid
+from typing import Any
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.core.exceptions.domain_exceptions import (
+ SystemSettingInvalidException,
+ SystemSettingMissingException,
+)
+from shared.models.database.webhook import WebhookEvent
+from shared.models.database.webhook_log import WebhookLog
+from shared.services.webhook.delivery_client import (
+ WebhookDeliveryClient,
+ WebhookDeliveryResult,
+)
+from shared.services.webhook.payload_enrichment import WebhookPayloadEnricher
+from shared.services.webhook.secret_resolver import WebhookSecretResolver
+from shared.services.webhook.signing import build_webhook_headers
+
+
+class WebhookEventDelivery:
+ """Send one direct webhook attempt and persist its delivery log."""
+
+ def __init__(
+ self,
+ *,
+ client: WebhookDeliveryClient | None = None,
+ enricher: WebhookPayloadEnricher | None = None,
+ secret_resolver: WebhookSecretResolver | None = None,
+ ) -> None:
+ self._client = client or WebhookDeliveryClient()
+ self._enricher = enricher or WebhookPayloadEnricher()
+ self._secret_resolver = secret_resolver or WebhookSecretResolver()
+
+ async def send(
+ self, *, db: AsyncSession, event: WebhookEvent, is_manual: bool = False
+ ) -> WebhookDeliveryResult:
+ attempt_id = str(uuid.uuid4())
+ target_validation = await self._client.validate_target(
+ event_id=event.id,
+ target_url=event.target_url,
+ )
+ if target_validation.failure:
+ return target_validation.failure
+ if not target_validation.target:
+ return WebhookDeliveryResult(
+ success=False,
+ status_code=400,
+ duration_ms=0,
+ error_message="Webhook target validation failed",
+ )
+
+ payload = await self._enricher.enrich(event)
+ if is_manual:
+ payload["trigger"] = "manual"
+
+ secret, secret_error = await self._resolve_secret(db, event)
+ if not secret:
+ logger.error(f"No secret found or created/resolved for event {event.id}")
+ return WebhookDeliveryResult(
+ success=False,
+ status_code=424,
+ duration_ms=0,
+ error_message=secret_error or "Secret resolution failed",
+ )
+
+ headers = build_webhook_headers(
+ payload=payload,
+ secret=secret,
+ attempt_id=attempt_id,
+ )
+ result = await self._client.post_json(
+ event_id=event.id,
+ target=target_validation.target,
+ payload=payload,
+ headers=headers,
+ )
+ await self._log_attempt(
+ db=db,
+ event=event,
+ is_manual=is_manual,
+ headers=headers,
+ payload=payload,
+ result=result,
+ )
+ return result
+
+ async def _resolve_secret(
+ self, db: AsyncSession, event: WebhookEvent
+ ) -> tuple[str | None, str | None]:
+ try:
+ return await self._secret_resolver.resolve_for_event(db, event), None
+ except (SystemSettingMissingException, SystemSettingInvalidException) as error:
+ logger.error(f"Configuration error during secret resolution: {error}")
+ return None, f"Configuration Error: {error}"
+
+ async def _log_attempt(
+ self,
+ *,
+ db: AsyncSession,
+ event: WebhookEvent,
+ is_manual: bool,
+ headers: dict[str, str],
+ payload: dict[str, Any],
+ result: WebhookDeliveryResult,
+ ) -> None:
+ try:
+ log = WebhookLog(
+ job_id=event.job_id,
+ event_id=None if is_manual else event.id,
+ webhook_url=event.target_url,
+ attempt_number=event.attempts + 1,
+ request_payload={"header": headers, "payload": payload},
+ signature=headers["X-Knowhere-Signature"],
+ idempotency_key=str(uuid.uuid4()),
+ response_status_code=result.status_code,
+ error_message=result.error_message,
+ duration_ms=result.duration_ms,
+ )
+ db.add(log)
+ await db.commit()
+
+ except Exception as error:
+ logger.error(f"Failed to log webhook delivery: {error}")
diff --git a/packages/shared-python/shared/services/webhook/payload_enrichment.py b/packages/shared-python/shared/services/webhook/payload_enrichment.py
new file mode 100644
index 000000000..2e8d1b73d
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/payload_enrichment.py
@@ -0,0 +1,50 @@
+"""Delivery-time webhook payload enrichment."""
+
+from collections.abc import Mapping
+from typing import Any, cast
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.orm import selectinload
+
+from shared.core.database import get_db_context
+from shared.models.database.job import Job
+from shared.models.database.webhook import WebhookEvent
+from shared.services.jobs.result_delivery import JobResultDeliveryResolver
+
+
+class WebhookPayloadEnricher:
+ """Add fresh Job Result delivery metadata to webhook payloads."""
+
+ def __init__(self, resolver: JobResultDeliveryResolver | None = None) -> None:
+ self._resolver = resolver or JobResultDeliveryResolver()
+
+ async def enrich(self, event: WebhookEvent) -> dict[str, Any]:
+ payload = dict(cast(Mapping[str, Any], event.payload))
+
+ if payload.get("event") != "job.completed":
+ return payload
+
+ try:
+ async with get_db_context() as db:
+ result = await db.execute(
+ select(Job)
+ .options(selectinload(Job.job_result))
+ .where(Job.job_id == event.job_id)
+ )
+ job = result.scalar_one_or_none()
+
+ if not job or not job.job_result:
+ logger.warning(
+ f"Job or result not found for enrichment: job_id={event.job_id}"
+ )
+ return payload
+
+ return self._resolver.enrich_payload(
+ payload,
+ job_result=job.job_result,
+ )
+
+ except Exception as error:
+ logger.error(f"Failed to enrich payload for event {event.id}: {error}")
+ return payload
diff --git a/packages/shared-python/shared/services/webhook/qstash_client.py b/packages/shared-python/shared/services/webhook/qstash_client.py
new file mode 100644
index 000000000..78547d04a
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/qstash_client.py
@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from loguru import logger
+
+from shared.core.config import app_config
+from shared.core.exceptions.domain_exceptions import QStashServiceException
+from shared.models.database.webhook import WebhookEventStatus
+
+
+@dataclass(frozen=True)
+class QStashDeliveryStatus:
+ """Terminal delivery status observed from QStash logs."""
+
+ status: str
+ response_status_code: Optional[int]
+ response_body: Optional[str]
+ error_message: Optional[str]
+
+
+class QStashClientAdapter:
+ """Upstash QStash client adapter for webhook publication and log lookup."""
+
+ def __init__(self) -> None:
+ self._client: Any = None
+
+ def get_client(self) -> Any:
+ """Lazily initialize the QStash client."""
+ if self._client is None:
+ try:
+ from qstash import QStash
+ except ImportError as exc:
+ raise QStashServiceException(
+ internal_message=(
+ "qstash package is required for QStash webhook delivery. "
+ "Install it with: pip install qstash"
+ ),
+ operation="initialize_client",
+ original_exception=exc,
+ ) from exc
+
+ token = app_config.QSTASH_TOKEN
+ if not token:
+ raise QStashServiceException(
+ internal_message="QSTASH_TOKEN is not configured",
+ operation="initialize_client",
+ )
+
+ self._client = QStash(token, base_url=app_config.QSTASH_BASE_URL)
+ return self._client
+
+ def publish_webhook(
+ self,
+ *,
+ target_url: str,
+ payload: dict[str, Any],
+ signature: str,
+ event_id: str,
+ ) -> Optional[str]:
+ """Call the QStash publish API."""
+ headers = {
+ "Content-Type": "application/json",
+ "X-Knowhere-Signature": signature,
+ "X-Knowhere-Event-ID": event_id,
+ "User-Agent": "Knowhere-Webhook/1.0",
+ }
+
+ callback_url = app_config.qstash_callback_url
+ failure_callback_url = app_config.qstash_failure_callback_url
+ if not callback_url or not failure_callback_url:
+ raise QStashServiceException(
+ internal_message=(
+ "QSTASH_CALLBACK_BASE_URL must be configured for QStash "
+ "webhook delivery"
+ ),
+ operation="publish_webhook",
+ )
+
+ publish_kwargs: dict[str, Any] = {
+ "url": target_url,
+ "body": json.dumps(payload, separators=(",", ":")),
+ "headers": headers,
+ "retries": app_config.QSTASH_MAX_RETRIES,
+ "content_type": "application/json",
+ "retry_delay": _get_retry_delay_expression(),
+ "callback": callback_url,
+ "failure_callback": failure_callback_url,
+ "deduplication_id": event_id,
+ "label": "knowhere-webhook",
+ }
+
+ response = self.get_client().message.publish(**publish_kwargs)
+
+ message_id = getattr(response, "message_id", None)
+ if message_id is None and isinstance(response, dict):
+ message_id = response.get("messageId") or response.get("message_id")
+
+ return message_id
+
+ def get_terminal_delivery_status(
+ self,
+ qstash_message_id: str,
+ ) -> Optional[QStashDeliveryStatus]:
+ """Read QStash logs for a terminal destination delivery state."""
+ try:
+ from qstash.log import LogState
+
+ response = self.get_client().log.list(
+ filter={"message_id": qstash_message_id},
+ count=20,
+ )
+ except Exception as exc:
+ logger.warning(
+ f"QStash delivery status lookup failed: "
+ f"message_id={qstash_message_id}, error={exc}"
+ )
+ return None
+
+ terminal_logs = sorted(response.logs, key=lambda log: log.time, reverse=True)
+ for log in terminal_logs:
+ if log.state == LogState.DELIVERED:
+ return QStashDeliveryStatus(
+ status=WebhookEventStatus.DELIVERED,
+ response_status_code=log.response_status,
+ response_body=log.response_body,
+ error_message=log.error,
+ )
+
+ if log.state == LogState.FAILED:
+ return QStashDeliveryStatus(
+ status=WebhookEventStatus.FAILED,
+ response_status_code=log.response_status,
+ response_body=log.response_body,
+ error_message=log.error,
+ )
+
+ return None
+
+
+def _get_retry_delay_expression() -> str:
+ # Approximate exponential backoff: 1m, 10m, ~100m, ~100m, ~100m.
+ return "pow(10, min(retried, 2)) * 60000"
diff --git a/packages/shared-python/shared/services/webhook/qstash_payload.py b/packages/shared-python/shared/services/webhook/qstash_payload.py
new file mode 100644
index 000000000..e454c076b
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/qstash_payload.py
@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.orm import selectinload
+
+from shared.models.database.job import Job
+from shared.services.jobs.result_delivery import JobResultDeliveryResolver
+
+
+class QStashPayloadEnricher:
+ """Sync payload enricher for QStash webhook publication."""
+
+ def __init__(self, resolver: JobResultDeliveryResolver | None = None) -> None:
+ self._resolver = resolver or JobResultDeliveryResolver()
+
+ def enrich(self, db: Any, event: Any) -> dict[str, Any]:
+ payload = dict(event.payload)
+ if payload.get("event") != "job.completed":
+ return payload
+
+ try:
+ result = db.execute(
+ select(Job)
+ .options(selectinload(Job.job_result))
+ .where(Job.job_id == event.job_id)
+ )
+ job = result.scalar_one_or_none()
+ if not job or not job.job_result:
+ return payload
+
+ return self._resolver.enrich_payload(
+ payload,
+ job_result=job.job_result,
+ )
+ except Exception as exc:
+ logger.error(f"Failed to enrich payload for event {event.id}: {exc}")
+ return payload
diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py
index 8150100b2..7a0580cb6 100644
--- a/packages/shared-python/shared/services/webhook/qstash_publisher.py
+++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py
@@ -11,63 +11,39 @@
from __future__ import annotations
-import hashlib
-import hmac
-import json
-import time
-from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from typing import Optional
from loguru import logger
-
-from shared.core.config import app_config
-from shared.core.exceptions.domain_exceptions import QStashServiceException
-from shared.models.database.webhook import WebhookEventStatus
+from sqlalchemy import select
+
+from shared.core.database_sync import get_sync_db_context
+from shared.models.database.job import Job
+from shared.models.database.webhook import WebhookEvent, WebhookEventStatus
+from shared.services.webhook.qstash_client import (
+ QStashClientAdapter,
+ QStashDeliveryStatus,
+)
+from shared.services.webhook.qstash_payload import QStashPayloadEnricher
+from shared.services.webhook.qstash_secret_resolver import QStashSecretResolver
+from shared.services.webhook.signing import sign_webhook_payload
from shared.utils.url_security import (
validate_http_url_and_resolve_ip,
)
-@dataclass(frozen=True)
-class QStashDeliveryStatus:
- """Terminal delivery status observed from QStash logs."""
-
- status: str
- response_status_code: Optional[int]
- response_body: Optional[str]
- error_message: Optional[str]
-
-
class QStashWebhookPublisher:
"""Publishes webhook events to customer endpoints via QStash."""
- def __init__(self) -> None:
- self._client: Any = None
-
- def _get_client(self) -> Any:
- """Lazily initialize the QStash client."""
- if self._client is None:
- try:
- from qstash import QStash
- except ImportError as exc:
- raise QStashServiceException(
- internal_message=(
- "qstash package is required for QStash webhook delivery. "
- "Install it with: pip install qstash"
- ),
- operation="initialize_client",
- original_exception=exc,
- ) from exc
-
- token = app_config.QSTASH_TOKEN
- if not token:
- raise QStashServiceException(
- internal_message="QSTASH_TOKEN is not configured",
- operation="initialize_client",
- )
-
- self._client = QStash(token, base_url=app_config.QSTASH_BASE_URL)
- return self._client
+ def __init__(
+ self,
+ *,
+ client_adapter: QStashClientAdapter | None = None,
+ payload_enricher: QStashPayloadEnricher | None = None,
+ secret_resolver: QStashSecretResolver | None = None,
+ ) -> None:
+ self._client_adapter = client_adapter or QStashClientAdapter()
+ self._payload_enricher = payload_enricher or QStashPayloadEnricher()
+ self._secret_resolver = secret_resolver or QStashSecretResolver()
def publish_event(self, event_id: str) -> Optional[str]:
"""Publish a webhook event via QStash.
@@ -77,12 +53,6 @@ def publish_event(self, event_id: str) -> Optional[str]:
Returns the QStash message_id on success, or None on failure.
"""
- from sqlalchemy import select
-
- from shared.core.database_sync import get_sync_db_context
- from shared.models.database.job import Job
- from shared.models.database.webhook import WebhookEvent
-
with get_sync_db_context() as db:
event = db.execute(
select(WebhookEvent).where(WebhookEvent.id == event_id)
@@ -109,10 +79,8 @@ def publish_event(self, event_id: str) -> Optional[str]:
db.commit()
return None
- # Enrich payload (presigned S3 URL for completed jobs)
- payload = self._enrich_payload(db, event)
+ payload = self._payload_enricher.enrich(db, event)
- # Resolve signing secret
user_id = db.execute(
select(Job.user_id).where(Job.job_id == event.job_id)
).scalar_one_or_none()
@@ -123,7 +91,11 @@ def publish_event(self, event_id: str) -> Optional[str]:
db.commit()
return None
- secret = self._resolve_secret(db, str(user_id), event.target_url)
+ secret = self._secret_resolver.resolve(
+ db,
+ user_id=str(user_id),
+ endpoint=event.target_url,
+ )
if not secret:
logger.error(
f"QStash publish: secret resolution failed for event {event_id}"
@@ -132,12 +104,10 @@ def publish_event(self, event_id: str) -> Optional[str]:
db.commit()
return None
- # Sign payload with our HMAC
- signature = self._sign_payload(payload, secret)
+ signature = sign_webhook_payload(payload, secret)
- # Publish to QStash
try:
- message_id = self._publish_to_qstash(
+ message_id = self._client_adapter.publish_webhook(
target_url=event.target_url,
payload=payload,
signature=signature,
@@ -159,210 +129,12 @@ def publish_event(self, event_id: str) -> Optional[str]:
)
return message_id
- def _publish_to_qstash(
- self,
- target_url: str,
- payload: Dict[str, Any],
- signature: str,
- event_id: str,
- ) -> Optional[str]:
- """Call the QStash publish API."""
- headers = {
- "Content-Type": "application/json",
- "X-Knowhere-Signature": signature,
- "X-Knowhere-Event-ID": event_id,
- "User-Agent": "Knowhere-Webhook/1.0",
- }
-
- # Approximate exponential backoff: 1m, 10m, ~100m, ~100m, ~100m
- # pow(10, min(retried, 2)) * 60000 → 60s, 600s, 6000s capped
- retry_delay_expression = "pow(10, min(retried, 2)) * 60000"
-
- callback_url = app_config.qstash_callback_url
- failure_callback_url = app_config.qstash_failure_callback_url
- if not callback_url or not failure_callback_url:
- raise QStashServiceException(
- internal_message=(
- "QSTASH_CALLBACK_BASE_URL must be configured for QStash "
- "webhook delivery"
- ),
- operation="publish_webhook",
- )
-
- client = self._get_client()
-
- publish_kwargs: Dict[str, Any] = {
- "url": target_url,
- "body": json.dumps(payload, separators=(",", ":")),
- "headers": headers,
- "retries": app_config.QSTASH_MAX_RETRIES,
- "content_type": "application/json",
- "retry_delay": retry_delay_expression,
- "callback": callback_url,
- "failure_callback": failure_callback_url,
- "deduplication_id": event_id,
- "label": "knowhere-webhook",
- }
-
- response = client.message.publish(**publish_kwargs)
-
- message_id = getattr(response, "message_id", None)
- if message_id is None and isinstance(response, dict):
- message_id = response.get("messageId") or response.get("message_id")
-
- return message_id
-
def get_terminal_delivery_status(
self,
qstash_message_id: str,
) -> Optional[QStashDeliveryStatus]:
"""Read QStash logs for a terminal destination delivery state."""
- try:
- from qstash.log import LogState
-
- response = self._get_client().log.list(
- filter={"message_id": qstash_message_id},
- count=20,
- )
- except Exception as exc:
- logger.warning(
- f"QStash delivery status lookup failed: "
- f"message_id={qstash_message_id}, error={exc}"
- )
- return None
-
- terminal_logs = sorted(response.logs, key=lambda log: log.time, reverse=True)
- for log in terminal_logs:
- if log.state == LogState.DELIVERED:
- return QStashDeliveryStatus(
- status=WebhookEventStatus.DELIVERED,
- response_status_code=log.response_status,
- response_body=log.response_body,
- error_message=log.error,
- )
-
- if log.state == LogState.FAILED:
- return QStashDeliveryStatus(
- status=WebhookEventStatus.FAILED,
- response_status_code=log.response_status,
- response_body=log.response_body,
- error_message=log.error,
- )
-
- return None
-
- def _enrich_payload(self, db: Any, event: Any) -> Dict[str, Any]:
- """Enrich the webhook payload (e.g., generate fresh presigned S3 URL)."""
- from sqlalchemy import select
- from sqlalchemy.orm import selectinload
-
- from shared.models.database.job import Job
-
- payload = dict(event.payload)
- if payload.get("event") != "job.completed":
- return payload
-
- try:
- result = db.execute(
- select(Job)
- .options(selectinload(Job.job_result))
- .where(Job.job_id == event.job_id)
- )
- job = result.scalar_one_or_none()
- if not job or not job.job_result:
- return payload
-
- job_result = job.job_result
- if job_result.result_s3_key:
- payload["result_url"] = app_config.get_storage_adapter().generate_presigned_url(
- job_result.result_s3_key,
- expiration=3600,
- method="GET",
- )
-
- if job_result.inline_payload:
- payload["result"] = job_result.inline_payload
- except Exception as exc:
- logger.error(f"Failed to enrich payload for event {event.id}: {exc}")
-
- return payload
-
- def _resolve_secret(self, db: Any, user_id: str, endpoint: str) -> Optional[str]:
- """Resolve the webhook signing secret for a user/endpoint."""
- from datetime import datetime, timezone
-
- from sqlalchemy import and_, select
-
- from shared.core.exceptions.domain_exceptions import (
- SystemSettingInvalidException,
- SystemSettingMissingException,
- )
- from shared.models.database.webhook_secret import (
- WebhookSecret,
- WebhookSecretStatus,
- )
- from shared.services.encryption import get_fernet_service
-
- try:
- fernet = get_fernet_service()
- except (SystemSettingMissingException, SystemSettingInvalidException) as exc:
- logger.error(f"Configuration error during secret resolution: {exc}")
- return None
-
- # Try endpoint-specific secret first, then global
- secret_obj = None
- if endpoint:
- result = db.execute(
- select(WebhookSecret).where(
- and_(
- WebhookSecret.user_id == user_id,
- WebhookSecret.endpoint == endpoint,
- WebhookSecret.status == WebhookSecretStatus.ACTIVE,
- )
- )
- )
- secret_obj = result.scalar_one_or_none()
-
- if secret_obj is None:
- result = db.execute(
- select(WebhookSecret).where(
- and_(
- WebhookSecret.user_id == user_id,
- WebhookSecret.endpoint.is_(None),
- WebhookSecret.status == WebhookSecretStatus.ACTIVE,
- )
- )
- )
- secret_obj = result.scalar_one_or_none()
-
- if secret_obj is None:
- raw_secret = fernet.generate_webhook_secret()
- secret_obj = WebhookSecret(
- user_id=user_id,
- endpoint=endpoint,
- secret_encrypted=fernet.encrypt(raw_secret),
- status=WebhookSecretStatus.ACTIVE,
- )
- db.add(secret_obj)
- db.commit()
- db.refresh(secret_obj)
-
- secret_obj.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None)
- db.add(secret_obj)
- return fernet.decrypt(secret_obj.secret_encrypted)
-
- @staticmethod
- def _sign_payload(payload: Dict[str, Any], secret: str) -> str:
- """Generate HMAC-SHA256 signature matching the existing Knowhere format."""
- timestamp = int(time.time())
- payload_str = json.dumps(payload, separators=(",", ":"))
- signed_content = f"{timestamp}.{payload_str}"
- sig = hmac.new(
- secret.encode("utf-8"),
- signed_content.encode("utf-8"),
- hashlib.sha256,
- ).hexdigest()
- return f"t={timestamp},v1={sig}"
+ return self._client_adapter.get_terminal_delivery_status(qstash_message_id)
# Module-level singleton
diff --git a/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py b/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py
new file mode 100644
index 000000000..fb8174883
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/qstash_secret_resolver.py
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any, Optional
+
+from loguru import logger
+from sqlalchemy import and_, select
+
+from shared.core.exceptions.domain_exceptions import (
+ SystemSettingInvalidException,
+ SystemSettingMissingException,
+)
+from shared.models.database.webhook_secret import (
+ WebhookSecret,
+ WebhookSecretStatus,
+)
+from shared.services.encryption import get_fernet_service
+
+
+class QStashSecretResolver:
+ """Sync webhook secret resolver for QStash publication."""
+
+ def resolve(self, db: Any, *, user_id: str, endpoint: str) -> Optional[str]:
+ try:
+ fernet = get_fernet_service()
+ except (SystemSettingMissingException, SystemSettingInvalidException) as exc:
+ logger.error(f"Configuration error during secret resolution: {exc}")
+ return None
+
+ secret = self._find_active_secret(db, user_id=user_id, endpoint=endpoint)
+ if secret is None:
+ raw_secret = fernet.generate_webhook_secret()
+ secret = WebhookSecret(
+ user_id=user_id,
+ endpoint=endpoint,
+ secret_encrypted=fernet.encrypt(raw_secret),
+ status=WebhookSecretStatus.ACTIVE,
+ )
+ db.add(secret)
+ db.commit()
+ db.refresh(secret)
+
+ secret.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None)
+ db.add(secret)
+ return fernet.decrypt(secret.secret_encrypted)
+
+ def _find_active_secret(
+ self,
+ db: Any,
+ *,
+ user_id: str,
+ endpoint: str,
+ ) -> WebhookSecret | None:
+ if endpoint:
+ result = db.execute(
+ select(WebhookSecret).where(
+ and_(
+ WebhookSecret.user_id == user_id,
+ WebhookSecret.endpoint == endpoint,
+ WebhookSecret.status == WebhookSecretStatus.ACTIVE,
+ )
+ )
+ )
+ secret = result.scalar_one_or_none()
+ if secret is not None:
+ return secret
+
+ result = db.execute(
+ select(WebhookSecret).where(
+ and_(
+ WebhookSecret.user_id == user_id,
+ WebhookSecret.endpoint.is_(None),
+ WebhookSecret.status == WebhookSecretStatus.ACTIVE,
+ )
+ )
+ )
+ return result.scalar_one_or_none()
diff --git a/packages/shared-python/shared/services/webhook/secret_resolver.py b/packages/shared-python/shared/services/webhook/secret_resolver.py
new file mode 100644
index 000000000..d0a648ef3
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/secret_resolver.py
@@ -0,0 +1,60 @@
+"""Webhook secret resolution for direct deliveries."""
+
+from datetime import datetime, timezone
+
+from loguru import logger
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.core.exceptions.domain_exceptions import (
+ SystemSettingInvalidException,
+ SystemSettingMissingException,
+)
+from shared.models.database.job import Job
+from shared.models.database.webhook import WebhookEvent
+from shared.repositories.webhook_secret_repository import WebhookSecretRepository
+
+
+class WebhookSecretResolver:
+ """Resolve the active endpoint secret for a WebhookEvent delivery."""
+
+ def __init__(self, repository: WebhookSecretRepository | None = None) -> None:
+ self._repository = repository or WebhookSecretRepository()
+
+ async def resolve_for_event(
+ self, db: AsyncSession, event: WebhookEvent
+ ) -> str | None:
+ user_id = await self._get_job_owner(db, event.job_id)
+ if not user_id:
+ logger.warning(f"Could not resolve secret: Job {event.job_id} has no user_id")
+ return None
+
+ return await self.resolve_for_endpoint(
+ db,
+ user_id=user_id,
+ endpoint=event.target_url,
+ )
+
+ async def resolve_for_endpoint(
+ self, db: AsyncSession, *, user_id: str, endpoint: str
+ ) -> str | None:
+ try:
+ secret = await self._repository.get_or_create_secret(
+ db, user_id, endpoint=endpoint
+ )
+
+ if secret:
+ secret.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None)
+ db.add(secret)
+
+ return self._repository.decrypt_secret(secret)
+
+ except (SystemSettingMissingException, SystemSettingInvalidException):
+ raise
+ except Exception as error:
+ logger.error(f"Failed to resolve/create secret for user {user_id}: {error}")
+ return None
+
+ async def _get_job_owner(self, db: AsyncSession, job_id: str) -> str | None:
+ result = await db.execute(select(Job.user_id).where(Job.job_id == job_id))
+ return result.scalar_one_or_none()
diff --git a/packages/shared-python/shared/services/webhook/signing.py b/packages/shared-python/shared/services/webhook/signing.py
new file mode 100644
index 000000000..15c521946
--- /dev/null
+++ b/packages/shared-python/shared/services/webhook/signing.py
@@ -0,0 +1,34 @@
+"""Webhook request signing."""
+
+import hashlib
+import hmac
+import json
+import time
+from collections.abc import Mapping
+from typing import Any
+
+
+def sign_webhook_payload(payload: Mapping[str, Any], secret: str) -> str:
+ """Generate the timestamped Knowhere webhook HMAC signature."""
+ timestamp = int(time.time())
+ payload_text = json.dumps(payload, separators=(",", ":"))
+ signed_content = f"{timestamp}.{payload_text}"
+ signature = hmac.new(
+ secret.encode("utf-8"),
+ signed_content.encode("utf-8"),
+ hashlib.sha256,
+ ).hexdigest()
+
+ return f"t={timestamp},v1={signature}"
+
+
+def build_webhook_headers(
+ *, payload: Mapping[str, Any], secret: str, attempt_id: str
+) -> dict[str, str]:
+ """Build signed HTTP headers for a direct webhook delivery attempt."""
+ return {
+ "Content-Type": "application/json",
+ "X-Knowhere-Signature": sign_webhook_payload(payload, secret),
+ "X-Knowhere-Attempt-ID": attempt_id,
+ "User-Agent": "Knowhere-Webhook/1.0",
+ }
diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py
index 11aa49efb..53c139510 100644
--- a/packages/shared-python/shared/testing/contract_runtime.py
+++ b/packages/shared-python/shared/testing/contract_runtime.py
@@ -358,6 +358,24 @@ def clear_application_modules() -> None:
sys.modules.pop(module_name, None)
continue
+ if module_name == "shared.services.storage" or module_name.startswith(
+ "shared.services.storage."
+ ):
+ sys.modules.pop(module_name, None)
+ continue
+
+ if module_name == "shared.services.jobs" or module_name.startswith(
+ "shared.services.jobs."
+ ):
+ sys.modules.pop(module_name, None)
+ continue
+
+ if module_name == "shared.services.webhook" or module_name.startswith(
+ "shared.services.webhook."
+ ):
+ sys.modules.pop(module_name, None)
+ continue
+
if module_name == "app" or module_name.startswith("app."):
sys.modules.pop(module_name, None)
diff --git a/packages/shared-python/shared/utils/CommonHelper.py b/packages/shared-python/shared/utils/CommonHelper.py
deleted file mode 100644
index 2790c7a6a..000000000
--- a/packages/shared-python/shared/utils/CommonHelper.py
+++ /dev/null
@@ -1,45 +0,0 @@
-from io import BytesIO
-from pathlib import Path
-
-import httpx
-import pandas as pd
-from starlette.datastructures import UploadFile as StarletteUploadFile
-
-from shared.utils.FileDownUpUtils import s3_upload_file
-
-
-def is_remote(path):
- """Check whether a path is a remote URL."""
- if path is None:
- return False
- if not isinstance(path, str):
- return False
- return path.startswith("http://") or path.startswith("https://")
-
-
-async def load_file_bytes(file_path, *, file_url="", timeout=None):
- if isinstance(file_path, str) and is_remote(file_path):
- # If file_path is already a full URL, use it directly.
- url_to_use = file_path
- if not isinstance(file_url, str):
- file_url = file_url.geturl()
- # Prefer file_url when provided; otherwise keep file_path.
- if file_url:
- url_to_use = file_url
- async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
- r = await client.get(url_to_use) # Fetch the resolved URL.
- r.raise_for_status()
- return r.content
- else:
- p = Path(file_path)
- return p.read_bytes()
-
-
-async def upload_dataframe_to_s3(df: pd.DataFrame, filename: str, prefix: str):
- # Write the DataFrame into an in-memory BytesIO buffer.
- buffer = BytesIO()
- df.to_csv(buffer, index=False)
- buffer.seek(0) # Reset the cursor to the buffer start.
-
- upload_file = StarletteUploadFile(file=buffer, filename=filename)
- s3_upload_file(upload_file, prefix)
diff --git a/packages/shared-python/shared/utils/FileDownUpUtils.py b/packages/shared-python/shared/utils/FileDownUpUtils.py
deleted file mode 100644
index 0f124ebe8..000000000
--- a/packages/shared-python/shared/utils/FileDownUpUtils.py
+++ /dev/null
@@ -1,253 +0,0 @@
-import os
-import uuid
-import zipfile
-from pathlib import Path
-from typing import Optional, Union
-from urllib.parse import urljoin
-
-import aiohttp
-import requests
-from botocore.exceptions import ClientError
-from starlette.datastructures import UploadFile
-
-from shared.core.config import settings
-from shared.core.config.storage import get_cached_storage_adapter
-from shared.core.exceptions.domain_exceptions import (
- KnowhereException,
- NotFoundException,
- StorageServiceException,
-)
-from shared.models.schemas.s3_file import FliesDownload
-
-
-def s3_upload_file(file: UploadFile, prefix: str):
- """
- Upload a file object to S3 storage.
- :param file: Input file such as ``abc15sa25ww.doc``
- :param prefix: Storage prefix such as ``upload/123``
- :return: Upload result payload
- """
- if prefix and not prefix.endswith("/"):
- prefix += "/"
- object_key = f"{prefix}{file.filename}"
- adapter = get_cached_storage_adapter()
- try:
- # ``upload_fileobj`` streams efficiently and avoids large in-memory copies.
- adapter.upload_fileobj(
- file.file, object_key, content_type="application/octet-stream"
- )
- public_url = (
- f"{settings.S3_PRIVATE_DOMAIN}/{object_key}"
- if settings.S3_PRIVATE_DOMAIN
- else f"storage/{object_key}"
- )
- content = {
- "message": "File uploaded successfully",
- "bucket": settings.S3_BUCKET_NAME,
- "file_key": object_key,
- "public_url_for_reference": public_url,
- }
- return content
-
- except KnowhereException:
- raise
- except Exception as e:
- # Wrap storage upload failures in a domain exception.
- raise StorageServiceException(
- internal_message=f"Storage upload failed: {str(e)}",
- operation="upload",
- original_exception=e,
- )
-
-
-def s3_download_extract_zip(
- url: str,
- dest_dir: Union[str, os.PathLike],
- *,
- filename: str = "parsed.zip",
- headers: Optional[dict] = None,
- timeout: int | None = None,
- chunk_size: int | None = None,
- keep_exts: tuple[str, ...] = (".md", ".json"),
- exclude_patterns: tuple[str, ...] = (),
- clean_empty_dirs: bool = True,
-):
- """
- Download and extract a zip file, keeping only specific file types.
-
- Args:
- exclude_patterns: Tuple of filename patterns to exclude (e.g., ("content_list", "middle.json"))
- """
- import fnmatch
-
- from shared.core.constants import APIConstants, ProcessingConstants
-
- # Use defaults when optional arguments are omitted.
- if timeout is None:
- timeout = APIConstants.S3_FILE_DOWNLOAD_TIMEOUT
- if chunk_size is None:
- chunk_size = ProcessingConstants.IMG_CHUNK_SIZE
-
- dest_dir = Path(dest_dir).expanduser().resolve()
- dest_dir.mkdir(parents=True, exist_ok=True)
- zip_path = dest_dir / filename
-
- # 1) Download to zip_path and extract.
- with requests.get(
- url, headers=headers or {}, timeout=timeout, stream=True, allow_redirects=True
- ) as r:
- r.raise_for_status()
- with open(zip_path, "wb") as f:
- for chunk in r.iter_content(chunk_size=chunk_size):
- if chunk:
- f.write(chunk)
-
- with zipfile.ZipFile(zip_path, "r") as zf:
- zf.extractall(dest_dir)
-
- # 2) Remove files outside keep_exts or matching exclude_patterns.
- kept_files = []
- for p in dest_dir.rglob("*"):
- if p.is_file():
- # Check if file should be excluded by pattern
- should_exclude = False
- for pattern in exclude_patterns:
- if pattern in p.name or fnmatch.fnmatch(p.name, pattern):
- should_exclude = True
- break
-
- if should_exclude:
- p.unlink()
- elif p.suffix.lower() in keep_exts:
- kept_files.append(p)
- else:
- p.unlink()
-
- # 4) Remove empty directories when requested.
- if clean_empty_dirs:
- for d in sorted(
- [p for p in dest_dir.rglob("*") if p.is_dir()],
- key=lambda x: len(x.parts),
- reverse=True,
- ):
- try:
- next(d.iterdir())
- except StopIteration:
- d.rmdir()
- # 5) Delete the downloaded zip file.
- zip_path.unlink(missing_ok=True)
-
-
-def s3_get_download_url(file_key: str, expires_in: int = 3600):
- """
- Get a file download URL from its storage key.
- :param file_key: Full file path and name
- :param expires_in: Desired URL lifetime
- :return: Signed download payload
- """
- s3_client = settings.get_s3_client()
- try:
- # Generate a pre-signed URL.
- presigned_url = s3_client.generate_presigned_url(
- "get_object",
- Params={"Bucket": settings.S3_BUCKET_NAME, "Key": file_key},
- ExpiresIn=expires_in, # URL lifetime.
- )
- fsdl = FliesDownload(
- message="URL signed successfully",
- file_key=file_key,
- download_url=presigned_url,
- expires_in_seconds=expires_in,
- )
- return fsdl
-
- except ClientError as e:
- # boto3 may still sign missing objects; the resulting URL can later 404.
- raise NotFoundException(
- resource="File",
- resource_id=file_key,
- internal_message=(
- f"Could not generate the URL. Check whether the file is correct "
- f"or the S3 configuration is valid: {str(e)}"
- ),
- )
-
-
-def get_url_file(path):
- file_sig = s3_get_download_url(path, expires_in=3600)
- # Assemble the final URL.
- file_url = file_sig.download_url
- response = requests.get(file_url, verify=True)
- response.raise_for_status() # Ensure the request succeeds.
- return response
-
-
-def get_pub_fileurl(path):
- """
- Build a public URL from a storage path.
- :param path:
- :return: Public URL
- """
- base_url = settings.S3_PRIVATE_DOMAIN.rstrip("/")
- clean_path = path.replace("\\", "/").strip()
- full_url = urljoin(base_url + "/", clean_path)
- return full_url
-
-
-def s3_public_file_url(file_key: str) -> str:
- permanent_url = f"{settings.S3_PRIVATE_DOMAIN}/{settings.S3_BUCKET_NAME}/{file_key}"
- return permanent_url
-
-
-async def download_and_upload_image(
- img_url: str, prefix: str = "images/", temp_store_path=None
-) -> dict:
- """
- Download an image, rename it, upload it to S3, and clean up locally.
- :param img_url: Image URL
- :param prefix: S3 storage prefix
- :return: Dict containing upload results and the new download reference
- """
- # Generate a unique filename.
- unique_filename = f"{uuid.uuid4()}.jpg"
- # Temporary directory.
- if temp_store_path is None:
- temp_store_path = r"/Volumes/U/temp/output/"
- local_file_path = Path(f"{settings.S3_TEMP_PATH or '/tmp'}{unique_filename}")
- # Path(f"{temp_store_path}{unique_filename}")
- try:
- # Download the image asynchronously.
- async with aiohttp.ClientSession() as session:
- async with session.get(img_url) as response:
- response.raise_for_status()
- with open(local_file_path, "wb") as f:
- f.write(await response.read())
-
- # Create a temporary UploadFile wrapper.
- from fastapi import UploadFile
-
- upload_file = UploadFile(
- filename=unique_filename, file=open(local_file_path, "rb")
- )
-
- # Upload to S3.
- result = s3_upload_file(upload_file, prefix)
-
- # Close the file handle and delete the local file.
- upload_file.file.close()
- os.remove(local_file_path)
- return result
-
- except KnowhereException:
- if local_file_path.exists():
- os.remove(local_file_path)
- raise
- except Exception as e:
- # Always remove the local file on failure as well.
- if local_file_path.exists():
- os.remove(local_file_path)
- raise StorageServiceException(
- internal_message=f"Failed to download and upload the image: {str(e)}",
- operation="download_and_upload",
- original_exception=e,
- )
diff --git a/packages/shared-python/shared/utils/CommonHelperSync.py b/packages/shared-python/shared/utils/file_loading.py
similarity index 73%
rename from packages/shared-python/shared/utils/CommonHelperSync.py
rename to packages/shared-python/shared/utils/file_loading.py
index 924501ddb..522516cea 100644
--- a/packages/shared-python/shared/utils/CommonHelperSync.py
+++ b/packages/shared-python/shared/utils/file_loading.py
@@ -1,22 +1,24 @@
-"""Sync helpers for gevent worker code paths.
-
-Keep API async helpers in `CommonHelper.py`; worker should import this module.
-"""
+"""Sync file-loading helpers for worker parsing paths."""
from pathlib import Path
-from typing import Optional
+from urllib.parse import ParseResult
import httpx
-def is_remote(path):
+def is_remote(path: object) -> bool:
"""Return True if `path` is an HTTP(S) URL."""
if path is None or not isinstance(path, str):
return False
return path.startswith("http://") or path.startswith("https://")
-def load_file_bytes(file_path, *, file_url: str = "", timeout: Optional[float] = None):
+def load_file_bytes(
+ file_path: str | Path,
+ *,
+ file_url: str | ParseResult = "",
+ timeout: float | None = None,
+) -> bytes:
"""Load bytes from local path or remote URL synchronously."""
if isinstance(file_path, str) and is_remote(file_path):
url_to_use = file_path
diff --git a/packages/shared-python/shared/utils/file_transfer.py b/packages/shared-python/shared/utils/file_transfer.py
deleted file mode 100644
index d307e7995..000000000
--- a/packages/shared-python/shared/utils/file_transfer.py
+++ /dev/null
@@ -1,171 +0,0 @@
-"""
-File Transfer Utilities
-
-Provides reliable file transfer operations for large files using temp files as buffers.
-Uses httpx for proper total timeout enforcement.
-"""
-
-import os
-import tempfile
-from typing import Dict, Optional
-from urllib.parse import urlparse
-
-import httpx
-from loguru import logger
-
-from shared.utils.http_clients import get_sync_client
-
-
-class FileTransferError(Exception):
- """Base exception for file transfer operations"""
-
- def __init__(self, message: str, status_code: Optional[int] = None):
- super().__init__(message)
- self.status_code = status_code
-
-
-class DownloadError(FileTransferError):
- """
- Download failed - typically a client error.
-
- The source file may be inaccessible, expired, or invalid.
- Worker should raise a 4xx (client error) when catching this.
- """
-
- pass
-
-
-class UploadError(FileTransferError):
- """
- Upload failed - typically a server/service error.
-
- The target service (e.g., MinerU) may be unavailable or experiencing issues.
- Worker should raise a 5xx (server error) when catching this.
- """
-
- pass
-
-
-def stream_download_and_upload(
- source_url: str,
- target_url: str,
- download_timeout: int = 300,
- upload_timeout: int = 300,
- chunk_size: int = 8192,
- upload_method: str = "PUT",
- upload_headers: Optional[Dict[str, str]] = None,
- upload_retries: int = 3,
-) -> httpx.Response:
- """
- Download a file from source_url and upload to target_url using a temp file buffer.
-
- Uses httpx for proper total timeout enforcement.
- Retries upload on failure since temp file is preserved on disk.
-
- Args:
- source_url: URL to download the file from
- target_url: URL to upload the file to
- download_timeout: Total timeout for download in seconds
- upload_timeout: Total timeout for upload in seconds
- chunk_size: Chunk size for streaming download
- upload_method: HTTP method for upload (PUT or POST)
- upload_headers: Additional headers for upload request
- upload_retries: Number of retry attempts for upload (default 3)
-
- Returns:
- httpx.Response: The upload response
-
- Raises:
- DownloadError: If download fails (source inaccessible)
- UploadError: If upload fails after all retries
- """
- # Create temp file manually for explicit cleanup control
- tmp_fd, tmp_path = tempfile.mkstemp(suffix=".tmp")
- source_host = urlparse(source_url).hostname or source_url[:60]
- target_host = urlparse(target_url).hostname or target_url[:60]
-
- try:
- # Phase 1: Download to temp file
- logger.debug(f"Downloading from {source_url[:100]}...")
- try:
- client = get_sync_client()
- with client.stream("GET", source_url, timeout=download_timeout) as response:
- response.raise_for_status()
-
- downloaded_bytes = 0
- with os.fdopen(tmp_fd, "wb") as tmp_file:
- for chunk in response.iter_bytes(chunk_size=chunk_size):
- tmp_file.write(chunk)
- downloaded_bytes += len(chunk)
-
- except httpx.TimeoutException as e:
- raise DownloadError(
- f"Download timed out: host={source_host}, timeout={download_timeout}s"
- ) from e
- except httpx.HTTPStatusError as e:
- raise DownloadError(
- f"Download failed: host={source_host}, status={e.response.status_code}"
- ) from e
- except httpx.RequestError as e:
- raise DownloadError(
- f"Download failed: host={source_host}, error={e}"
- ) from e
-
- # Get file size
- file_size = os.path.getsize(tmp_path)
- logger.info(f"Downloaded {file_size} bytes to temp file")
-
- # Phase 2: Upload from temp file (with retries)
- headers = upload_headers or {}
- headers["Content-Length"] = str(file_size)
-
- last_error = None
- for attempt in range(1, upload_retries + 1):
- try:
- logger.info(
- f"Uploading {file_size} bytes (attempt {attempt}/{upload_retries}, timeout={upload_timeout}s)..."
- )
-
- # Stream directly from file without loading to memory
- with open(tmp_path, "rb") as f:
- client = get_sync_client()
- if upload_method.upper() == "PUT":
- upload_response = client.put(
- target_url,
- content=f,
- headers=headers,
- timeout=upload_timeout,
- )
- else:
- upload_response = client.post(
- target_url,
- content=f,
- headers=headers,
- timeout=upload_timeout,
- )
-
- logger.info(f"Upload completed: status={upload_response.status_code}")
- return upload_response
-
- except (httpx.TimeoutException, httpx.RequestError) as e:
- last_error = e
- logger.warning(
- f"Upload attempt {attempt} failed: host={target_host}, error={e}"
- )
- if attempt < upload_retries:
- logger.info("Retrying upload...")
- continue
-
- # All retries exhausted
- raise UploadError(
- f"Upload failed: host={target_host}, attempts={upload_retries}, last_error={last_error}"
- ) from last_error
-
- finally:
- # Manual cleanup of temp file
- if os.path.exists(tmp_path):
- try:
- os.remove(tmp_path)
- logger.debug(f"Temp file cleaned up: {tmp_path}")
- except OSError as e:
- logger.warning(f"Failed to cleanup temp file {tmp_path}: {e}")
diff --git a/packages/shared-python/shared/utils/zip_download.py b/packages/shared-python/shared/utils/zip_download.py
new file mode 100644
index 000000000..630173c2f
--- /dev/null
+++ b/packages/shared-python/shared/utils/zip_download.py
@@ -0,0 +1,79 @@
+"""Download-and-extract helpers for remote ZIP artifacts."""
+
+import os
+import zipfile
+from pathlib import Path
+from collections.abc import Mapping
+
+import requests
+
+
+def download_and_extract_zip(
+ url: str,
+ dest_dir: str | os.PathLike[str],
+ *,
+ filename: str = "parsed.zip",
+ headers: Mapping[str, str] | None = None,
+ timeout: int | None = None,
+ chunk_size: int | None = None,
+ keep_exts: tuple[str, ...] = (".md", ".json"),
+ exclude_patterns: tuple[str, ...] = (),
+ clean_empty_dirs: bool = True,
+) -> None:
+ """Download a ZIP file, extract it, and keep only the requested artifacts."""
+ import fnmatch
+
+ from shared.core.constants import APIConstants, ProcessingConstants
+
+ if timeout is None:
+ timeout = APIConstants.S3_FILE_DOWNLOAD_TIMEOUT
+ if chunk_size is None:
+ chunk_size = ProcessingConstants.IMG_CHUNK_SIZE
+
+ destination = Path(dest_dir).expanduser().resolve()
+ destination.mkdir(parents=True, exist_ok=True)
+ zip_path = destination / filename
+
+ with requests.get(
+ url,
+ headers=headers or {},
+ timeout=timeout,
+ stream=True,
+ allow_redirects=True,
+ ) as response:
+ response.raise_for_status()
+ with open(zip_path, "wb") as zip_file:
+ for chunk in response.iter_content(chunk_size=chunk_size):
+ if chunk:
+ zip_file.write(chunk)
+
+ with zipfile.ZipFile(zip_path, "r") as extracted_zip:
+ extracted_zip.extractall(destination)
+
+ for extracted_path in destination.rglob("*"):
+ if not extracted_path.is_file():
+ continue
+
+ should_exclude = False
+ for pattern in exclude_patterns:
+ if pattern in extracted_path.name or fnmatch.fnmatch(extracted_path.name, pattern):
+ should_exclude = True
+ break
+
+ if should_exclude:
+ extracted_path.unlink()
+ elif extracted_path.suffix.lower() not in keep_exts:
+ extracted_path.unlink()
+
+ if clean_empty_dirs:
+ for directory in sorted(
+ [path for path in destination.rglob("*") if path.is_dir()],
+ key=lambda path: len(path.parts),
+ reverse=True,
+ ):
+ try:
+ next(directory.iterdir())
+ except StopIteration:
+ directory.rmdir()
+
+ zip_path.unlink(missing_ok=True)
|