iatoolkit 2.30.0__tar.gz → 2.31.0__tar.gz

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

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

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

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

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

@@ -7,6 +7,7 @@ import ast
 import json
 import logging
 import re
+from time import perf_counter
 from typing import Dict, List
 
 from sqlalchemy import text
@@ -101,36 +102,70 @@ class VSRepo:
         Returns:
             list of documents matching the query and filters
         """
+        query_started_at = perf_counter()
+        embedding_started_at = perf_counter()
         # Generate the embedding with the query text for the specific company
         try:
             query_embedding = self.embedding_service.embed_text(company_short_name, query_text)
         except Exception as e:
-            logging.error(f"error while creating text embedding: {str(e)}")
+            logging.error(
+                "RAG vector query embedding failed: company=%s duration_ms=%s error=%s",
+                company_short_name,
+                round((perf_counter() - embedding_started_at) * 1000),
+                str(e),
+            )
             raise IAToolkitException(IAToolkitException.ErrorType.EMBEDDING_ERROR,
                                f"embedding error: {str(e)}")
+        embedding_duration_ms = round((perf_counter() - embedding_started_at) * 1000)
+        try:
+            embedding_dimensions = len(query_embedding)
+        except TypeError as exc:
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.EMBEDDING_ERROR,
+                "embedding error: query embedding has no dimensions",
+            ) from exc
+        if embedding_dimensions < 1 or embedding_dimensions > 2000:
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.EMBEDDING_ERROR,
+                f"embedding error: unsupported query embedding dimensions ({embedding_dimensions})",
+            )
 
         sql_query, params = None, None
         try:
             # Get company ID from its short name for the SQL query
+            company_lookup_started_at = perf_counter()
             company = self.session.query(Company).filter(Company.short_name == company_short_name).one_or_none()
+            company_lookup_duration_ms = round((perf_counter() - company_lookup_started_at) * 1000)
             if not company:
                 raise IAToolkitException(IAToolkitException.ErrorType.VECTOR_STORE_ERROR,
                                    f"Company with short name '{company_short_name}' not found.")
 
+            distance_metric = "l2"
+            distance_expression = (
+                f"iat_vsdocs.embedding::vector({embedding_dimensions}) "
+                f"<-> CAST(:query_embedding AS vector({embedding_dimensions}))"
+            )
+
             # build the SQL query
-            sql_query_parts = ["""
+            sql_query_parts = [f"""
                                SELECT iat_vsdocs.id, \
                                       iat_documents.filename, \
                                       iat_vsdocs.text, \
                                       iat_documents.storage_key, \
                                       iat_documents.meta,
                                       iat_documents.id, \
-                                      iat_vsdocs.meta
+                                      iat_vsdocs.meta,
+                                      ({distance_expression}) AS distance
                                FROM iat_vsdocs, \
                                     iat_documents
                                WHERE iat_vsdocs.company_id = :company_id
                                  AND iat_vsdocs.document_id = iat_documents.id \
                                """]
+            # The literal dimension lets PostgreSQL prove that a matching partial
+            # pgvector expression index can satisfy the nearest-neighbor order.
+            sql_query_parts.append(
+                f" AND vector_dims(iat_vsdocs.embedding) = {embedding_dimensions}"
+            )
 
             # query parameters
             params = {
@@ -168,15 +203,18 @@ class VSRepo:
             sql_query = "".join(sql_query_parts)
 
             # add sorting and limit of results
-            sql_query += " ORDER BY embedding <-> CAST(:query_embedding AS VECTOR) LIMIT :n_results"
+            sql_query += " ORDER BY distance ASC LIMIT :n_results"
 
             logging.debug(f"Executing SQL query: {sql_query}")
             logging.debug(f"With parameters: {params}")
 
             # execute the query
+            vector_sql_started_at = perf_counter()
+            self.session.execute(text("SET LOCAL hnsw.iterative_scan = strict_order"))
+            self.session.execute(text("SET LOCAL hnsw.ef_search = 100"))
             result = self.session.execute(text(sql_query), params)
-
             rows = result.fetchall()
+            vector_sql_duration_ms = round((perf_counter() - vector_sql_started_at) * 1000)
             vs_documents = []
 
             for row in rows:
@@ -184,6 +222,8 @@ class VSRepo:
                 doc_meta = row[4] if len(row) > 4 and row[4] is not None else {}
                 chunk_meta = row[6] if len(row) > 6 and row[6] is not None else {}
 
+                distance = float(row[7]) if len(row) > 7 and row[7] is not None else None
+
                 # get the url of the document
                 storage_key = row[3] if len(row) > 3 and row[3] is not None else None
                 url = None
@@ -198,10 +238,26 @@ class VSRepo:
                         'text': row[2],
                         'meta': doc_meta,
                         'chunk_meta': chunk_meta,
-                        'url': url
+                        'url': url,
+                        'distance': distance,
+                        'distance_metric': distance_metric,
+                        'score': self._score_from_l2_distance(distance),
                     }
                 )
 
+            logging.info(
+                "RAG vector query completed: company=%s collection_ids=%s dimensions=%s n_results=%s rows=%s "
+                "embedding_ms=%s company_lookup_ms=%s vector_sql_ms=%s total_ms=%s",
+                company_short_name,
+                normalized_collection_ids,
+                embedding_dimensions,
+                n_results,
+                len(rows),
+                embedding_duration_ms,
+                company_lookup_duration_ms,
+                vector_sql_duration_ms,
+                round((perf_counter() - query_started_at) * 1000),
+            )
             return vs_documents
 
         except Exception as e:
@@ -388,6 +444,12 @@ class VSRepo:
 
         return sql_parts, params
 
+    @staticmethod
+    def _score_from_l2_distance(distance: float | None) -> float | None:
+        if distance is None:
+            return None
+        return 1 / (1 + distance)
+
     @staticmethod
     def _normalize_metadata_filter_input(metadata_filter):
         if metadata_filter is None:

{iatoolkit-2.30.0 → iatoolkit-2.31.0}/src/iatoolkit/services/tool_service.py

@@ -374,6 +374,9 @@ class ToolService:
                 "page": page,
                 "caption_text": caption,
                 "document_id": item.get("document_id"),
+                "distance": item.get("distance"),
+                "distance_metric": item.get("distance_metric"),
+                "score": item.get("score"),
             }
             lines.append(json.dumps(header, ensure_ascii=False))
             text_content = (item.get("text") or "").strip()

{iatoolkit-2.30.0 → iatoolkit-2.31.0}/src/iatoolkit/services/warmup_service.py

@@ -59,14 +59,38 @@ class WarmupService:
     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(
+            profiles = self._get_remote_text_embedding_profiles(company_short_name, startup_only=True)
+            if not profiles:
+                logging.info(
                     "Startup warm-up skipped for company='%s': no embedding provider enabled warmup_on_startup.",
                     company_short_name,
                 )
                 continue
 
+            logging.info(
+                "Startup warm-up starting for company='%s' profiles=%s trigger='%s'.",
+                company_short_name,
+                self._format_profile_labels(profiles),
+                trigger,
+            )
+            warmed_profiles = self._warmup_remote_text_embeddings(
+                company_short_name,
+                startup_only=True,
+                profiles=profiles,
+            )
+            if not warmed_profiles:
+                logging.warning(
+                    "Startup warm-up attempted but no embedding profiles warmed for company='%s' profiles=%s.",
+                    company_short_name,
+                    self._format_profile_labels(profiles),
+                )
+                continue
+
+            logging.info(
+                "Startup warm-up warmed company='%s' profiles=%s.",
+                company_short_name,
+                self._format_profile_labels(warmed_profiles),
+            )
             warmed_companies.append(company_short_name)
 
         return warmed_companies
@@ -74,8 +98,14 @@ class WarmupService:
     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)
+    def _warmup_remote_text_embeddings(
+            self,
+            company_short_name: str,
+            startup_only: bool = False,
+            profiles: list[tuple[str, str, str]] | None = None,
+    ) -> list[tuple[str, str, str]]:
+        if profiles is None:
+            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.",
@@ -102,7 +132,8 @@ class WarmupService:
                 )
                 warmed_profiles.append((model_type, config_section, tool_name))
             except Exception as e:
-                logging.debug(
+                log_fn = logging.warning if startup_only else logging.debug
+                log_fn(
                     "Warm-up failed for remote embedding profile company='%s' section='%s' model_type='%s' tool='%s': %s",
                     company_short_name,
                     config_section,
@@ -113,6 +144,13 @@ class WarmupService:
 
         return warmed_profiles
 
+    @staticmethod
+    def _format_profile_labels(profiles: list[tuple[str, str, str]]) -> list[str]:
+        return [
+            f"{config_section}:{model_type}:{tool_name}"
+            for model_type, config_section, tool_name in profiles
+        ]
+
     def _uses_remote_text_inference(self, company_short_name: str) -> bool:
         return bool(self._get_remote_text_embedding_profiles(company_short_name))
 

{iatoolkit-2.30.0 → iatoolkit-2.31.0}/src/iatoolkit/views/rag_api_view.py

@@ -499,6 +499,9 @@ class RagApiView(MethodView):
                 "text": item.get("text"),
                 "chunk_meta": item.get("chunk_meta") or {},
                 "document_meta": item.get("meta") or {},
+                "distance": item.get("distance"),
+                "distance_metric": item.get("distance_metric"),
+                "score": item.get("score"),
             })
         return normalized
 
@@ -537,6 +540,9 @@ class RagApiView(MethodView):
                 "chunk_id": item.get("chunk_id"),
                 "source_type": (item.get("chunk_meta") or {}).get("source_type"),
                 "page": (item.get("chunk_meta") or {}).get("page") or (item.get("chunk_meta") or {}).get("page_start"),
+                "distance": item.get("distance"),
+                "distance_metric": item.get("distance_metric"),
+                "score": item.get("score"),
             }
             lines.append(json.dumps(header, ensure_ascii=False))
             if item.get("text"):

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

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