diff --git a/agent-framework/prometheus_swarm/database/models.py b/agent-framework/prometheus_swarm/database/models.py index 5a4785aa..58d30260 100644 --- a/agent-framework/prometheus_swarm/database/models.py +++ b/agent-framework/prometheus_swarm/database/models.py @@ -2,8 +2,9 @@ from datetime import datetime from typing import Optional, List -from sqlmodel import SQLModel, Field, Relationship - +from sqlmodel import SQLModel, Field, Relationship, Column, UniqueConstraint +from sqlalchemy import CheckConstraint +from pydantic import field_validator, constr, ConfigDict class Conversation(SQLModel, table=True): """Conversation model.""" @@ -42,3 +43,34 @@ class Log(SQLModel, table=True): stack_trace: Optional[str] = None request_id: Optional[str] = None additional_data: Optional[str] = None + + +class Evidence(SQLModel, table=True): + """Evidence model to ensure uniqueness.""" + + __tablename__ = 'evidence' + + model_config = ConfigDict( + validate_assignment=True, + from_attributes=True + ) + + id: Optional[int] = Field(default=None, primary_key=True) + content: str = Field(min_length=1, max_length=10000) + type: str = Field(index=True) + created_at: datetime = Field(default_factory=datetime.utcnow) + + __table_args__ = ( + UniqueConstraint('content', 'type', name='uix_evidence_content_type'), + ) + + @field_validator('content') + @classmethod + def validate_non_empty(cls, v): + """Validate content is not just whitespace""" + if not isinstance(v, str): + raise ValueError("Content must be a string") + v = v.strip() + if not v: + raise ValueError("Content cannot be empty or just whitespace") + return v \ No newline at end of file diff --git a/agent-framework/prometheus_swarm/tests/__init__.py b/agent-framework/prometheus_swarm/tests/__init__.py new file mode 100644 index 00000000..739954cb --- /dev/null +++ b/agent-framework/prometheus_swarm/tests/__init__.py @@ -0,0 +1 @@ +# Tests package \ No newline at end of file diff --git a/agent-framework/prometheus_swarm/tests/integration/__init__.py b/agent-framework/prometheus_swarm/tests/integration/__init__.py new file mode 100644 index 00000000..e27cd7ab --- /dev/null +++ b/agent-framework/prometheus_swarm/tests/integration/__init__.py @@ -0,0 +1 @@ +# Integration tests package \ No newline at end of file diff --git a/agent-framework/prometheus_swarm/tests/integration/test_evidence_uniqueness.py b/agent-framework/prometheus_swarm/tests/integration/test_evidence_uniqueness.py new file mode 100644 index 00000000..2e22c1f0 --- /dev/null +++ b/agent-framework/prometheus_swarm/tests/integration/test_evidence_uniqueness.py @@ -0,0 +1,83 @@ +import pytest +from sqlmodel import SQLModel, create_engine, Session, select +from prometheus_swarm.database.models import Evidence +from pydantic import ValidationError + +def create_test_engine(): + """Create an in-memory SQLite engine for testing.""" + # Using SQLite in-memory database for testing + return create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False}) + +def test_evidence_uniqueness(): + """ + Integration test to verify evidence uniqueness across the system. + + This test ensures that: + 1. Duplicate evidence with same type cannot be inserted + 2. Each evidence entry has a unique identifier + 3. Constraints prevent duplicate evidence submissions + """ + # Create test engine and tables + engine = create_test_engine() + SQLModel.metadata.create_all(engine) + + def create_evidence(session: Session, content: str, type: str) -> Evidence: + """Helper function to create evidence entries""" + evidence = Evidence(content=content, type=type) + session.add(evidence) + session.commit() + return evidence + + with Session(engine) as session: + # First evidence submission should succeed + evidence1 = create_evidence(session, "Test Content 1", "test_type") + assert evidence1.id is not None + + # Attempt to insert identical evidence should fail + with pytest.raises(Exception) as excinfo: + create_evidence(session, "Test Content 1", "test_type") + + assert "UNIQUE constraint" in str(excinfo.value) + +def test_evidence_multiple_types(): + """ + Test that different types of evidence can have the same content + """ + engine = create_test_engine() + SQLModel.metadata.create_all(engine) + + with Session(engine) as session: + # Create evidence with same content, different types + evidence1 = Evidence(content="Shared Content", type="type1") + evidence2 = Evidence(content="Shared Content", type="type2") + + session.add(evidence1) + session.add(evidence2) + session.commit() + + # Verify both evidences were saved + saved_evidences = session.exec(select(Evidence).where(Evidence.content == "Shared Content")).all() + assert len(saved_evidences) == 2 + assert {e.type for e in saved_evidences} == {"type1", "type2"} + +def test_evidence_content_validation(): + """ + Test validation of evidence content + """ + # Test empty content + with pytest.raises(ValidationError) as excinfo: + Evidence(content="", type="test_type") + assert "String should have at least 1 character" in str(excinfo.value) + + # Whitespace-only content + with pytest.raises(ValidationError) as excinfo: + Evidence(content=" ", type="test_type") + assert "Content cannot be empty or just whitespace" in str(excinfo.value) + + # Test very long content + long_content = "a" * 10001 # Exceeding max length + with pytest.raises(ValidationError) as excinfo: + Evidence(content=long_content, type="test_type") + + # Validate content cannot be completely empty or too long + assert "String should have at most" in str(excinfo.value) \ No newline at end of file