topoteretes--cognee/cognee/cli/commands/cognify_command.py

import argparse
import asyncio
import os
from typing import Optional

from cognee.cli.reference import SupportsCliCommand
from cognee.cli import DEFAULT_DOCS_URL
from cognee.cli.config import CHUNKER_CHOICES
import cognee.cli.echo as fmt
from cognee.cli.exceptions import CliCommandException, CliCommandInnerException


class CognifyCommand(SupportsCliCommand):
    command_string = "cognify"
    help_string = "Transform ingested data into a structured knowledge graph"
    docs_url = DEFAULT_DOCS_URL
    description = """
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.

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 text summaries for navigation

After successful cognify processing, use `cognee search` to query the knowledge graph.
    """

    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "--datasets",
            "-d",
            nargs="*",
            help="Dataset name(s) to process. Processes all available data if not specified. Can be multiple: --datasets dataset1 dataset2",
        )
        parser.add_argument(
            "--chunk-size",
            type=int,
            help="Maximum tokens per chunk. Auto-calculated based on LLM if not specified (~512-8192 tokens)",
        )
        parser.add_argument(
            "--ontology-file",
            help="Path to RDF/OWL ontology file for domain-specific entity types. "
            "Multiple files can be given as a comma-separated list.",
        )
        parser.add_argument(
            "--chunker",
            choices=CHUNKER_CHOICES,
            default="TextChunker",
            help="Text chunking strategy (default: TextChunker)",
        )
        parser.add_argument(
            "--background",
            "-b",
            action="store_true",
            help="Run processing in background and return immediately (recommended for large datasets)",
        )
        parser.add_argument(
            "--verbose", "-v", action="store_true", help="Show detailed progress information"
        )
        parser.add_argument(
            "--chunks-per-batch",
            type=int,
            help="Number of chunks to process per task batch (try 50 for large single documents).",
        )

    def execute(self, args: argparse.Namespace) -> None:
        try:
            # Import cognee here to avoid circular imports
            import cognee

            if args.ontology_file:
                missing = [
                    path.strip()
                    for path in args.ontology_file.split(",")
                    if not os.path.isfile(path.strip())
                ]
                if missing:
                    raise CliCommandInnerException(f"Ontology file not found: {', '.join(missing)}")

            # Prepare datasets parameter
            datasets = args.datasets if args.datasets else None
            dataset_msg = f" for datasets {datasets}" if datasets else " for all available data"
            fmt.echo(f"Starting cognification{dataset_msg}...")

            if args.verbose:
                fmt.note("This process will analyze your data and build knowledge graphs.")
                fmt.note("Depending on data size, this may take several minutes.")
                if args.background:
                    fmt.note(
                        "Running in background mode - the process will continue after this command exits."
                    )

            # Prepare chunker parameter - will be handled in the async function

            # Run the async cognify function
            async def run_cognify():
                try:
                    from cognee.cli.user_resolution import resolve_cli_user

                    user = await resolve_cli_user(getattr(args, "user_id", None))

                    # Import chunker classes here
                    from cognee.modules.chunking.TextChunker import TextChunker

                    chunker_class = TextChunker  # Default
                    if args.chunker == "LangchainChunker":
                        try:
                            from cognee.modules.chunking.LangchainChunker import LangchainChunker

                            chunker_class = LangchainChunker
                        except ImportError:
                            fmt.warning("LangchainChunker not available, using TextChunker")
                    elif args.chunker == "CsvChunker":
                        try:
                            from cognee.modules.chunking.CsvChunker import CsvChunker

                            chunker_class = CsvChunker
                        except ImportError:
                            fmt.warning("CsvChunker not available, using TextChunker")

                    # Translate --ontology-file into the canonical ontology Config
                    # structure, built with the same factory cognify() uses for
                    # its env-based fallback.
                    config = None
                    if args.ontology_file:
                        from cognee.modules.ontology.get_default_ontology_resolver import (
                            get_ontology_resolver_from_env,
                        )

                        config = {
                            "ontology_config": {
                                "ontology_resolver": get_ontology_resolver_from_env(
                                    ontology_resolver="rdflib",
                                    matching_strategy="fuzzy",
                                    ontology_file_path=args.ontology_file,
                                )
                            }
                        }

                    result = await cognee.cognify(
                        datasets=datasets,
                        user=user,
                        chunker=chunker_class,
                        chunk_size=args.chunk_size,
                        config=config,
                        run_in_background=args.background,
                        chunks_per_batch=getattr(args, "chunks_per_batch", None),
                    )
                    return result
                except Exception as e:
                    raise CliCommandInnerException(f"Failed to cognify: {str(e)}") from e

            result = asyncio.run(run_cognify())

            if args.background:
                fmt.success("Cognification started in background!")
                if args.verbose and result:
                    fmt.echo(
                        "Background processing initiated. Use pipeline monitoring to track progress."
                    )
            else:
                fmt.success("Cognification completed successfully!")
                if args.verbose and result:
                    fmt.echo(f"Processing results: {result}")

        except Exception as e:
            if isinstance(e, CliCommandInnerException):
                raise CliCommandException(str(e), error_code=1) from e
            raise CliCommandException(f"Error during cognification: {str(e)}", error_code=1) from e