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

iatoolkit/services/prompt_service.py

@@ -4,12 +4,14 @@
 # IAToolkit is open source software.
 
 from injector import inject
+from iatoolkit import current_iatoolkit
 from iatoolkit.common.interfaces.asset_storage import AssetRepository, AssetType
 from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
 from iatoolkit.services.i18n_service import I18nService
 from iatoolkit.repositories.profile_repo import ProfileRepo
 from collections import defaultdict
-from iatoolkit.repositories.models import Prompt, PromptCategory, Company
+from iatoolkit.repositories.models import (Prompt, PromptCategory,
+                                           Company, PromptType)
 from iatoolkit.common.exceptions import IAToolkitException
 import importlib.resources
 import logging
@@ -17,9 +19,9 @@ import os
 
 # iatoolkit system prompts definitions
 _SYSTEM_PROMPTS = [
-    {'name': 'query_main', 'description': 'iatoolkit main prompt'},
-    {'name': 'format_styles', 'description': 'output format styles'},
-    {'name': 'sql_rules', 'description': 'instructions  for SQL queries'}
+    {'name': 'query_main', 'description': 'iatoolkit main prompt', 'order': 1},
+    {'name': 'format_styles', 'description': 'output format styles', 'order': 2},
+    {'name': 'sql_rules', 'description': 'instructions  for SQL queries', 'order': 3},
 ]
 
 class PromptService:
@@ -34,14 +36,197 @@ class PromptService:
         self.profile_repo = profile_repo
         self.i18n_service = i18n_service
 
-    def sync_company_prompts(self, company_short_name: str, prompts_config: list, categories_config: list):
+    def get_prompts(self, company_short_name: str, include_all: bool = False) -> dict:
+        try:
+            # validate company
+            company = self.profile_repo.get_company_by_short_name(company_short_name)
+            if not company:
+                return {"error": self.i18n_service.t('errors.company_not_found', company_short_name=company_short_name)}
+
+            # get all the company prompts
+            # If include_all is True, repo should return everything for the company
+            # Otherwise, it should return only active prompts
+            all_prompts = self.llm_query_repo.get_prompts(company, include_all=include_all)
+
+            # Deduplicate prompts by id
+            all_prompts = list({p.id: p for p in all_prompts}.values())
+
+            # group by category
+            prompts_by_category = defaultdict(list)
+            for prompt in all_prompts:
+                # Filter logic moved here or in repo.
+                # If include_all is False, we only want active prompts (and maybe only specific types)
+                if not include_all:
+
+                    # Standard user view: excludes system/agent hidden prompts if any?
+                    if prompt.prompt_type != PromptType.COMPANY.value:
+                        continue
+
+                # Grouping logic
+                cat_key = (0, "Uncategorized") # Default
+                if prompt.category:
+                    cat_key = (prompt.category.order, prompt.category.name)
+
+                prompts_by_category[cat_key].append(prompt)
+
+            # sort each category by order
+            for cat_key in prompts_by_category:
+                prompts_by_category[cat_key].sort(key=lambda p: p.order)
+
+            categorized_prompts = []
+
+            # sort categories by order
+            sorted_categories = sorted(prompts_by_category.items(), key=lambda item: item[0][0])
+
+            for (cat_order, cat_name), prompts in sorted_categories:
+                categorized_prompts.append({
+                    'category_name': cat_name,
+                    'category_order': cat_order,
+                    'prompts': [
+                        {
+                            'prompt': p.name,
+                            'description': p.description,
+                            'type': p.prompt_type,
+                            'active': p.active,
+                            'custom_fields': p.custom_fields,
+                            'order': p.order
+                        }
+                        for p in prompts
+                    ]
+                })
+
+            return {'message': categorized_prompts}
+
+        except Exception as e:
+            logging.error(f"error in get_prompts: {e}")
+            return {'error': str(e)}
+
+
+    def get_prompt_content(self, company: Company, prompt_name: str):
+        try:
+            # get the prompt from database
+            prompt = self.llm_query_repo.get_prompt_by_name(company, prompt_name)
+            if not prompt:
+                raise IAToolkitException(IAToolkitException.ErrorType.DOCUMENT_NOT_FOUND,
+                                   f"prompt not found '{prompt}' for company '{company.short_name}'")
+
+            try:
+                # read the prompt content from asset repository
+                user_prompt_content = self.asset_repo.read_text(
+                    company.short_name,
+                    AssetType.PROMPT,
+                    prompt.filename
+                )
+            except FileNotFoundError:
+                raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
+                                         f"prompt file '{prompt.filename}' does not exist for company '{company.short_name}'")
+            except Exception as e:
+                raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
+                                         f"error while reading prompt: '{prompt_name}': {e}")
+
+            return user_prompt_content
+
+        except IAToolkitException:
+            raise
+        except Exception as e:
+            logging.exception(
+                f"error loading prompt '{prompt_name}' content for '{company.short_name}': {e}")
+            raise IAToolkitException(IAToolkitException.ErrorType.PROMPT_ERROR,
+                               f'error loading prompt "{prompt_name}" content for company {company.short_name}: {str(e)}')
+
+    def save_prompt(self, company_short_name: str, prompt_name: str, data: dict):
+        """
+        Create or Update a prompt.
+        1. Saves the Jinja content to the .prompt asset file.
+        2. Updates the Database.
+        """
+        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")
+
+        # Validate category if present
+        category_id = None
+        if 'category' in data:
+             # simple lookup, assuming category names are unique per company
+             cat = self.llm_query_repo.get_category_by_name(company.id, data['category'])
+             if cat:
+                 category_id = cat.id
+
+        # 1. save the phisical part of the prompt (content)
+        if 'content' in data:
+            filename = f"{prompt_name}.prompt"
+            filename = filename.lower().replace(' ', '_')
+            self.asset_repo.write_text(company_short_name, AssetType.PROMPT, filename, data['content'])
+
+        # 2. update the prompt in the database
+        new_prompt = Prompt(
+            company_id=company.id,
+            name=prompt_name,
+            description=data.get('description', ''),
+            order=data.get('order', 1),
+            category_id=category_id,
+            active=data.get('active', True),
+            prompt_type=data.get('prompt_type', 'company'),
+            filename=f"{prompt_name.lower().replace(' ', '_')}.prompt",
+            custom_fields=data.get('custom_fields', [])
+        )
+        self.llm_query_repo.create_or_update_prompt(new_prompt)
+
+    def delete_prompt(self, company_short_name: str, prompt_name: str):
+        """
+        Deletes a prompt:
+        1. Removes from DB.
+        2. Removes from YAML config.
+        3. (Optional) Deletes/Archives physical file.
+        """
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME, f"Company not found")
+
+        prompt_db = self.llm_query_repo.get_prompt_by_name(company, prompt_name)
+        if not prompt_db:
+            raise IAToolkitException(IAToolkitException.ErrorType.DOCUMENT_NOT_FOUND, f"Prompt {prompt_name} not found")
+
+        # 1. Remove from DB
+        self.llm_query_repo.delete_prompt(prompt_db)
+
+    def get_system_prompt(self):
+        try:
+            system_prompt_content = []
+
+            # read all the system prompts from the database
+            system_prompts = self.llm_query_repo.get_system_prompts()
+
+            for prompt in system_prompts:
+                try:
+                    content = importlib.resources.read_text('iatoolkit.system_prompts', prompt.filename)
+                    system_prompt_content.append(content)
+                except FileNotFoundError:
+                    logging.warning(f"Prompt file does not exist in the package: {prompt.filename}")
+                except Exception as e:
+                    raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
+                                             f"error reading system prompt '{prompt.filename}': {e}")
+
+            # join the system prompts into a single string
+            return "\n".join(system_prompt_content)
+
+        except IAToolkitException:
+            raise
+        except Exception as e:
+            logging.exception(
+                f"Error al obtener el contenido del prompt de sistema: {e}")
+            raise IAToolkitException(IAToolkitException.ErrorType.PROMPT_ERROR,
+                               f'error reading the system prompts": {str(e)}')
+
+    def sync_company_prompts(self, company_short_name: str, prompt_list: list, categories_config: list):
         """
         Synchronizes prompt categories and prompts from YAML config to Database.
         Strategies:
         - Categories: Create or Update existing based on name.
         - Prompts: Create or Update existing based on name. Soft-delete or Delete unused.
         """
-        if not prompts_config:
+        if not prompt_list:
             return
 
         company = self.profile_repo.get_company_by_short_name(company_short_name)
@@ -49,6 +234,13 @@ class PromptService:
             raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
                                      f'Company {company_short_name} not found')
 
+        # Register system prompts
+        self._register_system_prompts(company)
+
+        # community edition has its own prompt management
+        if not current_iatoolkit().is_community:
+            return
+
         try:
             # 1. Sync Categories
             category_map = {}
@@ -66,7 +258,7 @@ class PromptService:
             # 2. Sync Prompts
             defined_prompt_names = set()
 
-            for prompt_data in prompts_config:
+            for prompt_data in prompt_list:
                 category_name = prompt_data.get('category')
                 if not category_name or category_name not in category_map:
                     logging.warning(
@@ -86,7 +278,7 @@ class PromptService:
                     order=prompt_data.get('order'),
                     category_id=category_obj.id,
                     active=prompt_data.get('active', True),
-                    is_system_prompt=False,
+                    prompt_type=prompt_data.get('prompt_type', PromptType.COMPANY.value).lower(),
                     filename=filename,
                     custom_fields=prompt_data.get('custom_fields', [])
                 )
@@ -106,30 +298,42 @@ class PromptService:
             self.llm_query_repo.rollback()
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))
 
-    def register_system_prompts(self):
+    def _register_system_prompts(self, company: Company):
         """
         Synchronizes system prompts defined in Dispatcher/Code to Database.
         """
+
+        # if there are system prompts already registered, skip
+        # if self.llm_query_repo.get_system_prompts(): return
+
+        sys_category = PromptCategory(company_id=company.id, name="System", order=0)
+        self.llm_query_repo.create_or_update_prompt_category(sys_category)
+
         try:
             defined_names = set()
 
             for i, prompt_data in enumerate(_SYSTEM_PROMPTS):
                 prompt_name = prompt_data['name']
                 defined_names.add(prompt_name)
+                prompt_filename = f"{prompt_name}.prompt"
 
                 new_prompt = Prompt(
-                    company_id=None,  # System prompts have no company
+                    company_id=company.id,
                     name=prompt_name,
                     description=prompt_data['description'],
-                    order=i + 1,
-                    category_id=None,
+                    order=prompt_data['order'],
+                    category_id=sys_category.id,
                     active=True,
-                    is_system_prompt=True,
-                    filename=f"{prompt_name}.prompt",
+                    prompt_type=PromptType.SYSTEM.value,
+                    filename=prompt_filename,
                     custom_fields=[]
                 )
                 self.llm_query_repo.create_or_update_prompt(new_prompt)
 
+                # add prompt to company assets
+                prompt_content = importlib.resources.read_text('iatoolkit.system_prompts', prompt_filename)
+                self.asset_repo.write_text(company.short_name, AssetType.PROMPT, prompt_filename, prompt_content)
+
             # Cleanup old system prompts
             existing_sys_prompts = self.llm_query_repo.get_system_prompts()
             for p in existing_sys_prompts:
@@ -142,162 +346,45 @@ class PromptService:
             self.llm_query_repo.rollback()
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))
 
-    def create_prompt(self,
-                      prompt_name: str,
-                      description: str,
-                      order: int,
-                      company: Company = None,
-                      category: PromptCategory = None,
-                      active: bool = True,
-                      is_system_prompt: bool = False,
-                      custom_fields: list = []
-                      ):
+    def sync_prompt_categories(self, company_short_name: str, categories_config: list):
         """
-            Direct creation method (used by sync or direct calls).
-            Validates file existence before creating DB entry.
+        Syncs only the prompt categories based on a simple list of names.
+        The order in the list determines the 'order' field in DB.
+        Removes categories not present in the list.
+        Finally, updates the YAML configuration.
         """
-        prompt_filename = prompt_name.lower() + '.prompt'
-        if is_system_prompt:
-            if not importlib.resources.files('iatoolkit.system_prompts').joinpath(prompt_filename).is_file():
-                raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                                f'missing system prompt file: {prompt_filename}')
-        else:
-            if not self.asset_repo.exists(company.short_name, AssetType.PROMPT, prompt_filename):
-                raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                               f'missing prompt file: {prompt_filename} in prompts/')
-
-        if custom_fields:
-            for f in custom_fields:
-                if ('data_key' not in f) or ('label' not in f):
-                    raise IAToolkitException(IAToolkitException.ErrorType.INVALID_PARAMETER,
-                               f'The field "custom_fields" must contain the following keys: data_key y label')
-
-                # add default value for data_type
-                if 'type' not in f:
-                    f['type'] = 'text'
-
-        prompt = Prompt(
-                company_id=company.id if company else None,
-                name=prompt_name,
-                description=description,
-                order=order,
-                category_id=category.id if category and not is_system_prompt else None,
-                active=active,
-                filename=prompt_filename,
-                is_system_prompt=is_system_prompt,
-                custom_fields=custom_fields
-            )
-
-        try:
-            self.llm_query_repo.create_or_update_prompt(prompt)
-        except Exception as e:
-            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
-                               f'error creating prompt "{prompt_name}": {str(e)}')
+        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')
 
-    def get_prompt_content(self, company: Company, prompt_name: str):
         try:
-            # get the user prompt
-            user_prompt = self.llm_query_repo.get_prompt_by_name(company, prompt_name)
-            if not user_prompt:
-                raise IAToolkitException(IAToolkitException.ErrorType.DOCUMENT_NOT_FOUND,
-                                   f"prompt not found '{prompt_name}' for company '{company.short_name}'")
+            processed_categories_ids = []
 
-            try:
-                user_prompt_content = self.asset_repo.read_text(
-                    company.short_name,
-                    AssetType.PROMPT,
-                    user_prompt.filename
+            # 1. Update/Create Categories
+            for idx, cat_name in enumerate(categories_config):
+                # Order is 0-based index or 1-based, consistent with current usage (seems 0 or 1 is fine, usually 0 for arrays)
+                new_cat = PromptCategory(
+                    company_id=company.id,
+                    name=cat_name,
+                    order=idx
                 )
-            except FileNotFoundError:
-                raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
-                                         f"prompt file '{user_prompt.filename}' does not exist for company '{company.short_name}'")
-            except Exception as e:
-                raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
-                                         f"error while reading prompt: '{prompt_name}': {e}")
-
-            return user_prompt_content
-
-        except IAToolkitException:
-            raise
-        except Exception as e:
-            logging.exception(
-                f"error loading prompt '{prompt_name}' content for '{company.short_name}': {e}")
-            raise IAToolkitException(IAToolkitException.ErrorType.PROMPT_ERROR,
-                               f'error loading prompt "{prompt_name}" content for company {company.short_name}: {str(e)}')
+                persisted_cat = self.llm_query_repo.create_or_update_prompt_category(new_cat)
+                processed_categories_ids.append(persisted_cat.id)
+
+            # 2. Delete missing categories
+            # We fetch all categories for the company and delete those not in processed_ids
+            all_categories = self.llm_query_repo.get_all_categories(company.id)
+            for cat in all_categories:
+                if cat.id not in processed_categories_ids:
+                    # Depending on logic, we might want to check if they have prompts assigned.
+                    # Usually, sync logic implies "force state", so we delete.
+                    # SQLAlchemy cascading might handle prompts or set them to null depending on model config.
+                    self.llm_query_repo.session.delete(cat)
 
-    def get_system_prompt(self):
-        try:
-            system_prompt_content = []
-
-            # read all the system prompts from the database
-            system_prompts = self.llm_query_repo.get_system_prompts()
-
-            for prompt in system_prompts:
-                try:
-                    content = importlib.resources.read_text('iatoolkit.system_prompts', prompt.filename)
-                    system_prompt_content.append(content)
-                except FileNotFoundError:
-                    logging.warning(f"Prompt file does not exist in the package: {prompt.filename}")
-                except Exception as e:
-                    raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
-                                             f"error reading system prompt '{prompt.filename}': {e}")
-
-            # join the system prompts into a single string
-            return "\n".join(system_prompt_content)
-
-        except IAToolkitException:
-            raise
-        except Exception as e:
-            logging.exception(
-                f"Error al obtener el contenido del prompt de sistema: {e}")
-            raise IAToolkitException(IAToolkitException.ErrorType.PROMPT_ERROR,
-                               f'error reading the system prompts": {str(e)}')
-
-    def get_user_prompts(self, company_short_name: str) -> dict:
-        try:
-            # validate company
-            company = self.profile_repo.get_company_by_short_name(company_short_name)
-            if not company:
-                return {"error": self.i18n_service.t('errors.company_not_found', company_short_name=company_short_name)}
-
-            # get all the prompts
-            all_prompts = self.llm_query_repo.get_prompts(company)
-
-            # group by category
-            prompts_by_category = defaultdict(list)
-            for prompt in all_prompts:
-                if prompt.active:
-                    if prompt.category:
-                        cat_key = (prompt.category.order, prompt.category.name)
-                        prompts_by_category[cat_key].append(prompt)
-
-            # sort each category by order
-            for cat_key in prompts_by_category:
-                prompts_by_category[cat_key].sort(key=lambda p: p.order)
-
-            categorized_prompts = []
-
-            # sort categories by order
-            sorted_categories = sorted(prompts_by_category.items(), key=lambda item: item[0][0])
-
-            for (cat_order, cat_name), prompts in sorted_categories:
-                categorized_prompts.append({
-                    'category_name': cat_name,
-                    'category_order': cat_order,
-                    'prompts': [
-                        {
-                            'prompt': p.name,
-                            'description': p.description,
-                            'custom_fields': p.custom_fields,
-                            'order': p.order
-                        }
-                        for p in prompts
-                    ]
-                })
-
-            return {'message': categorized_prompts}
+            self.llm_query_repo.commit()
 
         except Exception as e:
-            logging.error(f"error in get_prompts: {e}")
-            return {'error': str(e)}
-
+            self.llm_query_repo.rollback()
+            logging.exception(f"Error syncing prompt categories: {e}")
+            raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR, str(e))

iatoolkit/services/query_service.py

@@ -99,7 +99,7 @@ class QueryService:
         final_client_data.update(client_data)
 
         # Load attached files into the context
-        files_context = self.load_files_for_context(files)
+        files_context, images = self.load_files_for_context(files)
 
         # Initialize prompt_content. It will be an empty string for direct questions.
         main_prompt = ""
@@ -128,7 +128,7 @@ class QueryService:
         else:
             user_turn_prompt += f'\n### Contexto Adicional: El usuario ha aportado este contexto puede ayudar: {effective_question}'
 
-        return user_turn_prompt, effective_question
+        return user_turn_prompt, effective_question, images
 
     def _ensure_valid_history(self, company,
                               user_identifier: str,
@@ -356,7 +356,7 @@ class QueryService:
             effective_model = self._resolve_model(company_short_name, model)
 
             # --- Build User-Facing Prompt ---
-            user_turn_prompt, effective_question = self._build_user_facing_prompt(
+            user_turn_prompt, effective_question, images = self._build_user_facing_prompt(
                 company=company,
                 user_identifier=user_identifier,
                 client_data=client_data,
@@ -397,7 +397,8 @@ class QueryService:
                 question=effective_question,
                 context=user_turn_prompt,
                 tools=tools,
-                text=output_schema
+                text=output_schema,
+                images=images,
             )
 
             if not response.get('valid_response'):
@@ -421,24 +422,23 @@ class QueryService:
             return "unknown"
 
 
-    def load_files_for_context(self, files: list) -> str:
+    def load_files_for_context(self, files: list) -> tuple[str, list]:
         """
-        Processes a list of attached files, decodes their content,
-        and formats them into a string context for the LLM.
+        Processes a list of attached files.
+        Decodes text documents into context string and separates images for multimodal processing.
         """
         if not files:
-            return ''
+            return '', []
+
+        context_parts = []
+        images = []
+        text_files_count = 0
 
-        context = f"""
-            A continuación encontraras una lista de documentos adjuntos
-            enviados por el usuario que hace la pregunta, 
-            en total son: {len(files)} documentos adjuntos
-            """
         for document in files:
             # Support both 'file_id' and 'filename' for robustness
             filename = document.get('file_id') or document.get('filename') or document.get('name')
             if not filename:
-                context += "\n<error>Documento adjunto sin nombre ignorado.</error>\n"
+                context_parts.append("\n<error>Documento adjunto sin nombre ignorado.</error>\n")
                 continue
 
             # Support both 'base64' and 'content' for robustness
@@ -446,7 +446,12 @@ class QueryService:
 
             if not base64_content:
                 # Handles the case where a file is referenced but no content is provided
-                context += f"\n<error>El archivo '{filename}' no fue encontrado y no pudo ser cargado.</error>\n"
+                context_parts.append(f"\n<error>El archivo '{filename}' no fue encontrado y no pudo ser cargado.</error>\n")
+                continue
+
+            # Detect if the file is an image
+            if self._is_image(filename):
+                images.append({'name': filename, 'base64': base64_content})
                 continue
 
             try:
@@ -456,12 +461,26 @@ class QueryService:
 
                 file_content = base64.b64decode(base64_content)
                 document_text = self.document_service.file_to_txt(filename, file_content)
-                context += f"\n<document name='{filename}'>\n{document_text}\n</document>\n"
+                context_parts.append(f"\n<document name='{filename}'>\n{document_text}\n</document>\n")
+                text_files_count += 1
             except Exception as e:
                 # Catches errors from b64decode or file_to_txt
                 logging.error(f"Failed to process file {filename}: {e}")
-                context += f"\n<error>Error al procesar el archivo {filename}: {str(e)}</error>\n"
+                context_parts.append(f"\n<error>Error al procesar el archivo {filename}: {str(e)}</error>\n")
                 continue
 
-        return context
+        context = ""
+        if text_files_count > 0:
+            context = f"""
+            A continuación encontraras una lista de documentos adjuntos
+            enviados por el usuario que hace la pregunta, 
+            en total son: {text_files_count} documentos adjuntos
+            """ + "".join(context_parts)
+        elif context_parts:
+            # If only errors were collected
+            context = "".join(context_parts)
+
+        return context, images
 
+    def _is_image(self, filename: str) -> bool:
+        return filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.gif'))