iatoolkit 0.91.1__py3-none-any.whl → 1.7.0__py3-none-any.whl

iatoolkit/services/query_service.py

@@ -3,7 +3,7 @@
 #
 # IAToolkit is open source software.
 
-from iatoolkit.infra.llm_client import llmClient
+from iatoolkit.services.llm_client_service import llmClient
 from iatoolkit.services.profile_service import ProfileService
 from iatoolkit.repositories.profile_repo import ProfileRepo
 from iatoolkit.services.tool_service import ToolService
@@ -11,11 +11,11 @@ from iatoolkit.services.document_service import DocumentService
 from iatoolkit.services.company_context_service import CompanyContextService
 from iatoolkit.services.i18n_service import I18nService
 from iatoolkit.services.configuration_service import ConfigurationService
-from iatoolkit.repositories.models import Task
 from iatoolkit.services.dispatcher_service import Dispatcher
 from iatoolkit.services.prompt_service import PromptService
 from iatoolkit.services.user_session_context_service import UserSessionContextService
 from iatoolkit.services.history_manager_service import HistoryManagerService
+from iatoolkit.common.model_registry import ModelRegistry
 from iatoolkit.common.util import Utility
 from injector import inject
 import base64
@@ -33,6 +33,7 @@ class HistoryHandle:
     company_short_name: str
     user_identifier: str
     type: str
+    model: str | None = None
     request_params: dict = None
 
 
@@ -52,6 +53,7 @@ class QueryService:
                  configuration_service: ConfigurationService,
                  history_manager: HistoryManagerService,
                  util: Utility,
+                 model_registry: ModelRegistry
                  ):
         self.profile_service = profile_service
         self.company_context_service = company_context_service
@@ -66,6 +68,7 @@ class QueryService:
         self.configuration_service = configuration_service
         self.llm_client = llm_client
         self.history_manager = history_manager
+        self.model_registry = model_registry
 
 
     def _resolve_model(self, company_short_name: str, model: Optional[str]) -> str:
@@ -78,13 +81,16 @@ class QueryService:
         return effective_model
 
     def _get_history_type(self, model: str) -> str:
-        return HistoryManagerService.TYPE_SERVER_SIDE if self.util.is_openai_model(
-            model) else HistoryManagerService.TYPE_CLIENT_SIDE
+        history_type_str = self.model_registry.get_history_type(model)
+        if history_type_str == "server_side":
+            return HistoryManagerService.TYPE_SERVER_SIDE
+        else:
+            return HistoryManagerService.TYPE_CLIENT_SIDE
 
 
     def _build_user_facing_prompt(self, company, user_identifier: str,
                                   client_data: dict, files: list,
-                                  prompt_name: Optional[str], question: str) -> str:
+                                  prompt_name: Optional[str], question: str):
         # get the user profile data from the session context
         user_profile = self.profile_service.get_profile_by_identifier(company.short_name, user_identifier)
 
@@ -124,9 +130,12 @@ class QueryService:
 
         return user_turn_prompt, effective_question
 
-    def _ensure_valid_history(self, company, user_identifier: str,
-                              effective_model: str, user_turn_prompt: str,
-                              ignore_history: bool) -> tuple[Optional[HistoryHandle], Optional[dict]]:
+    def _ensure_valid_history(self, company,
+                              user_identifier: str,
+                              effective_model: str,
+                              user_turn_prompt: str,
+                              ignore_history: bool
+                              ) -> tuple[Optional[HistoryHandle], Optional[dict]]:
         """
             Manages the history strategy and rebuilds context if necessary.
             Returns: (HistoryHandle, error_response)
@@ -137,7 +146,8 @@ class QueryService:
         handle = HistoryHandle(
             company_short_name=company.short_name,
             user_identifier=user_identifier,
-            type=history_type
+            type=history_type,
+            model=effective_model
         )
 
         # pass the handle to populate request_params
@@ -197,22 +207,35 @@ class QueryService:
     def init_context(self, company_short_name: str,
                      user_identifier: str,
                      model: str = None) -> dict:
+        """
+        Forces a context rebuild for a given user and (optionally) model.
+
+        - Clears LLM-related context for the resolved model.
+        - Regenerates the static company/user context.
+        - Sends the context to the LLM for that model.
+        """
 
-        # 1. Execute the forced rebuild sequence using the unified identifier.
-        self.session_context.clear_all_context(company_short_name, user_identifier)
-        logging.info(f"Context for {company_short_name}/{user_identifier} has been cleared.")
+        # 1. Resolve the effective model for this user/company
+        effective_model = self._resolve_model(company_short_name, model)
 
-        # 2. LLM context is clean, now we can load it again
+        # 2. Clear only the LLM-related context for this model
+        self.session_context.clear_all_context(company_short_name, user_identifier,model=effective_model)
+        logging.info(
+            f"Context for {company_short_name}/{user_identifier} "
+            f"(model={effective_model}) has been cleared."
+        )
+
+        # 3. Static LLM context is now clean, we can prepare it again (model-agnostic)
         self.prepare_context(
             company_short_name=company_short_name,
             user_identifier=user_identifier
         )
 
-        # 3. communicate the new context to the LLM
+        # 4. Communicate the new context to the specific LLM model
         response = self.set_context_for_llm(
             company_short_name=company_short_name,
             user_identifier=user_identifier,
-            model=model
+            model=effective_model
         )
 
         return response
@@ -257,8 +280,10 @@ class QueryService:
                             company_short_name: str,
                             user_identifier: str,
                             model: str = ''):
-
-        # This service takes a pre-built context and send to the LLM
+        """
+        Takes a pre-built static context and sends it to the LLM for the given model.
+        Also initializes the model-specific history through HistoryManagerService.
+        """
         company = self.profile_repo.get_company_by_short_name(company_short_name)
         if not company:
             logging.error(f"Company not found: {company_short_name} in set_context_for_llm")
@@ -267,8 +292,8 @@ class QueryService:
         # --- Model Resolution ---
         effective_model = self._resolve_model(company_short_name, model)
 
-        # blocking logic to avoid multiple requests for the same user/company at the same time
-        lock_key = f"lock:context:{company_short_name}/{user_identifier}"
+        # Lock per (company, user, model) to avoid concurrent rebuilds for the same model
+        lock_key = f"lock:context:{company_short_name}/{user_identifier}/{effective_model}"
         if not self.session_context.acquire_lock(lock_key, expire_seconds=60):
             logging.warning(
                 f"try to rebuild context for user {user_identifier} while is still in process, ignored.")
@@ -310,13 +335,13 @@ class QueryService:
     def llm_query(self,
                   company_short_name: str,
                   user_identifier: str,
-                  task: Optional[Task] = None,
+                  model: Optional[str] = None,
                   prompt_name: str = None,
                   question: str = '',
                   client_data: dict = {},
                   ignore_history: bool = False,
-                  files: list = [],
-                  model: Optional[str] = None) -> dict:
+                  files: list = []
+                  ) -> dict:
         try:
             company = self.profile_repo.get_company_by_short_name(short_name=company_short_name)
             if not company:
@@ -378,10 +403,10 @@ class QueryService:
             if not response.get('valid_response'):
                 response['error'] = True
 
-                # save history using the manager passing the handle
-                self.history_manager.update_history(
-                    history_handle, user_turn_prompt, response
-                )
+            # save history using the manager passing the handle
+            self.history_manager.update_history(
+                history_handle, user_turn_prompt, response
+            )
 
             return response
         except Exception as e:

iatoolkit/services/sql_service.py

@@ -3,13 +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 sqlalchemy.exc import SQLAlchemyError
+from iatoolkit.common.util import Utility
 from injector import inject, singleton
+from typing import Callable
 import json
 import logging
 
@@ -28,91 +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_database(self, db_uri: str, db_name: str, schema: str | None = None):
+    def register_provider_factory(self, connection_type: str, factory: Callable[[dict], DatabaseProvider]):
         """
-        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.
+        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, company_short_name: str, db_name: str, config: dict):
         """
-        if db_name in self._db_connections:
+        Creates and caches a DatabaseProvider instance based on the configuration.
+        """
+        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.info(f"Registering and creating connection for database: '{db_name}' (schema: {schema})")
+        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, schema=schema, 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, company_short_name: str,
-                 database: str,
-                 query: str,
-                 format: str = 'json',
-                 commit: bool = False):
+    def exec_sql(self, company_short_name: str, **kwargs):
         """
-        Executes a raw SQL statement against a registered database.
-
-        Args:
-            company_short_name: The company identifier (for logging/context).
-            database: The logical name of the database to query.
-            query: The SQL statement to execute.
-            format: The output format ('json' or 'dict'). Only relevant for SELECT queries.
-            commit: Whether to commit the transaction immediately after execution.
-                    Use True for INSERT/UPDATE/DELETE statements.
-
-        Returns:
-            - A JSON string or list of dicts for SELECT queries.
-            - A dictionary {'rowcount': N} for non-returning statements (INSERT/UPDATE) if not using RETURNING.
+        Executes a raw SQL statement against a registered database provider.
+        Delegates the actual execution details to the provider implementation.
         """
-        try:
-            # 1. Get the database manager from the cache
-            db_manager = self.get_database_manager(database)
-            session = db_manager.get_session()
+        database_name = kwargs.get('database_key')
+        query = kwargs.get('query')
+        format = kwargs.get('format', 'json')
+        commit = kwargs.get('commit')
 
-            # 2. Execute the SQL statement
-            result = session.execute(text(query))
-
-            # 3. Handle Commit
-            if commit:
-                session.commit()
+        if not database_name:
+            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
+                                     'missing database_name in call to exec_sql')
 
-            # 4. Process Results
-            # Check if the query returns rows (e.g., SELECT or INSERT ... RETURNING)
-            if result.returns_rows:
-                cols = result.keys()
-                rows_context = [dict(zip(cols, row)) for row in result.fetchall()]
+        try:
+            # 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)]
 
-                if format == 'dict':
-                    return rows_context
+            # 2. Delegate execution
+            # The provider returns a clean List[Dict] or Dict result
+            result_data = provider.execute_query(query=query, commit=commit)
 
-                # serialize the result
-                return json.dumps(rows_context, default=self.util.serialize)
+            # 3. Handle Formatting (Service layer responsibility)
+            if format == 'dict':
+                return result_data
 
-            # For statements that don't return rows (standard UPDATE/DELETE)
-            return {'rowcount': result.rowcount}
+            # Serialize the result
+            return json.dumps(result_data, default=self.util.serialize)
 
         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):
@@ -122,22 +155,37 @@ class SqlService:
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
                                      error_message) from e
 
-    def commit(self, database: str):
+    def commit(self, company_short_name: str, database_name: str):
         """
-        Commits the current transaction for a registered database.
-        Useful when multiple exec_sql calls are part of a single transaction.
+        Commits the current transaction for a registered database provider.
         """
-
-        # Get the database manager from the cache
-        db_manager = self.get_database_manager(database)
+        provider = self.get_database_provider(company_short_name, database_name)
         try:
-            db_manager.get_session().commit()
-        except SQLAlchemyError as db_error:
-            db_manager.get_session().rollback()
-            logging.error(f"Error de base de datos: {str(db_error)}")
-            raise db_error
+            provider.commit()
         except Exception as e:
-            logging.error(f"error while commiting sql: '{str(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)
+            )
+
+    def get_database_structure(self, company_short_name: str, db_name: str) -> dict:
+        """
+        Introspects the specified database and returns its structure (Tables & Columns).
+        Used for the Schema Editor 2.0
+        """
+        try:
+            provider = self.get_database_provider(company_short_name, db_name)
+            return provider.get_database_structure()
+        except IAToolkitException:
+            raise
+        except Exception as e:
+            logging.error(f"Error introspecting database '{db_name}': {e}")
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.DATABASE_ERROR,
+                f"Failed to introspect database: {str(e)}"
             )

iatoolkit/services/tool_service.py

@@ -5,6 +5,7 @@
 
 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
@@ -98,20 +99,20 @@ _SYSTEM_TOOLS = [
     },
     {
         "function_name": "iat_sql_query",
-        "description": "Servicio SQL de IAToolkit: debes utilizar este servicio para todas las consultas a base de datos.",
+        "description": "Servicio SQL de IAToolkit: debes utilizar este servicio para todas las consultas SQL a bases de datos.",
         "parameters": {
             "type": "object",
             "properties": {
-                "database": {
+                "database_key": {
                     "type": "string",
-                    "description": "nombre de la base de datos a consultar: `database_name`"
+                    "description": "IMPORTANT: nombre de la base de datos a consultar."
                 },
                 "query": {
                     "type": "string",
                     "description": "string con la consulta en sql"
                 },
             },
-            "required": ["database", "query"]
+            "required": ["database_key", "query"]
         }
     }
 ]
@@ -121,10 +122,12 @@ 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
@@ -158,14 +161,22 @@ class ToolService:
             self.llm_query_repo.rollback()
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))
 
-    def sync_company_tools(self, company_instance, tools_config: list):
+    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_instance.company)
+                f.name: f for f in self.llm_query_repo.get_company_tools(company)
             }
             defined_tool_names = set()
 
@@ -177,7 +188,7 @@ class ToolService:
                 # 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_instance.company.id,
+                    company_id=company.id,
                     name=name,
                     description=tool_data['description'],
                     parameters=tool_data['params'],
@@ -206,11 +217,12 @@ class ToolService:
         Returns the list of tools (System + Company) formatted for the LLM (OpenAI Schema).
         """
         tools = []
-        # Obtiene tanto las de la empresa como las del sistema (la query del repo debería soportar esto con OR)
-        functions = self.llm_query_repo.get_company_tools(company)
 
-        for function in functions:
-            # Clonamos para no modificar el objeto de la sesión SQLAlchemy
+        # 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
 
@@ -221,6 +233,9 @@ class ToolService:
                 "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