required = ['OPENAI_API_KEY', 'GRAPH_URI', 'GRAPH_USER', 'GRAPH_PASSWORD']

missing = [v for v in required if not os.getenv(v)]

if missing:

click.echo(f"❌ Missing environment vars: {', '.join(missing)}")

return False

dataset_name: Optional[str], sample_chunks: Optional[int], sample_seed: Optional[int]):

"""Create chunk file(s) from documents (file or directory)."""

from file_chunk import chunk_documents, ChunkingStrategy, load_documents, load_documents_from_directory

import random

input_path_obj = Path(input_path)

if dataset_name is None:

dataset_name = input_path_obj.stem

out_dir = Path(output_dir) / dataset_name

out_dir.mkdir(parents=True, exist_ok=True)

chunk_file = out_dir / f"{dataset_name}_chunk.json"

# Load docs

if input_path_obj.is_dir():

documents = load_documents_from_directory(str(input_path_obj))

else:

documents = load_documents(str(input_path_obj))

if not documents:

click.echo("❌ No documents found")

return

click.echo(f"📄 Loaded {len(documents)} documents")

strategy_map = {

'semantic': ChunkingStrategy.SEMANTIC,

'hybrid': ChunkingStrategy.HYBRID,

'fixed_token': ChunkingStrategy.FIXED_TOKEN,

}

click.echo(f"🔄 Chunking (strategy={strategy}, size={chunk_size}, overlap={overlap})")

results = chunk_documents(documents, max_token_size=chunk_size, overlap_token_size=overlap,

strategy=strategy_map[strategy])

flattened: List[Dict[str, Any]] = []

for doc in results:

for ch in doc['chunks']:

flattened.append({

'hash_code': ch['hash_code'],

'text': ch['text'],

'tokens': ch['tokens'],

'quality_score': ch.get('quality_score', 0.0),

'strategy': ch['strategy']

})

original_total = len(flattened)

full_copy = None

if sample_chunks and 0 < sample_chunks < original_total:

full_copy = flattened

rng = random.Random(sample_seed)

flattened = rng.sample(flattened, sample_chunks)

click.echo(f"🎲 Sampled {len(flattened)} / {original_total} chunks (seed={sample_seed})")

elif sample_chunks and sample_chunks >= original_total:

click.echo("⚠️ sample-chunks >= total; using all chunks")

with open(chunk_file, 'w', encoding='utf-8') as f:

json.dump(flattened, f, ensure_ascii=False, indent=2)

if full_copy is not None:

with open(out_dir / f"{dataset_name}_chunk_full.json", 'w', encoding='utf-8') as f:

json.dump(full_copy, f, ensure_ascii=False, indent=2)

"""Extract triples and entities from chunks (auto-detect *_chunk.json if directory given).

if not validate_environment():

return

try:

path_obj = Path(input_path)

chosen_chunk: Optional[Path] = None

if path_obj.is_dir():

candidates = list(path_obj.glob('*_chunk.json')) or list(path_obj.rglob('*_chunk.json'))

if not candidates:

click.echo(f"❌ No *_chunk.json found in directory {input_path}")

return

candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)

chosen_chunk = candidates[0]

if len(candidates) > 1:

click.echo("⚠️ Multiple chunk files; using most recent: " + str(chosen_chunk))

else:

chosen_chunk = path_obj

if not chosen_chunk.name.endswith('_chunk.json'):

click.echo("⚠️ File does not end with _chunk.json – continuing")

if not chosen_chunk.exists():

click.echo("❌ Chunk file not found")

return

with open(chosen_chunk, 'r', encoding='utf-8') as f:

raw_chunks = json.load(f)

chunks_map = {c['hash_code']: c['text'] for c in raw_chunks}

click.echo(f"🔍 Extracting over {len(chunks_map)} chunks ...")

from GraphExtraction.chunk import triple_extraction, openai_generate_text

async def runner():

async def llm_func(prompt, history_messages=None):

if history_messages:

hist = '\n'.join(m['content'] for m in history_messages if m['role'] in ('user','assistant'))

prompt_full = hist + '\n' + prompt

else:

prompt_full = prompt

resp = await openai_generate_text(prompt_full)

return resp or ''

return await triple_extraction(chunks_map, llm_func)

import asyncio, sqlite3, json as _json

from database_utils import create_db_table_sqlite, get_db_path, get_embedding, store_text_units_records

create_db_table_sqlite(output_dir)

entities, relations = asyncio.run(runner())

# Merge entities by name

merged = {}

for e in entities:

name = e.get('entity_name')

if not name:

continue

entry = merged.setdefault(name, {'entity_name': name, 'description': '', 'source_id': '', 'degree': 0})

d = (e.get('description') or '').strip()

if d:

entry['description'] = (entry['description'] + ' | ' + d).strip(' |') if entry['description'] else d

sid = (e.get('source_id') or '').strip()

if sid and sid not in entry['source_id'].split('|'):

entry['source_id'] = (entry['source_id'] + '|' + sid).strip('|')

for v in merged.values():

if v['source_id']:

v['source_id'] = '|'.join(sorted(set(filter(None, v['source_id'].split('|')))))

# Normalize relations (dedupe by src,tgt,desc)

norm_rels = []

seen = set()

for r in relations:

if isinstance(r, list):

r = next((x for x in r if isinstance(x, dict)), None)

if not isinstance(r, dict):

continue

src = r.get('src_id') or r.get('source')

tgt = r.get('tgt_id') or r.get('target')

if not src or not tgt:

continue

desc = r.get('description') or ''

key = (src, tgt, desc)

if key in seen:

continue

seen.add(key)

try:

weight = int(r.get('weight', 1))

except Exception:

weight = 1

norm_rels.append({'src_tgt': src, 'tgt_src': tgt, 'description': desc, 'weight': weight})

db_path = get_db_path()

conn = sqlite3.connect(db_path)

cur = conn.cursor()

ent_count = 0

for ent in merged.values():

vec = []

if ent['description']:

vec = get_embedding(ent['description'])[0]

cur.execute(

"""INSERT OR REPLACE INTO entities (entity_name, description, source_id, degree, parent, level, vector, is_new)

VALUES (?, ?, ?, ?, ?, ?, ?, 1)""",

(ent['entity_name'], ent['description'], ent['source_id'], ent['degree'], None, 0, _json.dumps(vec) if vec else None)

)

ent_count += 1

rel_count = 0

for rel in norm_rels:

cur.execute(

"""INSERT INTO relations (src_tgt, tgt_src, description, weight, level, is_new)

VALUES (?, ?, ?, ?, ?, 1)""",

(rel['src_tgt'], rel['tgt_src'], rel['description'], rel['weight'], 0)

)

rel_count += 1

conn.commit()

conn.close()

inserted_units = store_text_units_records(raw_chunks) if raw_chunks else 0

click.echo("✅ Extraction complete")

click.echo(f" 🧬 Entities inserted: {ent_count}")

click.echo(f" 🔗 Relations inserted: {rel_count}")

click.echo(f" 📚 Text units stored: {inserted_units}")

click.echo(f" 💾 SQLite: {db_path}")

click.echo(" ➜ Next: leanrag build (working_dir optional)")

except Exception as e:

click.echo(f"❌ Extraction error: {e}", err=True)

"""Build graph & vector index from extracted (SQLite) data. WORKING_DIR optional."""

if not validate_environment():

return

try:

from database_utils import (

create_db_table_sqlite,

mark_entities_relations_processed,

store_text_units_in_db,

)

import build_graph as build_pipeline

from build_graph import hierarchical_clustering, embedding as pipeline_embedding

create_db_table_sqlite(working_dir)

chat_client = OpenAIChatClient()

build_pipeline.WORKING_DIR = working_dir

cpu_workers = os.cpu_count() or 4

max_workers = max(1, workers or cpu_workers)

json_format = {"type": "json_object"}

config = {

'max_workers': max_workers,

'working_dir': working_dir,

'use_llm_func': lambda prompt, **kwargs: chat_client.chat(prompt, **kwargs),

'embeddings_func': pipeline_embedding,

'aggregate_llm_kwargs': {'response_format': None},

'relation_llm_kwargs': {'response_format': None},

'clear_existing': refresh,

}

mode_label = "full refresh" if refresh else "incremental sync"

click.echo(f"📊 Building hierarchy with {max_workers} workers ({mode_label})")

graph_stats = hierarchical_clustering(config) or {}

mark_entities_relations_processed(working_dir)

if chunks_file:

try:

store_text_units_in_db(working_dir, chunks_file)

click.echo("📚 Text units stored from chunks file")

except Exception as te:

click.echo(f"⚠️ Text unit store failed: {te}")

cleared_nodes = graph_stats.get('cleared_nodes', 0)

entity_created = graph_stats.get('entities', 0)

relation_created = graph_stats.get('relationships', 0)

community_created = graph_stats.get('communities', 0)

sqlite_entities = graph_stats.get('sqlite_entities')

sqlite_relations = graph_stats.get('sqlite_relations')

sqlite_communities = graph_stats.get('sqlite_communities')

if refresh:

click.echo(f"🔄 Neo4j cleared ({cleared_nodes} nodes removed) and Qdrant collection recreated")

else:

click.echo("🔄 Hierarchy synced; Neo4j and Qdrant updated incrementally")

click.echo(f" 🧩 Entities created: {entity_created}")

click.echo(f" 🔗 Relationships created: {relation_created}")

click.echo(f" 🫂 Communities created: {community_created}")

if sqlite_entities is not None:

click.echo(f" 💾 SQLite entities stored: {sqlite_entities}")

if sqlite_relations is not None:

click.echo(f" 💾 SQLite relations stored: {sqlite_relations}")

if sqlite_communities is not None:

click.echo(f" 💾 SQLite communities stored: {sqlite_communities}")

summary_path = Path(working_dir) / "level_summary.txt"

if summary_path.exists():

click.echo("📐 Hierarchy level summary:")

click.echo(summary_path.read_text(encoding='utf-8').rstrip())

click.echo(f" ↳ Full summary: {summary_path}")

else:

click.echo("ℹ️ No hierarchical level summary found in working directory")

click.echo("✅ Build pipeline complete")

except Exception as e:

click.echo(f"❌ Build error: {e}")

"""Visualize a Neo4j subgraph (lightweight).

if not validate_environment():

return

try:

from visualize_graph import fetch_subgraph, build_nx_graph, draw_graph

except ImportError as ie:

click.echo(f"❌ Missing visualization dependencies: {ie}. Run: pip install networkx matplotlib")

return

try:

nodes, rels = fetch_subgraph(level, limit, rel_type, sample)

if not nodes:

click.echo("ℹ️ No nodes returned for given filters")

return

click.echo(f"📥 Retrieved {len(nodes)} nodes / {len(rels)} relationships")

import networkx as nx # noqa: F401 (ensure dependency present)

G = build_nx_graph(nodes, rels)

draw_graph(G, layout, output, not no_labels, node_size, font_size)

if output:

click.echo(f"✅ Visualization saved to {output}")

else:

click.echo("✅ Visualization complete (window may still be open)")

except Exception as e:

click.echo(f"❌ Visualization error: {e}")

"""Query the knowledge graph.

if not validate_environment():

return

try:

from query_graph import query_graph

click.echo(f"🔍 Querying: {query}")

# Create config from environment and parameters

from query_graph import embedding

config = {

'topk': top_k,

'level_mode': level_mode,

'chunks_file': chunks_file,

'embeddings_func': embedding

}

# Auto-detect chunks file if not provided

if not chunks_file:

# Try common locations for chunks file

potential_chunks = [

"test_chunk.json",

"datasets/cs/cs_chunk.json",

"datasets/legal/legal_chunk.json",

"datasets/agriculture/agriculture_chunk.json",

"datasets/mix/mix_chunk.json",

"output_science/science/science_chunk.json"

]

for chunk_path in potential_chunks:

if os.path.exists(chunk_path):

chunks_file = chunk_path

break

else:

# If user passed a directory, attempt auto-detect inside it

if os.path.isdir(chunks_file):

dir_path = Path(chunks_file)

candidates = list(dir_path.glob("*_chunk.json"))

if not candidates:

candidates = list(dir_path.rglob("*_chunk.json"))

if not candidates:

click.echo(f"❌ No *_chunk.json files found inside directory: {chunks_file}")

return

if len(candidates) > 1:

candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)

click.echo("⚠️ Multiple chunk files found; selecting most recent for query:")

for idx, c in enumerate(candidates, 1):

marker = "✅" if idx == 1 else " "

click.echo(f" {marker} {idx}. {c}")

chunks_file = str(candidates[0])

if not chunks_file:

click.echo("❌ Chunks file not found. Please specify with --chunks-file", err=True)

return

# Derive working_dir from chunks_file path

working_dir = os.path.dirname(chunks_file)

if not working_dir:

working_dir = "." # If chunks_file is in root directory

config['chunks_file'] = chunks_file

config['working_dir'] = working_dir

# Run query

click.echo("⏳ Processing query...")

ref, response = query_graph(config, None, query) # db parameter not used in new version

click.echo("\n" + "="*50)

click.echo("📋 REFERENCE INFORMATION:")

click.echo(ref)

click.echo("\n" + "="*50)

click.echo("🤖 GENERATED RESPONSE:")

click.echo(response)

click.echo("="*50)

except Exception as e:

click.echo(f"❌ Error during query: {str(e)}", err=True)

"""Run the complete LeanRAG pipeline: chunk → extract → build → query.

if not validate_environment():

return

try:

dataset_name = Path(input_file).stem

work_dir = Path(output_dir) / dataset_name

work_dir.mkdir(parents=True, exist_ok=True)

click.echo("🚀 Starting LeanRAG pipeline...")

click.echo(f" 📁 Working directory: {work_dir}")

click.echo(f" 📄 Input file: {input_file}")

# Step 1: Chunk documents

click.echo("\n1️⃣ Step 1: Chunking documents...")

chunk_file = work_dir / f"{dataset_name}_chunk.json"

# Import here to avoid circular imports

from file_chunk import chunk_documents, ChunkingStrategy

documents = []

with open(input_file, 'r', encoding='utf-8') as f:

for line in f:

if line.strip():

data = json.loads(line)

text = (data.get('context') or data.get('text') or

data.get('content') or data.get('input', ''))

if text:

documents.append(text)

results = chunk_documents(documents, strategy=ChunkingStrategy.SEMANTIC)

flattened_results = []

for doc_result in results:

for chunk in doc_result['chunks']:

flattened_results.append({

"hash_code": chunk["hash_code"],

"text": chunk["text"],

"tokens": chunk["tokens"],

"quality_score": chunk.get("quality_score", 0.0),

"strategy": chunk["strategy"]

})

with open(chunk_file, 'w', encoding='utf-8') as f:

json.dump(flattened_results, f, ensure_ascii=False, indent=2)

click.echo(f" ✅ Created {len(flattened_results)} chunks")

# Step 2: Extract triples (manual step notification)

click.echo("\n2️⃣ Step 2: Extract triples and entities (GraphRAG)")

click.echo(" ⚠️ Manual extraction step – run after configuring LLM env vars")

click.echo(f" 💡 Run: leanrag extract {chunk_file}")

click.echo(" 💡 Or manually configure and run the extraction scripts")

# Step 3: Build graph (would be manual)

click.echo("\n3️⃣ Step 3: Build knowledge graph")

click.echo(f" 💡 After extraction, run: leanrag build {work_dir}")

# Step 4: Query (if provided)

if query:

click.echo(f"\n4️⃣ Step 4: Query knowledge graph")

click.echo(" ⚠️ Requires completed graph building first")

click.echo(f" 💡 Run: leanrag query \"{query}\" {work_dir} --chunks-file {chunk_file}")

click.echo("\n🎉 Pipeline setup complete!")

click.echo(" 📋 Next steps:")

click.echo(" 1. Configure LLM endpoints in config files")

click.echo(" 2. Run extraction: leanrag extract ...")

click.echo(" 3. Build graph: leanrag build ...")

if query:

click.echo(f" 4. Query: leanrag query \"{query}\" ...")

except Exception as e:

click.echo(f"❌ Error in pipeline: {str(e)}", err=True)

"""Check system status and configuration."""

click.echo("🔍 Checking LeanRAG system status...")

# Check environment variables

env_status = {

'TOGETHER_API_KEY': bool(os.getenv('TOGETHER_API_KEY')),

'OPENAI_API_KEY': bool(os.getenv('OPENAI_API_KEY')),

'QDRANT_URL': bool(os.getenv('QDRANT_URL')),

'GRAPH_URI': bool(os.getenv('GRAPH_URI')),

'GRAPH_USER': bool(os.getenv('GRAPH_USER')),

'GRAPH_PASSWORD': bool(os.getenv('GRAPH_PASSWORD'))

}

click.echo("\n📋 Environment Variables:")

for var, status in env_status.items():

icon = "✅" if status else "❌"

click.echo(f" {icon} {var}: {'Set' if status else 'Missing'}")

# Check if databases are accessible

click.echo("\n🗄️ Database Connectivity:")

# Check Qdrant

try:

from qdrant_client import QdrantClient

qdrant_url = os.getenv('QDRANT_URL', 'http://localhost:6333')

client = QdrantClient(url=qdrant_url)

client.get_collections()

click.echo(" ✅ Qdrant: Connected")

except Exception as e:

click.echo(f" ❌ Qdrant: {str(e)}")

# Check Neo4j

try:

from neo4j import GraphDatabase

graph_uri = os.getenv('GRAPH_URI', 'bolt://localhost:7687')

graph_user = os.getenv('GRAPH_USER', 'neo4j')

graph_password = os.getenv('GRAPH_PASSWORD', 'test123456')

driver = GraphDatabase.driver(graph_uri, auth=(graph_user, graph_password))

with driver.session() as session:

session.run("RETURN 1")

click.echo(" ✅ Neo4j: Connected")

except Exception as e:

click.echo(f" ❌ Neo4j: {str(e)}")

# Check required files

click.echo("\n📁 Required Files:")

required_files = ['requirements.txt', '.env']

for file in required_files:

exists = os.path.exists(file)

icon = "✅" if exists else "❌"

click.echo(f" {icon} {file}: {'Found' if exists else 'Missing'}")

click.echo("\n💡 Tips:")

click.echo(" - Install missing dependencies: pip install -r requirements.txt")

click.echo(" - Set environment variables in .env file")

click.echo(" - Start databases: docker run -p 6333:6333 qdrant/qdrant")