unique_toolkit 0.8.31__py3-none-any.whl → 0.8.33__py3-none-any.whl

unique_toolkit/_common/chunk_relevancy_sorter/config.py

@@ -0,0 +1,45 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+from pydantic.json_schema import SkipJsonSchema
+
+from unique_toolkit._common.default_language_model import DEFAULT_GPT_35_TURBO
+from unique_toolkit._common.validators import LMI, get_LMI_default_field
+from unique_toolkit.evals.context_relevancy.schema import StructuredOutputConfig
+from unique_toolkit.tools.config import get_configuration_dict
+
+
+class ChunkRelevancySortConfig(BaseModel):
+    model_config = get_configuration_dict()
+    enabled: bool = Field(
+        default=False,
+        description="Whether to enable the chunk relevancy sort.",
+    )
+    relevancy_levels_to_consider: list[str] = Field(
+        default=["high", "medium", "low"],
+        description="The relevancy levels to consider.",
+    )
+    relevancy_level_order: dict[str, int] = Field(
+        default={"high": 0, "medium": 1, "low": 2},
+        description="The relevancy level order.",
+    )
+    language_model: LMI = get_LMI_default_field(
+        DEFAULT_GPT_35_TURBO,
+        description="The language model to use for the chunk relevancy sort.",
+    )
+    fallback_language_model: LMI = get_LMI_default_field(
+        DEFAULT_GPT_35_TURBO,
+        description="The language model to use as a fallback.",
+    )
+    additional_llm_options: dict[str, Any] = Field(
+        default={},
+        description="Additional options to pass to the language model.",
+    )
+    structured_output_config: StructuredOutputConfig = Field(
+        default_factory=StructuredOutputConfig,
+        description="The configuration for the structured output.",
+    )
+    max_tasks: int | SkipJsonSchema[None] = Field(
+        default=1000,
+        description="The maximum number of tasks to run in parallel.",
+    )

unique_toolkit/_common/chunk_relevancy_sorter/exception.py

@@ -0,0 +1,5 @@
+from unique_toolkit._common.exception import CommonException
+
+
+class ChunkRelevancySorterException(CommonException):
+    pass

unique_toolkit/_common/chunk_relevancy_sorter/schemas.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.evals.schemas import EvaluationMetricResult
+
+
+class ChunkRelevancy(BaseModel):
+    chunk: ContentChunk
+    relevancy: EvaluationMetricResult | None = None
+
+    def get_document_name(self):
+        title = self.chunk.key or self.chunk.title or "Unkown"
+        return title.split(":")[0]
+
+    def get_page_number(self):
+        start_page = self.chunk.start_page
+        end_page = self.chunk.end_page
+
+        if start_page is None or end_page is None:
+            return start_page or end_page or "Unknown Page"
+        elif start_page == end_page:
+            return str(start_page)
+        else:
+            return f"{start_page}-{end_page}"
+
+    def get_facts(self):
+        if self.relevancy is None:
+            return []
+        return self.relevancy.fact_list
+
+
+class ChunkRelevancySorterResult(BaseModel):
+    relevancies: list[ChunkRelevancy]
+    user_message: Optional[str] = None
+
+    @staticmethod
+    def from_chunks(chunks: list[ContentChunk]):
+        return ChunkRelevancySorterResult(
+            relevancies=[ChunkRelevancy(chunk=chunk) for chunk in chunks],
+        )
+
+    @property
+    def content_chunks(self):
+        return [chunk.chunk for chunk in self.relevancies]

unique_toolkit/_common/chunk_relevancy_sorter/service.py

@@ -0,0 +1,372 @@
+import logging
+import time
+from collections import Counter
+from typing import Any, overload
+
+from typing_extensions import deprecated
+
+from unique_toolkit._common.chunk_relevancy_sorter.config import (
+    ChunkRelevancySortConfig,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.exception import (
+    ChunkRelevancySorterException,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.schemas import (
+    ChunkRelevancy,
+    ChunkRelevancySorterResult,
+)
+from unique_toolkit._common.validate_required_values import validate_required_values
+from unique_toolkit.app.performance.async_tasks import run_async_tasks_parallel
+from unique_toolkit.app.schemas import BaseEvent, ChatEvent
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.evals.config import EvaluationMetricConfig
+from unique_toolkit.evals.context_relevancy.schema import (
+    EvaluationSchemaStructuredOutput,
+    StructuredOutputConfig,
+)
+from unique_toolkit.evals.context_relevancy.service import ContextRelevancyEvaluator
+from unique_toolkit.evals.exception import EvaluatorException
+from unique_toolkit.evals.schemas import (
+    EvaluationMetricInput,
+    EvaluationMetricName,
+    EvaluationMetricResult,
+)
+from unique_toolkit.language_model.infos import LanguageModelInfo
+
+
+class ChunkRelevancySorter:
+    @deprecated(
+        "Use __init__ with company_id and user_id instead or use the classmethod `from_event`"
+    )
+    @overload
+    def __init__(self, event: ChatEvent | BaseEvent):
+        """
+        Initialize the ChunkRelevancySorter with an event (deprecated)
+        """
+
+    @overload
+    def __init__(self, *, company_id: str, user_id: str):
+        """
+        Initialize the ChunkRelevancySorter with a company_id and user_id
+        """
+
+    def __init__(
+        self,
+        event: ChatEvent | BaseEvent | None = None,
+        company_id: str | None = None,
+        user_id: str | None = None,
+    ):
+        if isinstance(event, (ChatEvent, BaseEvent)):
+            self.chunk_relevancy_evaluator = ContextRelevancyEvaluator.from_event(event)
+        else:
+            [company_id, user_id] = validate_required_values([company_id, user_id])
+            self.chunk_relevancy_evaluator = ContextRelevancyEvaluator(
+                company_id=company_id, user_id=user_id
+            )
+        module_name = "ChunkRelevancySorter"
+        self.logger = logging.getLogger(f"{module_name}.{__name__}")
+
+    @classmethod
+    def from_event(cls, event: ChatEvent | BaseEvent):
+        return cls(company_id=event.company_id, user_id=event.user_id)
+
+    async def run(
+        self,
+        input_text: str,
+        chunks: list[ContentChunk],
+        config: ChunkRelevancySortConfig,
+    ) -> ChunkRelevancySorterResult:
+        """
+        Resorts the search chunks by classifying each chunk into High, Medium, Low depending on the relevancy to the user input, then
+        sorts the chunks based on the classification while preserving the orginial order.
+
+        Args:
+            chunks (list[ContentChunk]): The list of search chunks to be reranked.
+
+        Returns:
+            ChunkRelevancySorterResult: The result of the chunk relevancy sort.
+
+        Raises:
+            ChunkRelevancySorterException: If an error occurs while sorting the chunks.
+        """
+
+        if not config.enabled:
+            self.logger.info("Chunk relevancy sort is disabled.")
+            return ChunkRelevancySorterResult.from_chunks(chunks)
+
+        self.logger.info("Running chunk relevancy sort.")
+        return await self._run_chunk_relevancy_sort(input_text, chunks, config)
+
+    async def _run_chunk_relevancy_sort(
+        self,
+        input_text: str,
+        chunks: list[ContentChunk],
+        config: ChunkRelevancySortConfig,
+    ) -> ChunkRelevancySorterResult:
+        start_time = time.time()
+
+        resorted_relevancies = []
+
+        try:
+            self.logger.info(f"Resorting {len(chunks)} chunks based on relevancy...")
+            chunk_relevancies = await self._evaluate_chunks_relevancy(
+                input_text,
+                chunks,
+                config,
+            )
+            resorted_relevancies = await self._validate_and_sort_relevant_chunks(
+                config,
+                chunk_relevancies,
+            )
+        except ChunkRelevancySorterException as e:
+            self.logger.error(e.error_message)
+            raise e
+        except Exception as e:
+            unknown_error_msg = "Unknown error occurred while resorting search results."
+            raise ChunkRelevancySorterException(
+                user_message=f"{unknown_error_msg}. Fallback to original search results.",
+                error_message=f"{unknown_error_msg}: {e}",
+            )
+        finally:
+            end_time = time.time()
+            duration = end_time - start_time
+            total_chunks = len(resorted_relevancies)
+            success_msg = f"Resorted {total_chunks} chunks in {duration:.2f} seconds."
+            self.logger.info(success_msg)
+            return ChunkRelevancySorterResult(
+                relevancies=resorted_relevancies,
+                user_message=success_msg,
+            )
+
+    async def _evaluate_chunks_relevancy(
+        self,
+        input_text: str,
+        chunks: list[ContentChunk],
+        config: ChunkRelevancySortConfig,
+    ) -> list[ChunkRelevancy]:
+        """
+        Evaluates the relevancy of the chunks.
+        """
+        self.logger.info(
+            f"Processing chunk relevancy for {len(chunks)} chunks with {config.language_model.name}. "
+            f"(Structured output: {config.structured_output_config.enabled}. Extract fact list: {config.structured_output_config.extract_fact_list})",
+        )
+
+        # Evaluate the relevancy of each chunk
+        tasks = [
+            self._process_relevancy_evaluation(input_text, chunk=chunk, config=config)
+            for chunk in chunks
+        ]
+        chunk_relevancies = await run_async_tasks_parallel(
+            tasks=tasks,
+            max_tasks=config.max_tasks,
+            logger=self.logger,
+        )
+
+        # handle exceptions
+        for chunk_relevancy in chunk_relevancies:
+            if isinstance(chunk_relevancy, Exception):
+                error_msg = "Error occurred while evaluating context relevancy of a specific chunk"
+                raise ChunkRelevancySorterException(
+                    user_message=f"{error_msg}. Fallback to original search results.",
+                    error_message=f"{error_msg}: {chunk_relevancy}",
+                    exception=chunk_relevancy,
+                )
+
+        # This check is currently necessary for typing purposes only
+        # as the run_async_tasks_parallel function does not enforce the return type
+        # TODO fix return type in run_async_tasks_parallel
+        chunk_relevancies = [
+            chunk_relevancy
+            for chunk_relevancy in chunk_relevancies
+            if isinstance(chunk_relevancy, ChunkRelevancy)
+        ]
+
+        return chunk_relevancies
+
+    async def _evaluate_chunk_relevancy(
+        self,
+        input_text: str,
+        langugage_model: LanguageModelInfo,
+        chunk: ContentChunk,
+        structured_output_config: StructuredOutputConfig,
+        additional_llm_options: dict[str, Any],
+    ) -> EvaluationMetricResult | None:
+        """
+        Gets the relevancy score of the chunk compared to the user message txt.
+        """
+        structured_output_schema = (
+            (
+                EvaluationSchemaStructuredOutput.get_with_descriptions(
+                    structured_output_config
+                )
+            )
+            if structured_output_config.enabled
+            else None
+        )
+
+        metric_config = EvaluationMetricConfig(
+            enabled=True,
+            name=EvaluationMetricName.CONTEXT_RELEVANCY,
+            language_model=langugage_model,
+            additional_llm_options=additional_llm_options,
+        )
+        relevancy_input = EvaluationMetricInput(
+            input_text=input_text,
+            context_texts=[chunk.text],
+        )
+
+        return await self.chunk_relevancy_evaluator.analyze(
+            input=relevancy_input,
+            config=metric_config,
+            structured_output_schema=structured_output_schema,
+        )
+
+    async def _process_relevancy_evaluation(
+        self,
+        input_text: str,
+        chunk: ContentChunk,
+        config: ChunkRelevancySortConfig,
+    ):
+        model = config.language_model
+        fallback_model = config.fallback_language_model
+        try:
+            relevancy = await self._evaluate_chunk_relevancy(
+                input_text=input_text,
+                langugage_model=model,
+                chunk=chunk,
+                structured_output_config=config.structured_output_config,
+                additional_llm_options=config.additional_llm_options,
+            )
+            return ChunkRelevancy(
+                chunk=chunk,
+                relevancy=relevancy,
+            )
+        except EvaluatorException as e:
+            if e.exception:
+                self.logger.warning(
+                    "Error evaluating chunk ID %s with model %s. Trying fallback model %s.",
+                    chunk.chunk_id,
+                    model,
+                    e.exception,
+                )
+                relevancy = await self._evaluate_chunk_relevancy(
+                    input_text=input_text,
+                    langugage_model=fallback_model,
+                    chunk=chunk,
+                    structured_output_config=config.structured_output_config,
+                    additional_llm_options=config.additional_llm_options,
+                )
+            else:
+                raise e
+        except Exception as e:
+            raise ChunkRelevancySorterException(
+                user_message="Error occurred while evaluating context relevancy of a specific chunk.",
+                error_message=f"Error in _process_relevancy_evaluation: {e}",
+                exception=e,
+            )
+
+    async def _validate_and_sort_relevant_chunks(
+        self,
+        config: ChunkRelevancySortConfig,
+        chunk_relevancies: list[ChunkRelevancy],
+    ) -> list[ChunkRelevancy]:
+        """
+        Checks for error or no value in chunk relevancy.
+        """
+
+        # Check that all chunk relevancies have a relevancy level
+        await self._validate_chunk_relevancies(chunk_relevancies)
+
+        # Filter the chunks based on the relevancy levels to consider
+        chunk_relevancies = await self._filter_chunks_by_relevancy_levels(
+            config,
+            chunk_relevancies,
+        )
+
+        # Sort the chunks based on the relevancy levels
+        sorted_chunks = await self._sort_chunk_relevancies_by_relevancy_and_chunk(
+            config,
+            chunk_relevancies,
+        )
+
+        return sorted_chunks
+
+    async def _validate_chunk_relevancies(
+        self,
+        chunk_relevancies: list[ChunkRelevancy],
+    ):
+        for chunk_relevancy in chunk_relevancies:
+            if not chunk_relevancy.relevancy or not chunk_relevancy.relevancy.value:
+                raise ChunkRelevancySorterException(
+                    user_message="Error occurred while evaluating chunk relevancy.",
+                    error_message=f"No relevancy level returned for chunk ID {chunk_relevancy.chunk.chunk_id}.",
+                )
+
+    async def _sort_chunk_relevancies_by_relevancy_and_chunk(
+        self,
+        config: ChunkRelevancySortConfig,
+        chunk_relevancies: list[ChunkRelevancy],
+    ):
+        # Define the custom sorting order for relevancy
+        relevancy_level_order = config.relevancy_level_order
+
+        # Create a dictionary to map the chunk chunkId to its position in the original order
+        chunk_order = {
+            relevancy.chunk.chunk_id: index
+            for index, relevancy in enumerate(chunk_relevancies)
+        }
+
+        # Sort the chunk relevancies first by relevancy and then by original order within each relevancy level
+        sorted_chunk_relevancies = sorted(
+            chunk_relevancies,
+            key=lambda obj: (
+                relevancy_level_order[obj.relevancy.value.lower()],  # type: ignore
+                chunk_order[obj.chunk.chunk_id],
+            ),
+        )
+
+        # Count and print the distinct values of relevancy
+        self._count_distinct_values(sorted_chunk_relevancies)
+
+        # Return only the chunk in the sorted order
+        return sorted_chunk_relevancies
+
+    async def _filter_chunks_by_relevancy_levels(
+        self,
+        config: ChunkRelevancySortConfig,
+        chunk_relevancies: list[ChunkRelevancy],
+    ) -> list[ChunkRelevancy]:
+        levels_to_consider = [
+            relevancy_level.lower()
+            for relevancy_level in config.relevancy_levels_to_consider
+        ]
+        if not levels_to_consider:
+            self.logger.warning("No relevancy levels defined, including all levels.")
+            return chunk_relevancies
+
+        self.logger.info(
+            "Filtering chunks by relevancy levels: %s.", levels_to_consider
+        )
+        return [
+            chunk_relevancy
+            for chunk_relevancy in chunk_relevancies
+            if chunk_relevancy.relevancy.value.lower() in levels_to_consider  # type: ignore
+        ]
+
+    def _count_distinct_values(self, chunk_relevancies: list[ChunkRelevancy]):
+        # Extract the values from the relevancy field
+        values = [
+            cr.relevancy.value
+            for cr in chunk_relevancies
+            if cr.relevancy and cr.relevancy.value
+        ]
+
+        # Use Counter to count occurrences
+        value_counts = Counter(values)
+
+        self.logger.info("Count of distinct relevancy values:")
+        for value, count in value_counts.items():
+            self.logger.info(f"Relevancy: {value}, Count: {count}")
+
+        return value_counts

unique_toolkit/_common/chunk_relevancy_sorter/tests/test_service.py

@@ -0,0 +1,273 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from unique_toolkit._common.chunk_relevancy_sorter.config import (
+    ChunkRelevancySortConfig,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.exception import (
+    ChunkRelevancySorterException,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.schemas import (
+    ChunkRelevancy,
+    ChunkRelevancySorterResult,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.service import ChunkRelevancySorter
+from unique_toolkit._common.default_language_model import (
+    DEFAULT_GPT_35_TURBO,
+    DEFAULT_GPT_4o,
+)
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.evals.context_relevancy.schema import StructuredOutputConfig
+from unique_toolkit.evals.schemas import EvaluationMetricName, EvaluationMetricResult
+from unique_toolkit.language_model.infos import LanguageModelInfo
+
+
+@pytest.fixture
+def event():
+    event = MagicMock(spec=ChatEvent)
+    event.payload = MagicMock()
+    event.payload.user_message = MagicMock()
+    event.payload.user_message.text = "Test query"
+    event.user_id = "user_0"
+    event.company_id = "company_0"
+    return event
+
+
+@pytest.fixture
+def mock_chunks():
+    return [
+        ContentChunk(
+            id=f"chunk_{i}",
+            order=i,
+            chunk_id=f"chunk_{i}",
+            text=f"Test content {i}",
+        )
+        for i in range(3)
+    ]
+
+
+@pytest.fixture
+def config():
+    return ChunkRelevancySortConfig(
+        enabled=True,
+        relevancy_levels_to_consider=["high", "medium", "low"],
+        relevancy_level_order={"high": 0, "medium": 1, "low": 2},
+        language_model=LanguageModelInfo.from_name(DEFAULT_GPT_4o),
+        fallback_language_model=LanguageModelInfo.from_name(DEFAULT_GPT_35_TURBO),
+        structured_output_config=StructuredOutputConfig(
+            enabled=False,
+            extract_fact_list=False,
+        ),
+    )
+
+
+@pytest.fixture
+def chunk_relevancy_sorter(event):
+    return ChunkRelevancySorter(event)
+
+
+@pytest.mark.asyncio
+async def test_run_disabled_config(chunk_relevancy_sorter, mock_chunks, config):
+    config.enabled = False
+    result = await chunk_relevancy_sorter.run("test input", mock_chunks, config)
+
+    assert isinstance(result, ChunkRelevancySorterResult)
+    assert result.content_chunks == mock_chunks
+    assert len(result.content_chunks) == len(mock_chunks)
+
+
+@pytest.mark.asyncio
+async def test_run_enabled_config(chunk_relevancy_sorter, mock_chunks, config):
+    with patch.object(chunk_relevancy_sorter, "_run_chunk_relevancy_sort") as mock_sort:
+        mock_sort.return_value = ChunkRelevancySorterResult.from_chunks(mock_chunks)
+
+        result = await chunk_relevancy_sorter.run("test input", mock_chunks, config)
+
+        assert isinstance(result, ChunkRelevancySorterResult)
+        assert result.content_chunks == mock_chunks
+        mock_sort.assert_called_once_with("test input", mock_chunks, config)
+
+
+@pytest.mark.asyncio
+async def test_evaluate_chunks_relevancy(chunk_relevancy_sorter, mock_chunks, config):
+    mock_relevancy = EvaluationMetricResult(
+        value="high",
+        name=EvaluationMetricName.CONTEXT_RELEVANCY,
+        reason="Test reason",
+    )
+
+    with patch.object(
+        chunk_relevancy_sorter, "_process_relevancy_evaluation"
+    ) as mock_process:
+        mock_process.return_value = ChunkRelevancy(
+            chunk=mock_chunks[0], relevancy=mock_relevancy
+        )
+
+        result = await chunk_relevancy_sorter._evaluate_chunks_relevancy(
+            "test input", mock_chunks, config
+        )
+
+        assert len(result) == len(mock_chunks)
+        assert all(isinstance(r, ChunkRelevancy) for r in result)
+        assert mock_process.call_count == len(mock_chunks)
+
+
+@pytest.mark.asyncio
+async def test_evaluate_chunk_relevancy(chunk_relevancy_sorter, mock_chunks, config):
+    with patch(
+        "unique_toolkit._common.chunk_relevancy_sorter.service.ContextRelevancyEvaluator.analyze"
+    ) as mock_analyze:
+        mock_analyze.return_value = EvaluationMetricResult(
+            value="high",
+            name=EvaluationMetricName.CONTEXT_RELEVANCY,
+            reason="Test reason",
+        )
+
+        result = await chunk_relevancy_sorter._evaluate_chunk_relevancy(
+            input_text="test input",
+            chunk=mock_chunks[0],
+            langugage_model=config.language_model,
+            structured_output_config=config.structured_output_config,
+            additional_llm_options=config.additional_llm_options,
+        )
+
+        assert isinstance(result, EvaluationMetricResult)
+        assert result.value == "high"
+        mock_analyze.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_process_relevancy_evaluation_success(
+    chunk_relevancy_sorter, mock_chunks, config
+):
+    with patch.object(
+        chunk_relevancy_sorter, "_evaluate_chunk_relevancy"
+    ) as mock_evaluate:
+        mock_evaluate.return_value = EvaluationMetricResult(
+            value="high",
+            name=EvaluationMetricName.CONTEXT_RELEVANCY,
+            reason="Test reason",
+        )
+
+        result = await chunk_relevancy_sorter._process_relevancy_evaluation(
+            "test input", mock_chunks[0], config
+        )
+
+        assert isinstance(result, ChunkRelevancy)
+        assert result.chunk == mock_chunks[0]
+        assert result.relevancy is not None
+        assert result.relevancy.value == "high"
+
+
+@pytest.mark.asyncio
+async def test_process_relevancy_evaluation_fallback(
+    chunk_relevancy_sorter, mock_chunks, config
+):
+    with patch.object(
+        chunk_relevancy_sorter, "_evaluate_chunk_relevancy"
+    ) as mock_evaluate:
+        # First call raises exception, second call succeeds
+        mock_evaluate.side_effect = [
+            Exception("Test error"),
+            EvaluationMetricResult(
+                value="medium",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        ]
+
+        with pytest.raises(ChunkRelevancySorterException):
+            await chunk_relevancy_sorter._process_relevancy_evaluation(
+                "test input", mock_chunks[0], config
+            )
+
+
+@pytest.mark.asyncio
+async def test_validate_and_sort_relevant_chunks(
+    chunk_relevancy_sorter, mock_chunks, config
+):
+    chunk_relevancies = [
+        ChunkRelevancy(
+            chunk=mock_chunks[0],
+            relevancy=EvaluationMetricResult(
+                value="low",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    ]
+    chunk_relevancies.append(
+        ChunkRelevancy(
+            chunk=mock_chunks[1],
+            relevancy=EvaluationMetricResult(
+                value="medium",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    )
+    chunk_relevancies.append(
+        ChunkRelevancy(
+            chunk=mock_chunks[2],
+            relevancy=EvaluationMetricResult(
+                value="high",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    )
+
+    result = await chunk_relevancy_sorter._validate_and_sort_relevant_chunks(
+        config, chunk_relevancies
+    )
+
+    assert isinstance(result, list)
+    assert len(result) == len(mock_chunks)
+    assert all(isinstance(relevancy.chunk, ContentChunk) for relevancy in result)
+    assert result[0].chunk == mock_chunks[2]
+    assert result[1].chunk == mock_chunks[1]
+    assert result[2].chunk == mock_chunks[0]
+
+
+@pytest.mark.asyncio
+async def test_validate_chunk_relevancies_invalid(chunk_relevancy_sorter):
+    invalid_relevancies = [
+        ChunkRelevancy(
+            chunk=ContentChunk(chunk_id="test", text="test", id="test", order=0),
+            relevancy=None,
+        )
+    ]
+
+    with pytest.raises(ChunkRelevancySorterException):
+        await chunk_relevancy_sorter._validate_chunk_relevancies(invalid_relevancies)
+
+
+def test_count_distinct_values(chunk_relevancy_sorter, mock_chunks):
+    chunk_relevancies = [
+        ChunkRelevancy(
+            chunk=chunk,
+            relevancy=EvaluationMetricResult(
+                value="high",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+        for chunk in mock_chunks[:2]
+    ]
+    chunk_relevancies.append(
+        ChunkRelevancy(
+            chunk=mock_chunks[2],
+            relevancy=EvaluationMetricResult(
+                value="medium",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    )
+
+    value_counts = chunk_relevancy_sorter._count_distinct_values(chunk_relevancies)
+
+    assert value_counts["high"] == 2
+    assert value_counts["medium"] == 1

unique_toolkit/_common/feature_flags/schema.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel, Field
+
+from unique_toolkit.tools.config import get_configuration_dict
+
+
+class FeatureExtendedSourceSerialization(BaseModel):
+    """Mixin for experimental feature in Source serialization"""
+
+    model_config = get_configuration_dict()
+    full_sources_serialize_dump: bool = Field(
+        default=False,
+        description="Whether to include the full source object in the tool response. If True, includes the full Source object. If False, uses the old format with only source_number and content.",
+    )

unique_toolkit/_common/utils/structured_output/schema.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class StructuredOutputModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")

unique_toolkit/evals/context_relevancy/schema.py

@@ -1,13 +1,10 @@
-from pydantic import BaseModel, ConfigDict, Field, create_model
+from pydantic import BaseModel, Field, create_model
 from pydantic.json_schema import SkipJsonSchema
 
+from unique_toolkit._common.utils.structured_output.schema import StructuredOutputModel
 from unique_toolkit.tools.config import get_configuration_dict
 
 
-class StructuredOutputModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-
 class StructuredOutputConfig(BaseModel):
     model_config = get_configuration_dict()
 

unique_toolkit/evals/context_relevancy/service.py

@@ -1,9 +1,14 @@
 import logging
+from typing import overload
 
 from pydantic import BaseModel, ValidationError
+from typing_extensions import deprecated
 
-from unique_toolkit.app.schemas import ChatEvent
-from unique_toolkit.chat.service import ChatService
+from unique_toolkit._common.default_language_model import DEFAULT_GPT_35_TURBO
+from unique_toolkit._common.validate_required_values import (
+    validate_required_values,
+)
+from unique_toolkit.app.schemas import BaseEvent, ChatEvent
 from unique_toolkit.evals.config import EvaluationMetricConfig
 from unique_toolkit.evals.context_relevancy.schema import (
     EvaluationSchemaStructuredOutput,
@@ -21,7 +26,6 @@ from unique_toolkit.evals.schemas import (
 )
 from unique_toolkit.language_model.infos import (
     LanguageModelInfo,
-    LanguageModelName,
     ModelCapabilities,
 )
 from unique_toolkit.language_model.prompt import Prompt
@@ -45,9 +49,7 @@ USER_MSG_KEY = "userPrompt"
 default_config = EvaluationMetricConfig(
     enabled=False,
     name=EvaluationMetricName.CONTEXT_RELEVANCY,
-    language_model=LanguageModelInfo.from_name(
-        LanguageModelName.AZURE_GPT_4o_2024_1120
-    ),
+    language_model=LanguageModelInfo.from_name(DEFAULT_GPT_35_TURBO),
     custom_prompts={
         SYSTEM_MSG_KEY: CONTEXT_RELEVANCY_METRIC_SYSTEM_MSG,
         USER_MSG_KEY: CONTEXT_RELEVANCY_METRIC_USER_MSG,
@@ -61,13 +63,42 @@ relevancy_required_input_fields = [
 
 
 class ContextRelevancyEvaluator:
+    @deprecated(
+        "Use __init__ with company_id and user_id instead or use the classmethod `from_event`"
+    )
+    @overload
+    def __init__(self, event: ChatEvent | BaseEvent):
+        """
+        Initialize the ContextRelevancyEvaluator with an event (deprecated)
+        """
+
+    @overload
+    def __init__(self, *, company_id: str, user_id: str):
+        """
+        Initialize the ContextRelevancyEvaluator with a company_id and user_id
+        """
+
     def __init__(
         self,
-        event: ChatEvent,
+        event: ChatEvent | BaseEvent | None = None,
+        company_id: str | None = None,
+        user_id: str | None = None,
     ):
-        self.chat_service = ChatService(event)
-        self.language_model_service = LanguageModelService(event)
-        self.logger = logging.getLogger(f"ContextRelevancyEvaluator.{__name__}")
+        if isinstance(event, (ChatEvent, BaseEvent)):
+            self.language_model_service = LanguageModelService.from_event(event)
+        else:
+            [company_id, user_id] = validate_required_values([company_id, user_id])
+            self.language_model_service = LanguageModelService(
+                company_id=company_id, user_id=user_id
+            )
+
+        # Setup the logger
+        module_name = "ContextRelevancyEvaluator"
+        self.logger = logging.getLogger(f"{module_name}.{__name__}")
+
+    @classmethod
+    def from_event(cls, event: ChatEvent | BaseEvent):
+        return cls(company_id=event.company_id, user_id=event.user_id)
 
     async def analyze(
         self,
@@ -95,6 +126,7 @@ class ContextRelevancyEvaluator:
 
         input.validate_required_fields(relevancy_required_input_fields)
 
+        # TODO: Was already there in monorepo
         if len(input.context_texts) == 0:  # type: ignore
             error_message = "No context texts provided."
             raise EvaluatorException(

unique_toolkit/tools/agent_chunks_hanlder.py

@@ -0,0 +1,65 @@
+from typing_extensions import deprecated
+
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
+
+
+@deprecated("do not use this its only used in old tools")
+class AgentChunksHandler:
+    def __init__(self):
+        self._tool_chunks = {}
+        self._chunks: list[ContentChunk] = []
+        self._references: list[list[ContentReference]] = []
+
+    @property
+    def chunks(self) -> list[ContentChunk]:
+        return self._chunks
+
+    @property
+    def tool_chunks(self) -> dict:
+        return self._tool_chunks
+
+    def extend(self, chunks: list[ContentChunk]):
+        self._chunks.extend(chunks)
+
+    def replace(self, chunks: list[ContentChunk]):
+        self._chunks = chunks
+
+    def add_references(
+        self,
+        references: list[ContentReference],
+    ):
+        self._references.append(references)
+
+    @property
+    def all_references(
+        self,
+    ) -> list[list[ContentReference]]:
+        return self._references
+
+    @property
+    def latest_references(
+        self,
+    ) -> list[ContentReference]:
+        if not self._references:
+            return []
+        return self._references[-1]
+
+    @property
+    def latest_referenced_chunks(self) -> list[ContentChunk]:
+        if not self._references:
+            return []
+        return self._get_referenced_chunks_from_references(self._references[-1])
+
+    def _get_referenced_chunks_from_references(
+        self,
+        references: list[ContentReference],
+    ) -> list[ContentChunk]:
+        """
+        Get _referenced_chunks by matching sourceId from _references with merged id and chunk_id from _chunks.
+        """
+        referenced_chunks: list[ContentChunk] = []
+        for ref in references:
+            for chunk in self._chunks:
+                if ref.source_id == str(chunk.id) + "_" + str(chunk.chunk_id):
+                    referenced_chunks.append(chunk)
+        return referenced_chunks

unique_toolkit/tools/tool.py

@@ -15,12 +15,15 @@ from unique_toolkit.language_model.schemas import (
     LanguageModelMessage,
 )
 from unique_toolkit.language_model.service import LanguageModelService
+from unique_toolkit.tools.agent_chunks_hanlder import AgentChunksHandler
 from unique_toolkit.tools.config import ToolBuildConfig, ToolSelectionPolicy
 from unique_toolkit.tools.schemas import BaseToolConfig, ToolCallResponse, ToolPrompts
 from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
 
 ConfigType = TypeVar("ConfigType", bound=BaseToolConfig)
 
+ToolBuildConfig.model_rebuild()
+
 
 class Tool(ABC, Generic[ConfigType]):
     name: str
@@ -80,6 +83,7 @@ class Tool(ABC, Generic[ConfigType]):
     def get_tool_call_result_for_loop_history(
         self,
         tool_response: ToolCallResponse,
+        agent_chunks_handler: AgentChunksHandler,
     ) -> LanguageModelMessage:
         raise NotImplementedError
 
@@ -150,8 +154,7 @@ class Tool(ABC, Generic[ConfigType]):
     ):
         self.settings = ToolBuildConfig(
             name=self.name,
-            configuration=config,  # type: ignore
-            # the ToolBuildConfig has a wrong type in it to be fixed later.
+            configuration=config,
         )
 
         self.config = config

{unique_toolkit-0.8.31.dist-info → unique_toolkit-0.8.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.8.31
+Version: 0.8.33
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -116,6 +116,14 @@ 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.8.33] - 2025-08-31
+
+fixed tool for `web_search`
+
+## [0.8.32] - 2025-08-30
+
+moved over general packages for `web_search`
+
 ## [0.8.31] - 2025-08-29
 - Add various openai models to supported model list
   - o1

{unique_toolkit-0.8.31.dist-info → unique_toolkit-0.8.33.dist-info}/RECORD

@@ -1,11 +1,18 @@
 unique_toolkit/__init__.py,sha256=nbOYPIKERt-ITsgifrnJhatn1YNR38Ntumw-dCn_tsA,714
 unique_toolkit/_common/_base_service.py,sha256=S8H0rAebx7GsOldA7xInLp3aQJt9yEPDQdsGSFRJsGg,276
 unique_toolkit/_common/_time_utils.py,sha256=ztmTovTvr-3w71Ns2VwXC65OKUUh-sQlzbHdKTQWm-w,135
+unique_toolkit/_common/chunk_relevancy_sorter/config.py,sha256=v6Ljo-WIZCtYJgfaPfpzZegCV0DEw_nNhTzNtw0Jg7c,1744
+unique_toolkit/_common/chunk_relevancy_sorter/exception.py,sha256=1mY4zjbvnXsd5oIxwiVsma09bS2XRnHrxW8KJBGtgCM,126
+unique_toolkit/_common/chunk_relevancy_sorter/schemas.py,sha256=doAWPPx8d0zIqHMXmnJy47Z5_NlblJBhMqo8KE7fyyc,1329
+unique_toolkit/_common/chunk_relevancy_sorter/service.py,sha256=UxYn4xJMNEXQ1afMiT2sMwXgdmlFNPaglVhx6CRRtiM,13864
+unique_toolkit/_common/chunk_relevancy_sorter/tests/test_service.py,sha256=JRvLtJXPmz6bm1tFSSqt01HHVeanSD8zk70XVLJHOEM,8878
 unique_toolkit/_common/default_language_model.py,sha256=tmHSqg6e8G7RmKqmdE_tmLxkSN0x-aGoyUdy6Pl2oAE,334
 unique_toolkit/_common/endpoint_builder.py,sha256=oM6uDmxUqTAJut6MuJQj3bIX4yOccyErWD5bJ1d1lcY,4526
 unique_toolkit/_common/exception.py,sha256=caQIE1btsQnpKCHqL2cgWUSbHup06enQu_Pt7uGUTTE,727
+unique_toolkit/_common/feature_flags/schema.py,sha256=3JpTuld8kK-UQ5B0sbYTu0yqhyFPnChXG2Iv4BNqHdg,539
 unique_toolkit/_common/token/image_token_counting.py,sha256=VpFfZyY0GIH27q_Wy4YNjk2algqvbCtJyzuuROoFQPw,2189
 unique_toolkit/_common/token/token_counting.py,sha256=Jo5B11GLlnaZDi4u7xCzIboLl8zn5cY_dmrogHixVdk,6132
+unique_toolkit/_common/utils/structured_output/schema.py,sha256=Tp7kDYcmKtnUhcuRkH86TSYhylRff0ZZJYb2dLkISts,131
 unique_toolkit/_common/validate_required_values.py,sha256=Y_M1ub9gIKP9qZ45F6Zq3ZHtuIqhmOjl8Z2Vd3avg8w,588
 unique_toolkit/_common/validators.py,sha256=aZwbMho7XszN7lT5RtemaiXgC0WJ4u40oeVgsNGhF4U,2803
 unique_toolkit/app/__init__.py,sha256=ETxYDpEizg_PKmi4JPX_P76ySq-us-xypfAIdKQ1QZU,1284
@@ -39,8 +46,8 @@ unique_toolkit/embedding/service.py,sha256=2KjYlUKxeh-Je8S1mOPiQ735pqAmdbeztycVN
 unique_toolkit/embedding/utils.py,sha256=v86lo__bCJbxZBQ3OcLu5SuwT6NbFfWlcq8iyk6BuzQ,279
 unique_toolkit/evals/config.py,sha256=ywHIrJs5SFdKr1WXfrofWuFfzb0iPQw8iZDpq5oEug4,953
 unique_toolkit/evals/context_relevancy/prompts.py,sha256=EdHFUOB581yVxcOL8482KUv_LzaRjuiem71EF8udYMc,1331
-unique_toolkit/evals/context_relevancy/schema.py,sha256=lm9x0jExOinUk9itqC8ZpgReC7yj1VDwEMppxlZGqpY,2923
-unique_toolkit/evals/context_relevancy/service.py,sha256=txTWIhV65QGFhxG1jCb5TTqZc_c7K9I8pi2HVDTAfm8,8384
+unique_toolkit/evals/context_relevancy/schema.py,sha256=ILA0ClTBFJbtZavkVIrqPz2vj6oql8U50e0G4bh8jEI,2911
+unique_toolkit/evals/context_relevancy/service.py,sha256=xoe0ezA4xy-KaPjSO22xhiqo_WKMuYjlYFKEYJWl25A,9601
 unique_toolkit/evals/evaluation_manager.py,sha256=g-8qa_6_p53C9Okx8iNkuoIXYSJrf-6sQ-xku7bo9kI,7895
 unique_toolkit/evals/exception.py,sha256=7lcVbCyoN4Md1chNJDFxpUYyWbVrcr9dcc3TxWykJTc,115
 unique_toolkit/evals/hallucination/constants.py,sha256=FLcXl5XU07jCvS8YPX9l6UjTaqyQ8YvnSKpx4Z6wZ2Y,1997
@@ -105,6 +112,7 @@ unique_toolkit/tools/a2a/manager.py,sha256=-Vt0pL886zMaaG4sGQhXmbwOhOL928WbZzm4F
 unique_toolkit/tools/a2a/memory.py,sha256=F18kUA3m3NqoKdKAJSwDv8JQneHvZTGOkcZTLNMXAYs,1004
 unique_toolkit/tools/a2a/schema.py,sha256=1R7qIu2l2qnUJDGRuUnZLqPPkHnT3x4d3d4PALoBzcY,296
 unique_toolkit/tools/a2a/service.py,sha256=8yiG2zYqHva4fQJX4dfnE9Bm4vDnbK0k6YRyCSQy9RQ,5147
+unique_toolkit/tools/agent_chunks_hanlder.py,sha256=x32Dp1Z8cVW5i-XzXbaMwX2KHPcNGmqEU-FB4AV9ZGo,1909
 unique_toolkit/tools/config.py,sha256=nYwglwUSb3fxhqYxI83k-qoecuF5Zcol5FMOxHvTNeE,3827
 unique_toolkit/tools/factory.py,sha256=w3uNHuYBIJ330Xi8PTdAkr8G3OMbQH2cBgvk5UT16oE,1253
 unique_toolkit/tools/mcp/__init__.py,sha256=RLF_p-LDRC7GhiB3fdCi4u3bh6V9PY_w26fg61BLyco,122
@@ -114,14 +122,14 @@ unique_toolkit/tools/mcp/tool_wrapper.py,sha256=w7Fbo4FSMYvtgSq7Sqt1dmAPvqHjoBQS
 unique_toolkit/tools/schemas.py,sha256=rArQccbfIv7CWcozClAZ-BVlOwAsjpgL8KUab_WeO3k,4817
 unique_toolkit/tools/test/test_mcp_manager.py,sha256=dySiytBUfRjTOzwY_oGKi_jT6BNVgzZuh1du-EvbcJ4,15627
 unique_toolkit/tools/test/test_tool_progress_reporter.py,sha256=GTtmBqOUo0-4fh_q0lRgxDhwKeankc3FHFD5ULZAm4Y,6299
-unique_toolkit/tools/tool.py,sha256=E0qORpQBrdksAylf_RuPaEbXKi_AlswVGv-67t3DNZI,5879
+unique_toolkit/tools/tool.py,sha256=6kw8oadYfFlV0LEPDTHWiz0l5q_8gM-hLpLjmUn_VNs,5943
 unique_toolkit/tools/tool_manager.py,sha256=uW1uf8mYvbnBN_TVa4L79p_8Sf9HyY4V_nXZqrh3fPM,10206
 unique_toolkit/tools/tool_progress_reporter.py,sha256=ixud9VoHey1vlU1t86cW0-WTvyTwMxNSWBon8I11SUk,7955
 unique_toolkit/tools/utils/execution/execution.py,sha256=vjG2Y6awsGNtlvyQAGCTthQ5thWHYnn-vzZXaYLb3QE,7922
 unique_toolkit/tools/utils/source_handling/schema.py,sha256=vzAyf6ZWNexjMO0OrnB8y2glGkvAilmGGQXd6zcDaKw,870
 unique_toolkit/tools/utils/source_handling/source_formatting.py,sha256=C7uayNbdkNVJdEARA5CENnHtNY1SU6etlaqbgHNyxaQ,9152
 unique_toolkit/tools/utils/source_handling/tests/test_source_formatting.py,sha256=oM5ZxEgzROrnX1229KViCAFjRxl9wCTzWZoinYSHleM,6979
-unique_toolkit-0.8.31.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-0.8.31.dist-info/METADATA,sha256=Kcu6gzn5fnXVRnZIHaplzTJuHArVWkhOH1Z7kCgBeUk,29539
-unique_toolkit-0.8.31.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-0.8.31.dist-info/RECORD,,
+unique_toolkit-0.8.33.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-0.8.33.dist-info/METADATA,sha256=6ohnYw35vWjXiiVkqnq3oF0QVINLnuF1GYGZzntACf4,29666
+unique_toolkit-0.8.33.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-0.8.33.dist-info/RECORD,,