fix: fix UI add and cognify features

Dockerfile

@@ -31,7 +31,7 @@ COPY README.md pyproject.toml uv.lock entrypoint.sh ./
 
 # Install the project's dependencies using the lockfile and settings
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --extra api --extra postgres --extra weaviate --extra qdrant --extra neo4j --extra kuzu --extra llama-index --extra gemini --extra ollama --extra mistral --extra groq --extra anthropic --frozen --no-install-project --no-dev --no-editable
+    uv sync --extra debug --extra api --extra postgres --extra weaviate --extra qdrant --extra neo4j --extra kuzu --extra llama-index --extra gemini --extra ollama --extra mistral --extra groq --extra anthropic --frozen --no-install-project --no-dev --no-editable
 
 # Copy Alembic configuration
 COPY alembic.ini /app/alembic.ini
@@ -41,14 +41,15 @@ COPY alembic/ /app/alembic
 # Installing separately from its dependencies allows optimal layer caching
 COPY ./cognee /app/cognee
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --extra api --extra postgres --extra weaviate --extra qdrant --extra neo4j --extra kuzu --extra llama-index --extra gemini --extra ollama --extra mistral --extra groq --extra anthropic --frozen --no-dev --no-editable
+uv sync --extra debug --extra api --extra postgres --extra weaviate --extra qdrant --extra neo4j --extra kuzu --extra llama-index --extra gemini --extra ollama --extra mistral --extra groq --extra anthropic --frozen --no-dev --no-editable
 
 FROM python:3.12-slim-bookworm
 
 WORKDIR /app
 
-# COPY --from=uv /root/.local /root/.local
 COPY --from=uv /app /app
+# COPY --from=uv /app/.venv /app/.venv
+# COPY --from=uv /root/.local /root/.local
 
 RUN chmod +x /app/entrypoint.sh
 

cognee-frontend/src/app/page.tsx

@@ -43,7 +43,7 @@ export default function Home() {
   const onDataAdd = useCallback((dataset: { id: string }, files: File[]) => {
     return addData(dataset, files)
       .then(() => {
-        showNotification("Data added successfully.", 5000);
+        showNotification("Data added successfully. Please run \"Cognify\" when ready.", 5000);
         openDatasetData(dataset);
       });
   }, [showNotification])
@@ -60,6 +60,14 @@ export default function Home() {
       });
   }, [showNotification]);
 
+  const onCognify = useCallback(() => {
+    const dataset = datasets.find((dataset) => dataset.id === selectedDataset);
+    return onDatasetCognify({
+      id: dataset!.id,
+      name: dataset!.name,
+    });
+  }, [datasets, onDatasetCognify, selectedDataset]);
+
   const {
     value: isSettingsModalOpen,
     setTrue: openSettingsModal,
@@ -95,6 +103,7 @@ export default function Home() {
                 datasetId={selectedDataset}
                 onClose={closeDatasetData}
                 onDataAdd={onDataAdd}
+                onCognify={onCognify}
               />
             </div>
           )}

cognee-frontend/src/app/wizard/CognifyStep/CognifyStep.tsx

@@ -7,7 +7,7 @@ import cognifyDataset from '@/modules/datasets/cognifyDataset';
 
 interface ConfigStepProps {
   onNext: () => void;
-  dataset: { id: string }
+  dataset: { name: string }
 }
 
 export default function CognifyStep({ onNext, dataset }: ConfigStepProps) {

cognee-frontend/src/app/wizard/ExploreStep/ExploreStep.tsx

@@ -2,13 +2,13 @@ import { Explorer } from '@/ui/Partials';
 import { Spacer } from 'ohmy-ui';
 
 interface ExploreStepProps {
-  dataset: { id: string };
+  dataset: { name: string };
 }
 
 export default function ExploreStep({ dataset }: ExploreStepProps) {
   return (
     <Spacer horizontal="3">
-      <Explorer dataset={dataset!} />
+      <Explorer dataset={dataset} />
     </Spacer>
   )
 }

cognee-frontend/src/app/wizard/WizardPage.tsx

@@ -29,7 +29,7 @@ export default function WizardPage({
     setFalse: closeSettingsModal,
   } = useBoolean(false);
 
-  const dataset = { id: 'main' };
+  const dataset = { name: 'main' };
 
   return (
     <main className={styles.main}>

cognee-frontend/src/modules/datasets/cognifyDataset.ts

@@ -1,13 +1,13 @@
 import { fetch } from '@/utils';
 
-export default function cognifyDataset(dataset: { id: string, name: string }) {
+export default function cognifyDataset(dataset: { id?: string, name?: string }) {
   return fetch('/v1/cognify', {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
     },
     body: JSON.stringify({
-      datasets: [dataset.id],
+      datasets: [dataset.id || dataset.name],
     }),
   }).then((response) => response.json());
 }

cognee-frontend/src/modules/exploration/getExplorationGraphUrl.ts

@@ -1,6 +1,6 @@
 import { fetch } from '@/utils';
 
-export default function getExplorationGraphUrl(dataset: { id: string }) {
+export default function getExplorationGraphUrl(dataset: { name: string }) {
   return fetch('/v1/visualize')
       .then(async (response) => {
         if (response.status !== 200) {

cognee-frontend/src/modules/ingestion/DataView/DataView.tsx

@@ -6,6 +6,8 @@ import {
   Text,
   UploadInput,
   CloseIcon,
+  CTAButton,
+  useBoolean,
 } from "ohmy-ui";
 import { fetch } from '@/utils';
 import RawDataPreview from './RawDataPreview';
@@ -28,9 +30,10 @@ interface DataViewProps {
   datasetId: string;
   onClose: () => void;
   onDataAdd: (dataset: DatasetLike, files: File[]) => void;
+  onCognify: () => Promise<any>;
 }
 
-export default function DataView({ datasetId, data, onClose, onDataAdd }: DataViewProps) {
+export default function DataView({ datasetId, data, onClose, onDataAdd, onCognify }: DataViewProps) {
   // const handleDataDelete = () => {};
   const [rawData, setRawData] = useState<ArrayBuffer | null>(null);
   const [selectedData, setSelectedData] = useState<Data | null>(null);
@@ -52,7 +55,19 @@ export default function DataView({ datasetId, data, onClose, onDataAdd }: DataVi
 
   const handleDataAdd = (files: File[]) => {
     onDataAdd({ id: datasetId }, files);
-  }
+  };
+
+  const {
+    value: isCognifyButtonDisabled,
+    setTrue: disableCognifyButton,
+    setFalse: enableCognifyButton,
+  } = useBoolean(false);
+
+  const handleCognify = () => {
+    disableCognifyButton();
+    onCognify()
+      .finally(() => enableCognifyButton());
+  };
 
   return (
     <Stack orientation="vertical" gap="4">
@@ -62,6 +77,11 @@ export default function DataView({ datasetId, data, onClose, onDataAdd }: DataVi
             <Text>Add data</Text>
           </UploadInput>
         </div>
+        <div>
+          <CTAButton disabled={isCognifyButtonDisabled} onClick={handleCognify}>
+              <Text>Cognify</Text>
+          </CTAButton>
+        </div>
         <GhostButton hugContent onClick={onClose}>
           <CloseIcon />
         </GhostButton>

cognee-frontend/src/ui/Partials/Explorer/Explorer.tsx

@@ -7,7 +7,7 @@ import { getExplorationGraphUrl } from '@/modules/exploration';
 import styles from './Explorer.module.css';
 
 interface ExplorerProps {
-  dataset: { id: string };
+  dataset: { name: string };
   className?: string;
   style?: React.CSSProperties;
 }

cognee-frontend/src/ui/Partials/SearchView/SearchView.tsx

@@ -28,9 +28,6 @@ export default function SearchView() {
   }, []);
 
   const searchOptions = [{
-    value: 'INSIGHTS',
-    label: 'Query insights from documents',
-  }, {
     value: 'GRAPH_COMPLETION',
     label: 'Completion using Cognee\'s graph based memory',
   }, {

cognee-mcp/src/server.py

@@ -144,7 +144,7 @@ async def cognify_status():
     """Get status of cognify pipeline"""
     with redirect_stdout(sys.stderr):
         user = await get_default_user()
-        status = await get_pipeline_status([await get_unique_dataset_id("main_dataset", user)])
+        status = await get_pipeline_status([await get_unique_dataset_id("main_dataset", user)], "cognify_pipeline")
         return [types.TextContent(type="text", text=str(status))]
 
 
@@ -153,7 +153,7 @@ async def codify_status():
     """Get status of codify pipeline"""
     with redirect_stdout(sys.stderr):
         user = await get_default_user()
-        status = await get_pipeline_status([await get_unique_dataset_id("codebase", user)])
+        status = await get_pipeline_status([await get_unique_dataset_id("codebase", user)], "cognify_code_pipeline")
         return [types.TextContent(type="text", text=str(status))]
 
 

cognee/api/v1/add/add.py

@@ -1,4 +1,6 @@
 from typing import Union, BinaryIO, List, Optional
+from uuid import NAMESPACE_OID, uuid5
+from cognee.modules.pipelines.operations.log_pipeline_run_initiated import log_pipeline_run_initiated
 from cognee.modules.users.models import User
 from cognee.modules.pipelines import Task
 from cognee.tasks.ingestion import ingest_data, resolve_data_directories
@@ -13,6 +15,22 @@ async def add(
 ):
     tasks = [Task(resolve_data_directories), Task(ingest_data, dataset_name, user, node_set)]
 
+    # Refresh the add pipeline status so data is added to a dataset
+    dataset_id = uuid5(NAMESPACE_OID, f"{dataset_name}{str(user.id)}")
+    await log_pipeline_run_initiated(
+        pipeline_id=uuid5(NAMESPACE_OID, "add_pipeline"),
+        pipeline_name="add_pipeline",
+        dataset_id=dataset_id,
+    )
+
     await cognee_pipeline(
         tasks=tasks, datasets=dataset_name, data=data, user=user, pipeline_name="add_pipeline"
     )
+
+    # Refresh the cognify pipeline status so UI shows the correct status
+    dataset_id = uuid5(NAMESPACE_OID, f"{dataset_name}{str(user.id)}")
+    await log_pipeline_run_initiated(
+        pipeline_id=uuid5(NAMESPACE_OID, "cognify_pipeline"),
+        pipeline_name="cognify_pipeline",
+        dataset_id=dataset_id,
+    )

cognee/api/v1/cognify/cognify.py

@@ -34,7 +34,7 @@ async def cognify(
 ):
     tasks = await get_default_tasks(user, graph_model, chunker, chunk_size, ontology_file_path)
 
-    return await cognee_pipeline(tasks=tasks, datasets=datasets, user=user)
+    return await cognee_pipeline(tasks=tasks, datasets=datasets, user=user, pipeline_name="cognify_pipeline")
 
 
 async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's comment)

cognee/api/v1/datasets/datasets.py

@@ -28,7 +28,7 @@ async def list_data(dataset_id: str):
 
     @staticmethod
     async def get_status(dataset_ids: list[UUID]) -> dict:
-        return await get_pipeline_status(dataset_ids)
+        return await get_pipeline_status(dataset_ids, pipeline_name="cognify_pipeline")
 
     @staticmethod
     async def delete_dataset(dataset_id: str):

cognee/modules/pipelines/models/PipelineRun.py

@@ -6,6 +6,7 @@
 
 
 class PipelineRunStatus(enum.Enum):
+    DATASET_PROCESSING_INITIATED = "DATASET_PROCESSING_INITIATED"
     DATASET_PROCESSING_STARTED = "DATASET_PROCESSING_STARTED"
     DATASET_PROCESSING_COMPLETED = "DATASET_PROCESSING_COMPLETED"
     DATASET_PROCESSING_ERRORED = "DATASET_PROCESSING_ERRORED"

cognee/modules/pipelines/operations/__init__.py

@@ -1,3 +1,4 @@
+from .log_pipeline_run_initiated import log_pipeline_run_initiated
 from .log_pipeline_run_start import log_pipeline_run_start
 from .log_pipeline_run_complete import log_pipeline_run_complete
 from .log_pipeline_run_error import log_pipeline_run_error

cognee/modules/pipelines/operations/get_pipeline_status.py

@@ -5,7 +5,7 @@
 from sqlalchemy.orm import aliased
 
 
-async def get_pipeline_status(dataset_ids: list[UUID]):
+async def get_pipeline_status(dataset_ids: list[UUID], pipeline_name: str):
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:
@@ -20,6 +20,7 @@ async def get_pipeline_status(dataset_ids: list[UUID]):
                 .label("rn"),
             )
             .filter(PipelineRun.dataset_id.in_(dataset_ids))
+            .filter(PipelineRun.pipeline_name == pipeline_name)
             .subquery()
         )
 

cognee/modules/pipelines/operations/log_pipeline_run_initiated.py

@@ -0,0 +1,24 @@
+from uuid import UUID, uuid4
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.pipelines.models import PipelineRun, PipelineRunStatus
+
+
+async def log_pipeline_run_initiated(
+    pipeline_id: str, pipeline_name: str, dataset_id: UUID
+):
+    pipeline_run = PipelineRun(
+        pipeline_run_id=uuid4(),
+        pipeline_name=pipeline_name,
+        pipeline_id=pipeline_id,
+        status=PipelineRunStatus.DATASET_PROCESSING_INITIATED,
+        dataset_id=dataset_id,
+        run_info={},
+    )
+
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        session.add(pipeline_run)
+        await session.commit()
+
+    return pipeline_run

cognee/modules/pipelines/operations/pipeline.py

@@ -1,9 +1,8 @@
 import asyncio
 from cognee.shared.logging_utils import get_logger
 from typing import Union
-from uuid import uuid5, NAMESPACE_OID
 
-from cognee.modules.data.methods import get_datasets, get_datasets_by_name
+from cognee.modules.data.methods import get_datasets
 from cognee.modules.data.methods.get_dataset_data import get_dataset_data
 from cognee.modules.data.methods.get_unique_dataset_id import get_unique_dataset_id
 from cognee.modules.data.models import Data, Dataset
@@ -59,15 +58,16 @@ async def cognee_pipeline(
 
     # If no datasets are provided, work with all existing datasets.
     existing_datasets = await get_datasets(user.id)
-    if datasets is None or len(datasets) == 0:
+
+    if not datasets:
+        # Get datasets from database if none sent.
         datasets = existing_datasets
-        if isinstance(datasets[0], str):
-            datasets = await get_datasets_by_name(datasets, user.id)
     else:
-        # Try to get datasets objects from database, if they don't exist use dataset name
-        datasets_names = await get_datasets_by_name(datasets, user.id)
-        if datasets_names:
-            datasets = datasets_names
+        if existing_datasets:
+            # Filter out datasets that match with the ones from database owned by the user.
+            datasets = [dataset for dataset in existing_datasets if str(dataset.id) in datasets or dataset.name in datasets]
+        else:
+            datasets = [Dataset(id=await get_unique_dataset_id(dataset_name=dataset_name, user=user), name=dataset_name, owner_id=user.id) for dataset_name in datasets]
 
     awaitables = []
 
@@ -88,31 +88,28 @@ async def run_pipeline(
     data=None,
     pipeline_name: str = "custom_pipeline",
 ):
-    if isinstance(dataset, Dataset):
-        check_dataset_name(dataset.name)
-        dataset_id = dataset.id
-    elif isinstance(dataset, str):
-        check_dataset_name(dataset)
-        # Generate id based on unique dataset_id formula
-        dataset_id = await get_unique_dataset_id(dataset_name=dataset, user=user)
+    check_dataset_name(dataset.name)
+
+    dataset_id = dataset.id
 
     if not data:
         data: list[Data] = await get_dataset_data(dataset_id=dataset_id)
 
     # async with update_status_lock: TODO: Add UI lock to prevent multiple backend requests
     if isinstance(dataset, Dataset):
-        task_status = await get_pipeline_status([dataset_id])
+        task_status = await get_pipeline_status([dataset_id], pipeline_name)
     else:
         task_status = [
             PipelineRunStatus.DATASET_PROCESSING_COMPLETED
         ]  # TODO: this is a random assignment, find permanent solution
 
-    if (
-        str(dataset_id) in task_status
-        and task_status[str(dataset_id)] == PipelineRunStatus.DATASET_PROCESSING_STARTED
-    ):
-        logger.info("Dataset %s is already being processed.", dataset_id)
-        return
+    if (str(dataset_id) in task_status):
+        if task_status[str(dataset_id)] == PipelineRunStatus.DATASET_PROCESSING_STARTED:
+            logger.info("Dataset %s is already being processed.", dataset_id)
+            return
+        if task_status[str(dataset_id)] == PipelineRunStatus.DATASET_PROCESSING_COMPLETED:
+            logger.info("Dataset %s is already processed.", dataset_id)
+            return
 
     if not isinstance(tasks, list):
         raise ValueError("Tasks must be a list")