mirror of
https://github.com/microsoft/SkillOpt.git
synced 2026-07-07 00:15:39 +08:00
553 lines
19 KiB
Python
553 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import fnmatch
|
|
import html
|
|
import json
|
|
import os
|
|
import re
|
|
import socket
|
|
import time
|
|
from functools import lru_cache
|
|
from html.parser import HTMLParser
|
|
from pathlib import Path
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.parse import parse_qs, urlparse
|
|
from urllib.request import Request, urlopen
|
|
|
|
_MAX_READ_CHARS = 4000
|
|
_MAX_GREP_MATCHES = 20
|
|
_MAX_GLOB_MATCHES = 50
|
|
_MAX_ORACLE_PAGE_CHARS = 24000
|
|
_MAX_ORACLE_CONTEXT_CHARS = 80000
|
|
DEFAULT_USER_AGENT = (
|
|
"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
|
|
"(KHTML, like Gecko) Chrome/135.0 Safari/537.36"
|
|
)
|
|
DEFAULT_CUSTOM_SEARCH_URL = "http://apisix.westus2.cloudapp.azure.com/search_tool/search"
|
|
DEFAULT_CUSTOM_SEARCH_AUTH_ENV = "OFFICEQA_CUSTOM_SEARCH_AUTH"
|
|
DEFAULT_CUSTOM_SEARCH_PROVIDER = "duckduckgo"
|
|
DEFAULT_CUSTOM_SEARCH_MAX_RESULTS = 4
|
|
DEFAULT_CUSTOM_SEARCH_TIMEOUT = 20
|
|
DEFAULT_CUSTOM_SEARCH_MAX_RETRIES = 4
|
|
DEFAULT_CUSTOM_SEARCH_INITIAL_BACKOFF_SECONDS = 1.0
|
|
|
|
|
|
def _normalize_data_dirs(data_dirs: list[str] | tuple[str, ...] | str | None, project_root: Path) -> list[str]:
|
|
if data_dirs is None:
|
|
return []
|
|
if isinstance(data_dirs, str):
|
|
items = [part.strip() for chunk in data_dirs.split(os.pathsep) for part in chunk.split(",")]
|
|
else:
|
|
items = [str(item).strip() for item in data_dirs]
|
|
resolved: list[str] = []
|
|
for item in items:
|
|
if not item:
|
|
continue
|
|
path = Path(item).expanduser()
|
|
if not path.is_absolute():
|
|
path = project_root / path
|
|
resolved.append(str(path))
|
|
return resolved
|
|
|
|
|
|
def resolve_docs_roots(data_dirs: list[str] | tuple[str, ...] | str | None = None) -> list[str]:
|
|
project_root = Path(__file__).resolve().parents[3]
|
|
env_value = os.environ.get("OFFICEQA_DOCS_DIR", "").strip()
|
|
candidates = _normalize_data_dirs(data_dirs, project_root)
|
|
candidates.extend(_normalize_data_dirs(env_value, project_root))
|
|
candidates.extend([
|
|
str(project_root / "data" / "officeqa_docs_official"),
|
|
str(project_root / "data" / "officeqa_smoke_docs"),
|
|
os.path.expanduser("~/officeqa-sparse/treasury_bulletins_parsed"),
|
|
os.path.expanduser("~/officeqa/treasury_bulletins_parsed"),
|
|
])
|
|
roots: list[str] = []
|
|
seen: set[str] = set()
|
|
for candidate in candidates:
|
|
path = Path(candidate).expanduser()
|
|
if not path.is_dir():
|
|
continue
|
|
transformed = path / "transformed"
|
|
resolved = str((transformed if transformed.is_dir() else path).resolve())
|
|
if resolved in seen:
|
|
continue
|
|
seen.add(resolved)
|
|
roots.append(resolved)
|
|
if not roots:
|
|
raise FileNotFoundError("OfficeQA docs directory not found. Set OFFICEQA_DOCS_DIR or env.data_dirs.")
|
|
return roots
|
|
|
|
|
|
def _is_allowed(path: str, allowed_roots: list[str], allowed_files: list[str]) -> bool:
|
|
try:
|
|
resolved = str(Path(path).resolve())
|
|
except FileNotFoundError:
|
|
return False
|
|
if not any(resolved.startswith(root + os.sep) or resolved == root for root in allowed_roots):
|
|
return False
|
|
if not allowed_files:
|
|
return True
|
|
base = os.path.basename(resolved)
|
|
return base in allowed_files
|
|
|
|
|
|
def resolve_candidate_files(source_files: list[str], allowed_roots: list[str]) -> list[str]:
|
|
resolved: list[str] = []
|
|
seen: set[str] = set()
|
|
for root in allowed_roots:
|
|
for dirpath, _, filenames in os.walk(root):
|
|
for filename in filenames:
|
|
if source_files and filename not in source_files:
|
|
continue
|
|
full = str(Path(dirpath, filename).resolve())
|
|
if full in seen:
|
|
continue
|
|
seen.add(full)
|
|
resolved.append(full)
|
|
return resolved
|
|
|
|
|
|
def _as_list(value: object) -> list[str]:
|
|
if value is None:
|
|
return []
|
|
if isinstance(value, list):
|
|
return [str(item).strip() for item in value if str(item).strip()]
|
|
text = str(value).strip()
|
|
if not text:
|
|
return []
|
|
try:
|
|
loaded = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
loaded = None
|
|
if isinstance(loaded, list):
|
|
return [str(item).strip() for item in loaded if str(item).strip()]
|
|
if "\n" in text:
|
|
return [part.strip() for part in text.splitlines() if part.strip()]
|
|
return [text]
|
|
|
|
|
|
def _extract_page_number(source_doc: str) -> int | None:
|
|
text = str(source_doc or "").strip()
|
|
if not text:
|
|
return None
|
|
parsed = urlparse(text)
|
|
query = parse_qs(parsed.query)
|
|
for key in ("page", "pagenum", "page_id"):
|
|
for raw_value in query.get(key, []):
|
|
try:
|
|
return int(str(raw_value).strip())
|
|
except ValueError:
|
|
continue
|
|
match = re.search(r"(?:[?&]|^)page=(\d+)", text)
|
|
if match:
|
|
return int(match.group(1))
|
|
return None
|
|
|
|
|
|
def _iter_oracle_refs(source_files: object, source_docs: object) -> list[tuple[str, int, str]]:
|
|
files = _as_list(source_files)
|
|
docs = _as_list(source_docs)
|
|
refs: list[tuple[str, int, str]] = []
|
|
seen: set[tuple[str, int, str]] = set()
|
|
if not files or not docs:
|
|
return refs
|
|
for index, source_doc in enumerate(docs):
|
|
page_number = _extract_page_number(source_doc)
|
|
if page_number is None:
|
|
continue
|
|
if index < len(files):
|
|
source_file = files[index]
|
|
elif len(files) == 1:
|
|
source_file = files[0]
|
|
else:
|
|
continue
|
|
key = (source_file, page_number, source_doc)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
refs.append(key)
|
|
return refs
|
|
|
|
|
|
def _parsed_root_candidates(docs_roots: list[str]) -> list[Path]:
|
|
candidates: list[Path] = []
|
|
seen: set[str] = set()
|
|
for root in docs_roots:
|
|
path = Path(root).expanduser()
|
|
for candidate in (
|
|
path,
|
|
path.parent,
|
|
path / "treasury_bulletins_parsed",
|
|
path.parent / "treasury_bulletins_parsed",
|
|
):
|
|
resolved = str(candidate.resolve()) if candidate.exists() else str(candidate)
|
|
if resolved in seen:
|
|
continue
|
|
seen.add(resolved)
|
|
candidates.append(candidate)
|
|
return candidates
|
|
|
|
|
|
def _locate_parsed_json(source_file: str, docs_roots: list[str]) -> Path | None:
|
|
source_path = Path(str(source_file).strip())
|
|
stem = source_path.stem if source_path.suffix else source_path.name
|
|
if not stem:
|
|
return None
|
|
candidate_names = [stem + ".json"]
|
|
if source_path.suffix == ".json":
|
|
candidate_names.insert(0, source_path.name)
|
|
for root in _parsed_root_candidates(docs_roots):
|
|
for name in candidate_names:
|
|
path = root / "jsons" / name
|
|
if path.is_file():
|
|
return path
|
|
return None
|
|
|
|
|
|
class _TableMarkdownParser(HTMLParser):
|
|
def __init__(self) -> None:
|
|
super().__init__(convert_charrefs=True)
|
|
self.rows: list[list[str]] = []
|
|
self._row: list[str] | None = None
|
|
self._cell: list[str] | None = None
|
|
|
|
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
|
if tag.lower() == "tr":
|
|
self._row = []
|
|
elif tag.lower() in {"td", "th"} and self._row is not None:
|
|
self._cell = []
|
|
|
|
def handle_data(self, data: str) -> None:
|
|
if self._cell is not None:
|
|
self._cell.append(data)
|
|
|
|
def handle_endtag(self, tag: str) -> None:
|
|
normalized_tag = tag.lower()
|
|
if normalized_tag in {"td", "th"} and self._cell is not None and self._row is not None:
|
|
cell = re.sub(r"\s+", " ", "".join(self._cell)).strip()
|
|
self._row.append(cell)
|
|
self._cell = None
|
|
elif normalized_tag == "tr" and self._row is not None:
|
|
if any(cell for cell in self._row):
|
|
self.rows.append(self._row)
|
|
self._row = None
|
|
self._cell = None
|
|
|
|
|
|
def _escape_markdown_cell(value: str) -> str:
|
|
return str(value).replace("\n", " ").replace("|", "\\|").strip()
|
|
|
|
|
|
def _html_table_to_markdown(raw_html: str) -> str:
|
|
parser = _TableMarkdownParser()
|
|
try:
|
|
parser.feed(raw_html)
|
|
except Exception: # noqa: BLE001
|
|
parser.rows = []
|
|
rows = parser.rows
|
|
if not rows:
|
|
text = re.sub(r"(?is)<[^>]+>", " ", raw_html)
|
|
return re.sub(r"\s+", " ", html.unescape(text)).strip()
|
|
width = max(len(row) for row in rows)
|
|
normalized_rows = [row + [""] * (width - len(row)) for row in rows]
|
|
header = normalized_rows[0]
|
|
body = normalized_rows[1:]
|
|
lines = [
|
|
"| " + " | ".join(_escape_markdown_cell(cell) for cell in header) + " |",
|
|
"| " + " | ".join(["---"] * width) + " |",
|
|
]
|
|
lines.extend("| " + " | ".join(_escape_markdown_cell(cell) for cell in row) + " |" for row in body)
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _render_parsed_content(content: str) -> str:
|
|
text = content.strip()
|
|
if not text:
|
|
return ""
|
|
if "<table" in text.lower():
|
|
return _html_table_to_markdown(text)
|
|
text = html.unescape(text)
|
|
text = re.sub(r"\r\n?", "\n", text)
|
|
return re.sub(r"\n{3,}", "\n\n", text).strip()
|
|
|
|
|
|
def _element_page_ids(element: dict) -> set[int]:
|
|
page_ids: set[int] = set()
|
|
bbox = element.get("bbox")
|
|
if not isinstance(bbox, list):
|
|
return page_ids
|
|
for box in bbox:
|
|
if not isinstance(box, dict):
|
|
continue
|
|
raw_page_id = box.get("page_id")
|
|
try:
|
|
page_ids.add(int(raw_page_id))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
return page_ids
|
|
|
|
|
|
@lru_cache(maxsize=256)
|
|
def _load_parsed_elements(json_path: str) -> tuple[dict, ...]:
|
|
with open(json_path, encoding="utf-8") as f:
|
|
payload = json.load(f)
|
|
document = payload.get("document") if isinstance(payload, dict) else {}
|
|
elements = document.get("elements") if isinstance(document, dict) else []
|
|
if not isinstance(elements, list):
|
|
return ()
|
|
return tuple(element for element in elements if isinstance(element, dict))
|
|
|
|
|
|
@lru_cache(maxsize=2048)
|
|
def _render_parsed_page(json_path: str, page_number: int) -> str:
|
|
rendered: list[str] = []
|
|
for element in _load_parsed_elements(json_path):
|
|
if page_number not in _element_page_ids(element):
|
|
continue
|
|
content = element.get("content")
|
|
if not isinstance(content, str) or not content.strip():
|
|
continue
|
|
section = _render_parsed_content(content)
|
|
if section:
|
|
rendered.append(section)
|
|
return "\n\n".join(rendered).strip()
|
|
|
|
|
|
def build_oracle_parsed_pages_context(
|
|
source_files: object,
|
|
source_docs: object,
|
|
docs_roots: list[str],
|
|
*,
|
|
max_page_chars: int = _MAX_ORACLE_PAGE_CHARS,
|
|
max_total_chars: int = _MAX_ORACLE_CONTEXT_CHARS,
|
|
evidence_note: str = "Treat it as primary document evidence and combine it with custom web search results when useful.",
|
|
) -> str:
|
|
"""Render oracle parsed OfficeQA pages referenced by source_docs/source_files."""
|
|
refs = _iter_oracle_refs(source_files, source_docs)
|
|
if not refs:
|
|
return ""
|
|
|
|
blocks: list[str] = []
|
|
total_chars = 0
|
|
seen_pages: set[tuple[str, int]] = set()
|
|
for source_file, page_number, source_doc in refs:
|
|
json_path = _locate_parsed_json(source_file, docs_roots)
|
|
if json_path is None:
|
|
continue
|
|
page_key = (str(json_path), page_number)
|
|
if page_key in seen_pages:
|
|
continue
|
|
seen_pages.add(page_key)
|
|
page_text = _render_parsed_page(str(json_path), page_number)
|
|
if not page_text:
|
|
continue
|
|
if len(page_text) > max_page_chars:
|
|
omitted = len(page_text) - max_page_chars
|
|
page_text = page_text[:max_page_chars].rstrip() + f"\n\n[... {omitted} characters omitted from this parsed page ...]"
|
|
block = (
|
|
f"### {source_file} page {page_number}\n"
|
|
f"Source URL: {source_doc}\n\n"
|
|
f"{page_text}"
|
|
)
|
|
if total_chars + len(block) > max_total_chars:
|
|
remaining = max_total_chars - total_chars
|
|
if remaining <= 0:
|
|
break
|
|
block = block[:remaining].rstrip() + "\n\n[... oracle parsed page context truncated ...]"
|
|
blocks.append(block)
|
|
break
|
|
blocks.append(block)
|
|
total_chars += len(block)
|
|
if not blocks:
|
|
return ""
|
|
return (
|
|
"The following content is pre-parsed from the oracle OfficeQA source page(s). "
|
|
f"{evidence_note.strip()}\n\n"
|
|
+ "\n\n".join(blocks)
|
|
)
|
|
|
|
|
|
def _extract_search_items(payload: object) -> list[dict]:
|
|
if isinstance(payload, list):
|
|
return [item for item in payload if isinstance(item, dict)]
|
|
if not isinstance(payload, dict):
|
|
return []
|
|
candidate_keys = (
|
|
"results",
|
|
"items",
|
|
"data",
|
|
"organic",
|
|
"organic_results",
|
|
"search_results",
|
|
"webPages",
|
|
"value",
|
|
)
|
|
for key in candidate_keys:
|
|
value = payload.get(key)
|
|
if isinstance(value, list):
|
|
return [item for item in value if isinstance(item, dict)]
|
|
if isinstance(value, dict):
|
|
nested = _extract_search_items(value)
|
|
if nested:
|
|
return nested
|
|
return []
|
|
|
|
|
|
def _normalize_search_item(item: dict, index: int) -> str:
|
|
title = str(
|
|
item.get("title")
|
|
or item.get("name")
|
|
or item.get("headline")
|
|
or item.get("source")
|
|
or f"Result {index}"
|
|
).strip()
|
|
url = str(
|
|
item.get("url")
|
|
or item.get("link")
|
|
or item.get("href")
|
|
or item.get("display_url")
|
|
or ""
|
|
).strip()
|
|
snippet = str(
|
|
item.get("snippet")
|
|
or item.get("description")
|
|
or item.get("body")
|
|
or item.get("text")
|
|
or item.get("content")
|
|
or ""
|
|
).strip()
|
|
lines = [f"[{index}] {title}"]
|
|
if url:
|
|
lines.append(f"URL: {url}")
|
|
if snippet:
|
|
lines.append(f"Snippet: {snippet}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _format_search_payload(query: str, payload: object) -> str:
|
|
items = _extract_search_items(payload)
|
|
header = f"Query: {query}"
|
|
if not items:
|
|
body = json.dumps(payload, ensure_ascii=False) if payload else "[no results]"
|
|
return f"{header}\n{body}"
|
|
rendered = [_normalize_search_item(item, index) for index, item in enumerate(items, start=1)]
|
|
return f"{header}\n\n" + "\n\n".join(rendered)
|
|
|
|
|
|
def _is_retryable_search_http_error(status_code: int) -> bool:
|
|
return status_code in {408, 429} or status_code >= 500
|
|
|
|
|
|
def custom_search(
|
|
query: str,
|
|
*,
|
|
api_url: str = DEFAULT_CUSTOM_SEARCH_URL,
|
|
auth_token: str | None = None,
|
|
auth_env: str = DEFAULT_CUSTOM_SEARCH_AUTH_ENV,
|
|
provider: str = DEFAULT_CUSTOM_SEARCH_PROVIDER,
|
|
max_num_results: int = DEFAULT_CUSTOM_SEARCH_MAX_RESULTS,
|
|
timeout: int = DEFAULT_CUSTOM_SEARCH_TIMEOUT,
|
|
max_retries: int = DEFAULT_CUSTOM_SEARCH_MAX_RETRIES,
|
|
initial_backoff_seconds: float = DEFAULT_CUSTOM_SEARCH_INITIAL_BACKOFF_SECONDS,
|
|
) -> str:
|
|
query = str(query or "").strip()
|
|
if not query:
|
|
raise ValueError("custom_search query must be non-empty")
|
|
token = str(auth_token or os.environ.get(auth_env, "")).strip()
|
|
if not token:
|
|
raise ValueError(f"custom_search auth token missing; set {auth_env}")
|
|
payload = json.dumps(
|
|
{
|
|
"query": query,
|
|
"max_num_results": int(max_num_results),
|
|
"provider": provider,
|
|
},
|
|
ensure_ascii=False,
|
|
).encode("utf-8")
|
|
req = Request(
|
|
api_url,
|
|
data=payload,
|
|
headers={
|
|
"Authorization": token,
|
|
"Content-Type": "application/json",
|
|
"User-Agent": DEFAULT_USER_AGENT,
|
|
},
|
|
method="POST",
|
|
)
|
|
attempts = max(1, int(max_retries) + 1)
|
|
last_error: RuntimeError | None = None
|
|
for attempt in range(1, attempts + 1):
|
|
try:
|
|
with urlopen(req, timeout=timeout) as response:
|
|
raw_body = response.read().decode("utf-8", errors="ignore")
|
|
break
|
|
except HTTPError as exc:
|
|
detail = exc.read().decode("utf-8", errors="ignore")
|
|
last_error = RuntimeError(f"custom_search HTTP {exc.code}: {detail[:1000]}")
|
|
if attempt >= attempts or not _is_retryable_search_http_error(exc.code):
|
|
raise last_error from exc
|
|
except (URLError, TimeoutError, socket.timeout) as exc:
|
|
last_error = RuntimeError(f"custom_search connection error: {exc}")
|
|
if attempt >= attempts:
|
|
raise last_error from exc
|
|
backoff_seconds = max(0.0, float(initial_backoff_seconds)) * (2 ** (attempt - 1))
|
|
if backoff_seconds > 0:
|
|
time.sleep(backoff_seconds)
|
|
else:
|
|
raise last_error or RuntimeError("custom_search failed without a captured error")
|
|
try:
|
|
parsed = json.loads(raw_body)
|
|
except json.JSONDecodeError:
|
|
return f"Query: {query}\n\n{raw_body.strip() or '[empty response]'}"
|
|
return _format_search_payload(query, parsed)
|
|
|
|
|
|
def run_tool(name: str, arguments: dict, *, allowed_roots: list[str], allowed_files: list[str]) -> tuple[str, str]:
|
|
if name == "glob":
|
|
pattern = str(arguments.get("pattern") or "*")
|
|
matches: list[str] = []
|
|
for root in allowed_roots:
|
|
for dirpath, _, filenames in os.walk(root):
|
|
for filename in filenames:
|
|
if allowed_files and filename not in allowed_files:
|
|
continue
|
|
rel = os.path.relpath(os.path.join(dirpath, filename), root)
|
|
if fnmatch.fnmatch(rel, pattern) or fnmatch.fnmatch(filename, pattern):
|
|
matches.append(os.path.join(dirpath, filename))
|
|
if len(matches) >= _MAX_GLOB_MATCHES:
|
|
break
|
|
if len(matches) >= _MAX_GLOB_MATCHES:
|
|
break
|
|
return f"glob(pattern={pattern!r})", "\n".join(matches) if matches else "[no matches]"
|
|
|
|
if name == "read":
|
|
path = str(arguments.get("path") or "")
|
|
if not path:
|
|
return "read(path='')", "[read error: missing path]"
|
|
if not _is_allowed(path, allowed_roots, allowed_files):
|
|
return f"read(path={path!r})", "[read error: path not allowed]"
|
|
start = max(int(arguments.get("start") or 1), 1)
|
|
limit = max(int(arguments.get("limit") or 80), 1)
|
|
with open(path, encoding="utf-8") as f:
|
|
lines = f.readlines()
|
|
excerpt = "".join(lines[start - 1:start - 1 + limit])
|
|
return f"read(path={path!r}, start={start}, limit={limit})", excerpt[:_MAX_READ_CHARS] or "[empty file]"
|
|
|
|
if name == "grep":
|
|
pattern = str(arguments.get("pattern") or "").lower()
|
|
path = str(arguments.get("path") or "")
|
|
if not pattern or not path:
|
|
return f"grep(pattern={pattern!r}, path={path!r})", "[grep error: missing pattern or path]"
|
|
if not _is_allowed(path, allowed_roots, allowed_files):
|
|
return f"grep(pattern={pattern!r}, path={path!r})", "[grep error: path not allowed]"
|
|
matches: list[str] = []
|
|
with open(path, encoding="utf-8") as f:
|
|
for idx, line in enumerate(f, start=1):
|
|
if pattern in line.lower():
|
|
matches.append(f"{idx}: {line.rstrip()}")
|
|
if len(matches) >= _MAX_GREP_MATCHES:
|
|
break
|
|
return f"grep(pattern={pattern!r}, path={path!r})", "\n".join(matches) if matches else "[no matches]"
|
|
|
|
return name, f"[tool error: unknown tool {name}]"
|