diff --git a/src/models/database/prompts.py b/src/models/database/prompts.py index a10d97913..4368174a6 100644 --- a/src/models/database/prompts.py +++ b/src/models/database/prompts.py @@ -3,10 +3,9 @@ from datetime import datetime from sqlalchemy import JSON, DateTime, ForeignKey, String, Text, func -from sqlalchemy.orm import Mapped, mapped_column, validates +from sqlalchemy.orm import Mapped, mapped_column from models.database.base import Base -from prompts.exceptions import PromptValidationError class Prompt(Base): # pylint: disable=too-few-public-methods @@ -50,11 +49,3 @@ class PromptVersion(Base): # pylint: disable=too-few-public-methods created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now() ) - - @validates("variables") - def validate_unique_variables(self, _key: str, value: list[str]) -> list[str]: - """Ensure no duplicate variable names.""" - if len(value) != len(set(value)): - raise PromptValidationError("The variables contain duplicate names") - - return value diff --git a/src/prompts/exceptions.py b/src/prompts/exceptions.py index 8b23ce113..109bdec58 100644 --- a/src/prompts/exceptions.py +++ b/src/prompts/exceptions.py @@ -7,3 +7,7 @@ class PromptError(Exception): class PromptValidationError(PromptError): """Invalid prompt field values.""" + + +class PromptNotFoundError(PromptError): + """No prompt exists for the given identifier.""" diff --git a/src/prompts/store.py b/src/prompts/store.py new file mode 100644 index 000000000..1cd88e858 --- /dev/null +++ b/src/prompts/store.py @@ -0,0 +1,201 @@ +"""Store operations for the versioned prompt catalog.""" + +from typing import Optional +from uuid import uuid4 + +from app.database import get_session +from log import get_logger +from models.database.prompts import Prompt, PromptVersion +from prompts.exceptions import PromptNotFoundError, PromptValidationError + +logger = get_logger(__name__) + + +def _validate_variables(variables: list[str]) -> None: + """Reject duplicate variable names. + + Parameters: + variables: Declared variable names to validate. + + Raises: + PromptValidationError: If the list contains duplicate names. + + """ + if len(variables) != len(set(variables)): + raise PromptValidationError("The variables contain duplicate names") + + +def create_prompt(name: str, template: str, variables: list[str]) -> PromptVersion: + """Create a new prompt with its initial version. + + Parameters: + name: Display name for the prompt. + template: Mustache template text. + variables: Declared variable names for the template. + + Returns: + PromptVersion: The initial version row. + + Raises: + PromptValidationError: If variables contain duplicates. + + """ + # TODO: template validation # pylint: disable=fixme + + _validate_variables(variables) + + prompt_id = f"pmpt_{uuid4()}" + version = 1 + + with get_session() as session: + prompt = Prompt(id=prompt_id, name=name, default_version=version) + prompt_version = PromptVersion( + prompt_id=prompt_id, version=version, template=template, variables=variables + ) + + session.add(prompt) + session.add(prompt_version) + session.commit() + logger.debug("Created prompt id=%s with version %s", prompt_id, version) + session.refresh(prompt_version) + + return prompt_version + + +def get_prompt(prompt_id: str, version: Optional[int] = None) -> PromptVersion: + """Retrieve a prompt version by identifier. + + When version is None the prompt's default version is used. + + Parameters: + prompt_id: The prompt identifier. + version: Specific version number, or None for the default. + + Returns: + PromptVersion: The matching version row. + + Raises: + PromptNotFoundError: If the prompt or requested version does not exist. + + """ + with get_session() as session: + + # Grab default version from Prompt if version is not provided + if version is None: + prompt = session.query(Prompt).filter_by(id=prompt_id).first() + if prompt is None: + logger.debug("Prompt id=%s not found", prompt_id) + raise PromptNotFoundError(f"Prompt id={prompt_id} not found") + + version = prompt.default_version + + prompt_version = ( + session.query(PromptVersion) + .filter_by(prompt_id=prompt_id, version=version) + .first() + ) + + if prompt_version is None: + logger.debug("Version %s of prompt id=%s not found", version, prompt_id) + raise PromptNotFoundError( + f"Version {version} of prompt id={prompt_id} not found" + ) + + return prompt_version + + +def list_prompts() -> list[Prompt]: + """List all prompts in the catalog. + + Returns: + list[Prompt]: All prompt identity rows. + + """ + with get_session() as session: + prompts = session.query(Prompt).all() + return prompts + + +def update_prompt( + prompt_id: str, template: str, variables: list[str], *, default: bool = False +) -> PromptVersion: + """Create a new immutable version for an existing prompt. + + Parameters: + prompt_id: The prompt identifier. + template: Mustache template text for the new version. + variables: Declared variable names for the template. + default: If True, set the new version as the prompt's default. + + Returns: + PromptVersion: The newly created version row. + + Raises: + PromptNotFoundError: If the prompt does not exist. + PromptValidationError: If variables contain duplicates. + + """ + # TODO: validate template # pylint: disable=fixme + + _validate_variables(variables) + + with get_session() as session: + max_prompt_version = ( + session.query(PromptVersion) + .filter_by(prompt_id=prompt_id) + .order_by(PromptVersion.version.desc()) + .first() + ) + if max_prompt_version is None: + logger.debug("No available version found for prompt id=%s", prompt_id) + raise PromptNotFoundError( + f"No available version found for prompt id={prompt_id}" + ) + + max_version = max_prompt_version.version + + prompt_version = PromptVersion( + prompt_id=prompt_id, + version=max_version + 1, + template=template, + variables=variables, + ) + session.add(prompt_version) + logger.debug( + "Created version %s for prompt id=%s", prompt_version.version, prompt_id + ) + + if default: + prompt = session.query(Prompt).filter_by(id=prompt_id).first() + if prompt is None: + logger.debug("Prompt id=%s not found", prompt_id) + raise PromptNotFoundError(f"Prompt id={prompt_id} not found") + + prompt.default_version = prompt_version.version + + session.commit() + session.refresh(prompt_version) + + return prompt_version + + +def delete_prompt(prompt_id: str) -> None: + """Delete a prompt and all its versions. + + Parameters: + prompt_id: The prompt identifier. + + Raises: + PromptNotFoundError: If the prompt does not exist. + + """ + with get_session() as session: + prompt = session.query(Prompt).filter_by(id=prompt_id).first() + if prompt is None: + raise PromptNotFoundError(f"Prompt id={prompt_id} not found") + + # This cover the situation when FK enforcement is not on + session.query(PromptVersion).filter_by(prompt_id=prompt_id).delete() + session.delete(prompt) + session.commit() + logger.debug("Deleted prompt id=%s", prompt_id) diff --git a/tests/unit/models/database/test_prompts.py b/tests/unit/models/database/test_prompts.py index 6e998d20a..37a5c1665 100644 --- a/tests/unit/models/database/test_prompts.py +++ b/tests/unit/models/database/test_prompts.py @@ -11,7 +11,6 @@ from models.database.base import Base from models.database.prompts import Prompt, PromptVersion -from prompts.exceptions import PromptValidationError @pytest.fixture(name="sqlite_engine") @@ -265,47 +264,3 @@ def test_empty_variables_defaults_to_list( result = db_session.get(PromptVersion, (sample_prompt.id, 1)) assert result is not None assert result.variables == [] - - -class TestPromptVersionVariableValidation: - """Test cases for PromptVersion variable uniqueness validation.""" - - def test_unique_variables_accepted(self) -> None: - """Test that a list of unique variable names is accepted.""" - version = PromptVersion( - prompt_id="pmpt_001", - version=1, - template="{{a}} {{b}}", - variables=["a", "b"], - ) - assert version.variables == ["a", "b"] - - def test_duplicate_variables_raises_validation_error(self) -> None: - """Test that duplicate variable names raise PromptValidationError.""" - with pytest.raises(PromptValidationError, match="duplicate names"): - PromptVersion( - prompt_id="pmpt_001", - version=1, - template="{{a}} {{a}}", - variables=["a", "a"], - ) - - def test_empty_variables_accepted(self) -> None: - """Test that an empty variable list is accepted.""" - version = PromptVersion( - prompt_id="pmpt_001", - version=1, - template="No variables", - variables=[], - ) - assert version.variables == [] - - def test_single_variable_accepted(self) -> None: - """Test that a single variable is accepted.""" - version = PromptVersion( - prompt_id="pmpt_001", - version=1, - template="{{name}}", - variables=["name"], - ) - assert version.variables == ["name"] diff --git a/tests/unit/prompts/__init__.py b/tests/unit/prompts/__init__.py new file mode 100644 index 000000000..6fc051c4a --- /dev/null +++ b/tests/unit/prompts/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the prompts module.""" diff --git a/tests/unit/prompts/test_store.py b/tests/unit/prompts/test_store.py new file mode 100644 index 000000000..5885fdea6 --- /dev/null +++ b/tests/unit/prompts/test_store.py @@ -0,0 +1,356 @@ +"""Unit tests for prompts store operations.""" + +from collections.abc import Generator +from contextlib import contextmanager +from typing import Literal, Optional, overload + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.pool import StaticPool + +from models.database.base import Base +from models.database.prompts import Prompt, PromptVersion +from prompts.exceptions import PromptNotFoundError, PromptValidationError +from prompts.store import ( + create_prompt, + delete_prompt, + get_prompt, + list_prompts, + update_prompt, +) + +DUMMY_CREATE_PROMPT_PARAMS = ( + "test-prompt", + "Hello {{name}}! {{greeting}}!", + ["name", "greeting"], +) +DUMMY_UPDATE_PROMPT_PARAMS = ("{{a}} {{b}}", ["a", "b"]) +DUMMY_ID = "pmpt_testing" + + +@overload +def _setup_db( + session: Session, + with_v2: Literal[True], +) -> tuple[Prompt, PromptVersion, PromptVersion]: ... + + +@overload +def _setup_db( + session: Session, with_v2: Literal[False] = False +) -> tuple[Prompt, PromptVersion, None]: ... + + +def _setup_db( + session: Session, with_v2: bool = False +) -> tuple[Prompt, PromptVersion, Optional[PromptVersion]]: + """Insert a Prompt with one or two versions for test setup. + + Parameters: + session: Active database session. + with_v2: If True, also insert a second version. + + Returns: + A tuple of (Prompt, PromptVersion v1, PromptVersion v2 or None). + + """ + name, template_v1, variables_v1 = DUMMY_CREATE_PROMPT_PARAMS + template_v2, variables_v2 = DUMMY_UPDATE_PROMPT_PARAMS + version = 1 + + prompt = Prompt(id=DUMMY_ID, name=name, default_version=version) + prompt_version_v1 = PromptVersion( + prompt_id=DUMMY_ID, + version=version, + template=template_v1, + variables=variables_v1, + ) + + session.add(prompt) + session.add(prompt_version_v1) + session.commit() + + session.refresh(prompt) + session.refresh(prompt_version_v1) + + if with_v2: + prompt_version_v2 = PromptVersion( + prompt_id=DUMMY_ID, + version=version + 1, + template=template_v2, + variables=variables_v2, + ) + session.add(prompt_version_v2) + session.commit() + session.refresh(prompt_version_v2) + + return (prompt, prompt_version_v1, prompt_version_v2) + + return (prompt, prompt_version_v1, None) + + +@pytest.fixture(name="sqlite_engine") +def sqlite_engine_fixture() -> Generator[Engine, None, None]: + """Provide an in-memory SQLite engine with all tables created. + + Yields: + Engine: A SQLAlchemy Engine bound to an in-memory SQLite database with + tables created from Base.metadata. + + """ + engine = create_engine( + "sqlite:///:memory:", + echo=False, + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + yield engine + engine.dispose() + + +@pytest.fixture(name="db_session") +def db_session_fixture(mocker, sqlite_engine: Engine) -> Generator[Session, None, None]: + """Provide a database session for tests. + + Parameters: + sqlite_engine: In-memory SQLite engine fixture. + + Yields: + Session: A SQLAlchemy Session bound to the test engine. + + """ + session_factory = sessionmaker( + autocommit=False, autoflush=False, bind=sqlite_engine + ) + session = session_factory() + + @contextmanager + def mock_get_session(): + yield session + + mocker.patch("prompts.store.get_session", mock_get_session) + + try: + yield session + finally: + session.close() + + +class TestPromptStoreOperation: + """Test cases for prompt store CRUD operations.""" + + def test_create_prompt(self, db_session: Session) -> None: + """Test creating a prompt inserts both Prompt and PromptVersion rows.""" + name, template, variables = DUMMY_CREATE_PROMPT_PARAMS + + prompt_version = create_prompt(*DUMMY_CREATE_PROMPT_PARAMS) + + # Prompt + prompt_qy_result = ( + db_session.query(Prompt).filter_by(id=prompt_version.prompt_id).all() + ) + assert len(prompt_qy_result) == 1 + prompt_in_db = prompt_qy_result[0] + + assert prompt_in_db.name == name + assert prompt_in_db.default_version == 1 + + # PromptVersion + prompt_version_qy_result = ( + db_session.query(PromptVersion) + .filter_by(prompt_id=prompt_version.prompt_id) + .all() + ) + assert len(prompt_version_qy_result) == 1 + prompt_version_in_db = prompt_version_qy_result[0] + + assert prompt_version_in_db.created_at == prompt_version.created_at + assert prompt_version_in_db.template == template + assert prompt_version_in_db.variables == variables + assert prompt_version_in_db.version == 1 + + def test_create_prompt_with_invalid_variables(self, db_session: Session) -> None: + """Test creating a prompt with duplicate variables raises PromptValidationError.""" + _ = db_session + name, template, _ = DUMMY_CREATE_PROMPT_PARAMS + variables = ["name", "name"] + with pytest.raises(PromptValidationError, match="duplicate names"): + create_prompt(name, template, variables) + + def test_get_prompt(self, db_session: Session) -> None: + """Test retrieving a prompt by default version and by explicit version.""" + prompt, expected_pv_v1, expected_pv_v2 = _setup_db(db_session, True) + + # Version 1 + prompt_version_v1 = get_prompt(prompt.id) + assert prompt_version_v1.version == expected_pv_v1.version + assert prompt_version_v1.template == expected_pv_v1.template + assert prompt_version_v1.variables == expected_pv_v1.variables + assert prompt_version_v1.prompt_id == expected_pv_v1.prompt_id + assert prompt_version_v1.created_at == expected_pv_v1.created_at + + # Version 2 + prompt_version_v2 = get_prompt(prompt.id, expected_pv_v2.version) + assert prompt_version_v2.version == expected_pv_v2.version + assert prompt_version_v2.template == expected_pv_v2.template + assert prompt_version_v2.variables == expected_pv_v2.variables + assert prompt_version_v2.prompt_id == expected_pv_v2.prompt_id + + def test_get_prompt_with_invalid_version(self, db_session: Session) -> None: + """Test retrieving a nonexistent version raises PromptNotFoundError.""" + prompt, _, _ = _setup_db(db_session) + + with pytest.raises( + PromptNotFoundError, match=r"Version \d+ of prompt id=\bpmpt_\w+ not found" + ): + get_prompt(prompt.id, 101) + + def test_get_prompt_with_nonexistent_id(self, db_session: Session) -> None: + """Test retrieving a prompt with a nonexistent id raises PromptNotFoundError.""" + _setup_db(db_session) + + with pytest.raises( + PromptNotFoundError, match=r"Prompt id=\bpmpt_\w+ not found" + ): + get_prompt("pmpt_nonexistent") + + def test_update_prompt(self, db_session: Session) -> None: + """Test updating a prompt appends an immutable version and optionally sets default.""" + prompt, pv_1, _ = _setup_db(db_session) + template_v2, variables_v2 = DUMMY_UPDATE_PROMPT_PARAMS + + # Since the session is not closed during the entire test, we need to save the values + # before the operation for the comparison. Otherwise, we will always get the newest + # values by the lazily reload. + ori_default_version = prompt.default_version + + cur_pmpt_version_counts = ( + db_session.query(PromptVersion).filter_by(prompt_id=prompt.id).count() + ) + # Only one PromptVersion exist before the update + assert cur_pmpt_version_counts == 1 + + pv_2 = update_prompt(prompt.id, *DUMMY_UPDATE_PROMPT_PARAMS) + + # Without default=True, the default_version should stay the same + assert prompt.default_version == ori_default_version + + cur_pmpt_version_counts = ( + db_session.query(PromptVersion).filter_by(prompt_id=prompt.id).count() + ) + # update_prompt creates a new PromptVersion in the database + assert cur_pmpt_version_counts == 2 + + assert pv_2.version == pv_1.version + 1 + assert pv_2.prompt_id == prompt.id + assert pv_2.template == template_v2 + assert pv_2.variables == variables_v2 + + template_v3, variables_v3 = ("{{c}} {{d}}", ["c", "d"]) + pv_3 = update_prompt(prompt.id, template_v3, variables_v3, default=True) + # With default=True, the default version should be updated to the new version + assert prompt.default_version == pv_3.version + + assert pv_3.template == template_v3 + assert pv_3.variables == variables_v3 + + def test_update_prompt_with_no_prompt_version_exists( + self, db_session: Session + ) -> None: + """Test updating a prompt with no existing versions raises PromptNotFoundError.""" + prompt, _, _ = _setup_db(db_session) + + db_session.query(PromptVersion).filter_by(prompt_id=prompt.id).delete() + + with pytest.raises( + PromptNotFoundError, + match=r"No available version found for prompt id=\bpmpt_\w+", + ): + update_prompt(prompt.id, "{{x}} {{y}}", ["x", "y"]) + + def test_update_prompt_with_no_prompt_exists_when_default_is_set_to_true( + self, db_session: Session + ) -> None: + """Test updating with default=True when prompt is deleted raises PromptNotFoundError.""" + prompt, _, _ = _setup_db(db_session) + + db_session.query(Prompt).filter_by(id=prompt.id).delete() + + with pytest.raises( + PromptNotFoundError, match=r"Prompt id=\bpmpt_\w+ not found" + ): + update_prompt(prompt.id, "{{x}} {{y}}", ["x", "y"], default=True) + + def test_update_prompt_with_duplicate_variables(self, db_session: Session) -> None: + """Test updating a prompt with duplicate variables raises PromptValidationError.""" + prompt, _, _ = _setup_db(db_session) + + with pytest.raises(PromptValidationError, match="duplicate names"): + update_prompt(prompt.id, "{{x}} {{x}}", ["x", "x"]) + + def test_list_prompts(self, db_session: Session) -> None: + """Test listing prompts returns all prompt identities.""" + _setup_db(db_session, True) + + prompts = list_prompts() + + assert len(prompts) == 1 + + second_prompt_id = "pmpt_testing2" + second_prompt_name = "Second Prompt" + new_prompt = Prompt( + id=second_prompt_id, name=second_prompt_name, default_version=1 + ) + db_session.add(new_prompt) + db_session.commit() + + prompts = list_prompts() + + assert len(prompts) == 2 + filtered_prompts = [pv for pv in prompts if pv.id == second_prompt_id] + assert len(filtered_prompts) == 1 + second_prompt = filtered_prompts[0] + assert second_prompt.name == second_prompt_name + + def test_list_prompts_empty(self, db_session: Session) -> None: + """Test listing prompts returns an empty list when no prompts exist.""" + _ = db_session + prompts = list_prompts() + + assert len(prompts) == 0 + + def test_delete_prompt(self, db_session: Session) -> None: + """Test deleting a prompt removes the prompt and all its versions.""" + prompt, _, _ = _setup_db(db_session, True) + + prompt_counts = db_session.query(Prompt).filter_by(id=prompt.id).count() + prompt_version_counts = ( + db_session.query(PromptVersion).filter_by(prompt_id=prompt.id).count() + ) + + assert prompt_counts == 1 + assert prompt_version_counts == 2 + + delete_prompt(prompt.id) + + prompt_counts = db_session.query(Prompt).filter_by(id=prompt.id).count() + prompt_version_counts = ( + db_session.query(PromptVersion).filter_by(prompt_id=prompt.id).count() + ) + + assert prompt_counts == 0 + assert prompt_version_counts == 0 + + def test_delete_prompt_with_nonexistent_prompt_id( + self, db_session: Session + ) -> None: + """Test deleting a nonexistent prompt raises PromptNotFoundError.""" + _setup_db(db_session) + + with pytest.raises( + PromptNotFoundError, match=r"Prompt id=\bpmpt_\w+ not found" + ): + delete_prompt("pmpt_nonexistent")