From 7bb4110d28673a53b04da3dfab174b839c3d3a07 Mon Sep 17 00:00:00 2001 From: CTO Agent Date: Wed, 8 Jul 2026 11:49:31 -0400 Subject: [PATCH] perf: reuse extract's file list when writing the traces README write_traces_readme re-walked the whole traces directory with rglob even though the extract flow just wrote the files and already holds the exact list. Accept an optional trace_files argument (exclusion filters still apply) and pass extract's copied_files from both the CLI and studio extraction paths. Callers without a known list, like the upload flow, keep the rescan. Co-Authored-By: Claude Fable 5 --- src/teich/cli.py | 4 +++- src/teich/studio/extraction.py | 5 +++-- src/teich/trace_readme.py | 6 +++++- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/teich/cli.py b/src/teich/cli.py index 3aa0a8a..8df6fb7 100644 --- a/src/teich/cli.py +++ b/src/teich/cli.py @@ -373,6 +373,7 @@ def _write_extract_readme( *, model_filter: str | None = None, repo_id: str | None = None, + trace_files: list[Path] | None = None, ) -> Path: cfg = _extract_dataset_config(provider, output, model_filter=model_filter, repo_id=repo_id) return write_traces_readme( @@ -382,6 +383,7 @@ def _write_extract_readme( model_id=None, repo_id=cfg.get_publish_repo_id(), extraction_provider=provider, + trace_files=trace_files, ) @@ -465,7 +467,7 @@ def _run_extract_command( if stale_readme_path.exists() and stale_readme_path.is_file(): stale_readme_path.unlink() anonymize_report = None if skip_anonymize else anonymize_path(output, output, in_place=True) - readme_path = _write_extract_readme(provider, output, model_filter=model_filter) + readme_path = _write_extract_readme(provider, output, model_filter=model_filter, trace_files=result.copied_files) extracted_message = f"Extracted {result.count} {provider} trace{'s' if result.count != 1 else ''}" if model_filter: extracted_message += f" with {model_filter}" diff --git a/src/teich/studio/extraction.py b/src/teich/studio/extraction.py index 2d828b0..3b74847 100644 --- a/src/teich/studio/extraction.py +++ b/src/teich/studio/extraction.py @@ -152,7 +152,7 @@ def emit_progress(event: dict[str, Any]) -> None: "totals": self.anonymize_totals, } ) - self._write_readme() + self._write_readme(result.copied_files) self._emit_status( "completed", f"Extracted {self.result_rows} trace row(s) across {len(self.result_files)} file(s).", @@ -164,7 +164,7 @@ def emit_progress(event: dict[str, Any]) -> None: self.finished_at = datetime.now(timezone.utc) self.events.close() - def _write_readme(self) -> None: + def _write_readme(self, trace_files: list[Path] | None = None) -> None: cfg = _extract_dataset_config(self.provider, self.output_dir, model_filter=self.model_filter) readme_path = write_traces_readme( cfg.output.traces_dir, @@ -173,6 +173,7 @@ def _write_readme(self) -> None: model_id=None, repo_id=cfg.get_publish_repo_id(), extraction_provider=self.provider, + trace_files=trace_files, ) self.events.append({"kind": "extract_readme", "text": f"Wrote {readme_path.name}.", "path": str(readme_path)}) diff --git a/src/teich/trace_readme.py b/src/teich/trace_readme.py index ed39e88..8b83096 100644 --- a/src/teich/trace_readme.py +++ b/src/teich/trace_readme.py @@ -523,10 +523,14 @@ def write_traces_readme( tools: list[dict[str, Any]] | None = None, excluded_dirs: list[Path] | None = None, extraction_provider: str | None = None, + trace_files: list[Path] | None = None, ) -> Path: + # Callers that just wrote the traces can pass the file list to skip the + # directory rescan; the exclusion filters still apply either way. + candidates = trace_files if trace_files is not None else traces_dir.rglob("*.jsonl") trace_files = sorted( path - for path in traces_dir.rglob("*.jsonl") + for path in candidates if path.is_file() and not {"partials", "failures"}.intersection(path.relative_to(traces_dir).parts) and not any(_path_is_relative_to(path, excluded_dir) for excluded_dir in excluded_dirs or [])