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

unique_toolkit/app/schemas.py

@@ -1,4 +1,6 @@
+import json
 from enum import StrEnum
+from pathlib import Path
 from typing import Any, Optional
 
 from humps import camelize
@@ -27,6 +29,14 @@ class BaseEvent(BaseModel):
     user_id: str
     company_id: str
 
+    @classmethod
+    def from_json_file(cls, file_path: Path) -> "BaseEvent":
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
 
 ###
 # ChatEvent schemas
@@ -140,6 +150,14 @@ class ChatEvent(BaseEvent):
     created_at: Optional[int] = None
     version: Optional[str] = None
 
+    @classmethod
+    def from_json_file(cls, file_path: Path) -> "ChatEvent":
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
 
 @deprecated(
     """Use the more specific `ChatEvent` instead that has the same properties. \
@@ -150,3 +168,11 @@ class Event(ChatEvent):
     # The below should only affect type hints
     # event: EventName T
     # payload: EventPayload
+
+    @classmethod
+    def from_json_file(cls, file_path: Path) -> "Event":
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.model_validate(data)

unique_toolkit/chat/functions.py

@@ -3,6 +3,7 @@ import re
 from typing import Any, Dict, List, cast
 
 import unique_sdk
+from typing_extensions import deprecated
 from unique_sdk._list_object import ListObject
 
 from unique_toolkit._common import _time_utils
@@ -673,6 +674,7 @@ async def modify_message_assessment_async(
         raise e
 
 
+@deprecated("Use stream_complete_with_references instead")
 def stream_complete_to_chat(
     company_id: str,
     user_id: str,
@@ -682,7 +684,43 @@ def stream_complete_to_chat(
     assistant_id: str,
     messages: LanguageModelMessages,
     model_name: LanguageModelName | str,
-    content_chunks: list[ContentChunk] = [],
+    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 | None = None,
+) -> LanguageModelStreamResponse:
+    return stream_complete_with_references(
+        company_id=company_id,
+        user_id=user_id,
+        assistant_message_id=assistant_message_id,
+        user_message_id=user_message_id,
+        chat_id=chat_id,
+        assistant_id=assistant_id,
+        messages=messages,
+        model_name=model_name,
+        content_chunks=content_chunks,
+        debug_info=debug_info,
+        temperature=temperature,
+        timeout=timeout,
+        tools=tools,
+        start_text=start_text,
+        other_options=other_options,
+    )
+
+
+def stream_complete_with_references(
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    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,
@@ -719,7 +757,7 @@ def stream_complete_to_chat(
         temperature=temperature,
         tools=tools,
         other_options=other_options,
-        content_chunks=content_chunks,
+        content_chunks=content_chunks or [],
     )
 
     try:
@@ -747,6 +785,7 @@ def stream_complete_to_chat(
         raise e
 
 
+@deprecated("Use stream_complete_with_references_async instead")
 async def stream_complete_to_chat_async(
     company_id: str,
     user_id: str,
@@ -756,7 +795,43 @@ async def stream_complete_to_chat_async(
     assistant_id: str,
     messages: LanguageModelMessages,
     model_name: LanguageModelName | str,
-    content_chunks: list[ContentChunk] = [],
+    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 | None = None,
+) -> LanguageModelStreamResponse:
+    return await stream_complete_with_references_async(
+        company_id=company_id,
+        user_id=user_id,
+        assistant_message_id=assistant_message_id,
+        user_message_id=user_message_id,
+        chat_id=chat_id,
+        assistant_id=assistant_id,
+        messages=messages,
+        model_name=model_name,
+        content_chunks=content_chunks,
+        debug_info=debug_info,
+        temperature=temperature,
+        timeout=timeout,
+        tools=tools,
+        start_text=start_text,
+        other_options=other_options,
+    )
+
+
+async def stream_complete_with_references_async(
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    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,
@@ -778,7 +853,7 @@ async def stream_complete_to_chat_async(
         temperature=temperature,
         tools=tools,
         other_options=other_options,
-        content_chunks=content_chunks,
+        content_chunks=content_chunks or [],
     )
 
     try:

unique_toolkit/chat/service.py

@@ -48,8 +48,8 @@ from unique_toolkit.language_model.schemas import (
 )
 
 from .functions import (
-    stream_complete_to_chat,
-    stream_complete_to_chat_async,
+    stream_complete_with_references,
+    stream_complete_with_references_async,
 )
 
 logger = logging.getLogger(f"toolkit.{DOMAIN_NAME}.{__name__}")
@@ -1084,11 +1084,36 @@ class ChatService:
             label=label,
         )
 
+    @deprecated("Use complete_with_references instead")
     def stream_complete(
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
-        content_chunks: list[ContentChunk] = [],
+        content_chunks: list[ContentChunk] | None = None,
+        debug_info: dict = {},
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        timeout: int = DEFAULT_COMPLETE_TIMEOUT,
+        tools: Optional[list[LanguageModelTool | LanguageModelToolDescription]] = None,
+        start_text: Optional[str] = None,
+        other_options: Optional[dict] = None,
+    ) -> LanguageModelStreamResponse:
+        return self.complete_with_references(
+            messages=messages,
+            model_name=model_name,
+            content_chunks=content_chunks,
+            debug_info=debug_info,
+            temperature=temperature,
+            timeout=timeout,
+            tools=tools,
+            start_text=start_text,
+            other_options=other_options,
+        )
+
+    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,
@@ -1117,7 +1142,7 @@ class ChatService:
             ]
         )
 
-        return stream_complete_to_chat(
+        return stream_complete_with_references(
             company_id=company_id,
             user_id=user_id,
             assistant_message_id=assistant_message_id,
@@ -1139,7 +1164,7 @@ class ChatService:
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
-        content_chunks: list[ContentChunk] = [],
+        content_chunks: list[ContentChunk] | None = None,
         debug_info: dict = {},
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
@@ -1147,7 +1172,7 @@ class ChatService:
         start_text: Optional[str] = None,
         other_options: Optional[dict] = None,
     ) -> LanguageModelResponse:
-        response = self.stream_complete(
+        response = self.complete_with_references(
             messages=messages,
             model_name=model_name,
             content_chunks=content_chunks,
@@ -1161,11 +1186,12 @@ class ChatService:
 
         return LanguageModelResponse.from_stream_response(response)
 
+    @deprecated("use complete_with_references_async instead.")
     async def stream_complete_async(
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
-        content_chunks: list[ContentChunk] = [],
+        content_chunks: list[ContentChunk] | None = None,
         debug_info: dict = {},
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
@@ -1173,10 +1199,31 @@ class ChatService:
         start_text: Optional[str] = None,
         other_options: Optional[dict] = None,
     ) -> LanguageModelStreamResponse:
-        """
-        Streams a completion in the chat session asynchronously.
-        """
+        """Stream a completion in the chat session asynchronously."""
+        return await self.complete_with_references_async(
+            messages=messages,
+            model_name=model_name,
+            content_chunks=content_chunks,
+            debug_info=debug_info,
+            temperature=temperature,
+            timeout=timeout,
+            tools=tools,
+            start_text=start_text,
+            other_options=other_options,
+        )
 
+    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: Optional[list[LanguageModelTool | LanguageModelToolDescription]] = None,
+        start_text: Optional[str] = None,
+        other_options: Optional[dict] = None,
+    ) -> LanguageModelStreamResponse:
         [
             company_id,
             user_id,
@@ -1195,7 +1242,7 @@ class ChatService:
             ]
         )
 
-        return await stream_complete_to_chat_async(
+        return await stream_complete_with_references_async(
             company_id=company_id,
             user_id=user_id,
             assistant_message_id=assistant_message_id,
@@ -1217,7 +1264,7 @@ class ChatService:
         self,
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
-        content_chunks: list[ContentChunk] = [],
+        content_chunks: list[ContentChunk] | None,
         debug_info: dict = {},
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
@@ -1225,7 +1272,7 @@ class ChatService:
         start_text: Optional[str] = None,
         other_options: Optional[dict] = None,
     ) -> LanguageModelResponse:
-        response = self.stream_complete_async(
+        response = self.complete_with_references_async(
             messages=messages,
             model_name=model_name,
             content_chunks=content_chunks,

unique_toolkit/language_model/functions.py

@@ -1,23 +1,33 @@
+import copy
 import logging
-from typing import cast
+from datetime import datetime, timezone
+from typing import Any, cast
 
 import unique_sdk
 from pydantic import BaseModel
 
+from unique_toolkit.chat.schemas import ChatMessage, ChatMessageRole
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.evaluators import DOMAIN_NAME
-
-from .constants import (
-    DEFAULT_COMPLETE_TEMPERATURE,
-    DEFAULT_COMPLETE_TIMEOUT,
-)
-from .infos import LanguageModelName
-from .schemas import (
+from unique_toolkit.language_model import (
+    LanguageModelMessageRole,
     LanguageModelMessages,
     LanguageModelResponse,
+    LanguageModelStreamResponse,
+    LanguageModelStreamResponseMessage,
     LanguageModelTool,
     LanguageModelToolDescription,
 )
+from unique_toolkit.language_model.infos import LanguageModelName
+from unique_toolkit.language_model.reference import (
+    PotentialReference,
+    add_references_to_message,
+)
+
+from .constants import (
+    DEFAULT_COMPLETE_TEMPERATURE,
+    DEFAULT_COMPLETE_TIMEOUT,
+)
 
 logger = logging.getLogger(f"toolkit.{DOMAIN_NAME}.{__name__}")
 
@@ -36,6 +46,7 @@ def complete(
     """Call the completion endpoint synchronously without streaming the response.
 
     Args:
+    ----
         company_id (str): The company ID associated with the request.
         messages (LanguageModelMessages): The messages to complete.
         model_name (LanguageModelName | str): The model name to use for the completion.
@@ -45,6 +56,7 @@ def complete(
         other_options (Optional[dict]): Additional options to use. Defaults to None.
 
     Returns:
+    -------
         LanguageModelResponse: The response object containing the completed result.
 
     """
@@ -93,6 +105,7 @@ async def complete_async(
     the completed result.
 
     Args:
+    ----
         company_id (str): The company ID associated with the request.
         messages (LanguageModelMessages): The messages to complete.
         model_name (LanguageModelName | str): The model name to use for the completion.
@@ -102,9 +115,11 @@ async def complete_async(
         other_options (Optional[dict]): The other options to use. Defaults to None.
 
     Returns:
+    -------
         LanguageModelResponse: The response object containing the completed result.
 
     Raises:
+    ------
         Exception: If an error occurs during the request, an exception is raised
         and logged.
 
@@ -198,7 +213,8 @@ def _prepare_completion_params_util(
 ) -> tuple[dict, str, dict, dict | None]:
     """Prepare common parameters for completion requests.
 
-    Returns:
+    Returns
+    -------
         tuple containing:
         - options (dict): Combined options including tools and temperature
         - model (str): Resolved model name
@@ -232,3 +248,120 @@ def _prepare_completion_params_util(
     )
 
     return options, model, messages_dict, search_context
+
+
+def complete_with_references(
+    company_id: str,
+    messages: LanguageModelMessages,
+    model_name: LanguageModelName | str,
+    content_chunks: list[ContentChunk] | None = None,
+    debug_dict: 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:
+    # Use toolkit language model functions for chat completion
+    response = complete(
+        company_id=company_id,
+        model_name=model_name,
+        messages=messages,
+        temperature=temperature,
+        timeout=timeout,
+        tools=tools,
+        other_options=other_options,
+    )
+
+    return _create_language_model_stream_response_with_references(
+        response=response,
+        content_chunks=content_chunks,
+        start_text=start_text,
+    )
+
+
+async def complete_with_references_async(
+    company_id: str,
+    messages: LanguageModelMessages,
+    model_name: LanguageModelName | str,
+    content_chunks: list[ContentChunk] | None = None,
+    debug_dict: 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:
+    # Use toolkit language model functions for chat completion
+    response = await complete_async(
+        company_id=company_id,
+        model_name=model_name,
+        messages=messages,
+        temperature=temperature,
+        timeout=timeout,
+        tools=tools,
+        other_options=other_options,
+    )
+
+    return _create_language_model_stream_response_with_references(
+        response=response,
+        content_chunks=content_chunks,
+        start_text=start_text,
+    )
+
+
+def _create_language_model_stream_response_with_references(
+    response: LanguageModelResponse,
+    content_chunks: list[ContentChunk] | None = None,
+    start_text: str | None = None,
+):
+    content = response.choices[0].message.content
+    content_chunks = content_chunks or []
+
+    if content is None:
+        raise ValueError("Content is None, which is not supported")
+    elif isinstance(content, list):
+        raise ValueError("Content is a list, which is not supported")
+    else:
+        content = start_text or "" + str(content)
+
+    message = ChatMessage(
+        id="msg_unknown",
+        text=copy.deepcopy(content),
+        role=ChatMessageRole.ASSISTANT,
+        created_at=datetime.now(timezone.utc),
+        chat_id="chat_unknown",
+    )
+
+    search_context = [
+        PotentialReference(
+            id=source.id,
+            chunk_id=source.id,
+            title=source.title,
+            key=source.key or "",
+            url=source.url,
+        )
+        for source in content_chunks
+    ]
+
+    message, __ = add_references_to_message(
+        message=message,
+        search_context=search_context,
+    )
+
+    stream_response_message = LanguageModelStreamResponseMessage(
+        id="stream_unknown",
+        previous_message_id=None,
+        role=LanguageModelMessageRole.ASSISTANT,
+        text=message.content or "",
+        original_text=content,
+        references=[u.model_dump() for u in message.references or []],
+    )
+
+    tool_calls = [r.function for r in response.choices[0].message.tool_calls or []]
+    tool_calls = tool_calls if len(tool_calls) > 0 else None
+
+    return LanguageModelStreamResponse(
+        message=stream_response_message,
+        tool_calls=tool_calls,
+    )

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-0.7.26.dist-info → unique_toolkit-0.7.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.7.26
+Version: 0.7.27
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -112,6 +112,11 @@ 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).
 
 
+## [0.7.27] - 2025-06-16
+- Introduce a protocol for `complete_with_references` to enable testable services
+- Rename/Create functions `stream_complete` in chat service and llm service accordingly
+
+
 ## [0.7.26] - 2025-06-05
 - Add `scope_rules` to `ChatEventPayload`
 - Added `UniqueQL` compiler and pydantic classes for `UniqueQL`. Note this is functionally equivalent but not identical to `UQLOperator` or `UQLCombinator` in `unique_sdk`.

{unique_toolkit-0.7.26.dist-info → unique_toolkit-0.7.27.dist-info}/RECORD

@@ -9,13 +9,13 @@ unique_toolkit/app/init_logging.py,sha256=Sh26SRxOj8i8dzobKhYha2lLrkrMTHfB1V4jR3
 unique_toolkit/app/init_sdk.py,sha256=Nv4Now4pMfM0AgRhbtatLpm_39rKxn0WmRLwmPhRl-8,1285
 unique_toolkit/app/performance/async_tasks.py,sha256=H0l3OAcosLwNHZ8d2pd-Di4wHIXfclEvagi5kfqLFPA,1941
 unique_toolkit/app/performance/async_wrapper.py,sha256=yVVcRDkcdyfjsxro-N29SBvi-7773wnfDplef6-y8xw,1077
-unique_toolkit/app/schemas.py,sha256=c2Tu3woqc7gGuHYEspL0mOEyvmqOxTcavbnzZJhSYAs,3969
+unique_toolkit/app/schemas.py,sha256=_xQbgKxUtrFGoYb37aWQEZ2mi2gTyqKBGUeDWSLAj8Q,4935
 unique_toolkit/app/verification.py,sha256=GxFFwcJMy25fCA_Xe89wKW7bgqOu8PAs5y8QpHF0GSc,3861
 unique_toolkit/chat/__init__.py,sha256=LRs2G-JTVuci4lbtHTkVUiNcZcSR6uqqfnAyo7af6nY,619
 unique_toolkit/chat/constants.py,sha256=05kq6zjqUVB2d6_P7s-90nbljpB3ryxwCI-CAz0r2O4,83
-unique_toolkit/chat/functions.py,sha256=TP55fSVXWTO3OoGUuYBuK9cBHUw96wlQGbVfhhMalCI,27332
+unique_toolkit/chat/functions.py,sha256=QsJVhBXgK6jDWRYpEAt-22jy5NKGsYs4fmHrOEdHAyc,29865
 unique_toolkit/chat/schemas.py,sha256=ct3BbvdQmpcNeDmJdc-Iz33LX1_3O-SGVgrsuT0B99k,2881
-unique_toolkit/chat/service.py,sha256=K7XtB3IdKznNu1r4dy2dXiwZYQg_vKgRUD52RoKewQU,40799
+unique_toolkit/chat/service.py,sha256=jAgInqq7wU-erezl-BOS9BQ3DOR1nrbVFlMfACn26XI,42827
 unique_toolkit/chat/state.py,sha256=Cjgwv_2vhDFbV69xxsn7SefhaoIAEqLx3ferdVFCnOg,1445
 unique_toolkit/chat/utils.py,sha256=ihm-wQykBWhB4liR3LnwPVPt_qGW6ETq21Mw4HY0THE,854
 unique_toolkit/content/__init__.py,sha256=EdJg_A_7loEtCQf4cah3QARQreJx6pdz89Rm96YbMVg,940
@@ -47,13 +47,14 @@ unique_toolkit/evaluators/schemas.py,sha256=Jaue6Uhx75X1CyHKWj8sT3RE1JZXTqoLtfLt
 unique_toolkit/language_model/__init__.py,sha256=lRQyLlbwHbNFf4-0foBU13UGb09lwEeodbVsfsSgaCk,1971
 unique_toolkit/language_model/builder.py,sha256=69WCcmkm2rMP2-YEH_EjHiEp6OzwjwCs8VbhjVJaCe0,3168
 unique_toolkit/language_model/constants.py,sha256=B-topqW0r83dkC_25DeQfnPk3n53qzIHUCBS7YJ0-1U,119
-unique_toolkit/language_model/functions.py,sha256=koCAfhtkIGSiy8pSdDpIw9xRbwJ20EeLhDQMUXc8KZk,8049
+unique_toolkit/language_model/functions.py,sha256=J54mzKs-uOBejpiMDa1YviKvJDcYl0gvEF1CfLb7_S4,12208
 unique_toolkit/language_model/infos.py,sha256=peJ4cSJC__jGLWZoOZGRhoersmkwFmclsXTZi-KqYXc,30723
 unique_toolkit/language_model/prompt.py,sha256=JSawaLjQg3VR-E2fK8engFyJnNdk21zaO8pPIodzN4Q,3991
+unique_toolkit/language_model/reference.py,sha256=TrRfnxd-cjpdijmZ1dKII87aEACIbW7iU_cE9gtGuH4,8314
 unique_toolkit/language_model/schemas.py,sha256=DJD2aoMfs2Irnc4rzOrVuV4Fbt84LQAiDGG5rse1dgk,12770
-unique_toolkit/language_model/service.py,sha256=9LS3ouRNtzqZaKrMFagLZS9gBvNC5e46Ut86YWHBBHY,8470
+unique_toolkit/language_model/service.py,sha256=PJ0vOGPxzqsQXTQma8blLT4DfGP4gVZVcX2v1GVYozI,10610
 unique_toolkit/language_model/utils.py,sha256=bPQ4l6_YO71w-zaIPanUUmtbXC1_hCvLK0tAFc3VCRc,1902
-unique_toolkit/protocols/support.py,sha256=SD17M8jgjtzCh0bgDXgKrX96n6DizF1PT2SZIhyt4n8,888
+unique_toolkit/protocols/support.py,sha256=V15WEIFKVMyF1QCnR8vIi4GrJy4dfTCB6d6JlqPZ58o,2341
 unique_toolkit/short_term_memory/__init__.py,sha256=2mI3AUrffgH7Yt-xS57EGqnHf7jnn6xquoKEhJqk3Wg,185
 unique_toolkit/short_term_memory/constants.py,sha256=698CL6-wjup2MvU19RxSmQk3gX7aqW_OOpZB7sbz_Xg,34
 unique_toolkit/short_term_memory/functions.py,sha256=3WiK-xatY5nh4Dr5zlDUye1k3E6kr41RiscwtTplw5k,4484
@@ -61,7 +62,7 @@ unique_toolkit/short_term_memory/schemas.py,sha256=OhfcXyF6ACdwIXW45sKzjtZX_gkcJ
 unique_toolkit/short_term_memory/service.py,sha256=vEKFxP1SScPrFniso492fVthWR1sosdFibhiNF3zRvI,8081
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=44qDrrKD-bKCjjyUep9qa1IwNkneXoQezfFoVm1QToM,9558
-unique_toolkit-0.7.26.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-0.7.26.dist-info/METADATA,sha256=Au6JVHERLYsvDn8yH4tT_K_aB-fK3RhB83-GMIrIjCc,23803
-unique_toolkit-0.7.26.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-0.7.26.dist-info/RECORD,,
+unique_toolkit-0.7.27.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-0.7.27.dist-info/METADATA,sha256=Mh1lDoLWEiOU-0pEXDjd5kOBAgTJjqEagmhW2Hj2RLA,24000
+unique_toolkit-0.7.27.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-0.7.27.dist-info/RECORD,,