iatoolkit 0.63.1__py3-none-any.whl → 0.69.0__py3-none-any.whl

iatoolkit/services/profile_service.py

@@ -5,7 +5,7 @@
 
 from injector import inject
 from iatoolkit.repositories.profile_repo import ProfileRepo
-from iatoolkit.services.dispatcher_service import Dispatcher
+from iatoolkit.services.i18n_service import I18nService
 from iatoolkit.repositories.models import User, Company, ApiKey
 from flask_bcrypt import check_password_hash
 from iatoolkit.common.session_manager import SessionManager
@@ -16,16 +16,19 @@ import random
 import re
 import secrets
 import string
+import logging
 from iatoolkit.services.dispatcher_service import Dispatcher
 
 
 class ProfileService:
     @inject
     def __init__(self,
+                 i18n_service: I18nService,
                  profile_repo: ProfileRepo,
                  session_context_service: UserSessionContextService,
                  dispatcher: Dispatcher,
                  mail_app: MailApp):
+        self.i18n_service = i18n_service
         self.profile_repo = profile_repo
         self.dispatcher = dispatcher
         self.session_context = session_context_service
@@ -38,23 +41,23 @@ class ProfileService:
             # check if user exists
             user = self.profile_repo.get_user_by_email(email)
             if not user:
-                return {'success': False, "message": "Usuario no encontrado"}
+                return {'success': False, 'message': self.i18n_service.t('errors.auth.user_not_found')}
 
             # check the encrypted password
             if not check_password_hash(user.password, password):
-                return {'success': False, "message": "Contraseña inválida"}
+                return {'success': False, 'message': self.i18n_service.t('errors.auth.invalid_password')}
 
             company = self.profile_repo.get_company_by_short_name(company_short_name)
             if not company:
-                return {'success': False, "message": "Empresa no encontrada"}
+                return {'success': False, "message": "missing company"}
 
             # check that user belongs to company
             if company not in user.companies:
-                return {'success': False, "message": "Usuario no esta autorizado para esta empresa"}
+                return {'success': False, "message": self.i18n_service.t('errors.services.user_not_authorized')}
 
             if not user.verified:
                 return {'success': False,
-                        "message": "Tu cuenta no ha sido verificada. Por favor, revisa tu correo."}
+                        "message": self.i18n_service.t('errors.services.account_not_verified')}
 
             # 1. Build the local user profile dictionary here.
             # the user_profile variables are used on the LLM templates also (see in query_main.prompt)
@@ -71,8 +74,9 @@ class ProfileService:
 
             # 3. create the web session
             self.set_session_for_user(company.short_name, user_identifier)
-            return {'success': True, "user_identifier": user_identifier, "message": "Login exitoso"}
+            return {'success': True, "user_identifier": user_identifier, "message": "Login ok"}
         except Exception as e:
+            logging.error(f"Error in login: {e}")
             return {'success': False, "message": str(e)}
 
     def create_external_user_profile_context(self, company: Company, user_identifier: str):
@@ -91,6 +95,9 @@ class ProfileService:
             user_identifier=user_identifier,
             user_profile=external_user_profile)
 
+        # 3. make sure the flask session is clean
+        SessionManager.clear()
+
     def save_user_profile(self, company: Company, user_identifier: str, user_profile: dict):
         """
         Private helper: Takes a pre-built profile, saves it to Redis, and sets the Flask cookie.
@@ -131,6 +138,24 @@ class ProfileService:
             "profile": profile
         }
 
+    def update_user_language(self, user_identifier: str, new_lang: str) -> dict:
+        """
+        Business logic to update a user's preferred language.
+        It validates the language and then calls the generic update method.
+        """
+        # 1. Validate that the language is supported by checking the loaded translations.
+        if new_lang not in self.i18n_service.translations:
+            return {'success': False, 'error_message': self.i18n_service.t('errors.general.unsupported_language')}
+
+        try:
+            # 2. Call the generic update_user method, passing the specific field to update.
+            self.update_user(user_identifier, preferred_language=new_lang)
+            return {'success': True, 'message': 'Language updated successfully.'}
+        except Exception as e:
+            # Log the error and return a generic failure message.
+            logging.error(f"Failed to update language for {user_identifier}: {e}")
+            return {'success': False, 'error_message': self.i18n_service.t('errors.general.unexpected_error', error=str(e))}
+
 
     def get_profile_by_identifier(self, company_short_name: str, user_identifier: str) -> dict:
         """
@@ -155,7 +180,8 @@ class ProfileService:
             # get company info
             company = self.profile_repo.get_company_by_short_name(company_short_name)
             if not company:
-                return {"error": f"la empresa {company_short_name} no existe"}
+                return {
+                    "error": self.i18n_service.t('errors.signup.company_not_found', company_name=company_short_name)}
 
             # normalize  format's
             email = email.lower()
@@ -165,24 +191,25 @@ class ProfileService:
             if existing_user:
                 # validate password
                 if not self.bcrypt.check_password_hash(existing_user.password, password):
-                    return {"error": "La contraseña es incorrecta. No se puede agregar a la nueva empresa."}
+                    return {"error": self.i18n_service.t('errors.signup.incorrect_password_for_existing_user', email=email)}
 
                 # check if register
                 if company in existing_user.companies:
-                    return {"error": "Usuario ya registrado en esta empresa"}
+                    return {"error": self.i18n_service.t('errors.signup.user_already_registered', email=email)}
                 else:
                     # add new company to existing user
                     existing_user.companies.append(company)
                     self.profile_repo.save_user(existing_user)
-                    return {"message": "Usuario asociado a nueva empresa"}
+                    return {"message": self.i18n_service.t('flash_messages.user_associated_success')}
 
             # add the new user
             if password != confirm_password:
-                return {"error": "Las contraseñas no coinciden. Por favor, inténtalo de nuevo."}
+                return {"error": self.i18n_service.t('errors.signup.password_mismatch')}
 
             is_valid, message = self.validate_password(password)
             if not is_valid:
-                return {"error": message}
+                # Translate the key returned by validate_password
+                return {"error": self.i18n_service.t(message)}
 
             # encrypt the password
             hashed_password = self.bcrypt.generate_password_hash(password).decode('utf-8')
@@ -204,9 +231,9 @@ class ProfileService:
             # send email with verification
             self.send_verification_email(new_user, company_short_name)
 
-            return {"message": "Registro exitoso. Por favor, revisa tu correo para verificar tu cuenta."}
+            return {"message": self.i18n_service.t('flash_messages.signup_success')}
         except Exception as e:
-            return {"error": str(e)}
+            return {"error": self.i18n_service.t('errors.general.unexpected_error', error=str(e))}
 
     def update_user(self, email: str, **kwargs) -> User:
         return self.profile_repo.update_user(email, **kwargs)
@@ -216,14 +243,14 @@ class ProfileService:
             # check if user exist
             user = self.profile_repo.get_user_by_email(email)
             if not user:
-                return {"error": "El usuario no existe."}
+                return {"error": self.i18n_service.t('errors.verification.user_not_found')}
 
             # activate the user account
             self.profile_repo.verify_user(email)
-            return {"message": "Tu cuenta ha sido verificada exitosamente. Ahora puedes iniciar sesión."}
+            return {"message": self.i18n_service.t('flash_messages.account_verified_success')}
 
         except Exception as e:
-            return {"error": str(e)}
+            return {"error": self.i18n_service.t('errors.general.unexpected_error')}
 
     def change_password(self,
                          email: str,
@@ -232,28 +259,28 @@ class ProfileService:
                          confirm_password: str):
         try:
             if new_password != confirm_password:
-                return {"error": "Las contraseñas no coinciden. Por favor, inténtalo nuevamente."}
+                return {"error": self.i18n_service.t('errors.change_password.password_mismatch')}
 
             # check the temporary code
             user = self.profile_repo.get_user_by_email(email)
             if not user or user.temp_code != temp_code:
-                return {"error": "El código temporal no es válido. Por favor, verifica o solicita uno nuevo."}
+                return {"error": self.i18n_service.t('errors.change_password.invalid_temp_code')}
 
             # encrypt and save the password, make the temporary code invalid
             hashed_password = self.bcrypt.generate_password_hash(new_password).decode('utf-8')
             self.profile_repo.update_password(email, hashed_password)
             self.profile_repo.reset_temp_code(email)
 
-            return {"message": "La clave se cambio correctamente"}
+            return {"message": self.i18n_service.t('flash_messages.password_changed_success')}
         except Exception as e:
-            return {"error": str(e)}
+            return {"error": self.i18n_service.t('errors.general.unexpected_error')}
 
     def forgot_password(self, email: str, reset_url: str):
         try:
             # Verificar si el usuario existe
             user = self.profile_repo.get_user_by_email(email)
             if not user:
-                return {"error": "El usuario no existe."}
+                return {"error": self.i18n_service.t('errors.forgot_password.user_not_registered', email=email)}
 
             # Gen a temporary code and store in the repositories
             temp_code = ''.join(random.choices(string.ascii_letters + string.digits, k=6)).upper()
@@ -262,35 +289,31 @@ class ProfileService:
             # send email to the user
             self.send_forgot_password_email(user, reset_url)
 
-            return {"message": "se envio mail para cambio de clave"}
+            return {"message": self.i18n_service.t('flash_messages.forgot_password_success')}
         except Exception as e:
-            return {"error": str(e)}
+            return {"error": self.i18n_service.t('errors.general.unexpected_error')}
 
     def validate_password(self, password):
         """
-        Valida que una contraseña cumpla con los siguientes requisitos:
-        - Al menos 8 caracteres de longitud
-        - Contiene al menos una letra mayúscula
-        - Contiene al menos una letra minúscula
-        - Contiene al menos un número
-        - Contiene al menos un carácter especial
+        Validates that a password meets all requirements.
+        Returns (True, "...") on success, or (False, "translation.key") on failure.
         """
         if len(password) < 8:
-            return False, "La contraseña debe tener al menos 8 caracteres."
+            return False, "errors.validation.password_too_short"
 
         if not any(char.isupper() for char in password):
-            return False, "La contraseña debe tener al menos una letra mayúscula."
+            return False, "errors.validation.password_no_uppercase"
 
         if not any(char.islower() for char in password):
-            return False, "La contraseña debe tener al menos una letra minúscula."
+            return False, "errors.validation.password_no_lowercase"
 
         if not any(char.isdigit() for char in password):
-            return False, "La contraseña debe tener al menos un número."
+            return False, "errors.validation.password_no_digit"
 
         if not re.search(r'[!@#$%^&*(),.?":{}|<>]', password):
-            return False, "La contraseña debe tener al menos un carácter especial."
+            return False, "errors.validation.password_no_special_char"
 
-        return True, "La contraseña es válida."
+        return True, "Password is valid."
 
     def get_companies(self):
         return self.profile_repo.get_companies()
@@ -304,7 +327,7 @@ class ProfileService:
     def new_api_key(self, company_short_name: str):
         company = self.get_company_by_short_name(company_short_name)
         if not company:
-            return {"error": f"la empresa {company_short_name} no existe"}
+            return {"error": self.i18n_service.t('errors.company_not_found', company_short_name=company_short_name)}
 
         length = 40     # lenght of the api key
         alphabet = string.ascii_letters + string.digits

iatoolkit/services/prompt_manager_service.py

@@ -5,21 +5,25 @@
 
 from injector import inject
 from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
-
-import logging
+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
 import os
 from iatoolkit.common.exceptions import IAToolkitException
 import importlib.resources
+import logging
 
 
 class PromptService:
     @inject
-    def __init__(self, llm_query_repo: LLMQueryRepo, profile_repo: ProfileRepo):
+    def __init__(self,
+                 llm_query_repo: LLMQueryRepo,
+                 profile_repo: ProfileRepo,
+                 i18n_service: I18nService):
         self.llm_query_repo = llm_query_repo
         self.profile_repo = profile_repo
+        self.i18n_service = i18n_service
 
     def create_prompt(self,
                       prompt_name: str,
@@ -36,20 +40,20 @@ class PromptService:
         if is_system_prompt:
             if not importlib.resources.files('iatoolkit.system_prompts').joinpath(prompt_filename).is_file():
                 raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                                f'No existe el archivo de prompt de sistemas: {prompt_filename}')
+                                f'missing system prompt file: {prompt_filename}')
         else:
             template_dir = f'companies/{company.short_name}/prompts'
 
             relative_prompt_path = os.path.join(template_dir, prompt_filename)
             if not os.path.exists(relative_prompt_path):
                 raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                               f'No existe el archivo de prompt: {relative_prompt_path}')
+                               f'missing prompt file: {relative_prompt_path}')
 
         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'El campo custom_fields debe contener los campos: data_key y label')
+                               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:
@@ -82,20 +86,20 @@ class PromptService:
             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"No se encontró el prompt '{prompt_name}' para la empresa '{company.short_name}'")
+                                   f"prompt not found '{prompt_name}' for company '{company.short_name}'")
 
             prompt_file = f'companies/{company.short_name}/prompts/{user_prompt.filename}'
             absolute_filepath = os.path.join(execution_dir, prompt_file)
             if not os.path.exists(absolute_filepath):
                 raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
-                                   f"El archivo para el prompt '{prompt_name}' no existe: {absolute_filepath}")
+                                   f"prompt file '{prompt_name}' does not exist: {absolute_filepath}")
 
             try:
                 with open(absolute_filepath, 'r', encoding='utf-8') as f:
                     user_prompt_content = f.read()
             except Exception as e:
                 raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
-                                   f"Error leyendo el archivo de prompt '{prompt_name}' en {absolute_filepath}: {e}")
+                                   f"error while reading prompt: '{prompt_name}' in this pathname {absolute_filepath}: {e}")
 
             return user_prompt_content
 
@@ -105,9 +109,9 @@ class PromptService:
             raise
         except Exception as e:
             logging.exception(
-                f"Error al obtener el contenido del prompt para la empresa '{company.short_name}' y prompt '{prompt_name}': {e}")
+                f"error loading prompt '{prompt_name}' content for '{company.short_name}': {e}")
             raise IAToolkitException(IAToolkitException.ErrorType.PROMPT_ERROR,
-                               f'Error al obtener el contenido del prompt "{prompt_name}" para la empresa {company.short_name}: {str(e)}')
+                               f'error loading prompt "{prompt_name}" content for company {company.short_name}: {str(e)}')
 
     def get_system_prompt(self):
         try:
@@ -121,10 +125,10 @@ class PromptService:
                     content = importlib.resources.read_text('iatoolkit.system_prompts', prompt.filename)
                     system_prompt_content.append(content)
                 except FileNotFoundError:
-                    logging.warning(f"El archivo para el prompt de sistema no existe en el paquete: {prompt.filename}")
+                    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 leyendo el archivo de prompt del sistema '{prompt.filename}': {e}")
+                                             f"error reading system prompt '{prompt.filename}': {e}")
 
             # join the system prompts into a single string
             return "\n".join(system_prompt_content)
@@ -135,14 +139,14 @@ class PromptService:
             logging.exception(
                 f"Error al obtener el contenido del prompt de sistema: {e}")
             raise IAToolkitException(IAToolkitException.ErrorType.PROMPT_ERROR,
-                               f'Error al obtener el contenido de los prompts de sistema": {str(e)}')
+                               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': f'No existe la empresa: {company_short_name}'}
+                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)
@@ -183,6 +187,6 @@ class PromptService:
             return {'message': categorized_prompts}
 
         except Exception as e:
-            logging.error(f"Error en get_prompts: {e}")
+            logging.error(f"error in get_prompts: {e}")
             return {'error': str(e)}
 

iatoolkit/services/query_service.py

@@ -8,6 +8,8 @@ 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.repositories.llm_query_repo import LLMQueryRepo
 from iatoolkit.repositories.models import Task
 from iatoolkit.services.dispatcher_service import Dispatcher
@@ -32,21 +34,25 @@ class QueryService:
     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
                  ):
         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
@@ -56,7 +62,7 @@ class QueryService:
         self.model = os.getenv("LLM_MODEL", "")
         if not self.model:
             raise IAToolkitException(IAToolkitException.ErrorType.API_KEY,
-                               "La variable de entorno 'LLM_MODEL' no está configurada.")
+                               "missing ENV variable 'LLM_MODEL' configuration.")
 
     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
@@ -78,7 +84,7 @@ class QueryService:
         )
 
         # get the company context: schemas, database models, .md files
-        company_specific_context = self.dispatcher.get_company_context(company_name=company_short_name)
+        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}"
@@ -127,7 +133,7 @@ class QueryService:
         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"Intento de reconstruir contexto para {user_identifier} mientras ya estaba en progreso. Se omite.")
+                f"try to rebuild context for user {user_identifier} while is still in process, ignored.")
             return
 
         try:
@@ -138,11 +144,9 @@ class QueryService:
             prepared_context, version_to_save = self.session_context.get_and_clear_prepared_context(company_short_name,
                                                                                                     user_identifier)
             if not prepared_context:
-                logging.info(
-                    f"No se requiere reconstrucción de contexto para {company_short_name}/{user_identifier}. Finalización rápida.")
                 return
 
-            logging.info(f"Enviando contexto al LLM para {company_short_name}/{user_identifier}...")
+            logging.info(f"sending context to LLM for: {company_short_name}/{user_identifier}...")
 
             # Limpiar solo el historial de chat y el ID de respuesta anterior
             self.session_context.clear_llm_history(company_short_name, user_identifier)
@@ -161,9 +165,9 @@ class QueryService:
                 self.session_context.save_context_version(company_short_name, user_identifier, version_to_save)
 
             logging.info(
-                f"Contexto de {company_short_name}/{user_identifier} establecido en {int(time.time() - start_time)} seg.")
+                f"Context for: {company_short_name}/{user_identifier} settled in {int(time.time() - start_time)} sec.")
         except Exception as e:
-            logging.exception(f"Error en finalize_context_rebuild para {company_short_name}: {e}")
+            logging.exception(f"Error in finalize_context_rebuild for {company_short_name}: {e}")
             raise e
         finally:
             # --- Liberar el Bloqueo ---
@@ -181,11 +185,11 @@ class QueryService:
             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')}
 
             # get the previous response_id and context history
             previous_response_id = None
@@ -196,7 +200,7 @@ class QueryService:
                 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"No se encontró 'previous_response_id' para '{company.short_name}/{user_identifier}'. Reinicia el contexto para 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):
                 # check the length of the context_history and remove old messages
@@ -294,7 +298,7 @@ class QueryService:
                 return len(history) >= 1
             return False
         except Exception as e:
-            logging.warning(f"Error verificando caché de contexto: {e}")
+            logging.warning(f"error verifying context cache: {e}")
             return False
 
     def load_files_for_context(self, files: list) -> str:
@@ -353,7 +357,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)
@@ -364,8 +368,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_id=company.id,
                                            query_text=query,
                                            metadata_filter=metadata_filter)