iatoolkit 2.29.0__tar.gz → 2.30.0__tar.gz

{iatoolkit-2.29.0/src/iatoolkit.egg-info → iatoolkit-2.30.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iatoolkit
-Version: 2.29.0
+Version: 2.30.0
 Summary: IAToolkit
 Author: Fernando Libedinsky
 License-Expression: MIT

{iatoolkit-2.29.0 → iatoolkit-2.30.0}/src/iatoolkit/__init__.py

@@ -3,7 +3,7 @@
 #
 # IAToolkit is open source software.
 
-__version__ = "2.29.0"
+__version__ = "2.30.0"
 
 # Expose main classes and functions at the top level of the package
 

{iatoolkit-2.29.0 → iatoolkit-2.30.0}/src/iatoolkit/core.py

@@ -115,6 +115,7 @@ class IAToolkit:
         self._setup_request_globals()
         self._setup_context_processors()
         self._sync_system_tools_on_boot()
+        self._run_configured_startup_warmup()
 
         # register data sources
         if start:
@@ -149,6 +150,19 @@ class IAToolkit:
         except Exception as exc:
             logging.warning("⚠️ Unable to sync system tools on boot: %s", exc)
 
+    def _run_configured_startup_warmup(self):
+        from iatoolkit.services.warmup_service import WarmupService
+
+        try:
+            warmup_service = self._injector.get(WarmupService)
+            warmed_companies = warmup_service.warmup_startup_configured_companies(trigger="core_startup")
+            if warmed_companies:
+                logging.info("🔥 Startup warm-up completed for companies=%s", warmed_companies)
+            else:
+                logging.info("🔥 Startup warm-up skipped: no embedding provider enabled warmup_on_startup.")
+        except Exception:
+            logging.exception("⚠️ Startup warm-up failed.")
+
     def _get_config_value(self, key: str, default=None):
         # get a value from the config dict or the environment variable
         return self.config.get(key, os.getenv(key, default))

{iatoolkit-2.29.0 → iatoolkit-2.30.0}/src/iatoolkit/repositories/vs_repo.py

@@ -86,6 +86,7 @@ class VSRepo:
               n_results=5,
               metadata_filter=None,
               collection_ids: List[int] | None = None,
+              include_urls: bool = True,
               ) -> list[Dict]:
         """
         search documents similar to the query for a company
@@ -95,6 +96,7 @@ class VSRepo:
             query_text: query text
             n_results: max number of results to return
             metadata_filter:  (e.g., {"document_type": "certificate"})
+            include_urls: Generate storage URLs for returned chunks when true.
 
         Returns:
             list of documents matching the query and filters
@@ -185,7 +187,7 @@ class VSRepo:
                 # get the url of the document
                 storage_key = row[3] if len(row) > 3 and row[3] is not None else None
                 url = None
-                if storage_key:
+                if include_urls and storage_key:
                     url = self.storage_service.generate_presigned_url(company_short_name, storage_key)
 
                 vs_documents.append(

{iatoolkit-2.29.0 → iatoolkit-2.30.0}/src/iatoolkit/services/configuration_service.py

@@ -717,6 +717,23 @@ class ConfigurationService:
                 add_error("embedding_provider", "Missing required key: 'provider'")
             if not config.get("embedding_provider", {}).get("model"):
                 add_error("embedding_provider", "Missing required key: 'model'")
+            if (
+                    "warmup_on_startup" in config.get("embedding_provider", {})
+                    and not isinstance(config.get("embedding_provider", {}).get("warmup_on_startup"), bool)
+            ):
+                add_error("embedding_provider.warmup_on_startup", "Must be a boolean.")
+
+        embedding_providers_cfg = config.get("embedding_providers")
+        if embedding_providers_cfg is not None:
+            if not isinstance(embedding_providers_cfg, dict):
+                add_error("embedding_providers", "Section must be a dictionary.")
+            else:
+                for provider_key, provider_cfg in embedding_providers_cfg.items():
+                    if not isinstance(provider_cfg, dict):
+                        add_error(f"embedding_providers.{provider_key}", "Provider configuration must be a dictionary.")
+                        continue
+                    if "warmup_on_startup" in provider_cfg and not isinstance(provider_cfg.get("warmup_on_startup"), bool):
+                        add_error(f"embedding_providers.{provider_key}.warmup_on_startup", "Must be a boolean.")
 
         # 3b. Visual Embedding Provider (Optional)
         if config.get("visual_embedding_provider"):

{iatoolkit-2.29.0 → iatoolkit-2.30.0}/src/iatoolkit/services/inference_service.py

@@ -25,9 +25,12 @@ class InferenceService:
 
     DEFAULT_CONNECT_TIMEOUT_SECONDS = 5.0
     DEFAULT_READ_TIMEOUT_SECONDS = 300.0
+    DEFAULT_RETRY_BUDGET_SECONDS = 30.0
     DEFAULT_RETRY_INITIAL_DELAY_SECONDS = 5.0
     DEFAULT_RETRY_MAX_DELAY_SECONDS = 30.0
-    RETRYABLE_STATUS_CODES = {408, 429, 502, 503, 504}
+    # Inference POSTs are idempotent. Hosted model runtimes can emit transient
+    # generic 500 responses while recovering, so configured retry budgets may retry them.
+    RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}
 
     @inject
     def __init__(self,
@@ -251,7 +254,7 @@ class InferenceService:
         return self._resolve_float_config(
             config,
             "retry_budget_seconds",
-            0.0,
+            self.DEFAULT_RETRY_BUDGET_SECONDS,
             minimum=0.0,
         )
 
@@ -307,16 +310,47 @@ class InferenceService:
             attempt: int,
             delay_seconds: float,
             reason: str,
+            attempt_duration_seconds: float | None = None,
+            elapsed_seconds: float | None = None,
             suppress_error_logging: bool = False
     ) -> None:
         log_fn = logging.debug if suppress_error_logging else logging.warning
+        if attempt_duration_seconds is None or elapsed_seconds is None:
+            log_fn(
+                "Inference request retry %s scheduled in %.1fs: %s",
+                attempt,
+                delay_seconds,
+                reason,
+            )
+            return
+
         log_fn(
-            "Inference request retry %s scheduled in %.1fs: %s",
+            "Inference request retry %s scheduled in %.1fs after %.1fs attempt (elapsed %.1fs): %s",
             attempt,
             delay_seconds,
+            attempt_duration_seconds,
+            elapsed_seconds,
             reason,
         )
 
+    def _timeout_for_attempt(
+            self,
+            *,
+            timeout: tuple[float, float],
+            deadline: float | None,
+            attempt: int,
+    ) -> tuple[float, float] | None:
+        if deadline is None or attempt == 1:
+            return timeout
+
+        remaining_seconds = deadline - time.monotonic()
+        if remaining_seconds <= 0:
+            return None
+
+        connect_timeout, read_timeout = timeout
+        bounded_timeout = max(0.1, remaining_seconds)
+        return min(connect_timeout, bounded_timeout), min(read_timeout, bounded_timeout)
+
     def _call_endpoint(
             self,
             url: str,
@@ -324,35 +358,51 @@ class InferenceService:
             payload: dict,
             suppress_error_logging: bool = False,
             timeout: tuple[float, float] = (DEFAULT_CONNECT_TIMEOUT_SECONDS, DEFAULT_READ_TIMEOUT_SECONDS),
-            retry_budget_seconds: float = 0.0,
+            retry_budget_seconds: float = DEFAULT_RETRY_BUDGET_SECONDS,
     ) -> Any:
         """Performs the POST request to the HF Endpoint."""
         headers = {
             "Authorization": f"Bearer {api_key}",
             "Content-Type": "application/json",
         }
-        deadline = time.monotonic() + retry_budget_seconds if retry_budget_seconds > 0 else None
+        started_at = time.monotonic()
+        deadline = started_at + retry_budget_seconds if retry_budget_seconds > 0 else None
         attempt = 1
         delay_seconds = self.DEFAULT_RETRY_INITIAL_DELAY_SECONDS
 
         while True:
+            attempt_timeout = self._timeout_for_attempt(timeout=timeout, deadline=deadline, attempt=attempt)
+            if attempt_timeout is None:
+                error_msg = "Inference retry budget exhausted before next attempt."
+                if not suppress_error_logging:
+                    logging.error(error_msg)
+                raise ValueError(error_msg)
+
+            attempt_started_at = time.monotonic()
             try:
                 resp, status = self.call_service.post(
                     url,
                     json_dict=payload,
                     headers=headers,
-                    timeout=timeout
+                    timeout=attempt_timeout
                 )
             except Exception as exc:
+                attempt_duration = time.monotonic() - attempt_started_at
                 next_delay = self._next_retry_delay(delay_seconds=delay_seconds, deadline=deadline)
                 if self._is_retryable_request_exception(exc) and next_delay is not None:
                     self._log_retry_attempt(
                         attempt=attempt,
                         delay_seconds=next_delay,
                         reason=str(exc),
+                        attempt_duration_seconds=attempt_duration,
+                        elapsed_seconds=time.monotonic() - started_at,
                         suppress_error_logging=suppress_error_logging,
                     )
                     time.sleep(next_delay)
+                    if deadline is not None and time.monotonic() >= deadline:
+                        if not suppress_error_logging:
+                            logging.error(f"Failed to call inference endpoint: {exc}")
+                        raise
                     attempt += 1
                     delay_seconds = min(delay_seconds * 2, self.DEFAULT_RETRY_MAX_DELAY_SECONDS)
                     continue
@@ -364,6 +414,7 @@ class InferenceService:
             if status == 200:
                 return resp
 
+            attempt_duration = time.monotonic() - attempt_started_at
             error_msg = f"Inference Endpoint Error {status}"
             if isinstance(resp, dict) and 'error' in resp:
                 error_msg += f": {resp['error']}"
@@ -374,9 +425,16 @@ class InferenceService:
                     attempt=attempt,
                     delay_seconds=next_delay,
                     reason=error_msg,
+                    attempt_duration_seconds=attempt_duration,
+                    elapsed_seconds=time.monotonic() - started_at,
                     suppress_error_logging=suppress_error_logging,
                 )
                 time.sleep(next_delay)
+                if deadline is not None and time.monotonic() >= deadline:
+                    if not suppress_error_logging:
+                        logging.error(f"{error_msg} | Payload keys: {list(payload.keys())}")
+                        logging.error(f"Failed to call inference endpoint: {error_msg}")
+                    raise ValueError(error_msg)
                 attempt += 1
                 delay_seconds = min(delay_seconds * 2, self.DEFAULT_RETRY_MAX_DELAY_SECONDS)
                 continue

{iatoolkit-2.29.0 → iatoolkit-2.30.0}/src/iatoolkit/services/knowledge_base_service.py

@@ -488,7 +488,8 @@ class KnowledgeBaseService:
                query: str,
                n_results: int = 5,
                collection: str | list[str] | None = None,
-               metadata_filter: dict = None
+               metadata_filter: dict = None,
+               include_urls: bool = True,
                ):
         """
         Performs a semantic search and returns the list of Document objects (chunks).
@@ -500,6 +501,7 @@ class KnowledgeBaseService:
             n_results: Max number of chunks to retrieve.
             metadata_filter: Optional filter for document metadata.
             collection: Optional collection name or list of collection names.
+            include_urls: Generate storage URLs for returned chunks when true.
 
         Returns:
             List of Document objects found.
@@ -532,6 +534,7 @@ class KnowledgeBaseService:
             n_results=n_results,
             metadata_filter=metadata_filter,
             collection_ids=collection_ids,
+            include_urls=include_urls,
         )
 
         return chunk_list

iatoolkit-2.30.0/src/iatoolkit/services/warmup_service.py

@@ -0,0 +1,196 @@
+# Copyright (c) 2024 Fernando Libedinsky
+# Product: IAToolkit
+#
+# IAToolkit is open source software.
+
+from __future__ import annotations
+
+import logging
+import time
+from injector import inject
+
+from iatoolkit.common.interfaces.secret_provider import SecretProvider
+from iatoolkit.common.secret_resolver import resolve_secret
+from iatoolkit.company_registry import get_registered_companies
+from iatoolkit.services.configuration_service import ConfigurationService
+from iatoolkit.services.embedding_service import EmbeddingService
+
+
+class WarmupService:
+    """
+    Lightweight warm-up orchestrator.
+    Keep it simple: no shared state, no locking, no dedup.
+    """
+
+    @inject
+    def __init__(self,
+                 config_service: ConfigurationService,
+                 embedding_service: EmbeddingService,
+                 secret_provider: SecretProvider):
+        self.config_service = config_service
+        self.embedding_service = embedding_service
+        self.secret_provider = secret_provider
+
+    def warmup_company(self, company_short_name: str, trigger: str = "manual"):
+        start = time.perf_counter()
+        try:
+            self._warmup_remote_text_embeddings(company_short_name)
+            elapsed_ms = (time.perf_counter() - start) * 1000
+            logging.info(
+                "🔥 Warm-up done for company='%s' trigger='%s' in %.2f ms",
+                company_short_name,
+                trigger,
+                elapsed_ms
+            )
+        except Exception as e:
+            elapsed_ms = (time.perf_counter() - start) * 1000
+            logging.debug(
+                "⚠️ Warm-up failed for company='%s' trigger='%s' in %.2f ms: %s",
+                company_short_name,
+                trigger,
+                elapsed_ms,
+                e
+            )
+
+    def warmup_registered_companies(self, trigger: str = "manual"):
+        for company_short_name in get_registered_companies().keys():
+            self.warmup_company(company_short_name, trigger=trigger)
+
+    def warmup_startup_configured_companies(self, trigger: str = "startup") -> list[str]:
+        warmed_companies: list[str] = []
+        for company_short_name in get_registered_companies().keys():
+            warmed_profiles = self._warmup_remote_text_embeddings(company_short_name, startup_only=True)
+            if not warmed_profiles:
+                logging.debug(
+                    "Startup warm-up skipped for company='%s': no embedding provider enabled warmup_on_startup.",
+                    company_short_name,
+                )
+                continue
+
+            warmed_companies.append(company_short_name)
+
+        return warmed_companies
+
+    def is_startup_warmup_enabled(self, company_short_name: str) -> bool:
+        return bool(self._get_remote_text_embedding_profiles(company_short_name, startup_only=True))
+
+    def _warmup_remote_text_embeddings(self, company_short_name: str, startup_only: bool = False) -> list[tuple[str, str, str]]:
+        profiles = self._get_remote_text_embedding_profiles(company_short_name, startup_only=startup_only)
+        if not profiles:
+            logging.debug(
+                "Warm-up skipped for company='%s': no remote embedding inference configured.",
+                company_short_name
+            )
+            return []
+
+        warmed_profiles = []
+        for model_type, config_section, tool_name in profiles:
+            try:
+                # Prime the remote model/container and download/cache model weights.
+                self.embedding_service.embed_text(
+                    company_short_name,
+                    "hello",
+                    model_type=model_type,
+                    suppress_error_logging=True,
+                )
+                logging.debug(
+                    "Warm-up primed remote embedding profile company='%s' section='%s' model_type='%s' tool='%s'.",
+                    company_short_name,
+                    config_section,
+                    model_type,
+                    tool_name,
+                )
+                warmed_profiles.append((model_type, config_section, tool_name))
+            except Exception as e:
+                logging.debug(
+                    "Warm-up failed for remote embedding profile company='%s' section='%s' model_type='%s' tool='%s': %s",
+                    company_short_name,
+                    config_section,
+                    model_type,
+                    tool_name,
+                    e,
+                )
+
+        return warmed_profiles
+
+    def _uses_remote_text_inference(self, company_short_name: str) -> bool:
+        return bool(self._get_remote_text_embedding_profiles(company_short_name))
+
+    def _get_remote_text_embedding_profiles(
+            self,
+            company_short_name: str,
+            startup_only: bool = False,
+    ) -> list[tuple[str, str, str]]:
+        profiles: list[tuple[str, str, str, dict]] = []
+
+        embedding_cfg = self.config_service.get_configuration(company_short_name, "embedding_provider") or {}
+        if isinstance(embedding_cfg, dict):
+            profiles.append(("text", "embedding_provider", embedding_cfg.get("tool_name") or "text_embeddings", embedding_cfg))
+
+        embedding_providers = self.config_service.get_configuration(company_short_name, "embedding_providers") or {}
+        if isinstance(embedding_providers, dict):
+            for model_type, embedding_provider_cfg in embedding_providers.items():
+                if not isinstance(embedding_provider_cfg, dict):
+                    continue
+                normalized_model_type = str(model_type or "").strip()
+                if not normalized_model_type:
+                    continue
+                profiles.append((
+                    normalized_model_type,
+                    f"embedding_providers.{normalized_model_type}",
+                    embedding_provider_cfg.get("tool_name") or normalized_model_type,
+                    embedding_provider_cfg,
+                ))
+
+        inference_tools = self.config_service.get_configuration(company_short_name, "inference_tools") or {}
+        if not isinstance(inference_tools, dict):
+            return []
+
+        remote_profiles = [
+            (model_type, config_section, str(tool_name or "").strip())
+            for model_type, config_section, tool_name, embedding_cfg in profiles
+            if self._is_remote_embedding_profile(company_short_name, embedding_cfg, inference_tools, str(tool_name or "").strip())
+            and (not startup_only or bool(embedding_cfg.get("warmup_on_startup")))
+        ]
+
+        # Leave the default text profile warm last. It is the profile used by rag_search.
+        return sorted(remote_profiles, key=lambda item: 1 if item[0] == "text" else 0)
+
+    def _is_remote_embedding_profile(
+            self,
+            company_short_name: str,
+            embedding_cfg: dict,
+            inference_tools: dict,
+            tool_name: str,
+    ) -> bool:
+        provider = (embedding_cfg.get("provider") or "").strip().lower()
+        if provider != "huggingface":
+            return False
+
+        if not tool_name:
+            return False
+
+        defaults = inference_tools.get("_defaults") or {}
+        if not isinstance(defaults, dict):
+            defaults = {}
+
+        tool_cfg = inference_tools.get(tool_name) or {}
+        if not isinstance(tool_cfg, dict):
+            return False
+
+        resolved_cfg = {**defaults, **tool_cfg}
+        endpoint_url = (resolved_cfg.get("endpoint_url") or "").strip()
+        if not endpoint_url:
+            endpoint_url_secret_ref = (resolved_cfg.get("endpoint_url_secret_ref") or "").strip()
+            if endpoint_url_secret_ref:
+                endpoint_url = (
+                    resolve_secret(self.secret_provider, company_short_name, endpoint_url_secret_ref, default="") or ""
+                ).strip()
+        if not endpoint_url:
+            endpoint_url_env = (resolved_cfg.get("endpoint_url_env") or "").strip()
+            if endpoint_url_env:
+                endpoint_url = (
+                    resolve_secret(self.secret_provider, company_short_name, endpoint_url_env, default="") or ""
+                ).strip()
+
+        return bool(endpoint_url)

{iatoolkit-2.29.0 → iatoolkit-2.30.0/src/iatoolkit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: iatoolkit
-Version: 2.29.0
+Version: 2.30.0
 Summary: IAToolkit
 Author: Fernando Libedinsky
 License-Expression: MIT

iatoolkit-2.29.0/src/iatoolkit/services/warmup_service.py

@@ -1,108 +0,0 @@
-# Copyright (c) 2024 Fernando Libedinsky
-# Product: IAToolkit
-#
-# IAToolkit is open source software.
-
-from __future__ import annotations
-
-import logging
-import time
-from injector import inject
-
-from iatoolkit.common.interfaces.secret_provider import SecretProvider
-from iatoolkit.common.secret_resolver import resolve_secret
-from iatoolkit.company_registry import get_registered_companies
-from iatoolkit.services.configuration_service import ConfigurationService
-from iatoolkit.services.embedding_service import EmbeddingService
-
-
-class WarmupService:
-    """
-    Lightweight warm-up orchestrator.
-    Keep it simple: no shared state, no locking, no dedup.
-    """
-
-    @inject
-    def __init__(self,
-                 config_service: ConfigurationService,
-                 embedding_service: EmbeddingService,
-                 secret_provider: SecretProvider):
-        self.config_service = config_service
-        self.embedding_service = embedding_service
-        self.secret_provider = secret_provider
-
-    def warmup_company(self, company_short_name: str, trigger: str = "manual"):
-        start = time.perf_counter()
-        try:
-            self._warmup_remote_text_embeddings(company_short_name)
-            elapsed_ms = (time.perf_counter() - start) * 1000
-            logging.info(
-                "🔥 Warm-up done for company='%s' trigger='%s' in %.2f ms",
-                company_short_name,
-                trigger,
-                elapsed_ms
-            )
-        except Exception as e:
-            elapsed_ms = (time.perf_counter() - start) * 1000
-            logging.debug(
-                "⚠️ Warm-up failed for company='%s' trigger='%s' in %.2f ms: %s",
-                company_short_name,
-                trigger,
-                elapsed_ms,
-                e
-            )
-
-    def warmup_registered_companies(self, trigger: str = "manual"):
-        for company_short_name in get_registered_companies().keys():
-            self.warmup_company(company_short_name, trigger=trigger)
-
-    def _warmup_remote_text_embeddings(self, company_short_name: str):
-        if not self._uses_remote_text_inference(company_short_name):
-            logging.debug(
-                "Warm-up skipped for company='%s': no remote embedding inference configured.",
-                company_short_name
-            )
-            return
-
-        # Prime the remote model/container.
-        self.embedding_service.embed_text(
-            company_short_name,
-            "hello",
-            suppress_error_logging=True,
-        )
-
-    def _uses_remote_text_inference(self, company_short_name: str) -> bool:
-        embedding_cfg = self.config_service.get_configuration(company_short_name, "embedding_provider") or {}
-        provider = (embedding_cfg.get("provider") or "").strip().lower()
-        if provider != "huggingface":
-            return False
-
-        tool_name = (embedding_cfg.get("tool_name") or "text_embeddings").strip()
-        if not tool_name:
-            return False
-
-        inference_tools = self.config_service.get_configuration(company_short_name, "inference_tools") or {}
-        defaults = inference_tools.get("_defaults") or {}
-        if not isinstance(defaults, dict):
-            defaults = {}
-
-        tool_cfg = inference_tools.get(tool_name) or {}
-        if not isinstance(tool_cfg, dict):
-            return False
-
-        resolved_cfg = {**defaults, **tool_cfg}
-        endpoint_url = (resolved_cfg.get("endpoint_url") or "").strip()
-        if not endpoint_url:
-            endpoint_url_secret_ref = (resolved_cfg.get("endpoint_url_secret_ref") or "").strip()
-            if endpoint_url_secret_ref:
-                endpoint_url = (
-                    resolve_secret(self.secret_provider, company_short_name, endpoint_url_secret_ref, default="") or ""
-                ).strip()
-        if not endpoint_url:
-            endpoint_url_env = (resolved_cfg.get("endpoint_url_env") or "").strip()
-            if endpoint_url_env:
-                endpoint_url = (
-                    resolve_secret(self.secret_provider, company_short_name, endpoint_url_env, default="") or ""
-                ).strip()
-
-        return bool(endpoint_url)