117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
from audita.core.chunking import TokenEstimatorProtocol, chunk_transcript
|
|
from audita.core.schemas import parse_transcript_json
|
|
|
|
|
|
class FakeEstimator(TokenEstimatorProtocol):
|
|
def estimate_json(self, value):
|
|
if len(value) == 1:
|
|
return 4
|
|
return len(value) * 4
|
|
|
|
|
|
def test_chunk_transcript_batches_sections_by_token_limit():
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
|
|
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
|
|
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
|
|
]
|
|
"""
|
|
)
|
|
|
|
sections = chunk_transcript(transcript, max_section_tokens=8, estimator=FakeEstimator())
|
|
|
|
assert len(sections) == 2
|
|
assert [segment.segment.id for segment in sections[0].segments] == [1, 2]
|
|
assert [segment.segment.id for segment in sections[1].segments] == [3]
|
|
|
|
|
|
def test_chunk_transcript_targets_llm_concurrency_when_feasible():
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
|
|
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
|
|
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
|
|
{"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"}
|
|
]
|
|
"""
|
|
)
|
|
|
|
sections = chunk_transcript(
|
|
transcript,
|
|
max_section_tokens=8,
|
|
min_section_tokens=4,
|
|
target_section_count=2,
|
|
estimator=FakeEstimator(),
|
|
)
|
|
|
|
assert len(sections) == 2
|
|
assert [segment.segment.id for segment in sections[0].segments] == [1, 2]
|
|
assert [segment.segment.id for segment in sections[1].segments] == [3, 4]
|
|
|
|
|
|
def test_chunk_transcript_increases_section_count_when_target_sections_exceed_max_tokens():
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
|
|
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
|
|
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
|
|
{"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"},
|
|
{"id": 5, "speaker": "A", "start": 4.0, "end": 5.0, "text": "five"}
|
|
]
|
|
"""
|
|
)
|
|
|
|
sections = chunk_transcript(
|
|
transcript,
|
|
max_section_tokens=8,
|
|
min_section_tokens=4,
|
|
target_section_count=1,
|
|
estimator=FakeEstimator(),
|
|
)
|
|
|
|
assert len(sections) > 1
|
|
assert all(section.token_count <= 8 for section in sections)
|
|
|
|
|
|
def test_chunk_transcript_reduces_section_count_when_target_sections_fall_below_min_tokens():
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
|
|
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
|
|
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
|
|
]
|
|
"""
|
|
)
|
|
|
|
sections = chunk_transcript(
|
|
transcript,
|
|
max_section_tokens=12,
|
|
min_section_tokens=8,
|
|
target_section_count=3,
|
|
estimator=FakeEstimator(),
|
|
)
|
|
|
|
assert len(sections) == 1
|
|
assert sections[0].token_count >= 8
|
|
|
|
|
|
def test_chunk_transcript_prompt_payload_includes_categories_when_present():
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one", "categories": ["intro", "aside"]}
|
|
]
|
|
"""
|
|
)
|
|
|
|
sections = chunk_transcript(transcript, max_section_tokens=8, estimator=FakeEstimator())
|
|
|
|
assert sections[0].prompt_payload() == [
|
|
{"id": 1, "original_text": "one", "categories": ["intro", "aside"]}
|
|
]
|