iatoolkit 2.14.1__tar.gz → 2.14.2__tar.gz

{iatoolkit-2.14.1/src/iatoolkit.egg-info → iatoolkit-2.14.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iatoolkit
-Version: 2.14.1
+Version: 2.14.2
 Summary: IAToolkit
 Author: Fernando Libedinsky
 License-Expression: MIT
@@ -10,8 +10,8 @@ License-File: LICENSE
 License-File: LICENSE_COMMUNITY.md
 Requires-Dist: bcrypt==4.2.1
 Requires-Dist: anthropic>=0.83.0
-Requires-Dist: boto3==1.36.22
-Requires-Dist: botocore==1.36.22
+Requires-Dist: boto3==1.43.10
+Requires-Dist: botocore==1.43.10
 Requires-Dist: braintrust>=0.17.0
 Requires-Dist: build==1.2.2.post1
 Requires-Dist: click==8.1.8
@@ -74,9 +74,11 @@ Requires-Dist: regex==2024.11.6
 Requires-Dist: requests==2.32.3
 Requires-Dist: requests-oauthlib==2.0.0
 Requires-Dist: requests-toolbelt==1.0.0
-Requires-Dist: s3transfer==0.11.2
+Requires-Dist: s3transfer==0.17.0
 Requires-Dist: sib-api-v3-sdk==7.6.0
-Requires-Dist: SQLAlchemy==2.0.36
+Requires-Dist: SQLAlchemy==2.0.49
+Requires-Dist: sqlalchemy-redshift==1.0.0
+Requires-Dist: redshift_connector==2.1.14
 Requires-Dist: tiktoken==0.8.0
 Requires-Dist: tokenizers>=0.22.0
 Requires-Dist: websocket-client==1.8.0

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/requirements.txt

@@ -1,7 +1,7 @@
 bcrypt==4.2.1
 anthropic>=0.83.0
-boto3==1.36.22
-botocore==1.36.22
+boto3==1.43.10
+botocore==1.43.10
 braintrust>=0.17.0
 build==1.2.2.post1
 click==8.1.8
@@ -64,9 +64,11 @@ regex==2024.11.6
 requests==2.32.3
 requests-oauthlib==2.0.0
 requests-toolbelt==1.0.0
-s3transfer==0.11.2
+s3transfer==0.17.0
 sib-api-v3-sdk==7.6.0
-SQLAlchemy==2.0.36
+SQLAlchemy==2.0.49
+sqlalchemy-redshift==1.0.0
+redshift_connector==2.1.14
 tiktoken==0.8.0
 tokenizers>=0.22.0
 websocket-client==1.8.0

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/__init__.py

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

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/common/util.py

@@ -36,11 +36,13 @@ class Utility:
             env = Environment(loader=FileSystemLoader(template_dir))
             template = env.get_template(template_file)
 
-            # add all the keys in client_data to kwargs
-            kwargs.update(client_data)
+            # Explicit kwargs should win over client_data to preserve rich objects
+            # such as `company` passed by the caller.
+            render_context = dict(client_data)
+            render_context.update(kwargs)
 
             # render my dynamic prompt
-            prompt = template.render(**kwargs)
+            prompt = template.render(**render_context)
             return prompt
         except Exception as e:
             logging.exception(e)
@@ -72,9 +74,12 @@ class Utility:
             env = Environment(loader=loader)
             template = env.from_string(template_string)
 
-            kwargs.update(client_data)
+            # Explicit kwargs should win over client_data to preserve rich objects
+            # such as `company` passed by the caller.
+            render_context = dict(client_data)
+            render_context.update(kwargs)
 
-            prompt = template.render(**kwargs)
+            prompt = template.render(**render_context)
             return prompt
         except Exception as e:
             logging.exception(e)

iatoolkit-2.14.2/src/iatoolkit/config/system_prompts/sql_redshift_basics.prompt

@@ -0,0 +1,9 @@
+### SQL Redshift: reglas base
+
+- Si el contexto marca `dialect=redshift`, escribe SQL compatible con Amazon Redshift.
+- Prefiere SQL analitico clasico: `SELECT`, `JOIN`, `CTE`, agregaciones y funciones de ventana.
+- No asumas compatibilidad completa con PostgreSQL aunque la sintaxis se parezca.
+- No uses `DISTINCT ON`, `RETURNING`, `ON CONFLICT`, `ILIKE ANY` ni otras extensiones
+  de PostgreSQL que no sean necesarias para una consulta analitica.
+- Si el schema muestra tablas en un schema especifico, respeta ese schema o los alias
+  provistos por el contexto.

iatoolkit-2.14.2/src/iatoolkit/config/system_prompts/sql_redshift_types.prompt

@@ -0,0 +1,12 @@
+### SQL Redshift: tipos, casting y JSON
+
+- Usa `CAST(... AS INTEGER)`, `CAST(... AS DECIMAL(18, 4))`, `CAST(... AS DATE)` o el tipo
+  explicito que corresponda. Evita casts estilo PostgreSQL como `::NUMERIC`, `::INTEGER`
+  o `::DATE`.
+- No uses tipos ni operadores `jsonb` de PostgreSQL como `->`, `->>`, `jsonb_each`
+  o `jsonb_array_elements`.
+- Si hay columnas semiestructuradas, apóyate solo en las funciones o expresiones que el
+  schema y el contexto hagan evidentes. Si no aparecen detalles suficientes, prioriza
+  columnas relacionales normales antes que inventar acceso a JSON.
+- Para fechas y timestamps, usa funciones compatibles con Redshift y evita asumir helpers
+  exclusivos de PostgreSQL si no son necesarios.

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/config/system_prompts_pack.yaml

@@ -151,6 +151,24 @@ prompts:
         - agent
       all_capabilities:
         - can_query_sql
+  - key: sql_redshift_basics
+    filename: sql_redshift_basics.prompt
+    section: data_access_rules
+    include:
+      execution_modes:
+        - chat
+        - agent
+      all_capabilities:
+        - can_query_sql_redshift
+  - key: sql_redshift_types
+    filename: sql_redshift_types.prompt
+    section: data_access_rules
+    include:
+      execution_modes:
+        - chat
+        - agent
+      all_capabilities:
+        - can_query_sql_redshift
   - key: sql_mysql_json
     filename: sql_mysql_json.prompt
     section: data_access_rules

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/repositories/database_manager.py

@@ -5,6 +5,7 @@
 
 # database_manager.py
 from sqlalchemy import create_engine, event, inspect, text
+from sqlalchemy.dialects import registry
 from sqlalchemy.orm import sessionmaker, scoped_session
 from sqlalchemy.engine.url import make_url
 from iatoolkit.repositories.models import Base, ORM_SCHEMA
@@ -17,12 +18,14 @@ import logging
 class DatabaseManager(DatabaseProvider):
     _POSTGRES_BOOTSTRAP_PATCHES = (
     )
+    _DEFAULT_CONNECT_TIMEOUT = 10
 
     @inject
     def __init__(self,
                  database_url: str,
                  schema: str = 'public',
-                 register_pgvector: bool = True):
+                 register_pgvector: bool = True,
+                 timeout: int | None = None):
         """
         Inicializa el gestor de la base de datos.
         :param database_url: URL de la base de datos.
@@ -31,6 +34,7 @@ class DatabaseManager(DatabaseProvider):
         """
 
         self.schema = schema
+        self.timeout = self._normalize_timeout(timeout)
 
         # FIX HEROKU: replace postgres:// by postgresql:// for compatibility with SQLAlchemy 1.4+
         if database_url and database_url.startswith("postgres://"):
@@ -38,12 +42,15 @@ class DatabaseManager(DatabaseProvider):
 
         self.url = make_url(database_url)
         self.backend = self.url.get_backend_name()
+        if self._is_redshift():
+            self._ensure_redshift_dialect_registered()
+        self.engine_url = self._build_engine_url()
 
         if self.backend == 'sqlite':
-            raw_engine = create_engine(database_url, echo=False)
+            raw_engine = create_engine(self.engine_url, echo=False)
         else:
             raw_engine = create_engine(
-                database_url,
+                self.engine_url,
                 echo=False,
                 pool_size=10,  # per worker
                 max_overflow=20,
@@ -58,6 +65,7 @@ class DatabaseManager(DatabaseProvider):
         self._engine = raw_engine.execution_options(
             schema_translate_map={ORM_SCHEMA: translated_schema}
         )
+        self.engine = self._engine
         self.SessionFactory = sessionmaker(bind=self._engine,
                                            autoflush=False,
                                            autocommit=False,
@@ -76,24 +84,82 @@ class DatabaseManager(DatabaseProvider):
     def _is_postgres(self) -> bool:
         return self.backend in ('postgresql', 'postgres')
 
+    def _is_redshift(self) -> bool:
+        return self.backend == 'redshift'
+
     def _is_mysql(self) -> bool:
         return self.backend == 'mysql'
 
+    @classmethod
+    def _normalize_timeout(cls, timeout: int | str | None) -> int:
+        if timeout is None:
+            return cls._DEFAULT_CONNECT_TIMEOUT
+
+        try:
+            normalized = int(timeout)
+        except (TypeError, ValueError):
+            return cls._DEFAULT_CONNECT_TIMEOUT
+
+        return normalized if normalized > 0 else cls._DEFAULT_CONNECT_TIMEOUT
+
+    def _resolve_timeout(self) -> int:
+        if self.timeout != self._DEFAULT_CONNECT_TIMEOUT:
+            return self.timeout
+
+        for query_key in ("timeout", "connect_timeout"):
+            value = self.url.query.get(query_key)
+            if value is None:
+                continue
+            return self._normalize_timeout(value)
+
+        return self.timeout
+
+    def _build_engine_url(self):
+        normalized_url = self.url
+        if self._is_redshift():
+            normalized_url = self.url.difference_update_query(["timeout", "connect_timeout"])
+
+        return normalized_url.render_as_string(hide_password=False)
+
+    @staticmethod
+    def _ensure_redshift_dialect_registered():
+        try:
+            # Avoid relying solely on SQLAlchemy entrypoint discovery. In long-lived
+            # processes, the plugin registry may not notice a dialect installed later.
+            registry.register(
+                "redshift",
+                "sqlalchemy_redshift.dialect",
+                "RedshiftDialect_psycopg2",
+            )
+            registry.register(
+                "redshift.redshift_connector",
+                "sqlalchemy_redshift.dialect",
+                "RedshiftDialect_redshift_connector",
+            )
+        except Exception as e:
+            logging.debug("Unable to pre-register Redshift SQLAlchemy dialects: %s", e)
+
     def _build_connect_args(self) -> dict:
+        resolved_timeout = self._resolve_timeout()
+
         if self._is_postgres():
             return {
                 "keepalives": 1,
                 "keepalives_idle": 30,
                 "keepalives_interval": 10,
                 "keepalives_count": 5,
-                "connect_timeout": 10,
+                "connect_timeout": resolved_timeout,
+            }
+        if self._is_redshift():
+            return {
+                "timeout": resolved_timeout,
             }
         if self._is_mysql():
             return {
-                "connect_timeout": 10,
+                "connect_timeout": resolved_timeout,
             }
         return {
-            "connect_timeout": 10,
+            "connect_timeout": resolved_timeout,
         }
 
     def _effective_schema(self) -> str | None:
@@ -103,6 +169,10 @@ class DatabaseManager(DatabaseProvider):
         if self._is_postgres():
             return self.schema
 
+        if self._is_redshift():
+            normalized = str(self.schema or "").strip()
+            return normalized or "public"
+
         if self._is_mysql():
             normalized = str(self.schema or "").strip()
             if normalized and normalized.lower() != "public":
@@ -190,6 +260,10 @@ class DatabaseManager(DatabaseProvider):
         session = self.get_session()
         if self._is_postgres() and self.schema:
             session.execute(text(f"SET search_path TO {self.schema}"))
+        elif self._is_redshift():
+            normalized_schema = str(self.schema or "").strip()
+            if normalized_schema and normalized_schema.lower() != "public":
+                session.execute(text(f"SET search_path TO {normalized_schema}, public"))
 
         result = session.execute(text(query))
         if commit:

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/services/context_builder_service.py

@@ -295,6 +295,7 @@ class ContextBuilderService:
         if self.SQL_TOOL_NAME in enabled_tool_names and not sql_sources:
             prompt_capabilities.discard("can_query_sql")
             prompt_capabilities.discard("can_query_sql_postgres")
+            prompt_capabilities.discard("can_query_sql_redshift")
             prompt_capabilities.discard("can_query_sql_mysql")
         prompt_capabilities = self.prompt_service.resolve_system_prompt_capabilities(
             company.short_name,

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/services/prompt_service.py

@@ -729,6 +729,8 @@ class PromptService:
                 only_dialect = next(iter(dialects))
                 if only_dialect in {"postgresql", "postgres"}:
                     capabilities.add("can_query_sql_postgres")
+                elif only_dialect == "redshift":
+                    capabilities.add("can_query_sql_redshift")
                 elif only_dialect == "mysql":
                     capabilities.add("can_query_sql_mysql")
         except Exception as e:
@@ -760,6 +762,7 @@ class PromptService:
             capabilities.update(sql_capabilities)
         else:
             capabilities.discard("can_query_sql_postgres")
+            capabilities.discard("can_query_sql_redshift")
             capabilities.discard("can_query_sql_mysql")
 
         return capabilities

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/services/sql_service.py

@@ -52,10 +52,11 @@ class SqlService:
         """Default factory for standard SQLAlchemy connections."""
         uri = config.get('db_uri') or config.get('DATABASE_URI')
         schema = config.get('schema')
+        timeout = config.get('timeout')
         if not uri:
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
                                      "Missing db_uri for direct connection")
-        return DatabaseManager(uri, schema=schema, register_pgvector=False)
+        return DatabaseManager(uri, schema=schema, register_pgvector=False, timeout=timeout)
 
     def register_database(self, company_short_name: str, db_name: str, config: dict):
         """
@@ -111,12 +112,38 @@ class SqlService:
                 return value.strip().lower()
         return ""
 
+    def _hydrate_database_from_catalog(self, company_short_name: str, db_name: str) -> bool:
+        try:
+            from iatoolkit.core import current_iatoolkit
+            from iatoolkit.services.sql_source_service import SqlSourceService
+
+            injector = current_iatoolkit().get_injector()
+            sql_source_service = injector.get(SqlSourceService)
+            return bool(sql_source_service.ensure_runtime_registration(company_short_name, db_name))
+        except Exception as e:
+            logging.debug(
+                "Unable to hydrate SQL source '%s' for '%s' from catalog: %s",
+                db_name,
+                company_short_name,
+                e,
+            )
+            return False
+
     def get_database_provider(self, company_short_name: str, db_name: str) -> DatabaseProvider:
         """
         Retrieves a registered DatabaseProvider instance using the composite key.
         Replaces the old 'get_database_manager'.
         """
         key = (company_short_name, db_name)
+        provider = self._db_connections.get(key)
+        if provider is not None:
+            return provider
+
+        if self._hydrate_database_from_catalog(company_short_name, db_name):
+            provider = self._db_connections.get(key)
+            if provider is not None:
+                return provider
+
         try:
             return self._db_connections[key]
         except KeyError:

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit/services/sql_source_service.py

@@ -329,6 +329,67 @@ class SqlSourceService:
             "skipped": skipped,
         }
 
+    def _build_runtime_config(self, company_short_name: str, src: SqlSource) -> dict | None:
+        db_name = self._normalize_database(src.database)
+        connection_type = self._normalize_connection_type(src.connection_type)
+        schema_name = self._normalize_schema(src.schema)
+
+        db_config = {
+            "database": db_name,
+            "schema": schema_name,
+            "connection_type": connection_type,
+            "bridge_id": self._normalize_optional_text(src.bridge_id),
+        }
+
+        if connection_type == SqlSource.CONNECTION_DIRECT:
+            secret_ref = normalize_secret_ref(src.connection_string_env)
+            db_uri = resolve_secret(self.secret_provider, company_short_name, secret_ref)
+            if not db_uri:
+                logging.error(
+                    "Skipping SQL source '%s' for '%s': missing secret '%s'.",
+                    db_name,
+                    company_short_name,
+                    secret_ref,
+                )
+                return None
+            db_config["db_uri"] = db_uri
+            return db_config
+
+        if connection_type == SqlSource.CONNECTION_BRIDGE:
+            if not db_config.get("bridge_id"):
+                logging.error(
+                    "Skipping SQL source '%s' for '%s': missing bridge_id.",
+                    db_name,
+                    company_short_name,
+                )
+                return None
+            return db_config
+
+        logging.error(
+            "Skipping SQL source '%s' for '%s': unsupported connection_type '%s'.",
+            db_name,
+            company_short_name,
+            connection_type,
+        )
+        return None
+
+    def ensure_runtime_registration(self, company_short_name: str, db_name: str) -> bool:
+        company = self._get_company(company_short_name)
+        normalized_db_name = self._normalize_database(db_name)
+        if not normalized_db_name:
+            return False
+
+        source = self.sql_source_repo.get_by_database(company.id, normalized_db_name)
+        if not source or not bool(getattr(source, "is_active", False)):
+            return False
+
+        db_config = self._build_runtime_config(company_short_name, source)
+        if not db_config:
+            return False
+
+        self.sql_service.register_database(company_short_name, normalized_db_name, db_config)
+        return normalized_db_name in self.sql_service.get_db_names(company_short_name)
+
     def refresh_runtime(self, company_short_name: str) -> dict:
         company = self._get_company(company_short_name)
         self.sql_service.clear_company_connections(company_short_name)
@@ -339,40 +400,10 @@ class SqlSourceService:
 
         for src in active_sources:
             db_name = self._normalize_database(src.database)
-            connection_type = self._normalize_connection_type(src.connection_type)
-            schema_name = self._normalize_schema(src.schema)
-
-            db_config = {
-                "database": db_name,
-                "schema": schema_name,
-                "connection_type": connection_type,
-                "bridge_id": self._normalize_optional_text(src.bridge_id),
-            }
-
-            if connection_type == SqlSource.CONNECTION_DIRECT:
-                secret_ref = normalize_secret_ref(src.connection_string_env)
-                db_uri = resolve_secret(self.secret_provider, company_short_name, secret_ref)
-                if not db_uri:
-                    skipped += 1
-                    logging.error(
-                        "Skipping SQL source '%s' for '%s': missing secret '%s'.",
-                        db_name,
-                        company_short_name,
-                        secret_ref,
-                    )
-                    continue
-                db_config["db_uri"] = db_uri
-
-            elif connection_type == SqlSource.CONNECTION_BRIDGE:
-                if not db_config.get("bridge_id"):
-                    skipped += 1
-                    logging.error(
-                        "Skipping SQL source '%s' for '%s': missing bridge_id.",
-                        db_name,
-                        company_short_name,
-                    )
-                    continue
-
+            db_config = self._build_runtime_config(company_short_name, src)
+            if not db_config:
+                skipped += 1
+                continue
             self.sql_service.register_database(company_short_name, db_name, db_config)
             registered += 1
 

{iatoolkit-2.14.1 → iatoolkit-2.14.2/src/iatoolkit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iatoolkit
-Version: 2.14.1
+Version: 2.14.2
 Summary: IAToolkit
 Author: Fernando Libedinsky
 License-Expression: MIT
@@ -10,8 +10,8 @@ License-File: LICENSE
 License-File: LICENSE_COMMUNITY.md
 Requires-Dist: bcrypt==4.2.1
 Requires-Dist: anthropic>=0.83.0
-Requires-Dist: boto3==1.36.22
-Requires-Dist: botocore==1.36.22
+Requires-Dist: boto3==1.43.10
+Requires-Dist: botocore==1.43.10
 Requires-Dist: braintrust>=0.17.0
 Requires-Dist: build==1.2.2.post1
 Requires-Dist: click==8.1.8
@@ -74,9 +74,11 @@ Requires-Dist: regex==2024.11.6
 Requires-Dist: requests==2.32.3
 Requires-Dist: requests-oauthlib==2.0.0
 Requires-Dist: requests-toolbelt==1.0.0
-Requires-Dist: s3transfer==0.11.2
+Requires-Dist: s3transfer==0.17.0
 Requires-Dist: sib-api-v3-sdk==7.6.0
-Requires-Dist: SQLAlchemy==2.0.36
+Requires-Dist: SQLAlchemy==2.0.49
+Requires-Dist: sqlalchemy-redshift==1.0.0
+Requires-Dist: redshift_connector==2.1.14
 Requires-Dist: tiktoken==0.8.0
 Requires-Dist: tokenizers>=0.22.0
 Requires-Dist: websocket-client==1.8.0

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit.egg-info/SOURCES.txt

@@ -54,6 +54,8 @@ src/iatoolkit/config/system_prompts/sql_core.prompt
 src/iatoolkit/config/system_prompts/sql_jsonb.prompt
 src/iatoolkit/config/system_prompts/sql_mysql_casting.prompt
 src/iatoolkit/config/system_prompts/sql_mysql_json.prompt
+src/iatoolkit/config/system_prompts/sql_redshift_basics.prompt
+src/iatoolkit/config/system_prompts/sql_redshift_types.prompt
 src/iatoolkit/config/system_prompts/sql_rules.prompt
 src/iatoolkit/config/system_prompts/tool_html_passthrough.prompt
 src/iatoolkit/infra/__init__.py

{iatoolkit-2.14.1 → iatoolkit-2.14.2}/src/iatoolkit.egg-info/requires.txt

@@ -1,7 +1,7 @@
 bcrypt==4.2.1
 anthropic>=0.83.0
-boto3==1.36.22
-botocore==1.36.22
+boto3==1.43.10
+botocore==1.43.10
 braintrust>=0.17.0
 build==1.2.2.post1
 click==8.1.8
@@ -64,9 +64,11 @@ regex==2024.11.6
 requests==2.32.3
 requests-oauthlib==2.0.0
 requests-toolbelt==1.0.0
-s3transfer==0.11.2
+s3transfer==0.17.0
 sib-api-v3-sdk==7.6.0
-SQLAlchemy==2.0.36
+SQLAlchemy==2.0.49
+sqlalchemy-redshift==1.0.0
+redshift_connector==2.1.14
 tiktoken==0.8.0
 tokenizers>=0.22.0
 websocket-client==1.8.0