unique_toolkit 1.28.8__py3-none-any.whl → 1.33.3__py3-none-any.whl

unique_toolkit/__init__.py

@@ -1,6 +1,10 @@
 # Re-export commonly used classes for easier imports
 from unique_toolkit.chat import ChatService
 from unique_toolkit.content import ContentService
+from unique_toolkit.data_extraction import (
+    StructuredOutputDataExtractor,
+    StructuredOutputDataExtractorConfig,
+)
 from unique_toolkit.embedding import EmbeddingService
 from unique_toolkit.framework_utilities.openai.client import (
     get_async_openai_client,
@@ -26,17 +30,19 @@ except ImportError:
 # You can add other classes you frequently use here as well
 
 __all__ = [
-    "LanguageModelService",
-    "LanguageModelMessages",
-    "LanguageModelName",
-    "LanguageModelToolDescription",
     "ChatService",
     "ContentService",
     "EmbeddingService",
-    "ShortTermMemoryService",
-    "KnowledgeBaseService",
     "get_openai_client",
     "get_async_openai_client",
+    "KnowledgeBaseService",
+    "LanguageModelMessages",
+    "LanguageModelName",
+    "LanguageModelService",
+    "LanguageModelToolDescription",
+    "ShortTermMemoryService",
+    "StructuredOutputDataExtractor",
+    "StructuredOutputDataExtractorConfig",
 ]
 
 # Add langchain-specific exports if available

unique_toolkit/_common/docx_generator/service.py

@@ -14,7 +14,6 @@ from unique_toolkit._common.docx_generator.schemas import (
     RunField,
     RunsField,
 )
-from unique_toolkit.services import KnowledgeBaseService
 
 generator_dir_path = Path(__file__).resolve().parent
 
@@ -23,13 +22,9 @@ _LOGGER = logging.getLogger(__name__)
 
 
 class DocxGeneratorService:
-    def __init__(
-        self,
-        knowledge_base_service: KnowledgeBaseService,
-        config: DocxGeneratorConfig,
-    ):
-        self._knowledge_base_service = knowledge_base_service
+    def __init__(self, config: DocxGeneratorConfig, *, template: bytes | None = None):
         self._config = config
+        self._template = template
 
     @staticmethod
     def parse_markdown_to_list_content_fields(
@@ -190,10 +185,7 @@ class DocxGeneratorService:
             subdoc_content (list[HeadingField | ParagraphField | RunsField]): The content to be added to the docx file.
             fields (dict): Other fields to be added to the docx file. Defaults to None.
         """
-
-        docx_template_object = self._get_template(self._config.template_content_id)
-
-        doc = DocxTemplate(io.BytesIO(docx_template_object))
+        doc = DocxTemplate(io.BytesIO(self.template))
 
         try:
             content = {}
@@ -218,27 +210,7 @@ class DocxGeneratorService:
             _LOGGER.error(f"Error generating docx: {e}")
             return None
 
-    def _get_template(self, template_content_id: str):
-        try:
-            if template_content_id:
-                _LOGGER.info(
-                    f"Downloading template from content ID: {template_content_id}"
-                )
-                file_content = self._knowledge_base_service.download_content_to_bytes(
-                    content_id=template_content_id
-                )
-            else:
-                _LOGGER.info("No template content ID provided. Using default template.")
-                file_content = self._get_default_template()
-        except Exception as e:
-            _LOGGER.warning(
-                f"An error occurred while downloading the template {e}. Make sure the template content ID is valid. Falling back to default template."
-            )
-            file_content = self._get_default_template()
-
-        return file_content
-
-    def _get_default_template(self):
+    def _get_default_template(self) -> bytes:
         generator_dir_path = Path(__file__).resolve().parent
         path = generator_dir_path / "template" / "Doc Template.docx"
 
@@ -247,3 +219,7 @@ class DocxGeneratorService:
         _LOGGER.info("Template downloaded from default template")
 
         return file_content
+
+    @property
+    def template(self) -> bytes:
+        return self._template or self._get_default_template()

unique_toolkit/_common/utils/jinja/helpers.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+
+
+def get_parent_dir(file_path: str) -> Path:
+    return Path(file_path).parent
+
+
+def load_template(parent_dir: Path, template_name: str) -> str:
+    with open(parent_dir / template_name, "r") as file:
+        return file.read().strip()

unique_toolkit/_common/utils/jinja/render.py

@@ -0,0 +1,18 @@
+from typing import Any
+
+from jinja2 import Template
+
+from unique_toolkit._common.utils.jinja.schema import Jinja2PromptParams
+
+
+def render_template(
+    template: str, params: Jinja2PromptParams | dict[str, Any] | None = None, **kwargs
+) -> str:
+    params = params or {}
+
+    if isinstance(params, Jinja2PromptParams):
+        params = params.model_dump(exclude_none=True, mode="json")
+
+    params.update(kwargs)
+
+    return Template(template, lstrip_blocks=True).render(**params)

unique_toolkit/_common/utils/jinja/schema.py

@@ -0,0 +1,65 @@
+from datetime import date, datetime
+from typing import Annotated, Any
+
+from jinja2 import Template
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    SerializerFunctionWrapHandler,
+    WrapSerializer,
+)
+
+from unique_toolkit.agentic.tools.tool import Tool
+
+
+class Jinja2PromptParams(BaseModel):
+    model_config = ConfigDict(str_strip_whitespace=True)
+
+    def render_template(self, template: str) -> str:
+        params = self.model_dump(exclude_none=True, mode="json")
+
+        return Template(template, lstrip_blocks=True).render(**params)
+
+
+class ToolPromptParams(Jinja2PromptParams):
+    name: str
+    tool_description_for_system_prompt: str = ""
+    tool_format_information_for_system_prompt: str = ""
+    tool_format_reminder_for_user_prompt: str = ""
+
+    @classmethod
+    def from_tool(cls, tool: Tool) -> "ToolPromptParams":
+        return cls(
+            name=tool.name,
+            tool_description_for_system_prompt=tool.tool_description_for_system_prompt(),
+            tool_format_information_for_system_prompt=tool.tool_format_information_for_system_prompt(),
+            tool_format_reminder_for_user_prompt=tool.tool_format_reminder_for_user_prompt(),
+        )
+
+
+def serialize_iso8601_date(v: Any, handler: SerializerFunctionWrapHandler) -> str:
+    if isinstance(v, date):
+        return v.isoformat()
+    return handler(v)
+
+
+ISO8601Date = Annotated[
+    date,
+    WrapSerializer(serialize_iso8601_date, return_type=str),
+]
+
+
+class AgentSystemPromptParams(Jinja2PromptParams):
+    info_cutoff_at: ISO8601Date | None
+    current_date: ISO8601Date = Field(default_factory=lambda: datetime.now().date())
+    tools: list[ToolPromptParams]
+    used_tools: list[ToolPromptParams]
+    add_citation_appendix: bool = True
+    max_tools_per_iteration: int
+    max_loop_iterations: int
+    current_iteration: int
+
+
+class AgentUserPromptParams(Jinja2PromptParams):
+    user_prompt: str

unique_toolkit/_common/utils/jinja/utils.py

@@ -0,0 +1,80 @@
+from jinja2 import Environment
+from jinja2.nodes import Const, Getattr, Getitem, Name
+from pydantic import BaseModel
+
+
+class TemplateValidationResult(BaseModel):
+    is_valid: bool
+    missing_placeholders: list[str]
+    optional_placeholders: list[str]
+    unexpected_placeholders: list[str]
+
+
+def _get_nested_variables(node):
+    """Recursively extract all variable references from a Jinja2 AST node."""
+    variables = set()
+
+    if isinstance(node, Name):
+        variables.add(node.name)
+    elif isinstance(node, (Getattr, Getitem)):
+        # For nested attributes like example.category
+        if isinstance(node.node, Name):
+            if isinstance(node, Getattr):
+                variables.add(f"{node.node.name}.{node.attr}")
+            else:  # Getitem
+                if isinstance(node.arg, Const):
+                    variables.add(f"{node.node.name}.{node.arg.value}")
+                else:
+                    # For dynamic indices, just use the base variable
+                    variables.add(node.node.name)
+        # Recursively process nested nodes
+        variables.update(_get_nested_variables(node.node))
+
+    # Process child nodes
+    for child in node.iter_child_nodes():
+        variables.update(_get_nested_variables(child))
+
+    return variables
+
+
+def validate_template_placeholders(
+    template_content: str,
+    required_placeholders: set[str],
+    optional_placeholders: set[str],
+) -> TemplateValidationResult:
+    """
+    Validates that all required placeholders in the template are present.
+    Handles both top-level and nested variables (e.g. example.category).
+
+    Args:
+        template_content (str): The content of the Jinja template
+        required_placeholders (set[str]): Set of required placeholder names
+        optional_placeholders (set[str]): Set of optional placeholder names
+
+    Returns:
+        TemplateValidationResult: A result object containing validation information
+    """
+    # Create a Jinja environment
+    env = Environment()
+
+    # Parse the template and get all variables including nested ones
+    ast = env.parse(template_content)
+    template_vars = _get_nested_variables(ast)
+
+    # Check for missing required placeholders
+    missing_placeholders = required_placeholders - template_vars
+
+    # Check for optional placeholders present
+    present_optional = optional_placeholders & template_vars
+
+    # Check for any unexpected placeholders
+    unexpected_placeholders = template_vars - (
+        required_placeholders | optional_placeholders
+    )
+
+    return TemplateValidationResult(
+        is_valid=len(missing_placeholders) == 0,
+        missing_placeholders=sorted(list(missing_placeholders)),
+        optional_placeholders=sorted(list(present_optional)),
+        unexpected_placeholders=sorted(list(unexpected_placeholders)),
+    )

unique_toolkit/agentic/message_log_manager/service.py

@@ -5,6 +5,7 @@ Target of the method is to extend the step tracking on all levels of the tool.
 """
 
 from collections import defaultdict
+from logging import getLogger
 
 from unique_toolkit.chat.schemas import (
     MessageLogDetails,
@@ -14,6 +15,8 @@ from unique_toolkit.chat.schemas import (
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content.schemas import ContentReference
 
+_LOGGER = getLogger(__name__)
+
 # Per-request counters for message log ordering - keyed by message_id
 # This is a mandatory global variable since we have in the system a bug which makes it impossible to use it as a proper class variable.
 _request_counters = defaultdict(int)
@@ -71,6 +74,12 @@ class MessageStepLogger:
         """
 
         # Creating a new message log entry with the found hits.
+        if not self._chat_service._assistant_message_id:
+            _LOGGER.warning(
+                "Assistant message id is not set. Skipping message log entry creation."
+            )
+            return
+
         _ = self._chat_service.create_message_log(
             message_id=self._chat_service._assistant_message_id,
             text=text,

unique_toolkit/agentic/tools/a2a/postprocessing/_display_utils.py

@@ -1,6 +1,7 @@
 import re
-from typing import Literal
+from typing import Literal, NamedTuple
 
+from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
     SubAgentDisplayConfig,
     SubAgentResponseDisplayMode,
@@ -84,6 +85,12 @@ def _prepare_title_template(
     return display_title_template.replace("{}", "{%s}" % display_name_placeholder)
 
 
+def _clean_linebreaks(text: str) -> str:
+    text = text.strip()
+    text = re.sub(r"^(<br>)*|(<br>)*$", "", text)
+    return text
+
+
 def _get_display_template(
     mode: SubAgentResponseDisplayMode,
     add_quote_border: bool,
@@ -126,7 +133,7 @@ def _get_display_template(
     if add_block_border:
         template = _wrap_with_block_border(template)
 
-    return template
+    return _clean_linebreaks(template)
 
 
 def _get_display_removal_re(
@@ -150,10 +157,51 @@ def _get_display_removal_re(
     return re.compile(pattern, flags=re.DOTALL)
 
 
+class SubAgentAnswerPart(NamedTuple):
+    matching_text: str  # Matching text as found in the answer
+    formatted_text: str  # Formatted text to be displayed
+
+
+def get_sub_agent_answer_parts(
+    answer: str,
+    display_config: SubAgentDisplayConfig,
+) -> list[SubAgentAnswerPart]:
+    if display_config.mode == SubAgentResponseDisplayMode.HIDDEN:
+        return []
+
+    if len(display_config.answer_substrings_config) == 0:
+        return [SubAgentAnswerPart(matching_text=answer, formatted_text=answer)]
+
+    substrings = []
+    for config in display_config.answer_substrings_config:
+        for match in config.regexp.finditer(answer):
+            text = match.group(0)
+            substrings.append(
+                SubAgentAnswerPart(
+                    matching_text=text,
+                    formatted_text=config.display_template.format(text),
+                )
+            )
+
+    return substrings
+
+
+def get_sub_agent_answer_from_parts(
+    answer_parts: list[SubAgentAnswerPart],
+    config: SubAgentDisplayConfig,
+) -> str:
+    return render_template(
+        config.answer_substrings_jinja_template,
+        {
+            "substrings": [answer.formatted_text for answer in answer_parts],
+        },
+    )
+
+
 def get_sub_agent_answer_display(
     display_name: str,
     display_config: SubAgentDisplayConfig,
-    answer: str,
+    answer: str | list[SubAgentAnswerPart],
     assistant_id: str,
 ) -> str:
     template = _get_display_template(
@@ -162,6 +210,13 @@ def get_sub_agent_answer_display(
         add_block_border=display_config.add_block_border,
         display_title_template=display_config.display_title_template,
     )
+
+    if isinstance(answer, list):
+        answer = get_sub_agent_answer_from_parts(
+            answer_parts=answer,
+            config=display_config,
+        )
+
     return template.format(
         display_name=display_name, answer=answer, assistant_id=assistant_id
     )

unique_toolkit/agentic/tools/a2a/postprocessing/_ref_utils.py

@@ -1,3 +1,4 @@
+import re
 from typing import Callable, Iterable, Mapping, Sequence
 
 from unique_toolkit._common.referencing import get_reference_pattern
@@ -50,6 +51,16 @@ def add_content_refs(
     return message_refs
 
 
+def remove_unused_refs(
+    references: Sequence[ContentReference],
+    text: str,
+    ref_pattern_f: Callable[[int], str] = get_reference_pattern,
+) -> list[ContentReference]:
+    return [
+        ref for ref in references if re.search(ref_pattern_f(ref.sequence_number), text)
+    ]
+
+
 def add_content_refs_and_replace_in_text(
     message_text: str,
     message_refs: Sequence[ContentReference],

unique_toolkit/agentic/tools/a2a/postprocessing/config.py

@@ -1,3 +1,4 @@
+import re
 from enum import StrEnum
 from typing import Literal
 
@@ -13,6 +14,25 @@ class SubAgentResponseDisplayMode(StrEnum):
     PLAIN = "plain"
 
 
+class SubAgentAnswerSubstringConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    regexp: re.Pattern[str] = Field(
+        description="The regular expression to use to extract the substring. The first capture group will always be used.",
+    )
+    display_template: str = Field(
+        default="{}",
+        description="The template to use to display the substring. It should contain exactly one empty placeholder '{}' for the substring.",
+    )
+
+
+_ANSWER_SUBSTRINGS_JINJA_TEMPLATE = """
+{% for substring in substrings %}
+{{ substring }}
+{% endfor %}
+""".strip()
+
+
 class SubAgentDisplayConfig(BaseModel):
     model_config = get_configuration_dict()
 
@@ -43,3 +63,16 @@ class SubAgentDisplayConfig(BaseModel):
         default="before",
         description="The position of the sub agent response in the main agent response.",
     )
+    force_include_references: bool = Field(
+        default=False,
+        description="If set, the sub agent references will be added to the main agent response references even in not mentioned in the main agent response text.",
+    )
+
+    answer_substrings_config: list[SubAgentAnswerSubstringConfig] = Field(
+        default=[],
+        description="If set, only parts of the answer matching the provided regular expressions will be displayed.",
+    )
+    answer_substrings_jinja_template: str = Field(
+        default=_ANSWER_SUBSTRINGS_JINJA_TEMPLATE,
+        description="The template to use in order to format the different answer substrings, if any.",
+    )

unique_toolkit/agentic/tools/a2a/postprocessing/display.py

@@ -7,13 +7,18 @@ import unique_sdk
 from pydantic import BaseModel, Field
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
 from unique_toolkit.agentic.tools.a2a.postprocessing._display_utils import (
+    SubAgentAnswerPart,
     get_sub_agent_answer_display,
+    get_sub_agent_answer_from_parts,
+    get_sub_agent_answer_parts,
     remove_sub_agent_answer_from_text,
 )
 from unique_toolkit.agentic.tools.a2a.postprocessing._ref_utils import (
     add_content_refs_and_replace_in_text,
+    remove_unused_refs,
 )
 from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
     SubAgentDisplayConfig,
@@ -37,11 +42,26 @@ class SubAgentDisplaySpec(NamedTuple):
     display_config: SubAgentDisplayConfig
 
 
+_ANSWERS_JINJA_TEMPLATE = """
+{% for answer in answers %}
+{{ answer }}
+{% endfor %}
+""".strip()
+
+
 class SubAgentResponsesPostprocessorConfig(BaseModel):
     model_config = get_configuration_dict()
 
     sleep_time_before_update: float = Field(
-        default=1, description="Time to sleep before updating the main agent message."
+        default=0, description="Time to sleep before updating the main agent message."
+    )
+    answers_jinja_template: str = Field(
+        default=_ANSWERS_JINJA_TEMPLATE,
+        description="The template to use to display the sub agent answers.",
+    )
+    filter_duplicate_answers: bool = Field(
+        default=True,
+        description="If set, duplicate answers will be filtered out.",
     )
 
 
@@ -60,6 +80,8 @@ class SubAgentResponsesDisplayPostprocessor(Postprocessor):
             display_spec.assistant_id: display_spec
             for display_spec in display_specs
             if display_spec.display_config.mode != SubAgentResponseDisplayMode.HIDDEN
+            # We should keep track of these messages even if they are hidden
+            or display_spec.display_config.force_include_references
         }
 
     @override
@@ -93,19 +115,23 @@ class SubAgentResponsesDisplayPostprocessor(Postprocessor):
 
         answers_displayed_before = []
         answers_displayed_after = []
+        all_displayed_answers = set()
 
         for assistant_id, responses in displayed_sub_agent_responses.items():
+            tool_info = self._display_specs[assistant_id]
+            tool_name = tool_info.display_name
+
             for response in responses:
                 message = response.message
-                tool_info = self._display_specs[assistant_id]
-
-                _add_response_references_to_message_in_place(
-                    loop_response=loop_response, response=message
-                )
 
-                display_name = tool_info.display_name
-                if len(responses) > 1:
-                    display_name += f" {response.sequence_number}"
+                if tool_info.display_config.mode == SubAgentResponseDisplayMode.HIDDEN:
+                    # Add references and continue
+                    _add_response_references_to_message_in_place(
+                        loop_response=loop_response,
+                        response=message,
+                        remove_unused_references=False,
+                    )
+                    continue
 
                 if message["text"] is None:
                     logger.warning(
@@ -113,11 +139,44 @@ class SubAgentResponsesDisplayPostprocessor(Postprocessor):
                         assistant_id,
                         response.sequence_number,
                     )
+                    continue
+
+                answer_parts = get_sub_agent_answer_parts(
+                    answer=message["text"],
+                    display_config=tool_info.display_config,
+                )
+
+                if self._config.filter_duplicate_answers:
+                    answer_parts, all_displayed_answers = (
+                        _filter_and_update_duplicate_answers(
+                            answers=answer_parts,
+                            existing_answers=all_displayed_answers,
+                        )
+                    )
+
+                answer = get_sub_agent_answer_from_parts(
+                    answer_parts=answer_parts,
+                    config=tool_info.display_config,
+                )
+                message["text"] = answer
+
+                _add_response_references_to_message_in_place(
+                    loop_response=loop_response,
+                    response=message,
+                    remove_unused_references=not tool_info.display_config.force_include_references,
+                )
+
+                if len(answer_parts) == 0:
+                    continue
+
+                display_name = tool_name
+                if len(responses) > 1:
+                    display_name = tool_name + f" {response.sequence_number}"
 
                 answer = get_sub_agent_answer_display(
                     display_name=display_name,
                     display_config=tool_info.display_config,
-                    answer=message["text"] or "",
+                    answer=answer,
                     assistant_id=assistant_id,
                 )
 
@@ -130,6 +189,7 @@ class SubAgentResponsesDisplayPostprocessor(Postprocessor):
             text=loop_response.message.text,
             answers_before=answers_displayed_before,
             answers_after=answers_displayed_after,
+            template=self._config.answers_jinja_template,
         )
 
         return True
@@ -146,7 +206,9 @@ class SubAgentResponsesDisplayPostprocessor(Postprocessor):
 
 
 def _add_response_references_to_message_in_place(
-    loop_response: LanguageModelStreamResponse, response: unique_sdk.Space.Message
+    loop_response: LanguageModelStreamResponse,
+    response: unique_sdk.Space.Message,
+    remove_unused_references: bool = True,
 ) -> None:
     references = response["references"]
     text = response["text"]
@@ -156,6 +218,12 @@ def _add_response_references_to_message_in_place(
 
     content_refs = [ContentReference.from_sdk_reference(ref) for ref in references]
 
+    if remove_unused_references:
+        content_refs = remove_unused_refs(
+            references=content_refs,
+            text=text,
+        )
+
     text, refs = add_content_refs_and_replace_in_text(
         message_text=text,
         message_refs=loop_response.message.references,
@@ -170,11 +238,27 @@ def _get_final_answer_display(
     text: str,
     answers_before: list[str],
     answers_after: list[str],
-    sep: str = "<br>\n\n",
+    template: str = _ANSWERS_JINJA_TEMPLATE,
 ) -> str:
     if len(answers_before) > 0:
-        text = sep.join(answers_before) + sep + text
+        text = render_template(template, {"answers": answers_before}) + text
 
     if len(answers_after) > 0:
-        text = text + sep + sep.join(answers_after)
-    return text
+        text = text + render_template(template, {"answers": answers_after})
+
+    return text.strip()
+
+
+def _filter_and_update_duplicate_answers(
+    answers: list[SubAgentAnswerPart],
+    existing_answers: set[str],
+) -> tuple[list[SubAgentAnswerPart], set[str]]:
+    new_answers = []
+
+    for answer in answers:
+        if answer.matching_text in existing_answers:
+            continue
+        existing_answers.add(answer.matching_text)
+        new_answers.append(answer)
+
+    return new_answers, existing_answers