iatoolkit 0.67.0__py3-none-any.whl → 0.69.0__py3-none-any.whl

iatoolkit/__init__.py

@@ -10,10 +10,8 @@ from .iatoolkit import IAToolkit, current_iatoolkit, create_app
 # for registering the client companies
 from .company_registry import register_company
 from .base_company import BaseCompany
-from iatoolkit.repositories.database_manager import DatabaseManager
 
 # --- Services ---
-from iatoolkit.services.configuration_service import ConfigurationService
 from iatoolkit.services.query_service import QueryService
 from iatoolkit.services.sql_service import SqlService
 from iatoolkit.services.document_service import DocumentService
@@ -28,9 +26,7 @@ __all__ = [
     'current_iatoolkit',
     'register_company',
     'BaseCompany',
-    'DatabaseManager',
     'QueryService',
-    'ConfigurationService',
     'SqlService',
     'ExcelService',
     'DocumentService',

iatoolkit/base_company.py

@@ -68,12 +68,6 @@ class BaseCompany(ABC):
             **kwargs
         )
 
-
-    @abstractmethod
-    # get context specific for this company
-    def get_company_context(self, **kwargs) -> str:
-        raise NotImplementedError("La subclase debe implementar el método get_company_context()")
-
     @abstractmethod
     # get context specific for this company
     def get_user_info(self, user_identifier: str) -> dict:

iatoolkit/iatoolkit.py

@@ -19,7 +19,7 @@ from werkzeug.middleware.proxy_fix import ProxyFix
 from injector import Binder, Injector, singleton
 from importlib.metadata import version as _pkg_version, PackageNotFoundError
 
-IATOOLKIT_VERSION = "0.67.0"
+IATOOLKIT_VERSION = "0.69.0"
 
 # global variable for the unique instance of IAToolkit
 _iatoolkit_instance: Optional['IAToolkit'] = None
@@ -293,7 +293,6 @@ class IAToolkit:
         from iatoolkit.repositories.document_repo import DocumentRepo
         from iatoolkit.repositories.profile_repo import ProfileRepo
         from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
-
         from iatoolkit.repositories.vs_repo import VSRepo
         from iatoolkit.repositories.tasks_repo import TaskRepo
 
@@ -318,7 +317,6 @@ class IAToolkit:
         from iatoolkit.services.branding_service import BrandingService
         from iatoolkit.services.i18n_service import I18nService
         from iatoolkit.services.language_service import LanguageService
-        from iatoolkit.services.onboarding_service import OnboardingService
         from iatoolkit.services.configuration_service import ConfigurationService
 
         binder.bind(QueryService, to=QueryService)
@@ -333,7 +331,6 @@ class IAToolkit:
         binder.bind(JWTService, to=JWTService)
         binder.bind(Dispatcher, to=Dispatcher)
         binder.bind(BrandingService, to=BrandingService)
-        binder.bind(OnboardingService, to=OnboardingService)
         binder.bind(I18nService, to=I18nService)
         binder.bind(LanguageService, to=LanguageService)
         binder.bind(ConfigurationService, to=ConfigurationService)
@@ -363,7 +360,7 @@ class IAToolkit:
         # instantiate all the registered companies
         get_company_registry().instantiate_companies(self._injector)
 
-        # use the dispatcher to load the config and prepare the execution
+        # use the dispatcher to load the company config.yaml file and prepare the execution
         dispatcher = self._injector.get(Dispatcher)
         dispatcher.load_company_configs()
 

iatoolkit/repositories/database_manager.py

@@ -70,6 +70,11 @@ class DatabaseManager:
     def remove_session(self):
         self.scoped_session.remove()
 
+    def get_all_table_names(self) -> list[str]:
+        # Returns a list of all table names in the database
+        inspector = inspect(self._engine)
+        return inspector.get_table_names()
+
     def get_table_schema(self,
                          table_name: str,
                          schema_name: str | None = None,

iatoolkit/repositories/models.py

@@ -53,14 +53,12 @@ class Company(Base):
     id = Column(Integer, primary_key=True)
     short_name = Column(String(20), nullable=False, unique=True, index=True)
     name = Column(String(256), nullable=False)
-    default_language = Column(String(5), nullable=False, default='es')
 
     # encrypted api-key
     openai_api_key = Column(String, nullable=True)
     gemini_api_key = Column(String, nullable=True)
     parameters = Column(JSON, nullable=True)
     created_at = Column(DateTime, default=datetime.now)
-    allow_jwt = Column(Boolean, default=True, nullable=True)
 
     documents = relationship("Document",
                              back_populates="company",

iatoolkit/services/branding_service.py

@@ -69,16 +69,15 @@ class BrandingService:
             "send_button_color": "#212529"          # Gris oscuro/casi negro por defecto
         }
 
-    def get_company_branding(self, company: Company | None) -> dict:
+    def get_company_branding(self, company_short_name: str) -> dict:
         """
         Retorna los estilos de branding finales para una compañía,
         fusionando los valores por defecto con los personalizados.
         """
         final_branding_values = self._default_branding.copy()
+        branding_data = self.config_service.get_configuration(company_short_name, 'branding')
+        final_branding_values.update(branding_data)
 
-        if company:
-            branding_data = self.config_service.get_company_content(company.short_name, 'branding')
-            final_branding_values.update(branding_data)
 
         # Función para convertir HEX a RGB
         def hex_to_rgb(hex_color):
@@ -140,8 +139,11 @@ class BrandingService:
             }}
         """
 
+        # get the company name from configuration for the branding render
+        company_name = self.config_service.get_configuration(company_short_name, 'name')
+
         return {
-            "name": company.name if company else "IAToolkit",
+            "name": company_name,
             "primary_text_style": primary_text_style,
             "secondary_text_style": secondary_text_style,
             "tertiary_text_style": tertiary_text_style,

iatoolkit/services/company_context_service.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from iatoolkit.common.util import Utility
+from iatoolkit.services.configuration_service import ConfigurationService
+from iatoolkit.services.sql_service import SqlService
+from iatoolkit.common.exceptions import IAToolkitException
+import logging
+from injector import inject
+import os
+
+
+class CompanyContextService:
+    """
+    Responsible for building the complete context string for a given company
+    to be sent to the Language Model.
+    """
+
+    @inject
+    def __init__(self,
+                 sql_service: SqlService,
+                 utility: Utility,
+                 config_service: ConfigurationService):
+        self.sql_service = sql_service
+        self.utility = utility
+        self.config_service = config_service
+
+    def get_company_context(self, company_short_name: str) -> str:
+        """
+        Builds the full context by aggregating three sources:
+        1. Static context files (Markdown).
+        2. Static schema files (YAML for APIs, etc.).
+        3. Dynamic SQL database schema from the live connection.
+        """
+        context_parts = []
+
+        # 1. Context from Markdown (context/*.md) and yaml (schema/*.yaml) files
+        try:
+            md_context = self._get_static_file_context(company_short_name)
+            if md_context:
+                context_parts.append(md_context)
+        except Exception as e:
+            logging.warning(f"Could not load Markdown context for '{company_short_name}': {e}")
+
+        # 2. Context from company-specific Python logic (SQL schemas)
+        try:
+            sql_context = self._get_sql_schema_context(company_short_name)
+            if sql_context:
+                context_parts.append(sql_context)
+        except Exception as e:
+            logging.warning(f"Could not generate SQL context for '{company_short_name}': {e}")
+
+        # Join all parts with a clear separator
+        return "\n\n---\n\n".join(context_parts)
+
+    def _get_static_file_context(self, company_short_name: str) -> str:
+        # Get context from .md and .yaml schema files.
+        static_context = ''
+
+        # Part 1: Markdown context files
+        context_dir = f'companies/{company_short_name}/context'
+        if os.path.exists(context_dir):
+            context_files = self.utility.get_files_by_extension(context_dir, '.md', return_extension=True)
+            for file in context_files:
+                filepath = os.path.join(context_dir, file)
+                static_context += self.utility.load_markdown_context(filepath)
+
+        # Part 2: YAML schema files
+        schema_dir = f'companies/{company_short_name}/schema'
+        if os.path.exists(schema_dir):
+            schema_files = self.utility.get_files_by_extension(schema_dir, '.yaml', return_extension=True)
+            for file in schema_files:
+                schema_name = file.split('.')[0]  # Use full filename as entity name
+                filepath = os.path.join(schema_dir, file)
+                static_context += self.utility.generate_context_for_schema(schema_name, filepath)
+
+        return static_context
+
+    def _get_sql_schema_context(self, company_short_name: str) -> str:
+        """
+        Generates the SQL schema context by inspecting live database connections
+        based on the flexible company.yaml configuration.
+        It supports including all tables and providing specific overrides for a subset of them.
+        """
+        data_sources_config = self.config_service.get_configuration(company_short_name, 'data_sources')
+        if not data_sources_config or not data_sources_config.get('sql'):
+            return ''
+
+        sql_context = ''
+        for source in data_sources_config.get('sql', []):
+            db_name = source.get('database')
+            if not db_name:
+                continue
+
+            try:
+                db_manager = self.sql_service.get_database_manager(db_name)
+            except IAToolkitException as e:
+                logging.warning(f"Could not get DB manager for '{db_name}': {e}")
+                continue
+
+            db_description = source.get('description', '')
+            sql_context += f"{db_description}\n" if db_description else ""
+
+            # 1. get the list of tables to process.
+            tables_to_process = []
+            if source.get('include_all_tables', False):
+                all_tables = db_manager.get_all_table_names()
+                tables_to_exclude = set(source.get('exclude_tables', []))
+                tables_to_process = [t for t in all_tables if t not in tables_to_exclude]
+            elif 'tables' in source:
+                # if not include_all_tables, use the list of tables explicitly specified in the map.
+                tables_to_process = list(source['tables'].keys())
+
+            # 2. get the global list of columns to exclude.
+            global_exclude_columns = source.get('exclude_columns', [])
+
+            # 3. get the overrides for specific tables.
+            table_overrides = source.get('tables', {})
+
+            # 3. iterate over the tables.
+            for table_name in tables_to_process:
+                try:
+                    # 4. get the table specific configuration.
+                    table_config = table_overrides.get(table_name, {})
+
+                    # 5. define the schema name, using the override if it exists.
+                    schema_name = table_config.get('schema_name', table_name)
+
+                    # 6. define the list of columns to exclude, (local vs. global).
+                    local_exclude_columns = table_config.get('exclude_columns')
+                    final_exclude_columns = local_exclude_columns if local_exclude_columns is not None else global_exclude_columns
+
+                    # 7. get the table schema definition.
+                    table_definition = db_manager.get_table_schema(
+                        table_name=table_name,
+                        schema_name=schema_name,
+                        exclude_columns=final_exclude_columns
+                    )
+                    sql_context += table_definition
+                except (KeyError, RuntimeError) as e:
+                    logging.warning(f"Could not generate schema for table '{table_name}': {e}")
+
+        return sql_context

iatoolkit/services/configuration_service.py

@@ -3,8 +3,6 @@
 # Product: IAToolkit
 
 from pathlib import Path
-from iatoolkit import BaseCompany
-from iatoolkit.repositories.profile_repo import ProfileRepo
 from iatoolkit.repositories.models import Company
 from iatoolkit.common.util import Utility
 from injector import inject
@@ -18,13 +16,11 @@ class ConfigurationService:
 
     @inject
     def __init__(self,
-                 profile_repo: ProfileRepo,
                  utility: Utility):
-        self.profile_repo = profile_repo
         self.utility = utility
         self._loaded_configs = {}   # cache for store loaded configurations
 
-    def get_company_content(self, company_short_name: str, content_key: str) -> dict | list | None:
+    def get_configuration(self, company_short_name: str, content_key: str):
         """
         Public method to provide a specific section of a company's configuration.
         It uses a cache to avoid reading files from disk on every call.
@@ -32,31 +28,29 @@ class ConfigurationService:
         self._ensure_config_loaded(company_short_name)
         return self._loaded_configs[company_short_name].get(content_key)
 
-    def register_company(self, company_short_name: str, company_instance: BaseCompany):
+    def load_configuration(self, company_short_name: str, company_instance):
         """
         Main entry point for configuring a company instance.
         This method is invoked by the dispatcher for each registered company.
         """
         logging.info(f"⚙️  Starting configuration for company '{company_short_name}'...")
 
-        # 1. identify the instance with his name and load info from database
-        company_instance.id = company_short_name
-        company_instance.company = self.profile_repo.get_company_by_short_name(company_short_name)
-
-        # 2. Load the main configuration file and supplementary content files
+        # 1. Load the main configuration file and supplementary content files
         config = self._load_and_merge_configs(company_short_name)
 
-        # 3. Register core company details and get the database object
+        # 2. Register core company details and get the database object
         company_db_object = self._register_core_details(company_instance, config)
 
-        # 4. Register tools (functions)
+        # 3. Register tools (functions)
         self._register_tools(company_instance, config.get('tools', []))
 
-        # 5. Register prompt categories and prompts
+        # 4. Register prompt categories and prompts
         self._register_prompts(company_instance, config)
 
-        # 6. Link the persisted Company object back to the running instance
+        # 5. Link the persisted Company object back to the running instance
+        company_instance.company_short_name = company_short_name
         company_instance.company = company_db_object
+        company_instance.id = company_instance.company.id
 
         logging.info(f"✅ Company '{company_short_name}' configured successfully.")
 
@@ -92,15 +86,15 @@ class ConfigurationService:
 
         return config
 
-    def _register_core_details(self, company_instance: BaseCompany, config: dict) -> Company:
+    def _register_core_details(self, company_instance, config: dict) -> Company:
         """Calls _create_company with data from the merged YAML config."""
         return company_instance._create_company(
-            name=config['name'],
             short_name=config['id'],
+            name=config['name'],
             parameters=config.get('parameters', {})
         )
 
-    def _register_tools(self, company_instance: BaseCompany, tools_config: list):
+    def _register_tools(self, company_instance, tools_config: list):
         """Calls _create_function for each tool defined in the YAML."""
         for tool in tools_config:
             company_instance._create_function(
@@ -109,7 +103,7 @@ class ConfigurationService:
                 params=tool['params']
             )
 
-    def _register_prompts(self, company_instance: BaseCompany, config: dict):
+    def _register_prompts(self, company_instance, config: dict):
         """
         Creates prompt categories first, then creates each prompt and assigns
         it to its respective category.
@@ -129,7 +123,6 @@ class ConfigurationService:
                 continue
 
             category_obj = created_categories[category_name]
-
             company_instance._create_prompt(
                 prompt_name=prompt_data['name'],
                 description=prompt_data['description'],

iatoolkit/services/dispatcher_service.py

@@ -5,6 +5,7 @@
 
 from iatoolkit.common.exceptions import IAToolkitException
 from iatoolkit.services.prompt_manager_service import PromptService
+from iatoolkit.services.sql_service import SqlService
 from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
 from iatoolkit.services.configuration_service import ConfigurationService
 from iatoolkit.repositories.models import Company, Function
@@ -23,12 +24,14 @@ class Dispatcher:
                  prompt_service: PromptService,
                  llmquery_repo: LLMQueryRepo,
                  util: Utility,
+                 sql_service: SqlService,
                  excel_service: ExcelService,
                  mail_service: MailService):
         self.config_service = config_service
         self.prompt_service = prompt_service
         self.llmquery_repo = llmquery_repo
         self.util = util
+        self.sql_service = sql_service
         self.excel_service = excel_service
         self.mail_service = mail_service
         self.system_functions = _FUNCTION_LIST
@@ -64,14 +67,43 @@ class Dispatcher:
         """Loads the configuration of every company"""
         for company_name, company_instance in self.company_instances.items():
             try:
-                # register the company configuration
-                self.config_service.register_company(company_name, company_instance)
+                # read company configuration from company.yaml
+                self.config_service.load_configuration(company_name, company_instance)
+
+                # register the company databases
+                self._register_company_databases(company_name)
+
             except Exception as e:
                 logging.error(f"❌ Failed to register configuration for '{company_name}': {e}")
                 continue
 
         return True
 
+    def _register_company_databases(self, company_name: str):
+        """
+        Reads the data_sources config for a company and registers each
+        database with the central SqlService.
+        """
+        logging.info(f"  -> Registering databases for '{company_name}'...")
+        data_sources_config = self.config_service.get_configuration(company_name, 'data_sources')
+
+        if not data_sources_config or not data_sources_config.get('sql'):
+            logging.info(f"  -> No SQL data sources to register for '{company_name}'.")
+            return
+
+        for db_config in data_sources_config['sql']:
+            db_name = db_config.get('database')
+            db_env_var = db_config.get('connection_string_env')
+
+            # resolve the URI connection string from the environment variable
+            db_uri = os.getenv(db_env_var) if db_env_var else None
+            if not db_uri:
+                logging.warning(
+                    f"-> Skipping database registration for '{company_name}' due to missing 'database' name or connection URI.")
+                return
+
+            self.sql_service.register_database(db_name, db_uri)
+
     def setup_iatoolkit_system(self):
         # create system functions
         for function in self.system_functions:
@@ -123,37 +155,6 @@ class Dispatcher:
             raise IAToolkitException(IAToolkitException.ErrorType.EXTERNAL_SOURCE_ERROR,
                                f"Error en function call '{action}': {str(e)}") from e
 
-    def get_company_context(self, company_name: str, **kwargs) -> str:
-        if company_name not in self.company_instances:
-            raise IAToolkitException(IAToolkitException.ErrorType.EXTERNAL_SOURCE_ERROR,
-                               f"Empresa no configurada: {company_name}")
-
-        company_context = ''
-
-        # read the company context from this list of markdown files,
-        # company brief, credits, operation description, etc.
-        context_dir = os.path.join(os.getcwd(), f'companies/{company_name}/context')
-        context_files = self.util.get_files_by_extension(context_dir, '.md', return_extension=True)
-        for file in context_files:
-            filepath = os.path.join(context_dir, file)
-            company_context += self.util.load_markdown_context(filepath)
-
-        # add the schemas for every table or function call responses
-        schema_dir = os.path.join(os.getcwd(), f'companies/{company_name}/schema')
-        schema_files = self.util.get_files_by_extension(schema_dir, '.yaml', return_extension=True)
-        for file in schema_files:
-            schema_name = file.split('_')[0]
-            filepath = os.path.join(schema_dir, file)
-            company_context += self.util.generate_context_for_schema(schema_name, filepath)
-
-        company_instance = self.company_instances[company_name]
-        try:
-            return company_context + company_instance.get_company_context(**kwargs)
-        except Exception as e:
-            logging.exception(e)
-            raise IAToolkitException(IAToolkitException.ErrorType.EXTERNAL_SOURCE_ERROR,
-                               f"Error getting company context of: {company_name}: {str(e)}") from e
-
     def get_company_services(self, company: Company) -> list[dict]:
         # create the syntax with openai response syntax, for the company function list
         tools = []

iatoolkit/services/language_service.py

@@ -4,6 +4,7 @@ import logging
 from injector import inject, singleton
 from flask import g, request
 from iatoolkit.repositories.profile_repo import ProfileRepo
+from iatoolkit.services.configuration_service import ConfigurationService
 from iatoolkit.common.session_manager import SessionManager
 
 @singleton
@@ -17,7 +18,10 @@ class LanguageService:
     FALLBACK_LANGUAGE = 'es'
 
     @inject
-    def __init__(self, profile_repo: ProfileRepo):
+    def __init__(self,
+                 config_service: ConfigurationService,
+                 profile_repo: ProfileRepo):
+        self.config_service = config_service
         self.profile_repo = profile_repo
 
     def _get_company_short_name(self) -> str | None:
@@ -64,10 +68,10 @@ class LanguageService:
             # Priority 2: Company's default language
             company_short_name = self._get_company_short_name()
             if company_short_name:
-                company = self.profile_repo.get_company_by_short_name(company_short_name)
-                if company and company.default_language:
-                    logging.debug(f"Language determined by company default: {company.default_language}")
-                    g.lang = company.default_language
+                locale = self.config_service.get_configuration(company_short_name, 'locale')
+                if locale:
+                    company_language = locale.split('_')[0]
+                    g.lang = company_language
                     return g.lang
         except Exception as e:
             logging.info(f"Could not determine language, falling back to default. Reason: {e}")

iatoolkit/services/profile_service.py

@@ -95,6 +95,9 @@ class ProfileService:
             user_identifier=user_identifier,
             user_profile=external_user_profile)
 
+        # 3. make sure the flask session is clean
+        SessionManager.clear()
+
     def save_user_profile(self, company: Company, user_identifier: str, user_profile: dict):
         """
         Private helper: Takes a pre-built profile, saves it to Redis, and sets the Flask cookie.

iatoolkit/services/query_service.py

@@ -8,6 +8,7 @@ from iatoolkit.services.profile_service import ProfileService
 from iatoolkit.repositories.document_repo import DocumentRepo
 from iatoolkit.repositories.profile_repo import ProfileRepo
 from iatoolkit.services.document_service import DocumentService
+from iatoolkit.services.company_context_service import CompanyContextService
 from iatoolkit.services.i18n_service import I18nService
 from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
 from iatoolkit.repositories.models import Task
@@ -33,6 +34,7 @@ class QueryService:
     def __init__(self,
                  llm_client: llmClient,
                  profile_service: ProfileService,
+                 company_context_service: CompanyContextService,
                  document_service: DocumentService,
                  document_repo: DocumentRepo,
                  llmquery_repo: LLMQueryRepo,
@@ -44,6 +46,7 @@ class QueryService:
                  session_context: UserSessionContextService
                  ):
         self.profile_service = profile_service
+        self.company_context_service = company_context_service
         self.document_service = document_service
         self.document_repo = document_repo
         self.llmquery_repo = llmquery_repo
@@ -81,7 +84,7 @@ class QueryService:
         )
 
         # get the company context: schemas, database models, .md files
-        company_specific_context = self.dispatcher.get_company_context(company_name=company_short_name)
+        company_specific_context = self.company_context_service.get_company_context(company_short_name)
 
         # merge context: company + user
         final_system_context = f"{company_specific_context}\n{rendered_system_prompt}"

iatoolkit/services/search_service.py

@@ -5,19 +5,22 @@
 
 from iatoolkit.repositories.vs_repo import VSRepo
 from iatoolkit.repositories.document_repo import DocumentRepo
+from iatoolkit.repositories.profile_repo import ProfileRepo
+from iatoolkit.repositories.models import Company
 from injector import inject
 
 
 class SearchService:
     @inject
     def __init__(self,
+                 profile_repo: ProfileRepo,
                  doc_repo: DocumentRepo,
                  vs_repo: VSRepo):
-        super().__init__()
+        self.profile_repo = profile_repo
         self.vs_repo = vs_repo
         self.doc_repo = doc_repo
 
-    def search(self, company_id:  int, query: str, metadata_filter: dict = None) -> str:
+    def search(self, company_short_name: str, query: str, metadata_filter: dict = None) -> str:
         """
         Performs a semantic search for a given query within a company's documents.
 
@@ -26,7 +29,7 @@ class SearchService:
         content of the retrieved documents, which can be used as context for an LLM.
 
         Args:
-            company_id: The ID of the company to search within.
+            company_short_name: The  company to search within.
             query: The text query to search for.
             metadata_filter: An optional dictionary to filter documents by their metadata.
 
@@ -34,7 +37,11 @@ class SearchService:
             A string containing the concatenated content of the found documents,
             formatted to be used as a context.
         """
-        document_list = self.vs_repo.query(company_id=company_id,
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            return f"error: company {company_short_name} not found"
+
+        document_list = self.vs_repo.query(company_id=company.id,
                                            query_text=query,
                                            metadata_filter=metadata_filter)