70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import List, Protocol
|
|
|
|
from .chunking import TranscriptSection
|
|
from .config import AuditaConfig
|
|
from .prompts import build_glossary_correction_messages, build_grammar_correction_messages
|
|
from .schemas import CorrectionCandidate, CorrectionSet, Glossary
|
|
|
|
|
|
class LLMClient(Protocol):
|
|
def create_corrections(self, messages: List[dict], config: AuditaConfig) -> CorrectionSet:
|
|
...
|
|
|
|
|
|
class CorrectionPass(Protocol):
|
|
def run(
|
|
self,
|
|
section: TranscriptSection,
|
|
glossary: Glossary,
|
|
config: AuditaConfig,
|
|
run_dir: Path,
|
|
retry_pass: bool = False,
|
|
) -> List[CorrectionCandidate]:
|
|
...
|
|
|
|
|
|
class GlossaryCorrectionPass:
|
|
def __init__(self, llm_client: LLMClient) -> None:
|
|
self._llm_client = llm_client
|
|
|
|
def run(
|
|
self,
|
|
section: TranscriptSection,
|
|
glossary: Glossary,
|
|
config: AuditaConfig,
|
|
run_dir: Path,
|
|
retry_pass: bool = False,
|
|
) -> List[CorrectionCandidate]:
|
|
messages = build_glossary_correction_messages(section, glossary, retry_pass=retry_pass)
|
|
prompt_path = run_dir / f"prompt-{section.section_index:04d}.json"
|
|
prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
response = self._llm_client.create_corrections(messages, config)
|
|
response_path = run_dir / f"corrections-{section.section_index:04d}.json"
|
|
response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8")
|
|
return list(response.corrections)
|
|
|
|
|
|
class GrammarCorrectionPass:
|
|
def __init__(self, llm_client: LLMClient) -> None:
|
|
self._llm_client = llm_client
|
|
|
|
def run(
|
|
self,
|
|
section: TranscriptSection,
|
|
glossary: Glossary,
|
|
config: AuditaConfig,
|
|
run_dir: Path,
|
|
retry_pass: bool = False,
|
|
) -> List[CorrectionCandidate]:
|
|
messages = build_grammar_correction_messages(section, glossary, retry_pass=retry_pass)
|
|
prompt_path = run_dir / f"prompt-{section.section_index:04d}.json"
|
|
prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
response = self._llm_client.create_corrections(messages, config)
|
|
response_path = run_dir / f"corrections-{section.section_index:04d}.json"
|
|
response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8")
|
|
return list(response.corrections)
|