From 6e98d225458df0266c2aa497adbc87d5a3b0302c Mon Sep 17 00:00:00 2001 From: Aviraj <100823015+avirajsingh7@users.noreply.github.com> Date: Wed, 20 Aug 2025 15:07:50 +0530 Subject: [PATCH 1/2] take assistant id as input in create assistant api --- ...alter_unique_constraint_assistant_table.py | 33 +++++++++ backend/app/crud/assistants.py | 13 +++- backend/app/models/assistants.py | 14 +++- backend/app/tests/crud/test_assistants.py | 68 +++++++++++++++++++ 4 files changed, 125 insertions(+), 3 deletions(-) create mode 100644 backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py 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..6d194a26 --- /dev/null +++ b/backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py @@ -0,0 +1,33 @@ +"""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 From 4fb3bac612371fefd0bc53fe1f97d4698f3480e0 Mon Sep 17 00:00:00 2001 From: Aviraj <100823015+avirajsingh7@users.noreply.github.com> Date: Wed, 20 Aug 2025 15:31:59 +0530 Subject: [PATCH 2/2] pre commit --- ...alter_unique_constraint_assistant_table.py | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) 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 index 6d194a26..046aafbc 100644 --- a/backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py +++ b/backend/app/alembic/versions/38f0e8c8dc92_alter_unique_constraint_assistant_table.py @@ -11,23 +11,37 @@ # revision identifiers, used by Alembic. -revision = '38f0e8c8dc92' -down_revision = '5a59c6c29a82' +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']) + 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) + 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 ###