iatoolkit 0.11.0__py3-none-any.whl → 0.71.2__py3-none-any.whl

iatoolkit/services/query_service.py

@@ -4,11 +4,14 @@
 # IAToolkit is open source software.
 
 from iatoolkit.infra.llm_client import llmClient
+from iatoolkit.services.profile_service import ProfileService
 from iatoolkit.repositories.document_repo import DocumentRepo
 from iatoolkit.repositories.profile_repo import ProfileRepo
 from iatoolkit.services.document_service import DocumentService
+from iatoolkit.services.company_context_service import CompanyContextService
+from iatoolkit.services.i18n_service import I18nService
+from iatoolkit.services.configuration_service import ConfigurationService
 from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
-
 from iatoolkit.repositories.models import Task
 from iatoolkit.services.dispatcher_service import Dispatcher
 from iatoolkit.services.prompt_manager_service import PromptService
@@ -21,6 +24,7 @@ import logging
 from typing import Optional
 import json
 import time
+import hashlib
 import os
 
 
@@ -30,166 +34,248 @@ class QueryService:
     @inject
     def __init__(self,
                  llm_client: llmClient,
+                 profile_service: ProfileService,
+                 company_context_service: CompanyContextService,
                  document_service: DocumentService,
                  document_repo: DocumentRepo,
                  llmquery_repo: LLMQueryRepo,
                  profile_repo: ProfileRepo,
                  prompt_service: PromptService,
+                 i18n_service: I18nService,
                  util: Utility,
                  dispatcher: Dispatcher,
-                 session_context: UserSessionContextService
+                 session_context: UserSessionContextService,
+                 configuration_service: ConfigurationService
                  ):
+        self.profile_service = profile_service
+        self.company_context_service = company_context_service
         self.document_service = document_service
         self.document_repo = document_repo
         self.llmquery_repo = llmquery_repo
         self.profile_repo = profile_repo
         self.prompt_service = prompt_service
+        self.i18n_service = i18n_service
         self.util = util
         self.dispatcher = dispatcher
         self.session_context = session_context
+        self.configuration_service = configuration_service
         self.llm_client = llm_client
 
         # get the model from the environment variable
-        self.model = os.getenv("LLM_MODEL", "")
-        if not self.model:
+        self.default_model = os.getenv("LLM_MODEL", "")
+        if not self.default_model:
             raise IAToolkitException(IAToolkitException.ErrorType.API_KEY,
-                               "La variable de entorno 'LLM_MODEL' no está configurada.")
+                               "missing ENV variable 'LLM_MODEL' configuration.")
 
+    def init_context(self, company_short_name: str,
+                     user_identifier: str,
+                     model: str = None) -> dict:
 
-    def llm_init_context(self,
-                         company_short_name: str,
-                         external_user_id: str = None,
-                         local_user_id: int = 0,
-                         model: str = ''):
-        start_time = time.time()
-        if not model:
-            model = self.model
+        # 1. Execute the forced rebuild sequence using the unified identifier.
+        self.session_context.clear_all_context(company_short_name, user_identifier)
+        logging.info(f"Context for {company_short_name}/{user_identifier} has been cleared.")
 
-        # Validate the user and company
-        user_identifier, is_local_user = self.util.resolve_user_identifier(external_user_id, local_user_id)
-        if not user_identifier:
-            raise IAToolkitException(IAToolkitException.ErrorType.INVALID_USER,
-                        "No se pudo resolver el identificador del usuario")
+        # 2. LLM context is clean, now we can load it again
+        self.prepare_context(
+            company_short_name=company_short_name,
+            user_identifier=user_identifier
+        )
+
+        # 3. communicate the new context to the LLM
+        response = self.set_context_for_llm(
+            company_short_name=company_short_name,
+            user_identifier=user_identifier,
+            model=model
+        )
 
+        return response
+
+    def _build_context_and_profile(self, company_short_name: str, user_identifier: str) -> tuple:
+        # this method read the user/company context from the database and renders the system prompt
         company = self.profile_repo.get_company_by_short_name(company_short_name)
         if not company:
-            raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                               f"Empresa no encontrada: {company_short_name}")
+            return None, None
+
+        # Get the user profile from the single source of truth.
+        user_profile = self.profile_service.get_profile_by_identifier(company_short_name, user_identifier)
+
+        # render the iatoolkit main system prompt with the company/user information
+        system_prompt_template = self.prompt_service.get_system_prompt()
+        rendered_system_prompt = self.util.render_prompt_from_string(
+            template_string=system_prompt_template,
+            question=None,
+            client_data=user_profile,
+            company=company,
+            service_list=self.dispatcher.get_company_services(company)
+        )
+
+        # get the company context: schemas, database models, .md files
+        company_specific_context = self.company_context_service.get_company_context(company_short_name)
+
+        # merge context: company + user
+        final_system_context = f"{company_specific_context}\n{rendered_system_prompt}"
+
+        return final_system_context, user_profile
+
+    def prepare_context(self, company_short_name: str, user_identifier: str) -> dict:
+        # prepare the context and decide if it needs to be rebuilt
+        # save the generated context in the session context for later use
+        if not user_identifier:
+            return {'rebuild_needed': True, 'error': 'Invalid user identifier'}
+
+        # create the company/user context and compute its version
+        final_system_context, user_profile = self._build_context_and_profile(
+            company_short_name, user_identifier)
+
+        # save the user information in the session context
+        # it's needed for the jinja predefined prompts (filtering)
+        self.session_context.save_profile_data(company_short_name, user_identifier, user_profile)
+
+        # calculate the context version
+        current_version = self._compute_context_version_from_string(final_system_context)
 
-        logging.info(f"Inicializando contexto para {company_short_name}/{user_identifier} con modelo {model}  ...")
         try:
-            # 1. clean any previous context for company/user
-            self.session_context.clear_all_context(
-                company_short_name=company_short_name,
-                user_identifier=user_identifier
-            )
+            prev_version = self.session_context.get_context_version(company_short_name, user_identifier)
+        except Exception:
+            prev_version = None
 
-            # 2. get dictionary with user information from company DB
-            # user roles are read at this point from company db
-            user_profile = self.dispatcher.get_user_info(
-                company_name=company_short_name,
-                user_identifier=user_identifier,
-                is_local_user=is_local_user
-            )
+        rebuild_is_needed = not (prev_version and prev_version == current_version and
+                                 self._has_valid_cached_context(company_short_name, user_identifier))
 
-            # add the user logged in to the user_info dictionary
-            user_profile['user_id'] = user_identifier
+        if rebuild_is_needed:
+            # Guardar el contexto preparado y su versión para que `finalize_context_rebuild` los use.
+            self.session_context.save_prepared_context(company_short_name,
+                                                       user_identifier,
+                                                       final_system_context,
+                                                       current_version)
 
-            # save the user information in the session context
-            # it's needed for the jinja predefined prompts (filtering)
-            self.session_context.save_user_session_data(company_short_name, user_identifier, user_profile)
+        return {'rebuild_needed': rebuild_is_needed}
 
-            # 3. render the iatoolkit main system prompt with the company/user information
-            system_prompt_template = self.prompt_service.get_system_prompt()
-            rendered_system_prompt = self.util.render_prompt_from_string(
-                template_string=system_prompt_template,
-                question=None,
-                client_data=user_profile,
-                company=company,
-                service_list=self.dispatcher.get_company_services(company)
-            )
+    def set_context_for_llm(self,
+                            company_short_name: str,
+                            user_identifier: str,
+                            model: str = ''):
 
-            # 4. add more company context: schemas, database models, .md files
-            company_specific_context = self.dispatcher.get_company_context(company_name=company_short_name)
+        # This service takes a pre-built context and send to the LLM
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            logging.error(f"Company not found: {company_short_name} in set_context_for_llm")
+            return
 
-            # 5. merge contexts
-            final_system_context = f"{company_specific_context}\n{rendered_system_prompt}"
+        # --- Model Resolution ---
+        # Priority: 1. Explicit model -> 2. Company config -> 3. Global default
+        effective_model = model
+        if not effective_model:
+            llm_config = self.configuration_service.get_configuration(company_short_name, 'llm')
+            if llm_config and llm_config.get('model'):
+                effective_model = llm_config['model']
+
+        effective_model = effective_model or self.default_model
+
+        # blocking logic to avoid multiple requests for the same user/company at the same time
+        lock_key = f"lock:context:{company_short_name}/{user_identifier}"
+        if not self.session_context.acquire_lock(lock_key, expire_seconds=60):
+            logging.warning(
+                f"try to rebuild context for user {user_identifier} while is still in process, ignored.")
+            return
 
-            if self.util.is_gemini_model(model):
-                # save the initial context as `context_history` (list of messages)
-                context_history = [{"role": "user", "content": final_system_context}]
-                self.session_context.save_context_history(company_short_name, user_identifier, context_history)
-                logging.info(f"Contexto inicial para Gemini guardado en sesión")
-                return "gemini-context-initialized"
+        try:
+            start_time = time.time()
+            company = self.profile_repo.get_company_by_short_name(company_short_name)
+
+            # get the prepared context and version from the session cache
+            prepared_context, version_to_save = self.session_context.get_and_clear_prepared_context(company_short_name,
+                                                                                                    user_identifier)
+            if not prepared_context:
+                return
 
-            elif self.util.is_openai_model(model):
+            logging.info(f"sending context to LLM model {effective_model} for: {company_short_name}/{user_identifier}...")
 
-                # 6. set the company/user context as the initial context for the LLM
+            # clean only the chat history and the last response ID for this user/company
+            self.session_context.clear_llm_history(company_short_name, user_identifier)
+
+            response_id = ''
+            if self.util.is_gemini_model(effective_model):
+                context_history = [{"role": "user", "content": prepared_context}]
+                self.session_context.save_context_history(company_short_name, user_identifier, context_history)
+            elif self.util.is_openai_model(effective_model):
+                # Here is the call to the LLM client for settling the company/user context
                 response_id = self.llm_client.set_company_context(
                     company=company,
-                    company_base_context=final_system_context,
-                    model=model
+                    company_base_context=prepared_context,
+                    model=effective_model
                 )
-
-                # 7. save response_id in the session context
                 self.session_context.save_last_response_id(company_short_name, user_identifier, response_id)
 
-                logging.info(f"Contexto inicial de company '{company_short_name}/{user_identifier}' ha sido establecido en {int(time.time() - start_time)} seg.")
-                return response_id
+            if version_to_save:
+                self.session_context.save_context_version(company_short_name, user_identifier, version_to_save)
 
+            logging.info(
+                f"Context for: {company_short_name}/{user_identifier} settled in {int(time.time() - start_time)} sec.")
         except Exception as e:
-            logging.exception(f"Error al inicializar el contexto del LLM para {company_short_name}: {e}")
+            logging.exception(f"Error in finalize_context_rebuild for {company_short_name}: {e}")
             raise e
+        finally:
+            # release the lock
+            self.session_context.release_lock(lock_key)
+
+        return {'response_id': response_id }
 
     def llm_query(self,
                   company_short_name: str,
-                  external_user_id: Optional[str] = None,
-                  local_user_id: int = 0,
+                  user_identifier: str,
                   task: Optional[Task] = None,
                   prompt_name: str = None,
                   question: str = '',
                   client_data: dict = {},
-                  files: list = []) -> dict:
+                  response_id: str = '',
+                  files: list = [],
+                  model: Optional[str] = None) -> dict:
         try:
-            user_identifier, is_local_user = self.util.resolve_user_identifier(external_user_id, local_user_id)
-            if not user_identifier:
-                return {"error": True,
-                        "error_message": "No se pudo identificar al usuario"}
-
             company = self.profile_repo.get_company_by_short_name(short_name=company_short_name)
             if not company:
                 return {"error": True,
-                        "error_message": f'No existe Company ID: {company_short_name}'}
+                        "error_message": self.i18n_service.t('errors.company_not_found', company_short_name=company_short_name)}
 
             if not prompt_name and not question:
                 return {"error": True,
-                        "error_message": f'Hola, cual es tu pregunta?'}
+                        "error_message": self.i18n_service.t('services.start_query')}
+
+            # --- Model Resolution ---
+            # Priority: 1. Explicit model -> 2. Company config -> 3. Global default
+            effective_model = model
+            if not effective_model:
+                llm_config = self.configuration_service.get_configuration(company_short_name, 'llm')
+                if llm_config and llm_config.get('model'):
+                    effective_model = llm_config['model']
+
+            effective_model = effective_model or self.default_model
 
             # get the previous response_id and context history
             previous_response_id = None
             context_history = self.session_context.get_context_history(company.short_name, user_identifier) or []
 
-            if self.util.is_openai_model(self.model):
-                # get user context
-                previous_response_id = self.session_context.get_last_response_id(company.short_name, user_identifier)
-                if not previous_response_id:
-                    # try to initialize the company/user context
-                    previous_response_id = self.llm_init_context(company.short_name, external_user_id, local_user_id)
+            if self.util.is_openai_model(effective_model):
+                if response_id:
+                    # context is getting from this response_id
+                    previous_response_id = response_id
+                else:
+                    # use the full user history context
+                    previous_response_id = self.session_context.get_last_response_id(company.short_name, user_identifier)
                     if not previous_response_id:
                         return {'error': True,
-                                "error_message": f"FATAL: No se encontró 'previous_response_id' para '{company.short_name}/{user_identifier}'. La conversación no puede continuar."
+                                "error_message": self.i18n_service.t('errors.services.missing_response_id', company_short_name=company.short_name, user_identifier=user_identifier)
                                 }
-            elif self.util.is_gemini_model(self.model):
+            elif self.util.is_gemini_model(effective_model):
                 # check the length of the context_history and remove old messages
                 self._trim_context_history(context_history)
 
-            # get the user data from the session context
-            user_info_from_session = self.session_context.get_user_session_data(company.short_name, user_identifier)
+            # get the user profile data from the session context
+            user_profile = self.profile_service.get_profile_by_identifier(company.short_name, user_identifier)
 
-            # Combinar datos: los datos de la tarea/request tienen prioridad sobre los de la sesión
-            final_client_data = (user_info_from_session or {}).copy()
+            # combine client_data with user_profile
+            final_client_data = (user_profile or {}).copy()
             final_client_data.update(client_data)
-            final_client_data['user_id'] = user_identifier
 
             # Load attached files into the context
             files_context = self.load_files_for_context(files)
@@ -207,7 +293,7 @@ class QueryService:
                     template_string=prompt_content,
                     question=question,
                     client_data=final_client_data,
-                    external_user_id=external_user_id,
+                    user_identifier=user_identifier,
                     company=company,
                 )
 
@@ -219,7 +305,7 @@ class QueryService:
                 user_turn_prompt += f'\n### Contexto Adicional: El usuario ha aportado este contexto puede ayudar: {question}'
 
             # add to the history context
-            if self.util.is_gemini_model(self.model):
+            if self.util.is_gemini_model(effective_model):
                 context_history.append({"role": "user", "content": user_turn_prompt})
 
             # service list for the function calls
@@ -232,8 +318,9 @@ class QueryService:
             response = self.llm_client.invoke(
                 company=company,
                 user_identifier=user_identifier,
+                model=effective_model,
                 previous_response_id=previous_response_id,
-                context_history=context_history if self.util.is_gemini_model(self.model) else None,
+                context_history=context_history if self.util.is_gemini_model(effective_model) else None,
                 question=question,
                 context=user_turn_prompt,
                 tools=tools,
@@ -246,7 +333,7 @@ class QueryService:
             # save last_response_id for the history chain
             if "response_id" in response:
                 self.session_context.save_last_response_id(company.short_name, user_identifier, response["response_id"])
-            if self.util.is_gemini_model(self.model):
+            if self.util.is_gemini_model(effective_model):
                 self.session_context.save_context_history(company.short_name, user_identifier, context_history)
 
             return response
@@ -254,6 +341,31 @@ class QueryService:
             logging.exception(e)
             return {'error': True, "error_message": f"{str(e)}"}
 
+    def _compute_context_version_from_string(self, final_system_context: str) -> str:
+        # returns a hash of the context string
+        try:
+            return hashlib.sha256(final_system_context.encode("utf-8")).hexdigest()
+        except Exception:
+            return "unknown"
+
+    def _has_valid_cached_context(self, company_short_name: str, user_identifier: str) -> bool:
+        """
+        Verifica si existe un estado de contexto reutilizable en sesión.
+        - OpenAI: last_response_id presente.
+        - Gemini: context_history con al menos 1 mensaje.
+        """
+        try:
+            if self.util.is_openai_model(self.default_model):
+                prev_id = self.session_context.get_last_response_id(company_short_name, user_identifier)
+                return bool(prev_id)
+            if self.util.is_gemini_model(self.default_model):
+                history = self.session_context.get_context_history(company_short_name, user_identifier) or []
+                return len(history) >= 1
+            return False
+        except Exception as e:
+            logging.warning(f"error verifying context cache: {e}")
+            return False
+
     def load_files_for_context(self, files: list) -> str:
         """
         Processes a list of attached files, decodes their content,
@@ -310,7 +422,7 @@ class QueryService:
         try:
             total_tokens = sum(self.llm_client.count_tokens(json.dumps(message)) for message in context_history)
         except Exception as e:
-            logging.error(f"Error al calcular tokens del historial: {e}. No se pudo recortar el contexto.")
+            logging.error(f"error counting tokens for history: {e}.")
             return
 
         # Si se excede el límite, eliminar mensajes antiguos (empezando por el segundo)
@@ -321,8 +433,8 @@ class QueryService:
                 removed_tokens = self.llm_client.count_tokens(json.dumps(removed_message))
                 total_tokens -= removed_tokens
                 logging.warning(
-                    f"Historial de contexto ({total_tokens + removed_tokens} tokens) excedía el límite de {GEMINI_MAX_TOKENS_CONTEXT_HISTORY}. "
-                    f"Nuevo total: {total_tokens} tokens."
+                    f"history tokens ({total_tokens + removed_tokens} tokens) exceed the limit of: {GEMINI_MAX_TOKENS_CONTEXT_HISTORY}. "
+                    f"new context: {total_tokens} tokens."
                 )
             except IndexError:
                 # Se produce si solo queda el mensaje del sistema, el bucle debería detenerse.

iatoolkit/services/search_service.py

@@ -5,19 +5,22 @@
 
 from iatoolkit.repositories.vs_repo import VSRepo
 from iatoolkit.repositories.document_repo import DocumentRepo
+from iatoolkit.repositories.profile_repo import ProfileRepo
+from iatoolkit.repositories.models import Company
 from injector import inject
 
 
 class SearchService:
     @inject
     def __init__(self,
+                 profile_repo: ProfileRepo,
                  doc_repo: DocumentRepo,
                  vs_repo: VSRepo):
-        super().__init__()
+        self.profile_repo = profile_repo
         self.vs_repo = vs_repo
         self.doc_repo = doc_repo
 
-    def search(self, company_id:  int, query: str, metadata_filter: dict = None) -> str:
+    def search(self, company_short_name: str, query: str, metadata_filter: dict = None) -> str:
         """
         Performs a semantic search for a given query within a company's documents.
 
@@ -26,7 +29,7 @@ class SearchService:
         content of the retrieved documents, which can be used as context for an LLM.
 
         Args:
-            company_id: The ID of the company to search within.
+            company_short_name: The  company to search within.
             query: The text query to search for.
             metadata_filter: An optional dictionary to filter documents by their metadata.
 
@@ -34,7 +37,11 @@ class SearchService:
             A string containing the concatenated content of the found documents,
             formatted to be used as a context.
         """
-        document_list = self.vs_repo.query(company_id=company_id,
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            return f"error: company {company_short_name} not found"
+
+        document_list = self.vs_repo.query(company_short_name=company_short_name,
                                            query_text=query,
                                            metadata_filter=metadata_filter)
 

iatoolkit/services/sql_service.py

@@ -4,57 +4,89 @@
 # IAToolkit is open source software.
 
 from iatoolkit.repositories.database_manager import DatabaseManager
-
 from iatoolkit.common.util import Utility
+from iatoolkit.services.i18n_service import I18nService
+from iatoolkit.common.exceptions import IAToolkitException
 from sqlalchemy import text
-from injector import inject
+from injector import inject, singleton
 import json
-from iatoolkit.common.exceptions import IAToolkitException
+import logging
 
 
+@singleton
 class SqlService:
+    """
+    Manages database connections and executes SQL statements.
+    It maintains a cache of named DatabaseManager instances to avoid reconnecting.
+    """
+
     @inject
-    def __init__(self,util: Utility):
+    def __init__(self,
+                 util: Utility,
+                 i18n_service: I18nService):
         self.util = util
+        self.i18n_service = i18n_service
+
+        # Cache for database connections
+        self._db_connections: dict[str, DatabaseManager] = {}
 
-    def exec_sql(self, db_manager: DatabaseManager, sql_statement: str) -> str:
+    def register_database(self, db_name: str, db_uri: str):
+        """
+        Creates and caches a DatabaseManager instance for a given database name and URI.
+        If a database with the same name is already registered, it does nothing.
         """
-        Executes a raw SQL statement and returns the result as a JSON string.
+        if db_name in self._db_connections:
+            return
 
-        This method takes a DatabaseManager instance and a SQL query, executes it
-        against the database, and fetches all results. The results are converted
-        into a list of dictionaries, where each dictionary represents a row.
-        This list is then serialized to a JSON string.
-        If an exception occurs during execution, the transaction is rolled back,
-        and a custom IAToolkitException is raised.
+        logging.debug(f"Registering and creating connection for database: '{db_name}'")
 
-        Args:
-            db_manager: The DatabaseManager instance to get the database session from.
-            sql_statement: The raw SQL statement to be executed.
+        # create the database connection and save it on the cache
+        db_manager = DatabaseManager(db_uri, register_pgvector=False)
+        self._db_connections[db_name] = db_manager
 
-        Returns:
-            A JSON string representing the list of rows returned by the query.
+    def get_database_manager(self, db_name: str) -> DatabaseManager:
+        """
+        Retrieves a registered DatabaseManager instance from the cache.
         """
         try:
-            # here the SQL is executed
-            result = db_manager.get_session().execute(text(sql_statement))
+            return self._db_connections[db_name]
+        except KeyError:
+            logging.error(f"Attempted to access unregistered database: '{db_name}'")
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.DATABASE_ERROR,
+                f"Database '{db_name}' is not registered with the SqlService."
+            )
 
-            # get the column names
-            cols = result.keys()
+    def exec_sql(self, database: str, query: str) -> str:
+        """
+        Executes a raw SQL statement against a registered database and returns the result as a JSON string.
+        """
+        try:
+            # 1. Get the database manager from the cache
+            db_manager = self.get_database_manager(database)
 
-            # convert rows to dict
+            # 2. Execute the SQL statement
+            result = db_manager.get_session().execute(text(query))
+            cols = result.keys()
             rows_context = [dict(zip(cols, row)) for row in result.fetchall()]
 
-            # Serialize to JSON with type convertion
+            # seialize the result
             sql_result_json = json.dumps(rows_context, default=self.util.serialize)
 
             return sql_result_json
+        except IAToolkitException:
+            # Re-raise exceptions from get_database_manager to preserve the specific error
+            raise
         except Exception as e:
-            db_manager.get_session().rollback()
+            # Attempt to rollback if a session was active
+            db_manager = self._db_connections.get(database)
+            if db_manager:
+                db_manager.get_session().rollback()
 
             error_message = str(e)
             if 'timed out' in str(e):
-                error_message = 'Intentalo de nuevo, se agoto el tiempo de espera'
+                error_message = self.i18n_service.t('errors.timeout')
 
+            logging.error(f"Error executing SQL statement: {error_message}")
             raise IAToolkitException(IAToolkitException.ErrorType.DATABASE_ERROR,
                                      error_message) from e

iatoolkit/services/tasks_service.py

@@ -101,7 +101,7 @@ class TaskService:
         # call the IA
         response = self.query_service.llm_query(
             task=task,
-            local_user_id=0,
+            user_identifier='task-monitor',
             company_short_name=task.company.short_name,
             prompt_name=task.task_type.name,
             client_data=task.client_data,