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

iatoolkit/repositories/profile_repo.py

@@ -3,10 +3,11 @@
 #
 # IAToolkit is open source software.
 
-from iatoolkit.repositories.models import User, Company, ApiKey, UserFeedback
+from iatoolkit.repositories.models import (User, Company, user_company,
+                                           ApiKey, UserFeedback, AccessLog)
 from injector import inject
 from iatoolkit.repositories.database_manager import DatabaseManager
-from sqlalchemy.orm import joinedload # Para cargar la relación eficientemente
+from sqlalchemy import select, func, and_
 
 
 class ProfileRepo:
@@ -69,8 +70,63 @@ class ProfileRepo:
     def get_companies(self) -> list[Company]:
         return self.session.query(Company).all()
 
+    def get_user_role_in_company(self, company_id, user_id, ):
+        stmt = (
+            select(user_company.c.role)
+            .where(
+                user_company.c.user_id == user_id,
+                user_company.c.company_id == company_id,
+            )
+        )
+        result = self.session.execute(stmt).scalar_one_or_none()
+        return result
+
+    def get_companies_by_user_identifier(self, user_identifier: str) -> list:
+        """
+        Return all the companies to which the user belongs (by email),
+        and the role he has in each company.
+        """
+        return (
+            self.session.query(Company, user_company.c.role)
+            .join(user_company, Company.id == user_company.c.company_id)
+            .join(User, User.id == user_company.c.user_id)
+            .filter(User.email == user_identifier)
+            .all()
+        )
+
+    def get_company_users_with_details(self, company_short_name: str) -> list[dict]:
+        # returns the list of users in the company with their role and last access date
+
+        # subquery for last access date
+        last_access_sq = (
+            self.session.query(
+                AccessLog.user_identifier,
+                func.max(AccessLog.timestamp).label("max_ts")
+            )
+            .filter(AccessLog.company_short_name == company_short_name)
+            .group_by(AccessLog.user_identifier)
+            .subquery()
+        )
+
+        # main query
+        stmt = (
+            self.session.query(
+                User,
+                user_company.c.role,
+                last_access_sq.c.max_ts
+            )
+            .join(user_company, User.id == user_company.c.user_id)
+            .join(Company, Company.id == user_company.c.company_id)
+            .outerjoin(last_access_sq, User.email == last_access_sq.c.user_identifier)
+            .filter(Company.short_name == company_short_name)
+        )
+
+        results = stmt.all()
+
+        return results
+
     def create_company(self, new_company: Company):
-        company = self.session.query(Company).filter_by(name=new_company.name).first()
+        company = self.session.query(Company).filter_by(short_name=new_company.short_name).first()
         if company:
             if company.parameters != new_company.parameters:
                 company.parameters = new_company.parameters

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
 
 
@@ -38,8 +39,9 @@ class VSRepo:
               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
 
@@ -70,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
@@ -88,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):
@@ -117,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)}")
@@ -138,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/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,10 +22,12 @@ 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:
         """
@@ -35,7 +37,7 @@ class CompanyContextService:
         """
         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:
@@ -43,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 database schemas (schema/*.yaml files)
+        # 2. Context from company-specific SQL databases
         try:
             sql_context = self._get_sql_schema_context(company_short_name)
             if sql_context:
@@ -51,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
 
@@ -93,49 +103,63 @@ class CompanyContextService:
             if not db_name:
                 continue
 
+            # get database schema definition, for this source.
+            database_schema_name = source.get('schema', 'public')
+
             try:
-                db_manager = self.sql_service.get_database_manager(db_name)
+                # 1. Get the full database structure at once using the SQL service
+                db_structure = self.sql_service.get_database_structure(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 structure 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.
+            # 2. get the list of tables to process based on structure and config
             tables_to_process = []
             if source.get('include_all_tables', False):
-                all_tables = db_manager.get_all_table_names()
+                # Use keys from the fetched structure
+                all_tables = list(db_structure.keys())
                 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())
+                # Use keys from the config map, but check if they exist in DB structure
+                config_tables = list(source['tables'].keys())
+                tables_to_process = [t for t in config_tables if t in db_structure]
 
-            # 2. get the global settings and overrides.
+            # 3. get the global settings and overrides.
             global_exclude_columns = source.get('exclude_columns', [])
             table_prefix = source.get('table_prefix')
-
-            # get the global schema definition, for this source.
-            global_schema_name = source.get('schema')
-
             table_overrides = source.get('tables', {})
 
-            # 3. iterate over the tables.
+            # 4. iterate over the tables.
             for table_name in tables_to_process:
                 try:
-                    # 4. get the table specific configuration.
+                    table_data = db_structure[table_name]
+
+                    # 5. get the table specific configuration.
                     table_config = table_overrides.get(table_name, {})
 
-                    # 5. define the schema object name, using the override if it exists.
+                    # 6. 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_object_name')
-
-                    # Priority 2: Global schema defined in the source.
-                    if not schema_object_name and global_schema_name:
-                        schema_object_name = global_schema_name
+                    schema_object_name = table_config.get('schema_name')
 
                     if not schema_object_name:
                         # Priority 3: Automatic prefix stripping.
@@ -145,19 +169,68 @@ class CompanyContextService:
                             # Priority 4: Default to the table name itself.
                             schema_object_name = table_name
 
-                    # 6. define the list of columns to exclude, (local vs. global).
+                    # 7. define the list of columns to exclude, (local vs. global).
                     local_exclude_columns = table_config.get('exclude_columns')
                     final_exclude_columns = local_exclude_columns if local_exclude_columns is not None else global_exclude_columns
 
-                    # 7. get the table schema definition.
-                    table_definition = db_manager.get_table_schema(
-                        table_name=table_name,
-                        db_schema=db_manager.schema,
-                        schema_object_name=schema_object_name,
-                        exclude_columns=final_exclude_columns
-                    )
-                    sql_context += table_definition
+                    # 8. Build the table definition dictionary manually using the structure data
+                    json_dict = {
+                        "table": table_name,
+                        "schema": database_schema_name,
+                        "description": f"The table belongs to the **`{database_schema_name}`** schema.",
+                        "fields": []
+                    }
+
+                    if schema_object_name:
+                        json_dict["description"] += (
+                            f"The meaning of each field in this table is detailed in the **`{schema_object_name}`** object."
+                        )
+
+                    for col in table_data.get('columns', []):
+                        name = col["name"]
+                        if name in final_exclude_columns:
+                            continue
+
+                        json_dict["fields"].append({
+                            "name": name,
+                            "type": col["type"]
+                        })
+
+                    # Append as string representation of dict (consistent with previous behavior)
+                    sql_context += "\n\n" + str(json_dict)
+
                 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