unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

unique_toolkit/agentic/tools/utils/source_handling/source_formatting.py

@@ -0,0 +1,207 @@
+import re
+from string import Template
+
+from unique_toolkit.agentic.tools.utils.source_handling.schema import SourceFormatConfig
+from unique_toolkit.content.schemas import ContentChunk
+
+
+def _format_page_range(chunk: ContentChunk) -> str:
+    """Format page range string from chunk metadata."""
+    if not (
+        chunk.start_page
+        and chunk.end_page
+        and chunk.start_page > 0
+        and chunk.end_page > 0
+    ):
+        return ""
+    return (
+        str(chunk.start_page)
+        if chunk.start_page == chunk.end_page
+        else f"{chunk.start_page} - {chunk.end_page}"
+    )
+
+
+def _parse_chunk(
+    chunk: ContentChunk, section_templates: dict[str, str]
+) -> dict[str, str]:
+    """Extract sections from chunk text using regex patterns."""
+    text = chunk.text
+    result = dict()
+
+    for section, template in section_templates.items():
+        # Document and info are the only sections that are included in the text
+        if section in [
+            "document",
+            "info",
+        ]:  # Skip page as it's derived from metadata
+            pattern = SourceFormatConfig.template_to_pattern(template)
+            match = re.search(pattern, text, re.DOTALL)
+            result[section] = match.group(1) if match else ""
+            text = text.replace(match.group(0), "") if match else text
+
+    result["text"] = text.strip()
+    return result
+
+
+def format_chunk(index: int, chunk: ContentChunk, config: SourceFormatConfig) -> str:
+    """
+    This function formats a content chunk based on a given configuration template and its sections. Each chunk in the database includes a document section, an optional info section, and a text section, with the text section being the primary content. Typically, chunks are added to sources in search modules without any changes. However, certain scenarios necessitate extra formatting, such as incorporating page numbers or other metadata. This function enables the custom formatting of chunks when they are appended as sources.
+
+    Args:
+        index (int): The source index number to be used in the template.
+        chunk (ContentChunk): A ContentChunk object containing:
+            - text (str): The main content text
+            - start_page (int, optional): Starting page number
+            - end_page (int, optional): Ending page number
+            - metadata (dict, optional): Additional metadata key-value pairs
+        config (SourceFormatConfig): Configuration object containing:
+            - source_template (str): The overall template for the output
+            - sections (dict): Mapping of section names to their format templates
+
+    Returns:
+        str: Formatted string according to the template
+
+    Examples:
+        Using XML-style config without page numbers (default):
+        >>> config = SourceFormatConfig(
+        ...     source_template="<source${index}>${document}${info}${text}</source${index}>",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "info": "<|info|>{}<|/info|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content"
+        ... )
+        >>> format_chunk(1, chunk, config)
+        '<source1><|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content</source1>'
+
+        Using XML-style config with page numbers:
+        >>> config = SourceFormatConfig(
+        ...     source_template="<source${index}>${document}${page}${info}${text}</source${index}>",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "info": "<|info|>{}<|/info|>\n",
+        ...         "page": "<|page|>{}<|/page|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content",
+        ...     start_page=1,
+        ...     end_page=3,
+        ... )
+        >>> format_chunk(1, chunk, config)
+        '<source1><|document|>Sample Doc.pdf</|document|>\n<|page|>1 - 3</|page|>\n<|info|>Important info</|info|>\nMain content</source1>'
+
+        Using XML-style config with metadata:
+        >>> config = SourceFormatConfig(
+        ...     source_template="<source${index}>${document}${date}${text}</source${index}>",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "date": "<|DateFromMetaData|>{}<|/DateFromMetaData|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\nMain content",
+        ...     metadata={
+        ...         "key": "metadata-key",
+        ...         "mimeType": "text/plain",
+        ...         "date": "12.03.2025",
+        ...     },
+        ... )
+        >>> format_chunk(1, chunk, config)
+        '<source1><|document|>Sample Doc.pdf</|document|>\n<|DateFromMetaData|>12.03.2025</|DateFromMetaData|>\nMain content</source1>'
+
+        Using JSON-style config:
+        >>> config = SourceFormatConfig(
+        ...     source_template="{'source_number': ${index}, 'content': '${document}${page}${info}${text}'}",
+        ...     sections={
+        ...         "document": "<|document|>{}<|/document|>\n",
+        ...         "info": "<|info|>{}<|/info|>\n",
+        ...         "page": "<|page|>{}<|/page|>\n",
+        ...     },
+        ... )
+        >>> chunk = ContentChunk(
+        ...     text="<|document|>Sample Doc.pdf</|document|>\n<|info|>Important info</|info|>\nMain content",
+        ...     start_page=5,
+        ...     end_page=5,
+        ... )
+        >>> format_chunk(1, chunk, config)
+        "{'source_number': 1, 'content': '<|document|>Sample Doc.pdf</|document|>\n<|page|>5</|page|>\n<|info|>Important info</|info|>\nMain content'}"
+
+    Notes:
+        - The function extracts document and info sections from the chunk text using regex patterns
+        - Page numbers are formatted as single numbers when start_page equals end_page
+        - Page numbers are formatted as ranges (e.g., "1 - 3") when start_page differs from end_page
+        - If page numbers are not available (None or 0), the page section will be empty
+        - Metadata keys that match section names (except 'document' and 'text') will be included in the output
+        - Metadata is processed by the _process_metadata function to update the parsed dictionary
+        - When using custom metadata tags like '<|DateFromMetaData|>', the key in chunk.metadata must match
+          the key in the sections dictionary (e.g., 'date' in the example above), not the tag name
+    """
+    sections = config.sections
+    source_template = config.source_template
+
+    parsed = _parse_chunk(chunk, sections)
+    parsed["page"] = _format_page_range(chunk)
+
+    # Update parsed with metadata values
+    parsed = _process_metadata(chunk, parsed, sections)
+
+    # Create a new dictionary to hold the formatted sections
+    formatted_sections = {}
+
+    # Process each section
+    for section, template in sections.items():
+        if parsed.get(section):
+            formatted_sections[section] = template.format(parsed.get(section, ""))
+        else:
+            formatted_sections[section] = ""
+
+    # Add the text section
+    formatted_sections["text"] = parsed["text"]
+
+    return Template(source_template).substitute(index=index, **formatted_sections)
+
+
+def _process_metadata(
+    chunk: ContentChunk, parsed: dict[str, str], sections: dict[str, str]
+) -> dict[str, str]:
+    """
+    Process metadata from chunk and update the parsed dictionary.
+
+    This function extracts metadata from a ContentChunk object and updates the parsed
+    dictionary with values whose keys match section names defined in SourceFormatConfig.
+
+    Args:
+        chunk (ContentChunk): The content chunk containing metadata
+        parsed (dict): The dictionary of already parsed sections to update
+
+    Returns:
+        dict: The updated parsed dictionary with metadata values added
+
+    Notes:
+        - Keys 'document' and 'text' are explicitly excluded from metadata processing
+        - Only metadata keys that match section names in SourceFormatConfig will be processed
+        - If chunk.metadata is None or not iterable, the parsed dict is returned unchanged
+        - Metadata values are added directly to the parsed dictionary using their original keys
+    """
+    # Return unchanged parsed dict if metadata is None
+    if not hasattr(chunk, "metadata") or chunk.metadata is None:
+        return parsed
+
+    # Ensure metadata is a dictionary
+    metadata_dict = dict(chunk.metadata) if hasattr(chunk.metadata, "__iter__") else {}
+
+    # Define keys that should not be treated as metadata keys
+    excluded_keys = {"document", "info"}
+
+    # Get the keys from SourceFormatConfig.sections
+    valid_section_keys = set(sections.keys()) - excluded_keys
+
+    # Update parsed with valid metadata entries
+    for key, value in metadata_dict.items():
+        if key in valid_section_keys:
+            parsed[key] = value
+
+    return parsed

unique_toolkit/agentic/tools/utils/source_handling/tests/test_source_formatting.py

@@ -0,0 +1,216 @@
+import pytest
+
+from unique_toolkit.agentic.tools.utils.source_handling.schema import SourceFormatConfig
+from unique_toolkit.agentic.tools.utils.source_handling.source_formatting import (
+    _format_page_range,
+    format_chunk,
+)
+from unique_toolkit.content.schemas import ContentChunk
+
+
+@pytest.fixture
+def default_config():
+    return SourceFormatConfig()
+
+
+@pytest.fixture
+def xml_style_config_without_page_number():
+    return SourceFormatConfig(
+        source_template="<source${index}>${document}${info}${text}</source${index}>",
+        sections={
+            "document": "<|document|>{}<|/document|>\n",
+            "info": "<|info|>{}<|/info|>\n",
+        },
+    )
+
+
+@pytest.fixture
+def xml_style_config_with_page_number():
+    return SourceFormatConfig(
+        source_template="<source${index}>${document}${page}${info}${text}</source${index}>",
+        sections={
+            "document": "<|document|>{}<|/document|>\n",
+            "info": "<|info|>{}<|/info|>\n",
+            "page": "<|page|>{}<|/page|>\n",
+        },
+    )
+
+
+@pytest.fixture
+def xml_style_config_with_metadata():
+    return SourceFormatConfig(
+        source_template="<source${index}>${document}${date}${text}</source${index}>",
+        sections={
+            "document": "<|document|>{}<|/document|>\n",
+            "date": "<|DateFromMetaData|>{}<|/DateFromMetaData|>\n",
+        },
+    )
+
+
+@pytest.fixture
+def json_style_config():
+    return SourceFormatConfig(
+        source_template="{'source_number': ${index}, 'content': '${document}${page}${info}${text}'}",
+        sections={
+            "document": "<|document|>{}<|/document|>\n",
+            "info": "<|info|>{}<|/info|>\n",
+            "page": "<|page|>{}<|/page|>\n",
+        },
+    )
+
+
+def test_format_page_range():
+    # Test same start and end page
+    chunk = ContentChunk(id="1", order=1, text="test", start_page=1, end_page=1)
+    assert _format_page_range(chunk) == "1"
+
+    # Test page range
+    chunk = ContentChunk(id="1", order=1, text="test", start_page=1, end_page=3)
+    assert _format_page_range(chunk) == "1 - 3"
+
+    # Test invalid pages
+    chunk = ContentChunk(id="1", order=1, text="test", start_page=0, end_page=0)
+    assert _format_page_range(chunk) == ""
+
+
+def test_json_style_formatting(json_style_config):
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\n<|info|>Important<|/info|>\nContent text",
+        start_page=1,
+        end_page=2,
+    )
+
+    formatted = format_chunk(1, chunk, json_style_config)
+    expected = "{'source_number': 1, 'content': '<|document|>Doc1<|/document|>\n<|page|>1 - 2<|/page|>\n<|info|>Important<|/info|>\nContent text'}"
+    assert formatted == expected
+
+
+def test_metadata_handling(xml_style_config_with_metadata):
+    # Test with metadata that matches a section name
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\nContent text",
+        metadata={
+            "key": "metadata-key",
+            "mimeType": "text/plain",
+            "date": "12.03.2025",
+        },  # type: ignore
+    )
+
+    formatted = format_chunk(1, chunk, xml_style_config_with_metadata)
+    expected = "<source1><|document|>Doc1<|/document|>\n<|DateFromMetaData|>12.03.2025<|/DateFromMetaData|>\nContent text</source1>"
+    assert formatted == expected
+
+    # Test with metadata that doesn't match a section name
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\nContent text",
+        metadata={
+            "key": "metadata-key",
+            "mimeType": "text/plain",
+            "unrelated_key": "Some value",
+        },  # type: ignore
+    )
+
+    formatted = format_chunk(1, chunk, xml_style_config_with_metadata)
+    expected = "<source1><|document|>Doc1<|/document|>\nContent text</source1>"
+    assert formatted == expected
+
+    # Test with minimal metadata
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\nContent text",
+        metadata={"key": "metadata-key", "mimeType": "text/plain"},  # type: ignore
+    )
+
+    formatted = format_chunk(1, chunk, xml_style_config_with_metadata)
+    expected = "<source1><|document|>Doc1<|/document|>\nContent text</source1>"
+    assert formatted == expected
+
+
+def test_default_style(
+    default_config,
+):
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\n<|info|>Important<|/info|>\nContent text",
+        start_page=1,
+        end_page=2,
+    )
+
+    formatted = format_chunk(1, chunk, default_config)
+    expected = "<source1><|document|>Doc1<|/document|>\n<|info|>Important<|/info|>\nContent text</source1>"
+    assert formatted == expected
+
+
+def test_xml_style_without_page_number_formatting(
+    xml_style_config_without_page_number,
+):
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\n<|info|>Important<|/info|>\nContent text",
+        start_page=1,
+        end_page=2,
+    )
+
+    formatted = format_chunk(1, chunk, xml_style_config_without_page_number)
+    expected = "<source1><|document|>Doc1<|/document|>\n<|info|>Important<|/info|>\nContent text</source1>"
+    assert formatted == expected
+
+
+def test_xml_style_with_page_number_formatting(
+    xml_style_config_with_page_number,
+):
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc1<|/document|>\n<|info|>Important<|/info|>\nContent text",
+        start_page=1,
+        end_page=2,
+    )
+
+    formatted = format_chunk(1, chunk, xml_style_config_with_page_number)
+    expected = "<source1><|document|>Doc1<|/document|>\n<|page|>1 - 2<|/page|>\n<|info|>Important<|/info|>\nContent text</source1>"
+    assert formatted == expected
+
+
+def test_special_characters_handling(json_style_config):
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="<|document|>Doc's \"title\"<|/document|>\n<|info|>Info with {brackets}<|/info|>\nContent: with 'quotes'",
+        start_page=1,
+        end_page=1,
+    )
+
+    formatted = format_chunk(1, chunk, json_style_config)
+    expected = "{'source_number': 1, 'content': '<|document|>Doc's \"title\"<|/document|>\n<|page|>1<|/page|>\n<|info|>Info with {brackets}<|/info|>\nContent: with 'quotes''}"
+    assert formatted == expected
+
+
+def test_empty_sections(xml_style_config_without_page_number, json_style_config):
+    chunk = ContentChunk(
+        id="1",
+        order=1,
+        text="Just plain text without any sections",
+        start_page=None,
+        end_page=None,
+    )
+
+    # Test XML style
+    xml_formatted = format_chunk(1, chunk, xml_style_config_without_page_number)
+    assert xml_formatted == "<source1>Just plain text without any sections</source1>"
+
+    # Test JSON style
+    json_formatted = format_chunk(1, chunk, json_style_config)
+    assert (
+        json_formatted
+        == "{'source_number': 1, 'content': 'Just plain text without any sections'}"
+    )

unique_toolkit/app/__init__.py

@@ -38,6 +38,12 @@ from .schemas import (
 from .schemas import (
     EventUserMessage as EventUserMessage,
 )
+from .schemas import (
+    McpServer as McpServer,
+)
+from .schemas import (
+    McpTool as McpTool,
+)
 from .verification import (
     verify_signature_and_construct_event as verify_signature_and_construct_event,
 )

unique_toolkit/app/dev_util.py

@@ -0,0 +1,180 @@
+import asyncio
+import json
+from logging import getLogger
+from pathlib import Path
+from typing import (
+    Awaitable,
+    Callable,
+    Generator,
+    TypeVar,
+)
+
+from sseclient import SSEClient
+
+from unique_toolkit._common.exception import ConfigurationException
+from unique_toolkit.app import BaseEvent, ChatEvent, EventName
+from unique_toolkit.app.init_sdk import init_unique_sdk
+from unique_toolkit.app.unique_settings import UniqueSettings
+
+T = TypeVar("T", bound=BaseEvent)
+
+LOGGER = getLogger(__name__)
+
+
+def get_event_name_from_event_class(event_class: type[T]) -> EventName | None:
+    if event_class is ChatEvent:
+        return EventName.EXTERNAL_MODULE_CHOSEN
+
+    return None
+
+
+def get_sse_client(
+    unique_settings: UniqueSettings,
+    subscriptions: list[str],
+) -> SSEClient:
+    headers = {
+        "Authorization": f"Bearer {unique_settings.app.key.get_secret_value()}",
+        "x-app-id": unique_settings.app.id.get_secret_value(),
+        "x-company-id": unique_settings.auth.company_id.get_secret_value(),
+        "x-user-id": unique_settings.auth.user_id.get_secret_value(),
+        "x-api-version": unique_settings.api.version,
+    }
+    return SSEClient(url=unique_settings.api.sse_url(subscriptions), headers=headers)
+
+
+def get_event_generator(
+    unique_settings: UniqueSettings,
+    event_type: type[T],
+) -> Generator[T, None, None]:
+    """
+    Generator that updates the unique settings according to the events and
+    yields only events of the specified type from an SSE stream.
+
+    Args:
+        sse_client: The SSE client to read events from
+        event_type: The event class type to filter for
+
+    Yields:
+        Events matching the specified type
+    """
+    event_name = get_event_name_from_event_class(event_type)
+    if (
+        event_name is None
+        or not issubclass(event_type, BaseEvent)
+        or event_type is BaseEvent
+    ):
+        raise ValueError(f"Event model {event_type} is not a valid event model")
+
+    subscription = event_name.value
+
+    for sse_event in get_sse_client(unique_settings, [subscription]):
+        try:
+            payload = json.loads(sse_event.data)
+            parsed_event = event_type.model_validate(payload)
+            if parsed_event is None or parsed_event.filter_event(
+                filter_options=unique_settings.chat_event_filter_options
+            ):
+                continue
+
+            unique_settings.update_from_event(event=parsed_event)
+
+            yield parsed_event
+
+        except ConfigurationException as e:
+            # Re-raise ConfigurationException from filter_event (configuration errors)
+            raise e
+        except Exception as e:
+            LOGGER.error(f"Could not parse SSE event data as JSON: {e}")
+            continue
+
+
+def get_event_stream(
+    event_type: type[T] = BaseEvent,
+    settings_config: UniqueSettings | str | None = None,
+) -> Generator[T, None, None]:
+    """
+    Get an event stream from the SSE client.
+
+    Args:
+        event_type: The type of event to get
+        settings_or_filename: The settings or filename to use to setup the Unique settings object
+    """
+
+    if isinstance(settings_config, str):
+        unique_settings = UniqueSettings.from_env_auto_with_sdk_init(
+            filename=settings_config
+        )
+    elif isinstance(settings_config, UniqueSettings):
+        unique_settings = settings_config
+    else:
+        unique_settings = UniqueSettings.from_env_auto_with_sdk_init()
+
+    return get_event_generator(unique_settings, event_type)
+
+
+def run_demo_with_sse_client(
+    unique_settings: UniqueSettings,
+    handler: Callable[[BaseEvent], Awaitable[None] | None],
+    event_type: type[BaseEvent],
+) -> None:
+    """
+    Run a demo with an SSE client using sync handler.
+
+    Args:
+        unique_settings: The unique settings to use for the SSE client
+        handler: The sync handler to use for the SSE client
+        event_type: The type of event to use for the SSE client
+    """
+
+    event_name = get_event_name_from_event_class(event_type)
+    if event_name is None:
+        return
+
+    init_unique_sdk(unique_settings=unique_settings)
+    is_async_handler = asyncio.iscoroutinefunction(handler)
+
+    for event in get_event_generator(unique_settings, event_type):
+        if is_async_handler:
+            loop = asyncio.get_event_loop()
+            loop.run_until_complete(handler(event))
+        else:
+            handler(event)
+
+
+def load_event(file_path: Path, event_type: type[BaseEvent]) -> BaseEvent:
+    with file_path.open("r") as file:
+        event = json.load(file)
+
+    return event_type.model_validate(event)
+
+
+def run_demo_with_with_saved_event(
+    unique_settings: UniqueSettings,
+    handler: Callable[[BaseEvent], Awaitable[None] | None],
+    event_type: type[BaseEvent],
+    file_path: Path,
+) -> None:
+    """
+    Run a demo with an SSE client.
+
+    Note: event_type is the type of event that the handler expects.
+
+    Args:
+        unique_settings: The unique settings to use for the SSE client
+        handler: The handler to use for the SSE client
+        event_type: The type of event to use for the SSE client
+    """
+    init_unique_sdk(unique_settings=unique_settings)
+
+    event_name = get_event_name_from_event_class(event_type)
+    if event_name is None:
+        return
+
+    event = load_event(file_path, event_type)
+    if event is None:
+        raise ValueError(f"Event not found in {file_path}")
+
+    if asyncio.iscoroutinefunction(handler):
+        asyncio.run(handler(event))
+    else:
+        handler(event)

unique_toolkit/app/init_sdk.py

@@ -1,6 +1,11 @@
 import os
+from pathlib import Path
+from typing import overload
 
 import unique_sdk
+from typing_extensions import deprecated
+
+from unique_toolkit.app.unique_settings import UniqueSettings
 
 
 def get_env(var_name, default=None, strict=False):
@@ -24,12 +29,38 @@ def get_env(var_name, default=None, strict=False):
     return val or default
 
 
-def init_sdk(strict_all_vars=False):
+@overload
+def init_unique_sdk(*, env_file: Path | None = None): ...
+
+
+@overload
+def init_unique_sdk(*, unique_settings: UniqueSettings): ...
+
+
+def init_unique_sdk(
+    *, unique_settings: UniqueSettings | None = None, env_file: Path | None = None
+):
+    if unique_settings:
+        unique_sdk.api_key = unique_settings.app.key.get_secret_value()
+        unique_sdk.app_id = unique_settings.app.id.get_secret_value()
+        unique_sdk.api_base = unique_settings.api.sdk_url()
+    elif env_file:
+        unique_settings = UniqueSettings.from_env(env_file=env_file)
+        unique_sdk.api_key = unique_settings.app.key.get_secret_value()
+        unique_sdk.app_id = unique_settings.app.id.get_secret_value()
+        unique_sdk.api_base = unique_settings.api.sdk_url()
+
+
+@deprecated("Use init_unique_sdk instead")
+def init_sdk(
+    strict_all_vars: bool = False,
+):
     """Initialize the SDK.
 
     Args:
         strict_all_vars (bool, optional): This method raises a ValueError if strict and no value is found in the environment. Defaults to False.
     """
+
     unique_sdk.api_key = get_env("API_KEY", default="dummy", strict=strict_all_vars)
     unique_sdk.app_id = get_env("APP_ID", default="dummy", strict=strict_all_vars)
     unique_sdk.api_base = get_env("API_BASE", default=None, strict=strict_all_vars)