{
"event": "PreToolUse",
"tool_name": "Edit",
"tool_input": {
"file_path": "\/opt\/scripts\/pipeline\/model_registry.py",
"old_string": " def _get_connection(self):\n \"\"\"Get or create database connection.\"\"\"\n if self._conn is None or not self._conn.open:\n self._conn = mariadb.connect(\n host=\"localhost\",\n user=\"root\",\n password=get_db_password(),\n database=\"ki_dev\",\n autocommit=True,\n )\n return self._conn",
"new_string": " def _get_connection(self):\n \"\"\"Get or create database connection.\"\"\"\n if self._conn is None or not self._conn.is_connected():\n self._conn = mysql.connector.connect(\n host=\"localhost\",\n user=\"root\",\n password=get_db_password(),\n database=\"ki_dev\",\n autocommit=True,\n )\n return self._conn"
}
}
{
"tool_response": {
"filePath": "\/opt\/scripts\/pipeline\/model_registry.py",
"oldString": " def _get_connection(self):\n \"\"\"Get or create database connection.\"\"\"\n if self._conn is None or not self._conn.open:\n self._conn = mariadb.connect(\n host=\"localhost\",\n user=\"root\",\n password=get_db_password(),\n database=\"ki_dev\",\n autocommit=True,\n )\n return self._conn",
"newString": " def _get_connection(self):\n \"\"\"Get or create database connection.\"\"\"\n if self._conn is None or not self._conn.is_connected():\n self._conn = mysql.connector.connect(\n host=\"localhost\",\n user=\"root\",\n password=get_db_password(),\n database=\"ki_dev\",\n autocommit=True,\n )\n return self._conn",
"originalFile": "\"\"\"\nModel Registry - Single Source of Truth for AI Models.\n\nReads available models from ki_dev.ai_models database table.\nSynchronizes with PHP ModelRegistry for consistent model availability.\n\"\"\"\n\nimport mysql.connector\nfrom typing import Optional\nfrom config import get_db_password\n\n\nclass ModelRegistry:\n \"\"\"Central registry for all AI models.\"\"\"\n\n _cache: Optional[list] = None\n _instance: Optional[\"ModelRegistry\"] = None\n\n def __init__(self):\n self._conn = None\n\n @classmethod\n def get_instance(cls) -> \"ModelRegistry\":\n \"\"\"Get singleton instance.\"\"\"\n if cls._instance is None:\n cls._instance = cls()\n return cls._instance\n\n def _get_connection(self):\n \"\"\"Get or create database connection.\"\"\"\n if self._conn is None or not self._conn.open:\n self._conn = mariadb.connect(\n host=\"localhost\",\n user=\"root\",\n password=get_db_password(),\n database=\"ki_dev\",\n autocommit=True,\n )\n return self._conn\n\n @classmethod\n def clear_cache(cls):\n \"\"\"Clear cached models.\"\"\"\n cls._cache = None\n\n def get_chat_models(self) -> dict:\n \"\"\"Get all available chat models.\n\n Returns:\n dict: {full_key: display_name}\n \"\"\"\n return self._get_models(is_chat=True)\n\n def get_vision_models(self) -> dict:\n \"\"\"Get all vision-capable models.\n\n Returns:\n dict: {full_key: display_name}\n \"\"\"\n return self._get_models(is_vision=True)\n\n def get_embedding_models(self) -> dict:\n \"\"\"Get all embedding models.\n\n Returns:\n dict: {full_key: display_name}\n \"\"\"\n return self._get_models(is_embedding=True)\n\n def _get_models(\n self,\n is_chat: Optional[bool] = None,\n is_vision: Optional[bool] = None,\n is_embedding: Optional[bool] = None,\n provider: Optional[str] = None,\n ) -> dict:\n \"\"\"Get models with optional filters.\"\"\"\n all_models = self._load_all_models()\n result = {}\n\n for model in all_models:\n if not model[\"is_available\"]:\n continue\n if is_chat is not None and bool(model[\"is_chat\"]) != is_chat:\n continue\n if is_vision is not None and bool(model[\"is_vision\"]) != is_vision:\n continue\n if is_embedding is not None and bool(model[\"is_embedding\"]) != is_embedding:\n continue\n if provider is not None and model[\"provider\"] != provider:\n continue\n\n result[model[\"full_key\"]] = model[\"display_name\"]\n\n return result\n\n def get_model(self, full_key: str) -> Optional[dict]:\n \"\"\"Get a single model by full_key.\"\"\"\n all_models = self._load_all_models()\n\n for model in all_models:\n if model[\"full_key\"] == full_key:\n return model\n\n return None\n\n def get_label(self, full_key: str) -> str:\n \"\"\"Get display label for a model.\"\"\"\n model = self.get_model(full_key)\n return model[\"display_name\"] if model else full_key\n\n def is_valid(self, full_key: str) -> bool:\n \"\"\"Check if model exists and is available.\"\"\"\n model = self.get_model(full_key)\n return model is not None and model[\"is_available\"]\n\n def is_local(self, full_key: str) -> bool:\n \"\"\"Check if model is local (Ollama).\"\"\"\n return full_key.startswith(\"ollama:\")\n\n def get_default_chat_model(self) -> str:\n \"\"\"Get default chat model (first available by priority).\"\"\"\n chat_models = self.get_chat_models()\n if chat_models:\n return next(iter(chat_models.keys()))\n return \"ollama:mistral:latest\"\n\n def get_default_embedding_model(self) -> str:\n \"\"\"Get default embedding model.\"\"\"\n embed_models = self.get_embedding_models()\n if embed_models:\n return next(iter(embed_models.keys()))\n return \"ollama:mxbai-embed-large:latest\"\n\n def get_ollama_model_id(self, full_key: str) -> str:\n \"\"\"Extract Ollama model ID from full_key.\n\n Example: 'ollama:gemma3:27b-it-qat' -> 'gemma3:27b-it-qat'\n \"\"\"\n if full_key.startswith(\"ollama:\"):\n return full_key[7:] # Remove 'ollama:' prefix\n return full_key\n\n def _load_all_models(self) -> list:\n \"\"\"Load all models from database (with caching).\"\"\"\n if ModelRegistry._cache is not None:\n return ModelRegistry._cache\n\n conn = self._get_connection()\n cursor = conn.cursor(dictionary=True)\n\n cursor.execute(\n \"\"\"\n SELECT id, provider, model_id, display_name, full_key,\n is_available, is_chat, is_embedding, is_vision,\n context_length, parameters, priority\n FROM ai_models\n WHERE is_available = 1\n ORDER BY priority ASC\n \"\"\"\n )\n\n ModelRegistry._cache = cursor.fetchall()\n cursor.close()\n return ModelRegistry._cache\n\n\n# Convenience functions for backward compatibility\ndef get_chat_model() -> str:\n \"\"\"Get the default chat model from registry.\"\"\"\n registry = ModelRegistry.get_instance()\n # Prefer local Ollama models for pipeline\n ollama_models = registry._get_models(is_chat=True, provider=\"ollama\")\n if ollama_models:\n return next(iter(ollama_models.keys()))\n return registry.get_default_chat_model()\n\n\ndef get_embedding_model() -> str:\n \"\"\"Get the default embedding model from registry.\"\"\"\n return ModelRegistry.get_instance().get_default_embedding_model()\n\n\ndef get_model_id_for_ollama(full_key: str) -> str:\n \"\"\"Extract Ollama model ID from full_key.\"\"\"\n return ModelRegistry.get_instance().get_ollama_model_id(full_key)\n",
"structuredPatch": [
{
"oldStart": 28,
"oldLines": 8,
"newStart": 28,
"newLines": 8,
"lines": [
" ",
" def _get_connection(self):",
" \"\"\"Get or create database connection.\"\"\"",
"- if self._conn is None or not self._conn.open:",
"- self._conn = mariadb.connect(",
"+ if self._conn is None or not self._conn.is_connected():",
"+ self._conn = mysql.connector.connect(",
" host=\"localhost\",",
" user=\"root\",",
" password=get_db_password(),"
]
}
],
"userModified": false,
"replaceAll": false
}
}