unique_toolkit 0.7.25__py3-none-any.whl → 0.7.27__py3-none-any.whl

unique_toolkit/language_model/reference.py

@@ -0,0 +1,244 @@
+import re
+
+from pydantic import BaseModel
+
+from unique_toolkit.chat.schemas import ChatMessage, Reference
+
+
+class NodeReference(Reference):
+    original_index: list[int] = []
+    message_id: str | None = None
+
+
+class PotentialReference(BaseModel):
+    id: str
+    chunk_id: str | None = None
+    title: str | None = None
+    key: str
+    url: str | None = None
+    internally_stored_at: str | None = None
+
+
+def add_references_to_message(
+    message: ChatMessage,
+    search_context: list[PotentialReference],
+    model: str | None = None,
+) -> tuple[ChatMessage, bool]:
+    """Add references to a message and return the updated message with change status.
+
+    Returns:
+        Tuple[ChatMessage, bool]: (updated_message, references_changed)
+    """
+    if not message.content:
+        return message, False
+
+    if message.id is None:
+        raise ValueError("Message ID is required")
+
+    message.content = _preprocess_message(message.content)
+    text, ref_found = _add_references(
+        message.content, search_context, message.id, model
+    )
+    message.content = _postprocess_message(text)
+
+    message.references = [Reference(**ref.model_dump()) for ref in ref_found]
+    references_changed = len(ref_found) > 0
+    return message, references_changed
+
+
+def _add_references(
+    text: str,
+    search_context: list[PotentialReference],
+    message_id: str,
+    model: str | None = None,
+) -> tuple[str, list[NodeReference]]:
+    """Add references to text and return the processed text with reference status.
+
+    Returns:
+        Tuple[str, bool]: (processed_text, ref_found)
+    """
+    references = _find_references(
+        text=text,
+        search_context=search_context,
+        message_id=message_id,
+    )
+
+    # Only reference a source once, even if it is mentioned multiple times in the text.
+    with_footnotes = _add_footnotes_to_text(text=text, references=references)
+
+    # Gemini 2.5 flash model has tendency to add multiple references for the same fact
+    # This is a workaround to limit the number of references to 5
+    if model and model.startswith("litellm:gemini-2-5-flash"):
+        reduced_text = _limit_consecutive_source_references(with_footnotes)
+
+        # Get the references that remain after reduction
+        remaining_numbers = set()
+        sup_matches = re.findall(r"<sup>(\d+)</sup>", reduced_text)
+        remaining_numbers = {int(match) for match in sup_matches}
+
+        references = [
+            ref for ref in references if ref.sequence_number in remaining_numbers
+        ]
+        text = _remove_hallucinated_references(reduced_text)
+    else:
+        text = _remove_hallucinated_references(with_footnotes)
+
+    return text, references
+
+
+def _preprocess_message(text: str) -> str:
+    """Preprocess message text to normalize reference formats."""
+    # Remove user & assistant references: XML format '[<user>]', '[\<user>]', etc.
+    patterns = [
+        (r"\[(\\)?(<)?user(>)?\]", ""),
+        (r"\[(\\)?(<)?assistant(>)?\]", ""),
+        (r"source[\s]?\[(\\)?(<)?conversation(>)?\]", "the previous conversation"),
+        (r"\[(\\)?(<)?previous[_,\s]conversation(>)?\]", ""),
+        (r"\[(\\)?(<)?past[_,\s]conversation(>)?\]", ""),
+        (r"\[(\\)?(<)?previous[_,\s]?answer(>)?\]", ""),
+        (r"\[(\\)?(<)?previous[_,\s]question(>)?\]", ""),
+        (r"\[(\\)?(<)?conversation(>)?\]", ""),
+        (r"\[(\\)?(<)?none(>)?\]", ""),
+    ]
+
+    for pattern, replacement in patterns:
+        text = re.sub(pattern, replacement, text, flags=re.IGNORECASE)
+
+    # Replace XML format '[<source XX>]', '[<sourceXX>]' and '[\<sourceXX>]' with [XX]
+    text = re.sub(r"\[(\\)?<source[\s]?(\d+)>\]", r"[\2]", text)
+
+    # Replace format '[source XX]' and '[sourceXX]' with [XX]
+    text = re.sub(r"\[source[\s]?(\d+)\]", r"[\1]", text)
+
+    # Make all references non-bold
+    text = re.sub(r"\[\*\*(\d+)\*\*\]", r"[\1]", text)
+
+    # Replace 'SOURCEXX' and 'SOURCE XX' with [XX]
+    text = re.sub(r"source[\s]?(\d+)", r"[\1]", text, flags=re.IGNORECASE)
+
+    # Replace 'SOURCE n°X' with [XX]
+    text = re.sub(r"source[\s]?n°(\d+)", r"[\1]", text, flags=re.IGNORECASE)
+
+    # Replace '[<[XX]>]' and '[\<[XX]>]' with [XX]
+    text = re.sub(r"\[(\\)?\[?<\[(\d+)\]?\]>\]", r"[\2]", text)
+
+    # Replace '[[A], [B], ...]' or '[[A], B, C, ...]' with [A][B][C]...
+    def replace_combined_brackets(match):
+        numbers = re.findall(r"\d+", match.group(0))
+        return "".join(f"[{n}]" for n in numbers)
+
+    text = re.sub(
+        r"\[\[(\d+)\](?:,\s*(?:\[)?\d+(?:\])?)*\]", replace_combined_brackets, text
+    )
+
+    return text
+
+
+def _limit_consecutive_source_references(text: str) -> str:
+    """Limit consecutive source references to maximum 5 unique sources."""
+
+    def replace_consecutive(match):
+        # Extract all numbers from the match and get unique values
+        numbers = list(set(re.findall(r"\d+", match.group(0))))
+        # Take only the first five unique numbers
+        return "".join(f"<sup>{n}</sup>" for n in numbers[:5])
+
+    # Find sequences of 5+ consecutive sources
+    pattern = r"(?:<sup>\d+</sup>){5,}"
+    return re.sub(pattern, replace_consecutive, text)
+
+
+def _postprocess_message(text: str) -> str:
+    """Format superscript references to remove duplicates."""
+
+    def replace_sup_sequence(match):
+        # Extract unique numbers from the entire match
+        sup_numbers = set(re.findall(r"\d+", match.group(0)))
+        return "".join(f"<sup>{n}</sup>" for n in sup_numbers)
+
+    # Find sequences of 2+ superscripts including internal spaces
+    pattern = r"(<sup>\d+</sup>[ ]*)+<sup>\d+</sup>"
+    return re.sub(pattern, replace_sup_sequence, text)
+
+
+def _get_max_sub_count_in_text(text: str) -> int:
+    """Get the maximum superscript number in the text."""
+    matches = re.findall(r"<sup>(\d+)</sup>", text)
+    return max((int(match) for match in matches), default=0)
+
+
+def _find_references(
+    text: str,
+    search_context: list[PotentialReference],
+    message_id: str,
+) -> list[NodeReference]:
+    """Find references in text based on search context."""
+    references: list[NodeReference] = []
+    sequence_number = 1 + _get_max_sub_count_in_text(text)
+
+    # Find all numbers in brackets to ensure we get references in order of occurrence
+    numbers_in_brackets = _extract_numbers_in_brackets(text)
+
+    for number in numbers_in_brackets:
+        # Convert 1-based reference to 0-based index
+        index = number - 1
+        if index < 0 or index >= len(search_context):
+            continue
+
+        search = search_context[index]
+        if not search:
+            continue
+
+        # Don't put the reference twice
+        reference_name = search.title or search.key
+        found_reference = next(
+            (r for r in references if r.name == reference_name), None
+        )
+
+        if found_reference:
+            found_reference.original_index.append(number)
+            continue
+
+        url = (
+            search.url
+            if search.url and not search.internally_stored_at
+            else f"unique://content/{search.id}"
+        )
+
+        references.append(
+            NodeReference(
+                name=reference_name,
+                url=url,
+                sequence_number=sequence_number,
+                original_index=[number],
+                source_id=f"{search.id}_{search.chunk_id}"
+                if search.chunk_id
+                else search.id,
+                source="node-ingestion-chunks",
+                message_id=message_id,
+            )
+        )
+        sequence_number += 1
+
+    return references
+
+
+def _extract_numbers_in_brackets(text: str) -> list[int]:
+    """Extract numbers from [X] format in text."""
+    matches = re.findall(r"\[(\d+)\]", text)
+    return [int(match) for match in matches]
+
+
+def _add_footnotes_to_text(text: str, references: list[NodeReference]) -> str:
+    """Replace bracket references with superscript footnotes."""
+    for reference in references:
+        for original_index in reference.original_index:
+            text = text.replace(
+                f"[{original_index}]", f"<sup>{reference.sequence_number}</sup>"
+            )
+    return text
+
+
+def _remove_hallucinated_references(text: str) -> str:
+    """Remove any remaining bracket references that weren't converted."""
+    return re.sub(r"\[\d+\]", "", text).strip()

unique_toolkit/language_model/service.py

@@ -1,11 +1,12 @@
 import logging
-from typing import Optional, Type
+from typing import Any, Optional, Type
 
 from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from unique_toolkit._common.validate_required_values import validate_required_values
 from unique_toolkit.app.schemas import BaseEvent, ChatEvent, Event
+from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.language_model.constants import (
     DEFAULT_COMPLETE_TEMPERATURE,
     DEFAULT_COMPLETE_TIMEOUT,
@@ -14,11 +15,14 @@ from unique_toolkit.language_model.constants import (
 from unique_toolkit.language_model.functions import (
     complete,
     complete_async,
+    complete_with_references,
+    complete_with_references_async,
 )
 from unique_toolkit.language_model.infos import LanguageModelName
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessages,
     LanguageModelResponse,
+    LanguageModelStreamResponse,
     LanguageModelTool,
     LanguageModelToolDescription,
 )
@@ -260,3 +264,55 @@ class LanguageModelService:
             structured_output_model=structured_output_model,
             structured_output_enforce_schema=structured_output_enforce_schema,
         )
+
+    def complete_with_references(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        content_chunks: list[ContentChunk] | None = None,
+        debug_info: dict = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+        start_text: str | None = None,
+        other_options: dict[str, Any] | None = None,
+    ) -> LanguageModelStreamResponse:
+        [company_id] = validate_required_values([self._company_id])
+
+        return complete_with_references(
+            company_id=company_id,
+            messages=messages,
+            model_name=model_name,
+            content_chunks=content_chunks,
+            temperature=temperature,
+            timeout=timeout,
+            other_options=other_options,
+            tools=tools,
+            start_text=start_text,
+        )
+
+    async def complete_with_references_async(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        content_chunks: list[ContentChunk] | None = None,
+        debug_info: dict = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+        start_text: str | None = None,
+        other_options: dict[str, Any] | None = None,
+    ) -> LanguageModelStreamResponse:
+        [company_id] = validate_required_values([self._company_id])
+
+        return await complete_with_references_async(
+            company_id=company_id,
+            messages=messages,
+            model_name=model_name,
+            content_chunks=content_chunks,
+            temperature=temperature,
+            timeout=timeout,
+            other_options=other_options,
+            tools=tools,
+            start_text=start_text,
+        )

unique_toolkit/protocols/support.py

@@ -1,9 +1,11 @@
-from typing import Protocol
+from typing import Any, Awaitable, Protocol
 
+from unique_toolkit.content import ContentChunk
 from unique_toolkit.language_model import (
     LanguageModelMessages,
     LanguageModelName,
     LanguageModelResponse,
+    LanguageModelStreamResponse,
     LanguageModelTool,
     LanguageModelToolDescription,
 )
@@ -25,5 +27,37 @@ class SupportsComplete(Protocol):
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
-        **kwargs,
     ) -> LanguageModelResponse: ...
+
+    async def complete_async(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    ) -> Awaitable[LanguageModelResponse]: ...
+
+
+class SupportCompleteWithReferences(Protocol):
+    def complete_with_references(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        content_chunks: list[ContentChunk] | None = None,
+        debug_info: dict[str, Any] = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    ) -> LanguageModelStreamResponse: ...
+
+    def complete_with_references_async(
+        self,
+        messages: LanguageModelMessages,
+        model_name: LanguageModelName | str,
+        content_chunks: list[ContentChunk] | None = None,
+        debug_info: dict[str, Any] = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    ) -> Awaitable[LanguageModelStreamResponse]: ...

unique_toolkit/smart_rules/compile.py

@@ -0,0 +1,301 @@
+import re
+from datetime import datetime, timedelta, timezone
+from enum import Enum
+from typing import Any, Dict, List, 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: Dict[str, Union[str, int, bool]],
+        tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+        tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+        tool_parameters: Dict[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(
+        alias="and", validation_alias=AliasChoices("and", "and_list")
+    )
+
+    def _fill_in_variables(
+        self,
+        user_metadata: Dict[str, Union[str, int, bool]],
+        tool_parameters: Dict[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(
+        alias="or", validation_alias=AliasChoices("or", "or_list")
+    )
+
+    def _fill_in_variables(
+        self,
+        user_metadata: Dict[str, Union[str, int, bool]],
+        tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[str, Union[str, int, bool]],
+) -> Any:
+    return replace_variables(value, user_metadata, tool_parameters)
+
+
+def array_operator(
+    value: Any,
+    user_metadata: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[str, Union[str, int, bool]],
+) -> Any:
+    return value  # do nothing for now. No variables to replace
+
+
+def empty_operator(
+    operator: Operator,
+    user_metadata: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[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: Dict[str, Union[str, int, bool]],
+    tool_parameters: Dict[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")