feat: add support to start ui programmatically from cognee package
This commit is contained in:
parent
e0b4e4ee5e
commit
14c3d30853
3 changed files with 502 additions and 0 deletions
|
|
@ -27,6 +27,7 @@ from .api.v1.visualize import visualize_graph, start_visualization_server
|
|||
from cognee.modules.visualization.cognee_network_visualization import (
|
||||
cognee_network_visualization,
|
||||
)
|
||||
from .api.v1.ui import start_ui
|
||||
|
||||
# Pipelines
|
||||
from .modules import pipelines
|
||||
|
|
|
|||
1
cognee/api/v1/ui/__init__.py
Normal file
1
cognee/api/v1/ui/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .ui import start_ui, stop_ui, ui
|
||||
500
cognee/api/v1/ui/ui.py
Normal file
500
cognee/api/v1/ui/ui.py
Normal file
|
|
@ -0,0 +1,500 @@
|
|||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import webbrowser
|
||||
import zipfile
|
||||
import requests
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
import tempfile
|
||||
import shutil
|
||||
|
||||
from cognee.shared.logging_utils import get_logger
|
||||
from cognee.version import get_cognee_version
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def get_frontend_cache_dir() -> Path:
|
||||
"""
|
||||
Get the directory where downloaded frontend assets are cached.
|
||||
Uses user's home directory to persist across package updates.
|
||||
"""
|
||||
cache_dir = Path.home() / ".cognee" / "ui-cache"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
return cache_dir
|
||||
|
||||
|
||||
def get_frontend_download_info() -> Tuple[str, str]:
|
||||
"""
|
||||
Get the download URL and version for the actual cognee-frontend source.
|
||||
Downloads the real frontend from GitHub releases, matching the installed version.
|
||||
"""
|
||||
version = get_cognee_version()
|
||||
|
||||
# Clean up version string (remove -local suffix for development)
|
||||
clean_version = version.replace('-local', '')
|
||||
|
||||
# Download from specific release tag to ensure version compatibility
|
||||
download_url = f"https://github.com/topoteretes/cognee/archive/refs/tags/v{clean_version}.zip"
|
||||
|
||||
return download_url, version
|
||||
|
||||
|
||||
def download_frontend_assets(force: bool = False) -> bool:
|
||||
"""
|
||||
Download and cache frontend assets.
|
||||
|
||||
Args:
|
||||
force: If True, re-download even if already cached
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
cache_dir = get_frontend_cache_dir()
|
||||
frontend_dir = cache_dir / "frontend"
|
||||
version_file = cache_dir / "version.txt"
|
||||
|
||||
# Check if already downloaded and up to date
|
||||
if not force and frontend_dir.exists() and version_file.exists():
|
||||
try:
|
||||
cached_version = version_file.read_text().strip()
|
||||
current_version = get_cognee_version()
|
||||
if cached_version == current_version:
|
||||
logger.debug(f"Frontend assets already cached for version {current_version}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.debug(f"Error checking cached version: {e}")
|
||||
|
||||
download_url, version = get_frontend_download_info()
|
||||
|
||||
logger.info("Downloading cognee frontend assets...")
|
||||
logger.info("This is a one-time download and will be cached for future use.")
|
||||
|
||||
try:
|
||||
# Create a temporary directory for download
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
archive_path = temp_path / "cognee-main.zip"
|
||||
|
||||
# Download the actual cognee repository from releases
|
||||
logger.info(f"Downloading cognee v{version.replace('-local', '')} from GitHub releases...")
|
||||
logger.info(f"URL: {download_url}")
|
||||
response = requests.get(download_url, stream=True, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
with open(archive_path, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
# Extract the archive and find the cognee-frontend directory
|
||||
if frontend_dir.exists():
|
||||
shutil.rmtree(frontend_dir)
|
||||
|
||||
with zipfile.ZipFile(archive_path, 'r') as zip_file:
|
||||
# Extract to temp directory first
|
||||
extract_dir = temp_path / "extracted"
|
||||
zip_file.extractall(extract_dir)
|
||||
|
||||
# Find the cognee-frontend directory in the extracted content
|
||||
# The archive structure will be: cognee-{version}/cognee-frontend/
|
||||
cognee_frontend_source = None
|
||||
for root, dirs, files in os.walk(extract_dir):
|
||||
if "cognee-frontend" in dirs:
|
||||
cognee_frontend_source = Path(root) / "cognee-frontend"
|
||||
break
|
||||
|
||||
if not cognee_frontend_source or not cognee_frontend_source.exists():
|
||||
logger.error("Could not find cognee-frontend directory in downloaded release archive")
|
||||
logger.error("This might indicate a version mismatch or missing release.")
|
||||
return False
|
||||
|
||||
# Copy the cognee-frontend to our cache
|
||||
shutil.copytree(cognee_frontend_source, frontend_dir)
|
||||
logger.debug(f"Frontend extracted to: {frontend_dir}")
|
||||
|
||||
# Write version info
|
||||
version_file.write_text(version)
|
||||
|
||||
logger.info(f"✓ Cognee frontend v{version.replace('-local', '')} downloaded and cached successfully!")
|
||||
return True
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
if "404" in str(e):
|
||||
logger.error(f"Release v{version.replace('-local', '')} not found on GitHub.")
|
||||
logger.error("This version might not have been released yet, or you're using a development version.")
|
||||
logger.error("Try using a stable release version of cognee.")
|
||||
else:
|
||||
logger.error(f"Failed to download from GitHub: {str(e)}")
|
||||
logger.error("You can still use cognee without the UI functionality.")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to download frontend assets: {str(e)}")
|
||||
logger.error("You can still use cognee without the UI functionality.")
|
||||
return False
|
||||
|
||||
|
||||
|
||||
|
||||
def find_frontend_path() -> Optional[Path]:
|
||||
"""
|
||||
Find the cognee-frontend directory.
|
||||
Checks both development location and cached download location.
|
||||
"""
|
||||
current_file = Path(__file__)
|
||||
|
||||
# First, try development paths (for contributors/developers)
|
||||
dev_search_paths = [
|
||||
current_file.parents[4] / "cognee-frontend", # from cognee/api/v1/ui/ui.py to project root
|
||||
current_file.parents[3] / "cognee-frontend", # fallback path
|
||||
current_file.parents[2] / "cognee-frontend", # another fallback
|
||||
]
|
||||
|
||||
for path in dev_search_paths:
|
||||
if path.exists() and (path / "package.json").exists():
|
||||
logger.debug(f"Found development frontend at: {path}")
|
||||
return path
|
||||
|
||||
# Then try cached download location (for pip-installed users)
|
||||
cache_dir = get_frontend_cache_dir()
|
||||
cached_frontend = cache_dir / "frontend"
|
||||
|
||||
if cached_frontend.exists() and (cached_frontend / "package.json").exists():
|
||||
logger.debug(f"Found cached frontend at: {cached_frontend}")
|
||||
return cached_frontend
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def check_node_npm() -> tuple[bool, str]:
|
||||
"""
|
||||
Check if Node.js and npm are available.
|
||||
Returns (is_available, error_message)
|
||||
"""
|
||||
try:
|
||||
# Check Node.js
|
||||
result = subprocess.run(["node", "--version"],
|
||||
capture_output=True, text=True, timeout=10)
|
||||
if result.returncode != 0:
|
||||
return False, "Node.js is not installed or not in PATH"
|
||||
|
||||
node_version = result.stdout.strip()
|
||||
logger.debug(f"Found Node.js version: {node_version}")
|
||||
|
||||
# Check npm
|
||||
result = subprocess.run(["npm", "--version"],
|
||||
capture_output=True, text=True, timeout=10)
|
||||
if result.returncode != 0:
|
||||
return False, "npm is not installed or not in PATH"
|
||||
|
||||
npm_version = result.stdout.strip()
|
||||
logger.debug(f"Found npm version: {npm_version}")
|
||||
|
||||
return True, f"Node.js {node_version}, npm {npm_version}"
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return False, "Timeout checking Node.js/npm installation"
|
||||
except FileNotFoundError:
|
||||
return False, "Node.js/npm not found. Please install Node.js from https://nodejs.org/"
|
||||
except Exception as e:
|
||||
return False, f"Error checking Node.js/npm: {str(e)}"
|
||||
|
||||
|
||||
def install_frontend_dependencies(frontend_path: Path) -> bool:
|
||||
"""
|
||||
Install frontend dependencies if node_modules doesn't exist.
|
||||
This is needed for both development and downloaded frontends since both use npm run dev.
|
||||
"""
|
||||
node_modules = frontend_path / "node_modules"
|
||||
if node_modules.exists():
|
||||
logger.debug("Frontend dependencies already installed")
|
||||
return True
|
||||
|
||||
logger.info("Installing frontend dependencies (this may take a few minutes)...")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["npm", "install"],
|
||||
cwd=frontend_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300 # 5 minutes timeout
|
||||
)
|
||||
|
||||
if result.returncode == 0:
|
||||
logger.info("Frontend dependencies installed successfully")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Failed to install dependencies: {result.stderr}")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error("Timeout installing frontend dependencies")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Error installing frontend dependencies: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def is_development_frontend(frontend_path: Path) -> bool:
|
||||
"""
|
||||
Check if this is a development frontend (has Next.js) vs downloaded assets.
|
||||
"""
|
||||
package_json_path = frontend_path / "package.json"
|
||||
if not package_json_path.exists():
|
||||
return False
|
||||
|
||||
try:
|
||||
import json
|
||||
with open(package_json_path) as f:
|
||||
package_data = json.load(f)
|
||||
|
||||
# Development frontend has Next.js as dependency
|
||||
dependencies = package_data.get("dependencies", {})
|
||||
dev_dependencies = package_data.get("devDependencies", {})
|
||||
|
||||
return "next" in dependencies or "next" in dev_dependencies
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def start_python_server(frontend_path: Path, host: str, port: int) -> Optional[subprocess.Popen]:
|
||||
"""
|
||||
Start a Python HTTP server to serve downloaded frontend assets.
|
||||
This doesn't require Node.js and works with pre-built assets.
|
||||
"""
|
||||
try:
|
||||
# Change to the frontend directory
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(frontend_path)
|
||||
|
||||
# Use subprocess to run the server so we can return a process handle
|
||||
cmd = [
|
||||
"python", "-m", "http.server", str(port),
|
||||
"--bind", host
|
||||
]
|
||||
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
cwd=frontend_path,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
)
|
||||
|
||||
# Restore original directory
|
||||
os.chdir(original_cwd)
|
||||
|
||||
# Give it a moment to start
|
||||
time.sleep(2)
|
||||
|
||||
# Check if process is still running
|
||||
if process.poll() is not None:
|
||||
stdout, stderr = process.communicate()
|
||||
logger.error("Python HTTP server failed to start:")
|
||||
logger.error(f"stdout: {stdout}")
|
||||
logger.error(f"stderr: {stderr}")
|
||||
return None
|
||||
|
||||
return process
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start Python HTTP server: {str(e)}")
|
||||
# Restore original directory on error
|
||||
try:
|
||||
os.chdir(original_cwd)
|
||||
except OSError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def prompt_user_for_download() -> bool:
|
||||
"""
|
||||
Ask user if they want to download the frontend assets.
|
||||
Returns True if user consents, False otherwise.
|
||||
"""
|
||||
try:
|
||||
print("\n" + "="*60)
|
||||
print("🎨 Cognee UI Setup Required")
|
||||
print("="*60)
|
||||
print("The cognee frontend is not available on your system.")
|
||||
print("This is required to use the web interface.")
|
||||
print("\nWhat will happen:")
|
||||
print("• Download the actual cognee-frontend from GitHub")
|
||||
print("• Cache it in your home directory (~/.cognee/ui-cache/)")
|
||||
print("• Install dependencies with npm (requires Node.js)")
|
||||
print("• This is a one-time setup per cognee version")
|
||||
print("\nThe frontend will then be available offline for future use.")
|
||||
|
||||
response = input("\nWould you like to download the frontend now? (y/N): ").strip().lower()
|
||||
return response in ['y', 'yes']
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nOperation cancelled by user.")
|
||||
return False
|
||||
|
||||
|
||||
def start_ui(host: str = "localhost", port: int = 3000, open_browser: bool = True, auto_download: bool = False) -> Optional[subprocess.Popen]:
|
||||
"""
|
||||
Start the cognee frontend UI server.
|
||||
|
||||
This function will:
|
||||
1. Find the cognee-frontend directory (development) or download it (pip install)
|
||||
2. Check if Node.js and npm are available (for development mode)
|
||||
3. Install dependencies if needed (development mode)
|
||||
4. Start the appropriate server
|
||||
5. Optionally open the browser
|
||||
|
||||
Args:
|
||||
host: Host to bind the server to (default: localhost)
|
||||
port: Port to run the server on (default: 3000)
|
||||
open_browser: Whether to open the browser automatically (default: True)
|
||||
auto_download: If True, download frontend without prompting (default: False)
|
||||
|
||||
Returns:
|
||||
subprocess.Popen object representing the running server, or None if failed
|
||||
|
||||
Example:
|
||||
>>> import cognee
|
||||
>>> server = cognee.start_ui()
|
||||
>>> # UI will be available at http://localhost:3000
|
||||
>>> # To stop the server later:
|
||||
>>> server.terminate()
|
||||
"""
|
||||
logger.info("Starting cognee UI...")
|
||||
|
||||
# Find frontend directory
|
||||
frontend_path = find_frontend_path()
|
||||
|
||||
if not frontend_path:
|
||||
logger.info("Frontend not found locally. This is normal for pip-installed cognee.")
|
||||
|
||||
# Offer to download the frontend
|
||||
if auto_download or prompt_user_for_download():
|
||||
if download_frontend_assets():
|
||||
frontend_path = find_frontend_path()
|
||||
if not frontend_path:
|
||||
logger.error("Download succeeded but frontend still not found. This is unexpected.")
|
||||
return None
|
||||
else:
|
||||
logger.error("Failed to download frontend assets.")
|
||||
return None
|
||||
else:
|
||||
logger.info("Frontend download declined. UI functionality not available.")
|
||||
logger.info("You can still use all other cognee features without the web interface.")
|
||||
return None
|
||||
|
||||
# Check Node.js and npm
|
||||
node_available, node_message = check_node_npm()
|
||||
if not node_available:
|
||||
logger.error(f"Cannot start UI: {node_message}")
|
||||
logger.error("Please install Node.js from https://nodejs.org/ to use the UI functionality")
|
||||
return None
|
||||
|
||||
logger.debug(f"Environment check passed: {node_message}")
|
||||
|
||||
# Install dependencies if needed
|
||||
if not install_frontend_dependencies(frontend_path):
|
||||
logger.error("Failed to install frontend dependencies")
|
||||
return None
|
||||
|
||||
# Prepare environment variables
|
||||
env = os.environ.copy()
|
||||
env['HOST'] = host
|
||||
env['PORT'] = str(port)
|
||||
|
||||
# Start the development server
|
||||
logger.info(f"Starting frontend server at http://{host}:{port}")
|
||||
logger.info("This may take a moment to compile and start...")
|
||||
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
["npm", "run", "dev"],
|
||||
cwd=frontend_path,
|
||||
env=env,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
)
|
||||
|
||||
# Give it a moment to start up
|
||||
time.sleep(3)
|
||||
|
||||
# Check if process is still running
|
||||
if process.poll() is not None:
|
||||
stdout, stderr = process.communicate()
|
||||
logger.error("Frontend server failed to start:")
|
||||
logger.error(f"stdout: {stdout}")
|
||||
logger.error(f"stderr: {stderr}")
|
||||
return None
|
||||
|
||||
# Open browser if requested
|
||||
if open_browser:
|
||||
def open_browser_delayed():
|
||||
time.sleep(5) # Give Next.js time to fully start
|
||||
try:
|
||||
webbrowser.open(f"http://{host}:{port}") # TODO: use dashboard url?
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not open browser automatically: {e}")
|
||||
|
||||
browser_thread = threading.Thread(target=open_browser_delayed, daemon=True)
|
||||
browser_thread.start()
|
||||
|
||||
logger.info("✓ Cognee UI is starting up...")
|
||||
logger.info(f"✓ Open your browser to: http://{host}:{port}")
|
||||
logger.info("✓ The UI will be available once Next.js finishes compiling")
|
||||
|
||||
return process
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start frontend server: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def stop_ui(process: subprocess.Popen) -> bool:
|
||||
"""
|
||||
Stop a running UI server process.
|
||||
|
||||
Args:
|
||||
process: The subprocess.Popen object returned by start_ui()
|
||||
|
||||
Returns:
|
||||
bool: True if stopped successfully, False otherwise
|
||||
"""
|
||||
if not process:
|
||||
return False
|
||||
|
||||
try:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("Process didn't terminate gracefully, forcing kill")
|
||||
process.kill()
|
||||
process.wait()
|
||||
|
||||
logger.info("UI server stopped")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping UI server: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
# Convenience function similar to DuckDB's approach
|
||||
def ui() -> Optional[subprocess.Popen]:
|
||||
"""
|
||||
Convenient alias for start_ui() with default parameters.
|
||||
Similar to how DuckDB provides simple ui() function.
|
||||
"""
|
||||
return start_ui()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Test the UI startup
|
||||
server = start_ui()
|
||||
if server:
|
||||
try:
|
||||
input("Press Enter to stop the server...")
|
||||
finally:
|
||||
stop_ui(server)
|
||||
Loading…
Add table
Reference in a new issue