diff --git a/agent-framework/prometheus_swarm/database/__init__.py b/agent-framework/prometheus_swarm/database/__init__.py index f032618f..b24b77e4 100644 --- a/agent-framework/prometheus_swarm/database/__init__.py +++ b/agent-framework/prometheus_swarm/database/__init__.py @@ -2,6 +2,7 @@ from .database import get_db, get_session, initialize_database from .models import Conversation, Message, Log +from .evidence_validation import validate_evidence_uniqueness __all__ = [ "get_db", @@ -10,4 +11,5 @@ "Conversation", "Message", "Log", -] + "validate_evidence_uniqueness", +] \ No newline at end of file diff --git a/agent-framework/prometheus_swarm/database/evidence_validation.py b/agent-framework/prometheus_swarm/database/evidence_validation.py new file mode 100644 index 00000000..e397e8b9 --- /dev/null +++ b/agent-framework/prometheus_swarm/database/evidence_validation.py @@ -0,0 +1,49 @@ +from typing import Any, Dict, List +from .database import get_db + +def validate_evidence_uniqueness(evidence: Dict[str, Any], table_name: str = 'evidence') -> bool: + """ + Validate that an evidence entry is unique across the specified table. + + Args: + evidence (Dict[str, Any]): The evidence dictionary to validate + table_name (str, optional): The database table to check against. Defaults to 'evidence'. + + Returns: + bool: True if the evidence is unique, False otherwise + + Raises: + ValueError: If the evidence is invalid or missing key fields + """ + if not evidence: + raise ValueError("Evidence cannot be empty") + + # Define fields that must be unique + unique_fields = ['hash', 'source', 'type'] + + # Validate presence of unique fields + for field in unique_fields: + if field not in evidence: + raise ValueError(f"Missing required unique field: {field}") + + try: + # Get database session + db = get_db() + + # Construct query to check uniqueness + query = f"SELECT COUNT(*) as count FROM {table_name} WHERE " + conditions = [f"{field} = :{field}" for field in unique_fields] + query += " AND ".join(conditions) + + # Execute query + result = db.execute(query, evidence).scalar() + + # If count is 0, the evidence is unique + return result == 0 + + except Exception as e: + # Log the error in a real-world scenario + raise ValueError(f"Error validating evidence uniqueness: {e}") + finally: + if 'db' in locals(): + db.close() \ No newline at end of file diff --git a/agent-framework/tests/unit/test_evidence_validation.py b/agent-framework/tests/unit/test_evidence_validation.py new file mode 100644 index 00000000..a1caca4f --- /dev/null +++ b/agent-framework/tests/unit/test_evidence_validation.py @@ -0,0 +1,61 @@ +import pytest +from unittest.mock import patch, MagicMock +from prometheus_swarm.database.evidence_validation import validate_evidence_uniqueness + +def test_validate_evidence_uniqueness_valid(): + # Mock an empty database, so evidence should be unique + with patch('prometheus_swarm.database.evidence_validation.get_db') as mock_db: + mock_db_session = MagicMock() + mock_db.return_value = mock_db_session + mock_db_session.execute.return_value.scalar.return_value = 0 + + evidence = { + 'hash': 'unique_hash', + 'source': 'test_source', + 'type': 'test_type', + 'additional_data': 'some_info' + } + + assert validate_evidence_uniqueness(evidence) is True + +def test_validate_evidence_uniqueness_duplicate(): + # Mock a database with existing evidence + with patch('prometheus_swarm.database.evidence_validation.get_db') as mock_db: + mock_db_session = MagicMock() + mock_db.return_value = mock_db_session + mock_db_session.execute.return_value.scalar.return_value = 1 + + evidence = { + 'hash': 'duplicate_hash', + 'source': 'test_source', + 'type': 'test_type' + } + + assert validate_evidence_uniqueness(evidence) is False + +def test_validate_evidence_uniqueness_missing_fields(): + # Test missing unique fields + with pytest.raises(ValueError, match="Missing required unique field"): + validate_evidence_uniqueness({ + 'hash': 'some_hash', + 'source': 'test_source' + }) + +def test_validate_evidence_uniqueness_empty_evidence(): + # Test empty evidence + with pytest.raises(ValueError, match="Evidence cannot be empty"): + validate_evidence_uniqueness({}) + +def test_validate_evidence_uniqueness_database_error(): + # Test database connection error + with patch('prometheus_swarm.database.evidence_validation.get_db') as mock_db: + mock_db.side_effect = Exception("Database connection failed") + + evidence = { + 'hash': 'test_hash', + 'source': 'test_source', + 'type': 'test_type' + } + + with pytest.raises(ValueError, match="Error validating evidence uniqueness"): + validate_evidence_uniqueness(evidence) \ No newline at end of file