158 lines
5.3 KiB
Python
158 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Tests for Training Pair Provenance Tracking
|
|
"""
|
|
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
import pytest
|
|
|
|
from training_pair_provenance import ProvenanceTracker, load_jsonl, save_jsonl
|
|
|
|
|
|
class TestProvenanceTracker:
|
|
"""Test the ProvenanceTracker class."""
|
|
|
|
def test_init(self):
|
|
"""Test tracker initialization."""
|
|
tracker = ProvenanceTracker()
|
|
assert tracker.stats["total_pairs"] == 0
|
|
assert tracker.stats["pairs_with_provenance"] == 0
|
|
assert tracker.stats["pairs_without_provenance"] == 0
|
|
|
|
def test_generate_pair_id(self):
|
|
"""Test pair ID generation."""
|
|
tracker = ProvenanceTracker()
|
|
pair = {"prompt": "test", "chosen": "response", "rejected": "bad"}
|
|
|
|
id1 = tracker.generate_pair_id(pair)
|
|
id2 = tracker.generate_pair_id(pair)
|
|
|
|
# Same content should generate same ID
|
|
assert id1 == id2
|
|
assert len(id1) == 16
|
|
|
|
def test_add_provenance(self):
|
|
"""Test adding provenance to a pair."""
|
|
tracker = ProvenanceTracker()
|
|
pair = {"prompt": "test", "chosen": "response", "rejected": "bad"}
|
|
|
|
result = tracker.add_provenance(pair, source_session_id="session123", model="test-model")
|
|
|
|
assert "provenance" in result
|
|
assert result["provenance"]["source_session_id"] == "session123"
|
|
assert result["provenance"]["model"] == "test-model"
|
|
assert "timestamp" in result["provenance"]
|
|
assert result["provenance"]["source"] == "curated"
|
|
assert "content_hash" in result["provenance"]
|
|
|
|
def test_extract_provenance_from_existing(self):
|
|
"""Test extracting provenance from existing fields."""
|
|
tracker = ProvenanceTracker()
|
|
pair = {
|
|
"id": "session456",
|
|
"model": "claude-3-opus",
|
|
"started_at": "2024-01-01T00:00:00Z",
|
|
"conversations": [{"from": "human", "value": "test"}]
|
|
}
|
|
|
|
provenance = tracker.extract_provenance_from_existing(pair)
|
|
|
|
assert provenance["source_session_id"] == "session456"
|
|
assert provenance["model"] == "claude-3-opus"
|
|
assert provenance["timestamp"] == "2024-01-01T00:00:00Z"
|
|
assert provenance["source"] == "curated"
|
|
assert "content_hash" in provenance
|
|
|
|
def test_process_pair(self):
|
|
"""Test processing a pair."""
|
|
tracker = ProvenanceTracker()
|
|
pair = {"id": "test123", "model": "test-model", "conversations": []}
|
|
|
|
result = tracker.process_pair(pair)
|
|
|
|
assert tracker.stats["total_pairs"] == 1
|
|
assert tracker.stats["pairs_without_provenance"] == 1
|
|
assert "provenance" in result
|
|
|
|
def test_filter_by_provenance(self):
|
|
"""Test filtering pairs by provenance."""
|
|
tracker = ProvenanceTracker()
|
|
|
|
pairs = [
|
|
{"provenance": {"model": "anthropic/claude-3-opus"}},
|
|
{"provenance": {"model": "gpt-4"}},
|
|
{"provenance": {"model": "anthropic/claude-3-sonnet"}},
|
|
]
|
|
|
|
filtered = tracker.filter_by_provenance(pairs, exclude_models=["anthropic/claude-3-opus", "anthropic/claude-3-sonnet"])
|
|
|
|
assert len(filtered) == 1
|
|
assert filtered[0]["provenance"]["model"] == "gpt-4"
|
|
assert tracker.stats["excluded"] == 2
|
|
|
|
def test_generate_report(self):
|
|
"""Test report generation."""
|
|
tracker = ProvenanceTracker()
|
|
tracker.stats = {
|
|
"total_pairs": 10,
|
|
"pairs_with_provenance": 8,
|
|
"pairs_without_provenance": 2,
|
|
"by_model": {"gpt-4": 5, "claude-3": 3},
|
|
"by_source": {"curated": 8},
|
|
"excluded": 0
|
|
}
|
|
|
|
report = tracker.generate_report()
|
|
|
|
assert "Total pairs: 10" in report
|
|
assert "Pairs with provenance: 8" in report
|
|
assert "gpt-4: 5" in report
|
|
|
|
|
|
class TestJsonlFunctions:
|
|
"""Test JSONL load/save functions."""
|
|
|
|
def test_load_jsonl(self):
|
|
"""Test loading JSONL file."""
|
|
with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
|
|
f.write('{"id": "1", "value": "test1"}\n')
|
|
f.write('{"id": "2", "value": "test2"}\n')
|
|
f.write('{"id": "3", "value": "test3"}\n')
|
|
temp_path = Path(f.name)
|
|
|
|
try:
|
|
entries = load_jsonl(temp_path)
|
|
assert len(entries) == 3
|
|
assert entries[0]["id"] == "1"
|
|
assert entries[2]["value"] == "test3"
|
|
finally:
|
|
temp_path.unlink()
|
|
|
|
def test_save_jsonl(self):
|
|
"""Test saving JSONL file."""
|
|
entries = [
|
|
{"id": "1", "value": "test1"},
|
|
{"id": "2", "value": "test2"}
|
|
]
|
|
|
|
with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
|
|
temp_path = Path(f.name)
|
|
|
|
try:
|
|
save_jsonl(entries, temp_path)
|
|
|
|
with open(temp_path) as f:
|
|
lines = f.readlines()
|
|
|
|
assert len(lines) == 2
|
|
assert json.loads(lines[0])["id"] == "1"
|
|
assert json.loads(lines[1])["value"] == "test2"
|
|
finally:
|
|
temp_path.unlink()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|