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

iatoolkit/services/load_documents_service.py

@@ -1,17 +1,13 @@
 # Copyright (c) 2024 Fernando Libedinsky
 # Product: IAToolkit
 
-from iatoolkit.repositories.vs_repo import VSRepo
-from iatoolkit.repositories.document_repo import DocumentRepo
-from iatoolkit.repositories.models import Document, VSDoc, Company
-from iatoolkit.services.document_service import DocumentService
+from iatoolkit.repositories.models import Company
 from iatoolkit.services.configuration_service import ConfigurationService
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from iatoolkit.services.knowledge_base_service import KnowledgeBaseService
 from iatoolkit.infra.connectors.file_connector_factory import FileConnectorFactory
 from iatoolkit.services.file_processor_service import FileProcessorConfig, FileProcessor
 from iatoolkit.common.exceptions import IAToolkitException
 import logging
-import base64
 from injector import inject, singleton
 import os
 
@@ -19,31 +15,21 @@ import os
 @singleton
 class LoadDocumentsService:
     """
-    Orchestrates the process of loading, processing, and storing documents
-    from various sources defined in the company's configuration.
+    Orchestrates the discovery and loading of documents from configured sources.
+    Delegates the processing and ingestion logic to KnowledgeBaseService.
     """
     @inject
     def __init__(self,
                  config_service: ConfigurationService,
                  file_connector_factory: FileConnectorFactory,
-                 doc_service: DocumentService,
-                 doc_repo: DocumentRepo,
-                 vector_store: VSRepo,
+                 knowledge_base_service: KnowledgeBaseService
                  ):
         self.config_service = config_service
-        self.doc_service = doc_service
-        self.doc_repo = doc_repo
-        self.vector_store = vector_store
         self.file_connector_factory = file_connector_factory
+        self.knowledge_base_service = knowledge_base_service
 
         logging.getLogger().setLevel(logging.ERROR)
 
-        self.splitter = RecursiveCharacterTextSplitter(
-            chunk_size=1000,
-            chunk_overlap=100,
-            separators=["\n\n", "\n", "."]
-        )
-
     def load_sources(self,
                      company: Company,
                      sources_to_load: list[str] = None,
@@ -67,7 +53,7 @@ class LoadDocumentsService:
 
         if not sources_to_load:
             raise IAToolkitException(IAToolkitException.ErrorType.PARAM_NOT_FILLED,
-                                f"Missing sources to load for company '{company.short_name}'.")
+                                     f"Missing sources to load for company '{company.short_name}'.")
 
         base_connector_config = self._get_base_connector_config(knowledge_base_config)
         all_sources = knowledge_base_config.get('document_sources', {})
@@ -89,6 +75,7 @@ class LoadDocumentsService:
                 # Prepare the context for the callback function.
                 context = {
                     'company': company,
+                    'collection': source_config.get('metadata', {}).get('collection'),
                     'metadata': source_config.get('metadata', {})
                 }
 
@@ -130,45 +117,29 @@ class LoadDocumentsService:
 
     def _file_processing_callback(self, company: Company, filename: str, content: bytes, context: dict = None):
         """
-        Callback method to process a single file. It extracts text, merges metadata,
-        and saves the document to both relational and vector stores.
+        Callback method to process a single file.
+        Delegates the actual ingestion (storage, vectorization) to KnowledgeBaseService.
         """
         if not company:
             raise IAToolkitException(IAToolkitException.ErrorType.MISSING_PARAMETER, "Missing company object in callback.")
 
-        if self.doc_repo.get(company_id=company.id, filename=filename):
-            logging.debug(f"File '{filename}' already exists for company '{company.id}'. Skipping.")
-            return
-
         try:
-            document_content = self.doc_service.file_to_txt(filename, content)
-
             # Get predefined metadata from the context passed by the processor.
             predefined_metadata = context.get('metadata', {}) if context else {}
 
-            # Save the document to the relational database.
-            session = self.doc_repo.session
-            new_document = Document(
-                company_id=company.id,
+            # Delegate heavy lifting to KnowledgeBaseService
+            new_document = self.knowledge_base_service.ingest_document_sync(
+                company=company,
                 filename=filename,
-                content=document_content,
-                content_b64=base64.b64encode(content).decode('utf-8'),
-                meta=predefined_metadata
+                content=content,
+                collection=predefined_metadata.get('collection'),
+                metadata=predefined_metadata
             )
-            session.add(new_document)
-            session.flush()  # Flush to get the new_document.id without committing.
 
-            # Split into chunks and prepare for vector store.
-            chunks = self.splitter.split_text(document_content)
-            vs_docs = [VSDoc(company_id=company.id, document_id=new_document.id, text=text) for text in chunks]
-
-            # Add document chunks to the vector store.
-            self.vector_store.add_document(company.short_name, vs_docs)
-
-            session.commit()
             return new_document
+
         except Exception as e:
-            self.doc_repo.session.rollback()
+            # We log here but re-raise to let FileProcessor handle the error counting/continue logic
             logging.exception(f"Error processing file '{filename}': {e}")
             raise IAToolkitException(IAToolkitException.ErrorType.LOAD_DOCUMENT_ERROR,
                                      f"Error while processing file: {filename}")

iatoolkit/services/profile_service.py

@@ -9,6 +9,7 @@ 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
+from iatoolkit.services.dispatcher_service import Dispatcher
 from iatoolkit.services.language_service import LanguageService
 from iatoolkit.services.user_session_context_service import UserSessionContextService
 from iatoolkit.services.configuration_service import ConfigurationService
@@ -19,7 +20,7 @@ import re
 import secrets
 import string
 import logging
-from iatoolkit.services.dispatcher_service import Dispatcher
+from typing import List, Dict
 
 
 class ProfileService:
@@ -41,7 +42,6 @@ class ProfileService:
         self.mail_service = mail_service
         self.bcrypt = Bcrypt()
 
-
     def login(self, company_short_name: str, email: str, password: str) -> dict:
         try:
             # check if user exists
@@ -65,6 +65,8 @@ class ProfileService:
                 return {'success': False,
                         "message": self.i18n_service.t('errors.services.account_not_verified')}
 
+            user_role = self.profile_repo.get_user_role_in_company(company.id, user.id)
+
             # 1. Build the local user profile dictionary here.
             # the user_profile variables are used on the LLM templates also (see in query_main.prompt)
             user_identifier = user.email
@@ -73,6 +75,7 @@ class ProfileService:
                 "user_fullname": f'{user.first_name} {user.last_name}',
                 "user_is_local": True,
                 "user_id": user.id,
+                "user_role": user_role,
                 "extras": {}
             }
 
@@ -320,19 +323,44 @@ class ProfileService:
     def get_company_by_short_name(self, short_name: str) -> Company:
         return self.profile_repo.get_company_by_short_name(short_name)
 
+    def get_company_users(self, company_short_name: str) -> List[Dict]:
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            return []
+
+        # get the company users from the repo
+        company_users =  self.profile_repo.get_company_users_with_details(company_short_name)
+
+        users_data = []
+        for user, role, last_access in company_users:
+            users_data.append({
+                "first_name": user.first_name,
+                "last_name": user.last_name,
+                "email": user.email,
+                "created": user.created_at,
+                "verified": user.verified,
+                "role": role or "user",
+                "last_access": last_access
+            })
+
+        return users_data
+
     def get_active_api_key_entry(self, api_key_value: str) -> ApiKey | None:
         return self.profile_repo.get_active_api_key_entry(api_key_value)
 
-    def new_api_key(self, company_short_name: str):
+    def new_api_key(self, company_short_name: str, key_name: str):
         company = self.get_company_by_short_name(company_short_name)
         if not company:
             return {"error": self.i18n_service.t('errors.company_not_found', company_short_name=company_short_name)}
 
+        if not key_name:
+            return {"error": self.i18n_service.t('errors.auth.api_key_name_required')}
+
         length = 40     # lenght of the api key
         alphabet = string.ascii_letters + string.digits
         key = ''.join(secrets.choice(alphabet) for i in range(length))
 
-        api_key = ApiKey(key=key, company_id=company.id)
+        api_key = ApiKey(key=key, company_id=company.id, key_name=key_name)
         self.profile_repo.create_api_key(api_key)
         return {"api-key": key}
 

iatoolkit/services/prompt_service.py

@@ -4,15 +4,16 @@
 # IAToolkit is open source software.
 
 from injector import inject
+from iatoolkit.common.interfaces.asset_storage import AssetRepository, AssetType
 from iatoolkit.repositories.llm_query_repo import LLMQueryRepo
 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
+import os
 
 # iatoolkit system prompts definitions
 _SYSTEM_PROMPTS = [
@@ -24,27 +25,37 @@ _SYSTEM_PROMPTS = [
 class PromptService:
     @inject
     def __init__(self,
+                 asset_repo: AssetRepository,
                  llm_query_repo: LLMQueryRepo,
                  profile_repo: ProfileRepo,
                  i18n_service: I18nService):
+        self.asset_repo = asset_repo
         self.llm_query_repo = llm_query_repo
         self.profile_repo = profile_repo
         self.i18n_service = i18n_service
 
-    def sync_company_prompts(self, company_instance, prompts_config: list, categories_config: list):
+    def sync_company_prompts(self, company_short_name: str, prompts_config: list, categories_config: list):
         """
         Synchronizes prompt categories and prompts from YAML config to Database.
         Strategies:
         - Categories: Create or Update existing based on name.
         - Prompts: Create or Update existing based on name. Soft-delete or Delete unused.
         """
+        if not prompts_config:
+            return
+
+        company = self.profile_repo.get_company_by_short_name(company_short_name)
+        if not company:
+            raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
+                                     f'Company {company_short_name} not found')
+
         try:
             # 1. Sync Categories
             category_map = {}
 
             for i, category_name in enumerate(categories_config):
                 category_obj = PromptCategory(
-                    company_id=company_instance.company.id,
+                    company_id=company.id,
                     name=category_name,
                     order=i + 1
                 )
@@ -69,10 +80,10 @@ class PromptService:
                 filename = f"{prompt_name}.prompt"
 
                 new_prompt = Prompt(
-                    company_id=company_instance.company.id,
+                    company_id=company.id,
                     name=prompt_name,
-                    description=prompt_data['description'],
-                    order=prompt_data['order'],
+                    description=prompt_data.get('description'),
+                    order=prompt_data.get('order'),
                     category_id=category_obj.id,
                     active=prompt_data.get('active', True),
                     is_system_prompt=False,
@@ -83,7 +94,7 @@ class PromptService:
                 self.llm_query_repo.create_or_update_prompt(new_prompt)
 
             # 3. Cleanup: Delete prompts present in DB but not in Config
-            existing_prompts = self.llm_query_repo.get_prompts(company_instance.company)
+            existing_prompts = self.llm_query_repo.get_prompts(company)
             for p in existing_prompts:
                 if p.name not in defined_prompt_names:
                     # Using hard delete to keep consistent with previous "refresh" behavior
@@ -151,12 +162,9 @@ class PromptService:
                 raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
                                 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):
+            if not self.asset_repo.exists(company.short_name, AssetType.PROMPT, prompt_filename):
                 raise IAToolkitException(IAToolkitException.ErrorType.INVALID_NAME,
-                               f'missing prompt file: {relative_prompt_path}')
+                               f'missing prompt file: {prompt_filename} in prompts/')
 
         if custom_fields:
             for f in custom_fields:
@@ -188,33 +196,28 @@ class PromptService:
 
     def get_prompt_content(self, company: Company, prompt_name: str):
         try:
-            user_prompt_content = []
-            execution_dir = os.getcwd()
-
             # get the user prompt
             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"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"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()
+                user_prompt_content = self.asset_repo.read_text(
+                    company.short_name,
+                    AssetType.PROMPT,
+                    user_prompt.filename
+                )
+            except FileNotFoundError:
+                raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
+                                         f"prompt file '{user_prompt.filename}' does not exist for company '{company.short_name}'")
             except Exception as e:
                 raise IAToolkitException(IAToolkitException.ErrorType.FILE_IO_ERROR,
-                                   f"error while reading prompt: '{prompt_name}' in this pathname {absolute_filepath}: {e}")
+                                         f"error while reading prompt: '{prompt_name}': {e}")
 
             return user_prompt_content
 
         except IAToolkitException:
-            # Vuelve a lanzar las IAToolkitException que ya hemos manejado
-            # para que no sean capturadas por el siguiente bloque.
             raise
         except Exception as e:
             logging.exception(
@@ -260,7 +263,7 @@ class PromptService:
             # get all the prompts
             all_prompts = self.llm_query_repo.get_prompts(company)
 
-            # Agrupar prompts por categoría
+            # group by category
             prompts_by_category = defaultdict(list)
             for prompt in all_prompts:
                 if prompt.active:
@@ -268,14 +271,13 @@ class PromptService:
                         cat_key = (prompt.category.order, prompt.category.name)
                         prompts_by_category[cat_key].append(prompt)
 
-            # Ordenar los prompts dentro de cada categoría
+            # sort each category by order
             for cat_key in prompts_by_category:
                 prompts_by_category[cat_key].sort(key=lambda p: p.order)
 
-            # Crear la estructura de respuesta final, ordenada por la categoría
             categorized_prompts = []
 
-            # Ordenar las categorías por su 'order'
+            # sort categories by order
             sorted_categories = sorted(prompts_by_category.items(), key=lambda item: item[0][0])
 
             for (cat_order, cat_name), prompts in sorted_categories:

iatoolkit/services/query_service.py

@@ -3,7 +3,7 @@
 #
 # IAToolkit is open source software.
 
-from iatoolkit.infra.llm_client import llmClient
+from iatoolkit.services.llm_client_service import llmClient
 from iatoolkit.services.profile_service import ProfileService
 from iatoolkit.repositories.profile_repo import ProfileRepo
 from iatoolkit.services.tool_service import ToolService
@@ -11,11 +11,11 @@ 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.models import Task
 from iatoolkit.services.dispatcher_service import Dispatcher
 from iatoolkit.services.prompt_service import PromptService
 from iatoolkit.services.user_session_context_service import UserSessionContextService
 from iatoolkit.services.history_manager_service import HistoryManagerService
+from iatoolkit.common.model_registry import ModelRegistry
 from iatoolkit.common.util import Utility
 from injector import inject
 import base64
@@ -33,6 +33,7 @@ class HistoryHandle:
     company_short_name: str
     user_identifier: str
     type: str
+    model: str | None = None
     request_params: dict = None
 
 
@@ -52,6 +53,7 @@ class QueryService:
                  configuration_service: ConfigurationService,
                  history_manager: HistoryManagerService,
                  util: Utility,
+                 model_registry: ModelRegistry
                  ):
         self.profile_service = profile_service
         self.company_context_service = company_context_service
@@ -66,6 +68,7 @@ class QueryService:
         self.configuration_service = configuration_service
         self.llm_client = llm_client
         self.history_manager = history_manager
+        self.model_registry = model_registry
 
 
     def _resolve_model(self, company_short_name: str, model: Optional[str]) -> str:
@@ -78,13 +81,16 @@ class QueryService:
         return effective_model
 
     def _get_history_type(self, model: str) -> str:
-        return HistoryManagerService.TYPE_SERVER_SIDE if self.util.is_openai_model(
-            model) else HistoryManagerService.TYPE_CLIENT_SIDE
+        history_type_str = self.model_registry.get_history_type(model)
+        if history_type_str == "server_side":
+            return HistoryManagerService.TYPE_SERVER_SIDE
+        else:
+            return HistoryManagerService.TYPE_CLIENT_SIDE
 
 
     def _build_user_facing_prompt(self, company, user_identifier: str,
                                   client_data: dict, files: list,
-                                  prompt_name: Optional[str], question: str) -> str:
+                                  prompt_name: Optional[str], question: str):
         # get the user profile data from the session context
         user_profile = self.profile_service.get_profile_by_identifier(company.short_name, user_identifier)
 
@@ -124,9 +130,12 @@ class QueryService:
 
         return user_turn_prompt, effective_question
 
-    def _ensure_valid_history(self, company, user_identifier: str,
-                              effective_model: str, user_turn_prompt: str,
-                              ignore_history: bool) -> tuple[Optional[HistoryHandle], Optional[dict]]:
+    def _ensure_valid_history(self, company,
+                              user_identifier: str,
+                              effective_model: str,
+                              user_turn_prompt: str,
+                              ignore_history: bool
+                              ) -> tuple[Optional[HistoryHandle], Optional[dict]]:
         """
             Manages the history strategy and rebuilds context if necessary.
             Returns: (HistoryHandle, error_response)
@@ -137,7 +146,8 @@ class QueryService:
         handle = HistoryHandle(
             company_short_name=company.short_name,
             user_identifier=user_identifier,
-            type=history_type
+            type=history_type,
+            model=effective_model
         )
 
         # pass the handle to populate request_params
@@ -197,22 +207,35 @@ class QueryService:
     def init_context(self, company_short_name: str,
                      user_identifier: str,
                      model: str = None) -> dict:
+        """
+        Forces a context rebuild for a given user and (optionally) model.
+
+        - Clears LLM-related context for the resolved model.
+        - Regenerates the static company/user context.
+        - Sends the context to the LLM for that 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.")
+        # 1. Resolve the effective model for this user/company
+        effective_model = self._resolve_model(company_short_name, model)
 
-        # 2. LLM context is clean, now we can load it again
+        # 2. Clear only the LLM-related context for this model
+        self.session_context.clear_all_context(company_short_name, user_identifier,model=effective_model)
+        logging.info(
+            f"Context for {company_short_name}/{user_identifier} "
+            f"(model={effective_model}) has been cleared."
+        )
+
+        # 3. Static LLM context is now clean, we can prepare it again (model-agnostic)
         self.prepare_context(
             company_short_name=company_short_name,
             user_identifier=user_identifier
         )
 
-        # 3. communicate the new context to the LLM
+        # 4. Communicate the new context to the specific LLM model
         response = self.set_context_for_llm(
             company_short_name=company_short_name,
             user_identifier=user_identifier,
-            model=model
+            model=effective_model
         )
 
         return response
@@ -257,8 +280,10 @@ class QueryService:
                             company_short_name: str,
                             user_identifier: str,
                             model: str = ''):
-
-        # This service takes a pre-built context and send to the LLM
+        """
+        Takes a pre-built static context and sends it to the LLM for the given model.
+        Also initializes the model-specific history through HistoryManagerService.
+        """
         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")
@@ -267,8 +292,8 @@ class QueryService:
         # --- Model Resolution ---
         effective_model = self._resolve_model(company_short_name, 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}"
+        # Lock per (company, user, model) to avoid concurrent rebuilds for the same model
+        lock_key = f"lock:context:{company_short_name}/{user_identifier}/{effective_model}"
         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.")
@@ -310,13 +335,13 @@ class QueryService:
     def llm_query(self,
                   company_short_name: str,
                   user_identifier: str,
-                  task: Optional[Task] = None,
+                  model: Optional[str] = None,
                   prompt_name: str = None,
                   question: str = '',
                   client_data: dict = {},
                   ignore_history: bool = False,
-                  files: list = [],
-                  model: Optional[str] = None) -> dict:
+                  files: list = []
+                  ) -> dict:
         try:
             company = self.profile_repo.get_company_by_short_name(short_name=company_short_name)
             if not company:
@@ -378,10 +403,10 @@ class QueryService:
             if not response.get('valid_response'):
                 response['error'] = True
 
-                # save history using the manager passing the handle
-                self.history_manager.update_history(
-                    history_handle, user_turn_prompt, response
-                )
+            # save history using the manager passing the handle
+            self.history_manager.update_history(
+                history_handle, user_turn_prompt, response
+            )
 
             return response
         except Exception as e: