diff --git a/openai/backup-hindsight-adolf.sh b/openai/backup-hindsight-adolf.sh new file mode 100755 index 0000000..69d7ca3 --- /dev/null +++ b/openai/backup-hindsight-adolf.sh @@ -0,0 +1,82 @@ +#!/bin/bash +# Backup script for hindsight (Adolf's long-term memory bank) and the +# openai_adolf-state Docker volume (Matrix E2EE identity/sessions + config). +# Mirrors the seafile/vaultwarden/openai-llm-dbs backup.sh pattern (same repo): +# dump/tar via `docker exec`, gzip, retention of last 5, Zabbix freshness +# trapper item per target. +# +# hindsight is an embedded Postgres (pg0) instance living at +# /mnt/ssd/dbs/hindsight on the host, bind-mounted into the `hindsight` +# container at /home/hindsight/.pg0. We use pg_dump against the live, +# running instance (safe, no downtime/quiescing needed — same rationale as +# openai-llm-dbs). +# +# adolf-state is a named Docker volume (openai_adolf-state) owned by the +# container's `node` user, not readable directly from the host as this +# script's operator. We tar it from inside the `adolf` container instead +# (docker exec has access via the mount; no host-side permission needed). +# +# Run every 3 days via root crontab (same schedule as sibling backups), e.g.: +# 0 4 */3 * * /home/alvis/agap_git/openai/backup-hindsight-adolf.sh >> /mnt/backups/hindsight-adolf/backup.log 2>&1 +# +# Restore: +# # hindsight (drop+recreate the DB first if restoring into a fresh instance, +# # since the dump is a plain SQL dump, not --clean): +# gunzip -c /mnt/backups/hindsight-adolf//hindsight.sql.gz | \ +# docker exec -i -e PGPASSWORD=hindsight hindsight \ +# /home/hindsight/.pg0/installation/18.1.0/bin/psql -U hindsight -h 127.0.0.1 -p 5432 hindsight +# +# # adolf-state (container must be stopped first so files aren't overwritten +# # while in use; extract into the volume's mountpoint): +# docker stop adolf +# docker run --rm -v openai_adolf-state:/target -v /mnt/backups/hindsight-adolf/:/backup:ro \ +# alpine sh -c "rm -rf /target/* && tar xzf /backup/adolf-state.tar.gz -C /target" +# docker start adolf + +set -euo pipefail + +BACKUP_DIR="/mnt/backups/hindsight-adolf" +ZABBIX_TOKEN_FILE="/root/.zabbix_token" +ZABBIX_URL="http://192.168.1.4:81/api_jsonrpc.php" + +DATE=$(date '+%Y%m%d-%H%M') +DEST="$BACKUP_DIR/$DATE" + +mkdir -p "$DEST" + +notify_zabbix() { + local itemid="$1" label="$2" + if [[ -f "$ZABBIX_TOKEN_FILE" ]]; then + local token now_epoch + token=$(cat "$ZABBIX_TOKEN_FILE") + now_epoch=$(date '+%s') + env -u HTTPS_PROXY -u HTTP_PROXY -u ALL_PROXY -u https_proxy -u http_proxy -u all_proxy \ + curl -s -X POST "$ZABBIX_URL" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $token" \ + -d "{\"jsonrpc\":\"2.0\",\"method\":\"history.push\",\"id\":1,\"params\":{\"itemid\":\"$itemid\",\"value\":$now_epoch}}" > /dev/null \ + && echo "Zabbix notified ($label=$now_epoch)." + else + echo "WARNING: $ZABBIX_TOKEN_FILE not found -- skipped Zabbix freshness push for $label." >&2 + fi +} + +# --- hindsight (Postgres logical dump, live/read-only) --- +echo "Dumping hindsight..." +docker exec -e PGPASSWORD=hindsight hindsight \ + /home/hindsight/.pg0/installation/18.1.0/bin/pg_dump -U hindsight -h 127.0.0.1 -p 5432 hindsight \ + | gzip > "$DEST/hindsight.sql.gz" +echo "Dumped: hindsight -> $DEST/hindsight.sql.gz" +notify_zabbix "70639" "hindsight.backup.ts" + +# --- adolf-state (tar the volume from inside the adolf container) --- +echo "Archiving adolf-state..." +docker exec adolf tar czf - -C /home/node/.openclaw . > "$DEST/adolf-state.tar.gz" +echo "Archived: adolf-state -> $DEST/adolf-state.tar.gz" +notify_zabbix "70640" "adolf-state.backup.ts" + +echo "$(date): Backup complete: $DEST" +ls -la "$DEST/" + +# Rotate: keep last 5 backups +ls -1dt "$BACKUP_DIR"/[0-9]*-[0-9]* 2>/dev/null | tail -n +6 | xargs -r rm -rf diff --git a/openai/backup-llm-dbs.sh b/openai/backup-llm-dbs.sh new file mode 100755 index 0000000..debb736 --- /dev/null +++ b/openai/backup-llm-dbs.sh @@ -0,0 +1,65 @@ +#!/bin/bash +# Backup script for litellm-db and langfuse-db (openai stack postgres containers). +# litellm-db holds provisioned virtual keys + spend; langfuse-db holds all traces. +# Mirrors the seafile/vaultwarden backup.sh pattern (same repo): dump via +# `docker exec pg_dump`, gzip, retention of last 5, Zabbix freshness +# trapper item per DB. Uses pg_dump (safe against a live/running DB, no downtime +# needed — unlike gitea's stop-the-world dump). +# +# Run every 3 days via root crontab (same schedule as vaultwarden/seafile), e.g.: +# 0 3 */3 * * /home/alvis/agap_git/openai/backup-llm-dbs.sh >> /mnt/backups/openai-llm-dbs/backup.log 2>&1 +# +# Restore (litellm-db example, langfuse-db is identical with its own container/user/db): +# gunzip -c /mnt/backups/openai-llm-dbs//litellm-db.sql.gz | \ +# docker exec -i litellm-db psql -U litellm -d litellm +# # For langfuse-db: +# gunzip -c /mnt/backups/openai-llm-dbs//langfuse-db.sql.gz | \ +# docker exec -i langfuse-db psql -U langfuse -d langfuse +# # If restoring into a fresh/empty DB, first drop+recreate the DB (or restore +# # to a new container) since the dump is a plain SQL dump, not --clean. + +set -euo pipefail + +BACKUP_DIR="/mnt/backups/openai-llm-dbs" +ZABBIX_TOKEN_FILE="/root/.zabbix_token" +ZABBIX_URL="http://192.168.1.4:81/api_jsonrpc.php" + +DATE=$(date '+%Y%m%d-%H%M') +DEST="$BACKUP_DIR/$DATE" + +mkdir -p "$DEST" + +notify_zabbix() { + local itemid="$1" label="$2" + if [[ -f "$ZABBIX_TOKEN_FILE" ]]; then + local token now_epoch + token=$(cat "$ZABBIX_TOKEN_FILE") + now_epoch=$(date '+%s') + env -u HTTPS_PROXY -u HTTP_PROXY -u ALL_PROXY -u https_proxy -u http_proxy -u all_proxy \ + curl -s -X POST "$ZABBIX_URL" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $token" \ + -d "{\"jsonrpc\":\"2.0\",\"method\":\"history.push\",\"id\":1,\"params\":{\"itemid\":\"$itemid\",\"value\":$now_epoch}}" > /dev/null \ + && echo "Zabbix notified ($label=$now_epoch)." + else + echo "WARNING: $ZABBIX_TOKEN_FILE not found -- skipped Zabbix freshness push for $label." >&2 + fi +} + +# --- litellm-db --- +echo "Dumping litellm-db..." +docker exec litellm-db pg_dump -U litellm litellm | gzip > "$DEST/litellm-db.sql.gz" +echo "Dumped: litellm-db -> $DEST/litellm-db.sql.gz" +notify_zabbix "70637" "litellm.db.backup.ts" + +# --- langfuse-db --- +echo "Dumping langfuse-db..." +docker exec langfuse-db pg_dump -U langfuse langfuse | gzip > "$DEST/langfuse-db.sql.gz" +echo "Dumped: langfuse-db -> $DEST/langfuse-db.sql.gz" +notify_zabbix "70638" "langfuse.db.backup.ts" + +echo "$(date): Backup complete: $DEST" +ls -la "$DEST/" + +# Rotate: keep last 5 backups +ls -1dt "$BACKUP_DIR"/[0-9]*-[0-9]* 2>/dev/null | tail -n +6 | xargs -r rm -rf diff --git a/openai/cognee-mcp/Dockerfile b/openai/cognee-mcp/Dockerfile new file mode 100644 index 0000000..5bddfac --- /dev/null +++ b/openai/cognee-mcp/Dockerfile @@ -0,0 +1,23 @@ +# Adolf kb#70 — cognee-mcp deletion fix. +# +# Base: official upstream image (do not hand-roll cognee-mcp itself). +# Patches exactly two files to fix a real bug: the `forget` MCP tool (the +# only deletion-capable tool actually exposed to agents — `delete`, +# `delete_dataset`, and `prune` exist in src/server.py but are never +# registered with @mcp.tool(), so they're unreachable dead code) never +# exposed a `data_id` parameter, and its cognee_client.forget() wrapper +# never forwarded one either — even though cognee's own /api/v1/forget +# endpoint has always supported single-item deletion via dataset+data_id. +# Net effect: agents could delete an entire dataset but never a single +# entry/fact. Verified 2026-07-07 by calling the live /api/v1/forget +# endpoint directly with data_id — entry-level delete works fine +# server-side; the MCP bridge was just never wired up to use it. +# +# See src/cognee_client.py forget() and src/server.py forget() for the +# fix. Both files are full copies of the upstream 0.5.4 source with only +# the forget-related code changed (diff against the base image at +# /app/src/{cognee_client,server}.py to see the exact delta). +FROM cognee/cognee-mcp:1.2.2 + +COPY src/cognee_client.py /app/src/cognee_client.py +COPY src/server.py /app/src/server.py diff --git a/openai/cognee-mcp/src/cognee_client.py b/openai/cognee-mcp/src/cognee_client.py new file mode 100644 index 0000000..87b837f --- /dev/null +++ b/openai/cognee-mcp/src/cognee_client.py @@ -0,0 +1,629 @@ +""" +Cognee Client abstraction that supports both direct function calls and HTTP API calls. + +This module provides a unified interface for interacting with Cognee, supporting: +- Direct mode: Directly imports and calls cognee functions (default behavior) +- API mode: Makes HTTP requests to a running Cognee FastAPI server +""" + +import sys +import hashlib +from typing import Optional, Any, List, Dict +from uuid import UUID +from contextlib import redirect_stdout +import httpx +from cognee.shared.logging_utils import get_logger +import json + +try: + from .server_utils import normalize_delete_mode +except ImportError: + from server_utils import normalize_delete_mode + +try: + from .retrieval_utils import get_chunk_neighbors_from_graph, get_document_from_graph +except ImportError: + from retrieval_utils import get_chunk_neighbors_from_graph, get_document_from_graph + +logger = get_logger() + + +class CogneeClient: + """ + Unified client for interacting with Cognee via direct calls or HTTP API. + + Parameters + ---------- + api_url : str, optional + Base URL of the Cognee API server (e.g., "http://localhost:8000"). + If None, uses direct cognee function calls. + api_token : str, optional + Authentication token for the API (optional, required if API has authentication enabled). + """ + + def __init__(self, api_url: Optional[str] = None, api_token: Optional[str] = None): + self.api_url = api_url.rstrip("/") if api_url else None + self.api_token = api_token + self.use_api = bool(api_url) + + # Extract tenant ID from tenant URL pattern: tenant-.*.cognee.ai + self.tenant_id: Optional[str] = None + if self.api_url: + import re + + match = re.search(r"tenant-([0-9a-f-]{36})", self.api_url) + if match: + self.tenant_id = match.group(1) + + if self.use_api: + logger.info(f"Cognee client initialized in API mode: {self.api_url}") + if self.tenant_id: + logger.info(f"Tenant ID extracted from URL: {self.tenant_id}") + self.client = httpx.AsyncClient(timeout=300.0) # 5 minute timeout for long operations + else: + logger.info("Cognee client initialized in direct mode") + # Import cognee only if we're using direct mode + import cognee as _cognee + + self.cognee = _cognee + + def _get_headers(self, include_content_type: bool = True) -> Dict[str, str]: + """Get headers for API requests. + + Uses X-Api-Key + X-Tenant-Id for tenant APIs (cloud), + falls back to Bearer token for local/self-hosted backends. + """ + headers: Dict[str, str] = {} + if include_content_type: + headers["Content-Type"] = "application/json" + if self.api_token: + if self.tenant_id: + headers["X-Api-Key"] = self.api_token + headers["X-Tenant-Id"] = self.tenant_id + else: + headers["Authorization"] = f"Bearer {self.api_token}" + return headers + + @staticmethod + def _json_or_success(response: httpx.Response) -> Dict[str, Any]: + """Return a JSON body when present, otherwise a generic success shape.""" + if not response.content: + return {"status": "success"} + try: + parsed = response.json() + except ValueError: + return {"status": "success", "message": response.text} + if isinstance(parsed, dict): + return parsed + return {"status": "success", "result": parsed} + + @staticmethod + def _text_upload(data: Any) -> Dict[str, tuple[str, str, str]]: + """Create a content-addressed text upload for API-mode ingestion.""" + content = str(data) + digest = hashlib.md5(content.encode("utf-8")).hexdigest() + return {"data": (f"text_{digest}.txt", content, "text/plain")} + + async def add( + self, data: Any, dataset_name: str = "main_dataset", node_set: Optional[List[str]] = None + ) -> Dict[str, Any]: + """ + Add data to Cognee for processing. + + Parameters + ---------- + data : Any + Data to add (text, file path, etc.) + dataset_name : str + Name of the dataset to add data to + node_set : List[str], optional + List of node identifiers for graph organization + + Returns + ------- + Dict[str, Any] + Result of the add operation + """ + if self.use_api: + endpoint = f"{self.api_url}/api/v1/add" + + files = self._text_upload(data) + form_data = { + "datasetName": dataset_name, + } + if node_set is not None: + form_data["node_set"] = json.dumps(node_set) + + response = await self.client.post( + endpoint, + files=files, + data=form_data, + headers=self._get_headers(include_content_type=False), + ) + response.raise_for_status() + return response.json() + else: + with redirect_stdout(sys.stderr): + await self.cognee.add(data, dataset_name=dataset_name, node_set=node_set) + return {"status": "success", "message": "Data added successfully"} + + async def cognify( + self, + datasets: Optional[List[str]] = None, + custom_prompt: Optional[str] = None, + graph_model: Any = None, + ) -> Dict[str, Any]: + """ + Transform data into a knowledge graph. + + Parameters + ---------- + datasets : List[str], optional + List of dataset names to process + custom_prompt : str, optional + Custom prompt for entity extraction + graph_model : Any, optional + Custom graph model (only used in direct mode) + + Returns + ------- + Dict[str, Any] + Result of the cognify operation + """ + if self.use_api: + # API mode: Make HTTP request + endpoint = f"{self.api_url}/api/v1/cognify" + payload = { + "datasets": datasets or ["main_dataset"], + "run_in_background": False, + } + if custom_prompt: + payload["custom_prompt"] = custom_prompt + + response = await self.client.post(endpoint, json=payload, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + # Direct mode: Call cognee directly + with redirect_stdout(sys.stderr): + kwargs = {} + if datasets: + kwargs["datasets"] = datasets + if custom_prompt: + kwargs["custom_prompt"] = custom_prompt + if graph_model: + kwargs["graph_model"] = graph_model + + await self.cognee.cognify(**kwargs) + return {"status": "success", "message": "Cognify completed successfully"} + + async def search( + self, + query_text: str, + query_type: str, + datasets: Optional[List[str]] = None, + system_prompt: Optional[str] = None, + top_k: int = 15, + ) -> Any: + """ + Search the knowledge graph. + + Parameters + ---------- + query_text : str + The search query + query_type : str + Type of search (e.g., "GRAPH_COMPLETION", "INSIGHTS", etc.) + datasets : List[str], optional + List of datasets to search + system_prompt : str, optional + System prompt for completion searches + top_k : int + Maximum number of results + + Returns + ------- + Any + Search results + """ + if self.use_api: + # API mode: Make HTTP request + endpoint = f"{self.api_url}/api/v1/search" + payload = {"query": query_text, "search_type": query_type.upper(), "top_k": top_k} + if datasets: + payload["datasets"] = datasets + if system_prompt: + payload["system_prompt"] = system_prompt + + response = await self.client.post(endpoint, json=payload, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + # Direct mode: Call cognee directly + from cognee.modules.search.types import SearchType + + with redirect_stdout(sys.stderr): + search_kwargs = { + "query_type": SearchType[query_type.upper()], + "query_text": query_text, + "top_k": top_k, + } + if datasets: + search_kwargs["datasets"] = datasets + if system_prompt: + search_kwargs["system_prompt"] = system_prompt + results = await self.cognee.search(**search_kwargs) + return results + + async def delete(self, data_id: UUID, dataset_id: UUID, mode: str = "soft") -> Dict[str, Any]: + """ + Delete data from a dataset. + + Parameters + ---------- + data_id : UUID + ID of the data to delete + dataset_id : UUID + ID of the dataset containing the data + + Returns + ------- + Dict[str, Any] + Result of the deletion + """ + normalized_mode = normalize_delete_mode(mode) + + if self.use_api: + # The deprecated delete endpoint still carries the mode contract. + # Fall back to the datasets endpoint for older backends that removed it. + endpoint = f"{self.api_url}/api/v1/delete" + response = await self.client.delete( + endpoint, + params={ + "data_id": str(data_id), + "dataset_id": str(dataset_id), + "mode": normalized_mode, + }, + headers=self._get_headers(), + ) + if response.status_code in {404, 405}: + endpoint = f"{self.api_url}/api/v1/datasets/{str(dataset_id)}/data/{str(data_id)}" + response = await self.client.delete(endpoint, headers=self._get_headers()) + response.raise_for_status() + return self._json_or_success(response) + else: + # Direct mode: Call cognee directly + from cognee.modules.users.methods import get_default_user + + with redirect_stdout(sys.stderr): + user = await get_default_user() + result = await self.cognee.datasets.delete_data( + dataset_id=dataset_id, + data_id=data_id, + mode=normalized_mode, + user=user, + ) + return result or {"status": "success"} + + async def prune_data(self) -> Dict[str, Any]: + """ + Prune all data from the knowledge graph. + + Returns + ------- + Dict[str, Any] + Result of the prune operation + """ + if self.use_api: + # Note: The API doesn't expose a prune endpoint, so we'll need to handle this + # For now, raise an error + raise NotImplementedError("Prune operation is not available via API") + else: + # Direct mode: Call cognee directly + with redirect_stdout(sys.stderr): + await self.cognee.prune.prune_data() + return {"status": "success", "message": "Data pruned successfully"} + + async def prune_system(self, metadata: bool = True) -> Dict[str, Any]: + """ + Prune system data from the knowledge graph. + + Parameters + ---------- + metadata : bool + Whether to prune metadata + + Returns + ------- + Dict[str, Any] + Result of the prune operation + """ + if self.use_api: + # Note: The API doesn't expose a prune endpoint + raise NotImplementedError("Prune system operation is not available via API") + else: + # Direct mode: Call cognee directly + with redirect_stdout(sys.stderr): + await self.cognee.prune.prune_system(metadata=metadata) + return {"status": "success", "message": "System pruned successfully"} + + async def get_pipeline_status( + self, dataset_ids: List[UUID], pipeline_name: str + ) -> Dict[str, Any]: + """ + Get the status of a pipeline run. + + Parameters + ---------- + dataset_ids : List[UUID] + List of dataset IDs + pipeline_name : str + Name of the pipeline + + Returns + ------- + Dict[str, Any] + Status information keyed by dataset ID + """ + if self.use_api: + # API mode: query the server's dataset-status endpoint, which + # reports the pipeline run state keyed by dataset id. + endpoint = f"{self.api_url}/api/v1/datasets/status" + params = [("dataset", str(d)) for d in dataset_ids] + response = await self.client.get(endpoint, params=params, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + # Direct mode: Call cognee directly + from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status + + with redirect_stdout(sys.stderr): + status = await get_pipeline_status(dataset_ids, pipeline_name) + return status + + async def list_datasets(self) -> List[Dict[str, Any]]: + """ + List all datasets. + + Returns + ------- + List[Dict[str, Any]] + List of datasets + """ + if self.use_api: + # API mode: Make HTTP request + endpoint = f"{self.api_url}/api/v1/datasets" + response = await self.client.get(endpoint, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + # Direct mode: Call cognee directly + from cognee.modules.users.methods import get_default_user + from cognee.modules.data.methods import get_datasets + + with redirect_stdout(sys.stderr): + user = await get_default_user() + datasets = await get_datasets(user.id) + return [ + {"id": str(d.id), "name": d.name, "created_at": str(d.created_at)} + for d in datasets + ] + + async def get_document( + self, + document_id: str, + include_metadata: bool = True, + max_chunks: int = 0, + ) -> Dict[str, Any]: + """Retrieve a full document with its chunks from the graph database.""" + if self.use_api: + raise NotImplementedError("get_document is not available in API mode") + + from cognee.infrastructure.databases.unified import get_unified_engine + + with redirect_stdout(sys.stderr): + unified = await get_unified_engine() + return await get_document_from_graph( + unified.graph, + document_id, + include_metadata=include_metadata, + max_chunks=max_chunks, + ) + + async def get_chunk_neighbors( + self, + chunk_id: str, + neighbor_count: int = 2, + include_target: bool = True, + direction: str = "both", + ) -> Dict[str, Any]: + """Retrieve neighboring chunks around a target chunk from its parent document.""" + if self.use_api: + raise NotImplementedError("get_chunk_neighbors is not available in API mode") + + from cognee.infrastructure.databases.unified import get_unified_engine + + with redirect_stdout(sys.stderr): + unified = await get_unified_engine() + return await get_chunk_neighbors_from_graph( + unified.graph, + chunk_id, + neighbor_count=neighbor_count, + include_target=include_target, + direction=direction, + ) + + # -- V2 API methods ----------------------------------------------------- + + async def remember( + self, + data: Any, + dataset_name: str = "main_dataset", + session_id: Optional[str] = None, + custom_prompt: Optional[str] = None, + ) -> Dict[str, Any]: + """Store data in memory via remember(). + + With session_id: stores in session cache only (fast). + Without session_id: full add + cognify pipeline (permanent). + """ + if self.use_api: + if session_id: + if custom_prompt: + logger.warning( + "remember: custom_prompt is not supported with session_id in API mode " + "(the /remember/entry endpoint does not forward custom_prompt)" + ) + raise ValueError( + "custom_prompt is not supported when session_id is provided in API mode" + ) + # Session mode: POST a JSON QAEntry so the backend receives + # real text, not a multipart-file placeholder that triggers + # the _SESSION_PLACEHOLDER_PREFIXES skip in _add_to_session. + endpoint = f"{self.api_url}/api/v1/remember/entry" + payload = { + "entry": { + "type": "qa", + "question": "", + "answer": str(data), + "context": "", + }, + "dataset_name": dataset_name, + "session_id": session_id, + } + response = await self.client.post( + endpoint, + json=payload, + headers=self._get_headers(), + ) + response.raise_for_status() + return response.json() + + endpoint = f"{self.api_url}/api/v1/remember" + files = self._text_upload(data) + form_data = {"datasetName": dataset_name} + if custom_prompt: + form_data["custom_prompt"] = custom_prompt + response = await self.client.post( + endpoint, + files=files, + data=form_data, + headers=self._get_headers(include_content_type=False), + ) + response.raise_for_status() + return response.json() + else: + with redirect_stdout(sys.stderr): + kwargs = { + "data": data, + "dataset_name": dataset_name, + } + if session_id: + kwargs["session_id"] = session_id + if custom_prompt: + kwargs["custom_prompt"] = custom_prompt + result = await self.cognee.remember(**kwargs) + return { + "status": getattr(result, "status", "completed"), + "dataset_name": dataset_name, + "session_id": session_id, + } + + async def recall( + self, + query_text: str, + search_type: Optional[str] = None, + datasets: Optional[List[str]] = None, + session_id: Optional[str] = None, + top_k: int = 15, + ) -> Any: + """Search memory via recall() with auto-routing and session awareness.""" + if self.use_api: + endpoint = f"{self.api_url}/api/v1/recall" + payload = {"query": query_text, "top_k": top_k, "search_type": None} + if search_type: + payload["search_type"] = search_type.upper() + if datasets: + payload["datasets"] = datasets + if session_id: + payload["session_id"] = session_id + response = await self.client.post(endpoint, json=payload, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + with redirect_stdout(sys.stderr): + kwargs = {"top_k": top_k, "auto_route": True} + if search_type: + from cognee.modules.search.types import SearchType + + kwargs["query_type"] = SearchType[search_type.upper()] + if datasets: + kwargs["datasets"] = datasets + if session_id: + kwargs["session_id"] = session_id + return await self.cognee.recall(query_text=query_text, **kwargs) + + async def forget( + self, + dataset: Optional[str] = None, + data_id: Optional[UUID] = None, + dataset_id: Optional[UUID] = None, + everything: bool = False, + memory_only: bool = False, + ) -> Dict[str, Any]: + """Delete data via forget(). + + Bug fix (kb#70): this method previously dropped `data_id`, + `dataset_id`, and `memory_only` on the floor, so entry-level + deletion was impossible through the MCP surface even though the + cognee API's /api/v1/forget endpoint has always supported it + (dataset/datasetId + dataId). Forward all fields it accepts. + """ + if self.use_api: + endpoint = f"{self.api_url}/api/v1/forget" + payload = {"everything": everything, "memory_only": memory_only} + if dataset: + payload["dataset"] = dataset + if dataset_id: + payload["dataset_id"] = str(dataset_id) + if data_id: + payload["data_id"] = str(data_id) + response = await self.client.post(endpoint, json=payload, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + with redirect_stdout(sys.stderr): + return await self.cognee.forget( + dataset=dataset, + dataset_id=dataset_id, + data_id=data_id, + everything=everything, + memory_only=memory_only, + ) + + async def improve( + self, + dataset_name: str = "main_dataset", + session_ids: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """Enrich knowledge graph and bridge session data via improve().""" + if self.use_api: + endpoint = f"{self.api_url}/api/v1/improve" + payload = {"dataset_name": dataset_name} + if session_ids: + payload["session_ids"] = session_ids + response = await self.client.post(endpoint, json=payload, headers=self._get_headers()) + response.raise_for_status() + return response.json() + else: + with redirect_stdout(sys.stderr): + kwargs = {"dataset": dataset_name} + if session_ids: + kwargs["session_ids"] = session_ids + result = await self.cognee.improve(**kwargs) + return {"status": "success", "result": str(result)} + + async def close(self): + """Close the HTTP client if in API mode.""" + if self.use_api and hasattr(self, "client"): + await self.client.aclose() diff --git a/openai/cognee-mcp/src/server.py b/openai/cognee-mcp/src/server.py new file mode 100644 index 0000000..017effe --- /dev/null +++ b/openai/cognee-mcp/src/server.py @@ -0,0 +1,2071 @@ +import json +import os +import sys +import argparse +import asyncio +import subprocess +from collections import deque +from datetime import datetime, timezone +from pathlib import Path +from typing import Deque, List, Optional, Tuple +from cognee.modules.data.methods.get_datasets_by_name import get_datasets_by_name +from cognee.modules.data.methods.get_last_added_data import get_last_added_data +from cognee.modules.users.methods import get_default_user +from cognee.shared.logging_utils import get_logger, setup_logging, get_log_file_location +from cognee.shared.usage_logger import log_usage +import importlib.util +from contextlib import redirect_stdout +import mcp.types as types +from mcp.server import FastMCP +from mcp.server.transport_security import TransportSecuritySettings +from cognee.modules.storage.utils import JSONEncoder +from starlette.responses import JSONResponse +from starlette.middleware import Middleware +from starlette.middleware.cors import CORSMiddleware +import uvicorn + +try: + from .cognee_client import CogneeClient +except ImportError: + from cognee_client import CogneeClient + +try: + from .strip_vectors import strip_vectors +except ImportError: + from strip_vectors import strip_vectors + +try: + from .server_utils import ( + format_recall_results, + format_search_results, + normalize_delete_mode, + normalize_search_type, + parse_cognify_data, + parse_csv_list, + validate_cognify_file_paths, + validate_top_k, + ) +except ImportError: + from server_utils import ( + format_recall_results, + format_search_results, + normalize_delete_mode, + normalize_search_type, + parse_cognify_data, + parse_csv_list, + validate_cognify_file_paths, + validate_top_k, + ) + + +try: + from cognee.tasks.codingagents.coding_rule_associations import ( + add_rule_associations, + get_existing_rules, + ) +except ModuleNotFoundError: + from .codingagents.coding_rule_associations import ( + add_rule_associations, + get_existing_rules, + ) + + +mcp = FastMCP("Cognee") + +logger = get_logger() + +cognee_client: Optional[CogneeClient] = None + +# Per-dataset error ring buffer (bounded so long-running servers don't accumulate +# unbounded memory). Each entry is (iso_timestamp, error_message). +_TASK_ERROR_HISTORY = 50 +_task_errors: dict[str, Deque[Tuple[str, str]]] = {} + +# Strong references to in-flight background tasks. asyncio's event loop only keeps +# weak references to tasks, so a fire-and-forget task can be GC'd mid-execution if +# the only reference is a local that went out of scope. Adding here pins them; the +# done_callback removes them on completion. See: +# https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task +_background_tasks: set[asyncio.Task] = set() + + +def _track_background(coro) -> asyncio.Task: + """Spawn a background task and pin it so the event loop won't GC it.""" + task = asyncio.create_task(coro) + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) + return task + + +def _record_task_error(dataset: str, error: str) -> None: + """Append a background task error, bounded per-dataset.""" + bucket = _task_errors.setdefault(dataset, deque(maxlen=_TASK_ERROR_HISTORY)) + bucket.append((datetime.now(timezone.utc).isoformat(), error)) + + +def _configure_transport_security(host: str) -> None: + """Configure MCP transport security based on env vars and bind host. + + Must be called before run_sse_with_cors() or run_http_with_cors(), since + the SDK reads mcp.settings.transport_security lazily when creating the app. + + Env vars: + MCP_DISABLE_DNS_REBINDING_PROTECTION: Set to "true" to disable all + Host/Origin header validation. Useful for LAN or Docker deployments. + MCP_ALLOWED_HOSTS: Comma-separated additional Host header patterns + (e.g. "192.168.1.50:*,myserver.local:*"). Appended to the + localhost defaults. Requires the ":*" port glob suffix. + """ + disable = os.getenv("MCP_DISABLE_DNS_REBINDING_PROTECTION", "false").lower() == "true" + + if disable: + mcp.settings.transport_security = TransportSecuritySettings( + enable_dns_rebinding_protection=False, + ) + logger.info("MCP transport security: DNS rebinding protection disabled") + return + + extra_hosts = [h.strip() for h in os.getenv("MCP_ALLOWED_HOSTS", "").split(",") if h.strip()] + + # The SDK only auto-populates localhost defaults when transport_security is + # None AND host is a loopback address. When the user binds to 0.0.0.0 or a + # LAN IP, we must provide the full allowed list ourselves. + localhost_hosts = ["127.0.0.1:*", "localhost:*", "[::1]:*"] + localhost_origins = ["http://127.0.0.1:*", "http://localhost:*", "http://[::1]:*"] + + allowed_hosts = localhost_hosts + extra_hosts + # Derive origins from extra hosts so users don't need to set both. + allowed_origins = localhost_origins + [f"http://{h}" for h in extra_hosts] + + if host not in ("127.0.0.1", "localhost", "::1") or extra_hosts: + mcp.settings.transport_security = TransportSecuritySettings( + enable_dns_rebinding_protection=True, + allowed_hosts=allowed_hosts, + allowed_origins=allowed_origins, + ) + logger.info( + "MCP transport security: allowed_hosts=%s", + allowed_hosts, + ) + else: + # Loopback-only with no extra hosts — let the SDK use its own defaults. + logger.info("MCP transport security: using SDK defaults (localhost only)") + + +def _is_running_in_docker() -> bool: + """Check if the process is running inside a Docker container.""" + return os.path.exists("/.dockerenv") or os.path.isdir("/app") + + +def _get_cors_origins() -> list[str]: + """Parse CORS allowed origins from MCP_CORS_ALLOW_ORIGINS env var.""" + raw = os.getenv("MCP_CORS_ALLOW_ORIGINS", "http://localhost:3000") + return [o.strip() for o in raw.split(",") if o.strip()] + + +async def run_sse_with_cors(): + """Custom SSE transport with CORS middleware.""" + sse_app = mcp.sse_app() + sse_app.add_middleware( + CORSMiddleware, + allow_origins=_get_cors_origins(), + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + config = uvicorn.Config( + sse_app, + host=mcp.settings.host, + port=mcp.settings.port, + log_level=mcp.settings.log_level.lower(), + ) + server = uvicorn.Server(config) + await server.serve() + + +async def run_http_with_cors(): + """Custom HTTP transport with CORS middleware.""" + http_app = mcp.streamable_http_app() + http_app.add_middleware( + CORSMiddleware, + allow_origins=_get_cors_origins(), + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + config = uvicorn.Config( + http_app, + host=mcp.settings.host, + port=mcp.settings.port, + log_level=mcp.settings.log_level.lower(), + ) + server = uvicorn.Server(config) + await server.serve() + + +@mcp.custom_route("/health", methods=["GET"]) +async def health_check(request): + return JSONResponse({"status": "ok"}) + + +@log_usage(function_name="MCP cognify", log_type="mcp_tool") +async def cognify( + data: str, + dataset_name: str = None, + graph_model_file: str = None, + graph_model_name: str = None, + custom_prompt: str = None, +) -> list: + """ + Transform ingested data into a structured knowledge graph. + + This is the core processing step in Cognee that converts raw text and documents + into an intelligent knowledge graph. It analyzes content, extracts entities and + relationships, and creates semantic connections for enhanced search and reasoning. + + Prerequisites: + - **LLM_API_KEY**: Must be configured (required for entity extraction and graph generation) + - **Vector Database**: Must be accessible for embeddings storage + - **Graph Database**: Must be accessible for relationship storage + + Input Requirements: + - **Content Types**: Works with any text-extractable content including: + * Natural language documents + * Structured data (CSV, JSON) + * Code repositories + * Academic papers and technical documentation + * Mixed multimedia content (with text extraction) + + Processing Pipeline: + 1. **Document Classification**: Identifies document types and structures + 2. **Permission Validation**: Ensures user has processing rights + 3. **Text Chunking**: Breaks content into semantically meaningful segments + 4. **Entity Extraction**: Identifies key concepts, people, places, organizations + 5. **Relationship Detection**: Discovers connections between entities + 6. **Graph Construction**: Builds semantic knowledge graph with embeddings + 7. **Content Summarization**: Creates hierarchical summaries for navigation + + Parameters + ---------- + data : str + The data to be processed and transformed into structured knowledge. + This can include natural language, file location, or any text-based information + that should become part of the agent's memory. + + graph_model_file : str, optional + Path to a custom schema file that defines the structure of the generated knowledge graph. + If provided, this file will be loaded using importlib to create a custom graph model. + Default is None, which uses Cognee's built-in KnowledgeGraph model. + + graph_model_name : str, optional + Name of the class within the graph_model_file to instantiate as the graph model. + Required if graph_model_file is specified. + Default is None, which uses the default KnowledgeGraph class. + + custom_prompt : str, optional + Custom prompt string to use for entity extraction and graph generation. + If provided, this prompt will be used instead of the default prompts for + knowledge graph extraction. The prompt should guide the LLM on how to + extract entities and relationships from the text content. + + Returns + ------- + list + A list containing a single TextContent object with information about the + background task launch and how to check its status. + + Next Steps: + After successful cognify processing, use search functions to query the knowledge: + + ```python + import cognee + from cognee import SearchType + + # Process your data into knowledge graph + await cognee.cognify() + + # Query for insights using different search types: + + # 1. Natural language completion with graph context + insights = await cognee.search( + "What are the main themes?", + query_type=SearchType.GRAPH_COMPLETION + ) + + # 2. Get entity relationships and connections + relationships = await cognee.search( + "connections between concepts", + query_type=SearchType.GRAPH_COMPLETION + ) + + # 3. Find relevant document chunks + chunks = await cognee.search( + "specific topic", + query_type=SearchType.CHUNKS + ) + ``` + + Environment Variables: + Required: + - LLM_API_KEY: API key for your LLM provider + + Optional: + - LLM_PROVIDER, LLM_MODEL, VECTOR_DB_PROVIDER, GRAPH_DATABASE_PROVIDER + - LLM_RATE_LIMIT_ENABLED: Enable rate limiting (default: False) + - LLM_RATE_LIMIT_REQUESTS: Max requests per interval (default: 60) + + Notes + ----- + - The function launches a background task and returns immediately + - The actual cognify process may take significant time depending on text length + - Check the log file for progress + + """ + + dataset_name = dataset_name or _agent_scoped_default_dataset() + + try: + parsed_data = parse_cognify_data(data) + except ValueError as e: + return [ + types.TextContent( + type="text", + text=f"Error: {str(e)}", + ) + ] + + file_error = validate_cognify_file_paths( + parsed_data.items, + is_running_in_docker=_is_running_in_docker, + ) + if file_error: + return [ + types.TextContent( + type="text", + text=f"Error: {file_error}", + ) + ] + + async def cognify_task( + data_items: list[str], + dataset_name: str = "main_dataset", + graph_model_file: str = None, + graph_model_name: str = None, + custom_prompt: str = None, + ) -> str: + """Build knowledge graph from the input text""" + # NOTE: MCP uses stdout to communicate, we must redirect all output + # going to stdout ( like the print function ) to stderr. + with redirect_stdout(sys.stderr): + logger.info("Cognify process starting.") + + graph_model = None + if graph_model_file and graph_model_name: + if cognee_client.use_api: + logger.warning("Custom graph models are not supported in API mode, ignoring.") + else: + from cognee.shared.data_models import KnowledgeGraph + + graph_model = load_class(graph_model_file, graph_model_name) + + for data_item in data_items: + await cognee_client.add(data_item, dataset_name=dataset_name) + + try: + await cognee_client.cognify( + datasets=[dataset_name], custom_prompt=custom_prompt, graph_model=graph_model + ) + logger.info("Cognify process finished.") + except Exception as e: + logger.error("Cognify process failed.") + raise ValueError(f"Failed to cognify: {str(e)}") from e + + async def cognify_task_wrapper(**kwargs): + """Wrapper that captures errors from the background task.""" + try: + await cognify_task(**kwargs) + except Exception as e: + dataset = kwargs.get("dataset_name", "main_dataset") + _record_task_error(dataset, str(e)) + logger.error(f"Background cognify task failed for dataset '{dataset}': {e}") + + _track_background( + cognify_task_wrapper( + data_items=parsed_data.items, + dataset_name=dataset_name, + graph_model_file=graph_model_file, + graph_model_name=graph_model_name, + custom_prompt=custom_prompt, + ) + ) + + log_file = get_log_file_location() + text = ( + f"Background process launched due to MCP timeout limitations.\n" + f"Queued {len(parsed_data.items)} item(s) for dataset '{dataset_name}'.\n" + f"Check the log file at: {log_file}" + ) + + return [ + types.TextContent( + type="text", + text=text, + ) + ] + + +@log_usage(function_name="MCP save_interaction", log_type="mcp_tool") +async def save_interaction(data: str) -> list: + """ + Transform and save a user-agent interaction into structured knowledge. + + Parameters + ---------- + data : str + The input string containing user queries and corresponding agent answers. + + Returns + ------- + list + A list containing a single TextContent object with information about the background task launch. + """ + + async def save_user_agent_interaction(data: str) -> None: + """Build knowledge graph from the interaction data""" + with redirect_stdout(sys.stderr): + logger.info("Save interaction process starting.") + + await cognee_client.add(data, node_set=["user_agent_interaction"]) + + try: + await cognee_client.cognify() + + user = await get_default_user() + datasets = await get_datasets_by_name("main_dataset", user_id=user.id) + dataset = datasets[0] + added_data = await get_last_added_data(dataset.id) + + logger.info("Save interaction process finished.") + + # Rule associations only work in direct mode + if not cognee_client.use_api: + logger.info("Generating associated rules from interaction data.") + await add_rule_associations( + data=data, + rules_nodeset_name="coding_agent_rules", + context={ + "user": user, + "dataset": dataset, + "data": added_data, + }, + ) + logger.info("Associated rules generated from interaction data.") + else: + logger.warning("Rule associations are not available in API mode, skipping.") + + except Exception as e: + logger.error("Save interaction process failed.") + raise ValueError(f"Failed to Save interaction: {str(e)}") from e + + async def save_task_wrapper(**kwargs): + """Wrapper that captures errors from the background task.""" + try: + await save_user_agent_interaction(**kwargs) + except Exception as e: + _record_task_error("main_dataset", str(e)) + logger.error(f"Background save_interaction task failed: {e}") + + _track_background(save_task_wrapper(data=data)) + + log_file = get_log_file_location() + text = ( + f"Background process launched to process the user-agent interaction.\n" + f"Check the log file at: {log_file}" + ) + + return [ + types.TextContent( + type="text", + text=text, + ) + ] + + +@log_usage(function_name="MCP search", log_type="mcp_tool") +async def search( + search_query: str, search_type: str, top_k: int = 15, datasets: str = None +) -> list: + """ + Search and query the knowledge graph for insights, information, and connections. + + This is the final step in the Cognee workflow that retrieves information from the + processed knowledge graph. It supports multiple search modes optimized for different + use cases - from simple fact retrieval to complex reasoning and code analysis. + + Search Prerequisites: + - **LLM_API_KEY**: Required for GRAPH_COMPLETION and RAG_COMPLETION search types + - **Data Added**: Must have data previously added via `cognee.add()` + - **Knowledge Graph Built**: Must have processed data via `cognee.cognify()` + - **Vector Database**: Must be accessible for semantic search functionality + + Search Types & Use Cases: + + **GRAPH_COMPLETION** (Recommended): + Natural language Q&A using full graph context and LLM reasoning. + Best for: Complex questions, analysis, summaries, insights. + Returns: Conversational AI responses with graph-backed context. + + **RAG_COMPLETION**: + Traditional RAG using document chunks without graph structure. + Best for: Direct document retrieval, specific fact-finding. + Returns: LLM responses based on relevant text chunks. + + **CHUNKS**: + Raw text segments that match the query semantically. + Best for: Finding specific passages, citations, exact content. + Returns: Ranked list of relevant text chunks with metadata. + + **SUMMARIES**: + Pre-generated hierarchical summaries of content. + Best for: Quick overviews, document abstracts, topic summaries. + Returns: Multi-level summaries from detailed to high-level. + + **CODE**: + Code-specific search with syntax and semantic understanding. + Best for: Finding functions, classes, implementation patterns. + Returns: Structured code information with context and relationships. + + **CYPHER**: + Direct graph database queries using Cypher syntax. + Best for: Advanced users, specific graph traversals, debugging. + Returns: Raw graph query results. + + **FEELING_LUCKY**: + Intelligently selects and runs the most appropriate search type. + Best for: General-purpose queries or when you're unsure which search type is best. + Returns: The results from the automatically selected search type. + + Parameters + ---------- + search_query : str + Your question or search query in natural language. + Examples: + - "What are the main themes in this research?" + - "How do these concepts relate to each other?" + - "Find information about machine learning algorithms" + - "What functions handle user authentication?" + + search_type : str + The type of search to perform. Valid options include: + - "GRAPH_COMPLETION": Returns an LLM response based on the search query and Cognee's memory + - "RAG_COMPLETION": Returns an LLM response based on the search query and standard RAG data + - "CODE": Returns code-related knowledge in JSON format + - "CHUNKS": Returns raw text chunks from the knowledge graph + - "SUMMARIES": Returns pre-generated hierarchical summaries + - "CYPHER": Direct graph database queries + - "FEELING_LUCKY": Automatically selects best search type + + The search_type is case-insensitive and will be converted to uppercase. + + top_k : int, optional + Maximum number of results to return (default: 10). + Controls the amount of context retrieved from the knowledge graph. + - Lower values (3-5): Faster, more focused results + - Higher values (10-20): More comprehensive, but slower and more context-heavy + Helps manage response size and context window usage in MCP clients. + + Returns + ------- + list + A list containing a single TextContent object with the search results. + The format of the result depends on the search_type: + - **GRAPH_COMPLETION/RAG_COMPLETION**: Conversational AI response strings + - **CHUNKS**: Relevant text passages with source metadata + - **SUMMARIES**: Hierarchical summaries from general to specific + - **CODE**: Structured code information with context + - **FEELING_LUCKY**: Results in format of automatically selected search type + - **CYPHER**: Raw graph query results + + Performance & Optimization: + - **GRAPH_COMPLETION**: Slower but most intelligent, uses LLM + graph context + - **RAG_COMPLETION**: Medium speed, uses LLM + document chunks (no graph traversal) + - **CHUNKS**: Fastest, pure vector similarity search without LLM + - **SUMMARIES**: Fast, returns pre-computed summaries + - **CODE**: Medium speed, specialized for code understanding + - **FEELING_LUCKY**: Variable speed, uses LLM + search type selection intelligently + + Environment Variables: + Required for LLM-based search types (GRAPH_COMPLETION, RAG_COMPLETION): + - LLM_API_KEY: API key for your LLM provider + + Optional: + - LLM_PROVIDER, LLM_MODEL: Configure LLM for search responses + - VECTOR_DB_PROVIDER: Must match what was used during cognify + - GRAPH_DATABASE_PROVIDER: Must match what was used during cognify + + Notes + ----- + - Different search types produce different output formats + - The function handles the conversion between Cognee's internal result format and MCP's output format + + """ + + try: + normalized_search_type = normalize_search_type(search_type) + normalized_top_k = validate_top_k(top_k) + except ValueError as e: + return [types.TextContent(type="text", text=f"Error: {str(e)}")] + + async def search_task( + search_query: str, search_type: str, top_k: int, datasets_list: list = None + ) -> str: + """ + Internal task to execute knowledge graph search with result formatting. + + Handles the actual search execution and formats results appropriately + for MCP clients based on the search type and execution mode (API vs direct). + + Parameters + ---------- + search_query : str + The search query in natural language + search_type : str + Type of search to perform (GRAPH_COMPLETION, CHUNKS, etc.) + top_k : int + Maximum number of results to return + + Returns + ------- + str + Formatted search results as a string, with format depending on search_type + """ + # NOTE: MCP uses stdout to communicate, we must redirect all output + # going to stdout ( like the print function ) to stderr. + with redirect_stdout(sys.stderr): + search_results = await cognee_client.search( + query_text=search_query, + query_type=search_type, + top_k=top_k, + datasets=datasets_list, + ) + + # Strip embedding vectors from results to save LLM context + # text_vector contains raw floats (~92KB per result), useless for clients + search_results = strip_vectors(search_results) + + if not cognee_client.use_api and search_type == "INSIGHTS": + return retrieved_edges_to_string(search_results) + + return format_search_results( + search_results, + search_type, + json_encoder=JSONEncoder, + ) + + # Parse comma-separated datasets into list + datasets_list = parse_csv_list(datasets) + try: + search_results = await search_task( + search_query, + normalized_search_type, + normalized_top_k, + datasets_list, + ) + except Exception as e: + error_msg = f"Search failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + return [types.TextContent(type="text", text=search_results)] + + +@log_usage(function_name="MCP get_document", log_type="mcp_tool") +async def get_document( + document_id: str, + include_metadata: bool = True, + max_chunks: int = 0, +) -> list: + """ + Retrieve a complete source document and its chunks from the knowledge graph. + + Use this after a CHUNKS search or list_data lookup when you need the full source + context around a result. If a chunk ID is provided instead of a document ID, the + tool resolves the chunk's parent document and returns that document. + + Parameters + ---------- + document_id : str + Document ID to retrieve. A DocumentChunk ID is also accepted and resolves to + its parent document. + include_metadata : bool + Include document metadata fields in the response (default: True). + max_chunks : int + Maximum chunks to return. Use 0 to return all chunks. + """ + with redirect_stdout(sys.stderr): + try: + result = await cognee_client.get_document( + document_id=document_id, + include_metadata=include_metadata, + max_chunks=max_chunks, + ) + return [ + types.TextContent( + type="text", + text=json.dumps(result, indent=2, cls=JSONEncoder), + ) + ] + except Exception as e: + error_msg = f"get_document failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + + +@log_usage(function_name="MCP get_chunk_neighbors", log_type="mcp_tool") +async def get_chunk_neighbors( + chunk_id: str, + neighbor_count: int = 2, + include_target: bool = True, + direction: str = "both", +) -> list: + """ + Retrieve neighboring chunks around a target chunk from the same document. + + Use this after a CHUNKS search when the matching passage is too narrow and you + need local narrative context. Chunks are returned in reading order. + + Parameters + ---------- + chunk_id : str + Target DocumentChunk ID. + neighbor_count : int + Number of neighboring chunks to retrieve on each side. Must be 1-10. + include_target : bool + Include the target chunk in the returned chunk list (default: True). + direction : str + One of "both", "forward", or "backward". + """ + with redirect_stdout(sys.stderr): + try: + result = await cognee_client.get_chunk_neighbors( + chunk_id=chunk_id, + neighbor_count=neighbor_count, + include_target=include_target, + direction=direction, + ) + return [ + types.TextContent( + type="text", + text=json.dumps(result, indent=2, cls=JSONEncoder), + ) + ] + except Exception as e: + error_msg = f"get_chunk_neighbors failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + + +@log_usage(function_name="MCP list_data", log_type="mcp_tool") +async def list_data(dataset_id: str = None) -> list: + """ + List all datasets and their data items with IDs for deletion operations. + + This function helps users identify data IDs and dataset IDs that can be used + with the delete tool. It provides a comprehensive view of available data. + + Parameters + ---------- + dataset_id : str, optional + If provided, only list data items from this specific dataset. + If None, lists all datasets and their data items. + Should be a valid UUID string. + + Returns + ------- + list + A list containing a single TextContent object with formatted information + about datasets and data items, including their IDs for deletion. + + Notes + ----- + - Use this tool to identify data_id and dataset_id values for the delete tool + - The output includes both dataset information and individual data items + - UUIDs are displayed in a format ready for use with other tools + """ + from uuid import UUID + + with redirect_stdout(sys.stderr): + try: + output_lines = [] + + if dataset_id: + # Detailed data listing for specific dataset is only available in direct mode + if cognee_client.use_api: + return [ + types.TextContent( + type="text", + text="❌ Detailed data listing for specific datasets is not available in API mode.\nPlease use the API directly or use direct mode.", + ) + ] + + from cognee.modules.users.methods import get_default_user + from cognee.modules.data.methods import get_dataset, get_dataset_data + + logger.info(f"Listing data for dataset: {dataset_id}") + dataset_uuid = UUID(dataset_id) + user = await get_default_user() + + dataset = await get_dataset(user.id, dataset_uuid) + + if not dataset: + return [ + types.TextContent(type="text", text=f"❌ Dataset not found: {dataset_id}") + ] + + # Get data items in the dataset + data_items = await get_dataset_data(dataset.id) + + output_lines.append(f"📁 Dataset: {dataset.name}") + output_lines.append(f" ID: {dataset.id}") + output_lines.append(f" Created: {dataset.created_at}") + output_lines.append(f" Data items: {len(data_items)}") + output_lines.append("") + + if data_items: + for i, data_item in enumerate(data_items, 1): + output_lines.append(f" 📄 Data item #{i}:") + output_lines.append(f" Data ID: {data_item.id}") + output_lines.append(f" Name: {data_item.name or 'Unnamed'}") + output_lines.append(f" Created: {data_item.created_at}") + output_lines.append("") + else: + output_lines.append(" (No data items in this dataset)") + + else: + # List all datasets - works in both modes + logger.info("Listing all datasets") + datasets = await cognee_client.list_datasets() + + if not datasets: + return [ + types.TextContent( + type="text", + text="📂 No datasets found.\nUse the cognify tool to create your first dataset!", + ) + ] + + output_lines.append("📂 Available Datasets:") + output_lines.append("=" * 50) + output_lines.append("") + + for i, dataset in enumerate(datasets, 1): + # In API mode, dataset is a dict; in direct mode, it's formatted as dict + if isinstance(dataset, dict): + output_lines.append(f"{i}. 📁 {dataset.get('name', 'Unnamed')}") + output_lines.append(f" Dataset ID: {dataset.get('id')}") + output_lines.append(f" Created: {dataset.get('created_at', 'N/A')}") + else: + output_lines.append(f"{i}. 📁 {dataset.name}") + output_lines.append(f" Dataset ID: {dataset.id}") + output_lines.append(f" Created: {dataset.created_at}") + output_lines.append("") + + if not cognee_client.use_api: + output_lines.append("💡 To see data items in a specific dataset, use:") + output_lines.append(' list_data(dataset_id="your-dataset-id-here")') + output_lines.append("") + output_lines.append("🗑️ To delete specific data, use:") + output_lines.append(' delete(data_id="data-id", dataset_id="dataset-id")') + + result_text = "\n".join(output_lines) + logger.info("List data operation completed successfully") + + return [types.TextContent(type="text", text=result_text)] + + except ValueError as e: + error_msg = f"❌ Invalid UUID format: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=error_msg)] + + except Exception as e: + error_msg = f"❌ Failed to list data: {str(e)}" + logger.error(f"List data error: {str(e)}") + return [types.TextContent(type="text", text=error_msg)] + + +@log_usage(function_name="MCP delete_dataset", log_type="mcp_tool") +async def delete_dataset(dataset_name: str) -> list: + """ + Delete an entire dataset and all its data from the knowledge graph. + + This removes the dataset completely: graph data, vector indices, + and metadata in the relational database. This operation cannot be undone. + + Parameters + ---------- + dataset_name : str + The name of the dataset to delete (e.g. 'main_dataset'). + + Returns + ------- + list + A list containing a TextContent with deletion status. + """ + with redirect_stdout(sys.stderr): + try: + if cognee_client.use_api: + return [ + types.TextContent( + type="text", + text="❌ delete_dataset is not available in API mode. Use the API directly.", + ) + ] + + from cognee.modules.users.methods import get_default_user + from cognee.modules.data.methods import delete_dataset as _delete_dataset + from cognee.modules.data.methods import get_datasets + + user = await get_default_user() + datasets = await get_datasets(user.id) + matching = [ds for ds in datasets if ds.name == dataset_name] + + if not matching: + return [types.TextContent(type="text", text=f"Dataset '{dataset_name}' not found.")] + + if len(matching) > 1: + ids = ", ".join(str(ds.id) for ds in matching) + return [ + types.TextContent( + type="text", + text=f"Multiple datasets named '{dataset_name}' found (IDs: {ids}). Please delete by ID instead.", + ) + ] + + await _delete_dataset(matching[0]) + return [ + types.TextContent( + type="text", + text=f"Dataset '{dataset_name}' deleted successfully. Graph, vectors, and metadata removed.", + ) + ] + except Exception as e: + return [types.TextContent(type="text", text=f"Error deleting dataset: {str(e)}")] + + +@log_usage(function_name="MCP delete", log_type="mcp_tool") +async def delete(data_id: str, dataset_id: str, mode: str = "soft") -> list: + """ + Delete specific data from a dataset in the Cognee knowledge graph. + + This function removes a specific data item from a dataset while keeping the + dataset itself intact. It supports both soft and hard deletion modes. + + Parameters + ---------- + data_id : str + The UUID of the data item to delete from the knowledge graph. + This should be a valid UUID string identifying the specific data item. + + dataset_id : str + The UUID of the dataset containing the data to be deleted. + This should be a valid UUID string identifying the dataset. + + mode : str, optional + The deletion mode to use. Options are: + - "soft" (default): Removes the data but keeps related entities that might be shared + - "hard": Also removes degree-one entity nodes that become orphaned after deletion + Default is "soft" for safer deletion that preserves shared knowledge. + + Returns + ------- + list + A list containing a single TextContent object with the deletion results, + including status, deleted node counts, and confirmation details. + + Notes + ----- + - This operation cannot be undone. The specified data will be permanently removed. + - Hard mode may remove additional entity nodes that become orphaned + - The function provides detailed feedback about what was deleted + - Use this for targeted deletion instead of the prune tool which removes everything + """ + from uuid import UUID + + with redirect_stdout(sys.stderr): + try: + normalized_mode = normalize_delete_mode(mode) + logger.info( + f"Starting delete operation for data_id: {data_id}, dataset_id: {dataset_id}, mode: {normalized_mode}" + ) + + # Convert string UUIDs to UUID objects + data_uuid = UUID(data_id) + dataset_uuid = UUID(dataset_id) + + # Call the cognee delete function via client + result = await cognee_client.delete( + data_id=data_uuid, dataset_id=dataset_uuid, mode=normalized_mode + ) + + logger.info(f"Delete operation completed successfully: {result}") + + # Format the result for MCP response + formatted_result = json.dumps(result, indent=2, cls=JSONEncoder) + + return [ + types.TextContent( + type="text", + text=f"✅ Delete operation completed successfully!\n\n{formatted_result}", + ) + ] + + except ValueError as e: + error_msg = f"❌ Invalid delete request: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=error_msg)] + + except Exception as e: + # Handle all other errors (DocumentNotFoundError, DatasetNotFoundError, etc.) + error_msg = f"❌ Delete operation failed: {str(e)}" + logger.error(f"Delete operation error: {str(e)}") + return [types.TextContent(type="text", text=error_msg)] + + +@log_usage(function_name="MCP prune", log_type="mcp_tool") +async def prune(): + """ + Reset the Cognee knowledge graph by removing all stored information. + + This function performs a complete reset of both the data layer and system layer + of the Cognee knowledge graph, removing all nodes, edges, and associated metadata. + It is typically used during development or when needing to start fresh with a new + knowledge base. + + Returns + ------- + list + A list containing a single TextContent object with confirmation of the prune operation. + + Notes + ----- + - This operation cannot be undone. All memory data will be permanently deleted. + - The function prunes both data content (using prune_data) and system metadata (using prune_system) + - This operation is not available in API mode + """ + with redirect_stdout(sys.stderr): + try: + await cognee_client.prune_data() + await cognee_client.prune_system(metadata=True) + return [types.TextContent(type="text", text="Pruned")] + except NotImplementedError: + error_msg = "❌ Prune operation is not available in API mode" + logger.error(error_msg) + return [types.TextContent(type="text", text=error_msg)] + except Exception as e: + error_msg = f"❌ Prune operation failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=error_msg)] + + +# --------------------------------------------------------------------------- +# Session-aware memory operations (remember, recall, forget) +# --------------------------------------------------------------------------- + + +@mcp.tool() +@log_usage(function_name="MCP remember", log_type="mcp_tool") +async def remember( + data: str, + dataset_name: str = None, + session_id: str = None, + custom_prompt: str = None, +) -> list: + """Store data in memory. + + Two modes depending on whether session_id is provided: + + Without session_id (permanent memory): Runs the full add + cognify + pipeline to ingest data and build the knowledge graph. + + With session_id (session memory): Stores the data in the session + cache only. Fast, no entity extraction. Omit session_id when the + content should be stored as permanent graph memory. + + Parameters + ---------- + data : str + The data to store (text content). + dataset_name : str, optional + Target dataset name. Defaults to the current MCP client's + agent-scoped dataset (e.g. "cursor_vscode_memory"), or + "main_dataset" if no client identity is detected. + session_id : str, optional + Session ID. When set, stores in session cache only. + custom_prompt : str, optional + Custom prompt for entity extraction (permanent mode only). + """ + dataset_name = dataset_name or _agent_scoped_default_dataset() + with redirect_stdout(sys.stderr): + try: + result = await cognee_client.remember( + data=data, + dataset_name=dataset_name, + session_id=session_id, + custom_prompt=custom_prompt, + ) + status = result.get("status", "completed") + if session_id: + text = f"Stored in session cache (session_id={session_id}, status={status})." + else: + text = f"Stored permanently in knowledge graph (dataset={dataset_name}, status={status})." + return [types.TextContent(type="text", text=text)] + except Exception as e: + error_msg = f"Remember failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + + +@mcp.tool() +@log_usage(function_name="MCP recall", log_type="mcp_tool") +async def recall( + query: str, + search_type: str = None, + datasets: str = None, + session_id: str = None, + top_k: int = 15, +) -> list: + """Search memory with auto-routing and session awareness. + + When session_id is provided without datasets or search_type, + searches session cache first by keyword matching. Falls through + to the permanent knowledge graph if no session results match. + + Auto-routing picks the best search strategy when search_type + is not specified. + + Parameters + ---------- + query : str + Natural language query to search for. + search_type : str, optional + Override auto-routing. Options: GRAPH_COMPLETION, + GRAPH_COMPLETION_COT, RAG_COMPLETION, CHUNKS, SUMMARIES, + TEMPORAL, FEELING_LUCKY, etc. + datasets : str, optional + Comma-separated dataset names to search within. + session_id : str, optional + Session ID for session-first search. + top_k : int + Maximum results to return (default: 10). + """ + with redirect_stdout(sys.stderr): + try: + normalized_top_k = validate_top_k(top_k) + dataset_list = parse_csv_list(datasets) + results = await cognee_client.recall( + query_text=query, + search_type=search_type, + datasets=dataset_list, + session_id=session_id, + top_k=normalized_top_k, + ) + return [ + types.TextContent( + type="text", + text=format_recall_results(results, json_encoder=JSONEncoder), + ) + ] + except Exception as e: + error_msg = f"Recall failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + + +@mcp.tool() +@log_usage(function_name="MCP forget", log_type="mcp_tool") +async def forget( + dataset: str = None, + data_id: str = None, + everything: bool = False, + memory_only: bool = False, +) -> list: + """Delete data from memory. + + Can target a single data item, an entire dataset, or delete everything + the user owns. Removes data from the relational DB, graph DB, and + vector DB (unless memory_only is set). + + Parameters + ---------- + dataset : str, optional + Dataset name. Alone, deletes the entire dataset. Combined with + data_id, scopes the single-item delete to this dataset. + data_id : str, optional + UUID of a single data item to remove (e.g. from + list_dataset_data_json). Requires 'dataset' to also be set. + everything : bool + If true, delete ALL data across all datasets. Ignores dataset/data_id. + memory_only : bool + If true, delete only the knowledge graph + vector embeddings for + the target (the dataset, or the single data item when data_id is + set), preserving the raw file/data record so it can be + re-cognified later. + + Notes + ----- + - Bug fix (kb#70): this tool previously had no `data_id` parameter at + all, so individual entries/facts could never be deleted -- only + whole datasets. The underlying cognee API always supported + dataset+data_id deletion; this tool just never exposed it. + """ + with redirect_stdout(sys.stderr): + try: + if not dataset and not everything: + return [ + types.TextContent( + type="text", + text="Error: Specify 'dataset' name (optionally with 'data_id') or set 'everything' to true.", + ) + ] + if data_id and not dataset: + return [ + types.TextContent( + type="text", + text="Error: 'data_id' requires 'dataset' to also be set.", + ) + ] + + data_uuid = None + if data_id: + from uuid import UUID + + try: + data_uuid = UUID(data_id) + except ValueError as e: + return [ + types.TextContent( + type="text", text=f"Error: 'data_id' is not a valid UUID: {e}" + ) + ] + + result = await cognee_client.forget( + dataset=dataset, + data_id=data_uuid, + everything=everything, + memory_only=memory_only, + ) + status = result.get("status", "unknown") if isinstance(result, dict) else "completed" + ok = status == "success" + prefix = "✅" if ok else "⚠️" + if everything: + text = f"{prefix} All data deleted (status={status})." + elif data_id: + text = ( + f"{prefix} Data item '{data_id}' deleted from dataset '{dataset}' " + f"(status={status})." + ) + else: + text = f"{prefix} Dataset '{dataset}' deleted (status={status})." + return [types.TextContent(type="text", text=text)] + except Exception as e: + error_msg = f"Forget failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + + +@log_usage(function_name="MCP improve", log_type="mcp_tool") +async def improve( + dataset_name: str = None, + session_ids: str = None, +) -> list: + """Enrich the knowledge graph and bridge session data to the permanent graph. + + When session_ids is provided, runs a 4-stage pipeline: + 1. Apply feedback weights from session scores to graph nodes/edges + 2. Persist session Q&A text into the permanent knowledge graph + 3. Enrich graph with triplet embeddings (memify) + 4. Sync enriched graph knowledge back into session caches + + Without session_ids, only stage 3 runs (triplet enrichment). + + Parameters + ---------- + dataset_name : str, optional + Dataset to process. Defaults to the current MCP client's + agent-scoped dataset, or "main_dataset" if no client identity is + detected. + session_ids : str, optional + Comma-separated session IDs to bridge into the permanent graph. + """ + dataset_name = dataset_name or _agent_scoped_default_dataset() + with redirect_stdout(sys.stderr): + try: + session_list = parse_csv_list(session_ids) + result = await cognee_client.improve( + dataset_name=dataset_name, + session_ids=session_list, + ) + status = result.get("status", "completed") if isinstance(result, dict) else "completed" + if session_list: + text = ( + f"Improve completed (status={status}). " + f"Bridged {len(session_list)} session(s) into permanent graph." + ) + else: + text = f"Graph enrichment completed (status={status})." + return [types.TextContent(type="text", text=text)] + except Exception as e: + error_msg = f"Improve failed: {str(e)}" + logger.error(error_msg) + return [types.TextContent(type="text", text=f"Error: {error_msg}")] + + +# --------------------------------------------------------------------------- +# V1 pipeline status tool +# --------------------------------------------------------------------------- + + +@log_usage(function_name="MCP cognify_status", log_type="mcp_tool") +async def cognify_status( + dataset_name: str = None, + pipelines: List[str] = None, +) -> list: + """ + Get the current status of selected pipelines. + + This function retrieves information about current and recently completed + pipeline operations in the selected dataset. When `dataset_name` is omitted + it defaults to the current MCP client's agent-scoped dataset (e.g. + "cursor_vscode_memory") so each agent sees its own status. + + Returns + ------- + list + A list containing a single TextContent object with the status information as a string. + The status includes information about active and completed jobs for the + requested pipelines. + + Notes + ----- + - By default this checks "cognify_pipeline" (backward compatible) + - Use `pipelines` to restrict to specific pipeline names + - Status information includes job progress, execution time, and completion status + - The status is returned in string format for easy reading + - In API mode the dataset id is resolved over HTTP and status is read + from the server's `GET /api/v1/datasets/status` endpoint + """ + dataset_name = dataset_name or _agent_scoped_default_dataset() + with redirect_stdout(sys.stderr): + try: + if cognee_client.use_api: + # API mode: resolve the dataset id over HTTP (no local cognee + # instance exists in this process) before querying status. + datasets = await cognee_client.list_datasets() + dataset_id = next( + (d["id"] for d in datasets if d.get("name") == dataset_name), None + ) + if dataset_id is None: + return [ + types.TextContent( + type="text", + text=f"❌ Dataset '{dataset_name}' not found via API", + ) + ] + else: + from cognee.modules.data.methods.get_unique_dataset_id import get_unique_dataset_id + from cognee.modules.users.methods import get_default_user + + user = await get_default_user() + dataset_id = await get_unique_dataset_id(dataset_name, user) + + requested_pipelines = list(dict.fromkeys(pipelines or ["cognify_pipeline"])) + + if len(requested_pipelines) == 1: + status = await cognee_client.get_pipeline_status( + [dataset_id], requested_pipelines[0] + ) + else: + status: dict[str, dict] = {str(dataset_id): {}} + for pipeline_name in requested_pipelines: + pipeline_status = await cognee_client.get_pipeline_status( + [dataset_id], pipeline_name + ) + if str(dataset_id) in pipeline_status: + status[str(dataset_id)][pipeline_name] = pipeline_status[str(dataset_id)] + + # Append any background task errors + status_text = str(status) + dataset_errors = _task_errors.get(dataset_name, []) + if dataset_errors: + error_lines = ["\n\nBackground task errors:"] + for ts, err in sorted(dataset_errors, reverse=True): + error_lines.append(f" [{ts}] {err}") + status_text += "\n".join(error_lines) + + return [types.TextContent(type="text", text=status_text)] + except NotImplementedError: + error_msg = "❌ Pipeline status is not available in API mode" + logger.error(error_msg) + return [types.TextContent(type="text", text=error_msg)] + except Exception as e: + error_msg = f"❌ Failed to get cognify status: {str(e)}" + # Still report background errors even if pipeline status fails + dataset_errors = _task_errors.get(dataset_name, []) + if dataset_errors: + error_lines = ["\n\nBackground task errors:"] + for ts, err in sorted(dataset_errors, reverse=True): + error_lines.append(f" [{ts}] {err}") + error_msg += "\n".join(error_lines) + logger.error(error_msg) + return [types.TextContent(type="text", text=error_msg)] + + +# MCP App: interactive graph visualization UI. Rendered by MCP Apps-capable +# hosts (Cursor, Claude Desktop) via the _meta.ui.resourceUri contract. +_VISUALIZE_APP_URI = "ui://cognee-visualize/graph.html" + + +@mcp.resource( + _VISUALIZE_APP_URI, + name="Cognee Graph Visualization UI", + description="Interactive MCP App UI that renders a Cognee knowledge graph.", + mime_type="text/html;profile=mcp-app", +) +def _visualize_graph_ui_resource() -> str: + # The bundle path is resolved as a sibling of this file. In a Docker / + # PyPI install, that's site-packages/src/app_bundles/. In from-source + # dev (running `python src/server.py` directly), it's cognee-mcp/src/ + # app_bundles/. Both resolutions only work because we read via __file__ + # rather than a hardcoded `/app/...` or repo-relative path, so the bundle + # lookup follows wherever this module was loaded from. + bundle = Path(__file__).parent / "app_bundles" / "visualize-graph.html" + if not bundle.is_file(): + raise FileNotFoundError( + f"MCP App bundle not found at {bundle}. " + "Build it with: cd cognee-mcp/apps-src && npm install && npm run build" + ) + return bundle.read_text(encoding="utf-8") + + +# CSS overrides appended to cognee's graph HTML so it fits the MCP App +# iframe better: the floating bottom control bar can wrap to multiple +# rows when the iframe is narrow, and the standalone "Light mode" +# toggle is hidden (the workspace owns theming). +# +# Note: d3 is loaded from a CDN by cognee's HTML, which the MCP App iframe +# blocks via CSP. The workspace bundles d3 from its npm dependency and +# substitutes the CDN