iatoolkit 2.31.0__tar.gz → 2.32.1__tar.gz

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

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

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

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

{iatoolkit-2.31.0 → iatoolkit-2.32.1}/src/iatoolkit/repositories/document_repo.py

@@ -27,6 +27,18 @@ class DocumentRepo:
         self.session.commit()
         return document_image
 
+    def get_document_images(self, document_id: int, limit: int | None = None) -> list[DocumentImage]:
+        if not document_id:
+            return []
+        query = (
+            self.session.query(DocumentImage)
+            .filter(DocumentImage.document_id == document_id)
+            .order_by(DocumentImage.page.asc(), DocumentImage.image_index.asc(), DocumentImage.id.asc())
+        )
+        if limit is not None:
+            query = query.limit(max(int(limit), 0))
+        return list(query.all())
+
     def get(self, company_id, filename: str ) -> Document:
         if not company_id or not filename:
             raise IAToolkitException(IAToolkitException.ErrorType.PARAM_NOT_FILLED,

{iatoolkit-2.31.0 → iatoolkit-2.32.1}/src/iatoolkit/repositories/llm_query_repo.py

@@ -10,7 +10,6 @@ from iatoolkit.repositories.models import (
     Prompt,
     PromptAgentRole,
     PromptCategory,
-    PromptExecutionMode,
 )
 from injector import inject
 from iatoolkit.repositories.database_manager import DatabaseManager
@@ -147,16 +146,34 @@ class LLMQueryRepo:
             return self.session.query(Prompt).filter(
                 Prompt.company_id == company.id,
             ).all()
-        else:
-            # Only active prompts explicitly exposed in chat for end users.
-            return self.session.query(Prompt).filter(
-                Prompt.company_id == company.id,
-                Prompt.active.is_(True),
-                Prompt.agent_role.in_([
-                    PromptAgentRole.WORKSPACE_CHAT.value,
-                    PromptAgentRole.WORKSPACE_AGENT.value,
-                ]),
-            ).all()
+
+        # Only active prompts explicitly exposed in chat for end users.
+        prompts = self.session.query(Prompt).filter(
+            Prompt.company_id == company.id,
+            Prompt.active.is_(True),
+        ).all()
+        return [
+            prompt for prompt in prompts
+            if self._get_prompt_runtime_role(prompt) in {
+                PromptAgentRole.WORKSPACE_CHAT.value,
+                PromptAgentRole.WORKSPACE_AGENT.value,
+            }
+        ]
+
+    @staticmethod
+    def _get_prompt_runtime_role(prompt: Prompt) -> str:
+        runtime_policy = getattr(prompt, "runtime_policy", None)
+        if not isinstance(runtime_policy, dict):
+            runtime_policy = {}
+        role = str(runtime_policy.get("role") or "").strip().lower()
+        if role in {
+            PromptAgentRole.WORKSPACE_CHAT.value,
+            PromptAgentRole.WORKSPACE_AGENT.value,
+            PromptAgentRole.CHANNELS.value,
+            PromptAgentRole.OPERATIONS.value,
+        }:
+            return role
+        return PromptAgentRole.WORKSPACE_CHAT.value
 
     def create_or_update_prompt(self, new_prompt: Prompt):
         prompt = self.session.query(Prompt).filter_by(company_id=new_prompt.company_id,
@@ -167,16 +184,6 @@ class LLMQueryRepo:
             prompt.order = new_prompt.order
             if new_prompt.active is not None:
                 prompt.active = new_prompt.active
-            prompt.agent_role = (
-                getattr(new_prompt, "agent_role", None)
-                or prompt.agent_role
-                or PromptAgentRole.WORKSPACE_CHAT.value
-            )
-            prompt.execution_mode = (
-                PromptExecutionMode.CONVERSATIONAL.value
-                if prompt.agent_role == PromptAgentRole.WORKSPACE_CHAT.value
-                else PromptExecutionMode.AGENTIC.value
-            )
             prompt.filename = new_prompt.filename
             prompt.custom_fields = new_prompt.custom_fields
             prompt.output_schema = new_prompt.output_schema
@@ -191,20 +198,15 @@ class LLMQueryRepo:
             )
             prompt.attachment_fallback = new_prompt.attachment_fallback or prompt.attachment_fallback or "extract"
             prompt.llm_model = new_prompt.llm_model
-            prompt.queue_tier = getattr(new_prompt, "queue_tier", None) or prompt.queue_tier or "default"
             prompt.llm_request_options = dict(new_prompt.llm_request_options or {})
             prompt.tool_policy = dict(new_prompt.tool_policy or {})
+            prompt.runtime_policy = (
+                dict(getattr(new_prompt, "runtime_policy", None) or {})
+                or dict(getattr(prompt, "runtime_policy", None) or {})
+            )
         else:
             if new_prompt.active is None:
                 new_prompt.active = True
-            if not new_prompt.agent_role:
-                new_prompt.agent_role = PromptAgentRole.WORKSPACE_CHAT.value
-            if not new_prompt.execution_mode:
-                new_prompt.execution_mode = (
-                    PromptExecutionMode.CONVERSATIONAL.value
-                    if new_prompt.agent_role == PromptAgentRole.WORKSPACE_CHAT.value
-                    else PromptExecutionMode.AGENTIC.value
-                )
             if not new_prompt.output_schema_mode:
                 new_prompt.output_schema_mode = "best_effort"
             if not new_prompt.output_response_mode:
@@ -215,12 +217,12 @@ class LLMQueryRepo:
                 new_prompt.attachment_parser_provider = "basic"
             if not new_prompt.attachment_fallback:
                 new_prompt.attachment_fallback = "extract"
-            if not getattr(new_prompt, "queue_tier", None):
-                new_prompt.queue_tier = "default"
             if new_prompt.llm_request_options is None:
                 new_prompt.llm_request_options = {}
             if new_prompt.tool_policy is None:
                 new_prompt.tool_policy = {}
+            if getattr(new_prompt, "runtime_policy", None) is None:
+                new_prompt.runtime_policy = {}
             self.session.add(new_prompt)
             prompt = new_prompt
 

{iatoolkit-2.31.0 → iatoolkit-2.32.1}/src/iatoolkit/repositories/models.py

@@ -724,8 +724,6 @@ class Prompt(Base):
     description = Column(String, nullable=False)
     filename = Column(String, nullable=False)
     active = Column(Boolean, default=True)
-    agent_role = Column(String, default=PromptAgentRole.WORKSPACE_CHAT.value, nullable=False)
-    execution_mode = Column(String, default=PromptExecutionMode.CONVERSATIONAL.value, nullable=False)
     order = Column(Integer, nullable=True, default=0)
     category_id = Column(Integer, ForeignKey(f'{ORM_SCHEMA}.iat_prompt_categories.id', ondelete='SET NULL'), nullable=True)
     custom_fields = Column(JSON, nullable=False, default=[])
@@ -737,9 +735,9 @@ class Prompt(Base):
     attachment_parser_provider = Column(String, nullable=False, default="basic")
     attachment_fallback = Column(String, nullable=False, default="extract")
     llm_model = Column(String, nullable=True, default=None)
-    queue_tier = Column(String, nullable=False, default="default")
     llm_request_options = Column(JSON_NATIVE, nullable=False, default=dict)
     tool_policy = Column(JSON_NATIVE, nullable=False, default=dict)
+    runtime_policy = Column(JSON_NATIVE, nullable=False, default=dict)
     created_at = Column(DateTime, default=datetime.now)
 
     def to_dict(self):
@@ -754,6 +752,59 @@ class Prompt(Base):
         order_by="PromptResourceBinding.binding_order",
     )
 
+    def _runtime_policy_dict(self) -> dict:
+        return dict(self.runtime_policy or {})
+
+    def _set_runtime_policy_value(self, key: str, value):
+        policy = self._runtime_policy_dict()
+        policy["version"] = 1
+        policy[key] = value
+        self.runtime_policy = policy
+
+    @property
+    def agent_role(self) -> str:
+        role = str((self.runtime_policy or {}).get("role") or "").strip().lower()
+        if role in {item.value for item in PromptAgentRole}:
+            return role
+        return PromptAgentRole.WORKSPACE_CHAT.value
+
+    @agent_role.setter
+    def agent_role(self, value: str | None):
+        role = str(value or "").strip().lower()
+        if role not in {item.value for item in PromptAgentRole}:
+            role = PromptAgentRole.WORKSPACE_CHAT.value
+        self._set_runtime_policy_value("role", role)
+
+    @property
+    def execution_mode(self) -> str:
+        if self.agent_role == PromptAgentRole.WORKSPACE_CHAT.value:
+            return PromptExecutionMode.CONVERSATIONAL.value
+        return PromptExecutionMode.AGENTIC.value
+
+    @execution_mode.setter
+    def execution_mode(self, value: str | None):
+        # execution_mode is derived from runtime_policy.role and is intentionally not persisted.
+        return
+
+    @property
+    def queue_tier(self) -> str:
+        tier = str((self.runtime_policy or {}).get("queue_tier") or "default").strip().lower()
+        return "low" if tier == "low" else "default"
+
+    @queue_tier.setter
+    def queue_tier(self, value: str | None):
+        tier = str(value or "default").strip().lower()
+        self._set_runtime_policy_value("queue_tier", "low" if tier == "low" else "default")
+
+    @property
+    def context_policy(self) -> dict:
+        context = (self.runtime_policy or {}).get("context")
+        return dict(context or {}) if isinstance(context, dict) else {}
+
+    @context_policy.setter
+    def context_policy(self, value: dict | None):
+        self._set_runtime_policy_value("context", dict(value or {}))
+
 
 class PromptResourceBinding(Base):
     __tablename__ = "iat_prompt_resource_bindings"

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

@@ -81,6 +81,21 @@ class VSRepo:
             self.session.rollback()
             raise e
 
+    def get_document_tables(self, document_id: int, limit: int | None = None) -> list[VSDoc]:
+        if not document_id:
+            return []
+        query = (
+            self.session.query(VSDoc)
+            .filter(
+                VSDoc.document_id == document_id,
+                VSDoc.meta["source_type"].as_string() == "table",
+            )
+            .order_by(VSDoc.id.asc())
+        )
+        if limit is not None:
+            query = query.limit(max(int(limit), 0))
+        return list(query.all())
+
     def query(self,
               company_short_name: str,
               query_text: str,

{iatoolkit-2.31.0 → iatoolkit-2.32.1}/src/iatoolkit/services/configuration_service.py

@@ -777,13 +777,13 @@ class ConfigurationService:
         prompt_list, categories_config = self._get_prompt_config(config)
 
         category_set = set(categories_config)
-        allowed_execution_modes = {"conversational", "agentic"}
         allowed_agent_roles = {
             "workspace_chat",
             "workspace_agent",
             "channels",
             "operations",
         }
+        allowed_queue_tiers = {"default", "low"}
         allowed_output_schema_modes = {"best_effort", "strict"}
         allowed_output_response_modes = {"chat_compatible", "structured_only"}
         allowed_attachment_modes = {"extracted_only", "native_only", "native_plus_extracted"}
@@ -816,19 +816,24 @@ class ConfigurationService:
                     add_error(f"prompts[{i}]", "Missing required key: 'description'")
 
             prompt_cat = prompt.get("category")
-            execution_mode = str(prompt.get("execution_mode", "") or "").strip().lower()
-            if execution_mode and execution_mode not in allowed_execution_modes:
+            runtime_policy = prompt.get("runtime_policy") or {}
+            if not isinstance(runtime_policy, dict):
+                add_error(f"prompts[{i}].runtime_policy", "Must be a dictionary.")
+                continue
+
+            agent_role = str(runtime_policy.get("role", "") or "").strip().lower()
+            if agent_role and agent_role not in allowed_agent_roles:
                 add_error(
-                    f"prompts[{i}]",
-                    f"Unsupported execution_mode '{execution_mode}'. Must be one of: {sorted(allowed_execution_modes)}."
+                    f"prompts[{i}].runtime_policy.role",
+                    f"Unsupported runtime_policy.role '{agent_role}'. Must be one of: {sorted(allowed_agent_roles)}.",
                 )
                 continue
 
-            agent_role = str(prompt.get("agent_role", "") or "").strip().lower()
-            if agent_role and agent_role not in allowed_agent_roles:
+            queue_tier = str(runtime_policy.get("queue_tier", "") or "").strip().lower()
+            if queue_tier and queue_tier not in allowed_queue_tiers:
                 add_error(
-                    f"prompts[{i}].agent_role",
-                    f"Unsupported agent_role '{agent_role}'. Must be one of: {sorted(allowed_agent_roles)}.",
+                    f"prompts[{i}].runtime_policy.queue_tier",
+                    f"Unsupported runtime_policy.queue_tier '{queue_tier}'. Must be one of: {sorted(allowed_queue_tiers)}.",
                 )
                 continue
 

{iatoolkit-2.31.0 → iatoolkit-2.32.1}/src/iatoolkit/services/context_builder_service.py

@@ -129,15 +129,17 @@ class ContextBuilderService:
             except Exception:
                 schema = None
 
-        raw_execution_mode = str(getattr(prompt_obj, "execution_mode", "") or "").strip().lower()
-        execution_mode = raw_execution_mode if raw_execution_mode in {"conversational", "agentic"} else "conversational"
-        raw_agent_role = str(getattr(prompt_obj, "agent_role", "") or "").strip().lower()
-        agent_role = (
-            raw_agent_role
-            if raw_agent_role in {"workspace_chat", "workspace_agent", "channels", "operations"}
-            else "workspace_chat"
-        )
-        execution_mode = "conversational" if agent_role == "workspace_chat" else "agentic"
+        raw_runtime_policy = getattr(prompt_obj, "runtime_policy", None)
+        if isinstance(raw_runtime_policy, dict):
+            runtime_policy = PromptService.normalize_runtime_policy(raw_runtime_policy)
+        else:
+            runtime_policy = PromptService.normalize_runtime_policy({
+                "role": getattr(prompt_obj, "agent_role", None),
+                "queue_tier": getattr(prompt_obj, "queue_tier", None),
+                "context": getattr(prompt_obj, "context_policy", None),
+            })
+        agent_role = runtime_policy["role"]
+        execution_mode = PromptService.execution_mode_for_agent_role(agent_role)
 
         resource_bindings: list[dict] = []
         for binding in getattr(prompt_obj, "resource_bindings", []) or []:
@@ -163,6 +165,8 @@ class ContextBuilderService:
             "prompt_name": prompt_obj.name,
             "agent_role": agent_role,
             "execution_mode": execution_mode,
+            "runtime_policy": runtime_policy,
+            "context_policy": runtime_policy["context"],
             "schema": schema,
             "schema_yaml": prompt_obj.output_schema_yaml,
             "schema_mode": prompt_obj.output_schema_mode or "best_effort",
@@ -354,15 +358,25 @@ class ContextBuilderService:
                 collection_names=rag_collections,
             )
 
+        context_policy = self.prompt_service.normalize_context_policy(
+            resolved_contract.get("context_policy"),
+            resolved_contract.get("agent_role"),
+        )
+        company_context_policy = context_policy.get("company_context_blocks") or {}
         company_context_blocks = self.company_context_service.get_company_context_blocks(
             company_short_name,
-            include_sql_context=False,
-            include_yaml_context=False,
+            include_sql_context=bool(company_context_policy.get("sql_context")),
+            include_yaml_context=bool(company_context_policy.get("yaml_context")),
+        )
+        markdown_context = (
+            company_context_blocks.get("markdown_context")
+            if bool(company_context_policy.get("markdown_context"))
+            else ""
         )
         final_system_context = self._join_context_sections(
             rendered_sections.get("identity"),
             rendered_sections.get("business_context"),
-            company_context_blocks.get("markdown_context"),
+            markdown_context,
             rendered_sections.get("conversation_rules"),
             rendered_sections.get("retrieval_guidance"),
             collection_context,

{iatoolkit-2.31.0 → iatoolkit-2.32.1}/src/iatoolkit/services/knowledge_base_service.py

@@ -759,6 +759,99 @@ class KnowledgeBaseService:
             raise IAToolkitException(IAToolkitException.ErrorType.FILE_FORMAT_ERROR,
                         f"Error reading file content: {e}")
 
+    def get_document_images(
+        self,
+        company_short_name: str,
+        document_id: int,
+        *,
+        limit: int | None = None,
+    ) -> list[dict[str, Any]]:
+        document = self.document_repo.get_by_id(document_id)
+        if not document:
+            return []
+        document_company = str(getattr(getattr(document, "company", None), "short_name", "") or "").strip()
+        if document_company != str(company_short_name or "").strip():
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.INVALID_OPERATION,
+                "Document does not belong to the requested company.",
+            )
+
+        images = self.document_repo.get_document_images(document_id, limit=limit)
+        payloads: list[dict[str, Any]] = []
+        for image in images:
+            storage_key = str(getattr(image, "storage_key", "") or "").strip()
+            if not storage_key:
+                continue
+            content = self.storage_service.get_document_content(document_company, storage_key)
+            if not content:
+                continue
+            metadata = dict(getattr(image, "meta", None) or {})
+            filename = os.path.basename(storage_key) or f"document-{document_id}-image-{image.id}.png"
+            mime_type, _ = mimetypes.guess_type(filename)
+            if not mime_type:
+                image_format = str(metadata.get("format") or "").strip().lower()
+                mime_type = {
+                    "jpeg": "image/jpeg",
+                    "jpg": "image/jpeg",
+                    "png": "image/png",
+                    "webp": "image/webp",
+                    "gif": "image/gif",
+                }.get(image_format, "application/octet-stream")
+            payloads.append(
+                {
+                    "image_id": image.id,
+                    "document_id": document_id,
+                    "page": image.page,
+                    "image_index": image.image_index,
+                    "filename": filename,
+                    "mime_type": mime_type,
+                    "width": metadata.get("width"),
+                    "height": metadata.get("height"),
+                    "metadata": metadata,
+                    "content": content,
+                }
+            )
+        return payloads
+
+    def get_document_tables(
+        self,
+        company_short_name: str,
+        document_id: int,
+        *,
+        limit: int | None = None,
+    ) -> list[dict[str, Any]]:
+        document = self.document_repo.get_by_id(document_id)
+        if not document:
+            return []
+        document_company = str(getattr(getattr(document, "company", None), "short_name", "") or "").strip()
+        if document_company != str(company_short_name or "").strip():
+            raise IAToolkitException(
+                IAToolkitException.ErrorType.INVALID_OPERATION,
+                "Document does not belong to the requested company.",
+            )
+
+        payloads: list[dict[str, Any]] = []
+        for table in self.vs_repo.get_document_tables(document_id, limit=limit):
+            metadata = dict(getattr(table, "meta", None) or {})
+            table_text = str(getattr(table, "text", "") or "").strip()
+            if not table_text:
+                continue
+            payloads.append(
+                {
+                    "chunk_id": table.id,
+                    "document_id": document_id,
+                    "table_index": metadata.get("table_index"),
+                    "block_index": metadata.get("block_index"),
+                    "page": metadata.get("page"),
+                    "page_start": metadata.get("page_start"),
+                    "page_end": metadata.get("page_end"),
+                    "title": metadata.get("title") or metadata.get("caption_text"),
+                    "text": table_text,
+                    "table_json": metadata.get("table_json"),
+                }
+            )
+        return payloads
+
     def delete_document(self, document_id: int) -> bool:
         """
         Deletes a document and its associated vectors.