cognee 0.3.0.dev0__py3-none-any.whl → 0.3.2__py3-none-any.whl

cognee/cli/_cognee.py

@@ -1,6 +1,8 @@
 import sys
 import os
 import argparse
+import signal
+import subprocess
 from typing import Any, Sequence, Dict, Type, cast, List
 import click
 
@@ -51,6 +53,31 @@ class DebugAction(argparse.Action):
         fmt.note("Debug mode enabled. Full stack traces will be shown.")
 
 
+class UiAction(argparse.Action):
+    def __init__(
+        self,
+        option_strings: Sequence[str],
+        dest: Any = argparse.SUPPRESS,
+        default: Any = argparse.SUPPRESS,
+        help: str = None,
+    ) -> None:
+        super(UiAction, self).__init__(
+            option_strings=option_strings, dest=dest, default=default, nargs=0, help=help
+        )
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Any,
+        option_string: str = None,
+    ) -> None:
+        # Set a flag to indicate UI should be started
+        global ACTION_EXECUTED
+        ACTION_EXECUTED = True
+        namespace.start_ui = True
+
+
 # Debug functionality is now in cognee.cli.debug module
 
 
@@ -97,6 +124,11 @@ def _create_parser() -> tuple[argparse.ArgumentParser, Dict[str, SupportsCliComm
         action=DebugAction,
         help="Enable debug mode to show full stack traces on exceptions",
     )
+    parser.add_argument(
+        "-ui",
+        action=UiAction,
+        help="Start the cognee web UI interface",
+    )
 
     subparsers = parser.add_subparsers(title="Available commands", dest="command")
 
@@ -140,6 +172,76 @@ def main() -> int:
     parser, installed_commands = _create_parser()
     args = parser.parse_args()
 
+    # Handle UI flag
+    if hasattr(args, "start_ui") and args.start_ui:
+        spawned_pids = []
+
+        def signal_handler(signum, frame):
+            """Handle Ctrl+C and other termination signals"""
+            nonlocal spawned_pids
+            fmt.echo("\nShutting down UI server...")
+
+            for pid in spawned_pids:
+                try:
+                    pgid = os.getpgid(pid)
+                    os.killpg(pgid, signal.SIGTERM)
+                    fmt.success(f"✓ Process group {pgid} (PID {pid}) terminated.")
+                except (OSError, ProcessLookupError) as e:
+                    fmt.warning(f"Could not terminate process {pid}: {e}")
+
+            sys.exit(0)
+
+        signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
+        signal.signal(signal.SIGTERM, signal_handler)  # Termination request
+
+        try:
+            from cognee import start_ui
+
+            fmt.echo("Starting cognee UI...")
+
+            # Callback to capture PIDs of all spawned processes
+            def pid_callback(pid):
+                nonlocal spawned_pids
+                spawned_pids.append(pid)
+
+            server_process = start_ui(
+                host="localhost",
+                port=3000,
+                open_browser=True,
+                start_backend=True,
+                auto_download=True,
+                pid_callback=pid_callback,
+            )
+
+            if server_process:
+                fmt.success("UI server started successfully!")
+                fmt.echo("The interface is available at: http://localhost:3000")
+                fmt.echo("The API backend is available at: http://localhost:8000")
+                fmt.note("Press Ctrl+C to stop the server...")
+
+                try:
+                    # Keep the server running
+                    import time
+
+                    while server_process.poll() is None:  # While process is still running
+                        time.sleep(1)
+                except KeyboardInterrupt:
+                    # This shouldn't happen now due to signal handler, but kept for safety
+                    signal_handler(signal.SIGINT, None)
+
+                return 0
+            else:
+                fmt.error("Failed to start UI server. Check the logs above for details.")
+                signal_handler(signal.SIGTERM, None)
+                return 1
+
+        except Exception as ex:
+            fmt.error(f"Error starting UI: {str(ex)}")
+            signal_handler(signal.SIGTERM, None)
+            if debug.is_debug_enabled():
+                raise ex
+            return 1
+
     if cmd := installed_commands.get(args.command):
         try:
             cmd.execute(args)

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -128,4 +128,4 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
                 question=query, answer=completion, context=context_text, triplets=triplets
             )
 
-        return completion
+        return [completion]

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -138,4 +138,4 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                 question=query, answer=completion, context=context_text, triplets=triplets
             )
 
-        return completion
+        return [completion]

cognee/modules/retrieval/graph_completion_retriever.py

@@ -171,7 +171,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
                 question=query, answer=completion, context=context_text, triplets=triplets
             )
 
-        return completion
+        return [completion]
 
     async def save_qa(self, question: str, answer: str, context: str, triplets: List) -> None:
         """

cognee/modules/retrieval/insights_retriever.py

@@ -96,17 +96,18 @@ class InsightsRetriever(BaseGraphRetriever):
                 unique_node_connections_map[unique_id] = True
                 unique_node_connections.append(node_connection)
 
-        return [
-            Edge(
-                node1=Node(node_id=connection[0]["id"], attributes=connection[0]),
-                node2=Node(node_id=connection[2]["id"], attributes=connection[2]),
-                attributes={
-                    **connection[1],
-                    "relationship_type": connection[1]["relationship_name"],
-                },
-            )
-            for connection in unique_node_connections
-        ]
+        return unique_node_connections
+        # return [
+        #     Edge(
+        #         node1=Node(node_id=connection[0]["id"], attributes=connection[0]),
+        #         node2=Node(node_id=connection[2]["id"], attributes=connection[2]),
+        #         attributes={
+        #             **connection[1],
+        #             "relationship_type": connection[1]["relationship_name"],
+        #         },
+        #     )
+        #     for connection in unique_node_connections
+        # ]
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
         """

cognee/modules/retrieval/temporal_retriever.py

@@ -149,4 +149,4 @@ class TemporalRetriever(GraphCompletionRetriever):
                 system_prompt_path=self.system_prompt_path,
             )
 
-        return completion
+        return [completion]

cognee/modules/search/methods/search.py

@@ -132,14 +132,37 @@ async def search(
             ],
         )
     else:
-        return [
-            SearchResult(
-                search_result=result,
-                dataset_id=datasets[min(index, len(datasets) - 1)].id if datasets else None,
-                dataset_name=datasets[min(index, len(datasets) - 1)].name if datasets else None,
-            )
-            for index, (result, _, datasets) in enumerate(search_results)
-        ]
+        # This is for maintaining backwards compatibility
+        if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
+            return_value = []
+            for search_result in search_results:
+                result, context, datasets = search_result
+                return_value.append(
+                    {
+                        "search_result": result,
+                        "dataset_id": datasets[0].id,
+                        "dataset_name": datasets[0].name,
+                    }
+                )
+            return return_value
+        else:
+            return_value = []
+            for search_result in search_results:
+                result, context, datasets = search_result
+                return_value.append(result)
+            # For maintaining backwards compatibility
+            if len(return_value) == 1 and isinstance(return_value[0], list):
+                return return_value[0]
+            else:
+                return return_value
+        # return [
+        #     SearchResult(
+        #         search_result=result,
+        #         dataset_id=datasets[min(index, len(datasets) - 1)].id if datasets else None,
+        #         dataset_name=datasets[min(index, len(datasets) - 1)].name if datasets else None,
+        #     )
+        #     for index, (result, _, datasets) in enumerate(search_results)
+        # ]
 
 
 async def authorized_search(

cognee/tests/test_permissions.py

@@ -79,7 +79,7 @@ async def main():
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
-    assert search_results[0].dataset_name == "NLP", (
+    assert search_results[0]["dataset_name"] == "NLP", (
         f"Dict must contain dataset name 'NLP': {search_results[0]}"
     )
 
@@ -93,7 +93,7 @@ async def main():
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
-    assert search_results[0].dataset_name == "QUANTUM", (
+    assert search_results[0]["dataset_name"] == "QUANTUM", (
         f"Dict must contain dataset name 'QUANTUM': {search_results[0]}"
     )
 
@@ -170,7 +170,7 @@ async def main():
     for result in search_results:
         print(f"{result}\n")
 
-    assert search_results[0].dataset_name == "QUANTUM", (
+    assert search_results[0]["dataset_name"] == "QUANTUM", (
         f"Dict must contain dataset name 'QUANTUM': {search_results[0]}"
     )
 

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_save_export_path.py

@@ -0,0 +1,116 @@
+import os
+import asyncio
+from uuid import uuid4
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_save_uses_custom_export_path(tmp_path, monkeypatch):
+    # Import target after tmp fixtures are ready
+    from cognee.api.v1.save import save as save_mod
+
+    # Prepare two mock datasets
+    class Dataset:
+        def __init__(self, id_, name):
+            self.id = id_
+            self.name = name
+
+    ds1 = Dataset(uuid4(), "dataset_alpha")
+    ds2 = Dataset(uuid4(), "dataset_beta")
+
+    # Mock dataset discovery
+    async def mock_get_authorized_existing_datasets(datasets, permission_type, user):
+        return [ds1, ds2]
+
+    monkeypatch.setattr(
+        save_mod, "get_authorized_existing_datasets", mock_get_authorized_existing_datasets
+    )
+
+    # Mock data items (with filename collision in ds1)
+    class DataItem:
+        def __init__(self, id_, name, original_path=None):
+            self.id = id_
+            self.name = name
+            self.original_data_location = original_path
+
+    ds1_items = [
+        DataItem(uuid4(), "report.txt", "/root/a/report.txt"),
+        DataItem(uuid4(), "report.txt", "/root/b/report.txt"),  # collision
+    ]
+    ds2_items = [
+        DataItem(uuid4(), "notes.md", "/root/x/notes.md"),
+    ]
+
+    async def mock_get_dataset_data(dataset_id):
+        if dataset_id == ds1.id:
+            return ds1_items
+        if dataset_id == ds2.id:
+            return ds2_items
+        return []
+
+    monkeypatch.setattr(save_mod, "get_dataset_data", mock_get_dataset_data)
+
+    # Mock summary retrieval
+    async def mock_get_document_summaries_text(data_id: str) -> str:
+        return "This is a summary."
+
+    monkeypatch.setattr(save_mod, "_get_document_summaries_text", mock_get_document_summaries_text)
+
+    # Mock questions
+    async def mock_generate_questions(file_name: str, summary_text: str):
+        return ["Q1?", "Q2?", "Q3?"]
+
+    monkeypatch.setattr(save_mod, "_generate_questions", mock_generate_questions)
+
+    # Mock searches per question
+    async def mock_run_searches_for_question(question, dataset_id, search_types, top_k):
+        return {st.value: [f"{question} -> ok"] for st in search_types}
+
+    monkeypatch.setattr(save_mod, "_run_searches_for_question", mock_run_searches_for_question)
+
+    # Use custom export path
+    export_dir = tmp_path / "my_exports"
+    export_dir_str = str(export_dir)
+
+    # Run
+    result = await save_mod.save(
+        datasets=None,
+        export_root_directory=export_dir_str,
+        max_questions=3,
+        search_types=["GRAPH_COMPLETION", "INSIGHTS", "CHUNKS"],
+        top_k=2,
+        include_summary=True,
+        include_ascii_tree=True,
+        concurrency=2,
+        timeout=None,
+    )
+
+    # Verify returned mapping points to our custom path
+    assert str(ds1.id) in result and str(ds2.id) in result
+    assert result[str(ds1.id)].startswith(export_dir_str)
+    assert result[str(ds2.id)].startswith(export_dir_str)
+
+    # Verify directories and files exist
+    ds1_dir = result[str(ds1.id)]
+    ds2_dir = result[str(ds2.id)]
+
+    assert os.path.isdir(ds1_dir)
+    assert os.path.isdir(ds2_dir)
+
+    # index.md present
+    assert os.path.isfile(os.path.join(ds1_dir, "index.md"))
+    assert os.path.isfile(os.path.join(ds2_dir, "index.md"))
+
+    # File markdowns exist; collision handling: two files with similar base
+    ds1_files = [f for f in os.listdir(ds1_dir) if f.endswith(".md") and f != "index.md"]
+    assert len(ds1_files) == 2
+    assert any(f == "report.txt.md" for f in ds1_files)
+    assert any(f.startswith("report.txt__") and f.endswith(".md") for f in ds1_files)
+
+    # Content sanity: ensure question headers exist in one file
+    sample_md_path = os.path.join(ds1_dir, ds1_files[0])
+    with open(sample_md_path, "r", encoding="utf-8") as fh:
+        content = fh.read()
+        assert "## Question ideas" in content
+        assert "## Searches" in content

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):

{cognee-0.3.0.dev0.dist-info → cognee-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.3.0.dev0
+Version: 0.3.2
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee
@@ -44,7 +44,7 @@ Requires-Dist: pydantic-settings<3,>=2.2.1
 Requires-Dist: pydantic<3.0.0,>=2.10.5
 Requires-Dist: pylance<1.0.0,>=0.22.0
 Requires-Dist: pympler<2.0.0,>=1.1
-Requires-Dist: pypdf<6.0.0,>=4.1.0
+Requires-Dist: pypdf<7.0.0,>=4.1.0
 Requires-Dist: python-dotenv<2.0.0,>=1.0.1
 Requires-Dist: python-magic-bin<0.5; platform_system == 'Windows'
 Requires-Dist: python-multipart<1.0.0,>=0.0.20