From c7ca4c6782cc8052ff3e6c5b741b35996a569b85 Mon Sep 17 00:00:00 2001
From: Daulet Amirkhanov <damirkhanov01@gmail.com>
Date: Fri, 10 Oct 2025 16:02:58 +0100
Subject: [PATCH] chore: update Docker image comment and improve import
 handling for CogneeClient

---
 cognee-mcp/src/__init__.py | 5 ++++-
 cognee-mcp/src/server.py   | 7 +++++--
 cognee/api/v1/ui/ui.py     | 6 +++---
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/cognee-mcp/src/__init__.py b/cognee-mcp/src/__init__.py
index 933bcd73f..f1c3706c6 100644
--- a/cognee-mcp/src/__init__.py
+++ b/cognee-mcp/src/__init__.py
@@ -1,4 +1,7 @@
-from .server import main as server_main
+try:
+    from .server import main as server_main
+except ImportError:
+    from server import main as server_main
 import warnings
 import sys
 
diff --git a/cognee-mcp/src/server.py b/cognee-mcp/src/server.py
index f98800b18..f2b309fa5 100755
--- a/cognee-mcp/src/server.py
+++ b/cognee-mcp/src/server.py
@@ -18,7 +18,10 @@ from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
 import uvicorn
 
-from .cognee_client import CogneeClient
+try:
+    from .cognee_client import CogneeClient
+except ImportError:
+    from cognee_client import CogneeClient
 
 
 try:
@@ -1088,7 +1091,7 @@ async def main():
     parser.add_argument(
         "--api-token",
         default=None,
-        help="Authentication token for the Cognee API. Required if --api-url is provided.",
+        help="Authentication token for the API (optional, required if API has authentication enabled).",
     )
 
     args = parser.parse_args()
diff --git a/cognee/api/v1/ui/ui.py b/cognee/api/v1/ui/ui.py
index 2a5746f8c..51088c3e1 100644
--- a/cognee/api/v1/ui/ui.py
+++ b/cognee/api/v1/ui/ui.py
@@ -503,7 +503,7 @@ def start_ui(
     if start_mcp:
         logger.info("Starting Cognee MCP server with Docker...")
         try:
-            image = "cognee/cognee-mcp:feature-standalone-mcp"  # TODO: change to main right before merging into main
+            image = "cognee/cognee-mcp:feature-standalone-mcp"  # TODO: change to "cognee/cognee-mcp:main" right before merging into main
             subprocess.run(["docker", "pull", image], check=True)
 
             import uuid
@@ -539,8 +539,8 @@ def start_ui(
                 docker_cmd.extend(["--env-file", env_file])
 
             docker_cmd.append(
-                "cognee/cognee-mcp:feature-standalone-mcp"
-            )  # TODO: change to main right before merging into main
+                image
+            )  # TODO: change to "cognee/cognee-mcp:main" right before merging into main
 
             mcp_process = subprocess.Popen(
                 docker_cmd,