Fix failing search tests (#1387)

<!-- .github/pull_request_template.md -->

## Description
<!-- 
Please provide a clear, human-generated description of the changes in
this PR.
DO NOT use AI-generated descriptions. We want to understand your thought
process and reasoning.
-->

## Type of Change
<!-- Please check the relevant option -->
- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Code refactoring
- [ ] Performance improvement
- [ ] Other (please specify):

## Changes Made
<!-- List the specific changes made in this PR -->
- 
- 
- 

## Testing
<!-- Describe how you tested your changes -->

## Screenshots/Videos (if applicable)
<!-- Add screenshots or videos to help explain your changes -->

## Pre-submission Checklist
<!-- Please check all boxes that apply before submitting your PR -->
- [ ] **I have tested my changes thoroughly before submitting this PR**
- [ ] **This PR contains minimal changes necessary to address the
issue/feature**
- [ ] My code follows the project's coding standards and style
guidelines
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have added necessary documentation (if applicable)
- [ ] All new and existing tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been
submitted already
- [ ] I have linked any relevant issues in the description
- [ ] My commits have clear and descriptive messages

## Related Issues
<!-- Link any related issues using "Fixes #issue_number" or "Relates to
#issue_number" -->

## Additional Notes
<!-- Add any additional notes, concerns, or context for reviewers -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.

.github/workflows/cli_tests.yml

@@ -60,7 +60,7 @@ jobs:
           python-version: ${{ inputs.python-version }}
 
       - name: Run CLI Unit Tests
-        run: uv run pytest cognee/tests/unit/cli/ -v
+        run: uv run pytest cognee/tests/cli_tests/cli_unit_tests/ -v
 
   cli-integration-tests:
     name: CLI Integration Tests
@@ -87,7 +87,7 @@ jobs:
           python-version: ${{ inputs.python-version }}
 
       - name: Run CLI Integration Tests
-        run: uv run pytest cognee/tests/integration/cli/ -v
+        run: uv run pytest cognee/tests/cli_tests/cli_integration_tests/ -v
 
   cli-functionality-tests:
     name: CLI Functionality Tests
@@ -135,12 +135,12 @@ jobs:
         run: |
           # Test invalid command (should fail gracefully)
           ! uv run python -m cognee.cli._cognee invalid_command
-          
+
           # Test missing required arguments (should fail gracefully)
           ! uv run python -m cognee.cli._cognee search
-          
+
           # Test invalid search type (should fail gracefully)
           ! uv run python -m cognee.cli._cognee search "test" --query-type INVALID_TYPE
-          
+
           # Test invalid chunker (should fail gracefully)
           ! uv run python -m cognee.cli._cognee cognify --chunker InvalidChunker

.github/workflows/test_suites.yml

@@ -34,49 +34,49 @@ jobs:
 
   docker-compose-test:
     name: Docker Compose Test
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/docker_compose.yml
     secrets: inherit
 
   docker-ci-test:
     name: Docker CI test
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/backend_docker_build_test.yml
     secrets: inherit
 
   graph-db-tests:
     name: Graph Database Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/graph_db_tests.yml
     secrets: inherit
 
   temporal-graph-tests:
     name: Temporal Graph Test
-    needs: [ basic-tests, e2e-tests, cli-tests, graph-db-tests ]
+    needs: [ basic-tests, e2e-tests, graph-db-tests ]
     uses: ./.github/workflows/temporal_graph_tests.yml
     secrets: inherit
 
   search-db-tests:
     name: Search Test on Different DBs
-    needs: [basic-tests, e2e-tests, cli-tests, graph-db-tests]
+    needs: [basic-tests, e2e-tests, graph-db-tests]
     uses: ./.github/workflows/search_db_tests.yml
     secrets: inherit
 
   relational-db-migration-tests:
     name: Relational DB Migration Tests
-    needs: [basic-tests, e2e-tests, cli-tests, graph-db-tests]
+    needs: [basic-tests, e2e-tests, graph-db-tests]
     uses: ./.github/workflows/relational_db_migration_tests.yml
     secrets: inherit
 
   notebook-tests:
     name: Notebook Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/notebooks_tests.yml
     secrets: inherit
 
   different-operating-systems-tests:
     name: Operating System and Python Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/test_different_operating_systems.yml
     with:
       python-versions: '["3.10.x", "3.11.x", "3.12.x"]'
@@ -85,20 +85,20 @@ jobs:
   # Matrix-based vector database tests
   vector-db-tests:
     name: Vector DB Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/vector_db_tests.yml
     secrets: inherit
 
   # Matrix-based example tests
   example-tests:
     name: Example Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/examples_tests.yml
     secrets: inherit
 
   mcp-test:
     name: MCP Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/test_mcp.yml
     secrets: inherit
 
@@ -110,20 +110,20 @@ jobs:
 
   s3-file-storage-test:
     name: S3 File Storage Test
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/test_s3_file_storage.yml
     secrets: inherit
 
   # Additional LLM tests
   gemini-tests:
     name: Gemini Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/test_gemini.yml
     secrets: inherit
 
   openrouter-tests:
     name: OpenRouter Tests
-    needs: [basic-tests, e2e-tests, cli-tests]
+    needs: [basic-tests, e2e-tests]
     uses: ./.github/workflows/test_openrouter.yml
     secrets: inherit
 
@@ -133,7 +133,6 @@ jobs:
     needs: [
       basic-tests,
       e2e-tests,
-      cli-tests,
       graph-db-tests,
       notebook-tests,
       different-operating-systems-tests,

cognee/tests/test_relational_db_migration.py

@@ -45,15 +45,13 @@ async def relational_db_migration():
     await migrate_relational_database(graph_engine, schema=schema)
 
     # 1. Search the graph
-    search_results: List[SearchResult] = await cognee.search(
+    search_results = await cognee.search(
         query_type=SearchType.GRAPH_COMPLETION, query_text="Tell me about the artist AC/DC"
-    )  # type: ignore
+    )
     print("Search results:", search_results)
 
     # 2. Assert that the search results contain "AC/DC"
-    assert any("AC/DC" in r.search_result for r in search_results), (
-        "AC/DC not found in search results!"
-    )
+    assert any("AC/DC" in r for r in search_results), "AC/DC not found in search results!"
 
     migration_db_provider = migration_engine.engine.dialect.name
     if migration_db_provider == "postgresql":

cognee/tests/test_search_db.py

@@ -144,13 +144,16 @@ async def main():
         ("GRAPH_COMPLETION_CONTEXT_EXTENSION", completion_ext),
         ("GRAPH_SUMMARY_COMPLETION", completion_sum),
     ]:
-        for search_result in search_results:
-            completion = search_result.search_result
-            assert isinstance(completion, str), f"{name}: should return a string"
-            assert completion.strip(), f"{name}: string should not be empty"
-            assert "netherlands" in completion.lower(), (
-                f"{name}: expected 'netherlands' in result, got: {completion!r}"
-            )
+        assert isinstance(search_results, list), f"{name}: should return a list"
+        assert len(search_results) == 1, (
+            f"{name}: expected single-element list, got {len(search_results)}"
+        )
+        text = search_results[0]
+        assert isinstance(text, str), f"{name}: element should be a string"
+        assert text.strip(), f"{name}: string should not be empty"
+        assert "netherlands" in text.lower(), (
+            f"{name}: expected 'netherlands' in result, got: {text!r}"
+        )
 
     graph_engine = await get_graph_engine()
     graph = await graph_engine.get_graph_data()

cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py

@@ -59,8 +59,10 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         answer = await retriever.get_completion("Who works at Canva?")
 
-        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
-        assert answer.strip(), "Answer must contain only non-empty strings"
+        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+        assert all(isinstance(item, str) and item.strip() for item in answer), (
+            "Answer must contain only non-empty strings"
+        )
 
     @pytest.mark.asyncio
     async def test_graph_completion_extension_context_complex(self):
@@ -140,8 +142,10 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
-        assert answer.strip(), "Answer must contain only non-empty strings"
+        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+        assert all(isinstance(item, str) and item.strip() for item in answer), (
+            "Answer must contain only non-empty strings"
+        )
 
     @pytest.mark.asyncio
     async def test_get_graph_completion_extension_context_on_empty_graph(self):
@@ -171,5 +175,7 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
-        assert answer.strip(), "Answer must contain only non-empty strings"
+        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+        assert all(isinstance(item, str) and item.strip() for item in answer), (
+            "Answer must contain only non-empty strings"
+        )

cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py

@@ -55,8 +55,10 @@ class TestGraphCompletionCoTRetriever:
 
         answer = await retriever.get_completion("Who works at Canva?")
 
-        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
-        assert answer.strip(), "Answer must contain only non-empty strings"
+        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+        assert all(isinstance(item, str) and item.strip() for item in answer), (
+            "Answer must contain only non-empty strings"
+        )
 
     @pytest.mark.asyncio
     async def test_graph_completion_cot_context_complex(self):
@@ -133,8 +135,10 @@ class TestGraphCompletionCoTRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
-        assert answer.strip(), "Answer must contain only non-empty strings"
+        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+        assert all(isinstance(item, str) and item.strip() for item in answer), (
+            "Answer must contain only non-empty strings"
+        )
 
     @pytest.mark.asyncio
     async def test_get_graph_completion_cot_context_on_empty_graph(self):
@@ -164,5 +168,7 @@ class TestGraphCompletionCoTRetriever:
 
         answer = await retriever.get_completion("Who works at Figma?")
 
-        assert isinstance(answer, str), f"Expected string, got {type(answer).__name__}"
-        assert answer.strip(), "Answer must contain only non-empty strings"
+        assert isinstance(answer, list), f"Expected list, got {type(answer).__name__}"
+        assert all(isinstance(item, str) and item.strip() for item in answer), (
+            "Answer must contain only non-empty strings"
+        )

cognee/tests/unit/modules/retrieval/insights_retriever_test.py

@@ -82,7 +82,7 @@ class TestInsightsRetriever:
 
         context = await retriever.get_context("Mike")
 
-        assert context[0].node1.attributes["name"] == "Mike Broski", "Failed to get Mike Broski"
+        assert context[0][0]["name"] == "Mike Broski", "Failed to get Mike Broski"
 
     @pytest.mark.asyncio
     async def test_insights_context_complex(self):
@@ -222,9 +222,7 @@ class TestInsightsRetriever:
 
         context = await retriever.get_context("Christina")
 
-        assert context[0].node1.attributes["name"] == "Christina Mayer", (
-            "Failed to get Christina Mayer"
-        )
+        assert context[0][0]["name"] == "Christina Mayer", "Failed to get Christina Mayer"
 
     @pytest.mark.asyncio
     async def test_insights_context_on_empty_graph(self):

poetry.lock

@@ -11728,4 +11728,4 @@ posthog = ["posthog"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.10,<=3.13"
-content-hash = "576318d370b89d128a7c3e755fe3c898fef4e359acdd3f05f952ae497751fb04"
+content-hash = "1e8cdbf6919cea9657d51b7839630dac7a0d8a2815eca0bd811838a282051625"

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "cognee"
 
-version = "0.3.0"
+version = "0.3.2"
 description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
 authors = [
     { name = "Vasilije Markovic" },

uv.lock

@@ -807,7 +807,7 @@ wheels = [
 
 [[package]]
 name = "cognee"
-version = "0.2.4"
+version = "0.3.2"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles" },
@@ -1029,7 +1029,7 @@ requires-dist = [
     { name = "pylance", specifier = ">=0.22.0,<1.0.0" },
     { name = "pylint", marker = "extra == 'dev'", specifier = ">=3.0.3,<4" },
     { name = "pympler", specifier = ">=1.1,<2.0.0" },
-    { name = "pypdf", specifier = ">=4.1.0,<6.0.0" },
+    { name = "pypdf", specifier = ">=4.1.0,<7.0.0" },
     { name = "pypika", marker = "extra == 'chromadb'", specifier = "==0.48.8" },
     { name = "pyside6", marker = "extra == 'gui'", specifier = ">=6.8.3,<7" },
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.4.0,<8" },