iatoolkit 0.3.9__py3-none-any.whl → 0.107.4__py3-none-any.whl

iatoolkit/services/company_context_service.py

@@ -0,0 +1,214 @@
+# 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.common.interfaces.asset_storage import AssetRepository, AssetType
+from iatoolkit.services.sql_service import SqlService
+from iatoolkit.common.exceptions import IAToolkitException
+import logging
+from injector import inject
+
+
+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,
+                 asset_repo: AssetRepository):
+        self.sql_service = sql_service
+        self.utility = utility
+        self.config_service = config_service
+        self.asset_repo = asset_repo
+
+    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 files for SQL data sources).
+        """
+        context_parts = []
+
+        # 1. Context from Markdown (context/*.md)  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 SQL databases
+        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}")
+
+        # 3. Context from yaml (schema/*.yaml) files
+        try:
+            yaml_schema_context = self._get_yaml_schema_context(company_short_name)
+            if yaml_schema_context:
+                context_parts.append(yaml_schema_context)
+        except Exception as e:
+            logging.warning(f"Could not load Yaml 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 files using the repository
+        static_context = ''
+
+        try:
+            # 1. List markdown files in the context "folder"
+            # Note: The repo handles where this folder actually is (FS or DB)
+            md_files = self.asset_repo.list_files(company_short_name, AssetType.CONTEXT, extension='.md')
+
+            for filename in md_files:
+                try:
+                    # 2. Read content
+                    content = self.asset_repo.read_text(company_short_name, AssetType.CONTEXT, filename)
+                    static_context += content + "\n"  # Append content
+                except Exception as e:
+                    logging.warning(f"Error reading context file {filename}: {e}")
+
+        except Exception as e:
+            # If listing fails (e.g. folder doesn't exist), just log and return empty
+            logging.warning(f"Error listing context files for {company_short_name}: {e}")
+
+        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_provider = self.sql_service.get_database_provider(company_short_name, db_name)
+            except IAToolkitException as e:
+                logging.warning(f"Could not get DB provider for '{db_name}': {e}")
+                continue
+
+            db_description = source.get('description', '')
+            sql_context = f"***Database (`database_key`)***: {db_name}\n"
+
+            if db_description:
+                sql_context += (
+                    f"**Description:** : {db_description}\n"
+                )
+
+            sql_context += (
+                f"IMPORTANT: To query this database you MUST use the service/tool "
+                f"**iat_sql_query**, with `database_key={db_name}`.\n"
+            )
+
+            sql_context += (
+                f"IMPORTANT: The value of **database_key** is ALWAYS the literal string "
+                f"'{db_name}'. Do not invent or infer alternative names. "
+                f"Use exactly: `database_key='{db_name}'`.\n"
+            )
+
+            # 1. get the list of tables to process.
+            tables_to_process = []
+            if source.get('include_all_tables', False):
+                all_tables = db_provider.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 settings and overrides.
+            global_exclude_columns = source.get('exclude_columns', [])
+            table_prefix = source.get('table_prefix')
+
+            # get database schema definition, for this source.
+            database_schema_name = source.get('schema')
+
+            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 object name, using the override if it exists.
+                    # Priority 1: Explicit override from the 'tables' map.
+                    schema_object_name = table_config.get('schema_name')
+
+                    if not schema_object_name:
+                        # Priority 3: Automatic prefix stripping.
+                        if table_prefix and table_name.startswith(table_prefix):
+                            schema_object_name = table_name[len(table_prefix):]
+                        else:
+                            # Priority 4: Default to the table name itself.
+                            schema_object_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_provider.get_table_schema(
+                        table_name=table_name,
+                        db_schema=db_provider.schema,
+                        schema_object_name=schema_object_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}")
+
+        if sql_context:
+            sql_context = "These are the SQL databases you can query using the **`iat_sql_service`**: \n" + sql_context
+        return sql_context
+
+    def _get_yaml_schema_context(self, company_short_name: str) -> str:
+        # Get context from .yaml schema files using the repository
+        yaml_schema_context = ''
+
+        try:
+            # 1. List yaml files in the schema "folder"
+            schema_files = self.asset_repo.list_files(company_short_name, AssetType.SCHEMA, extension='.yaml')
+
+            for filename in schema_files:
+                try:
+                    # 2. Read content
+                    content = self.asset_repo.read_text(company_short_name, AssetType.SCHEMA, filename)
+
+                    # 3. Parse YAML content into a dict
+                    schema_dict = self.utility.load_yaml_from_string(content)
+
+                    # 4. Generate markdown description from the dict
+                    if schema_dict:
+                        # We use generate_schema_table which accepts a dict directly
+                        yaml_schema_context += self.utility.generate_schema_table(schema_dict)
+
+                except Exception as e:
+                    logging.warning(f"Error processing schema file {filename}: {e}")
+
+        except Exception as e:
+            logging.warning(f"Error listing schema files for {company_short_name}: {e}")
+
+        return yaml_schema_context

iatoolkit/services/configuration_service.py

@@ -0,0 +1,375 @@
+# iatoolkit/services/configuration_service.py
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+
+from iatoolkit.repositories.models import Company
+from iatoolkit.common.interfaces.asset_storage import AssetRepository, AssetType
+from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
+from iatoolkit.repositories.profile_repo import ProfileRepo
+from iatoolkit.common.util import Utility
+from injector import inject
+import logging
+import os
+
+
+class ConfigurationService:
+    """
+    Orchestrates the configuration of a Company by reading its YAML files
+    and using the BaseCompany's protected methods to register settings.
+    """
+
+    @inject
+    def __init__(self,
+                 asset_repo: AssetRepository,
+                 llm_query_repo: LLMQueryRepo,
+                 profile_repo: ProfileRepo,
+                 utility: Utility):
+        self.asset_repo = asset_repo
+        self.llm_query_repo = llm_query_repo
+        self.profile_repo = profile_repo
+        self.utility = utility
+        self._loaded_configs = {}   # cache for store loaded configurations
+
+    def _ensure_config_loaded(self, company_short_name: str):
+        """
+        Checks if the configuration for a company is in the cache.
+        If not, it loads it from files and stores it.
+        """
+        if company_short_name not in self._loaded_configs:
+            self._loaded_configs[company_short_name] = self._load_and_merge_configs(company_short_name)
+
+    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.
+        """
+        self._ensure_config_loaded(company_short_name)
+        return self._loaded_configs[company_short_name].get(content_key)
+
+    def get_llm_configuration(self, company_short_name: str):
+        """
+        Convenience helper to obtain the 'llm' configuration block for a company.
+        Kept separate from get_configuration() to avoid coupling tests that
+        assert the number of calls to get_configuration().
+        """
+        default_llm_model = None
+        available_llm_models = []
+        self._ensure_config_loaded(company_short_name)
+        llm_config = self._loaded_configs[company_short_name].get("llm")
+        if llm_config:
+            default_llm_model = llm_config.get("model")
+            available_llm_models = llm_config.get('available_models') or []
+
+        # fallback: if no explicit list of models is provided, use the default model
+        if not available_llm_models and default_llm_model:
+            available_llm_models = [{
+                "id": default_llm_model,
+                "label": default_llm_model,
+                "description": "Modelo por defecto configurado para esta compañía."
+            }]
+        return default_llm_model, available_llm_models
+
+    def load_configuration(self, company_short_name: str):
+        """
+        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. Load the main configuration file and supplementary content files
+        config = self._load_and_merge_configs(company_short_name)
+        if config:
+            # 2. create/update company in database
+            self._register_company_database(config)
+
+            # 3. Register config databases with db manager
+            # self.register_data_sources(company_short_name, config=config)
+
+            # 4. Register tools
+            self._register_tools(company_short_name, config)
+
+            # 5. Register prompt categories and prompts
+            self._register_prompts(company_short_name, config)
+
+        # Final step: validate the configuration against platform
+        errors = self._validate_configuration(company_short_name, config)
+
+        logging.info(f"✅ Company '{company_short_name}' configured successfully.")
+        return config, errors
+
+    def _load_and_merge_configs(self, company_short_name: str) -> dict:
+        """
+        Loads the main company.yaml and merges data from supplementary files
+        specified in the 'content_files' section using AssetRepository.
+        """
+        main_config_filename = "company.yaml"
+
+        # verify existence of the main configuration file
+        if not self.asset_repo.exists(company_short_name, AssetType.CONFIG, main_config_filename):
+            # raise FileNotFoundError(f"Main configuration file not found: {main_config_filename}")
+            logging.exception(f"Main configuration file not found: {main_config_filename}")
+
+            # return the minimal configuration needed for starting the IAToolkit
+            # this is a for solving a chicken/egg problem when trying to migrate the configuration
+            # from filesystem to database in enterprise installation
+            # see create_assets cli command in enterprise-iatoolkit)
+            return {
+                'id': company_short_name,
+                'name': company_short_name,
+                'llm': {'model': 'gpt-5', 'provider_api_keys': {'openai':''} },
+                }
+
+        # read text and parse
+        yaml_content = self.asset_repo.read_text(company_short_name, AssetType.CONFIG, main_config_filename)
+        config = self.utility.load_yaml_from_string(yaml_content)
+        if not config:
+            return {}
+
+        # Load and merge supplementary content files (e.g., onboarding_cards)
+        for key, filename in config.get('help_files', {}).items():
+            if self.asset_repo.exists(company_short_name, AssetType.CONFIG, filename):
+                supp_content = self.asset_repo.read_text(company_short_name, AssetType.CONFIG, filename)
+                config[key] = self.utility.load_yaml_from_string(supp_content)
+            else:
+                logging.warning(f"⚠️  Warning: Content file not found: {filename}")
+                config[key] = None
+
+        return config
+
+    def _register_company_database(self, config: dict) -> Company:
+        # register the company in the database: create_or_update logic
+        if not config:
+            return None
+
+        # create or update the company in database
+        company_obj = Company(short_name=config.get('id'),
+                              name=config.get('name'),
+                              parameters=config.get('parameters', {}))
+        company = self.profile_repo.create_company(company_obj)
+
+        # save company object with the configuration
+        config['company'] = company
+
+        return company
+
+    def register_data_sources(self,
+                              company_short_name: str,
+                              config: dict = None):
+        """
+        Reads the data_sources config and registers databases with SqlService.
+        Uses Lazy Loading to avoid circular dependency.
+
+        Public method: Can be called externally after initialization (e.g. by Enterprise)
+        to re-register sources once new factories (like 'bridge') are available.
+        """
+
+        # If config is not provided, try to load it from cache
+        if config is None:
+            self._ensure_config_loaded(company_short_name)
+            config = self._loaded_configs.get(company_short_name)
+
+        if not config:
+            return
+
+        from iatoolkit import current_iatoolkit
+        from iatoolkit.services.sql_service import SqlService
+        sql_service = current_iatoolkit().get_injector().get(SqlService)
+
+        data_sources = config.get('data_sources', {})
+        sql_sources = data_sources.get('sql', [])
+
+        if not sql_sources:
+            return
+
+        logging.info(f"🛢️ Registering databases  for '{company_short_name}'...")
+
+        for source in sql_sources:
+            db_name = source.get('database')
+            if not db_name:
+                continue
+
+            # Prepare the config dictionary for the factory
+            db_config = {
+                'schema': source.get('schema', 'public'),
+                'connection_type': source.get('connection_type', 'direct'),
+                # Pass through keys needed for Bridge or other plugins
+                'bridge_id': source.get('bridge_id'),
+                'timeout': source.get('timeout')
+            }
+
+            # Resolve URI if env var is present (Required for 'direct', optional for others)
+            db_env_var = source.get('connection_string_env')
+            if db_env_var:
+                db_uri = os.getenv(db_env_var)
+                if db_uri:
+                    db_config['db_uri'] = db_uri
+
+            # Validation: 'direct' connections MUST have a URI
+            if db_config['connection_type'] == 'direct' and not db_config.get('db_uri'):
+                logging.error(
+                    f"-> Skipping DB '{db_name}' for '{company_short_name}': missing URI in env '{db_env_var}'.")
+                continue
+
+            # Register with the SQL service
+            sql_service.register_database(company_short_name, db_name, db_config)
+
+    def _register_tools(self, company_short_name: str, config: dict):
+        """creates in the database each tool defined in the YAML."""
+        # Lazy import and resolve ToolService locally
+        from iatoolkit import current_iatoolkit
+        from iatoolkit.services.tool_service import ToolService
+        tool_service = current_iatoolkit().get_injector().get(ToolService)
+
+        tools_config = config.get('tools', [])
+        tool_service.sync_company_tools(company_short_name, tools_config)
+
+    def _register_prompts(self, company_short_name: str, config: dict):
+        """
+         Delegates prompt synchronization to PromptService.
+         """
+        # Lazy import to avoid circular dependency
+        from iatoolkit import current_iatoolkit
+        from iatoolkit.services.prompt_service import PromptService
+        prompt_service = current_iatoolkit().get_injector().get(PromptService)
+
+        prompts_config = config.get('prompts', [])
+        categories_config = config.get('prompt_categories', [])
+
+        prompt_service.sync_company_prompts(
+            company_short_name=company_short_name,
+            prompts_config=prompts_config,
+            categories_config=categories_config
+        )
+
+    def _validate_configuration(self, company_short_name: str, config: dict):
+        """
+        Validates the structure and consistency of the company.yaml configuration.
+        It checks for required keys, valid values, and existence of related files.
+        Raises IAToolkitException if any validation error is found.
+        """
+        errors = []
+
+        # Helper to collect errors
+        def add_error(section, message):
+            errors.append(f"[{section}] {message}")
+
+        if not config:
+            add_error("General", "Configuration file missing or with errors, check the application logs.")
+            return errors
+
+        # 1. Top-level keys
+        if not config.get("id"):
+            add_error("General", "Missing required key: 'id'")
+        elif config["id"] != company_short_name:
+            add_error("General",
+                      f"'id' ({config['id']}) does not match the company short name ('{company_short_name}').")
+        if not config.get("name"):
+            add_error("General", "Missing required key: 'name'")
+
+        # 2. LLM section
+        if not isinstance(config.get("llm"), dict):
+            add_error("llm", "Missing or invalid 'llm' section.")
+        else:
+            if not config.get("llm", {}).get("model"):
+                add_error("llm", "Missing required key: 'model'")
+            if not config.get("llm", {}).get("provider_api_keys"):
+                add_error("llm", "Missing required key: 'provider_api_keys'")
+
+        # 3. Embedding Provider
+        if isinstance(config.get("embedding_provider"), dict):
+            if not config.get("embedding_provider", {}).get("provider"):
+                add_error("embedding_provider", "Missing required key: 'provider'")
+            if not config.get("embedding_provider", {}).get("model"):
+                add_error("embedding_provider", "Missing required key: 'model'")
+            if not config.get("embedding_provider", {}).get("api_key_name"):
+                add_error("embedding_provider", "Missing required key: 'api_key_name'")
+
+        # 4. Data Sources
+        for i, source in enumerate(config.get("data_sources", {}).get("sql", [])):
+            if not source.get("database"):
+                add_error(f"data_sources.sql[{i}]", "Missing required key: 'database'")
+
+            connection_type = source.get("connection_type")
+            if connection_type == 'direct' and not source.get("connection_string_env"):
+                add_error(f"data_sources.sql[{i}]", "Missing required key: 'connection_string_env'")
+            elif connection_type == 'bridge' and not source.get("bridge_id"):
+                add_error(f"data_sources.sql[{i}]", "Missing bridge_id'")
+
+        # 5. Tools
+        for i, tool in enumerate(config.get("tools", [])):
+            function_name = tool.get("function_name")
+            if not function_name:
+                add_error(f"tools[{i}]", "Missing required key: 'function_name'")
+
+            # check that function exist in dispatcher
+            if not tool.get("description"):
+                add_error(f"tools[{i}]", "Missing required key: 'description'")
+            if not isinstance(tool.get("params"), dict):
+                add_error(f"tools[{i}]", "'params' key must be a dictionary.")
+
+        # 6. Prompts
+        category_set = set(config.get("prompt_categories", []))
+        for i, prompt in enumerate(config.get("prompts", [])):
+            prompt_name = prompt.get("name")
+            if not prompt_name:
+                add_error(f"prompts[{i}]", "Missing required key: 'name'")
+            else:
+                prompt_filename = f"{prompt_name}.prompt"
+                if not self.asset_repo.exists(company_short_name, AssetType.PROMPT, prompt_filename):
+                    add_error(f"prompts/{prompt_name}:", f"Prompt file not found: {prompt_filename}")
+
+                prompt_description = prompt.get("description")
+                if not prompt_description:
+                    add_error(f"prompts[{i}]", "Missing required key: 'description'")
+
+            prompt_cat = prompt.get("category")
+            if not prompt_cat:
+                add_error(f"prompts[{i}]", "Missing required key: 'category'")
+            elif prompt_cat not in category_set:
+                add_error(f"prompts[{i}]", f"Category '{prompt_cat}' is not defined in 'prompt_categories'.")
+
+        # 7. User Feedback
+        feedback_config = config.get("parameters", {}).get("user_feedback", {})
+        if feedback_config.get("channel") == "email" and not feedback_config.get("destination"):
+            add_error("parameters.user_feedback", "When channel is 'email', a 'destination' is required.")
+
+        # 8. Knowledge Base
+        kb_config = config.get("knowledge_base", {})
+        if kb_config and not isinstance(kb_config, dict):
+            add_error("knowledge_base", "Section must be a dictionary.")
+        elif kb_config:
+            prod_connector = kb_config.get("connectors", {}).get("production", {})
+            if prod_connector.get("type") == "s3":
+                for key in ["bucket", "prefix", "aws_access_key_id_env", "aws_secret_access_key_env", "aws_region_env"]:
+                    if not prod_connector.get(key):
+                        add_error("knowledge_base.connectors.production", f"S3 connector is missing '{key}'.")
+
+        # 9. Mail Provider
+        mail_config = config.get("mail_provider", {})
+        if mail_config:
+            provider = mail_config.get("provider")
+            if not provider:
+                add_error("mail_provider", "Missing required key: 'provider'")
+            elif provider not in ["brevo_mail", "smtplib"]:
+                add_error("mail_provider", f"Unsupported provider: '{provider}'. Must be 'brevo_mail' or 'smtplib'.")
+
+            if not mail_config.get("sender_email"):
+                add_error("mail_provider", "Missing required key: 'sender_email'")
+
+        # 10. Help Files
+        for key, filename in config.get("help_files", {}).items():
+            if not filename:
+                add_error(f"help_files.{key}", "Filename cannot be empty.")
+                continue
+            if not self.asset_repo.exists(company_short_name, AssetType.CONFIG, filename):
+                add_error(f"help_files.{key}", f"Help file not found: {filename}")
+
+
+        # If any errors were found, log all messages and raise an exception
+        if errors:
+            error_summary = f"Configuration file '{company_short_name}/config/company.yaml' for '{company_short_name}' has validation errors:\n" + "\n".join(
+                f" - {e}" for e in errors)
+            logging.error(error_summary)
+
+        return errors
+