iatoolkit 2.31.0__tar.gz → 2.32.0__tar.gz

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

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

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

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

{iatoolkit-2.31.0 → iatoolkit-2.32.0}/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.0}/src/iatoolkit/repositories/llm_query_repo.py

@@ -194,6 +194,7 @@ class LLMQueryRepo:
             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.context_policy = dict(getattr(new_prompt, "context_policy", None) or {})
         else:
             if new_prompt.active is None:
                 new_prompt.active = True
@@ -221,6 +222,8 @@ class LLMQueryRepo:
                 new_prompt.llm_request_options = {}
             if new_prompt.tool_policy is None:
                 new_prompt.tool_policy = {}
+            if getattr(new_prompt, "context_policy", None) is None:
+                new_prompt.context_policy = {}
             self.session.add(new_prompt)
             prompt = new_prompt
 

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

@@ -740,6 +740,7 @@ class Prompt(Base):
     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)
+    context_policy = Column(JSON_NATIVE, nullable=False, default=dict)
     created_at = Column(DateTime, default=datetime.now)
 
     def to_dict(self):

{iatoolkit-2.31.0 → iatoolkit-2.32.0}/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.0}/src/iatoolkit/services/context_builder_service.py

@@ -163,6 +163,10 @@ class ContextBuilderService:
             "prompt_name": prompt_obj.name,
             "agent_role": agent_role,
             "execution_mode": execution_mode,
+            "context_policy": self.prompt_service.normalize_context_policy(
+                getattr(prompt_obj, "context_policy", None),
+                agent_role,
+            ),
             "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.0}/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.

{iatoolkit-2.31.0 → iatoolkit-2.32.0}/src/iatoolkit/services/prompt_service.py

@@ -56,6 +56,11 @@ class PromptService:
     AGENT_ROLE_CHANNELS = PromptAgentRole.CHANNELS.value
     AGENT_ROLE_OPERATIONS = PromptAgentRole.OPERATIONS.value
     DEFAULT_CATEGORY_LABEL = "General"
+    COMPANY_CONTEXT_BLOCK_KEYS = {
+        "markdown_context",
+        "sql_context",
+        "yaml_context",
+    }
 
     @inject
     def __init__(self,
@@ -142,6 +147,41 @@ class PromptService:
             cls.AGENT_ROLE_CHANNELS,
         }
 
+    @classmethod
+    def default_context_policy_for_agent_role(cls, agent_role: str | None) -> dict:
+        normalized = str(agent_role or "").strip().lower()
+        if normalized == cls.AGENT_ROLE_WORKSPACE_CHAT:
+            blocks = {
+                "markdown_context": True,
+                "sql_context": True,
+                "yaml_context": True,
+            }
+        else:
+            blocks = {
+                "markdown_context": True,
+                "sql_context": False,
+                "yaml_context": False,
+            }
+        return {"company_context_blocks": blocks}
+
+    @classmethod
+    def normalize_context_policy(cls, context_policy: dict | None, agent_role: str | None = None) -> dict:
+        defaults = cls.default_context_policy_for_agent_role(agent_role)
+        if not isinstance(context_policy, dict):
+            return defaults
+
+        raw_blocks = context_policy.get("company_context_blocks")
+        if raw_blocks is None:
+            raw_blocks = context_policy.get("company_context")
+        if not isinstance(raw_blocks, dict):
+            return defaults
+
+        blocks = dict(defaults["company_context_blocks"])
+        for key in cls.COMPANY_CONTEXT_BLOCK_KEYS:
+            if key in raw_blocks:
+                blocks[key] = bool(raw_blocks.get(key))
+        return {"company_context_blocks": blocks}
+
     def _normalize_output_schema_mode(self, output_schema_mode: str | None) -> str:
         candidate = str(output_schema_mode or self.OUTPUT_SCHEMA_MODE_BEST_EFFORT).strip().lower()
         allowed = {
@@ -557,6 +597,10 @@ class PromptService:
                             ),
                             'llm_request_options': dict(getattr(p, 'llm_request_options', None) or {}),
                             'tool_policy': dict(getattr(p, 'tool_policy', None) or {}),
+                            'context_policy': self.normalize_context_policy(
+                                getattr(p, 'context_policy', None),
+                                getattr(p, 'agent_role', None),
+                            ),
                         }
                         for p in prompts
                     ]
@@ -687,6 +731,10 @@ class PromptService:
             ),
             llm_request_options=llm_request_options,
             tool_policy=tool_policy,
+            context_policy=self.normalize_context_policy(
+                data.get("context_policy"),
+                agent_role,
+            ),
         )
         self.llm_query_repo.create_or_update_prompt(new_prompt)
 
@@ -963,6 +1011,10 @@ class PromptService:
                         self._extract_llm_request_options_payload(prompt_data)
                     ),
                     tool_policy=self._normalize_tool_policy(prompt_data.get("tool_policy")),
+                    context_policy=self.normalize_context_policy(
+                        prompt_data.get("context_policy"),
+                        agent_role,
+                    ),
                 )
 
                 self.llm_query_repo.create_or_update_prompt(new_prompt)

{iatoolkit-2.31.0 → iatoolkit-2.32.0}/src/iatoolkit/services/query_service.py

@@ -15,6 +15,7 @@ from iatoolkit.services.context_builder_service import ContextBuilderService
 from iatoolkit.services.attachment_policy_service import AttachmentPolicyService
 from iatoolkit.services.memory_lookup_policy_service import MemoryLookupPolicyService
 from iatoolkit.services.telemetry_service import NoopTelemetryService
+from iatoolkit.services.prompt_service import PromptService
 from iatoolkit.common.model_registry import ModelRegistry
 from injector import inject
 import logging
@@ -341,6 +342,10 @@ class QueryService:
                     })
             contract["resource_bindings"] = normalized_resource_bindings
             contract["tool_policy"] = self._normalize_prompt_tool_policy(contract.get("tool_policy"))
+            contract["context_policy"] = PromptService.normalize_context_policy(
+                contract.get("context_policy"),
+                contract.get("agent_role"),
+            )
             contract["provider"] = self._get_provider(company.short_name, contract.get("llm_model") or "")
             return contract
         except Exception as e:
@@ -1875,6 +1880,10 @@ class QueryService:
                 "provider": provider,
                 "effective_model": effective_model,
                 "tool_policy": self._normalize_prompt_tool_policy(prompt_output_contract.get("tool_policy")),
+                "context_policy": PromptService.normalize_context_policy(
+                    prompt_output_contract.get("context_policy"),
+                    prompt_output_contract.get("agent_role"),
+                ),
                 "resource_bindings": list(prompt_output_contract.get("resource_bindings") or []),
                 "selected_system_prompt_keys": selected_system_prompt_keys,
                 "tool_names": self._get_tool_names(tools),

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

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