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

unique_toolkit/content/smart_rules.py

@@ -0,0 +1,301 @@
+import re
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+from typing import Any, Dict, List, Mapping, Self, Union
+
+from pydantic import AliasChoices, BaseModel, Field
+from pydantic.config import ConfigDict
+
+
+class Operator(str, Enum):
+    EQUALS = "equals"
+    NOT_EQUALS = "notEquals"
+    GREATER_THAN = "greaterThan"
+    GREATER_THAN_OR_EQUAL = "greaterThanOrEqual"
+    LESS_THAN = "lessThan"
+    LESS_THAN_OR_EQUAL = "lessThanOrEqual"
+    IN = "in"
+    NOT_IN = "notIn"
+    CONTAINS = "contains"
+    NOT_CONTAINS = "notContains"
+    IS_NULL = "isNull"
+    IS_NOT_NULL = "isNotNull"
+    IS_EMPTY = "isEmpty"
+    IS_NOT_EMPTY = "isNotEmpty"
+    NESTED = "nested"
+
+
+class BaseStatement(BaseModel):
+    model_config = ConfigDict(serialize_by_alias=True)
+
+    def with_variables(
+        self,
+        user_metadata: Mapping[str, Union[str, int, bool]],
+        tool_parameters: Mapping[str, Union[str, int, bool]],
+    ) -> Self:
+        return self._fill_in_variables(user_metadata, tool_parameters)
+
+    def is_compiled(self) -> bool:
+        # Serialize the object to json string
+        json_str = self.model_dump_json()
+        # Check if the json string has <T> or <T+> or <T-> or <toolParameters or <userMetadata
+        return (
+            "<T>" in json_str
+            or "<T+" in json_str
+            or "<T-" in json_str
+            or "<toolParameters" in json_str
+            or "<userMetadata" in json_str
+        )
+
+    def _fill_in_variables(
+        self,
+        user_metadata: Mapping[str, Union[str, int, bool]],
+        tool_parameters: Mapping[str, Union[str, int, bool]],
+    ) -> Self:
+        return self.model_copy()
+
+
+class Statement(BaseStatement):
+    operator: Operator
+    value: Union[str, int, bool, list[str], "AndStatement", "OrStatement"]
+    path: List[str] = Field(default_factory=list)
+
+    def _fill_in_variables(
+        self,
+        user_metadata: Mapping[str, Union[str, int, bool]],
+        tool_parameters: Mapping[str, Union[str, int, bool]],
+    ) -> Self:
+        new_stmt = self.model_copy()
+        new_stmt.value = eval_operator(self, user_metadata, tool_parameters)
+        return new_stmt
+
+
+class AndStatement(BaseStatement):
+    and_list: List[Union["Statement", "AndStatement", "OrStatement"]] = Field(
+        validation_alias=AliasChoices("and", "and_list"), serialization_alias="and"
+    )
+
+    def _fill_in_variables(
+        self,
+        user_metadata: Mapping[str, Union[str, int, bool]],
+        tool_parameters: Mapping[str, Union[str, int, bool]],
+    ) -> Self:
+        new_stmt = self.model_copy()
+        new_stmt.and_list = [
+            sub_query._fill_in_variables(user_metadata, tool_parameters)
+            for sub_query in self.and_list
+        ]
+        return new_stmt
+
+
+class OrStatement(BaseStatement):
+    or_list: List[Union["Statement", "AndStatement", "OrStatement"]] = Field(
+        validation_alias=AliasChoices("or", "or_list"), serialization_alias="or"
+    )
+
+    def _fill_in_variables(
+        self,
+        user_metadata: Mapping[str, Union[str, int, bool]],
+        tool_parameters: Mapping[str, Union[str, int, bool]],
+    ) -> Self:
+        new_stmt = self.model_copy()
+        new_stmt.or_list = [
+            sub_query._fill_in_variables(user_metadata, tool_parameters)
+            for sub_query in self.or_list
+        ]
+        return new_stmt
+
+
+# Update the forward references
+Statement.model_rebuild()
+AndStatement.model_rebuild()
+OrStatement.model_rebuild()
+
+
+UniqueQL = Union[Statement, AndStatement, OrStatement]
+
+
+def is_array_of_strings(value: Any) -> bool:
+    return isinstance(value, list) and all(isinstance(item, str) for item in value)
+
+
+def eval_operator(
+    query: Statement,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    if query.operator in [
+        Operator.EQUALS,
+        Operator.NOT_EQUALS,
+        Operator.GREATER_THAN,
+        Operator.GREATER_THAN_OR_EQUAL,
+        Operator.LESS_THAN,
+        Operator.LESS_THAN_OR_EQUAL,
+        Operator.CONTAINS,
+        Operator.NOT_CONTAINS,
+    ]:
+        return binary_operator(query.value, user_metadata, tool_parameters)
+    elif query.operator in [Operator.IS_NULL, Operator.IS_NOT_NULL]:
+        return null_operator(query.value, user_metadata, tool_parameters)
+    elif query.operator in [Operator.IS_EMPTY, Operator.IS_NOT_EMPTY]:
+        return empty_operator(query.operator, user_metadata, tool_parameters)
+    elif query.operator == Operator.NESTED:
+        return eval_nested_operator(query.value, user_metadata, tool_parameters)
+    elif query.operator in [Operator.IN, Operator.NOT_IN]:
+        return array_operator(query.value, user_metadata, tool_parameters)
+    else:
+        raise ValueError(f"Operator {query.operator} not supported")
+
+
+def eval_nested_operator(
+    value: Any,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Union[AndStatement, OrStatement]:
+    if not isinstance(value, (AndStatement, OrStatement)):
+        raise ValueError("Nested operator must be an AndStatement or OrStatement")
+    return value._fill_in_variables(user_metadata, tool_parameters)
+
+
+def binary_operator(
+    value: Any,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    return replace_variables(value, user_metadata, tool_parameters)
+
+
+def array_operator(
+    value: Any,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    if is_array_of_strings(value):
+        return [
+            replace_variables(item, user_metadata, tool_parameters) for item in value
+        ]
+    return value
+
+
+def null_operator(
+    value: Any,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    return value  # do nothing for now. No variables to replace
+
+
+def empty_operator(
+    operator: Operator,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    """Handle IS_EMPTY and IS_NOT_EMPTY operators."""
+    if operator == Operator.IS_EMPTY:
+        return ""
+    elif operator == Operator.IS_NOT_EMPTY:
+        return "not_empty"
+    return None
+
+
+def calculate_current_date() -> str:
+    """Calculate current date in UTC with seconds precision."""
+    return datetime.now(timezone.utc).isoformat(timespec="seconds")
+
+
+def calculate_earlier_date(input_str: str) -> str:
+    match = re.search(r"<T-(\d+)>", input_str)
+    if not match:
+        return calculate_current_date()  # Return current date if no match
+    days = int(match.group(1))
+    return (datetime.now(timezone.utc) - timedelta(days=days)).isoformat(
+        timespec="seconds"
+    )
+
+
+def calculate_later_date(input_str: str) -> str:
+    match = re.search(r"<T\+(\d+)>", input_str)  # Note: escaped + in regex
+    if not match:
+        return calculate_current_date()  # Return current date if no match
+    days = int(match.group(1))
+    return (datetime.now(timezone.utc) + timedelta(days=days)).isoformat(
+        timespec="seconds"
+    )
+
+
+def replace_variables(
+    value: Any,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    if isinstance(value, str):
+        if "||" in value:
+            return get_fallback_values(value, user_metadata, tool_parameters)
+        elif value == "<T>":
+            return calculate_current_date()
+        elif "<T-" in value:
+            return calculate_earlier_date(value)
+        elif "<T+" in value:
+            return calculate_later_date(value)
+
+        value = replace_tool_parameters_patterns(value, tool_parameters)
+        value = replace_user_metadata_patterns(value, user_metadata)
+
+        if value == "":
+            return value
+        try:
+            return int(value)
+        except ValueError:
+            if value.lower() in ["true", "false"]:
+                return value.lower() == "true"
+            return value
+    return value
+
+
+def replace_tool_parameters_patterns(
+    value: str, tool_parameters: Dict[str, Union[str, int, bool]]
+) -> str:
+    def replace_match(match):
+        param_name = match.group(1)
+        return str(tool_parameters.get(param_name, ""))
+
+    return re.sub(r"<toolParameters\.(\w+)>", replace_match, value)
+
+
+def replace_user_metadata_patterns(
+    value: str, user_metadata: Dict[str, Union[str, int, bool]]
+) -> str:
+    def replace_match(match):
+        param_name = match.group(1)
+        return str(user_metadata.get(param_name, ""))
+
+    return re.sub(r"<userMetadata\.(\w+)>", replace_match, value)
+
+
+def get_fallback_values(
+    value: str,
+    user_metadata: Mapping[str, Union[str, int, bool]],
+    tool_parameters: Mapping[str, Union[str, int, bool]],
+) -> Any:
+    values = value.split("||")
+    for val in values:
+        data = replace_variables(val, user_metadata, tool_parameters)
+        if data != "":
+            return data
+    return values
+
+
+# Example usage:
+def parse_uniqueql(json_data: Dict[str, Any]) -> UniqueQL:
+    if "operator" in json_data:
+        return Statement.model_validate(json_data)
+    elif "or" in json_data:
+        return OrStatement.model_validate(
+            {"or": [parse_uniqueql(item) for item in json_data["or"]]}
+        )
+    elif "and" in json_data:
+        return AndStatement.model_validate(
+            {"and": [parse_uniqueql(item) for item in json_data["and"]]}
+        )
+    else:
+        raise ValueError("Invalid UniqueQL format")

unique_toolkit/content/utils.py

@@ -190,9 +190,11 @@ def count_tokens(text: str, encoding_model="cl100k_base") -> int:
     return len(encoding.encode(text))
 
 
-def map_content_chunk(content_chunk: dict):
+def map_content_chunk(content_id: str, content_key: str, content_chunk: dict):
     return ContentChunk(
-        id=content_chunk["id"],
+        id=content_id,
+        key=content_key,
+        chunk_id=content_chunk["id"],
         text=content_chunk["text"],
         start_page=content_chunk["startPage"],
         end_page=content_chunk["endPage"],
@@ -206,9 +208,14 @@ def map_content(content: dict):
         key=content["key"],
         title=content["title"],
         url=content["url"],
-        chunks=[map_content_chunk(chunk) for chunk in content["chunks"]],
+        chunks=[
+            map_content_chunk(content["id"], content["key"], chunk)
+            for chunk in content["chunks"]
+        ],
         created_at=content["createdAt"],
         updated_at=content["updatedAt"],
+        ingestion_state=content.get("ingestionState"),
+        expired_at=content.get("expiredAt"),
     )
 
 

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/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/data_extraction/augmented/__init__.py

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

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/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/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/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/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/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),
+        )