unique_toolkit 1.32.1__tar.gz → 1.33.1__tar.gz

{unique_toolkit-1.32.1 → unique_toolkit-1.33.1}/CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.33.1] - 2025-12-01
+- Add `data_extraction` to unique_toolkit
+
+## [1.33.0] - 2025-11-28
+- Add support for system reminders in sub agent responses.
+
 ## [1.32.1] - 2025-12-01
 - Added documentation for the toolkit,some missing type hints and doc string fixes.
 

{unique_toolkit-1.32.1 → unique_toolkit-1.33.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.32.1
+Version: 1.33.1
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -121,6 +121,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.33.1] - 2025-12-01
+- Add `data_extraction` to unique_toolkit
+
+## [1.33.0] - 2025-11-28
+- Add support for system reminders in sub agent responses.
+
 ## [1.32.1] - 2025-12-01
 - Added documentation for the toolkit,some missing type hints and doc string fixes.
 

{unique_toolkit-1.32.1 → unique_toolkit-1.33.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "1.32.1"
+version = "1.33.1"
 description = ""
 authors = [
     "Cedric Klinkert <cedric.klinkert@unique.ch>",

{unique_toolkit-1.32.1 → unique_toolkit-1.33.1}/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-1.32.1 → unique_toolkit-1.33.1}/unique_toolkit/_common/utils/jinja/render.py

@@ -5,8 +5,14 @@ from jinja2 import Template
 from unique_toolkit._common.utils.jinja.schema import Jinja2PromptParams
 
 
-def render_template(template: str, params: Jinja2PromptParams | dict[str, Any]) -> str:
+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-1.32.1 → unique_toolkit-1.33.1}/unique_toolkit/agentic/tools/a2a/tool/config.py

@@ -1,10 +1,74 @@
-from typing import Literal
+import re
+from enum import StrEnum
+from typing import Annotated, Generic, Literal, TypeVar
 
 from pydantic import Field
+from pydantic.main import BaseModel
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
 from unique_toolkit.agentic.tools.schemas import BaseToolConfig
 
+
+class SubAgentSystemReminderType(StrEnum):
+    FIXED = "fixed"
+    REGEXP = "regexp"
+    REFERENCE = "reference"
+
+
+T = TypeVar("T", bound=SubAgentSystemReminderType)
+
+
+class SystemReminderConfig(BaseModel, Generic[T]):
+    model_config = get_configuration_dict()
+
+    type: T
+
+
+_SYSTEM_REMINDER_FIELD_DESCRIPTION = """
+The reminder to add to the tool response. The reminder can be a Jinja template and can contain the following placeholders:
+- {{ display_name }}: The display name of the sub agent.
+- {{ tool_name }}: The tool name. 
+""".strip()
+
+
+class ReferenceSystemReminderConfig(SystemReminderConfig):
+    type: Literal[SubAgentSystemReminderType.REFERENCE] = (
+        SubAgentSystemReminderType.REFERENCE
+    )
+    reminder: str = Field(
+        default="Rememeber to properly reference EACH fact from sub agent {{ display_name }}'s response with the correct format INLINE.",
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+class FixedSystemReminderConfig(SystemReminderConfig):
+    type: Literal[SubAgentSystemReminderType.FIXED] = SubAgentSystemReminderType.FIXED
+    reminder: str = Field(
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+_REGEXP_DETECTED_REMINDER_FIELD_DESCRIPTION = """
+The reminder to add to the tool response. The reminder can be a Jinja template and can contain the following placeholders:
+- {{ display_name }}: The display name of the sub agent.
+- {{ tool_name }}: The tool name. 
+- {{ text_matches }}: Will be replaced with the portions of the text that triggered the reminder.
+""".strip()
+
+
+class RegExpDetectedSystemReminderConfig(SystemReminderConfig):
+    """A system reminder that is only added if the sub agent response matches a regular expression."""
+
+    type: Literal[SubAgentSystemReminderType.REGEXP] = SubAgentSystemReminderType.REGEXP
+
+    regexp: re.Pattern[str] = Field(
+        description="The regular expression to use to detect whether the system reminder should be added.",
+    )
+    reminder: str = Field(
+        description=_REGEXP_DETECTED_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
 DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
 This is the message that will be sent to the sub-agent.
 """.strip()
@@ -80,3 +144,15 @@ class SubAgentToolConfig(BaseToolConfig):
         default=False,
         description="If set, the sub-agent response will be interpreted as a list of content chunks.",
     )
+
+    system_reminders_config: list[
+        Annotated[
+            FixedSystemReminderConfig
+            | RegExpDetectedSystemReminderConfig
+            | ReferenceSystemReminderConfig,
+            Field(discriminator="type"),
+        ]
+    ] = Field(
+        default=[],
+        description="Configuration for the system reminders to add to the tool response.",
+    )

{unique_toolkit-1.32.1 → unique_toolkit-1.33.1}/unique_toolkit/agentic/tools/a2a/tool/service.py

@@ -4,7 +4,7 @@ import json
 import logging
 import re
 from datetime import datetime
-from typing import override
+from typing import cast, override
 
 import unique_sdk
 from pydantic import Field, TypeAdapter, create_model
@@ -15,6 +15,7 @@ from unique_toolkit._common.referencing import (
     remove_all_refs,
     replace_ref_number,
 )
+from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.tools.a2a.response_watcher import SubAgentResponseWatcher
 from unique_toolkit.agentic.tools.a2a.tool._memory import (
@@ -25,6 +26,8 @@ from unique_toolkit.agentic.tools.a2a.tool._schema import (
     SubAgentToolInput,
 )
 from unique_toolkit.agentic.tools.a2a.tool.config import (
+    RegExpDetectedSystemReminderConfig,
+    SubAgentSystemReminderType,
     SubAgentToolConfig,
 )
 from unique_toolkit.agentic.tools.factory import ToolFactory
@@ -215,12 +218,63 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             )
 
             return ToolCallResponse(
-                id=tool_call.id,  # type: ignore
+                id=tool_call.id,
                 name=tool_call.name,
-                content=content,
+                content=_format_response(
+                    tool_name=self.name,
+                    text=content,
+                    system_reminders=self._get_system_reminders(response),
+                ),
                 content_chunks=content_chunks,
             )
 
+    def _get_system_reminders(self, message: unique_sdk.Space.Message) -> list[str]:
+        reminders = []
+        for reminder_config in self.config.system_reminders_config:
+            if reminder_config.type == SubAgentSystemReminderType.FIXED:
+                reminders.append(
+                    render_template(
+                        reminder_config.reminder,
+                        display_name=self.display_name(),
+                        tool_name=self.name,
+                    )
+                )
+            elif (
+                reminder_config.type == SubAgentSystemReminderType.REFERENCE
+                and self.config.use_sub_agent_references
+                and message["references"] is not None
+                and len(message["references"]) > 0
+            ):
+                reminders.append(
+                    render_template(
+                        reminder_config.reminder,
+                        display_name=self.display_name(),
+                        tool_name=self.name,
+                    )
+                )
+            elif (
+                reminder_config.type == SubAgentSystemReminderType.REGEXP
+                and message["text"] is not None
+            ):
+                reminder_config = cast(
+                    RegExpDetectedSystemReminderConfig, reminder_config
+                )
+                text_matches = [
+                    match.group(0)
+                    for match in reminder_config.regexp.finditer(message["text"])
+                ]
+                if len(text_matches) > 0:
+                    reminders.append(
+                        render_template(
+                            reminder_config.reminder,
+                            display_name=self.display_name(),
+                            tool_name=self.name,
+                            text_matches=text_matches,
+                        )
+                    )
+
+        return reminders
+
     async def _get_chat_id(self) -> str | None:
         if not self.config.reuse_chat:
             return None
@@ -326,4 +380,14 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             ) from e
 
 
+def _format_response(tool_name: str, text: str, system_reminders: list[str]) -> str:
+    if len(system_reminders) == 0:
+        return text
+
+    reponse_key = f"{tool_name} response"
+    response = {reponse_key: text, "SYSTEM_REMINDERS": system_reminders}
+
+    return json.dumps(response, indent=2)
+
+
 ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)

unique_toolkit-1.33.1/unique_toolkit/data_extraction/README.md

@@ -0,0 +1,96 @@
+# Data Extraction Module
+
+This module provides a flexible framework for extracting structured data from text using language models. It supports both basic and augmented data extraction capabilities.
+
+## Overview
+
+The module consists of two main components:
+
+1. **Basic Data Extraction**: Uses language models to extract structured data from text based on a provided schema.
+2. **Augmented Data Extraction**: Extends basic extraction by adding extra fields to the output schema while maintaining the original data structure.
+
+## Components
+
+### Base Classes
+
+- `BaseDataExtractor`: Abstract base class that defines the interface for data extraction
+- `BaseDataExtractionResult`: Generic base class for extraction results
+
+### Basic Extraction
+
+- `StructuredOutputDataExtractor`: Implements basic data extraction using language models
+- `StructuredOutputDataExtractorConfig`: Configuration for the basic extractor
+
+### Augmented Extraction
+
+- `AugmentedDataExtractor`: Extends basic extraction with additional fields
+- `AugmentedDataExtractionResult`: Result type for augmented extraction
+
+## Usage Examples
+
+### Basic Data Extraction
+
+```python
+from pydantic import BaseModel
+from unique_toolkit._common.data_extraction import StructuredOutputDataExtractor, StructuredOutputDataExtractorConfig
+from unique_toolkit import LanguageModelService
+
+# Define your schema
+class PersonInfo(BaseModel):
+    name: str
+    age: int
+    occupation: str
+
+# Create the extractor
+config = StructuredOutputDataExtractorConfig()
+lm_service = LanguageModelService()  # Configure as needed
+extractor = StructuredOutputDataExtractor(config, lm_service)
+
+# Extract data
+text = "John is 30 years old and works as a software engineer."
+result = await extractor.extract_data_from_text(text, PersonInfo)
+print(result.data)  # PersonInfo(name="John", age=30, occupation="software engineer")
+```
+
+### Augmented Data Extraction
+
+```python
+from pydantic import BaseModel, Field
+from _common.data_extraction import AugmentedDataExtractor, StructuredOutputDataExtractor
+
+# Define your base schema
+class PersonInfo(BaseModel):
+    name: str
+    age: int
+
+# Create base extractor
+base_extractor = StructuredOutputDataExtractor(...)
+
+# Create augmented extractor with confidence scores
+augmented_extractor = AugmentedDataExtractor(
+    base_extractor,
+    confidence=float,
+    source=("extracted", Field(description="Source of the information"))
+)
+
+# Extract data
+text = "John is 30 years old."
+result = await augmented_extractor.extract_data_from_text(text, PersonInfo)
+print(result.data)  # Original PersonInfo
+print(result.augmented_data)  # Contains additional fields
+```
+
+## Configuration
+
+The `StructuredOutputDataExtractorConfig` allows customization of:
+
+- Language model selection
+- System and user prompt templates
+- Schema enforcement settings
+
+## Best Practices
+
+1. Always define clear Pydantic models for your extraction schemas
+2. Use augmented extraction when you need additional metadata
+3. Consider using strict mode for augmented extraction when you want to enforce schema compliance
+4. Customize prompts for better extraction results in specific domains 

unique_toolkit-1.33.1/unique_toolkit/data_extraction/__init__.py

@@ -0,0 +1,11 @@
+from unique_toolkit.data_extraction.augmented import AugmentedDataExtractor
+from unique_toolkit.data_extraction.basic import (
+    StructuredOutputDataExtractor,
+    StructuredOutputDataExtractorConfig,
+)
+
+__all__ = [
+    "StructuredOutputDataExtractor",
+    "StructuredOutputDataExtractorConfig",
+    "AugmentedDataExtractor",
+]

unique_toolkit-1.33.1/unique_toolkit/data_extraction/augmented/__init__.py

@@ -0,0 +1,5 @@
+from unique_toolkit.data_extraction.augmented.service import (
+    AugmentedDataExtractor,
+)
+
+__all__ = ["AugmentedDataExtractor"]

unique_toolkit-1.33.1/unique_toolkit/data_extraction/augmented/service.py

@@ -0,0 +1,93 @@
+from docxtpl.template import Any
+from pydantic import BaseModel, create_model
+from pydantic.alias_generators import to_pascal
+from pydantic.fields import FieldInfo
+from typing_extensions import override
+
+from unique_toolkit.data_extraction.base import (
+    BaseDataExtractionResult,
+    BaseDataExtractor,
+    ExtractionSchema,
+)
+
+
+def _build_augmented_model_for_field(
+    field_name: str,
+    field_type: Any | tuple[Any, FieldInfo],
+    strict: bool = False,
+    **extra_fields: Any | tuple[Any, FieldInfo],
+) -> type[BaseModel]:
+    camelized_field_name = to_pascal(field_name)
+
+    fields = {
+        **extra_fields,
+        field_name: field_type,
+    }
+
+    return create_model(
+        f"{camelized_field_name}Value",
+        **fields,  # type: ignore
+        __config__={"extra": "forbid" if strict else "ignore"},
+    )
+
+
+class AugmentedDataExtractionResult(BaseDataExtractionResult[ExtractionSchema]):
+    """
+    Result of data extraction from text using an augmented schema.
+    """
+
+    augmented_data: BaseModel
+
+
+class AugmentedDataExtractor(BaseDataExtractor):
+    def __init__(
+        self,
+        base_data_extractor: BaseDataExtractor,
+        strict: bool = False,
+        **extra_fields: Any | tuple[Any, FieldInfo],
+    ):
+        self._base_data_extractor = base_data_extractor
+        self._extra_fields = extra_fields
+        self._strict = strict
+
+    def _prepare_schema(self, schema: type[ExtractionSchema]) -> type[BaseModel]:
+        fields = {}
+
+        for field_name, field_type in schema.model_fields.items():
+            wrapped_field = _build_augmented_model_for_field(
+                field_name,
+                (field_type.annotation, field_type),
+                strict=self._strict,
+                **self._extra_fields,
+            )
+            fields[field_name] = wrapped_field
+
+        return create_model(
+            schema.__name__,
+            **fields,
+            __config__={"extra": "forbid" if self._strict else "ignore"},
+            __doc__=schema.__doc__,
+        )
+
+    def _extract_output(
+        self, llm_output: BaseModel, schema: type[ExtractionSchema]
+    ) -> ExtractionSchema:
+        output_data = {
+            field_name: getattr(value, field_name) for field_name, value in llm_output
+        }
+        return schema.model_validate(output_data)
+
+    @override
+    async def extract_data_from_text(
+        self, text: str, schema: type[ExtractionSchema]
+    ) -> AugmentedDataExtractionResult[ExtractionSchema]:
+        model_with_extra_fields = self._prepare_schema(schema)
+        augmented_data = (
+            await self._base_data_extractor.extract_data_from_text(
+                text, model_with_extra_fields
+            )
+        ).data
+        return AugmentedDataExtractionResult(
+            data=self._extract_output(augmented_data, schema),
+            augmented_data=augmented_data,
+        )

unique_toolkit-1.33.1/unique_toolkit/data_extraction/base.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel
+
+ExtractionSchema = TypeVar("ExtractionSchema", bound=BaseModel)
+
+
+class BaseDataExtractionResult(BaseModel, Generic[ExtractionSchema]):
+    """
+    Base class for data extraction results.
+    """
+
+    data: ExtractionSchema
+
+
+class BaseDataExtractor(ABC):
+    """
+    Extract structured data from text.
+    """
+
+    @abstractmethod
+    async def extract_data_from_text(
+        self, text: str, schema: type[ExtractionSchema]
+    ) -> BaseDataExtractionResult[ExtractionSchema]: ...

unique_toolkit-1.33.1/unique_toolkit/data_extraction/basic/__init__.py

@@ -0,0 +1,11 @@
+from unique_toolkit.data_extraction.basic.config import (
+    StructuredOutputDataExtractorConfig,
+)
+from unique_toolkit.data_extraction.basic.service import (
+    StructuredOutputDataExtractor,
+)
+
+__all__ = [
+    "StructuredOutputDataExtractorConfig",
+    "StructuredOutputDataExtractor",
+]

unique_toolkit-1.33.1/unique_toolkit/data_extraction/basic/config.py

@@ -0,0 +1,18 @@
+from pydantic import BaseModel
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit._common.validators import LMI, get_LMI_default_field
+from unique_toolkit.data_extraction.basic.prompt import (
+    DEFAULT_DATA_EXTRACTION_SYSTEM_PROMPT,
+    DEFAULT_DATA_EXTRACTION_USER_PROMPT,
+)
+from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
+
+
+class StructuredOutputDataExtractorConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    language_model: LMI = get_LMI_default_field(DEFAULT_GPT_4o)
+    structured_output_enforce_schema: bool = False
+    system_prompt_template: str = DEFAULT_DATA_EXTRACTION_SYSTEM_PROMPT
+    user_prompt_template: str = DEFAULT_DATA_EXTRACTION_USER_PROMPT

unique_toolkit-1.33.1/unique_toolkit/data_extraction/basic/prompt.py

@@ -0,0 +1,13 @@
+DEFAULT_DATA_EXTRACTION_SYSTEM_PROMPT = """
+You are a thorough and accurate expert in data processing.
+
+You will be given some text and an output schema, describing what needs to be extracted from the text.
+You will need to extract the data from the text and return it in the output schema.
+""".strip()
+
+DEFAULT_DATA_EXTRACTION_USER_PROMPT = """
+Here is the text to extract data from:
+{{ text }}
+
+Please thoroughly extract the data from the text and return it in the output schema.
+""".strip()

unique_toolkit-1.33.1/unique_toolkit/data_extraction/basic/service.py

@@ -0,0 +1,55 @@
+from typing_extensions import override
+
+from unique_toolkit._common.utils.jinja.render import render_template
+from unique_toolkit.data_extraction.base import (
+    BaseDataExtractionResult,
+    BaseDataExtractor,
+    ExtractionSchema,
+)
+from unique_toolkit.data_extraction.basic.config import (
+    StructuredOutputDataExtractorConfig,
+)
+from unique_toolkit.language_model import LanguageModelService
+from unique_toolkit.language_model.builder import MessagesBuilder
+
+
+class StructuredOutputDataExtractor(BaseDataExtractor):
+    """
+    Basic Structured Output Data Extraction.
+    """
+
+    def __init__(
+        self,
+        config: StructuredOutputDataExtractorConfig,
+        language_model_service: LanguageModelService,
+    ):
+        self._config = config
+        self._language_model_service = language_model_service
+
+    @override
+    async def extract_data_from_text(
+        self, text: str, schema: type[ExtractionSchema]
+    ) -> BaseDataExtractionResult[ExtractionSchema]:
+        messages_builder = (
+            MessagesBuilder()
+            .system_message_append(self._config.system_prompt_template)
+            .user_message_append(
+                render_template(
+                    self._config.user_prompt_template,
+                    {
+                        "text": text,
+                    },
+                )
+            )
+        )
+        response = await self._language_model_service.complete_async(
+            messages=messages_builder.build(),
+            model_name=self._config.language_model.name,
+            structured_output_model=schema,
+            temperature=0.0,
+            structured_output_enforce_schema=self._config.structured_output_enforce_schema,
+        )
+
+        return BaseDataExtractionResult(
+            data=schema.model_validate(response.choices[0].message.parsed),
+        )