feat: fixes and updates to MCP, retrievers, general fixes

Dockerfile

@@ -1,60 +1,61 @@
-FROM python:3.11-slim
-
-# Define Poetry extras to install
-ARG POETRY_EXTRAS="\
-# API \
-api \
-# Storage & Databases \
-postgres weaviate qdrant neo4j falkordb milvus kuzu chromadb \
-# Notebooks & Interactive Environments \
-notebook \
-# LLM & AI Frameworks \
-langchain llama-index gemini huggingface ollama mistral groq anthropic \
-# Evaluation & Monitoring \
-deepeval evals posthog \
-# Graph Processing & Code Analysis \
-codegraph graphiti \
-# Document Processing \
-docs"
+# Use a Python image with uv pre-installed
+FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim AS uv
+
+# Install the project into `/app`
+WORKDIR /app
+
+# Enable bytecode compilation
+# ENV UV_COMPILE_BYTECODE=1
+
+# Copy from the cache instead of linking since it's a mounted volume
+ENV UV_LINK_MODE=copy
 
 # Set build argument
 ARG DEBUG
 
 # Set environment variable based on the build argument
 ENV DEBUG=${DEBUG}
-ENV PIP_NO_CACHE_DIR=true
-ENV PATH="${PATH}:/root/.poetry/bin"
 
-RUN apt-get update
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    gcc \
+    libpq-dev \
+    git \
+    curl \
+    clang \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
 
-RUN apt-get install -y \
-  gcc \
-  build-essential \
-  libpq-dev
+# Copy pyproject.toml and lockfile first for better caching
+COPY README.md pyproject.toml uv.lock entrypoint.sh ./
 
-WORKDIR /app
-COPY pyproject.toml poetry.lock /app/
-
-RUN pip install poetry
+# Install the project's dependencies using the lockfile and settings
+RUN --mount=type=cache,target=/root/.cache/uv \
+    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
 
-# Don't create virtualenv since Docker is already isolated
-RUN poetry config virtualenvs.create false
+# Copy Alembic configuration
+COPY alembic.ini /app/alembic.ini
+COPY alembic/ /app/alembic
 
-# Install the dependencies using the defined extras
-RUN poetry install --extras "${POETRY_EXTRAS}" --no-root
+# Then, add the rest of the project source code and install it
+# Installing separately from its dependencies allows optimal layer caching
+COPY ./cognee /app/cognee
+RUN --mount=type=cache,target=/root/.cache/uv \
+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
 
-# Set the PYTHONPATH environment variable to include the /app directory
-ENV PYTHONPATH=/app
+FROM python:3.12-slim-bookworm
 
-COPY cognee/ /app/cognee
+WORKDIR /app
 
-# Copy Alembic configuration
-COPY alembic.ini /app/alembic.ini
-COPY alembic/ /app/alembic
+COPY --from=uv /app /app
+# COPY --from=uv /app/.venv /app/.venv
+# COPY --from=uv /root/.local /root/.local
 
-COPY entrypoint.sh /app/entrypoint.sh
 RUN chmod +x /app/entrypoint.sh
 
-RUN sed -i 's/\r$//' /app/entrypoint.sh
+# Place executables in the environment at the front of the path
+ENV PATH="/app/.venv/bin:$PATH"
+
+ENV PYTHONPATH=/app
 
 ENTRYPOINT ["/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',
   }, {
@@ -81,6 +78,8 @@ export default function SearchView() {
 
     scrollToBottom();
 
+    setInputValue('');
+
     const searchTypeValue = searchType.value;
 
     fetch('/v1/search', {
@@ -103,10 +102,12 @@ export default function SearchView() {
             text: convertToSearchTypeOutput(systemMessage, searchTypeValue),
           },
         ]);
-        setInputValue('');
 
         scrollToBottom();
       })
+      .catch(() => {
+        setInputValue(inputValue);
+      });
   }, [inputValue, scrollToBottom, searchType.value]);
 
   const {

cognee-frontend/src/utils/fetch.ts

@@ -1,7 +1,7 @@
 import handleServerErrors from './handleServerErrors';
 
 export default function fetch(url: string, options: RequestInit = {}): Promise<Response> {
-  return global.fetch('http://127.0.0.1:8000/api' + url, {
+  return global.fetch('http://localhost:8000/api' + url, {
     ...options,
     headers: {
       ...options.headers,

cognee-mcp/src/server.py

@@ -144,7 +144,9 @@ 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 +155,9 @@ 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,8 +1,9 @@
 from typing import Union, BinaryIO, List, Optional
-from cognee.modules.users.models import User
+
 from cognee.modules.pipelines import Task
-from cognee.tasks.ingestion import ingest_data, resolve_data_directories
+from cognee.modules.users.models import User
 from cognee.modules.pipelines import cognee_pipeline
+from cognee.tasks.ingestion import ingest_data, resolve_data_directories
 
 
 async def add(

cognee/api/v1/cognify/cognify.py

@@ -34,7 +34,9 @@ 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/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -1,8 +1,8 @@
 import asyncio
-from typing import Generic, List, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
 import lancedb
-from lancedb.pydantic import LanceModel, Vector
 from pydantic import BaseModel
+from lancedb.pydantic import LanceModel, Vector
+from typing import Generic, List, Optional, TypeVar, Union, get_args, get_origin, get_type_hints
 
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.engine import DataPoint
@@ -16,8 +16,6 @@
 from ..utils import normalize_distances
 from ..vector_db_interface import VectorDBInterface
 
-from tenacity import retry, stop_after_attempt, wait_exponential
-
 
 class IndexSchema(DataPoint):
     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