"""Database seeder — Postgres-native. Post-Part-2 cutover: writes template prompts/tools/agents/sources directly into Postgres via the repository layer. No MongoDB dependencies. The seeder is invoked by the ``python -m application.seed.commands init`` CLI (not at Flask app startup). All template rows are owned by the sentinel user id ``__system__`` — kept in sync with the migration backfill/cleanup-trigger sentinel so template ownership is predictable. """ import logging import os from typing import Dict, List, Optional import yaml from dotenv import load_dotenv from application.agents.tools.tool_manager import ToolManager from application.api.user.tasks import ingest_remote from application.storage.db.repositories.agents import AgentsRepository from application.storage.db.repositories.prompts import PromptsRepository from application.storage.db.repositories.sources import SourcesRepository from application.storage.db.repositories.user_tools import UserToolsRepository from application.storage.db.session import db_readonly, db_session load_dotenv() tool_config = {} tool_manager = ToolManager(config=tool_config) # Sentinel user id for template rows (agents/prompts/sources/tools). # Kept in sync with the Postgres backfill / cleanup-trigger sentinel so # template ownership is predictable across the cutover. SYSTEM_USER_ID = "__system__" class DatabaseSeeder: """Postgres-backed seeder. The constructor accepts an optional positional argument for back compatibility with legacy callers that used to pass a Mongo ``db`` handle. The value is ignored — all persistence goes through the Postgres repositories. """ def __init__(self, db=None): self._legacy_db = db # unused; retained for call-site compatibility self.system_user_id = SYSTEM_USER_ID self.logger = logging.getLogger(__name__) def seed_initial_data(self, config_path: str = None, force=False): """Main entry point for seeding all initial data.""" if not force and self._is_already_seeded(): self.logger.info("Database already seeded. Use force=True to reseed.") return config_path = config_path or os.path.join( os.path.dirname(__file__), "config", "premade_agents.yaml" ) try: with open(config_path, "r") as f: config = yaml.safe_load(f) self._seed_from_config(config) except Exception as e: self.logger.error(f"Failed to load seeding config: {str(e)}") raise def _seed_from_config(self, config: Dict): """Seed all data from configuration.""" self.logger.info("Starting seeding...") if not config.get("agents"): self.logger.warning("No agents found in config") return for agent_config in config["agents"]: try: self.logger.info(f"Processing agent: {agent_config['name']}") # 1. Handle Source source_result = self._handle_source(agent_config) if source_result is False: self.logger.error( f"Skipping agent {agent_config['name']} due to source ingestion failure" ) continue source_id = source_result # 2. Handle Tools tool_ids = self._handle_tools(agent_config) if len(tool_ids) == 0: self.logger.warning( f"No valid tools for agent {agent_config['name']}" ) # 3. Handle Prompt prompt_id = self._handle_prompt(agent_config) # 4. Create or update Agent self._upsert_agent(agent_config, source_id, tool_ids, prompt_id) except Exception as e: self.logger.error( f"Error processing agent {agent_config['name']}: {str(e)}" ) continue self.logger.info("Database seeding completed") @staticmethod def _coerce_uuid_fk(raw) -> Optional[str]: """Coerce sentinel/blank values to ``None`` for nullable UUID FK columns. Mirrors the route-side handling in ``application/api/user/agents/routes.py``: the literal string ``"default"``, empty string, and ``None`` all map to ``None`` so the repository layer skips the column and Postgres keeps the FK NULL (FKs are ``ON DELETE SET NULL``). """ if raw in (None, "", "default"): return None return str(raw) def _upsert_agent( self, agent_config: Dict, source_id: Optional[str], tool_ids: List[str], prompt_id: Optional[str], ) -> None: """Create or update a template agent owned by ``__system__``.""" name = agent_config["name"] prompt_id_val = self._coerce_uuid_fk( prompt_id if prompt_id is not None else agent_config.get("prompt_id") ) folder_id_val = self._coerce_uuid_fk(agent_config.get("folder_id")) workflow_id_val = self._coerce_uuid_fk(agent_config.get("workflow_id")) source_id_val = self._coerce_uuid_fk(source_id) agent_fields = { "description": agent_config["description"], "image": agent_config.get("image", ""), "tools": [str(tid) for tid in tool_ids], "agent_type": agent_config["agent_type"], "prompt_id": prompt_id_val, "chunks": agent_config.get("chunks", "0"), "retriever": agent_config.get("retriever", ""), } if folder_id_val is not None: agent_fields["folder_id"] = folder_id_val if workflow_id_val is not None: agent_fields["workflow_id"] = workflow_id_val if source_id_val is not None: agent_fields["source_id"] = source_id_val with db_session() as conn: repo = AgentsRepository(conn) existing = self._find_system_agent_by_name(repo, name) if existing: self.logger.info(f"Updating existing agent: {name}") repo.update(str(existing["id"]), self.system_user_id, agent_fields) self.logger.info(f"Successfully updated agent: {name} (ID: {existing['id']})") else: self.logger.info(f"Creating new agent: {name}") created = repo.create( user_id=self.system_user_id, name=name, status="template", **agent_fields, ) self.logger.info( f"Successfully created agent: {name} (ID: {created.get('id')})" ) @staticmethod def _find_system_agent_by_name(repo: AgentsRepository, name: str) -> Optional[dict]: """Find a system-owned agent by name among the template rows.""" for row in repo.list_for_user(SYSTEM_USER_ID): if row.get("name") == name: return row return None def _handle_source(self, agent_config: Dict): """Handle source ingestion and return a source id (UUID string) or ``None``/``False``.""" if not agent_config.get("source"): self.logger.info( "No source provided for agent - will create agent without source" ) return None source_config = agent_config["source"] self.logger.info(f"Ingesting source: {source_config['url']}") try: with db_readonly() as conn: existing = self._find_system_source_by_remote_url( SourcesRepository(conn), source_config["url"] ) if existing: self.logger.info(f"Source already exists: {existing['id']}") return existing["id"] # Ingest new source using worker task = ingest_remote.delay( source_data=source_config["url"], job_name=source_config["name"], user=self.system_user_id, loader=source_config.get("loader", "url"), ) result = task.get(timeout=300) if not task.successful(): raise Exception(f"Source ingestion failed: {result}") source_id = None if isinstance(result, dict) and "id" in result: source_id = result["id"] else: raise Exception(f"Source ingestion result missing 'id': {result}") self.logger.info(f"Source ingested successfully: {source_id}") return source_id except Exception as e: self.logger.error(f"Failed to ingest source: {str(e)}") return False @staticmethod def _find_system_source_by_remote_url( repo: SourcesRepository, url: str ) -> Optional[dict]: """Scan system-owned sources for a row whose remote_data matches ``url``.""" # TODO(migration-postgres): push this into SourcesRepository once a # remote_data search helper exists; today we keep the scan here to # stay within this slice's boundaries. try: rows = repo.list_for_user(SYSTEM_USER_ID) # type: ignore[attr-defined] except AttributeError: return None for row in rows: remote = row.get("remote_data") if remote == url: return row if isinstance(remote, dict) and remote.get("url") == url: return row return None def _handle_tools(self, agent_config: Dict) -> List[str]: """Handle tool creation and return list of tool ids (UUID strings).""" tool_ids: List[str] = [] if not agent_config.get("tools"): return tool_ids for tool_config in agent_config["tools"]: try: tool_name = tool_config["name"] processed_config = self._process_config(tool_config.get("config", {})) self.logger.info(f"Processing tool: {tool_name}") with db_session() as conn: repo = UserToolsRepository(conn) existing = self._find_system_tool( repo, tool_name, processed_config ) if existing: self.logger.info(f"Tool already exists: {existing['id']}") tool_ids.append(existing["id"]) continue created = repo.create( user_id=self.system_user_id, name=tool_name, display_name=tool_config.get("display_name", tool_name), description=tool_config.get("description", ""), actions=tool_manager.tools[tool_name].get_actions_metadata(), config=processed_config, status=True, ) tool_ids.append(created["id"]) self.logger.info(f"Created new tool: {created['id']}") except Exception as e: self.logger.error(f"Failed to process tool {tool_name}: {str(e)}") continue return tool_ids @staticmethod def _find_system_tool( repo: UserToolsRepository, name: str, config: dict ) -> Optional[dict]: """Locate a system-owned tool by (name, config) among existing rows.""" existing = repo.find_by_user_and_name(SYSTEM_USER_ID, name) if existing and existing.get("config") == config: return existing return None def _handle_prompt(self, agent_config: Dict) -> Optional[str]: """Handle prompt creation and return prompt id (UUID string).""" if not agent_config.get("prompt"): return None prompt_config = agent_config["prompt"] prompt_name = prompt_config.get("name", f"{agent_config['name']} Prompt") prompt_content = prompt_config.get("content", "") if not prompt_content: self.logger.warning( f"No prompt content provided for agent {agent_config['name']}" ) return None self.logger.info(f"Processing prompt: {prompt_name}") try: with db_session() as conn: repo = PromptsRepository(conn) row = repo.find_or_create( self.system_user_id, prompt_name, prompt_content ) prompt_id = str(row["id"]) self.logger.info(f"Prompt ready: {prompt_id}") return prompt_id except Exception as e: self.logger.error(f"Failed to process prompt {prompt_name}: {str(e)}") return None def _process_config(self, config: Dict) -> Dict: """Process config values to replace environment variables.""" processed = {} for key, value in config.items(): if ( isinstance(value, str) and value.startswith("${") and value.endswith("}") ): env_var = value[2:-1] processed[key] = os.getenv(env_var, "") else: processed[key] = value return processed def _is_already_seeded(self) -> bool: """Check if premade (system-owned) agents already exist in Postgres.""" with db_readonly() as conn: repo = AgentsRepository(conn) return len(repo.list_for_user(SYSTEM_USER_ID)) > 0 @classmethod def initialize_from_env(cls, worker=None): """Factory method to create seeder from environment. Retained for back compatibility with existing call sites. The Postgres connection is resolved lazily via the repository layer (``application.storage.db.engine``), so no explicit wiring is required here. """ return cls()