unique_toolkit 0.7.29__py3-none-any.whl → 0.7.31__py3-none-any.whl

unique_toolkit/chat/schemas.py

@@ -10,6 +10,8 @@ from pydantic import (
     model_validator,
 )
 
+from unique_toolkit.content.schemas import ContentReference
+
 # set config to convert camelCase to snake_case
 model_config = ConfigDict(
     alias_generator=camelize, populate_by_name=True, arbitrary_types_allowed=True
@@ -37,15 +39,6 @@ class ToolCall(BaseModel):
     function: Function
 
 
-class Reference(BaseModel):
-    model_config = model_config
-    name: str
-    url: str | None = None
-    sequence_number: int
-    source_id: str
-    source: str
-
-
 class ChatMessage(BaseModel):
     # This model should strictly meets https://github.com/Unique-AG/monorepo/blob/master/node/apps/node-chat/src/public-api/2023-12-06/dtos/message/public-message.dto.ts
     model_config = model_config
@@ -63,7 +56,7 @@ class ChatMessage(BaseModel):
     created_at: datetime | None = None
     completed_at: datetime | None = None
     updated_at: datetime | None = None
-    references: list[Reference] | None = None
+    references: list[ContentReference] | None = None
 
     # TODO make sdk return role consistently in lowercase
     # Currently needed as sdk returns role in uppercase

unique_toolkit/content/schemas.py

@@ -60,13 +60,23 @@ class Content(BaseModel):
 
 class ContentReference(BaseModel):
     model_config = model_config
-    id: str
-    message_id: str
+    id: str = Field(
+        default="",
+        description="The id of the content reference. Can be empty on the ChatMessage Object",
+    )
+    message_id: str = Field(
+        default="",
+        description="The id of the message that this reference belongs to. Can be empty on the ChatMessage Object",
+    )
     name: str
     sequence_number: int
     source: str
     source_id: str
     url: str
+    original_index: list[int] = Field(
+        default=[],
+        description="List of indices in the ChatMessage original_content this reference refers to. This is usually the id in the functionCallResponse. List type due to implementation in node-chat",
+    )
 
 
 class ContentSearchType(StrEnum):

unique_toolkit/language_model/functions.py

@@ -7,7 +7,7 @@ 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.content.schemas import ContentChunk, ContentReference
 from unique_toolkit.evaluators import DOMAIN_NAME
 from unique_toolkit.language_model import (
     LanguageModelMessageRole,
@@ -20,7 +20,6 @@ from unique_toolkit.language_model import (
 )
 from unique_toolkit.language_model.infos import LanguageModelName
 from unique_toolkit.language_model.reference import (
-    PotentialReference,
     add_references_to_message,
 )
 
@@ -333,20 +332,9 @@ def _create_language_model_stream_response_with_references(
         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,
+        search_context=content_chunks,
     )
 
     stream_response_message = LanguageModelStreamResponseMessage(
@@ -355,7 +343,9 @@ def _create_language_model_stream_response_with_references(
         role=LanguageModelMessageRole.ASSISTANT,
         text=message.content or "",
         original_text=content,
-        references=[u.model_dump() for u in message.references or []],
+        references=[
+            ContentReference(**u.model_dump()) for u in message.references or []
+        ],
     )
 
     tool_calls = [r.function for r in response.choices[0].message.tool_calls or []]

unique_toolkit/language_model/infos.py

@@ -34,9 +34,13 @@ class LanguageModelName(StrEnum):
     ANTHROPIC_CLAUDE_SONNET_4 = "litellm:anthropic-claude-sonnet-4"
     ANTHROPIC_CLAUDE_OPUS_4 = "litellm:anthropic-claude-opus-4"
     GEMINI_2_0_FLASH = "litellm:gemini-2-0-flash"
+    GEMINI_2_5_FLASH = "litellm:gemini-2-5-flash"
+    GEMINI_2_5_FLASH_LITE_PREVIEW_0617 = "litellm:gemini-2-5-flash-lite-preview-06-17"
     GEMINI_2_5_FLASH_PREVIEW_0417 = "litellm:gemini-2-5-flash-preview-04-17"
     GEMINI_2_5_FLASH_PREVIEW_0520 = "litellm:gemini-2-5-flash-preview-05-20"
+    GEMINI_2_5_PRO = "litellm:gemini-2-5-pro"
     GEMINI_2_5_PRO_EXP_0325 = "litellm:gemini-2-5-pro-exp-03-25"
+    GEMINI_2_5_PRO_PREVIEW_0605 = "litellm:gemini-2-5-pro-preview-06-05"
 
 
 class EncoderName(StrEnum):
@@ -506,6 +510,44 @@ class LanguageModelInfo(BaseModel):
                     info_cutoff_at=date(2024, 8, 1),
                     published_at=date(2025, 2, 1),
                 )
+            case LanguageModelName.GEMINI_2_5_FLASH:
+                return cls(
+                    name=model_name,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.REASONING,
+                    ],
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gemini-2-5-flash",
+                    encoder_name=EncoderName.O200K_BASE,  # TODO:Replace with LLM tokenizer
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=1_048_576, token_limit_output=65_536
+                    ),
+                    info_cutoff_at=date(2025, 1, day=1),
+                    published_at=date(2025, 4, 1),
+                )
+            case LanguageModelName.GEMINI_2_5_FLASH_LITE_PREVIEW_0617:
+                return cls(
+                    name=model_name,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.REASONING,
+                    ],
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gemini-2-5-flash-lite-preview-06-17",
+                    encoder_name=EncoderName.O200K_BASE,  # TODO:Replace with LLM tokenizer
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=1_000_000, token_limit_output=64_000
+                    ),
+                    info_cutoff_at=date(2025, 1, day=1),
+                    published_at=date(2025, 6, 17),
+                )
             case LanguageModelName.GEMINI_2_5_FLASH_PREVIEW_0417:
                 return cls(
                     name=model_name,
@@ -544,6 +586,25 @@ class LanguageModelInfo(BaseModel):
                     info_cutoff_at=date(2025, 1, day=1),
                     published_at=date(2025, 4, 1),
                 )
+            case LanguageModelName.GEMINI_2_5_PRO:
+                return cls(
+                    name=model_name,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.REASONING,
+                    ],
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gemini-2-5-pro",
+                    encoder_name=EncoderName.O200K_BASE,  # TODO: Update encoder with litellm
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=1_048_576, token_limit_output=65_536
+                    ),
+                    info_cutoff_at=date(2025, 1, day=1),
+                    published_at=date(2025, 6, 17),
+                )
             case LanguageModelName.GEMINI_2_5_PRO_EXP_0325:
                 return cls(
                     name=model_name,
@@ -563,6 +624,25 @@ class LanguageModelInfo(BaseModel):
                     info_cutoff_at=date(2025, 1, day=1),
                     published_at=date(2025, 3, 1),
                 )
+            case LanguageModelName.GEMINI_2_5_PRO_PREVIEW_0605:
+                return cls(
+                    name=model_name,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.REASONING,
+                    ],
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gemini-2-5-pro-preview-06-05",
+                    encoder_name=EncoderName.O200K_BASE,  # TODO: Update encoder with litellm
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=1_048_576, token_limit_output=65_536
+                    ),
+                    info_cutoff_at=date(2025, 1, day=1),
+                    published_at=date(2025, 6, 5),
+                )
             case _:
                 if isinstance(model_name, LanguageModelName):
                     raise ValueError(

unique_toolkit/language_model/reference.py

@@ -1,27 +1,12 @@
 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
+from unique_toolkit.chat.schemas import ChatMessage
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
 
 
 def add_references_to_message(
     message: ChatMessage,
-    search_context: list[PotentialReference],
+    search_context: list[ContentChunk],
     model: str | None = None,
 ) -> tuple[ChatMessage, bool]:
     """Add references to a message and return the updated message with change status.
@@ -41,21 +26,21 @@ def add_references_to_message(
     )
     message.content = _postprocess_message(text)
 
-    message.references = [Reference(**ref.model_dump()) for ref in ref_found]
+    message.references = ref_found
     references_changed = len(ref_found) > 0
     return message, references_changed
 
 
 def _add_references(
     text: str,
-    search_context: list[PotentialReference],
+    search_context: list[ContentChunk],
     message_id: str,
     model: str | None = None,
-) -> tuple[str, list[NodeReference]]:
+) -> tuple[str, list[ContentReference]]:
     """Add references to text and return the processed text with reference status.
 
     Returns:
-        Tuple[str, bool]: (processed_text, ref_found)
+        Tuple[str, list[Reference]]: (processed_text, ref_found)
     """
     references = _find_references(
         text=text,
@@ -169,11 +154,11 @@ def _get_max_sub_count_in_text(text: str) -> int:
 
 def _find_references(
     text: str,
-    search_context: list[PotentialReference],
+    search_context: list[ContentChunk],
     message_id: str,
-) -> list[NodeReference]:
+) -> list[ContentReference]:
     """Find references in text based on search context."""
-    references: list[NodeReference] = []
+    references: list[ContentReference] = []
     sequence_number = 1 + _get_max_sub_count_in_text(text)
 
     # Find all numbers in brackets to ensure we get references in order of occurrence
@@ -190,7 +175,7 @@ def _find_references(
             continue
 
         # Don't put the reference twice
-        reference_name = search.title or search.key
+        reference_name = search.title or search.key or f"Content {search.id}"
         found_reference = next(
             (r for r in references if r.name == reference_name), None
         )
@@ -206,7 +191,7 @@ def _find_references(
         )
 
         references.append(
-            NodeReference(
+            ContentReference(
                 name=reference_name,
                 url=url,
                 sequence_number=sequence_number,
@@ -216,6 +201,7 @@ def _find_references(
                 else search.id,
                 source="node-ingestion-chunks",
                 message_id=message_id,
+                id=search.id,
             )
         )
         sequence_number += 1
@@ -229,7 +215,7 @@ def _extract_numbers_in_brackets(text: str) -> list[int]:
     return [int(match) for match in matches]
 
 
-def _add_footnotes_to_text(text: str, references: list[NodeReference]) -> str:
+def _add_footnotes_to_text(text: str, references: list[ContentReference]) -> str:
     """Replace bracket references with superscript footnotes."""
     for reference in references:
         for original_index in reference.original_index:

unique_toolkit/language_model/schemas.py

@@ -18,6 +18,7 @@ from pydantic import (
 )
 from typing_extensions import deprecated
 
+from unique_toolkit.content.schemas import ContentReference
 from unique_toolkit.language_model.utils import format_message
 
 # set config to convert camelCase to snake_case
@@ -48,7 +49,7 @@ class LanguageModelStreamResponseMessage(BaseModel):
     role: LanguageModelMessageRole
     text: str
     original_text: str | None = None
-    references: list[dict[str, list | dict | str | int | float | bool]] = []  # type: ignore
+    references: list[ContentReference] = []
 
     # TODO make sdk return role in lowercase
     # Currently needed as sdk returns role in uppercase

{unique_toolkit-0.7.29.dist-info → unique_toolkit-0.7.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.7.29
+Version: 0.7.31
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -111,14 +111,17 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.7.31] - 2025-06-19
+- Add typings to references in payload from `LanguageModelStreamResponseMessage` 
+- Add `original_index` to the base reference to reflect updated api
 
+## [0.7.30] - 2025-06-20
+- Adding litellm models `litellm:gemini-2-5-flash`, `gemini-2-5-flash-lite-preview-06-17`, `litellm:gemini-2-5-pro`, `litellm:gemini-2-5-pro-preview-06-05`
 
 ## [0.7.29] - 2025-06-19
 - Fix typehintin in services
 - Error on invalid initialization
 
-#
-
 ## [0.7.28] - 2025-06-17
 - Revert default factory change on `ChatEventPayload` for attribute `metadata_filter` due to error in `backend-ingestion` on empty dict
 
@@ -126,7 +129,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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.29.dist-info → unique_toolkit-0.7.31.dist-info}/RECORD

@@ -14,14 +14,14 @@ unique_toolkit/app/verification.py,sha256=GxFFwcJMy25fCA_Xe89wKW7bgqOu8PAs5y8QpH
 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=QsJVhBXgK6jDWRYpEAt-22jy5NKGsYs4fmHrOEdHAyc,29865
-unique_toolkit/chat/schemas.py,sha256=ct3BbvdQmpcNeDmJdc-Iz33LX1_3O-SGVgrsuT0B99k,2881
+unique_toolkit/chat/schemas.py,sha256=abPPeDtUIeEyKDnLhIqgIyqTsFANxh3j44EYrITBlHw,2786
 unique_toolkit/chat/service.py,sha256=dRydaABJspU1KlFUQ_Q-rXl9lDkhgeh99SxAguY4s_M,37432
 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
 unique_toolkit/content/constants.py,sha256=1iy4Y67xobl5VTnJB6SxSyuoBWbdLl9244xfVMUZi5o,60
 unique_toolkit/content/functions.py,sha256=0ELepm3_sl0SD_SYzvQVQ-jTdrcUqK5mVJZv0nQBuAw,18367
-unique_toolkit/content/schemas.py,sha256=iednldGojmoy9iPZwuIW23BfhWLE9G-fG_qAhxSgr7k,2451
+unique_toolkit/content/schemas.py,sha256=KJ604BOx0vBh2AwlTCZkOo55aHsI6yj8vxDAARKKqEo,2995
 unique_toolkit/content/service.py,sha256=jGRTIt0JQKL6qRyZ9-_njWJejKKvIvDOzEQrf8CdMaU,19491
 unique_toolkit/content/utils.py,sha256=GUVPrkZfMoAj4MRoBs5BD_7vSuLZTZx69hyWzYFrI50,7747
 unique_toolkit/embedding/__init__.py,sha256=uUyzjonPvuDCYsvXCIt7ErQXopLggpzX-MEQd3_e2kE,250
@@ -47,11 +47,11 @@ 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=J54mzKs-uOBejpiMDa1YviKvJDcYl0gvEF1CfLb7_S4,12208
-unique_toolkit/language_model/infos.py,sha256=peJ4cSJC__jGLWZoOZGRhoersmkwFmclsXTZi-KqYXc,30723
+unique_toolkit/language_model/functions.py,sha256=WhgHbJgz4Z2aZt9TLdOpI0PGyYWA5R90tdwkwdDeT8c,11987
+unique_toolkit/language_model/infos.py,sha256=w5__BVG-IiiEYKG1FwM838wzqNbYI3eCCEDocKezc0I,34801
 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/reference.py,sha256=ntoWdJ_hly8FntTgg1qpj0ta55SOCCVsC5L49VqOiAQ,8067
+unique_toolkit/language_model/schemas.py,sha256=_fJ9rckKaS082wd2x9DrJ9N6B7iTJ9ytFnq4eIFDNAQ,12781
 unique_toolkit/language_model/service.py,sha256=VRkUk2XbijqGlnTTvqU7uCue6qtT7lpLd_Y8f3bWv1I,10486
 unique_toolkit/language_model/utils.py,sha256=bPQ4l6_YO71w-zaIPanUUmtbXC1_hCvLK0tAFc3VCRc,1902
 unique_toolkit/protocols/support.py,sha256=V15WEIFKVMyF1QCnR8vIi4GrJy4dfTCB6d6JlqPZ58o,2341
@@ -62,7 +62,7 @@ unique_toolkit/short_term_memory/schemas.py,sha256=OhfcXyF6ACdwIXW45sKzjtZX_gkcJ
 unique_toolkit/short_term_memory/service.py,sha256=cqpXA0nMbi4PhFweg-Cql3u0RvaTi5c8Xjv0uHMiSGc,8112
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=cxWjb2dxEI2HGsakKdVCkSNi7VK9mr08w5sDcFCQyWI,9553
-unique_toolkit-0.7.29.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-0.7.29.dist-info/METADATA,sha256=9-GvYOZhYI_KjxC9SVHK2lU5cniY1ZPAPoY9CGKkdrw,24255
-unique_toolkit-0.7.29.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-0.7.29.dist-info/RECORD,,
+unique_toolkit-0.7.31.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-0.7.31.dist-info/METADATA,sha256=aNS24gdydQe6icQg_M3EFprEMJ2o8qcR-95-lya8tj8,24607
+unique_toolkit-0.7.31.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-0.7.31.dist-info/RECORD,,