Initial commit

This commit is contained in:
Cuzyoung
2026-05-08 18:16:18 +00:00
commit 9fda7311ea
243 changed files with 31492 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
"""OfficeQA environment package for ReflACT."""
+133
View File
@@ -0,0 +1,133 @@
from __future__ import annotations
import os
from skillopt.datasets.base import BatchSpec
from skillopt.envs.base import EnvAdapter
from skillopt.envs.deep_reflect import run_no_reference_deep_reflect
from skillopt.envs.officeqa.dataloader import OfficeQADataLoader
from skillopt.envs.officeqa.rollout import run_batch
from skillopt.gradient.reflect import run_minibatch_reflect
class OfficeQAAdapter(EnvAdapter):
def __init__(
self,
split_dir: str = "",
workers: int = 8,
analyst_workers: int = 8,
failure_only: bool = False,
minibatch_size: int = 8,
edit_budget: int = 4,
seed: int = 42,
limit: int = 0,
max_tool_turns: int = 12,
data_dirs: list[str] | str | None = None,
docs_dirs: list[str] | str | None = None,
use_deep_reflect: bool = False,
deep_reflect_failures: int = 4,
deep_reflect_successes: int = 2,
) -> None:
self.workers = workers
self.analyst_workers = analyst_workers
self.failure_only = failure_only
self.minibatch_size = minibatch_size
self.edit_budget = edit_budget
self.max_tool_turns = max_tool_turns
self.data_dirs = data_dirs if data_dirs is not None else docs_dirs
self.use_deep_reflect = use_deep_reflect
self.deep_reflect_failures = deep_reflect_failures
self.deep_reflect_successes = deep_reflect_successes
self.dataloader = OfficeQADataLoader(split_dir=split_dir, seed=seed, limit=limit)
def setup(self, cfg: dict) -> None:
super().setup(cfg)
self.dataloader.setup(cfg)
def get_dataloader(self):
return self.dataloader
def build_env_from_batch(self, batch: BatchSpec, **kwargs):
return list(batch.payload or [])
def build_train_env(self, batch_size: int, seed: int, **kwargs):
batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
return self.build_env_from_batch(batch, **kwargs)
def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
return self.build_env_from_batch(batch, **kwargs)
def rollout(self, env_manager, skill_content: str, out_dir: str, **kwargs) -> list[dict]:
items: list[dict] = env_manager
return run_batch(
items=items,
out_root=out_dir,
skill_content=skill_content,
workers=self.workers,
max_tool_turns=self.max_tool_turns,
data_dirs=self.data_dirs,
diagnostic_mode=kwargs.get("diagnostic_mode", False),
diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
)
def reflect(self, results: list[dict], skill_content: str, out_dir: str, **kwargs) -> list[dict | None]:
prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
random_seed = kwargs.get("random_seed")
step_buffer_context = kwargs.get("step_buffer_context", "")
return run_minibatch_reflect(
results=results,
skill_content=skill_content,
prediction_dir=prediction_dir,
patches_dir=patches_dir,
workers=self.analyst_workers,
failure_only=self.failure_only,
minibatch_size=self.minibatch_size,
edit_budget=self.edit_budget,
random_seed=random_seed,
error_system=self.get_error_minibatch_prompt(),
success_system=self.get_success_minibatch_prompt(),
step_buffer_context=step_buffer_context,
update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
)
def deep_reflect(
self,
results: list[dict],
skill_content: str,
out_dir: str,
**kwargs,
) -> list[dict | None]:
return run_no_reference_deep_reflect(
self,
results,
skill_content,
out_dir,
env_manager=kwargs.get("env_manager"),
prediction_dir=kwargs.get("prediction_dir"),
random_seed=kwargs.get("random_seed"),
step_buffer_context=kwargs.get("step_buffer_context", ""),
output_requirements=[
"- There is no hidden reference block. Use only the question, candidate files, tool trace, student output, and evaluation result to infer what intermediate state is worth probing.",
"- The instruction must explicitly request a short <analysis>...</analysis> block before the final <answer>...</answer>.",
"- The readout should focus on selected document/file, evidence span or table, extracted value, units, and any date or fiscal-period normalization.",
"- Do not ask for exhaustive copying of source text or a full chain-of-thought.",
"- The instruction text should be ready to append directly to the student's prompt.",
],
metadata_builder=lambda item: {
"id": str(item.get("id")),
"task_type": str(item.get("task_type") or "officeqa"),
"question_preview": str(item.get("question") or "")[:200],
"source_files": item.get("source_files", []),
"source_docs": item.get("source_docs", []),
},
)
def get_task_types(self) -> list[str]:
seen: list[str] = []
for item in self.dataloader.train_items + self.dataloader.val_items + self.dataloader.test_items:
task_type = str(item.get("task_type") or "officeqa")
if task_type not in seen:
seen.append(task_type)
return seen or ["officeqa"]
+71
View File
@@ -0,0 +1,71 @@
from __future__ import annotations
import csv
import json
import os
from pathlib import Path
from skillopt.datasets.base import SplitDataLoader
def _parse_list_field(value: str | list[str] | None) -> 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()]
if "," in text and not text.lower().endswith(".txt"):
return [part.strip() for part in text.split(",") if part.strip()]
return [text]
def _normalize_row(row: dict[str, str]) -> dict:
item_id = str(row.get("uid") or row.get("id") or "").strip()
question = str(row.get("question") or "").strip()
ground_truth = str(row.get("ground_truth") or row.get("answer") or "").strip()
task_type = str(row.get("category") or row.get("difficulty") or "officeqa").strip() or "officeqa"
source_files = _parse_list_field(row.get("source_files"))
source_docs = _parse_list_field(row.get("source_docs"))
split = str(row.get("split") or "").strip()
return {
"id": item_id,
"uid": item_id,
"question": question,
"ground_truth": ground_truth,
"answers": [ground_truth] if ground_truth else [],
"task_type": task_type,
"category": task_type,
"source_files": source_files,
"source_docs": source_docs,
"split": split,
}
class OfficeQADataLoader(SplitDataLoader):
def load_split_items(self, split_path: str) -> list[dict]:
path = Path(split_path)
csv_files = sorted(path.glob("*.csv"))
if csv_files:
with csv_files[0].open(encoding="utf-8", newline="") as f:
reader = csv.DictReader(f)
return [_normalize_row(row) for row in reader]
json_files = sorted(path.glob("*.json"))
if json_files:
with json_files[0].open(encoding="utf-8") as f:
data = json.load(f)
if not isinstance(data, list):
raise ValueError(f"Expected JSON array in {json_files[0]}")
return [_normalize_row(item) for item in data]
raise FileNotFoundError(f"No .csv or .json file found in {split_path}")
+46
View File
@@ -0,0 +1,46 @@
from __future__ import annotations
import re
import string
from collections import Counter
_NUMERIC_CHARS = set("0123456789.-")
def normalize_answer(text: str) -> str:
text = text.lower().strip()
text = text.replace(",", "")
text = "".join(ch for ch in text if ch not in string.punctuation or ch in _NUMERIC_CHARS or ch == "%")
text = re.sub(r"\b(million|millions|billion|billions|dollars|dollar|nominal)\b", " ", text)
text = " ".join(text.split())
return text
def exact_match(prediction: str, gold: str) -> float:
return 1.0 if normalize_answer(prediction) == normalize_answer(gold) else 0.0
def token_f1(prediction: str, gold: str) -> float:
pred_tokens = normalize_answer(prediction).split()
gold_tokens = normalize_answer(gold).split()
if not pred_tokens or not gold_tokens:
return 1.0 if pred_tokens == gold_tokens else 0.0
common = Counter(pred_tokens) & Counter(gold_tokens)
n_common = sum(common.values())
if n_common == 0:
return 0.0
precision = n_common / len(pred_tokens)
recall = n_common / len(gold_tokens)
return 2 * precision * recall / (precision + recall)
def evaluate(prediction: str, gold: str) -> dict:
em = exact_match(prediction, gold)
f1 = token_f1(prediction, gold)
return {
"em": em,
"f1": f1,
"predicted_answer": prediction.strip(),
"gold_answer": gold,
}
@@ -0,0 +1,37 @@
You are an expert failure-analysis agent for OfficeQA document-retrieval question answering tasks.
You will be given MULTIPLE failed OfficeQA trajectories from a single minibatch and the current skill document. The trajectories may include local document tool calls such as file search, grep, and partial reads.
Your job is to identify COMMON failure patterns across the batch and propose concise skill edits.
## Failure Type Categories
- retrieval_miss: the agent searched the wrong file or failed to narrow to the right file
- evidence_miss: the agent read documents but missed the decisive evidence span
- operand_error: the agent extracted the wrong value or the wrong operands
- calculation_error: the agent identified the right evidence but computed the result incorrectly
- answer_format: the agent reached the right result but formatted it wrong
- other: none of the above
## Rules
- Focus on patterns common across multiple trajectories.
- Prefer general retrieval and evidence-grounding rules over task-specific hacks.
- Only patch gaps in the skill; do not duplicate rules already present.
- Do not hardcode file names, years, or question-specific constants unless the pattern truly requires a reusable retrieval heuristic.
Respond ONLY with a valid JSON object (no markdown fences, no extra text):
{
"batch_size": <number of trajectories analysed>,
"failure_summary": [
{"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
],
"patch": {
"reasoning": "<why these edits address the batch's common failures>",
"edits": [
{"op": "append", "content": "<markdown to add at end of skill>"},
{"op": "insert_after", "target": "<exact heading/text to insert after>", "content": "<markdown>"},
{"op": "replace", "target": "<exact text to replace>", "content": "<replacement>"},
{"op": "delete", "target": "<exact text to remove>"}
]
}
}
Only include edits that are needed. "edits" can be an empty list if no patch is warranted.
@@ -0,0 +1,25 @@
You are an expert success-pattern analyst for OfficeQA document-retrieval question answering tasks.
You will be given MULTIPLE successful OfficeQA trajectories from a single minibatch and the current skill document. Your job is to identify common retrieval, evidence-selection, and numeric-grounding behaviors worth encoding in the skill.
## Rules
- Focus on patterns shared across multiple successful trajectories.
- Prefer reusable retrieval and extraction discipline over question-specific tips.
- Reinforce compact, high-value behaviors such as narrowing files early, reading only the relevant span, building a clean operand ledger, and copying the final answer from checked evidence.
- Only propose patches for patterns not already captured in the current skill.
Respond ONLY with a valid JSON object:
{
"batch_size": <number of trajectories analysed>,
"success_patterns": ["<pattern 1>", "<pattern 2>"],
"patch": {
"reasoning": "<why these patterns are worth encoding>",
"edits": [
{"op": "append", "content": "<markdown>"},
{"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
{"op": "replace", "target": "<old text>", "content": "<new text>"},
{"op": "delete", "target": "<exact text to remove>"}
]
}
}
"edits" may be empty if the skill already covers all observed patterns.
@@ -0,0 +1,15 @@
You are an expert OfficeQA agent working over local Treasury bulletin text files.
{skill_section}## Rules
1. Use only the provided local document tools to inspect candidate files.
2. Narrow to the most relevant file before reading long passages.
3. Prefer short targeted searches, then small reads around matching evidence.
4. Do not invent values that are not grounded in the retrieved text.
5. When the question requires arithmetic, compute only after extracting the exact operands.
6. If you have enough evidence, return the final answer inside <answer>...</answer>.
## Tool Use
Use the provided function tools directly when you need them. Prefer searching and small reads before answering. Do not ask the user for permission to use tools; just call the tools.
## Final Answer Format
When you are ready to answer, emit the final answer inside <answer>...</answer> and do not request another tool.
+363
View File
@@ -0,0 +1,363 @@
from __future__ import annotations
import json
import os
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
from skillopt.envs.officeqa.evaluator import evaluate
from skillopt.envs.officeqa.tool_runtime import resolve_candidate_files, resolve_docs_roots, run_tool
from skillopt.model import chat_student_messages, get_student_backend, is_student_exec_backend
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
from skillopt.prompts import load_prompt
_TOOL_SCHEMAS = [
{
"type": "function",
"function": {
"name": "glob",
"description": "Find candidate local document files by filename or relative-path glob pattern.",
"parameters": {
"type": "object",
"properties": {"pattern": {"type": "string"}},
"required": ["pattern"],
},
},
},
{
"type": "function",
"function": {
"name": "read",
"description": "Read a local text document excerpt by path and line window.",
"parameters": {
"type": "object",
"properties": {
"path": {"type": "string"},
"start": {"type": "integer"},
"limit": {"type": "integer"},
},
"required": ["path"],
},
},
},
{
"type": "function",
"function": {
"name": "grep",
"description": "Search a local text document for a literal pattern and return matching lines.",
"parameters": {
"type": "object",
"properties": {
"pattern": {"type": "string"},
"path": {"type": "string"},
},
"required": ["pattern", "path"],
},
},
},
]
_FINAL_RE = re.compile(r"<answer>(.*?)</answer>", re.IGNORECASE | re.DOTALL)
def _build_system(skill_content: str) -> str:
if skill_content.strip():
skill_section = f"## Skill\n{skill_content.strip()}\n\n"
else:
skill_section = ""
return load_prompt("rollout_system", env="officeqa").format(skill_section=skill_section)
def _build_user(
item: dict,
candidate_files: list[str],
*,
diagnostic_mode: bool = False,
diagnostic_instruction: str = "",
corpus_note: str = "",
) -> str:
file_block = "\n".join(f"- {path}" for path in candidate_files[:20]) or "- none resolved"
parts = [f"## Question\n{item['question']}"]
if corpus_note.strip():
parts.append(f"## Document Corpus\n{corpus_note.strip()}")
parts.append(f"## Candidate Files\n{file_block}")
if item.get("source_docs"):
parts.append("## Source Hints\n" + "\n".join(f"- {hint}" for hint in item["source_docs"]))
if diagnostic_mode and diagnostic_instruction.strip():
parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
return "\n\n".join(parts)
def _extract_answer(text: str) -> str:
match = _FINAL_RE.search(text)
if match:
return match.group(1).strip()
lines = [line.strip() for line in text.splitlines() if line.strip()]
return lines[-1] if lines else text.strip()
def _docs_link_targets(docs_roots: list[str]) -> list[tuple[str, str]]:
return [(root, os.path.join("docs", f"root_{idx}")) for idx, root in enumerate(docs_roots, start=1)]
def _workspace_doc_path(path: str, docs_roots: list[str]) -> str:
resolved_path = os.path.realpath(path)
for idx, root in enumerate(docs_roots, start=1):
resolved_root = os.path.realpath(root)
if resolved_path == resolved_root or resolved_path.startswith(resolved_root + os.sep):
rel_path = os.path.relpath(resolved_path, resolved_root)
return os.path.join("docs", f"root_{idx}", rel_path)
return path
def _build_codex_skill(skill_content: str) -> str:
return render_skill_md(
skill_content,
description="Dynamic ReflACT skill for solving the current OfficeQA local-document question.",
preamble=(
"Use this skill when answering the current OfficeQA question.\n"
"Inspect the provided local document excerpts or files, ground the answer in the evidence,\n"
"and return the final answer inside <answer>...</answer>."
),
)
def _run_codex_once(
*,
pred_dir: str,
item: dict,
skill_content: str,
candidate_files: list[str],
docs_roots: list[str],
model: str,
timeout: int,
diagnostic_mode: bool = False,
diagnostic_instruction: str = "",
previous_response: str = "",
) -> tuple[str, str, str, str]:
rel_files = [_workspace_doc_path(path, docs_roots) for path in candidate_files[:20]]
corpus_note = (
"The full OfficeQA document corpus is available under `docs/`. "
"The candidate files below are source hints or likely starting points; search the full corpus if needed."
)
user = _build_user(
item,
rel_files,
diagnostic_mode=diagnostic_mode,
diagnostic_instruction=diagnostic_instruction,
corpus_note=corpus_note,
)
task_parts = [user]
if previous_response:
task_parts.append(
"## Previous Attempt\n"
f"{previous_response}\n\n"
"Review the local documents again and correct the answer if needed."
)
task_text = "\n\n".join(task_parts)
skill_md = _build_codex_skill(skill_content)
work_dir = os.path.join(pred_dir, "codex_exec")
prepare_workspace(
work_dir=work_dir,
skill_md=skill_md,
task_text=task_text,
link_dirs=_docs_link_targets(docs_roots),
)
prompt = (
"Use the `skillopt-student` skill available in this workspace.\n"
"Read `task.md`, inspect or search the full OfficeQA corpus under `docs/`, and answer the question.\n"
"Treat candidate files in `task.md` as hints, not an access limit.\n"
"Return the final answer inside <answer>...</answer>."
)
final_message, raw = run_student_exec(
work_dir=work_dir,
prompt=prompt,
model=model,
timeout=timeout,
data_dirs=docs_roots,
)
return final_message or raw, raw, skill_md, task_text
def process_one(
item: dict,
out_root: str,
skill_content: str,
*,
max_tool_turns: int = 12,
data_dirs: list[str] | str | None = None,
diagnostic_mode: bool = False,
diagnostic_instruction: str = "",
) -> dict:
item_id = str(item["id"])
pred_dir = os.path.join(out_root, "predictions", item_id)
os.makedirs(pred_dir, exist_ok=True)
docs_roots = resolve_docs_roots(data_dirs)
candidate_files = resolve_candidate_files(item.get("source_files", []), docs_roots)
system = _build_system(skill_content)
user = _build_user(item, candidate_files, diagnostic_mode=diagnostic_mode, diagnostic_instruction=diagnostic_instruction)
messages: list[dict] = [
{"role": "system", "content": system},
{"role": "user", "content": user},
]
conversation: list[dict] = [{"role": "user", "content": user}]
final_response = ""
final_answer = ""
fail_reason = ""
allowed_files = [os.path.basename(path) for path in candidate_files]
try:
if is_student_exec_backend():
from skillopt.model import azure_openai as _llm
response = ""
system = ""
user = ""
for turn in range(1, max_tool_turns + 1):
response, _raw, system, user = _run_codex_once(
pred_dir=pred_dir,
item=item,
skill_content=skill_content,
candidate_files=candidate_files,
docs_roots=docs_roots,
model=_llm.STUDENT_DEPLOYMENT,
timeout=180,
diagnostic_mode=diagnostic_mode if turn == 1 else False,
diagnostic_instruction=diagnostic_instruction if turn == 1 else "",
previous_response=response if turn > 1 else "",
)
final_response = response
conversation.append({"type": "message", "turn": turn, "content": response})
if "<answer>" in response.lower():
final_answer = _extract_answer(response)
break
if not final_answer:
fail_reason = f"Exceeded codex turn budget ({max_tool_turns})"
system = system or _build_codex_skill(skill_content)
user = user or _build_user(item, [_workspace_doc_path(path, docs_roots) for path in candidate_files])
else:
for turn in range(1, max_tool_turns + 1):
message, _ = chat_student_messages(
messages=messages,
max_completion_tokens=768,
retries=5,
stage="rollout",
tools=_TOOL_SCHEMAS,
tool_choice="auto",
return_message=True,
)
response = message.content or ""
final_response = response
assistant_message = {"role": "assistant", "content": response}
if getattr(message, "tool_calls", None):
assistant_message["tool_calls"] = [tool_call.model_dump(mode="json") for tool_call in message.tool_calls]
messages.append(assistant_message)
conversation.append({"type": "message", "content": response})
if getattr(message, "tool_calls", None):
for tool_call in message.tool_calls:
tool_name = tool_call.function.name
arguments = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
cmd, obs = run_tool(tool_name, arguments, allowed_roots=docs_roots, allowed_files=allowed_files)
conversation.append({"type": "tool_call", "cmd": cmd, "obs": obs})
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": obs,
})
continue
if "<answer>" in response.lower():
final_answer = _extract_answer(response)
break
if turn == max_tool_turns:
fail_reason = f"Exceeded tool-turn budget ({max_tool_turns})"
else:
fail_reason = "Model neither produced a tool request nor a final answer"
break
except Exception as e: # noqa: BLE001
fail_reason = f"error: {e}"
with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
f.write(system)
with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w", encoding="utf-8") as f:
f.write(user)
with open(os.path.join(pred_dir, "conversation.json"), "w", encoding="utf-8") as f:
json.dump(conversation, f, ensure_ascii=False, indent=2)
eval_result = evaluate(final_answer, item.get("ground_truth", "")) if final_answer else {"em": 0.0, "f1": 0.0, "predicted_answer": "", "gold_answer": item.get("ground_truth", "")}
result = {
"id": item_id,
"question": item.get("question", ""),
"task_type": item.get("task_type", "officeqa"),
"task_description": item.get("question", ""),
"predicted_answer": eval_result["predicted_answer"],
"response": final_response,
"ground_truth": item.get("ground_truth", ""),
"source_files": item.get("source_files", []),
"resolved_source_paths": candidate_files,
"hard": int(eval_result["em"]),
"soft": eval_result["f1"],
"fail_reason": fail_reason or ("" if eval_result["em"] else f"predicted '{eval_result['predicted_answer']}' but expected '{item.get('ground_truth', '')}'"),
"agent_ok": not fail_reason,
"n_turns": len(conversation),
"student_system_prompt": system,
"student_user_prompt": user,
}
return result
def run_batch(
items: list[dict],
out_root: str,
skill_content: str,
*,
workers: int = 8,
max_tool_turns: int = 12,
data_dirs: list[str] | str | None = None,
diagnostic_mode: bool = False,
diagnostic_instruction: str = "",
) -> list[dict]:
results_path = os.path.join(out_root, "results.jsonl")
os.makedirs(out_root, exist_ok=True)
done_ids: set[str] = set()
existing: list[dict] = []
if os.path.exists(results_path):
with open(results_path, encoding="utf-8") as f:
for line in f:
try:
row = json.loads(line)
except json.JSONDecodeError:
continue
done_ids.add(str(row.get("id")))
existing.append(row)
pending = [item for item in items if str(item["id"]) not in done_ids]
if not pending:
return existing
results = list(existing)
with open(results_path, "a", encoding="utf-8") as outf, ThreadPoolExecutor(max_workers=workers) as ex:
futs = {
ex.submit(
process_one,
item,
out_root,
skill_content,
max_tool_turns=max_tool_turns,
data_dirs=data_dirs,
diagnostic_mode=diagnostic_mode,
diagnostic_instruction=diagnostic_instruction,
): item
for item in pending
}
for fut in as_completed(futs):
res = fut.result()
results.append(res)
outf.write(json.dumps(res, ensure_ascii=False) + "\n")
outf.flush()
return results
+15
View File
@@ -0,0 +1,15 @@
# OfficeQA Skill
## Retrieval Discipline
- Start by narrowing to the most likely candidate file before reading long passages.
- Prefer targeted search terms that name the exact entity, period, measure, or table concept from the question.
- After a promising match, read only a small surrounding span and verify it matches the requested year, basis, and unit.
## Evidence Discipline
- Extract the exact value from the retrieved text before doing any arithmetic.
- Keep track of each operand's period, unit, and semantic role so nearby proxy values are not mixed in.
- If the question asks for a transformed or derived quantity, compute only after confirming every operand.
## Final Answer Discipline
- Return the final answer only after one last consistency check against the retrieved evidence.
- Copy the final answer from a checked value, not from an unverified intermediate guess.
+134
View File
@@ -0,0 +1,134 @@
from __future__ import annotations
import fnmatch
import os
from pathlib import Path
_MAX_READ_CHARS = 4000
_MAX_GREP_MATCHES = 20
_MAX_GLOB_MATCHES = 50
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 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}]"