iatoolkit 0.71.4__py3-none-any.whl → 1.4.2__py3-none-any.whl

iatoolkit/repositories/vs_repo.py

@@ -9,6 +9,7 @@ from iatoolkit.common.exceptions import IAToolkitException
 from iatoolkit.repositories.database_manager import DatabaseManager
 from iatoolkit.services.embedding_service import EmbeddingService
 from iatoolkit.repositories.models import Document, VSDoc, Company
+from typing import Dict
 import logging
 
 
@@ -29,17 +30,18 @@ class VSRepo:
                 self.session.add(doc)
             self.session.commit()
         except Exception as e:
-            logging.error(f"Error inserting documents into PostgreSQL: {str(e)}")
+            logging.error(f"Error while inserting embedding chunk list: {str(e)}")
             self.session.rollback()
             raise IAToolkitException(IAToolkitException.ErrorType.VECTOR_STORE_ERROR,
-                               f"Error inserting documents into PostgreSQL: {str(e)}")
+                               f"Error while inserting embedding chunk list: {str(e)}")
 
     def query(self,
               company_short_name: str,
               query_text: str,
               n_results=5,
-              metadata_filter=None
-              ) -> list[Document]:
+              metadata_filter=None,
+              collection_id: int = None
+              ) -> list[Dict]:
         """
         search documents similar to the query for a company
 
@@ -53,7 +55,12 @@ class VSRepo:
             list of documents matching the query and filters
         """
         # Generate the embedding with the query text for the specific company
-        query_embedding = self.embedding_service.embed_text(company_short_name, query_text)
+        try:
+            query_embedding = self.embedding_service.embed_text(company_short_name, query_text)
+        except Exception as e:
+            logging.error(f"error while creating text embedding: {str(e)}")
+            raise IAToolkitException(IAToolkitException.ErrorType.EMBEDDING_ERROR,
+                               f"embedding error: {str(e)}")
 
         sql_query, params = None, None
         try:
@@ -65,11 +72,12 @@ class VSRepo:
 
             # build the SQL query
             sql_query_parts = ["""
-                               SELECT iat_documents.id, \
+                               SELECT iat_vsdocs.id, \
                                       iat_documents.filename, \
-                                      iat_documents.content, \
+                                      iat_vsdocs.text, \
                                       iat_documents.content_b64, \
-                                      iat_documents.meta
+                                      iat_documents.meta,
+                                      iat_documents.id
                                FROM iat_vsdocs, \
                                     iat_documents
                                WHERE iat_vsdocs.company_id = :company_id
@@ -83,6 +91,10 @@ class VSRepo:
                 "n_results": n_results
             }
 
+            # Filter by Collection ID
+            if collection_id:
+                sql_query_parts.append(" AND iat_documents.collection_type_id = :collection_id")
+                params['collection_id'] = collection_id
 
             # add metadata filter, if exists
             if metadata_filter and isinstance(metadata_filter, dict):
@@ -98,7 +110,7 @@ class VSRepo:
             sql_query = "".join(sql_query_parts)
 
             # add sorting and limit of results
-            sql_query += " ORDER BY embedding <-> :query_embedding LIMIT :n_results"
+            sql_query += " ORDER BY embedding <-> CAST(:query_embedding AS VECTOR) LIMIT :n_results"
 
             logging.debug(f"Executing SQL query: {sql_query}")
             logging.debug(f"With parameters: {params}")
@@ -112,17 +124,18 @@ class VSRepo:
             for row in rows:
                 # create the document object with the data
                 meta_data = row[4] if len(row) > 4 and row[4] is not None else {}
-                doc = Document(
-                    id=row[0],
-                    company_id=company.id,
-                    filename=row[1],
-                    content=row[2],
-                    content_b64=row[3],
-                    meta=meta_data
+                vs_documents.append(
+                    {
+                        'id': row[0],
+                        'document_id': row[5],
+                        'filename': row[1],
+                        'text': row[2],
+                        'meta': meta_data,
+                    }
                 )
-                vs_documents.append(doc)
 
-            return self.remove_duplicates_by_id(vs_documents)
+            return vs_documents
+
 
         except Exception as e:
             logging.error(f"Error en la consulta de documentos: {str(e)}")
@@ -133,13 +146,3 @@ class VSRepo:
         finally:
             self.session.close()
 
-    def remove_duplicates_by_id(self, objects):
-        unique_by_id = {}
-        result = []
-
-        for obj in objects:
-            if obj.id not in unique_by_id:
-                unique_by_id[obj.id] = True
-                result.append(obj)
-
-        return result

iatoolkit/services/auth_service.py

@@ -132,7 +132,7 @@ class AuthService:
         # check if the api-key is valid and active
         api_key_entry = self.profile_service.get_active_api_key_entry(api_key)
         if not api_key_entry:
-            logging.info(f"Invalid or inactive API Key {api_key}")
+            logging.error(f"Invalid or inactive IAToolkit API Key: {api_key}")
             return {"success": False,
                     "error_message": self.i18n_service.t('errors.auth.invalid_api_key'),
                     "status_code": 402}

iatoolkit/services/branding_service.py

@@ -42,7 +42,7 @@ class BrandingService:
             # Estilos para Alertas de Error ---
             "brand_danger_color": "#dc3545",  # Rojo principal para alertas
             "brand_danger_bg": "#f8d7da",  # Fondo rojo pálido
-            "brand_danger_text": "#842029",  # Texto rojo oscuro
+            "brand_danger_text": "#000000",
             "brand_danger_border": "#f5c2c7",  # Borde rojo intermedio
 
             # Estilos para Alertas Informativas ---

iatoolkit/services/company_context_service.py

@@ -5,11 +5,11 @@
 
 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
-import os
 
 
 class CompanyContextService:
@@ -22,21 +22,22 @@ class CompanyContextService:
     def __init__(self,
                  sql_service: SqlService,
                  utility: Utility,
-                 config_service: ConfigurationService):
+                 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 for APIs, etc.).
-        3. Dynamic SQL database schema from the live connection.
+        2. Static schema files (YAML files for SQL data sources).
         """
         context_parts = []
 
-        # 1. Context from Markdown (context/*.md) and yaml (schema/*.yaml) files
+        # 1. Context from Markdown (context/*.md)  files
         try:
             md_context = self._get_static_file_context(company_short_name)
             if md_context:
@@ -44,7 +45,7 @@ class CompanyContextService:
         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)
+        # 2. Context from company-specific SQL databases
         try:
             sql_context = self._get_sql_schema_context(company_short_name)
             if sql_context:
@@ -52,29 +53,37 @@ class CompanyContextService:
         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 and .yaml schema files.
+        # Get context from .md files using the repository
         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)
+        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
 
@@ -94,21 +103,38 @@ class CompanyContextService:
             if not db_name:
                 continue
 
+            # get database schema definition, for this source.
+            database_schema_name = source.get('schema')
+
             try:
-                db_manager = self.sql_service.get_database_manager(db_name)
+                db_provider = self.sql_service.get_database_provider(company_short_name, db_name)
             except IAToolkitException as e:
-                logging.warning(f"Could not get DB manager for '{db_name}': {e}")
+                logging.warning(f"Could not get DB provider for '{db_name}': {e}")
                 continue
 
             db_description = source.get('description', '')
-            sql_context = f'***Base de datos (database_name)***: {db_name}\n'
-            sql_context += f"**Descripción:**: {db_description}\n" if db_description else ""
-            sql_context += "Para consultar esta base de datos debes utilizar el servicio ***iat_sql_query***.\n"
+            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_manager.get_all_table_names()
+                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:
@@ -126,30 +152,61 @@ class CompanyContextService:
                     # 4. get the table specific configuration.
                     table_config = table_overrides.get(table_name, {})
 
-                    # 5. define the schema name, using the override if it exists.
+                    # 5. define the schema object name, using the override if it exists.
                     # Priority 1: Explicit override from the 'tables' map.
-                    schema_name = table_config.get('schema_name')
+                    schema_object_name = table_config.get('schema_name')
 
-                    if not schema_name:
-                        # Priority 2: Automatic prefix stripping.
+                    if not schema_object_name:
+                        # Priority 3: Automatic prefix stripping.
                         if table_prefix and table_name.startswith(table_prefix):
-                            schema_name = table_name[len(table_prefix):]
+                            schema_object_name = table_name[len(table_prefix):]
                         else:
-                            # Priority 3: Default to the table name itself.
-                            schema_name = table_name
+                            # 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_manager.get_table_schema(
+                    table_definition = db_provider.get_table_description(
                         table_name=table_name,
-                        schema_name=schema_name,
+                        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}")
 
-        return sql_context
+        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