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

unique_toolkit/_common/chunk_relevancy_sorter/service.py

@@ -0,0 +1,374 @@
+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.agentic.evaluation.config import EvaluationMetricConfig
+from unique_toolkit.agentic.evaluation.context_relevancy.schema import (
+    EvaluationSchemaStructuredOutput,
+    StructuredOutputConfig,
+)
+from unique_toolkit.agentic.evaluation.context_relevancy.service import (
+    ContextRelevancyEvaluator,
+)
+from unique_toolkit.agentic.evaluation.exception import EvaluatorException
+from unique_toolkit.agentic.evaluation.schemas import (
+    EvaluationMetricInput,
+    EvaluationMetricName,
+    EvaluationMetricResult,
+)
+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.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,275 @@
+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.agentic.evaluation.context_relevancy.schema import (
+    StructuredOutputConfig,
+)
+from unique_toolkit.agentic.evaluation.schemas import (
+    EvaluationMetricName,
+    EvaluationMetricResult,
+)
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
+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_4o),
+        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/default_language_model.py

@@ -0,0 +1,12 @@
+import warnings
+
+from unique_toolkit.language_model.infos import LanguageModelName
+
+warnings.warn(
+    "unique_toolkit._common.default_language_model is deprecated. "
+    "Import DEFAULT_GPT_4o from unique_toolkit.language_model instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+DEFAULT_GPT_4o = LanguageModelName.AZURE_GPT_4o_2024_1120

unique_toolkit/_common/docx_generator/__init__.py

@@ -0,0 +1,7 @@
+from unique_toolkit._common.docx_generator.config import DocxGeneratorConfig
+from unique_toolkit._common.docx_generator.service import DocxGeneratorService
+
+__all__ = [
+    "DocxGeneratorService",
+    "DocxGeneratorConfig",
+]

unique_toolkit/_common/docx_generator/config.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+
+
+class DocxGeneratorConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    template_content_id: str = Field(
+        default="",
+        description="The content id of the template file uploaded to the knowledge base.",
+    )