diff --git a/backend/app/seed_data/seed_data.json b/backend/app/seed_data/seed_data.json index 6eb4860a..bfaa929f 100644 --- a/backend/app/seed_data/seed_data.json +++ b/backend/app/seed_data/seed_data.json @@ -53,77 +53,4 @@ "is_deleted": false, "deleted_at": null } - ], - "credentials": [ - { - "is_active": true, - "provider": "openai", - "credential": "{\"api_key\": \"sk-proj-GlificI3i5SCxN\"}", - "project_name": "Glific", - "organization_name": "Project Tech4dev", - "deleted_at": null - }, - { - "is_active": true, - "provider": "openai", - "credential": "{\"api_key\": \"sk-proj-DalgoI3i5SCxN\"}", - "project_name": "Dalgo", - "organization_name": "Project Tech4dev", - "deleted_at": null - }, - { - "is_active": true, - "provider": "langfuse", - "credential": "{\"secret_key\": \"sk-lf-test-glific-secret\", \"public_key\": \"pk-lf-test-glific-public\", \"host\": \"https://cloud.langfuse.com\"}", - "project_name": "Glific", - "organization_name": "Project Tech4dev", - "deleted_at": null - }, - { - "is_active": true, - "provider": "langfuse", - "credential": "{\"secret_key\": \"sk-lf-test-dalgo-secret\", \"public_key\": \"pk-lf-test-dalgo-public\", \"host\": \"https://cloud.langfuse.com\"}", - "project_name": "Dalgo", - "organization_name": "Project Tech4dev", - "deleted_at": null - } - ], - "assistants": [ - { - "assistant_id": "assistant_glific", - "name": "Test Assistant Glific", - "instructions": "Test instructions", - "model": "gpt-4o", - "vector_store_ids": ["vs_glific"], - "temperature": 0.1, - "max_num_results": 20, - "project_name": "Glific", - "organization_name": "Project Tech4dev" - }, - { - "assistant_id": "assistant_dalgo", - "name": "Test Assistant Dalgo", - "instructions": "Test instructions", - "model": "gpt-4o", - "vector_store_ids": ["vs_dalgo"], - "temperature": 0.1, - "max_num_results": 20, - "project_name": "Dalgo", - "organization_name": "Project Tech4dev" - } - ], - "documents": [ - { - "fname": "glific_sample.json", - "object_store_url": "s3://test-bucket/glific_sample.json", - "project_name": "Glific", - "organization_name": "Project Tech4dev" - }, - { - "fname": "dalgo_sample.json", - "object_store_url": "s3://test-bucket/dalgo_sample.json", - "project_name": "Dalgo", - "organization_name": "Project Tech4dev" - } - ] -} + ]} diff --git a/backend/app/seed_data/seed_data.py b/backend/app/seed_data/seed_data.py index 2f890625..58e96e3c 100644 --- a/backend/app/seed_data/seed_data.py +++ b/backend/app/seed_data/seed_data.py @@ -9,15 +9,12 @@ from app.core.db import engine from app.core import settings -from app.core.security import get_password_hash, encrypt_credentials +from app.core.security import get_password_hash from app.models import ( APIKey, Organization, Project, User, - Credential, - Assistant, - Document, ) @@ -51,34 +48,6 @@ class APIKeyData(BaseModel): deleted_at: Optional[str] = None -class CredentialData(BaseModel): - is_active: bool - provider: str - credential: str - organization_name: str - project_name: str - deleted_at: Optional[str] = None - - -class AssistantData(BaseModel): - assistant_id: str - name: str - instructions: str - model: str - vector_store_ids: list[str] - temperature: float - max_num_results: int - project_name: str - organization_name: str - - -class DocumentData(BaseModel): - fname: str - object_store_url: str - organization_name: str - project_name: str - - def load_seed_data() -> dict: """Load seed data from JSON file.""" json_path = Path(__file__).parent / "seed_data.json" @@ -217,158 +186,13 @@ def create_api_key(session: Session, api_key_data_raw: dict) -> APIKey: raise -def create_credential(session: Session, credential_data_raw: dict) -> Credential: - """Create a credential from data.""" - try: - credential_data = CredentialData.model_validate(credential_data_raw) - logging.info(f"Creating credential for provider: {credential_data.provider}") - - # Query organization ID by name - organization = session.exec( - select(Organization).where( - Organization.name == credential_data.organization_name - ) - ).first() - if not organization: - raise ValueError( - f"Organization '{credential_data.organization_name}' not found" - ) - - # Query organization ID by name - project = session.exec( - select(Project).where(Project.name == credential_data.project_name) - ).first() - if not project: - raise ValueError(f"Project '{credential_data.project_name}' not found") - - # Encrypt the credential data - convert string to dict first, then encrypt - credential_dict = json.loads(credential_data.credential) - encrypted_credential = encrypt_credentials(credential_dict) - - credential = Credential( - is_active=credential_data.is_active, - provider=credential_data.provider, - credential=encrypted_credential, - organization_id=organization.id, - project_id=project.id, - deleted_at=credential_data.deleted_at, - ) - session.add(credential) - session.flush() # Ensure ID is assigned - return credential - except Exception as e: - logging.error(f"Error creating credential: {e}") - raise - - -def create_assistant(session: Session, assistant_data_raw: dict) -> Assistant: - """Create an assistant from data.""" - try: - assistant_data = AssistantData.model_validate(assistant_data_raw) - logging.info(f"Creating assistant: {assistant_data.name}") - - # Query organization ID by name - organization = session.exec( - select(Organization).where( - Organization.name == assistant_data.organization_name - ) - ).first() - if not organization: - raise ValueError( - f"Organization '{assistant_data.organization_name}' not found" - ) - - # Query project ID by name - project = session.exec( - select(Project).where(Project.name == assistant_data.project_name) - ).first() - if not project: - raise ValueError(f"Project '{assistant_data.project_name}' not found") - - assistant = Assistant( - assistant_id=assistant_data.assistant_id, - name=assistant_data.name, - instructions=assistant_data.instructions, - model=assistant_data.model, - vector_store_ids=assistant_data.vector_store_ids, - temperature=assistant_data.temperature, - max_num_results=assistant_data.max_num_results, - organization_id=organization.id, - project_id=project.id, - ) - session.add(assistant) - session.flush() # Ensure ID is assigned - return assistant - except Exception as e: - logging.error(f"Error creating assistant: {e}") - raise - - -def create_document(session: Session, document_data_raw: dict) -> Document: - """Create a document from seed data.""" - try: - document_data = DocumentData.model_validate(document_data_raw) - logging.info(f"Creating document: {document_data.fname}") - - # Get the organization - organization = session.exec( - select(Organization).where( - Organization.name == document_data.organization_name - ) - ).first() - - if not organization: - raise ValueError( - f"Organization '{document_data.organization_name}' not found" - ) - - # Get the project - project = session.exec( - select(Project).where( - Project.name == document_data.project_name, - Project.organization_id == organization.id, - ) - ).first() - if not project: - raise ValueError( - f"Project '{document_data.project_name}' not found in organization '{organization.name}'" - ) - - users = session.exec( - select(User) - .join(APIKey, APIKey.user_id == User.id) - .where(APIKey.organization_id == organization.id) - ).all() - - user = users[1] - if not user: - raise ValueError(f"No user found in organization '{organization.name}'") - - # Create and store document - document = Document( - fname=document_data.fname, - object_store_url=document_data.object_store_url, - project_id=project.id, - ) - - session.add(document) - session.flush() - return document - - except Exception as e: - logging.error(f"Error creating document: {e}") - raise - - def clear_database(session: Session) -> None: """Clear all seeded data from the database.""" logging.info("Clearing existing data...") - session.exec(delete(Assistant)) session.exec(delete(APIKey)) session.exec(delete(Project)) session.exec(delete(Organization)) session.exec(delete(User)) - session.exec(delete(Credential)) session.commit() logging.info("Existing data cleared.") @@ -382,13 +206,8 @@ def seed_database(session: Session) -> None: - Projects (Glific, Dalgo) - Users (superuser, test user) - API Keys for both users - - OpenAI Credentials for both projects (ensures all tests have credentials) - - Langfuse Credentials for both projects (for tracing and observability tests) - - Test Assistants for both projects - - Sample Documents - This seed data is used by the test suite and ensures that all tests - can rely on both OpenAI and Langfuse credentials being available without manual setup. + This seed data is used by the test suite. """ logging.info("Starting database seeding...") @@ -441,28 +260,6 @@ def seed_database(session: Session) -> None: api_keys.append(api_key) logging.info(f"Created API key (ID: {api_key.id})") - # Create credentials - credentials = [] - for credential_data in seed_data["credentials"]: - credential = create_credential(session, credential_data) - credentials.append(credential) - logging.info( - f"Created credential for provider: {credential.provider} (ID: {credential.id})" - ) - - # Create assistants - assistants = [] - for assistant_data in seed_data.get("assistants", []): - assistant = create_assistant(session, assistant_data) - assistants.append(assistant) - logging.info(f"Created assistant: {assistant.name} (ID: {assistant.id})") - - documents = [] - for document_data in seed_data.get("documents", []): - document = create_document(session, document_data) - documents.append(document) - logging.info(f"Created document: {document.fname} (ID: {document.id})") - logging.info("Database seeding completed successfully!") session.commit() except Exception as e: diff --git a/backend/app/tests/conftest.py b/backend/app/tests/conftest.py index d396a435..3cbe2591 100644 --- a/backend/app/tests/conftest.py +++ b/backend/app/tests/conftest.py @@ -21,7 +21,7 @@ get_user_test_auth_context, TestAuthContext, ) -from app.seed_data.seed_data import seed_database +from app.tests.seed_data.seed_data import seed_database @pytest.fixture(scope="function") diff --git a/backend/app/tests/seed_data/seed_data.json b/backend/app/tests/seed_data/seed_data.json new file mode 100644 index 00000000..9888cc84 --- /dev/null +++ b/backend/app/tests/seed_data/seed_data.json @@ -0,0 +1,129 @@ +{ + "_comment":"This data will be used in testing, modifying this will affect tests. Please ensure to update tests/utils/auth.py if you change this data.", + "organization": [ + { + "name": "Project Tech4dev", + "is_active": true + } + ], + "projects": [ + { + "name": "Glific", + "description": "Two way communication platform", + "is_active": true, + "organization_name": "Project Tech4dev" + }, + { + "name": "Dalgo", + "description": "Data platform for the social sector", + "is_active": true, + "organization_name": "Project Tech4dev" + } + ], + "users": [ + { + "email": "{{SUPERUSER_EMAIL}}", + "password": "changethis", + "full_name": "SUPERUSER", + "is_active": true, + "is_superuser": true + }, + { + "email": "{{ADMIN_EMAIL}}", + "password": "admin123", + "full_name": "ADMIN", + "is_active": true, + "is_superuser": false + } + ], + "apikeys": [ + { + "organization_name": "Project Tech4dev", + "user_email": "{{SUPERUSER_EMAIL}}", + "project_name": "Glific", + "api_key": "ApiKey No3x47A5qoIGhm0kVKjQ77dhCqEdWRIQZlEPzzzh7i8", + "is_deleted": false, + "deleted_at": null + }, + { + "organization_name": "Project Tech4dev", + "user_email": "{{ADMIN_EMAIL}}", + "project_name": "Dalgo", + "api_key": "ApiKey Px8y47B6roJHin1lWLkR88eiDrFdXSJRZmFQazzai8j", + "is_deleted": false, + "deleted_at": null + } + ], + "credentials": [ + { + "is_active": true, + "provider": "openai", + "credential": "{\"api_key\": \"sk-proj-GlificI3i5SCxN\"}", + "project_name": "Glific", + "organization_name": "Project Tech4dev", + "deleted_at": null + }, + { + "is_active": true, + "provider": "openai", + "credential": "{\"api_key\": \"sk-proj-DalgoI3i5SCxN\"}", + "project_name": "Dalgo", + "organization_name": "Project Tech4dev", + "deleted_at": null + }, + { + "is_active": true, + "provider": "langfuse", + "credential": "{\"secret_key\": \"sk-lf-test-glific-secret\", \"public_key\": \"pk-lf-test-glific-public\", \"host\": \"https://cloud.langfuse.com\"}", + "project_name": "Glific", + "organization_name": "Project Tech4dev", + "deleted_at": null + }, + { + "is_active": true, + "provider": "langfuse", + "credential": "{\"secret_key\": \"sk-lf-test-dalgo-secret\", \"public_key\": \"pk-lf-test-dalgo-public\", \"host\": \"https://cloud.langfuse.com\"}", + "project_name": "Dalgo", + "organization_name": "Project Tech4dev", + "deleted_at": null + } + ], + "assistants": [ + { + "assistant_id": "assistant_glific", + "name": "Test Assistant Glific", + "instructions": "Test instructions", + "model": "gpt-4o", + "vector_store_ids": ["vs_glific"], + "temperature": 0.1, + "max_num_results": 20, + "project_name": "Glific", + "organization_name": "Project Tech4dev" + }, + { + "assistant_id": "assistant_dalgo", + "name": "Test Assistant Dalgo", + "instructions": "Test instructions", + "model": "gpt-4o", + "vector_store_ids": ["vs_dalgo"], + "temperature": 0.1, + "max_num_results": 20, + "project_name": "Dalgo", + "organization_name": "Project Tech4dev" + } + ], + "documents": [ + { + "fname": "glific_sample.json", + "object_store_url": "s3://test-bucket/glific_sample.json", + "project_name": "Glific", + "organization_name": "Project Tech4dev" + }, + { + "fname": "dalgo_sample.json", + "object_store_url": "s3://test-bucket/dalgo_sample.json", + "project_name": "Dalgo", + "organization_name": "Project Tech4dev" + } + ] +} diff --git a/backend/app/tests/seed_data/seed_data.py b/backend/app/tests/seed_data/seed_data.py new file mode 100644 index 00000000..2f890625 --- /dev/null +++ b/backend/app/tests/seed_data/seed_data.py @@ -0,0 +1,483 @@ +import json +import logging +from datetime import datetime +from pathlib import Path +from typing import Optional + +from pydantic import BaseModel, EmailStr +from sqlmodel import Session, delete, select + +from app.core.db import engine +from app.core import settings +from app.core.security import get_password_hash, encrypt_credentials +from app.models import ( + APIKey, + Organization, + Project, + User, + Credential, + Assistant, + Document, +) + + +# Pydantic models for data validation +class OrgData(BaseModel): + name: str + is_active: bool + + +class ProjectData(BaseModel): + name: str + description: str + is_active: bool + organization_name: str + + +class UserData(BaseModel): + email: EmailStr + full_name: str + is_superuser: bool + is_active: bool + password: str + + +class APIKeyData(BaseModel): + organization_name: str + project_name: str + user_email: EmailStr + api_key: str + is_deleted: bool + deleted_at: Optional[str] = None + + +class CredentialData(BaseModel): + is_active: bool + provider: str + credential: str + organization_name: str + project_name: str + deleted_at: Optional[str] = None + + +class AssistantData(BaseModel): + assistant_id: str + name: str + instructions: str + model: str + vector_store_ids: list[str] + temperature: float + max_num_results: int + project_name: str + organization_name: str + + +class DocumentData(BaseModel): + fname: str + object_store_url: str + organization_name: str + project_name: str + + +def load_seed_data() -> dict: + """Load seed data from JSON file.""" + json_path = Path(__file__).parent / "seed_data.json" + try: + with open(json_path, "r") as f: + return json.load(f) + except FileNotFoundError: + logging.error(f"Error: Seed data file not found at {json_path}") + raise + except json.JSONDecodeError as e: + logging.error(f"Error: Failed to decode JSON from {json_path}: {e}") + raise + + +def create_organization(session: Session, org_data_raw: dict) -> Organization: + """Create an organization from data.""" + try: + org_data = OrgData.model_validate(org_data_raw) + logging.info(f"Creating organization: {org_data.name}") + organization = Organization(name=org_data.name, is_active=org_data.is_active) + session.add(organization) + session.flush() # Ensure ID is assigned + return organization + except Exception as e: + logging.error(f"Error creating organization: {e}") + raise + + +def create_project(session: Session, project_data_raw: dict) -> Project: + """Create a project from data.""" + try: + project_data = ProjectData.model_validate(project_data_raw) + logging.info(f"Creating project: {project_data.name}") + # Query organization ID by name + organization = session.exec( + select(Organization).where( + Organization.name == project_data.organization_name + ) + ).first() + if not organization: + raise ValueError( + f"Organization '{project_data.organization_name}' not found" + ) + project = Project( + name=project_data.name, + description=project_data.description, + is_active=project_data.is_active, + organization_id=organization.id, + ) + session.add(project) + session.flush() # Ensure ID is assigned + return project + except Exception as e: + logging.error(f"Error creating project: {e}") + raise + + +def create_user(session: Session, user_data_raw: dict) -> User: + """Create a user from data.""" + try: + user_data = UserData.model_validate(user_data_raw) + logging.info(f"Creating user: {user_data.email}") + hashed_password = get_password_hash(user_data.password) + user = User( + email=user_data.email, + full_name=user_data.full_name, + is_superuser=user_data.is_superuser, + is_active=user_data.is_active, + hashed_password=hashed_password, + ) + session.add(user) + session.flush() # Ensure ID is assigned + return user + except Exception as e: + logging.error(f"Error creating user: {e}") + raise + + +def create_api_key(session: Session, api_key_data_raw: dict) -> APIKey: + """Create an API key from data.""" + try: + api_key_data = APIKeyData.model_validate(api_key_data_raw) + logging.info(f"Creating API key for user {api_key_data.user_email}") + # Query organization ID by name + organization = session.exec( + select(Organization).where( + Organization.name == api_key_data.organization_name + ) + ).first() + if not organization: + raise ValueError( + f"Organization '{api_key_data.organization_name}' not found" + ) + project = session.exec( + select(Project).where(Project.name == api_key_data.project_name) + ).first() + if not project: + raise ValueError(f"Project '{api_key_data.project_name}' not found") + # Query user ID by email + user = session.exec( + select(User).where(User.email == api_key_data.user_email) + ).first() + if not user: + raise ValueError(f"User '{api_key_data.user_email}' not found") + + # Extract key_prefix from the provided API key and hash the full key + # API key format: "ApiKey {key_prefix}{random_key}" where key_prefix is 16 chars + raw_key = api_key_data.api_key + if not raw_key.startswith("ApiKey "): + raise ValueError(f"Invalid API key format: {raw_key}") + + # Extract the key_prefix (first 16 characters after "ApiKey ") + key_portion = raw_key[7:] # Remove "ApiKey " prefix + + key_prefix = key_portion[:12] # First 12 characters as prefix + + from passlib.context import CryptContext + + pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") + key_hash = pwd_context.hash(key_portion[12:]) + + api_key = APIKey( + organization_id=organization.id, + project_id=project.id, + user_id=user.id, + key_prefix=key_prefix, + key_hash=key_hash, + is_deleted=api_key_data.is_deleted, + deleted_at=api_key_data.deleted_at, + ) + session.add(api_key) + session.flush() + return api_key + except Exception as e: + logging.error(f"Error creating API key: {e}") + raise + + +def create_credential(session: Session, credential_data_raw: dict) -> Credential: + """Create a credential from data.""" + try: + credential_data = CredentialData.model_validate(credential_data_raw) + logging.info(f"Creating credential for provider: {credential_data.provider}") + + # Query organization ID by name + organization = session.exec( + select(Organization).where( + Organization.name == credential_data.organization_name + ) + ).first() + if not organization: + raise ValueError( + f"Organization '{credential_data.organization_name}' not found" + ) + + # Query organization ID by name + project = session.exec( + select(Project).where(Project.name == credential_data.project_name) + ).first() + if not project: + raise ValueError(f"Project '{credential_data.project_name}' not found") + + # Encrypt the credential data - convert string to dict first, then encrypt + credential_dict = json.loads(credential_data.credential) + encrypted_credential = encrypt_credentials(credential_dict) + + credential = Credential( + is_active=credential_data.is_active, + provider=credential_data.provider, + credential=encrypted_credential, + organization_id=organization.id, + project_id=project.id, + deleted_at=credential_data.deleted_at, + ) + session.add(credential) + session.flush() # Ensure ID is assigned + return credential + except Exception as e: + logging.error(f"Error creating credential: {e}") + raise + + +def create_assistant(session: Session, assistant_data_raw: dict) -> Assistant: + """Create an assistant from data.""" + try: + assistant_data = AssistantData.model_validate(assistant_data_raw) + logging.info(f"Creating assistant: {assistant_data.name}") + + # Query organization ID by name + organization = session.exec( + select(Organization).where( + Organization.name == assistant_data.organization_name + ) + ).first() + if not organization: + raise ValueError( + f"Organization '{assistant_data.organization_name}' not found" + ) + + # Query project ID by name + project = session.exec( + select(Project).where(Project.name == assistant_data.project_name) + ).first() + if not project: + raise ValueError(f"Project '{assistant_data.project_name}' not found") + + assistant = Assistant( + assistant_id=assistant_data.assistant_id, + name=assistant_data.name, + instructions=assistant_data.instructions, + model=assistant_data.model, + vector_store_ids=assistant_data.vector_store_ids, + temperature=assistant_data.temperature, + max_num_results=assistant_data.max_num_results, + organization_id=organization.id, + project_id=project.id, + ) + session.add(assistant) + session.flush() # Ensure ID is assigned + return assistant + except Exception as e: + logging.error(f"Error creating assistant: {e}") + raise + + +def create_document(session: Session, document_data_raw: dict) -> Document: + """Create a document from seed data.""" + try: + document_data = DocumentData.model_validate(document_data_raw) + logging.info(f"Creating document: {document_data.fname}") + + # Get the organization + organization = session.exec( + select(Organization).where( + Organization.name == document_data.organization_name + ) + ).first() + + if not organization: + raise ValueError( + f"Organization '{document_data.organization_name}' not found" + ) + + # Get the project + project = session.exec( + select(Project).where( + Project.name == document_data.project_name, + Project.organization_id == organization.id, + ) + ).first() + if not project: + raise ValueError( + f"Project '{document_data.project_name}' not found in organization '{organization.name}'" + ) + + users = session.exec( + select(User) + .join(APIKey, APIKey.user_id == User.id) + .where(APIKey.organization_id == organization.id) + ).all() + + user = users[1] + if not user: + raise ValueError(f"No user found in organization '{organization.name}'") + + # Create and store document + document = Document( + fname=document_data.fname, + object_store_url=document_data.object_store_url, + project_id=project.id, + ) + + session.add(document) + session.flush() + return document + + except Exception as e: + logging.error(f"Error creating document: {e}") + raise + + +def clear_database(session: Session) -> None: + """Clear all seeded data from the database.""" + logging.info("Clearing existing data...") + session.exec(delete(Assistant)) + session.exec(delete(APIKey)) + session.exec(delete(Project)) + session.exec(delete(Organization)) + session.exec(delete(User)) + session.exec(delete(Credential)) + session.commit() + logging.info("Existing data cleared.") + + +def seed_database(session: Session) -> None: + """ + Seed the database with initial test data. + + This function creates a complete test environment including: + - Organizations (Project Tech4dev) + - Projects (Glific, Dalgo) + - Users (superuser, test user) + - API Keys for both users + - OpenAI Credentials for both projects (ensures all tests have credentials) + - Langfuse Credentials for both projects (for tracing and observability tests) + - Test Assistants for both projects + - Sample Documents + + This seed data is used by the test suite and ensures that all tests + can rely on both OpenAI and Langfuse credentials being available without manual setup. + """ + logging.info("Starting database seeding...") + + try: + # Clear existing data first + clear_database(session) + + # Load seed data from JSON + seed_data = load_seed_data() + + # Create organizations + organizations = [] + for org_data in seed_data["organization"]: + organization = create_organization(session, org_data) + organizations.append(organization) + logging.info( + f"Created organization: {organization.name} (ID: {organization.id})" + ) + + for user_data in seed_data["users"]: + if user_data["email"] == "{{SUPERUSER_EMAIL}}": + user_data["email"] = settings.FIRST_SUPERUSER + elif user_data["email"] == "{{ADMIN_EMAIL}}": + user_data["email"] = settings.EMAIL_TEST_USER + + # Create users + users = [] + for user_data in seed_data["users"]: + user = create_user(session, user_data) + users.append(user) + logging.info(f"Created user: {user.email} (ID: {user.id})") + + # Create projects + projects = [] + for project_data in seed_data["projects"]: + project = create_project(session, project_data) + projects.append(project) + logging.info(f"Created project: {project.name} (ID: {project.id})") + + for api_key_data in seed_data["apikeys"]: + if api_key_data["user_email"] == "{{SUPERUSER_EMAIL}}": + api_key_data["user_email"] = settings.FIRST_SUPERUSER + elif api_key_data["user_email"] == "{{ADMIN_EMAIL}}": + api_key_data["user_email"] = settings.EMAIL_TEST_USER + + # Create API keys + api_keys = [] + for api_key_data in seed_data["apikeys"]: + api_key = create_api_key(session, api_key_data) + api_keys.append(api_key) + logging.info(f"Created API key (ID: {api_key.id})") + + # Create credentials + credentials = [] + for credential_data in seed_data["credentials"]: + credential = create_credential(session, credential_data) + credentials.append(credential) + logging.info( + f"Created credential for provider: {credential.provider} (ID: {credential.id})" + ) + + # Create assistants + assistants = [] + for assistant_data in seed_data.get("assistants", []): + assistant = create_assistant(session, assistant_data) + assistants.append(assistant) + logging.info(f"Created assistant: {assistant.name} (ID: {assistant.id})") + + documents = [] + for document_data in seed_data.get("documents", []): + document = create_document(session, document_data) + documents.append(document) + logging.info(f"Created document: {document.fname} (ID: {document.id})") + + logging.info("Database seeding completed successfully!") + session.commit() + except Exception as e: + logging.error(f"Error during seeding: {e}") + session.rollback() + raise + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + logging.info("Initializing database session...") + with Session(engine) as session: + try: + seed_database(session) + logging.info("Database seeded successfully!") + except Exception as e: + logging.error(f"Error seeding database: {e}") + session.rollback()