chore: import upstream snapshot with attribution

tests/unit/query/context_builder/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License

tests/unit/query/context_builder/dynamic_community_selection.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""Tests for dynamic community selection with type handling."""
+
+from unittest.mock import MagicMock
+
+from graphrag.data_model.community import Community
+from graphrag.data_model.community_report import CommunityReport
+from graphrag.query.context_builder.dynamic_community_selection import (
+    DynamicCommunitySelection,
+)
+
+
+def create_mock_tokenizer() -> MagicMock:
+    """Create a mock tokenizer."""
+    tokenizer = MagicMock()
+    tokenizer.encode.return_value = [1, 2, 3]
+    return tokenizer
+
+
+def create_mock_model() -> MagicMock:
+    """Create a mock chat model."""
+    return MagicMock()
+
+
+def test_dynamic_community_selection_handles_int_children():
+    """Test that DynamicCommunitySelection correctly handles children IDs as integers.
+
+    This tests the fix for issue #2004 where children IDs could be integers
+    while self.reports keys are strings, causing child communities to be skipped.
+    """
+    # Create communities with integer children (simulating the bug scenario)
+    # Note: Even though the type annotation says list[str], actual data may have ints
+    communities = [
+        Community(
+            id="comm-0",
+            short_id="0",
+            title="Root Community",
+            level="0",
+            parent="",
+            children=[1, 2],  # type: ignore[list-item]  # Integer children - testing bug fix
+        ),
+        Community(
+            id="comm-1",
+            short_id="1",
+            title="Child Community 1",
+            level="1",
+            parent="0",
+            children=[],
+        ),
+        Community(
+            id="comm-2",
+            short_id="2",
+            title="Child Community 2",
+            level="1",
+            parent="0",
+            children=[],
+        ),
+    ]
+
+    # Create community reports with string community_id
+    reports = [
+        CommunityReport(
+            id="report-0",
+            short_id="0",
+            title="Report 0",
+            community_id="0",
+            summary="Root community summary",
+            full_content="Root community full content",
+            rank=1.0,
+        ),
+        CommunityReport(
+            id="report-1",
+            short_id="1",
+            title="Report 1",
+            community_id="1",
+            summary="Child 1 summary",
+            full_content="Child 1 full content",
+            rank=1.0,
+        ),
+        CommunityReport(
+            id="report-2",
+            short_id="2",
+            title="Report 2",
+            community_id="2",
+            summary="Child 2 summary",
+            full_content="Child 2 full content",
+            rank=1.0,
+        ),
+    ]
+
+    model = create_mock_model()
+    tokenizer = create_mock_tokenizer()
+
+    selector = DynamicCommunitySelection(
+        community_reports=reports,
+        communities=communities,
+        model=model,
+        tokenizer=tokenizer,
+        threshold=1,
+        keep_parent=False,
+        max_level=2,
+    )
+
+    # Verify that reports are keyed by string
+    assert "0" in selector.reports
+    assert "1" in selector.reports
+    assert "2" in selector.reports
+
+    # Verify that communities are keyed by string short_id
+    assert "0" in selector.communities
+    assert "1" in selector.communities
+    assert "2" in selector.communities
+
+    # Verify that the children are properly accessible
+    # Before the fix, int children would fail the `in self.reports` check
+    root_community = selector.communities["0"]
+    for child in root_community.children:
+        child_id = str(child)
+        # This should now work with the fix
+        assert child_id in selector.reports, (
+            f"Child {child} (as '{child_id}') should be found in reports"
+        )
+
+
+def test_dynamic_community_selection_handles_str_children():
+    """Test that DynamicCommunitySelection works correctly with string children IDs."""
+    communities = [
+        Community(
+            id="comm-0",
+            short_id="0",
+            title="Root Community",
+            level="0",
+            parent="",
+            children=["1", "2"],  # String children - expected type
+        ),
+        Community(
+            id="comm-1",
+            short_id="1",
+            title="Child Community 1",
+            level="1",
+            parent="0",
+            children=[],
+        ),
+        Community(
+            id="comm-2",
+            short_id="2",
+            title="Child Community 2",
+            level="1",
+            parent="0",
+            children=[],
+        ),
+    ]
+
+    reports = [
+        CommunityReport(
+            id="report-0",
+            short_id="0",
+            title="Report 0",
+            community_id="0",
+            summary="Root community summary",
+            full_content="Root community full content",
+            rank=1.0,
+        ),
+        CommunityReport(
+            id="report-1",
+            short_id="1",
+            title="Report 1",
+            community_id="1",
+            summary="Child 1 summary",
+            full_content="Child 1 full content",
+            rank=1.0,
+        ),
+        CommunityReport(
+            id="report-2",
+            short_id="2",
+            title="Report 2",
+            community_id="2",
+            summary="Child 2 summary",
+            full_content="Child 2 full content",
+            rank=1.0,
+        ),
+    ]
+
+    model = create_mock_model()
+    tokenizer = create_mock_tokenizer()
+
+    selector = DynamicCommunitySelection(
+        community_reports=reports,
+        communities=communities,
+        model=model,
+        tokenizer=tokenizer,
+        threshold=1,
+        keep_parent=False,
+        max_level=2,
+    )
+
+    # Verify that children can be found in reports
+    root_community = selector.communities["0"]
+    for child in root_community.children:
+        child_id = str(child)
+        assert child_id in selector.reports, (
+            f"Child {child} (as '{child_id}') should be found in reports"
+        )

tests/unit/query/context_builder/test_entity_extraction.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+from typing import Any
+
+from graphrag.data_model.entity import Entity
+from graphrag.query.context_builder.entity_extraction import (
+    EntityVectorStoreKey,
+    map_query_to_entities,
+)
+from graphrag_llm.config import LLMProviderType, ModelConfig
+from graphrag_llm.embedding import create_embedding
+from graphrag_vectors import (
+    TextEmbedder,
+    VectorStore,
+    VectorStoreDocument,
+    VectorStoreSearchResult,
+)
+
+embedding_model = create_embedding(
+    ModelConfig(
+        type=LLMProviderType.MockLLM,
+        model_provider="openai",
+        model="text-embedding-3-small",
+        mock_responses=[1.0, 1.0, 1.0],
+    )
+)
+
+
+class MockVectorStore(VectorStore):
+    def __init__(self, documents: list[VectorStoreDocument]) -> None:
+        super().__init__(index_name="mock")
+        self.documents = documents
+
+    def connect(self, **kwargs: Any) -> None:
+        raise NotImplementedError
+
+    def create_index(self) -> None:
+        raise NotImplementedError
+
+    def load_documents(self, documents: list[VectorStoreDocument]) -> None:
+        raise NotImplementedError
+
+    def insert(self, document: VectorStoreDocument) -> None:
+        raise NotImplementedError
+
+    def similarity_search_by_vector(
+        self,
+        query_embedding: list[float],
+        k: int = 10,
+        select: list[str] | None = None,
+        filters: Any = None,
+        include_vectors: bool = True,
+    ) -> list[VectorStoreSearchResult]:
+        return [
+            VectorStoreSearchResult(document=document, score=1)
+            for document in self.documents[:k]
+        ]
+
+    def similarity_search_by_text(
+        self,
+        text: str,
+        text_embedder: TextEmbedder,
+        k: int = 10,
+        select: list[str] | None = None,
+        filters: Any = None,
+        include_vectors: bool = True,
+    ) -> list[VectorStoreSearchResult]:
+        return sorted(
+            [
+                VectorStoreSearchResult(
+                    document=document,
+                    score=abs(len(text) - len(str(document.id) or "")),
+                )
+                for document in self.documents
+            ],
+            key=lambda x: x.score,
+        )[:k]
+
+    def search_by_id(
+        self, id: str, select: list[str] | None = None, include_vectors: bool = True
+    ) -> VectorStoreDocument:
+        result = self.documents[0]
+        result.id = id
+        return result
+
+    def count(self) -> int:
+        return len(self.documents)
+
+    def remove(self, ids: list[str]) -> None:
+        raise NotImplementedError
+
+    def update(self, document: VectorStoreDocument) -> None:
+        raise NotImplementedError
+
+
+def test_map_query_to_entities():
+    entities = [
+        Entity(
+            id="2da37c7a-50a8-44d4-aa2c-fd401e19976c",
+            short_id="sid1",
+            title="t1",
+            rank=2,
+        ),
+        Entity(
+            id="c4f93564-4507-4ee4-b102-98add401a965",
+            short_id="sid2",
+            title="t22",
+            rank=4,
+        ),
+        Entity(
+            id="7c6f2bc9-47c9-4453-93a3-d2e174a02cd9",
+            short_id="sid3",
+            title="t333",
+            rank=1,
+        ),
+        Entity(
+            id="8fd6d72a-8e9d-4183-8a97-c38bcc971c83",
+            short_id="sid4",
+            title="t4444",
+            rank=3,
+        ),
+    ]
+
+    assert map_query_to_entities(
+        query="t22",
+        text_embedding_vectorstore=MockVectorStore([
+            VectorStoreDocument(id=entity.title, vector=None) for entity in entities
+        ]),
+        text_embedder=embedding_model,
+        all_entities_dict={entity.id: entity for entity in entities},
+        embedding_vectorstore_key=EntityVectorStoreKey.TITLE,
+        k=1,
+        oversample_scaler=1,
+    ) == [
+        Entity(
+            id="c4f93564-4507-4ee4-b102-98add401a965",
+            short_id="sid2",
+            title="t22",
+            rank=4,
+        )
+    ]
+
+    assert map_query_to_entities(
+        query="",
+        text_embedding_vectorstore=MockVectorStore([
+            VectorStoreDocument(id=entity.id, vector=None) for entity in entities
+        ]),
+        text_embedder=embedding_model,
+        all_entities_dict={entity.id: entity for entity in entities},
+        embedding_vectorstore_key=EntityVectorStoreKey.TITLE,
+        k=2,
+    ) == [
+        Entity(
+            id="c4f93564-4507-4ee4-b102-98add401a965",
+            short_id="sid2",
+            title="t22",
+            rank=4,
+        ),
+        Entity(
+            id="8fd6d72a-8e9d-4183-8a97-c38bcc971c83",
+            short_id="sid4",
+            title="t4444",
+            rank=3,
+        ),
+    ]