unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

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.",
+    )

unique_toolkit/_common/docx_generator/schemas.py

@@ -0,0 +1,80 @@
+from docx.document import Document as DocumentObject
+from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
+from docxtpl import DocxTemplate
+from pydantic import BaseModel
+
+
+class HeadingField(BaseModel):
+    text: str
+    level: int = 4
+    alignment: WD_PARAGRAPH_ALIGNMENT = WD_PARAGRAPH_ALIGNMENT.LEFT
+
+    def add(self, doc: DocumentObject):
+        p = doc.add_heading(self.text, level=self.level)
+        p.alignment = self.alignment
+        return p
+
+    def __str__(self):
+        return f"HeadingField(text={self.text}, level={self.level}, alignment={self.alignment})"
+
+
+class ParagraphField(BaseModel):
+    text: str
+    style: str | None = None
+    alignment: WD_PARAGRAPH_ALIGNMENT = WD_PARAGRAPH_ALIGNMENT.LEFT
+
+    def add(self, doc: DocumentObject):
+        p = doc.add_paragraph(self.text, style=self.style)
+        p.alignment = self.alignment
+        return p
+
+    def __str__(self):
+        return f"ParagraphField(text={self.text}, style={self.style}, alignment={self.alignment})"
+
+
+class RunField(BaseModel):
+    text: str
+    italic: bool | None = False
+    bold: bool | None = False
+    alignment: WD_PARAGRAPH_ALIGNMENT = WD_PARAGRAPH_ALIGNMENT.LEFT
+
+    def __str__(self):
+        return f"RunField(text={self.text}, italic={self.italic}, alignment={self.alignment})"
+
+
+class RunsField(BaseModel):
+    runs: list[RunField]
+    style: str | None = None
+    alignment: WD_PARAGRAPH_ALIGNMENT = WD_PARAGRAPH_ALIGNMENT.LEFT
+
+    def add(self, doc: DocumentObject):
+        if not self.runs:
+            return None
+        p = doc.add_paragraph(style=self.style)
+        for run in self.runs:
+            r = p.add_run(run.text)
+            if run.italic:
+                r.italic = True
+            if run.bold:
+                r.bold = True
+        return p
+
+    def __str__(self):
+        return f"RunsField(runs={self.runs}, style={self.style}, alignment={self.alignment})"
+
+
+class ContentField(BaseModel):
+    contents: list[HeadingField | ParagraphField | RunsField]
+
+    def add(self, doc: DocxTemplate):
+        sd = doc.new_subdoc()
+        for content in self.contents:
+            # if isinstance(content, ImageField):
+            #     content.download_image(self.download_path)
+            #     content.add(sd) # type: ignore
+            # else:
+            content.add(sd)  # type: ignore
+        return sd
+
+    def __str__(self):
+        return f"ContentField(contents={self.contents})"

unique_toolkit/_common/docx_generator/service.py

@@ -0,0 +1,252 @@
+import io
+import logging
+import re
+from pathlib import Path
+
+from docxtpl import DocxTemplate
+from markdown_it import MarkdownIt
+
+from unique_toolkit._common.docx_generator.config import DocxGeneratorConfig
+from unique_toolkit._common.docx_generator.schemas import (
+    ContentField,
+    HeadingField,
+    # ImageField,
+    ParagraphField,
+    RunField,
+    RunsField,
+)
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.services import KnowledgeBaseService
+
+generator_dir_path = Path(__file__).resolve().parent
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class DocxGeneratorService:
+    def __init__(
+        self,
+        chat_service: ChatService,
+        knowledge_base_service: KnowledgeBaseService,
+        config: DocxGeneratorConfig,
+    ):
+        self._knowledge_base_service = knowledge_base_service
+        self._config = config
+
+    @staticmethod
+    def parse_markdown_to_list_content_fields(
+        markdown: str, offset_header_lvl: int = 0
+    ) -> list[HeadingField | ParagraphField | RunsField]:
+        # Initialize markdown-it parser
+        md = MarkdownIt()
+
+        # Preprocess markdown.
+        # - Replace all headings with the correct heading level
+        # - Remove "Relevant sources" heading
+        # - Replace "# Proposed answer" with "#### Proposed answer"
+        markdown = re.sub(r"(?m)^\s*## ", "#### ", markdown)
+        markdown = re.sub(r"(?m)^\s*### ", "##### ", markdown)
+        markdown = markdown.replace("# Relevant sources", "")
+        markdown = markdown.replace("# Proposed answer", "#### Proposed answer")
+
+        tokens = md.parse(markdown)
+
+        elements = []
+        current_section = {}
+        in_list = False
+        bullet_list_indent_level = 0
+        list_item_open = False
+
+        for token in tokens:
+            if token.type == "bullet_list_open":
+                in_list = True
+                bullet_list_indent_level = int(token.level / 2)
+
+            elif token.type == "bullet_list_close":
+                in_list = False
+                bullet_list_indent_level = 0
+
+            elif token.type == "list_item_open":
+                if list_item_open:
+                    elements.append(current_section)
+                list_item_open = True
+                list_level = token.level - bullet_list_indent_level
+                current_section = {
+                    "type": RunsField,
+                    "runs": [],
+                    "is_list_item": True,
+                    "level": list_level,
+                }
+
+            elif token.type == "list_item_close":
+                if current_section and current_section.get("runs"):
+                    elements.append(current_section)
+                current_section = {}
+                list_item_open = False
+
+            if token.type == "heading_open":
+                # Heading start, token.tag gives the level (e.g., 'h1', 'h2', etc.)
+                header_lvl = int(token.tag[1])  # Extract the level number from tag
+                current_section = {
+                    "type": HeadingField,
+                    "text": "",
+                    "level": header_lvl + offset_header_lvl,
+                }
+
+            elif token.type == "heading_close":
+                if current_section:
+                    elements.append(current_section)
+                current_section = {}
+
+            elif token.type == "paragraph_open":
+                if not in_list:  # Only create new paragraph if not in a list
+                    current_section = {"type": RunsField, "runs": []}
+
+            elif token.type == "paragraph_close":
+                if not in_list and current_section:  # Only append if not in a list
+                    elements.append(current_section)
+                    current_section = {}
+
+            elif token.type == "inline":
+                if current_section.get("type") == HeadingField:
+                    content = token.content
+                    if content.startswith("_page"):
+                        # replace "_pageXXXX_" with "PageXXXX", where XXXX can be any characters and numbers
+                        content = re.sub(
+                            r"^_page([a-zA-Z0-9\s-]+)_(.*?)",
+                            r"Page\1",
+                            content,
+                        )
+                        bold = True
+                    current_section["text"] += content
+                elif "runs" in current_section:
+                    bold = False
+                    italic = False
+                    runs = []
+                    if token.children:
+                        for child in token.children:
+                            content = child.content
+                            if child.type == "strong_open":
+                                bold = True
+                            elif child.type == "strong_close":
+                                bold = False
+                            elif child.type == "em_open":
+                                italic = True
+                            elif child.type == "em_close":
+                                italic = False
+                            if child.type == "softbreak":
+                                content += "\n"
+                            if content:  # Only add non-empty content
+                                runs.append(
+                                    RunField(
+                                        text=content,
+                                        bold=bold,
+                                        italic=italic,
+                                    )
+                                )
+                    else:
+                        runs.append(
+                            RunField(
+                                text=token.content,
+                                bold=bold,
+                                italic=italic,
+                            )
+                        )
+
+                    current_section["runs"].extend(runs)
+
+        # Process remaining elements
+        contents = []
+        for element in elements:
+            if not element:
+                continue
+            if element["type"] == HeadingField:
+                contents.append(
+                    HeadingField(
+                        text=element["text"],
+                        level=element["level"],
+                    )
+                )
+            elif element["type"] == RunsField:
+                if element.get("is_list_item", False):
+                    level: int = min(element.get("level", 1), 5)
+                    if level > 1:
+                        style = "List Bullet " + str(level)
+                    else:
+                        style = "List Bullet"
+                    contents.append(RunsField(style=style, runs=element["runs"]))
+                else:
+                    contents.append(RunsField(runs=element["runs"]))
+
+        return contents
+
+    def generate_from_template(
+        self,
+        subdoc_content: list[HeadingField | ParagraphField | RunsField],
+        fields: dict | None = None,
+    ):
+        """
+        Generate a docx file from a template with the given content.
+
+        Args:
+            subdoc_content (list[HeadingField | ParagraphField | RunsField]): The content to be added to the docx file.
+            fields (dict): Other fields to be added to the docx file. Defaults to None.
+        """
+
+        docx_template_object = self._get_template(self._config.template_content_id)
+
+        doc = DocxTemplate(io.BytesIO(docx_template_object))
+
+        try:
+            content = {}
+            content["body"] = ContentField(contents=subdoc_content)
+
+            if fields:
+                content.update(fields)
+
+            for key, value in content.items():
+                if isinstance(value, ContentField):
+                    content[key] = value.add(doc)
+
+            doc.render(content)
+            docx_rendered_object = io.BytesIO()
+
+            doc.save(docx_rendered_object)
+            docx_rendered_object.seek(0)
+
+            return docx_rendered_object.getvalue()
+
+        except Exception as e:
+            _LOGGER.error(f"Error generating docx: {e}")
+            return None
+
+    def _get_template(self, template_content_id: str):
+        try:
+            if template_content_id:
+                _LOGGER.info(
+                    f"Downloading template from content ID: {template_content_id}"
+                )
+                file_content = self._knowledge_base_service.download_content_to_bytes(
+                    content_id=template_content_id
+                )
+            else:
+                _LOGGER.info("No template content ID provided. Using default template.")
+                file_content = self._get_default_template()
+        except Exception as e:
+            _LOGGER.warning(
+                f"An error occurred while downloading the template {e}. Make sure the template content ID is valid. Falling back to default template."
+            )
+            file_content = self._get_default_template()
+
+        return file_content
+
+    def _get_default_template(self):
+        generator_dir_path = Path(__file__).resolve().parent
+        path = generator_dir_path / "template" / "Doc Template.docx"
+
+        file_content = path.read_bytes()
+
+        _LOGGER.info("Template downloaded from default template")
+
+        return file_content