diff --git a/backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py b/backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py new file mode 100644 index 00000000..046aafbc --- /dev/null +++ b/backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py @@ -0,0 +1,47 @@ +"""Alter unique constraint assistant table + +Revision ID: 38f0e8c8dc92 +Revises: 5a59c6c29a82 +Create Date: 2025-08-20 15:02:39.151977 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes + + +# revision identifiers, used by Alembic. +revision = "38f0e8c8dc92" +down_revision = "5a59c6c29a82" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index("ix_openai_assistant_assistant_id", table_name="openai_assistant") + op.create_index( + op.f("ix_openai_assistant_assistant_id"), + "openai_assistant", + ["assistant_id"], + unique=False, + ) + op.create_unique_constraint( + "uq_project_assistant_id", "openai_assistant", ["project_id", "assistant_id"] + ) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint("uq_project_assistant_id", "openai_assistant", type_="unique") + op.drop_index( + op.f("ix_openai_assistant_assistant_id"), table_name="openai_assistant" + ) + op.create_index( + "ix_openai_assistant_assistant_id", + "openai_assistant", + ["assistant_id"], + unique=True, + ) + # ### end Alembic commands ### diff --git a/backend/app/crud/assistants.py b/backend/app/crud/assistants.py index 6a8e8e5b..62478aa5 100644 --- a/backend/app/crud/assistants.py +++ b/backend/app/crud/assistants.py @@ -173,8 +173,19 @@ def create_assistant( ) -> Assistant: verify_vector_store_ids_exist(openai_client, assistant.vector_store_ids) + assistant.assistant_id = assistant.assistant_id or str(uuid.uuid4()) + + existing = get_assistant_by_id(session, assistant.assistant_id, project_id) + if existing: + logger.error( + f"[create_assistant] Assistant with ID {mask_string(assistant.assistant_id)} already exists. | project_id: {project_id}" + ) + raise HTTPException( + status_code=409, + detail=f"Assistant with ID {assistant.assistant_id} already exists.", + ) + assistant = Assistant( - assistant_id=str(uuid.uuid4()), **assistant.model_dump(exclude_unset=True), project_id=project_id, organization_id=organization_id, diff --git a/backend/app/models/assistants.py b/backend/app/models/assistants.py index b97e25b1..a67ae6be 100644 --- a/backend/app/models/assistants.py +++ b/backend/app/models/assistants.py @@ -3,13 +3,17 @@ from sqlalchemy import Column, String, Text from sqlalchemy.dialects.postgresql import ARRAY -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field, Relationship, SQLModel, UniqueConstraint from app.core.util import now class AssistantBase(SQLModel): - assistant_id: str = Field(index=True, unique=True) + __table_args__ = ( + UniqueConstraint("project_id", "assistant_id", name="uq_project_assistant_id"), + ) + + assistant_id: str = Field(index=True) name: str instructions: str = Field(sa_column=Column(Text, nullable=False)) model: str @@ -45,6 +49,12 @@ class AssistantCreate(SQLModel): instructions: str = Field( description="Instructions for the assistant", min_length=10 ) + assistant_id: str | None = Field( + default=None, + description="Unique identifier for the assistant", + min_length=3, + max_length=50, + ) model: str = Field( default="gpt-4o", description="Model name for the assistant", diff --git a/backend/app/tests/crud/test_assistants.py b/backend/app/tests/crud/test_assistants.py index 322eed77..227fdb31 100644 --- a/backend/app/tests/crud/test_assistants.py +++ b/backend/app/tests/crud/test_assistants.py @@ -157,6 +157,74 @@ def test_create_assistant_success(self, mock_vector_store_ids_exist, db: Session assert result.temperature == assistant_create.temperature assert result.max_num_results == assistant_create.max_num_results + @patch("app.crud.assistants.verify_vector_store_ids_exist") + def test_create_assistant_with_id_success( + self, mock_vector_store_ids_exist, db: Session + ): + """Assistant is created with a specific ID when vector store IDs are valid""" + project = get_project(db) + assistant_create = AssistantCreate( + name="Test Assistant", + instructions="Test instructions", + model="gpt-4o", + vector_store_ids=["vs_1", "vs_2"], + temperature=0.7, + max_num_results=10, + assistant_id="test_assistant_id", + ) + client = OpenAI(api_key="test_key") + mock_vector_store_ids_exist.return_value = None + result = create_assistant( + db, client, assistant_create, project.id, project.organization_id + ) + + assert result.name == assistant_create.name + assert result.instructions == assistant_create.instructions + assert result.model == assistant_create.model + assert result.vector_store_ids == assistant_create.vector_store_ids + assert result.temperature == assistant_create.temperature + assert result.max_num_results == assistant_create.max_num_results + assert result.assistant_id == assistant_create.assistant_id + + @patch("app.crud.assistants.verify_vector_store_ids_exist") + def test_create_assistant_duplicate_assistant_id( + self, mock_vector_store_ids_exist, db: Session + ): + """Creating an assistant with a duplicate assistant_id should raise 409 Conflict""" + project = get_project(db) + + assistant_id = "duplicate_id" + assistant_create_1 = AssistantCreate( + name="Assistant One", + instructions="First assistant instructions", + model="gpt-4o", + vector_store_ids=[], + assistant_id=assistant_id, + ) + client = OpenAI(api_key="test_key") + mock_vector_store_ids_exist.return_value = None + create_assistant( + db, client, assistant_create_1, project.id, project.organization_id + ) + + assistant_create_2 = AssistantCreate( + name="Assistant Two", + instructions="Second assistant instructions", + model="gpt-4o", + vector_store_ids=[], + assistant_id=assistant_id, + ) + + with pytest.raises(HTTPException) as exc_info: + create_assistant( + db, client, assistant_create_2, project.id, project.organization_id + ) + + assert exc_info.value.status_code == 409 + assert f"Assistant with ID {assistant_id} already exists." in str( + exc_info.value.detail + ) + @patch("app.crud.assistants.verify_vector_store_ids_exist") def test_create_assistant_vector_store_invalid( self, mock_vector_store_ids_exist, db: Session