Cognee Walkthrough

class NotebookUtils:
    """Utility class for cognee demo - helper methods to keep the main notebook clean and focused."""
    
    def __init__(self):
        """Initialize the NotebookUtils with default configurations."""
        self.artifacts_dir = None
        self.assets_config = self._initialize_assets_config()
    
    def _initialize_assets_config(self) -> Dict[str, Tuple[str, str]]:
        """Initialize configuration mapping for remote assets to download from cognee repository."""
        return {
            "human_agent_conversations": (
                "/content/copilot_conversations.json",
                "https://raw.githubusercontent.com/topoteretes/cognee/main/notebooks/data/copilot_conversations.json",
            ),
            "python_zen_principles": (
                "/content/zen_principles.md",
                "https://raw.githubusercontent.com/topoteretes/cognee/main/notebooks/data/zen_principles.md",
            ),
            "ontology": (
                "/content/basic_ontology.owl",
                "https://raw.githubusercontent.com/topoteretes/cognee/main/examples/demos/comprehensive_example/data/basic_ontology.owl",
            ),
        }
    
    def download_remote_file_if_not_exists(self, local_path: str, remote_url: str) -> str:
        """Download remote file if it doesn't exist locally to avoid unnecessary re-downloads."""
        file_path = Path(local_path)
        if not file_path.exists():
            file_path.parent.mkdir(parents=True, exist_ok=True)
            urllib.request.urlretrieve(remote_url, file_path)
            print(f"Downloaded: {file_path.name}")
        else:
            print(f"File already exists: {file_path.name}")
        return str(file_path)
    
    def load_json_file_content(self, file_path: str) -> Dict[str, Any]:
        """Load and parse JSON file content into a Python dictionary."""
        with open(file_path, "r", encoding="utf-8") as file:
            return json.load(file)
    
    def load_text_file_content(self, file_path: str) -> str:
        """Load and return raw text content from a text file."""
        with open(file_path, "r", encoding="utf-8") as file:
            return file.read()
    
    def preview_json_structure(self, json_data: Dict[str, Any], max_keys: int = 3) -> None:
        """Display formatted preview of JSON data structure and sample content."""
        print("JSON Structure Preview:")
        pprint.pp(list(json_data.keys())[:max_keys])
        if "conversations" in json_data and json_data["conversations"]:
            print("Sample conversation:")
            pprint.pp(json_data["conversations"][0])
    
    def preview_text_content(self, text_content: str, max_chars: int = 200) -> None:
        """Display formatted preview of text content to show its format."""
        print("Text Content Preview:")
        print(text_content[:max_chars])
        if len(text_content) > max_chars:
            print(f"... (truncated, total length: {len(text_content)} characters)")
    
    def create_notebook_artifacts_directory(self, dir_name: str = "artifacts") -> Path:
        """Create and return artifacts directory for storing notebook outputs like graph visualizations."""
        notebook_dir = Path.cwd()
        self.artifacts_dir = notebook_dir / dir_name
        self.artifacts_dir.mkdir(exist_ok=True)
        print(f"Artifacts directory created/verified at: {self.artifacts_dir}")
        return self.artifacts_dir
    
    def download_remote_assets(self) -> Dict[str, str]:
        """Download all remote assets from cognee repository and return their local file paths."""
        downloaded_assets = {}
        
        print("Downloading remote assets...")
        print("-" * 40)
        
        for asset_name, (local_path, remote_url) in self.assets_config.items():
            downloaded_assets[asset_name] = self.download_remote_file_if_not_exists(
                local_path, remote_url
            )
        
        print("-" * 40)
        print(f"Successfully processed {len(downloaded_assets)} assets")
        return downloaded_assets
    
    def preview_downloaded_assets(self, asset_paths: Dict[str, str]) -> None:
        """Display comprehensive preview of all downloaded assets."""
        print("=== ASSET PREVIEWS ===\n")
        
        # Preview JSON files
        for asset_name, file_path in asset_paths.items():
            if file_path.endswith('.json'):
                print(f"--- {asset_name.upper()} ---")
                json_data = self.load_json_file_content(file_path)
                self.preview_json_structure(json_data)
                print()
        
        # Preview text files
        for asset_name, file_path in asset_paths.items():
            if file_path.endswith(('.md', '.txt')):
                print(f"--- {asset_name.upper()} ---")
                text_content = self.load_text_file_content(file_path)
                self.preview_text_content(text_content)
                print()
        
        # Preview OWL files
        for asset_name, file_path in asset_paths.items():
            if file_path.endswith('.owl'):
                print(f"--- {asset_name.upper()} ---")
                print(f"OWL ontology file: {Path(file_path).name}")
                text_content = self.load_text_file_content(file_path)
                self.preview_text_content(text_content, max_chars=300)
                print()

# Initialize the utility class
utils = NotebookUtils()