534 lines
17 KiB
Python
534 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
D&D 5e SRD Context Expansion Tool
|
|
|
|
This tool expands the context around search results from search_with_positions.py,
|
|
providing larger, structured views of the source material. Designed for LLM consumption.
|
|
|
|
Usage:
|
|
# Expand specific result from a search
|
|
python expand_context.py "fireball" --result 3 --mode section --all
|
|
|
|
# Expand multiple results
|
|
python expand_context.py "fireball" --results 1,3,5 --mode paragraph --all
|
|
|
|
# Direct expansion from file position
|
|
python expand_context.py --file "DND5eSRD_121-137.md" --position 1234 --mode section
|
|
|
|
# Get full document structure with position
|
|
python expand_context.py --file "DND5eSRD_121-137.md" --position 1234 --mode document
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
# Import from search tool in the same directory
|
|
import sys
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
script_dir = Path(__file__).parent
|
|
sys.path.insert(0, str(script_dir))
|
|
from search_with_positions import (
|
|
get_all_reference_files,
|
|
get_files_by_page_range,
|
|
get_references_dir,
|
|
search_files,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class HeadingNode:
|
|
"""Represents a heading in the document hierarchy."""
|
|
|
|
level: int
|
|
text: str
|
|
start_pos: int
|
|
end_pos: int
|
|
content: str = ""
|
|
children: List["HeadingNode"] = None
|
|
|
|
def __post_init__(self):
|
|
if self.children is None:
|
|
self.children = []
|
|
|
|
|
|
@dataclass
|
|
class ExpandedContext:
|
|
"""Represents expanded context around a position."""
|
|
|
|
file_path: str
|
|
original_match: str
|
|
match_start: int
|
|
match_end: int
|
|
expanded_text: str
|
|
expansion_start: int
|
|
expansion_end: int
|
|
mode: str
|
|
heading_path: List[str]
|
|
metadata: Dict = None
|
|
|
|
def __post_init__(self):
|
|
if self.metadata is None:
|
|
self.metadata = {}
|
|
|
|
|
|
class DocumentParser:
|
|
"""Parses markdown documents with structure awareness."""
|
|
|
|
def __init__(self, file_path: Path):
|
|
self.file_path = file_path
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
self.content = f.read()
|
|
self.headings = self._parse_headings()
|
|
self.paragraphs = self._parse_paragraphs()
|
|
|
|
def _parse_headings(self) -> List[Dict]:
|
|
"""Parse all headings with their positions."""
|
|
headings: List[Dict] = []
|
|
pattern = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
|
|
|
|
for match in pattern.finditer(self.content):
|
|
level = len(match.group(1))
|
|
text = match.group(2).strip()
|
|
start_pos = match.start()
|
|
|
|
headings.append(
|
|
{
|
|
"level": level,
|
|
"text": text,
|
|
"start_pos": start_pos,
|
|
"end_pos": match.end(),
|
|
}
|
|
)
|
|
|
|
return headings
|
|
|
|
def _parse_paragraphs(self) -> List[Tuple[int, int]]:
|
|
"""Parse paragraph boundaries (text blocks separated by blank lines)."""
|
|
paragraphs: List[Tuple[int, int]] = []
|
|
|
|
in_paragraph = False
|
|
para_start = 0
|
|
|
|
lines = self.content.split("\n")
|
|
line_pos = 0
|
|
|
|
for line in lines:
|
|
line_start = line_pos
|
|
line_end = line_pos + len(line)
|
|
line_pos = line_end + 1 # +1 for newline
|
|
|
|
if line.strip():
|
|
if not in_paragraph:
|
|
para_start = line_start
|
|
in_paragraph = True
|
|
else:
|
|
if in_paragraph:
|
|
paragraphs.append((para_start, line_start))
|
|
in_paragraph = False
|
|
|
|
if in_paragraph:
|
|
paragraphs.append((para_start, len(self.content)))
|
|
|
|
return paragraphs
|
|
|
|
def get_heading_path(self, position: int) -> List[str]:
|
|
"""Get the breadcrumb trail of headings for a position."""
|
|
path: List[str] = []
|
|
current_levels: Dict[int, str] = {}
|
|
|
|
for heading in self.headings:
|
|
if heading["start_pos"] > position:
|
|
break
|
|
|
|
level = heading["level"]
|
|
keys_to_remove = [k for k in current_levels.keys() if k >= level]
|
|
for k in keys_to_remove:
|
|
del current_levels[k]
|
|
|
|
current_levels[level] = heading["text"]
|
|
|
|
for level in sorted(current_levels.keys()):
|
|
path.append(current_levels[level])
|
|
|
|
return path
|
|
|
|
def get_section_bounds(
|
|
self, position: int, include_subsections: bool = True
|
|
) -> Tuple[int, int]:
|
|
"""Get the bounds of the section containing the position."""
|
|
containing_heading = None
|
|
containing_level = None
|
|
|
|
for i, heading in enumerate(self.headings):
|
|
if heading["start_pos"] <= position:
|
|
containing_heading = i
|
|
containing_level = heading["level"]
|
|
else:
|
|
break
|
|
|
|
if containing_heading is None:
|
|
if self.headings:
|
|
return (0, self.headings[0]["start_pos"])
|
|
return (0, len(self.content))
|
|
|
|
section_start = self.headings[containing_heading]["start_pos"]
|
|
section_end = len(self.content)
|
|
|
|
for i in range(containing_heading + 1, len(self.headings)):
|
|
next_heading = self.headings[i]
|
|
if include_subsections:
|
|
if next_heading["level"] <= containing_level:
|
|
section_end = next_heading["start_pos"]
|
|
break
|
|
else:
|
|
section_end = next_heading["start_pos"]
|
|
break
|
|
|
|
return (section_start, section_end)
|
|
|
|
def get_paragraph_bounds(self, position: int) -> Tuple[int, int]:
|
|
"""Get the bounds of the paragraph containing the position."""
|
|
for para_start, para_end in self.paragraphs:
|
|
if para_start <= position < para_end:
|
|
return (para_start, para_end)
|
|
|
|
return (max(0, position - 500), min(len(self.content), position + 500))
|
|
|
|
def get_document_structure(self, focus_position: Optional[int] = None) -> Dict:
|
|
"""Get the full document structure with optional focus on a position."""
|
|
structure: Dict = {
|
|
"file": os.path.basename(self.file_path),
|
|
"total_chars": len(self.content),
|
|
"headings": [],
|
|
}
|
|
|
|
for heading in self.headings:
|
|
heading_info = {
|
|
"level": heading["level"],
|
|
"text": heading["text"],
|
|
"position": heading["start_pos"],
|
|
"is_focus": False,
|
|
}
|
|
|
|
if focus_position is not None:
|
|
start = heading["start_pos"]
|
|
end = len(self.content)
|
|
for next_h in self.headings:
|
|
if (
|
|
next_h["start_pos"] > start
|
|
and next_h["level"] <= heading["level"]
|
|
):
|
|
end = next_h["start_pos"]
|
|
break
|
|
|
|
if start <= focus_position < end:
|
|
heading_info["is_focus"] = True
|
|
|
|
structure["headings"].append(heading_info)
|
|
|
|
return structure
|
|
|
|
|
|
def expand_context(
|
|
file_path: Path,
|
|
position: int,
|
|
match_text: str,
|
|
mode: str = "paragraph",
|
|
match_end: Optional[int] = None,
|
|
) -> ExpandedContext:
|
|
"""Expand context around a position with various modes."""
|
|
if match_end is None:
|
|
match_end = position + len(match_text)
|
|
|
|
parser = DocumentParser(file_path)
|
|
heading_path = parser.get_heading_path(position)
|
|
|
|
if mode == "char":
|
|
char_context = 1000
|
|
start = max(0, position - char_context)
|
|
end = min(len(parser.content), match_end + char_context)
|
|
expanded_text = parser.content[start:end]
|
|
elif mode == "paragraph":
|
|
start, end = parser.get_paragraph_bounds(position)
|
|
expanded_text = parser.content[start:end]
|
|
elif mode == "section":
|
|
start, end = parser.get_section_bounds(position, include_subsections=True)
|
|
expanded_text = parser.content[start:end]
|
|
elif mode == "section-only":
|
|
start, end = parser.get_section_bounds(position, include_subsections=False)
|
|
expanded_text = parser.content[start:end]
|
|
elif mode == "document":
|
|
start = 0
|
|
end = len(parser.content)
|
|
expanded_text = parser.content
|
|
else:
|
|
raise ValueError(f"Unknown mode: {mode}")
|
|
|
|
return ExpandedContext(
|
|
file_path=str(file_path),
|
|
original_match=match_text,
|
|
match_start=position,
|
|
match_end=match_end,
|
|
expanded_text=expanded_text,
|
|
expansion_start=start,
|
|
expansion_end=end,
|
|
mode=mode,
|
|
heading_path=heading_path,
|
|
metadata={
|
|
"file_size": len(parser.content),
|
|
"expansion_size": len(expanded_text),
|
|
"num_headings_in_path": len(heading_path),
|
|
},
|
|
)
|
|
|
|
|
|
def format_expanded_context(
|
|
context: ExpandedContext, format_type: str = "human"
|
|
) -> str:
|
|
"""Format expanded context for output."""
|
|
|
|
if format_type == "json":
|
|
return json.dumps(asdict(context), indent=2)
|
|
elif format_type == "human":
|
|
output: List[str] = []
|
|
output.append("=" * 80)
|
|
output.append(f"File: {os.path.basename(context.file_path)}")
|
|
output.append(f"Mode: {context.mode}")
|
|
output.append(f"Original Match: '{context.original_match}'")
|
|
output.append(f"Match Position: {context.match_start}-{context.match_end}")
|
|
output.append(
|
|
f"Expanded Range: {context.expansion_start}-{context.expansion_end}"
|
|
)
|
|
output.append(f"Expansion Size: {context.metadata['expansion_size']} chars")
|
|
|
|
if context.heading_path:
|
|
output.append("\nHeading Path:")
|
|
for i, heading in enumerate(context.heading_path, 1):
|
|
output.append(f" {' ' * (i - 1)}→ {heading}")
|
|
|
|
output.append("\n" + "-" * 80)
|
|
output.append("EXPANDED CONTENT:")
|
|
output.append("-" * 80)
|
|
output.append(context.expanded_text)
|
|
output.append("=" * 80)
|
|
|
|
return "\n".join(output)
|
|
elif format_type == "llm":
|
|
output: List[str] = []
|
|
output.append("### EXPANDED CONTEXT")
|
|
output.append(f"**Source**: {os.path.basename(context.file_path)}")
|
|
output.append(
|
|
f"**Character Range**: {context.expansion_start}-{context.expansion_end}"
|
|
)
|
|
output.append(
|
|
f"**Original Match** at position {context.match_start}: `{context.original_match}`"
|
|
)
|
|
|
|
if context.heading_path:
|
|
breadcrumb = " → ".join(context.heading_path)
|
|
output.append(f"**Location**: {breadcrumb}")
|
|
|
|
output.append("\n---\n")
|
|
output.append(context.expanded_text)
|
|
output.append("\n---\n")
|
|
|
|
return "\n".join(output)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description=("Expand context around search results from D&D 5e SRD references"),
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog=(
|
|
"""
|
|
Examples:
|
|
# Expand the 3rd search result by section
|
|
python expand_context.py "fireball" --result 3 --mode section --all
|
|
|
|
# Expand multiple results by paragraph
|
|
python expand_context.py "wizard" --results 1,3,5 --mode paragraph --all
|
|
|
|
# Direct expansion from file position
|
|
python expand_context.py --file "DND5eSRD_121-137.md" --position 1234 --mode section
|
|
|
|
# Get document structure
|
|
python expand_context.py --file "DND5eSRD_121-137.md" --position 1234 --mode document
|
|
|
|
# Output as JSON for machine processing
|
|
python expand_context.py "fireball" --result 1 --all --format json
|
|
|
|
Expansion Modes:
|
|
char - Expand by ±1000 characters (simple)
|
|
paragraph - Expand to paragraph boundaries (blank line separated)
|
|
section - Expand to full section including subsections
|
|
section-only - Expand to section excluding subsections
|
|
document - Show entire document (use with caution)
|
|
|
|
Output Formats:
|
|
human - Human-readable formatted output
|
|
llm - Optimized for LLM consumption (default)
|
|
json - Machine-readable JSON
|
|
"""
|
|
),
|
|
)
|
|
|
|
# Search parameters (for finding results to expand)
|
|
parser.add_argument("search_term", nargs="?", help="Term to search for")
|
|
parser.add_argument("--result", type=int, help="Result number to expand (1-based)")
|
|
parser.add_argument(
|
|
"--results", help="Comma-separated result numbers to expand (e.g., '1,3,5')"
|
|
)
|
|
parser.add_argument(
|
|
"--all-search",
|
|
action="store_true",
|
|
dest="search_all",
|
|
help="Search all reference files",
|
|
)
|
|
parser.add_argument("--pages", help="Search files in page range (e.g., 001-120)")
|
|
parser.add_argument("--files", nargs="+", help="Specific files to search")
|
|
|
|
# Direct expansion parameters
|
|
parser.add_argument("--file", help="Direct file path for expansion")
|
|
parser.add_argument(
|
|
"--position", type=int, help="Character position for direct expansion"
|
|
)
|
|
|
|
# Expansion options
|
|
parser.add_argument(
|
|
"--mode",
|
|
choices=["char", "paragraph", "section", "section-only", "document"],
|
|
default="paragraph",
|
|
help="Expansion mode (default: paragraph)",
|
|
)
|
|
parser.add_argument(
|
|
"--format",
|
|
choices=["human", "llm", "json"],
|
|
default="llm",
|
|
help="Output format (default: llm)",
|
|
)
|
|
|
|
# Search options
|
|
parser.add_argument(
|
|
"--case-sensitive", action="store_true", help="Case-sensitive search"
|
|
)
|
|
parser.add_argument(
|
|
"--max-search-results",
|
|
type=int,
|
|
default=50,
|
|
help="Maximum search results to find (default: 50)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Direct expansion mode
|
|
if args.file and args.position is not None:
|
|
file_path = Path(args.file)
|
|
if not file_path.exists():
|
|
refs_dir = get_references_dir()
|
|
file_path = refs_dir / args.file
|
|
|
|
if not file_path.exists():
|
|
print(f"Error: File not found: {args.file}")
|
|
return
|
|
|
|
context = expand_context(
|
|
file_path=file_path,
|
|
position=args.position,
|
|
match_text="[direct position]",
|
|
mode=args.mode,
|
|
)
|
|
|
|
print(format_expanded_context(context, args.format))
|
|
return
|
|
|
|
# Search-based expansion mode
|
|
if not args.search_term:
|
|
print("Error: Provide either search_term + result, or --file + --position")
|
|
parser.print_help()
|
|
return
|
|
|
|
if not args.result and not args.results:
|
|
print("Error: Specify which result(s) to expand with --result or --results")
|
|
return
|
|
|
|
# Determine files to search
|
|
files_to_search: List[Path] = []
|
|
refs_dir = get_references_dir()
|
|
|
|
if args.search_all:
|
|
files_to_search = get_all_reference_files()
|
|
elif args.pages:
|
|
files_to_search = get_files_by_page_range(args.pages)
|
|
elif args.files:
|
|
for filename in args.files:
|
|
file_path = refs_dir / filename
|
|
if file_path.exists():
|
|
files_to_search.append(file_path)
|
|
else:
|
|
print(f"Warning: File not found: {filename}")
|
|
else:
|
|
files_to_search = get_all_reference_files()
|
|
|
|
if not files_to_search:
|
|
print("Error: No files to search")
|
|
return
|
|
|
|
# Perform search
|
|
print(f"Searching for '{args.search_term}'...", file=sys.stderr)
|
|
search_results = search_files(
|
|
files_to_search,
|
|
args.search_term,
|
|
case_sensitive=args.case_sensitive,
|
|
max_results=args.max_search_results,
|
|
context_chars=50,
|
|
)
|
|
|
|
if not search_results:
|
|
print(f"No results found for '{args.search_term}'")
|
|
return
|
|
|
|
print(f"Found {len(search_results)} results\n", file=sys.stderr)
|
|
|
|
# Determine which results to expand
|
|
results_to_expand: List[int] = []
|
|
if args.results:
|
|
result_nums = [int(x.strip()) for x in args.results.split(",")]
|
|
results_to_expand = result_nums
|
|
elif args.result:
|
|
results_to_expand = [args.result]
|
|
|
|
# Expand specified results
|
|
for result_num in results_to_expand:
|
|
if result_num < 1 or result_num > len(search_results):
|
|
print(
|
|
f"Warning: Result {result_num} out of range (1-{len(search_results)})"
|
|
)
|
|
continue
|
|
|
|
search_result = search_results[result_num - 1]
|
|
|
|
print(f"\n{'=' * 80}", file=sys.stderr)
|
|
print(f"EXPANDING RESULT {result_num}/{len(search_results)}", file=sys.stderr)
|
|
print(f"{'=' * 80}\n", file=sys.stderr)
|
|
|
|
context = expand_context(
|
|
file_path=Path(search_result.file_path),
|
|
position=search_result.start_pos,
|
|
match_text=search_result.match_text,
|
|
mode=args.mode,
|
|
match_end=search_result.end_pos,
|
|
)
|
|
|
|
print(format_expanded_context(context, args.format))
|
|
|
|
if len(results_to_expand) > 1 and result_num != results_to_expand[-1]:
|
|
print("\n" + "=" * 80 + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|