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

iatoolkit/services/configuration_service.py

@@ -38,41 +38,11 @@ class ConfigurationService:
         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.
+        And for the configurator, for editing the configuration of a company.
         """
         logging.info(f"⚙️  Starting configuration for company '{company_short_name}'...")
 
@@ -97,6 +67,14 @@ class ConfigurationService:
         logging.info(f"✅ Company '{company_short_name}' configured successfully.")
         return config, errors
 
+    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 update_configuration_key(self, company_short_name: str, key: str, value) -> tuple[dict, list[str]]:
         """
         Updates a specific key in the company's configuration file, validates the result,
@@ -186,7 +164,6 @@ class ConfigurationService:
 
         return config, []
 
-
     def validate_configuration(self, company_short_name: str) -> list[str]:
         """
         Public method to trigger validation of the current configuration.
@@ -194,83 +171,6 @@ class ConfigurationService:
         config = self._load_and_merge_configs(company_short_name)
         return self._validate_configuration(company_short_name, config)
 
-    def _set_nested_value(self, data: dict, key: str, value):
-        """
-        Helper to set a value in a nested dictionary or list using dot notation (e.g. 'llm.model', 'tools.0.name').
-        Handles traversal through both dictionaries and lists.
-        """
-        keys = key.split('.')
-        current = data
-
-        # Traverse up to the parent of the target key
-        for i, k in enumerate(keys[:-1]):
-            if isinstance(current, dict):
-                # If it's a dict, we can traverse or create the path
-                current = current.setdefault(k, {})
-            elif isinstance(current, list):
-                # If it's a list, we MUST use an integer index
-                try:
-                    idx = int(k)
-                    current = current[idx]
-                except (ValueError, IndexError) as e:
-                    raise ValueError(
-                        f"Invalid path: cannot access index '{k}' in list at '{'.'.join(keys[:i + 1])}'") from e
-            else:
-                raise ValueError(
-                    f"Invalid path: '{k}' is not a container (got {type(current)}) at '{'.'.join(keys[:i + 1])}'")
-
-        # Set the final value
-        last_key = keys[-1]
-        if isinstance(current, dict):
-            current[last_key] = value
-        elif isinstance(current, list):
-            try:
-                idx = int(last_key)
-                current[idx] = value
-            except (ValueError, IndexError) as e:
-                raise ValueError(f"Invalid path: cannot assign to index '{last_key}' in list") from e
-        else:
-            raise ValueError(f"Cannot assign value to non-container type {type(current)} at '{key}'")
-
-    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:
@@ -374,12 +274,11 @@ class ConfigurationService:
         from iatoolkit.services.prompt_service import PromptService
         prompt_service = current_iatoolkit().get_injector().get(PromptService)
 
-        prompts_config = config.get('prompts', {})
-
+        prompt_list, categories_config = self._get_prompt_config(config)
         prompt_service.sync_company_prompts(
             company_short_name=company_short_name,
-            prompts_config=prompts_config.get('prompt_list', []),
-            categories_config=prompts_config.get('prompt_categories', []),
+            prompt_list=prompt_list,
+            categories_config=categories_config,
         )
 
     def _register_knowledge_base(self, company_short_name: str, config: dict):
@@ -394,7 +293,6 @@ class ConfigurationService:
         # sync collection types in database
         knowledge_base.sync_collection_types(company_short_name, categories_config)
 
-
     def _validate_configuration(self, company_short_name: str, config: dict):
         """
         Validates the structure and consistency of the company.yaml configuration.
@@ -462,10 +360,9 @@ class ConfigurationService:
                 add_error(f"tools[{i}]", "'params' key must be a dictionary.")
 
         # 6. Prompts
-        prompt_list = config.get("prompts", {}).get("prompt_list", [])
-        prompt_categories = config.get("prompts", {}).get("prompt_categories", [])
+        prompt_list, categories_config = self._get_prompt_config(config)
 
-        category_set = set(prompt_categories)
+        category_set = set(categories_config)
         for i, prompt in enumerate(prompt_list):
             prompt_name = prompt.get("name")
             if not prompt_name:
@@ -480,10 +377,12 @@ class ConfigurationService:
                     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'.")
+            prompt_type = prompt.get("prompt_type", 'company').lower()
+            if prompt_type == 'company':
+                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", {})
@@ -530,3 +429,123 @@ class ConfigurationService:
 
         return errors
 
+
+    def _set_nested_value(self, data: dict, key: str, value):
+        """
+        Helper to set a value in a nested dictionary or list using dot notation (e.g. 'llm.model', 'tools.0.name').
+        Handles traversal through both dictionaries and lists.
+        """
+        keys = key.split('.')
+        current = data
+
+        # Traverse up to the parent of the target key
+        for i, k in enumerate(keys[:-1]):
+            if isinstance(current, dict):
+                # If it's a dict, we can traverse or create the path
+                current = current.setdefault(k, {})
+            elif isinstance(current, list):
+                # If it's a list, we MUST use an integer index
+                try:
+                    idx = int(k)
+                    # Allow accessing existing index
+                    current = current[idx]
+                except (ValueError, IndexError) as e:
+                    raise ValueError(
+                        f"Invalid path: cannot access index '{k}' in list at '{'.'.join(keys[:i + 1])}'") from e
+            else:
+                raise ValueError(
+                    f"Invalid path: '{k}' is not a container (got {type(current)}) at '{'.'.join(keys[:i + 1])}'")
+
+        # Set the final value
+        last_key = keys[-1]
+        if isinstance(current, dict):
+            current[last_key] = value
+        elif isinstance(current, list):
+            try:
+                idx = int(last_key)
+                # If index equals length, it means append
+                if idx == len(current):
+                    current.append(value)
+                elif 0 <= idx < len(current):
+                    current[idx] = value
+                else:
+                    raise IndexError(f"Index {idx} out of range for list of size {len(current)}")
+            except (ValueError, IndexError) as e:
+                raise ValueError(f"Invalid path: cannot assign to index '{last_key}' in list") from e
+        else:
+            raise ValueError(f"Cannot assign value to non-container type {type(current)} at '{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_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 _get_prompt_config(self, config):
+        prompts_config = config.get('prompts', {})
+        if isinstance(prompts_config, dict):
+            prompt_list = prompts_config.get('prompt_list', [])
+            categories_config = prompts_config.get('prompt_categories', [])
+        else:
+            prompt_list = config.get('prompts', [])
+            categories_config = config.get('prompt_categories', [])
+
+        return prompt_list, categories_config
+

iatoolkit/services/dispatcher_service.py

@@ -87,9 +87,6 @@ class Dispatcher:
             # system tools registration
             self.tool_service.register_system_tools()
 
-            # system prompts registration
-            self.prompt_service.register_system_prompts()
-
         except Exception as e:
             self.llmquery_repo.rollback()
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))
@@ -109,7 +106,7 @@ class Dispatcher:
         if self.tool_service.is_system_tool(function_name):
             # this is the system function to be executed.
             handler = self.tool_service.get_system_handler(function_name)
-            logging.info(
+            logging.debug(
                 f"Calling system handler [{function_name}] "
                 f"with company_short_name={company_short_name} "
                 f"and kwargs={kwargs}"

iatoolkit/services/knowledge_base_service.py

@@ -382,27 +382,49 @@ class KnowledgeBaseService:
     def sync_collection_types(self, company_short_name: str, categories_config: list):
         """
         This should be called during company initialization or configuration reload.
+        Syncs DB collection types with the provided list.
+        Also updates the configuration YAML.
         """
         company = self.profile_service.get_company_by_short_name(company_short_name)
         if not company:
             raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                            f'Company {company_short_name} not found')
-
+                                     f'Company {company_short_name} not found')
 
         session = self.document_repo.session
+
+        # 1. Get existing types
         existing_types = session.query(CollectionType).filter_by(company_id=company.id).all()
         existing_names = {ct.name: ct for ct in existing_types}
 
+        # 2. Add new types
+        current_config_names = set()
         for cat_name in categories_config:
+            current_config_names.add(cat_name)
             if cat_name not in existing_names:
                 new_type = CollectionType(company_id=company.id, name=cat_name)
                 session.add(new_type)
 
-        # Opcional: Eliminar los que ya no están en el config?
-        # Por seguridad de datos, mejor no borrar automáticamente, o marcarlos inactivos.
+        # 3. Delete types not in config
+        # Note: This might cascade delete documents depending on FK setup.
+        # Assuming safe deletion is desired here to match "Sync" behavior.
+        for existing_ct in existing_types:
+            if existing_ct.name not in current_config_names:
+                session.delete(existing_ct)
 
         session.commit()
 
+        # 4. Update Configuration YAML
+        # Lazy import to avoid circular dependency
+        from iatoolkit import current_iatoolkit
+        from iatoolkit.services.configuration_service import ConfigurationService
+        config_service = current_iatoolkit().get_injector().get(ConfigurationService)
+
+        config_service.update_configuration_key(
+            company_short_name,
+            "knowledge_base.collections",
+            categories_config
+        )
+
     def get_collection_names(self, company_short_name: str) -> List[str]:
         """
         Retrieves the names of all collections defined for a specific company.

iatoolkit/services/llm_client_service.py

@@ -21,6 +21,7 @@ import re
 import tiktoken
 from typing import Dict, Optional, List
 from iatoolkit.services.dispatcher_service import Dispatcher
+from iatoolkit.services.storage_service import StorageService
 
 CONTEXT_ERROR_MESSAGE = 'Tu consulta supera el límite de contexto, utiliza el boton de recarga de contexto.'
 
@@ -33,11 +34,13 @@ class llmClient:
                  llmquery_repo: LLMQueryRepo,
                  llm_proxy: LLMProxy,
                  model_registry: ModelRegistry,
+                 storage_service: StorageService,
                  util: Utility
                  ):
         self.llmquery_repo = llmquery_repo
         self.llm_proxy = llm_proxy
         self.model_registry = model_registry
+        self.storage_service = storage_service
         self.util = util
         self._dispatcher = None # Cache for the lazy-loaded dispatcher
 
@@ -69,8 +72,10 @@ class llmClient:
                text: dict,
                model: str,
                context_history: Optional[List[Dict]] = None,
+               images: list = None,
                ) -> dict:
 
+        images = images or []
         f_calls = []  # keep track of the function calls executed by the LLM
         f_call_time = 0
         response = None
@@ -84,7 +89,7 @@ class llmClient:
 
         try:
             start_time = time.time()
-            logging.info(f"calling llm model '{model}' with {self.count_tokens(context, context_history)} tokens...")
+            logging.info(f"calling llm model '{model}' with {self.count_tokens(context, context_history)} tokens...and {len(images)} images...")
 
             # this is the first call to the LLM on the iteration
             try:
@@ -102,6 +107,7 @@ class llmClient:
                     tools=tools,
                     text=text_payload,
                     reasoning=reasoning,
+                    images=images,
                 )
                 stats = self.get_stats(response)
 
@@ -163,7 +169,7 @@ class llmClient:
                         error_message = f"Dispatch error en {function_name} con args {args} -******- {str(e)}"
                         raise IAToolkitException(IAToolkitException.ErrorType.CALL_ERROR, error_message)
 
-                    # add  the return value into the list of messages
+                    # add the return value into the list of messages
                     input_messages.append({
                         "type": "function_call_output",
                         "call_id": tool_call.call_id,
@@ -198,9 +204,14 @@ class llmClient:
                     tool_choice=tool_choice_value,
                     tools=tools,
                     text=text_payload,
+                    images=images,
                 )
                 stats_fcall = self.add_stats(stats_fcall, self.get_stats(response))
 
+            # --- IMAGE PROCESSING ---
+            # before save or respond, upload the images to S3 and clean content_parts
+            self._process_generated_images(response, company.short_name)
+
             # save the statistices
             stats['response_time']=int(time.time() - start_time)
             stats['sql_retry_count'] = sql_retry_count
@@ -239,6 +250,7 @@ class llmClient:
                 'query_id': query.id,
                 'model': model,
                 'reasoning_content': final_reasoning,
+                'content_parts': response.content_parts
             }
         except SQLAlchemyError as db_error:
             # rollback
@@ -295,6 +307,50 @@ class llmClient:
 
         return response.id
 
+    def _process_generated_images(self, response, company_short_name: str):
+        """
+        Traverse content_parts, detect images in Base64, upload to S3 and update content_parts.
+        """
+        if not response.content_parts:
+            return
+
+        for part in response.content_parts:
+            if part.get('type') == 'image':
+                source = part.get('source', {})
+                if source.get('type') in ['base64', 'url']:
+                    try:
+                        if source.get('type') == 'url':
+                            url = source.get('url')
+                            storage_key = None
+                        else:
+                            # upload image to S3
+                            result = self.storage_service.store_generated_image(
+                                company_short_name,
+                                source.get('data'),
+                                source.get('media_type', 'image/png')
+                            )
+                            url = result['url']
+                            storage_key = result['storage_key']
+
+                        # Update content_part: Now it's a remote reference, not base64 anymore.
+                        # We keep 'url' for the frontend to display it itself, and storage_key for internal reference.
+                        part['source'] = {
+                            'type': 'url',
+                            'url': url,
+                            'storage_key': storage_key,
+                            'media_type': source.get('media_type')
+                        }
+
+                        # clean data
+                        logging.info(f"Imagen procesada y subida: {url}")
+
+                    except Exception as e:
+                        logging.error(f"Fallo al subir imagen generada: {e}")
+
+                        # Fallback: keep the base64 and signal the error
+                        part['error'] = "Failed to upload image"
+
+
     def decode_response(self, response) -> dict:
         message = response.output_text
         decoded_response = {