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

iatoolkit/locales/en.yaml

@@ -0,0 +1,167 @@
+# Language: English
+ui:
+  login_widget:
+    title: "Sign In"
+    welcome_message: "Enter your credentials or register to access Sample Company’s AI platform."
+    email_placeholder: "Email address"
+    password_placeholder: "Password"
+    login_button: "Login"
+    forgot_password_link: "Forgot your password?"
+    no_account_prompt: "Don't have an account?"
+    signup_link: "Sign Up"
+
+  signup:
+    title: "Create an Account"
+    subtitle: "Start your journey with us today."
+    first_name_label: "First Name"
+    last_name_label: "Last Name"
+    email_label: "Email Address"
+    password_label: "Password"
+    confirm_password_label: "Confirm Password"
+    signup_button: "Create Account"
+    already_have_account: "Already have an account?"
+    login_link: "Log In"
+    disclaimer: "🔒We value your privacy. Your data will be used exclusively for the operation of the platform."
+
+  forgot_password:
+    title: "Recover Password"
+    subtitle: "Enter your email address and we will send you a link to reset your password."
+    submit_button: "Send Recovery Link"
+    back_to_login: "Back to Login"
+
+  change_password:
+    title: "Create New Password"
+    subtitle: "You are changing the password for <strong>{email}</strong>."
+    temp_code_label: "Temporary Code"
+    temp_code_placeholder: "Check your email"
+    new_password_label: "New Password"
+    password_instructions: "Must contain at least 8 characters, an uppercase letter, a lowercase letter, a number, and a special character."
+    confirm_password_label: "Confirm New Password"
+    save_button: "Save Password"
+    back_to_home: "Back to home"
+
+  chat:
+    welcome_message: "Hello! How can I help you today?"
+    input_placeholder: "Type your query here..."
+    prompts_available: "Available prompts"
+
+  tooltips:
+    history: "History of my queries"
+    reload_context: "Force Context Reload"
+    feedback: "Your feedback is very important"
+    usage_guide: "Usage Guide"
+    onboarding: "How to ask better questions"
+    logout: "Log out"
+    use_prompt_assistant: "Use Prompt Assistant"
+    attach_files: "Attach files"
+    view_attached_files: "View attached files"
+    send: "Send"
+    stop: "Stop"
+
+  modals:
+    files_title: "Uploaded Files"
+    feedback_title: "Your Opinion is Important"
+    feedback_prompt: "How useful was the assistant's response?"
+    feedback_comment_label: "Your feedback helps us improve:"
+    feedback_comment_placeholder: "Write your opinion, suggestions, or comments here..."
+    history_title: "Query History"
+    history_table_date: "Date"
+    history_table_query: "Query"
+    loading_history: "Loading history..."
+    no_history_found: "No query history found."
+    help_title: "AI Assistant User Guide"
+
+  buttons:
+    cancel: "Close"
+    send: "Send"
+    stop: "Stop"
+
+errors:
+  company_not_found: "The company {company_short_name} does not exist."
+  timeout: "timeout expired."
+  auth:
+    invalid_password: "The provided password is incorrect."
+    user_not_found: "A user with that email address was not found."
+    invalid_or_expired_token: "Invalid or expired token."
+    session_creation_failed: "Could not create user session."
+    authentication_required: "Authentication required. No session cookie or API Key provided."
+    invalid_api_key: "Invalid or inactive API Key."
+    no_user_identifier_api: "No user_identifier provided for API call."
+  templates:
+    company_not_found: "Company not found."
+    home_template_not_found: "The home page template for the company '{company_short_name}' is not configured."
+    template_not_found: "Template not found: '{template_name}'."
+
+    processing_error: "An error occurred while processing the custom home page template: {error}"
+  general:
+    unexpected_error: "An unexpected error has occurred. Please contact support."
+    unsupported_language: "The selected language is not supported."
+  signup:
+    company_not_found: "The company {company_short_name} does not exist."
+    incorrect_password_for_existing_user: "The password for the user {email} is incorrect."
+    user_already_registered: "The user with email '{email}' already exists in this company."
+    password_mismatch: "The passwords do not match. Please try again."
+  change_password:
+    token_expired: "The password change link has expired. Please request a new one."
+    password_mismatch: "The passwords do not match. Please try again."
+    invalid_temp_code: "The temporary code is not valid. Please check it or request a new one."
+  forgot_password:
+    user_not_registered: "The user with email {email} is not registered."
+  verification:
+    token_expired: "The verification link has expired. Please contact support if you need a new one."
+    user_not_found: "The user you are trying to verify does not exist."
+  validation:
+    password_too_short: "Password must be at least 8 characters long."
+    password_no_uppercase: "Password must contain at least one uppercase letter."
+    password_no_lowercase: "Password must contain at least one lowercase letter."
+    password_no_digit: "Password must contain at least one number."
+    password_no_special_char: "Password must contain at least one special character."
+
+  services:
+    no_text_file: "The file is not text or the encoding is not UTF-8"
+    no_output_file: "Missing output file name"
+    no_data_for_excel: "Missing data or it is not a list of dictionaries"
+    no_download_directory: "Temporary directory for saving Excel files is not configured"
+    cannot_create_excel: "Could not create the Excel file"
+    invalid_filename: "Invalid filename"
+    file_not_exist: "File not found"
+    path_is_not_a_file: "The path does not correspond to a file"
+    file_validation_error: "Error validating file"
+    user_not_authorized: "user is not authorized for this company"
+    account_not_verified: "Your account has not been verified. Please check your email."
+    missing_response_id: "Can not found 'previous_response_id' for '{company_short_name}/{user_identifier}'. Reinit context"
+
+
+api_responses:
+  context_reloaded_success: "The context has been successfully reloaded."
+
+services:
+  mail_sent: "Email sent successfully."
+  start_query: "Hello, what can I help you with today?"
+
+
+flash_messages:
+  password_changed_success: "Your password has been successfully reset. You can now log in."
+  login_required: "Please log in to continue."
+  signup_success: "Registration successful. Please check your email to verify your account."
+  user_associated_success: "Existing user successfully associated with the new company."
+  account_verified_success: "Your account has been successfully verified. Welcome!"
+  forgot_password_success: "If your email is registered, you will receive a link to reset your password."
+
+# Keys specifically for JavaScript
+js_messages:
+  feedback_sent_success_title: "Feedback Sent"
+  feedback_sent_success_body: "Thank you for your feedback!"
+  feedback_sent_error: "Could not send feedback, please try again."
+  feedback_rating_error: "Please rate the assistant using the stars."
+  feedback_comment_error: "Please write your comment before sending."
+  context_reloaded: "Context has been reloaded."
+  error_loading_history: "An error occurred while loading the history."
+  request_aborted: "The response generation has been stopped."
+  processing_error: "An error occurred while processing the request."
+  server_comm_error: "Server communication error ({status}). Please try again later."
+  network_error: "A network error occurred. Please try again in a few moments."
+  unknown_server_error: "Unknown server error."
+  loading: "Loading..."
+  reload_init: "init reloading context in background..."
+  no_history_found: "No query history found."

iatoolkit/locales/es.yaml

@@ -0,0 +1,163 @@
+    # locales/es.yaml
+    ui:
+      login_widget:
+        title: "Iniciar Sesión"
+        welcome_message: "Ingresa tus credenciales o registrate para acceder a la plataforma IA de Sample Company."
+        email_placeholder: "Correo electrónico"
+        login_button: "Acceder"
+        forgot_password_link: "¿Olvidaste tu contraseña?"
+        no_account_prompt: "¿No tienes una cuenta?"
+        signup_link: "Regístrate"
+
+      signup:
+        title: "Crear una Cuenta"
+        subtitle: "Comienza tu viaje con nosotros hoy."
+        first_name_label: "Nombre"
+        last_name_label: "Apellido"
+        email_label: "Correo Electrónico"
+        password_label: "Contraseña"
+        confirm_password_label: "Confirmar Contraseña"
+        signup_button: "Crear Cuenta"
+        already_have_account: "¿Ya tienes una cuenta?"
+        login_link: "Inicia Sesión"
+        disclaimer: "🔒 Valoramos tu privacidad. Tus datos se usarán exclusivamente para el funcionamiento de la plataforma."
+
+      forgot_password:
+        title: "Recuperar Contraseña"
+        subtitle: "Ingresa tu correo electrónico y te enviaremos un enlace para restablecer tu contraseña."
+        submit_button: "Enviar Enlace de Recuperación"
+        back_to_login: "Volver a Iniciar Sesión"
+
+      change_password:
+        title: "Crear Nueva Contraseña"
+        subtitle: "Estás cambiando la contraseña para <strong>{email}</strong>."
+        temp_code_label: "Código Temporal"
+        temp_code_placeholder: "Revisa tu correo electrónico"
+        new_password_label: "Nueva Contraseña"
+        password_instructions: "Debe contener al menos 8 caracteres, mayúscula, minúscula, número y un carácter especial."
+        confirm_password_label: "Confirmar Nueva Contraseña"
+        save_button: "Guardar Contraseña"
+        back_to_home: "Volver al inicio"
+
+      chat:
+        welcome_message: "¡Hola! ¿En qué te puedo ayudar hoy?"
+        input_placeholder: "Escribe tu consulta aquí..."
+        prompts_available: "Prompts disponibles"
+
+      tooltips:
+        history: "Historial con mis consultas"
+        reload_context: "Forzar Recarga de Contexto"
+        feedback: "Tu feedback es muy importante"
+        usage_guide: "Guía de Uso"
+        onboarding: "Cómo preguntar mejor"
+        logout: "Cerrar sesión"
+        use_prompt_assistant: "Usar Asistente de Prompts"
+        attach_files: "Adjuntar archivos"
+        view_attached_files: "Ver archivos adjuntos"
+      modals:
+        files_title: "Archivos Cargados"
+        feedback_title: "Tu Opinión es Importante"
+        feedback_prompt: "¿Qué tan útil fue la respuesta del asistente?"
+        feedback_comment_label: "Tu comentario nos ayuda a mejorar:"
+        feedback_comment_placeholder: "Escribe aquí tu opinión, sugerencias o comentarios..."
+        history_title: "Historial de Consultas"
+        history_table_date: "Fecha"
+        history_table_query: "Consulta"
+        loading_history: "Cargando historial..."
+        no_history_found: "No se encontró historial de consultas."
+        help_title: "Guía de uso del Asistente IA"
+
+      buttons:
+        cancel: "Cerrar"
+        send: "Enviar"
+        stop: "Detener"
+
+    errors:
+      company_not_found: "La empresa {company_short_name} no existe."
+      timeout: "El tiempo de espera ha expirado."
+
+
+      auth:
+        invalid_password: "La contraseña proporcionada es incorrecta."
+        user_not_found: "No se encontró un usuario con ese correo."
+        invalid_or_expired_token: "Token inválido o expirado."
+        session_creation_failed: "No se pudo crear la sesión del usuario."
+        authentication_required: "Autenticación requerida. No se proporcionó cookie de sesión o clave de API."
+        invalid_api_key: "Clave de API inválida o inactiva."
+        no_user_identifier_api: "No se proporcionó user_identifier para la llamada a la API."
+      templates:
+        company_not_found: "Empresa no encontrada."
+        home_template_not_found: "La plantilla de la página de inicio para la empresa '{company_short_name}' no está configurada."
+        processing_error: "Error al procesar el template: {error}"
+        template_not_found: "No se encontro el template: '{template_name}'."
+
+      general:
+        unexpected_error: "Ha ocurrido un error inesperado: {error}."
+        unsupported_language: "El idioma seleccionado no es válido."
+      signup:
+        company_not_found: "La empresa {company_short_name} no existe."
+        incorrect_password_for_existing_user: "La contraseña para el usuario {email} es incorrecta."
+        user_already_registered: "El usuario con email '{email}' ya existe en esta empresa."
+        password_mismatch: "Las contraseñas no coinciden. Por favor, inténtalo de nuevo."
+      change_password:
+        token_expired: "El enlace de cambio de contraseña ha expirado. Por favor, solicita uno nuevo."
+        password_mismatch: "Las contraseñas no coinciden. Por favor, inténtalo nuevamente."
+        invalid_temp_code: "El código temporal no es válido. Por favor, verifica o solicita uno nuevo."
+      forgot_password:
+        user_not_registered: "El usuario con correo {email} no está registrado."
+      verification:
+        token_expired: "El enlace de verificación ha expirado. Por favor, contacta a soporte si necesitas uno nuevo."
+        user_not_found: "El usuario que intentas verificar no existe."
+      validation:
+        password_too_short: "La contraseña debe tener al menos 8 caracteres."
+        password_no_uppercase: "La contraseña debe tener al menos una letra mayúscula."
+        password_no_lowercase: "La contraseña debe tener al menos una letra minúscula."
+        password_no_digit: "La contraseña debe tener al menos un número."
+        password_no_special_char: "La contraseña debe tener al menos un carácter especial."
+
+      services:
+        no_text_file: "El archivo no es texto o la codificación no es UTF-8"
+        no_output_file: "falta el nombre del archivo de salida"
+        no_data_for_excel: "faltan los datos  o no es una lista de diccionarios"
+        no_download_directory: "no esta configurado el directorio temporal para guardar excels"
+        cannot_create_excel: "no se pudo crear el archivo excel"
+        invalid_filename: "Nombre de archivo inválido"
+        file_not_exist : "Archivo no encontrado"
+        path_is_not_a_file : "La ruta no corresponde a un archivo"
+        file_validation_error : "Error validando archivo"
+        user_not_authorized: "Usuario no esta autorizado para esta empresa"
+        account_not_verified: "Tu cuenta no ha sido verificada. Por favor, revisa tu correo."
+        missing_response_id: "No se encontró 'previous_response_id' para '{company_short_name}/{user_identifier}'. Reinicia el contexto."
+
+    api_responses:
+      context_reloaded_success: "El contexto se ha recargado con éxito."
+
+    services:
+      mail_sent: "mail enviado exitosamente."
+      start_query: "Hola, cual es tu pregunta?"
+
+    flash_messages:
+      password_changed_success: "Tu contraseña ha sido restablecida exitosamente. Ahora puedes iniciar sesión."
+      signup_success: "Registro exitoso. Por favor, revisa tu correo para verificar tu cuenta."
+      user_associated_success: "Usuario existente asociado exitosamente a la nueva empresa."
+      account_verified_success: "Tu cuenta ha sido verificada exitosamente. ¡Bienvenido!"
+      forgot_password_success: "Si tu correo está registrado, recibirás un enlace para restablecer tu contraseña."
+
+    # Claves específicas para JavaScript
+    js_messages:
+      feedback_sent_success_title: "Feedback Enviado"
+      feedback_sent_success_body: "¡Gracias por tu comentario!"
+      feedback_sent_error: "No se pudo enviar el feedback, por favor intente nuevamente."
+      feedback_rating_error: "Por favor, califica al asistente con las estrellas."
+      feedback_comment_error: "Por favor, escribe tu comentario antes de enviar."
+      context_reloaded: "El contexto ha sido recargado."
+      error_loading_history: "Ocurrió un error al cargar el historial."
+      request_aborted: "La generación de la respuesta ha sido detenida."
+      processing_error: "Ocurrió un error al procesar la solicitud."
+      server_comm_error: "Error de comunicación con el servidor ({status}). Por favor, intente de nuevo más tarde."
+      network_error: "Ocurrió un error de red. Por favor, inténtalo de nuevo en unos momentos."
+      unknown_server_error: "Error desconocido del servidor."
+      loading: "Cargando..."
+      reload_init: "Iniciando recarga de contexto en segundo plano..."
+      no_history_found: "No existe historial de consultas."
+

iatoolkit/repositories/database_manager.py

@@ -21,8 +21,23 @@ class DatabaseManager:
         :param echo: Si True, habilita logs de SQL.
         """
         self.url = make_url(database_url)
-        self._engine = create_engine(database_url, echo=False)
-        self.SessionFactory = sessionmaker(bind=self._engine)
+        if database_url.startswith('sqlite'):       # for tests
+            self._engine = create_engine(database_url, echo=False)
+        else:
+            self._engine = create_engine(
+                database_url,
+                echo=False,
+                pool_size=10,  # per worker
+                max_overflow=20,
+                pool_timeout=30,
+                pool_recycle=1800,
+                pool_pre_ping=True,
+                future=True,
+            )
+        self.SessionFactory = sessionmaker(bind=self._engine,
+                                           autoflush=False,
+                                           autocommit=False,
+                                           expire_on_commit=False)
         self.scoped_session = scoped_session(self.SessionFactory)
 
         # REGISTRAR pgvector para cada nueva conexión solo en postgres
@@ -55,6 +70,11 @@ class DatabaseManager:
     def remove_session(self):
         self.scoped_session.remove()
 
+    def get_all_table_names(self) -> list[str]:
+        # Returns a list of all table names in the database
+        inspector = inspect(self._engine)
+        return inspector.get_table_names()
+
     def get_table_schema(self,
                          table_name: str,
                          schema_name: str | None = None,
@@ -62,7 +82,7 @@ class DatabaseManager:
         inspector = inspect(self._engine)
 
         if table_name not in inspector.get_table_names():
-            raise RuntimeError(f"La tabla '{table_name}' no existe en la BD.")
+            raise RuntimeError(f"Table '{table_name}' does not exist.")
 
         if exclude_columns is None:
             exclude_columns = []

iatoolkit/repositories/document_repo.py

@@ -22,7 +22,7 @@ class DocumentRepo:
     def get(self, company_id, filename: str ) -> Document:
         if not company_id or not filename:
             raise IAToolkitException(IAToolkitException.ErrorType.PARAM_NOT_FILLED,
-                               'Falta empresa o filename')
+                               'missing company_id or filename')
 
         return self.session.query(Document).filter_by(company_id=company_id, filename=filename).first()
 

iatoolkit/repositories/models.py

@@ -3,9 +3,10 @@
 #
 # IAToolkit is open source software.
 
-from sqlalchemy import Column, Integer, String, DateTime, Enum, Text, JSON, Boolean, ForeignKey, Table
+from sqlalchemy import Column, Integer, BigInteger, String, DateTime, Enum, Text, JSON, Boolean, ForeignKey, Table
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import relationship, class_mapper, declarative_base
+from sqlalchemy.sql import func
 from datetime import datetime
 from pgvector.sqlalchemy import Vector
 from enum import Enum as PyEnum
@@ -56,11 +57,8 @@ class Company(Base):
     # encrypted api-key
     openai_api_key = Column(String, nullable=True)
     gemini_api_key = Column(String, nullable=True)
-
-    branding = Column(JSON, nullable=True)
-    parameters = Column(JSON, nullable=True, default={})
+    parameters = Column(JSON, nullable=True)
     created_at = Column(DateTime, default=datetime.now)
-    allow_jwt = Column(Boolean, default=True, nullable=True)
 
     documents = relationship("Document",
                              back_populates="company",
@@ -105,6 +103,7 @@ class User(Base):
     created_at = Column(DateTime, default=datetime.now)
     password = Column(String, nullable=False)
     verified = Column(Boolean, nullable=False, default=False)
+    preferred_language = Column(String(5), nullable=True)
     verification_url = Column(String, nullable=True)
     temp_code = Column(String, nullable=True)
 
@@ -157,10 +156,10 @@ class Document(Base):
     company_id = Column(Integer, ForeignKey('iat_companies.id',
                     ondelete='CASCADE'), nullable=False)
     filename = Column(String(256), nullable=False, index=True)
-    content = Column(Text, nullable=False)
-    content_b64 = Column(Text, nullable=False)
     meta = Column(JSON, nullable=True)
     created_at = Column(DateTime, default=datetime.now)
+    content = Column(Text, nullable=False)
+    content_b64 = Column(Text, nullable=False)
 
     company = relationship("Company", back_populates="documents")
 
@@ -203,7 +202,10 @@ class VSDoc(Base):
     document_id = Column(Integer, ForeignKey('iat_documents.id',
                         ondelete='CASCADE'), nullable=False)
     text = Column(Text, nullable=False)
-    embedding = Column(Vector(384), nullable=False)  # Ajusta la dimensión si es necesario
+
+    # the size of this vector should be set depending on the embedding model used
+    # for OpenAI is 1536, and for huggingface is 384
+    embedding = Column(Vector(1536), nullable=False)
 
     company = relationship("Company", back_populates="vsdocs")
 
@@ -264,8 +266,7 @@ class UserFeedback(Base):
     id = Column(Integer, primary_key=True)
     company_id = Column(Integer, ForeignKey('iat_companies.id',
                                             ondelete='CASCADE'), nullable=False)
-    local_user_id = Column(Integer, default=0, nullable=True)
-    external_user_id = Column(String(128), default='', nullable=True)
+    user_identifier = Column(String(128), default='', nullable=True)
     message = Column(Text, nullable=False)
     rating = Column(Integer, nullable=False)
     created_at = Column(DateTime, default=datetime.now)
@@ -307,3 +308,27 @@ class Prompt(Base):
 
     company = relationship("Company", back_populates="prompts")
     category = relationship("PromptCategory", back_populates="prompts")
+
+class AccessLog(Base):
+    # Modelo ORM para registrar cada intento de acceso a la plataforma.
+    __tablename__ = 'iat_access_log'
+
+    id = Column(BigInteger, primary_key=True)
+
+    timestamp = Column(DateTime(timezone=True), server_default=func.now(), nullable=False, index=True)
+    company_short_name = Column(String(100), nullable=False, index=True)
+    user_identifier = Column(String(255), index=True)
+
+    # Cómo y el Resultado
+    auth_type = Column(String(20), nullable=False) # 'local', 'external_api', 'redeem_token', etc.
+    outcome = Column(String(10), nullable=False)   # 'success' o 'failure'
+    reason_code = Column(String(50))               # Causa de fallo, ej: 'INVALID_CREDENTIALS'
+
+    # Contexto de la Petición
+    source_ip = Column(String(45), nullable=False)
+    user_agent_hash = Column(String(16))           # Hash corto del User-Agent
+    request_path = Column(String(255), nullable=False)
+
+    def __repr__(self):
+        return (f"<AccessLog(id={self.id}, company='{self.company_short_name}', "
+                f"user='{self.user_identifier}', outcome='{self.outcome}')>")

iatoolkit/repositories/profile_repo.py

@@ -72,9 +72,10 @@ class ProfileRepo:
     def create_company(self, new_company: Company):
         company = self.session.query(Company).filter_by(name=new_company.name).first()
         if company:
-            company.parameters = new_company.parameters
-            company.branding = new_company.branding
+            if company.parameters != new_company.parameters:
+                company.parameters = new_company.parameters
         else:
+            # Si la compañía no existe, la añade a la sesión.
             self.session.add(new_company)
             company = new_company
 

iatoolkit/repositories/vs_repo.py

@@ -4,40 +4,38 @@
 # IAToolkit is open source software.
 
 from sqlalchemy import  text
-from huggingface_hub import InferenceClient
 from injector import inject
 from iatoolkit.common.exceptions import IAToolkitException
 from iatoolkit.repositories.database_manager import DatabaseManager
-from iatoolkit.repositories.models import Document, VSDoc
-import os
+from iatoolkit.services.embedding_service import EmbeddingService
+from iatoolkit.repositories.models import Document, VSDoc, Company
 import logging
 
+
 class VSRepo:
     @inject
-    def __init__(self, db_manager: DatabaseManager):
+    def __init__(self,
+                 db_manager: DatabaseManager,
+                 embedding_service: EmbeddingService):
         self.session = db_manager.get_session()
-
-        # Inicializar el modelo de embeddings
-        self.embedder = InferenceClient(
-            model="sentence-transformers/all-MiniLM-L6-v2",
-            token=os.getenv('HF_TOKEN'))
+        self.embedding_service = embedding_service
 
 
-    def add_document(self, vs_chunk_list: list[VSDoc]):
+    def add_document(self, company_short_name, vs_chunk_list: list[VSDoc]):
         try:
             for doc in vs_chunk_list:
                 # calculate the embedding for the text
-                doc.embedding = self.embedder.feature_extraction(doc.text)
+                doc.embedding = self.embedding_service.embed_text(company_short_name, doc.text)
                 self.session.add(doc)
             self.session.commit()
         except Exception as e:
-            logging.error(f"Error insertando documentos en PostgreSQL: {str(e)}")
+            logging.error(f"Error inserting documents into PostgreSQL: {str(e)}")
             self.session.rollback()
             raise IAToolkitException(IAToolkitException.ErrorType.VECTOR_STORE_ERROR,
-                               f"Error insertando documentos en PostgreSQL: {str(e)}")
+                               f"Error inserting documents into PostgreSQL: {str(e)}")
 
     def query(self,
-              company_id: int,
+              company_short_name: str,
               query_text: str,
               n_results=5,
               metadata_filter=None
@@ -46,18 +44,25 @@ class VSRepo:
         search documents similar to the query for a company
 
         Args:
-            company_id:
+            company_short_name: The company's unique short name.
             query_text: query text
             n_results: max number of results to return
-            metadata_filter:  (ej: {"document_type": "certificate"})
+            metadata_filter:  (e.g., {"document_type": "certificate"})
 
         Returns:
             list of documents matching the query and filters
         """
-        # Generate the embedding with the query text
-        query_embedding = self.embedder.feature_extraction([query_text])[0]
+        # Generate the embedding with the query text for the specific company
+        query_embedding = self.embedding_service.embed_text(company_short_name, query_text)
 
+        sql_query, params = None, None
         try:
+            # Get company ID from its short name for the SQL query
+            company = self.session.query(Company).filter(Company.short_name == company_short_name).one_or_none()
+            if not company:
+                raise IAToolkitException(IAToolkitException.ErrorType.VECTOR_STORE_ERROR,
+                                   f"Company with short name '{company_short_name}' not found.")
+
             # build the SQL query
             sql_query_parts = ["""
                                SELECT iat_documents.id, \
@@ -73,11 +78,12 @@ class VSRepo:
 
             # query parameters
             params = {
-                "company_id": company_id,
+                "company_id": company.id,
                 "query_embedding": query_embedding,
                 "n_results": n_results
             }
 
+
             # add metadata filter, if exists
             if metadata_filter and isinstance(metadata_filter, dict):
                 for key, value in metadata_filter.items():
@@ -108,7 +114,7 @@ class VSRepo:
                 meta_data = row[4] if len(row) > 4 and row[4] is not None else {}
                 doc = Document(
                     id=row[0],
-                    company_id=company_id,
+                    company_id=company.id,
                     filename=row[1],
                     content=row[2],
                     content_b64=row[3],