Files
DEMO-AGENT/review_agent/regulatory_info_package/services/package_generate.py

71 lines
2.9 KiB
Python

from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from review_agent.models import RegulatoryInfoPackageBatch
from review_agent.regulatory_info_package.constants import GENERATED_FILE_FAILED
from review_agent.regulatory_info_package.schemas import GeneratedFileResult, MergedField, TemplateSpec
from review_agent.regulatory_info_package.services.docx_document import write_docx_from_template
from review_agent.regulatory_info_package.services.legacy_doc_document import write_legacy_doc_or_fallback
from review_agent.regulatory_info_package.services.template_repository import copy_template_to_batch, template_specs
from review_agent.regulatory_info_package.storage import ensure_batch_subdir
def generate_package_documents(
batch: RegulatoryInfoPackageBatch,
config: dict,
merged_fields: dict[str, MergedField],
) -> list[GeneratedFileResult]:
specs = template_specs(config)
with ThreadPoolExecutor(max_workers=min(4, len(specs) or 1)) as executor:
futures = [executor.submit(_generate_one, batch, config, spec, merged_fields) for spec in specs]
return [future.result() for future in as_completed(futures)]
def _generate_one(
batch: RegulatoryInfoPackageBatch,
config: dict,
spec: TemplateSpec,
merged_fields: dict[str, MergedField],
) -> GeneratedFileResult:
try:
template_path = copy_template_to_batch(batch, config, spec)
generated_dir = ensure_batch_subdir(batch, "generated")
output_path = generated_dir / spec.output_name
adapter_summary = {}
if spec.file_format == "doc":
actual_path, status, adapter_summary = write_legacy_doc_or_fallback(template_path, output_path, merged_fields)
actual_format = actual_path.suffix.lower().lstrip(".")
highlight_count = missing_count = llm_only_count = 0
else:
highlight_count, missing_count, llm_only_count = write_docx_from_template(
template_path,
output_path,
merged_fields,
template_code=spec.code,
)
actual_path = output_path
actual_format = "docx"
status = "success"
return GeneratedFileResult(
template_code=spec.code,
file_name=actual_path.name,
requested_format=spec.file_format,
actual_format=actual_format,
status=status,
path=str(actual_path),
highlight_count=highlight_count,
missing_count=missing_count,
llm_only_count=llm_only_count,
)
except Exception as exc:
return GeneratedFileResult(
template_code=spec.code,
file_name=spec.output_name,
requested_format=spec.file_format,
actual_format=spec.file_format,
status=GENERATED_FILE_FAILED,
error_message=str(exc),
)