codegraphcontext 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

codegraphcontext/cli/cli_helpers.py

@@ -2,10 +2,12 @@ import asyncio
 import json
 import uuid
 import urllib.parse
+from collections import Counter
 from pathlib import Path
 import time
 import os
 from typing import Optional, List, Dict, Any
+import typer
 from rich.console import Console
 from rich.table import Table
 from rich.progress import (
@@ -31,6 +33,31 @@ from .config_manager import resolve_context, ResolvedContext, register_repo_in_c
 console = Console()
 
 
+def _print_call_resolution_diagnostics(graph_builder: GraphBuilder, limit: int = 5) -> None:
+    diagnostics = getattr(graph_builder, "last_call_resolution_diagnostics", [])
+    if not diagnostics:
+        return
+
+    reason_counts = Counter(d.get("reason", "unknown") for d in diagnostics)
+    summary = ", ".join(
+        f"{reason}={count}" for reason, count in reason_counts.most_common()
+    )
+    console.print(
+        f"[yellow]Skipped {len(diagnostics)} unresolved call relationship(s): {summary}[/yellow]"
+    )
+    table = Table(show_header=True, header_style="bold magenta")
+    table.add_column("Call", style="cyan", overflow="fold")
+    table.add_column("Reason", style="yellow")
+    table.add_column("Location", style="dim", overflow="fold")
+    for diagnostic in diagnostics[:limit]:
+        table.add_row(
+            str(diagnostic.get("full_call_name") or ""),
+            str(diagnostic.get("reason") or ""),
+            f"{diagnostic.get('caller_file_path')}:{diagnostic.get('line_number')}",
+        )
+    console.print(table)
+
+
 def _initialize_services(cli_context_flag: Optional[str] = None) -> tuple[Any, Any, Any, ResolvedContext]:
     """
     Initializes and returns core service managers based on the resolved context.
@@ -243,6 +270,7 @@ def index_helper(path: str, context: Optional[str] = None):
         asyncio.run(_run_index_with_progress(graph_builder, path_obj, is_dependency=False, cgcignore_path=ctx.cgcignore_path))
         time_end = time.time()
         elapsed = time_end - time_start
+        _print_call_resolution_diagnostics(graph_builder)
         console.print(f"[green]Successfully finished indexing: {path} in {elapsed:.2f} seconds[/green]")
         
         # Check if auto-watch is enabled
@@ -259,6 +287,7 @@ def index_helper(path: str, context: Optional[str] = None):
             
     except Exception as e:
         console.print(f"[bold red]An error occurred during indexing:[/bold red] {e}")
+        raise typer.Exit(code=1)
     finally:
         db_manager.close_driver()
 
@@ -289,6 +318,7 @@ def add_package_helper(package_name: str, language: str, context: Optional[str]
 
     try:
         asyncio.run(_run_index_with_progress(graph_builder, package_path, is_dependency=True, cgcignore_path=ctx.cgcignore_path))
+        _print_call_resolution_diagnostics(graph_builder)
         console.print(f"[green]Successfully finished indexing package: {package_name}[/green]")
     except Exception as e:
         console.print(f"[bold red]An error occurred during package indexing:[/bold red] {e}")
@@ -354,9 +384,11 @@ def cypher_helper(query: str, context: Optional[str] = None):
 
     db_manager, _, _, ctx = services
     
-    # Replicating safety checks from MCPServer
+    # Replicating safety checks from MCPServer (using word boundaries to avoid false positives like 'createEmail')
+    import re
     forbidden_keywords = ['CREATE', 'MERGE', 'DELETE', 'SET', 'REMOVE', 'DROP', 'CALL apoc']
-    if any(keyword in query.upper() for keyword in forbidden_keywords):
+    pattern = r'\b(' + '|'.join(forbidden_keywords) + r')\b'
+    if re.search(pattern, query, re.IGNORECASE):
         console.print("[bold red]Error: This command only supports read-only queries.[/bold red]")
         db_manager.close_driver()
         return
@@ -382,9 +414,11 @@ def cypher_helper_visual(query: str, context: Optional[str] = None):
 
     db_manager, _, _, ctx = services
     
-    # Replicating safety checks from MCPServer
+    # Replicating safety checks from MCPServer (using word boundaries to avoid false positives like 'createEmail')
+    import re
     forbidden_keywords = ['CREATE', 'MERGE', 'DELETE', 'SET', 'REMOVE', 'DROP', 'CALL apoc']
-    if any(keyword in query.upper() for keyword in forbidden_keywords):
+    pattern = r'\b(' + '|'.join(forbidden_keywords) + r')\b'
+    if re.search(pattern, query, re.IGNORECASE):
         console.print("[bold red]Error: This command only supports read-only queries.[/bold red]")
         db_manager.close_driver()
         return
@@ -540,9 +574,11 @@ def reindex_helper(path: str, context: Optional[str] = None):
         asyncio.run(_run_index_with_progress(graph_builder, path_obj, is_dependency=False, cgcignore_path=ctx.cgcignore_path))
         time_end = time.time()
         elapsed = time_end - time_start
+        _print_call_resolution_diagnostics(graph_builder)
         console.print(f"[green]Successfully re-indexed: {path} in {elapsed:.2f} seconds[/green]")
     except Exception as e:
         console.print(f"[bold red]An error occurred during re-indexing:[/bold red] {e}")
+        raise typer.Exit(code=1)
     finally:
         db_manager.close_driver()
 
@@ -669,6 +705,12 @@ def stats_helper(path: str = None, context: Optional[str] = None):
                     class_count = session.run("MATCH (c:Class) RETURN count(c) as c").single()["c"]
                     module_count = session.run("MATCH (m:Module) RETURN count(m) as c").single()["c"]
                     
+                    # Extended node types (PHP, Rust, Go, etc.)
+                    interface_count = session.run("MATCH (i:Interface) RETURN count(i) as c").single()["c"]
+                    trait_count = session.run("MATCH (t:Trait) RETURN count(t) as c").single()["c"]
+                    struct_count = session.run("MATCH (s:Struct) RETURN count(s) as c").single()["c"]
+                    enum_count = session.run("MATCH (e:Enum) RETURN count(e) as c").single()["c"]
+                    
                     table = Table(show_header=True, header_style="bold magenta")
                     table.add_column("Metric", style="cyan")
                     table.add_column("Count", style="green", justify="right")
@@ -677,6 +719,14 @@ def stats_helper(path: str = None, context: Optional[str] = None):
                     table.add_row("Files", str(file_count))
                     table.add_row("Functions", str(func_count))
                     table.add_row("Classes", str(class_count))
+                    if interface_count > 0:
+                        table.add_row("Interfaces", str(interface_count))
+                    if trait_count > 0:
+                        table.add_row("Traits", str(trait_count))
+                    if struct_count > 0:
+                        table.add_row("Structs", str(struct_count))
+                    if enum_count > 0:
+                        table.add_row("Enums", str(enum_count))
                     table.add_row("Modules", str(module_count))
                     
                     console.print(table)

codegraphcontext/cli/config_manager.py

@@ -48,7 +48,7 @@ DEFAULT_CONFIG = {
     "INDEX_SOURCE": "true",
     # SCIP indexer feature flag (default off — existing Tree-sitter behaviour unchanged)
     "SCIP_INDEXER": "false",
-    "SCIP_LANGUAGES": "python,typescript,go,rust,java",
+    "SCIP_LANGUAGES": "python,typescript,javascript,go,rust,java,dart,cpp,c,csharp",
     "SKIP_EXTERNAL_RESOLUTION": "false",
     # 0 = unlimited; any positive integer caps MCP tool response size.
     "MAX_TOOL_RESPONSE_TOKENS": "0",
@@ -85,7 +85,7 @@ CONFIG_DESCRIPTIONS = {
     "IGNORE_DIRS": "Comma-separated list of directory names to ignore during indexing",
     "INDEX_SOURCE": "Store full source code in graph database (for faster indexing use false, for better performance use true)",
     "SCIP_INDEXER": "Use SCIP-based indexing for higher accuracy call/inheritance resolution (requires scip-<lang> tools installed)",
-    "SCIP_LANGUAGES": "Comma-separated languages to index via SCIP when SCIP_INDEXER=true (python,typescript,go,rust,java)",
+    "SCIP_LANGUAGES": "Comma-separated languages to index via SCIP when SCIP_INDEXER=true (python,typescript,javascript,go,rust,java,dart,cpp,c,csharp)",
     "SKIP_EXTERNAL_RESOLUTION": "Skip resolution attempts for external library method calls (recommended for enterprise large Java/Spring codebases)",
     "MAX_TOOL_RESPONSE_TOKENS": "Maximum tokens per MCP tool response (0 = unlimited). Truncates oversized payloads and appends a notice.",
     "TOOL_RESULT_LIMITS": "JSON object mapping tool names to max result counts, e.g. {\"find_code\": 20, \"analyze_code_relationships\": 10}. Missing keys use built-in defaults.",
@@ -859,7 +859,7 @@ def resolve_context(
             )
 
     # --- 4. Global fallback ---
-    db = load_config().get("DEFAULT_DATABASE", "falkordb")
+    db = os.getenv("CGC_RUNTIME_DB_TYPE") or load_config().get("DEFAULT_DATABASE", "falkordb")
     return ResolvedContext(
         mode="global",
         context_name="",

codegraphcontext/cli/main.py

@@ -300,7 +300,7 @@ def _load_credentials():
     Step 2 skips duplicate loading when that file is the same path as the global file.
     Arbitrary repo-root `.env` files are not loaded—only CodeGraphContext config paths.
     """
-    from dotenv import dotenv_values
+    from dotenv import dotenv_values, find_dotenv
     from codegraphcontext.cli.config_manager import (
         ensure_config_dir,
         codegraphcontext_dotenv_at_cwd,
@@ -586,7 +586,7 @@ def bundle_export(
     services = _initialize_services(context)
     if not all(services[:3]):
         return
-    db_manager, graph_builder, code_finder = services[:3]
+    db_manager, _, code_finder = services[:3]
     
     try:
         output_path = Path(output)
@@ -753,9 +753,19 @@ def load_shortcut(
 # REGISTRY COMMAND GROUP - Browse and Download Bundles
 # ============================================================================
 
-registry_app = typer.Typer(help="Browse and download bundles from the registry")
+registry_app = typer.Typer(
+    help="Browse and download bundles from the registry",
+    invoke_without_command=True,
+)
 app.add_typer(registry_app, name="registry")
 
+
+@registry_app.callback()
+def registry_callback(ctx: typer.Context):
+    """Browse and download bundles from the registry."""
+    if ctx.invoked_subcommand is None:
+        typer.echo(ctx.get_help())
+
 @registry_app.command("list")
 def registry_list(
     verbose: bool = typer.Option(False, "--verbose", "-v", help="Show detailed information including download URLs"),
@@ -940,11 +950,11 @@ def doctor():
         elif default_db == "kuzudb":
             from importlib.util import find_spec
 
-            if find_spec("real_ladybug") is not None:
+            if find_spec("kuzu") is not None:
                 console.print(f"   [green]✓[/green] KuzuDB is installed")
             else:
                 console.print(f"   [red]✗[/red] KuzuDB is not installed")
-                console.print(f"       Run: pip install real_ladybug")
+                console.print(f"       Run: pip install kuzu")
                 all_checks_passed = False
         else:
             # FalkorDB
@@ -1199,7 +1209,7 @@ def report(
     """
     _load_credentials()
     output_path = Path(output) if output else Path.cwd() / "CGC_REPORT.md"
-    db_manager, _, _ = _initialize_services(context)
+    db_manager, _, _, _ = _initialize_services(context)
     try:
         from codegraphcontext.tools.report_generator import generate_report
         report_text = generate_report(db_manager, output_path=output_path, include_java=java)
@@ -1374,6 +1384,18 @@ def find_by_name(
             results.extend(variables)
             results.extend(modules)
             results.extend(imports)
+
+            # Also search Interface, Trait, Struct, Enum nodes (PHP, Rust, Go, etc.)
+            with db_manager.get_driver().session() as session:
+                for label in ['Interface', 'Trait', 'Struct', 'Enum']:
+                    res = session.run(
+                        f"MATCH (n:{label}) WHERE n.name = $name RETURN n.name as name, n.path as path, n.line_number as line_number",
+                        name=name
+                    )
+                    for record in res:
+                        row = dict(record)
+                        row['type'] = label
+                        results.append(row)
         
         elif type.lower() == 'function':
             results = code_finder.find_by_function_name(name, fuzzy_search=False)
@@ -1460,7 +1482,7 @@ def find_by_pattern(
             if not case_sensitive:
                 query = """
                     MATCH (n)
-                    WHERE (n:Function OR n:Class OR n:Module OR n:Variable) AND toLower(n.name) CONTAINS toLower($pattern)
+                    WHERE (n:Function OR n:Class OR n:Module OR n:Variable OR n:Interface OR n:Trait OR n:Struct OR n:Enum) AND toLower(n.name) CONTAINS toLower($pattern)
                     RETURN 
                         labels(n)[0] as type,
                         n.name as name,
@@ -1473,7 +1495,7 @@ def find_by_pattern(
             else:
                  query = """
                     MATCH (n)
-                    WHERE (n:Function OR n:Class OR n:Module OR n:Variable) AND n.name CONTAINS $pattern
+                    WHERE (n:Function OR n:Class OR n:Module OR n:Variable OR n:Interface OR n:Trait OR n:Struct OR n:Enum) AND n.name CONTAINS $pattern
                     RETURN 
                         labels(n)[0] as type,
                         n.name as name,
@@ -1978,6 +2000,65 @@ def analyze_chain(
     finally:
         db_manager.close_driver()
 
+@analyze_app.command("kotlin-call-audit")
+def analyze_kotlin_call_audit(
+    repo_path: Optional[str] = typer.Option(None, "--repo-path", "-r", help="Limit audit to paths under this repository root"),
+    limit: int = typer.Option(20, "--limit", "-n", help="Maximum examples/top names to show"),
+    json_output: bool = typer.Option(False, "--json", help="Print machine-readable JSON"),
+    fail_on_ambiguity: bool = typer.Option(False, "--fail-on-ambiguity", help="Exit non-zero if any ambiguous Kotlin call groups are found"),
+    context: Optional[str] = typer.Option(None, "--context", "-c", help="Specific context to use"),
+):
+    """
+    Audit Kotlin function call edges for multi-target callsite ambiguity.
+
+    Example:
+        cgc analyze kotlin-call-audit --context elrond-stable --fail-on-ambiguity
+        cgc analyze kotlin-call-audit --json
+    """
+    _load_credentials()
+    services = _initialize_services(context)
+    if not all(services[:3]):
+        return
+    db_manager, _, code_finder = services[:3]
+
+    try:
+        result = code_finder.audit_kotlin_call_ambiguity(repo_path=repo_path, limit=limit)
+        if json_output:
+            console.print_json(json.dumps(result))
+        else:
+            console.print("\n[bold cyan]Kotlin CALLS ambiguity audit[/bold cyan]")
+            summary = Table(show_header=True, header_style="bold magenta", box=box.ROUNDED)
+            summary.add_column("Metric", style="cyan")
+            summary.add_column("Value", style="green")
+            summary.add_row("Kotlin fn→fn CALLS edges", str(result["kotlin_fn_to_fn_edges"]))
+            summary.add_row("Ambiguous groups", str(result["ambiguous_groups"]))
+            summary.add_row("Ambiguous edges", str(result["ambiguous_edges"]))
+            console.print(summary)
+
+            if result["examples"]:
+                examples = Table(show_header=True, header_style="bold magenta", box=box.ROUNDED)
+                examples.add_column("Callsite", style="cyan", overflow="fold")
+                examples.add_column("Call", style="yellow", overflow="fold")
+                examples.add_column("Targets", style="green", overflow="fold")
+                for example in result["examples"]:
+                    targets = "\n".join(
+                        f"{target['context'] or ''}:{target['line_number']} {target['path']}"
+                        for target in example["targets"]
+                    )
+                    examples.add_row(
+                        f"{example.get('caller_name')} {example.get('caller_path')}:{example.get('call_line')}",
+                        str(example.get("full_call_name") or ""),
+                        targets,
+                    )
+                console.print(examples)
+            else:
+                console.print("[green]No ambiguous Kotlin call groups found.[/green]")
+
+        if fail_on_ambiguity and result["ambiguous_groups"]:
+            raise typer.Exit(1)
+    finally:
+        db_manager.close_driver()
+
 @analyze_app.command("deps")
 def analyze_dependencies(
     ctx: typer.Context,
@@ -2654,4 +2735,4 @@ def _write_datasource_graph(ingested: dict) -> None:
 
 
 if __name__ == "__main__":
-    app()
+    app()

codegraphcontext/core/__init__.py

@@ -22,7 +22,7 @@ import importlib.util
 def _is_kuzudb_available() -> bool:
     """Check if KùzuDB is installed."""
     try:
-        return importlib.util.find_spec("real_ladybug") is not None
+        return importlib.util.find_spec("kuzu") is not None
     except ImportError:
         return False
 
@@ -78,7 +78,7 @@ def get_database_manager(db_path: Optional[str] = None) -> Union['DatabaseManage
         db_type = db_type.lower()
         if db_type == 'kuzudb':
             if not _is_kuzudb_available():
-                raise ValueError("Database set to 'kuzudb' but Kùzu is not installed.\nRun 'pip install real_ladybug'")
+                raise ValueError("Database set to 'kuzudb' but Kùzu is not installed.\nRun 'pip install kuzu'")
             from .database_kuzu import KuzuDBManager
             info_logger(f"Using KùzuDB (explicit) at {db_path or 'default path'}")
             return KuzuDBManager(db_path=db_path)
@@ -168,7 +168,7 @@ def get_database_manager(db_path: Optional[str] = None) -> Union['DatabaseManage
         return NornicDBManager()
 
     error_msg = "No database backend available.\n"
-    error_msg += "Recommended: Install KùzuDB for zero-config ('pip install real_ladybug')\n"
+    error_msg += "Recommended: Install KùzuDB for zero-config ('pip install kuzu')\n"
 
     if platform.system() != "Windows":
         error_msg += "Alternative: Install FalkorDB Lite ('pip install falkordblite')\n"

codegraphcontext/core/database_kuzu.py

@@ -81,7 +81,7 @@ class KuzuDBManager:
                             info_logger("KùzuDB connection established and schema verified")
                             break
                         except ImportError:
-                            error_logger("KùzuDB is not installed. Run 'pip install real_ladybug'")
+                            error_logger("KùzuDB is not installed. Run 'pip install kuzu'")
                             raise ValueError("KùzuDB missing.")
                         except Exception as e:
                             if "lock" in str(e).lower() and attempt < max_retries - 1:
@@ -108,8 +108,8 @@ class KuzuDBManager:
             ("Directory", "path STRING, name STRING, PRIMARY KEY (path)"),
             ("Module", "name STRING, lang STRING, full_import_name STRING, PRIMARY KEY (name)"),
             # For types with composite keys (name, path, line_number), we use a 'uid'
-            ("Function", "uid STRING, name STRING, path STRING, line_number INT64, end_line INT64, source STRING, docstring STRING, lang STRING, cyclomatic_complexity INT64, context STRING, context_type STRING, class_context STRING, is_dependency BOOLEAN, decorators STRING[], args STRING[], http_method STRING, http_path STRING, PRIMARY KEY (uid)"),
-            ("Class", "uid STRING, name STRING, path STRING, line_number INT64, end_line INT64, source STRING, docstring STRING, lang STRING, is_dependency BOOLEAN, decorators STRING[], PRIMARY KEY (uid)"),
+            ("Function", "uid STRING, name STRING, path STRING, line_number INT64, end_line INT64, source STRING, docstring STRING, lang STRING, cyclomatic_complexity INT64, context STRING, context_type STRING, class_context STRING, class_context_line INT64, is_dependency BOOLEAN, decorators STRING[], args STRING[], http_method STRING, http_path STRING, PRIMARY KEY (uid)"),
+            ("Class", "uid STRING, name STRING, path STRING, line_number INT64, end_line INT64, source STRING, docstring STRING, lang STRING, node_type STRING, is_dependency BOOLEAN, decorators STRING[], PRIMARY KEY (uid)"),
             ("Variable", "uid STRING, name STRING, path STRING, line_number INT64, source STRING, docstring STRING, lang STRING, value STRING, context STRING, is_dependency BOOLEAN, PRIMARY KEY (uid)"),
             ("Trait", "uid STRING, name STRING, path STRING, line_number INT64, end_line INT64, source STRING, docstring STRING, lang STRING, is_dependency BOOLEAN, PRIMARY KEY (uid)"),
             ("Interface", "uid STRING, name STRING, path STRING, line_number INT64, end_line INT64, source STRING, docstring STRING, lang STRING, is_dependency BOOLEAN, PRIMARY KEY (uid)"),
@@ -132,9 +132,21 @@ class KuzuDBManager:
             # or the rel table creation will fail silently, leading to runtime
             # "Binder exception: Table CONTAINS does not exist".
             ("CONTAINS", "FROM File TO Function, FROM File TO Class, FROM File TO Variable, FROM File TO Trait, FROM File TO Interface, FROM `Macro` TO `Macro`, FROM File TO `Macro`, FROM File TO Struct, FROM File TO Enum, FROM File TO `Union`, FROM File TO Annotation, FROM File TO Record, FROM File TO `Property`, FROM Repository TO Directory, FROM Directory TO Directory, FROM Directory TO File, FROM Repository TO File, FROM Class TO Function, FROM Function TO Function", True),
-            ("CALLS", "FROM Function TO Function, FROM Function TO Class, FROM File TO Function, FROM File TO Class, FROM Class TO Function, FROM Class TO Class, line_number INT64, args STRING[], full_call_name STRING, confidence DOUBLE, resolution_tier INT64, confidence_label STRING, source STRING, resolution_method STRING, called_name STRING", True),
+            ("CALLS", """
+                FROM Function TO Function, FROM Function TO Class, FROM Function TO Interface, FROM Function TO Trait, FROM Function TO Struct, FROM Function TO Enum, FROM Function TO `Record`, FROM Function TO `Union`,
+                FROM Class TO Function, FROM Class TO Class, FROM Class TO Interface, FROM Class TO Trait, FROM Class TO Struct, FROM Class TO Enum, FROM Class TO `Record`, FROM Class TO `Union`,
+                FROM Interface TO Function, FROM Interface TO Class, FROM Interface TO Interface,
+                FROM File TO Function, FROM File TO Class, FROM File TO Interface, FROM File TO Trait, FROM File TO Struct, FROM File TO Enum, FROM File TO `Record`, FROM File TO `Union`,
+                line_number INT64, args STRING[], full_call_name STRING, confidence DOUBLE, resolution_tier INT64, confidence_label STRING, source STRING, resolution_method STRING, called_name STRING
+            """, True),
             ("IMPORTS", "FROM File TO Module, alias STRING, full_import_name STRING, imported_name STRING, line_number INT64", False),
-            ("INHERITS", "FROM Class TO Class, FROM Record TO Record, FROM Interface TO Interface, confidence_label STRING", True),
+            ("INHERITS", """
+                FROM Class TO Class, FROM Class TO Interface, FROM Class TO Trait,
+                FROM `Record` TO `Record`, FROM `Record` TO Interface,
+                FROM Interface TO Interface, FROM Interface TO Trait,
+                FROM Struct TO Interface, FROM Struct TO Trait,
+                confidence_label STRING
+            """, True),
             ("HAS_PARAMETER", "FROM Function TO Parameter", False),
             ("INCLUDES", "FROM Class TO Module", False),
             ("IMPLEMENTS", "FROM Class TO Interface, FROM Struct TO Interface, FROM Record TO Interface", True),
@@ -171,12 +183,15 @@ class KuzuDBManager:
             ("Module", "full_import_name", "STRING"),
             ("IMPORTS", "full_import_name", "STRING"),
             ("IMPORTS", "imported_name", "STRING"),
-            # Freshness properties added to Repository in 0.4.7
+            # Freshness properties added to Repository in 0.4.8
             ("Repository", "indexed_at", "STRING"),
             ("Repository", "commit_hash", "STRING"),
             # Spring endpoint properties on Function
             ("Function", "http_method", "STRING"),
             ("Function", "http_path", "STRING"),
+            # Kotlin/JVM precision improvements
+            ("Function", "class_context_line", "INT64"),
+            ("Class", "node_type", "STRING"),
         ]
 
         # REL TABLE GROUP migrations: KuzuDB creates sub-tables named
@@ -222,6 +237,7 @@ class KuzuDBManager:
                     continue
                 warning_logger(f"Kuzu Schema Migration Error ({table_name}.{column_name}): {e}")
                 debug_log(f"Kuzu Schema Migration Error ({table_name}.{column_name}): {e}")
+                raise RuntimeError("Kuzu Schema Migration Failed") from e
 
     def close_driver(self):
         """Closes the connection."""
@@ -259,7 +275,7 @@ class KuzuDBManager:
             import kuzu
             return True, None
         except ImportError:
-            return False, "KùzuDB is not installed. Run 'pip install real_ladybug'"
+            return False, "KùzuDB is not installed. Run 'pip install kuzu'"
 
 class KuzuDriverWrapper:
     def __init__(self, conn, query_lock: Optional[threading.RLock] = None):
@@ -470,8 +486,8 @@ class KuzuSessionWrapper:
             'File': {'path', 'name', 'relative_path', 'package_name', 'is_dependency'},
             'Directory': {'path', 'name'},
             'Module': {'name', 'lang', 'full_import_name'},
-            'Function': {'uid', 'name', 'path', 'line_number', 'end_line', 'source', 'docstring', 'lang', 'cyclomatic_complexity', 'context', 'context_type', 'class_context', 'is_dependency', 'decorators', 'args', 'http_method', 'http_path'},
-            'Class': {'uid', 'name', 'path', 'line_number', 'end_line', 'source', 'docstring', 'lang', 'is_dependency', 'decorators'},
+            'Function': {'uid', 'name', 'path', 'line_number', 'end_line', 'source', 'docstring', 'lang', 'cyclomatic_complexity', 'context', 'context_type', 'class_context', 'class_context_line', 'is_dependency', 'decorators', 'args', 'http_method', 'http_path'},
+            'Class': {'uid', 'name', 'path', 'line_number', 'end_line', 'source', 'docstring', 'lang', 'node_type', 'is_dependency', 'decorators'},
             'Variable': {'uid', 'name', 'path', 'line_number', 'source', 'docstring', 'lang', 'value', 'context', 'is_dependency'},
             'Trait': {'uid', 'name', 'path', 'line_number', 'end_line', 'source', 'docstring', 'lang', 'is_dependency'},
             'Interface': {'uid', 'name', 'path', 'line_number', 'end_line', 'source', 'docstring', 'lang', 'is_dependency'},
@@ -701,6 +717,38 @@ class KuzuSessionWrapper:
         for label in labels_to_escape:
             query = re.sub(rf':{label}\b', f':`{label}`', query)
 
+        # Translate (n:Label1 OR n:Label2 ...) to label(n) IN ['Label1', 'Label2', ...]
+        def poly_replacer(match):
+            full_match = match.group(0)
+            var_name = match.group(1)
+            # Find all labels associated with this variable in the OR chain
+            labels = re.findall(rf'{var_name}:([a-zA-Z0-9_`]+)', full_match)
+            # Strip backticks from labels
+            labels = [l.strip('`') for l in labels]
+            return f"label({var_name}) IN {json.dumps(labels)}"
+        
+        # Regex to match (n:Label1 OR n:Label2 OR n:Label3)
+        query = re.sub(r'\((\w+):[a-zA-Z0-9_`]+(?:\s+OR\s+\1:[a-zA-Z0-9_`]+)+\)', poly_replacer, query)
+        
+        # Translate single WHERE n:Label to label(n) = 'Label'
+        # This is more complex because we don't want to match MATCH/MERGE
+        # For now, we only target where it appears after WHERE or AND/OR
+        def single_label_replacer(match):
+            prefix = match.group(1)
+            var_name = match.group(2)
+            label = match.group(3).strip('`')
+            return f"{prefix}label({var_name}) = '{label}'"
+            
+        query = re.sub(r'(WHERE\s+|AND\s+|OR\s+|WHEN\s+)(\w+):([a-zA-Z0-9_`]+)', single_label_replacer, query, flags=re.IGNORECASE)
+
+        # Handle NOT n:Label → NOT label(n) = 'Label'
+        def not_label_replacer(match):
+            prefix = match.group(1)
+            var_name = match.group(2)
+            label_name = match.group(3).strip('`')
+            return f"{prefix}NOT label({var_name}) = '{label_name}'"
+        query = re.sub(r'(WHERE\s+|AND\s+|OR\s+)NOT\s+(\w+):([a-zA-Z0-9_`]+)', not_label_replacer, query, flags=re.IGNORECASE)
+
         # 4. Polymorphic matches and label access
         query = query.replace("labels(n)[0]", "label(n)")