iatoolkit 0.71.4__py3-none-any.whl → 1.4.2__py3-none-any.whl

iatoolkit/services/sql_service.py

@@ -3,12 +3,13 @@
 #
 # IAToolkit is open source software.
 
+from iatoolkit.common.interfaces.database_provider import DatabaseProvider
 from iatoolkit.repositories.database_manager import DatabaseManager
-from iatoolkit.common.util import Utility
 from iatoolkit.services.i18n_service import I18nService
 from iatoolkit.common.exceptions import IAToolkitException
-from sqlalchemy import text
+from iatoolkit.common.util import Utility
 from injector import inject, singleton
+from typing import Callable
 import json
 import logging
 
@@ -27,61 +28,124 @@ class SqlService:
         self.util = util
         self.i18n_service = i18n_service
 
-        # Cache for database connections
-        self._db_connections: dict[str, DatabaseManager] = {}
+        # Cache for database providers. Key is tuple: (company_short_name, db_name)
+        # Value is the abstract interface DatabaseProvider
+        self._db_connections: dict[tuple[str, str], DatabaseProvider] = {}
+
+        # cache for database schemas. Key is tuple: (company_short_name, db_name)
+        self._db_schemas: dict[tuple[str, str], str] = {}
+
+        # Registry of factory functions.
+        # Format: {'connection_type': function(config_dict) -> DatabaseProvider}
+        self._provider_factories: dict[str, Callable[[dict], DatabaseProvider]] = {}
+
+        # Register the default 'direct' strategy (SQLAlchemy)
+        self.register_provider_factory('direct', self._create_direct_connection)
+
+    def register_provider_factory(self, connection_type: str, factory: Callable[[dict], DatabaseProvider]):
+        """
+        Allows plugins (Enterprise) to register new connection types.
+        """
+        self._provider_factories[connection_type] = factory
+
+    def _create_direct_connection(self, config: dict) -> DatabaseProvider:
+        """Default factory for standard SQLAlchemy connections."""
+        uri = config.get('db_uri') or config.get('DATABASE_URI')
+        schema = config.get('schema')
+        if not uri:
+            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
+                                     "Missing db_uri for direct connection")
+        return DatabaseManager(uri, schema=schema, register_pgvector=False)
 
-    def register_database(self, db_name: str, db_uri: str):
+    def register_database(self, company_short_name: str, db_name: str, config: dict):
         """
-        Creates and caches a DatabaseManager instance for a given database name and URI.
-        If a database with the same name is already registered, it does nothing.
+        Creates and caches a DatabaseProvider instance based on the configuration.
         """
-        if db_name in self._db_connections:
+        key = (company_short_name, db_name)
+
+        # Determine connection type (default to 'direct')
+        conn_type = config.get('connection_type', 'direct')
+        logging.info(f"Registering DB '{db_name}' ({conn_type}) for company '{company_short_name}'")
+
+        factory = self._provider_factories.get(conn_type)
+        if not factory:
+            logging.error(f"Unknown connection type '{conn_type}' for DB '{db_name}'. Skipping.")
             return
 
-        logging.debug(f"Registering and creating connection for database: '{db_name}'")
+        try:
+            # Create the provider using the appropriate factory
+            provider_instance = factory(config)
+            self._db_connections[key] = provider_instance
+
+            # save the db_schema
+            self._db_schemas[key] = config.get('schema', 'public')
+        except Exception as e:
+            logging.error(f"Failed to register DB '{db_name}': {e}")
+            # We don't raise here to allow other DBs to load if one fails
 
-        # create the database connection and save it on the cache
-        db_manager = DatabaseManager(db_uri, register_pgvector=False)
-        self._db_connections[db_name] = db_manager
+    def get_db_names(self, company_short_name: str) -> list[str]:
+        """
+        Returns list of logical database names available ONLY for the specified company.
+        """
+        return [db for (co, db) in self._db_connections.keys() if co == company_short_name]
 
-    def get_database_manager(self, db_name: str) -> DatabaseManager:
+    def get_database_provider(self, company_short_name: str, db_name: str) -> DatabaseProvider:
         """
-        Retrieves a registered DatabaseManager instance from the cache.
+        Retrieves a registered DatabaseProvider instance using the composite key.
+        Replaces the old 'get_database_manager'.
         """
+        key = (company_short_name, db_name)
         try:
-            return self._db_connections[db_name]
+            return self._db_connections[key]
         except KeyError:
-            logging.error(f"Attempted to access unregistered database: '{db_name}'")
+            logging.error(
+                f"Attempted to access unregistered database: '{db_name}' for company '{company_short_name}'"
+            )
             raise IAToolkitException(
                 IAToolkitException.ErrorType.DATABASE_ERROR,
-                f"Database '{db_name}' is not registered with the SqlService."
+                f"Database '{db_name}' is not registered for this company."
             )
 
-    def exec_sql(self, database: str, query: str) -> str:
+    def exec_sql(self, company_short_name: str, **kwargs):
         """
-        Executes a raw SQL statement against a registered database and returns the result as a JSON string.
+        Executes a raw SQL statement against a registered database provider.
+        Delegates the actual execution details to the provider implementation.
         """
+        database_name = kwargs.get('database_key')
+        query = kwargs.get('query')
+        format = kwargs.get('format', 'json')
+        commit = kwargs.get('commit')
+
+        if not database_name:
+            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
+                                     'missing database_name in call to exec_sql')
+
         try:
-            # 1. Get the database manager from the cache
-            db_manager = self.get_database_manager(database)
+            # 1. Get the abstract provider (could be Direct or Bridge)
+            provider = self.get_database_provider(company_short_name, database_name)
+            db_schema = self._db_schemas[(company_short_name, database_name)]
 
-            # 2. Execute the SQL statement
-            result = db_manager.get_session().execute(text(query))
-            cols = result.keys()
-            rows_context = [dict(zip(cols, row)) for row in result.fetchall()]
+            # 2. Delegate execution
+            # The provider returns a clean List[Dict] or Dict result
+            result_data = provider.execute_query(query=query, commit=commit)
 
-            # seialize the result
-            sql_result_json = json.dumps(rows_context, default=self.util.serialize)
+            # 3. Handle Formatting (Service layer responsibility)
+            if format == 'dict':
+                return result_data
+
+            # Serialize the result
+            return json.dumps(result_data, default=self.util.serialize)
 
-            return sql_result_json
         except IAToolkitException:
-            # Re-raise exceptions from get_database_manager to preserve the specific error
             raise
         except Exception as e:
-            # Attempt to rollback if a session was active
-            db_manager = self._db_connections.get(database)
-            if db_manager:
-                db_manager.get_session().rollback()
+            # Attempt rollback if supported/needed
+            try:
+                provider = self.get_database_provider(company_short_name, database_name)
+                if provider:
+                    provider.rollback()
+            except Exception:
+                pass
 
             error_message = str(e)
             if 'timed out' in str(e):
@@ -89,4 +153,22 @@ class SqlService:
 
             logging.error(f"Error executing SQL statement: {error_message}")
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
-                                     error_message) from e
+                                     error_message) from e
+
+    def commit(self, company_short_name: str, database_name: str):
+        """
+        Commits the current transaction for a registered database provider.
+        """
+        provider = self.get_database_provider(company_short_name, database_name)
+        try:
+            provider.commit()
+        except Exception as e:
+            # Try rollback
+            try:
+                provider.rollback()
+            except:
+                pass
+            logging.error(f"Error while committing sql: '{str(e)}'")
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.DATABASE_ERROR, str(e)
+            )

iatoolkit/services/tool_service.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from injector import inject
+from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
+from iatoolkit.repositories.profile_repo import ProfileRepo
+from iatoolkit.repositories.models import Company, Tool
+from iatoolkit.common.exceptions import IAToolkitException
+from iatoolkit.services.sql_service import SqlService
+from iatoolkit.services.excel_service import ExcelService
+from iatoolkit.services.mail_service import MailService
+
+
+_SYSTEM_TOOLS = [
+    {
+        "function_name": "iat_generate_excel",
+        "description": "Generador de Excel."
+                       "Genera un archivo Excel (.xlsx) a partir de una lista de diccionarios. "
+                       "Cada diccionario representa una fila del archivo. "
+                       "el archivo se guarda en directorio de descargas."
+                       "retorna diccionario con filename, attachment_token (para enviar archivo por mail)"
+                       "content_type y download_link",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "filename": {
+                    "type": "string",
+                    "description": "Nombre del archivo de salida (ejemplo: 'reporte.xlsx')",
+                    "pattern": "^.+\\.xlsx?$"
+                },
+                "sheet_name": {
+                    "type": "string",
+                    "description": "Nombre de la hoja dentro del Excel",
+                    "minLength": 1
+                },
+                "data": {
+                    "type": "array",
+                    "description": "Lista de diccionarios. Cada diccionario representa una fila.",
+                    "minItems": 1,
+                    "items": {
+                        "type": "object",
+                        "properties": {},
+                        "additionalProperties": {
+                            "anyOf": [
+                                {"type": "string"},
+                                {"type": "number"},
+                                {"type": "boolean"},
+                                {"type": "null"},
+                                {
+                                    "type": "string",
+                                    "format": "date"
+                                }
+                            ]
+                        }
+                    }
+                }
+            },
+            "required": ["filename", "sheet_name", "data"]
+        }
+    },
+    {
+        'function_name': "iat_send_email",
+        'description': "iatoolkit mail system. "
+                       "envia mails cuando un usuario lo solicita.",
+        'parameters': {
+            "type": "object",
+            "properties": {
+                "recipient": {"type": "string", "description": "email del destinatario"},
+                "subject": {"type": "string", "description": "asunto del email"},
+                "body": {"type": "string", "description": "HTML del email"},
+                "attachments": {
+                    "type": "array",
+                    "description": "Lista de archivos adjuntos codificados en base64",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "filename": {
+                                "type": "string",
+                                "description": "Nombre del archivo con su extensión (ej. informe.pdf)"
+                            },
+                            "content": {
+                                "type": "string",
+                                "description": "Contenido del archivo en b64."
+                            },
+                            "attachment_token": {
+                                "type": "string",
+                                "description": "token para descargar el archivo."
+                            }
+                        },
+                        "required": ["filename", "content", "attachment_token"],
+                        "additionalProperties": False
+                    }
+                }
+            },
+            "required": ["recipient", "subject", "body", "attachments"]
+        }
+    },
+    {
+        "function_name": "iat_sql_query",
+        "description": "Servicio SQL de IAToolkit: debes utilizar este servicio para todas las consultas SQL a bases de datos.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "database_key": {
+                    "type": "string",
+                    "description": "IMPORTANT: nombre de la base de datos a consultar."
+                },
+                "query": {
+                    "type": "string",
+                    "description": "string con la consulta en sql"
+                },
+            },
+            "required": ["database_key", "query"]
+        }
+    }
+]
+
+
+class ToolService:
+    @inject
+    def __init__(self,
+                 llm_query_repo: LLMQueryRepo,
+                 profile_repo: ProfileRepo,
+                 sql_service: SqlService,
+                 excel_service: ExcelService,
+                 mail_service: MailService):
+        self.llm_query_repo = llm_query_repo
+        self.profile_repo = profile_repo
+        self.sql_service = sql_service
+        self.excel_service = excel_service
+        self.mail_service = mail_service
+
+        # execution mapper for system tools
+        self.system_handlers = {
+            "iat_generate_excel": self.excel_service.excel_generator,
+            "iat_send_email": self.mail_service.send_mail,
+            "iat_sql_query": self.sql_service.exec_sql
+        }
+
+    def register_system_tools(self):
+        """Creates or updates system functions in the database."""
+        try:
+            # delete all system tools
+            self.llm_query_repo.delete_system_tools()
+
+            # create new system tools
+            for function in _SYSTEM_TOOLS:
+                new_tool = Tool(
+                    company_id=None,
+                    system_function=True,
+                    name=function['function_name'],
+                    description=function['description'],
+                    parameters=function['parameters']
+                )
+                self.llm_query_repo.create_or_update_tool(new_tool)
+
+            self.llm_query_repo.commit()
+        except Exception as e:
+            self.llm_query_repo.rollback()
+            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))
+
+    def sync_company_tools(self, company_short_name: str, tools_config: list):
+        """
+        Synchronizes tools from YAML config to Database (Create/Update/Delete strategy).
+        """
+        if not tools_config:
+            return
+
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
+                                     f'Company {company_short_name} not found')
+
+        try:
+            # 1. Get existing tools map for later cleanup
+            existing_tools = {
+                f.name: f for f in self.llm_query_repo.get_company_tools(company)
+            }
+            defined_tool_names = set()
+
+            # 2. Sync (Create or Update) from Config
+            for tool_data in tools_config:
+                name = tool_data['function_name']
+                defined_tool_names.add(name)
+
+                # Construct the tool object with current config values
+                # We create a new transient object and let the repo merge it
+                tool_obj = Tool(
+                    company_id=company.id,
+                    name=name,
+                    description=tool_data['description'],
+                    parameters=tool_data['params'],
+                    system_function=False
+                )
+
+                # Always call create_or_update. The repo handles checking for existence by name.
+                self.llm_query_repo.create_or_update_tool(tool_obj)
+
+            # 3. Cleanup: Delete tools present in DB but not in Config
+            for name, tool in existing_tools.items():
+                # Ensure we don't delete system functions or active tools accidentally if logic changes,
+                # though get_company_tools filters by company_id so system functions shouldn't be here usually.
+                if not tool.system_function and (name not in defined_tool_names):
+                    self.llm_query_repo.delete_tool(tool)
+
+            self.llm_query_repo.commit()
+
+        except Exception as e:
+            self.llm_query_repo.rollback()
+            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))
+
+
+    def get_tools_for_llm(self, company: Company) -> list[dict]:
+        """
+        Returns the list of tools (System + Company) formatted for the LLM (OpenAI Schema).
+        """
+        tools = []
+
+        # get all the tools for the company and system
+        company_tools = self.llm_query_repo.get_company_tools(company)
+
+        for function in company_tools:
+            # clone for no modify the SQLAlchemy session object
+            params = function.parameters.copy() if function.parameters else {}
+            params["additionalProperties"] = False
+
+            ai_tool = {
+                "type": "function",
+                "name": function.name,
+                "description": function.description,
+                "parameters": params,
+                "strict": True
+            }
+            if function.name == 'iat_sql_query':
+                params['properties']['database_key']['enum'] = self.sql_service.get_db_names(company.short_name)
+
+            tools.append(ai_tool)
+        return tools
+
+    def get_system_handler(self, function_name: str):
+        return self.system_handlers.get(function_name)
+
+    def is_system_tool(self, function_name: str) -> bool:
+        return function_name in self.system_handlers

iatoolkit/services/user_feedback_service.py

@@ -8,7 +8,7 @@ from injector import inject
 from iatoolkit.repositories.profile_repo import ProfileRepo
 from iatoolkit.services.i18n_service import I18nService
 from iatoolkit.infra.google_chat_app import GoogleChatApp
-from iatoolkit.infra.mail_app import MailApp
+from iatoolkit.services.mail_service import MailService
 import logging
 
 
@@ -18,11 +18,11 @@ class UserFeedbackService:
                  profile_repo: ProfileRepo,
                  i18n_service: I18nService,
                  google_chat_app: GoogleChatApp,
-                 mail_app: MailApp):
+                 mail_service: MailService):
         self.profile_repo = profile_repo
         self.i18n_service = i18n_service
         self.google_chat_app = google_chat_app
-        self.mail_app = mail_app
+        self.mail_service = mail_service
 
     def _send_google_chat_notification(self, space_name: str, message_text: str):
         """Envía una notificación de feedback a un espacio de Google Chat."""
@@ -38,13 +38,21 @@ class UserFeedbackService:
         except Exception as e:
             logging.exception(f"error sending notification to Google Chat: {e}")
 
-    def _send_email_notification(self, destination_email: str, company_name: str, message_text: str):
+    def _send_email_notification(self,
+                                 company_short_name: str,
+                                 destination_email: str,
+                                 company_name: str,
+                                 message_text: str):
         """Envía una notificación de feedback por correo electrónico."""
         try:
             subject = f"Nuevo Feedback de {company_name}"
             # Convertir el texto plano a un HTML simple para mantener los saltos de línea
             html_body = message_text.replace('\n', '<br>')
-            self.mail_app.send_email(to=destination_email, subject=subject, body=html_body)
+            self.mail_service.send_mail(
+                company_short_name=company_short_name,
+                to=destination_email,
+                subject=subject,
+                body=html_body)
         except Exception as e:
             logging.exception(f"error sending email de feedback: {e}")
 
@@ -65,7 +73,11 @@ class UserFeedbackService:
         if channel == 'google_chat':
             self._send_google_chat_notification(space_name=destination, message_text=message_text)
         elif channel == 'email':
-            self._send_email_notification(destination_email=destination, company_name=company.short_name, message_text=message_text)
+            self._send_email_notification(
+                company_short_name=company.short_name,
+                destination_email=destination,
+                company_name=company.short_name,
+                message_text=message_text)
         else:
             logging.warning(f"unknown feedback channel: '{channel}' for company {company.short_name}.")