iatoolkit 2.10.0__tar.gz → 2.11.0__tar.gz

{iatoolkit-2.10.0/src/iatoolkit.egg-info → iatoolkit-2.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iatoolkit
-Version: 2.10.0
+Version: 2.11.0
 Summary: IAToolkit
 Author: Fernando Libedinsky
 License-Expression: MIT
@@ -57,6 +57,7 @@ Requires-Dist: pgvector==0.3.6
 Requires-Dist: pillow==11.0.0
 Requires-Dist: psutil==7.0.0
 Requires-Dist: psycopg2-binary==2.9.10
+Requires-Dist: PyMySQL==1.1.1
 Requires-Dist: PyJWT==2.10.1
 Requires-Dist: PyMuPDF==1.25.0
 Requires-Dist: python-dotenv==1.0.1

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/requirements.txt

@@ -47,6 +47,7 @@ pgvector==0.3.6
 pillow==11.0.0
 psutil==7.0.0
 psycopg2-binary==2.9.10
+PyMySQL==1.1.1
 PyJWT==2.10.1
 PyMuPDF==1.25.0
 python-dotenv==1.0.1

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/__init__.py

@@ -3,7 +3,7 @@
 #
 # IAToolkit is open source software.
 
-__version__ = "2.10.0"
+__version__ = "2.11.0"
 
 # Expose main classes and functions at the top level of the package
 

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/common/interfaces/database_provider.py

@@ -24,6 +24,14 @@ class DatabaseProvider(abc.ABC):
         """
         pass
 
+    @abc.abstractmethod
+    def get_dialect(self) -> str:
+        """
+        Returns the normalized SQL dialect/backend name used by this provider
+        (for example: 'postgresql', 'mysql', 'sqlite').
+        """
+        pass
+
     # --- Execution Methods ---
     @abc.abstractmethod
     def execute_query(self, query: str, commit: bool = False) -> Union[List[Dict[str, Any]], Dict[str, int]]:
@@ -40,4 +48,4 @@ class DatabaseProvider(abc.ABC):
 
     @abc.abstractmethod
     def rollback(self) -> None:
-        pass
+        pass

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/config/system_prompts/sql_casting.prompt

@@ -1,4 +1,4 @@
-### SQL: casting y tipos
+### SQL PostgreSQL: casting y tipos
 
 - Si un numero puede venir con decimales, usa `NUMERIC`, no `INTEGER`.
 - Para numeros extraidos desde JSONB que pueden venir vacios, usa

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/config/system_prompts/sql_core.prompt

@@ -1,9 +1,12 @@
 ### SQL: reglas base
 
-- Trabaja siempre en PostgreSQL y genera consultas sintacticamente validas.
+- Trabaja siempre con el dialecto SQL indicado para la base seleccionada.
+- Antes de escribir la consulta, revisa el contexto de la base y respeta su marcador
+  `dialect=...` (por ejemplo `postgresql` o `mysql`).
 - Usa solo tablas, columnas y relaciones presentes en el schema provisto.
   No inventes columnas ni supongas nombres parecidos.
 - Si la consulta requiere una base especifica, usa `iat_sql_query` con el
   `database_key` correcto.
 - Si hay mas de una tabla, subconsulta o CTE, califica las columnas con alias
   en `SELECT`, `JOIN`, `ON`, `WHERE`, `GROUP BY`, `HAVING` y `ORDER BY`.
+- No mezcles sintaxis de distintos dialectos en la misma consulta.

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/config/system_prompts/sql_jsonb.prompt

@@ -1,4 +1,4 @@
-### SQL: columnas JSONB
+### SQL PostgreSQL: columnas JSONB
 
 - Nunca trates una clave JSONB como si fuera una columna fisica.
 - Usa `jsonb_col->>'campo'` para texto y `jsonb_col->'campo'` para objetos o arrays.

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/config/system_prompts_pack.yaml

@@ -106,7 +106,7 @@ prompts:
         - chat
         - agent
       all_capabilities:
-        - can_query_sql
+        - can_query_sql_postgres
   - key: sql_casting
     filename: sql_casting.prompt
     section: data_access_rules
@@ -115,7 +115,7 @@ prompts:
         - chat
         - agent
       all_capabilities:
-        - can_query_sql
+        - can_query_sql_postgres
   - key: sql_aggregation_scope
     filename: sql_aggregation_scope.prompt
     section: data_access_rules
@@ -125,3 +125,21 @@ prompts:
         - agent
       all_capabilities:
         - can_query_sql
+  - key: sql_mysql_json
+    filename: sql_mysql_json.prompt
+    section: data_access_rules
+    include:
+      execution_modes:
+        - chat
+        - agent
+      all_capabilities:
+        - can_query_sql_mysql
+  - key: sql_mysql_casting
+    filename: sql_mysql_casting.prompt
+    section: data_access_rules
+    include:
+      execution_modes:
+        - chat
+        - agent
+      all_capabilities:
+        - can_query_sql_mysql

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/repositories/database_manager.py

@@ -37,8 +37,9 @@ class DatabaseManager(DatabaseProvider):
             database_url = database_url.replace("postgres://", "postgresql://", 1)
 
         self.url = make_url(database_url)
+        self.backend = self.url.get_backend_name()
 
-        if database_url.startswith('sqlite'):
+        if self.backend == 'sqlite':
             raw_engine = create_engine(database_url, echo=False)
         else:
             raw_engine = create_engine(
@@ -50,16 +51,10 @@ class DatabaseManager(DatabaseProvider):
                 pool_recycle=1800,
                 pool_pre_ping=True,
                 pool_use_lifo=True,
-                connect_args={
-                    "keepalives": 1,
-                    "keepalives_idle": 30,
-                    "keepalives_interval": 10,
-                    "keepalives_count": 5,
-                    "connect_timeout": 10,
-                },
+                connect_args=self._build_connect_args(),
                 future=True,
             )
-        translated_schema = None if self.url.get_backend_name() == 'sqlite' else self.schema
+        translated_schema = self.schema if self._is_postgres() else None
         self._engine = raw_engine.execution_options(
             schema_translate_map={ORM_SCHEMA: translated_schema}
         )
@@ -70,8 +65,7 @@ class DatabaseManager(DatabaseProvider):
         self.scoped_session = scoped_session(self.SessionFactory)
 
         # Register pgvector for each new connection
-        backend = self.url.get_backend_name()
-        if backend == 'postgresql' or backend == 'postgres':
+        if self._is_postgres():
             if register_pgvector:
                 event.listen(raw_engine, 'connect', self.on_connect)
 
@@ -79,6 +73,45 @@ class DatabaseManager(DatabaseProvider):
             if self.schema:
                 event.listen(raw_engine, 'checkout', self.set_search_path)
 
+    def _is_postgres(self) -> bool:
+        return self.backend in ('postgresql', 'postgres')
+
+    def _is_mysql(self) -> bool:
+        return self.backend == 'mysql'
+
+    def _build_connect_args(self) -> dict:
+        if self._is_postgres():
+            return {
+                "keepalives": 1,
+                "keepalives_idle": 30,
+                "keepalives_interval": 10,
+                "keepalives_count": 5,
+                "connect_timeout": 10,
+            }
+        if self._is_mysql():
+            return {
+                "connect_timeout": 10,
+            }
+        return {
+            "connect_timeout": 10,
+        }
+
+    def _effective_schema(self) -> str | None:
+        if self.backend == 'sqlite':
+            return None
+
+        if self._is_postgres():
+            return self.schema
+
+        if self._is_mysql():
+            normalized = str(self.schema or "").strip()
+            if normalized and normalized.lower() != "public":
+                return normalized
+            return self.url.database
+
+        normalized = str(self.schema or "").strip()
+        return normalized or None
+
     def set_search_path(self, dbapi_connection, connection_record, connection_proxy):
         # Configure the search_path for this connection
         cursor = dbapi_connection.cursor()
@@ -110,10 +143,12 @@ class DatabaseManager(DatabaseProvider):
     def get_connection(self):
         return self._engine.connect()
 
+    def get_dialect(self) -> str:
+        return str(self.backend or "").strip().lower()
+
     def create_all(self):
         # if there is a schema defined, make sure it exists before creating tables
-        backend = self.url.get_backend_name()
-        if self.schema and (backend == 'postgresql' or backend == 'postgres'):
+        if self.schema and self._is_postgres():
             with self._engine.begin() as conn:
                 conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {self.schema}"))
 
@@ -126,8 +161,7 @@ class DatabaseManager(DatabaseProvider):
             )
 
     def _apply_bootstrap_patches(self) -> list[str]:
-        backend = self.url.get_backend_name()
-        if backend not in ('postgresql', 'postgres'):
+        if not self._is_postgres():
             return []
 
         applied = []
@@ -154,7 +188,7 @@ class DatabaseManager(DatabaseProvider):
         Implementation for Direct SQLAlchemy connection.
         """
         session = self.get_session()
-        if self.schema:
+        if self._is_postgres() and self.schema:
             session.execute(text(f"SET search_path TO {self.schema}"))
 
         result = session.execute(text(query))
@@ -178,14 +212,15 @@ class DatabaseManager(DatabaseProvider):
     def get_database_structure(self) -> dict:
         inspector = inspect(self._engine)
         structure = {}
-        for table in inspector.get_table_names(schema=self.schema):
+        effective_schema = self._effective_schema()
+        for table in inspector.get_table_names(schema=effective_schema):
             columns_data = []
 
             # get columns
             try:
-                columns = inspector.get_columns(table, schema=self.schema)
+                columns = inspector.get_columns(table, schema=effective_schema)
                 # Obtener PKs para marcarlas
-                pks = inspector.get_pk_constraint(table, schema=self.schema).get('constrained_columns', [])
+                pks = inspector.get_pk_constraint(table, schema=effective_schema).get('constrained_columns', [])
 
                 for col in columns:
                     columns_data.append({

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/services/company_context_service.py

@@ -128,7 +128,9 @@ class CompanyContextService:
                     continue
 
                 # 2. Build Header for this Database
-                db_context = f"### Base SQL (`database_key`): {db_name}\n"
+                dialect = self.sql_service.get_database_dialect(company_short_name, db_name)
+                dialect_suffix = f" [dialect={dialect}]" if dialect else ""
+                db_context = f"### Base SQL (`database_key`): {db_name}{dialect_suffix}\n"
 
                 # Optional: Add DB description from config if available (useful context)
                 db_desc = source.get('description', '')

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/services/context_builder_service.py

@@ -73,7 +73,10 @@ class ContextBuilderService:
             company_id=company.id,
             company_short_name=company.short_name,
             query_text=query_text,
-            capabilities_override=self._resolve_system_prompt_capabilities(available_tools),
+            capabilities_override=self.prompt_service.resolve_system_prompt_capabilities(
+                company.short_name,
+                self._resolve_system_prompt_capabilities(available_tools),
+            ),
             execution_mode="chat",
             response_mode="chat_compatible",
         )
@@ -179,7 +182,10 @@ class ContextBuilderService:
             company_id=company.id,
             company_short_name=company.short_name,
             query_text=query_text,
-            capabilities_override=self._resolve_system_prompt_capabilities(available_tools),
+            capabilities_override=self.prompt_service.resolve_system_prompt_capabilities(
+                company.short_name,
+                self._resolve_system_prompt_capabilities(available_tools),
+            ),
             execution_mode="chat",
             response_mode="chat_compatible",
         )
@@ -282,15 +288,22 @@ class ContextBuilderService:
             PromptResourceType.RAG_COLLECTION.value,
         )
 
-        capabilities_override = self._resolve_system_prompt_capabilities(enabled_tools)
+        prompt_capabilities = self._resolve_system_prompt_capabilities(enabled_tools)
         if self.SQL_TOOL_NAME in enabled_tool_names and not sql_sources:
-            capabilities_override.discard("can_query_sql")
+            prompt_capabilities.discard("can_query_sql")
+            prompt_capabilities.discard("can_query_sql_postgres")
+            prompt_capabilities.discard("can_query_sql_mysql")
+        prompt_capabilities = self.prompt_service.resolve_system_prompt_capabilities(
+            company.short_name,
+            prompt_capabilities,
+            allowed_sql_databases=sql_sources if self.SQL_TOOL_NAME in enabled_tool_names else None,
+        )
 
         system_prompt_payload = self.prompt_service.get_system_prompt_payload(
             company_id=company.id,
             company_short_name=company.short_name,
             query_text=query_text,
-            capabilities_override=capabilities_override,
+            capabilities_override=prompt_capabilities,
             execution_mode="agent",
             response_mode=str(resolved_contract.get("response_mode") or "chat_compatible").strip().lower() or "chat_compatible",
         )

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/services/prompt_service.py

@@ -633,14 +633,50 @@ class PromptService:
         self.llm_query_repo.delete_prompt(prompt_db)
 
     def _resolve_system_prompt_capabilities(self, company_short_name: str) -> set[str]:
+        return self._resolve_sql_prompt_capabilities(company_short_name)
+
+    def _resolve_sql_prompt_capabilities(
+        self,
+        company_short_name: str,
+        allowed_sql_databases: list[str] | None = None,
+    ) -> set[str]:
         capabilities: set[str] = set()
         if not company_short_name:
             return capabilities
 
         try:
             db_names = self.sql_service.get_db_names(company_short_name)
-            if isinstance(db_names, list) and db_names:
+            if not isinstance(db_names, list):
+                return capabilities
+
+            allowed_database_set = None
+            if allowed_sql_databases is not None:
+                allowed_database_set = {
+                    str(item or "").strip()
+                    for item in allowed_sql_databases
+                    if str(item or "").strip()
+                }
+                db_names = [db_name for db_name in db_names if db_name in allowed_database_set]
+
+            if db_names:
                 capabilities.add("can_query_sql")
+
+            dialects: set[str] = set()
+            for db_name in db_names:
+                try:
+                    dialect = self.sql_service.get_database_dialect(company_short_name, db_name)
+                except Exception:
+                    continue
+                normalized = str(dialect or "").strip().lower()
+                if normalized:
+                    dialects.add(normalized)
+
+            if len(dialects) == 1:
+                only_dialect = next(iter(dialects))
+                if only_dialect in {"postgresql", "postgres"}:
+                    capabilities.add("can_query_sql_postgres")
+                elif only_dialect == "mysql":
+                    capabilities.add("can_query_sql_mysql")
         except Exception as e:
             logging.debug(
                 "Could not resolve SQL capabilities for company '%s': %s",
@@ -650,6 +686,30 @@ class PromptService:
 
         return capabilities
 
+    def resolve_system_prompt_capabilities(
+        self,
+        company_short_name: str,
+        base_capabilities: set[str] | list[str] | tuple[str, ...] | None = None,
+        allowed_sql_databases: list[str] | None = None,
+    ) -> set[str]:
+        capabilities = {
+            str(item or "").strip()
+            for item in (base_capabilities or set())
+            if str(item or "").strip()
+        }
+        sql_capabilities = self._resolve_sql_prompt_capabilities(
+            company_short_name,
+            allowed_sql_databases=allowed_sql_databases,
+        )
+
+        if "can_query_sql" in capabilities:
+            capabilities.update(sql_capabilities)
+        else:
+            capabilities.discard("can_query_sql_postgres")
+            capabilities.discard("can_query_sql_mysql")
+
+        return capabilities
+
     def get_system_prompt_payload(
         self,
         company_id: int,

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit/services/sql_service.py

@@ -102,6 +102,15 @@ class SqlService:
         """
         return [db for (co, db) in self._db_connections.keys() if co == company_short_name]
 
+    def get_database_dialect(self, company_short_name: str, db_name: str) -> str:
+        provider = self.get_database_provider(company_short_name, db_name)
+        dialect = getattr(provider, "get_dialect", None)
+        if callable(dialect):
+            value = dialect()
+            if isinstance(value, str):
+                return value.strip().lower()
+        return ""
+
     def get_database_provider(self, company_short_name: str, db_name: str) -> DatabaseProvider:
         """
         Retrieves a registered DatabaseProvider instance using the composite key.

{iatoolkit-2.10.0 → iatoolkit-2.11.0/src/iatoolkit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iatoolkit
-Version: 2.10.0
+Version: 2.11.0
 Summary: IAToolkit
 Author: Fernando Libedinsky
 License-Expression: MIT
@@ -57,6 +57,7 @@ Requires-Dist: pgvector==0.3.6
 Requires-Dist: pillow==11.0.0
 Requires-Dist: psutil==7.0.0
 Requires-Dist: psycopg2-binary==2.9.10
+Requires-Dist: PyMySQL==1.1.1
 Requires-Dist: PyJWT==2.10.1
 Requires-Dist: PyMuPDF==1.25.0
 Requires-Dist: python-dotenv==1.0.1

{iatoolkit-2.10.0 → iatoolkit-2.11.0}/src/iatoolkit.egg-info/requires.txt

@@ -47,6 +47,7 @@ pgvector==0.3.6
 pillow==11.0.0
 psutil==7.0.0
 psycopg2-binary==2.9.10
+PyMySQL==1.1.1
 PyJWT==2.10.1
 PyMuPDF==1.25.0
 python-dotenv==1.0.1