src/cli/ontology.py

import argparse
import json
import os
import sys
from typing import List
from src.ontology.matcher import ConceptMatcher
from src.embeddings.engine import EmbeddingEngine
from src.ontology.models import OntologyRecord

def main():
    parser = argparse.ArgumentParser(description="Ontology Operations CLI")
    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # stats
    subparsers.add_parser("stats", help="Show ontology statistics")

    # inspect
    inspect_parser = subparsers.add_parser("inspect", help="Inspect a specific concept")
    inspect_parser.add_argument("name", help="Canonical name or alias to inspect")

    # search
    search_parser = subparsers.add_parser("search", help="Search for concepts")
    search_parser.add_argument("query", help="Search query")
    search_parser.add_argument("--top-k", type=int, default=5, help="Number of results")

    # validate
    subparsers.add_parser("validate", help="Validate ontology files")

    # rebuild-faiss
    subparsers.add_parser("rebuild-faiss", help="Rebuild FAISS index")

    args = parser.parse_args()

    ontology_dir = "data/ontology"
    
    if args.command == "stats":
        show_stats(ontology_dir)
    elif args.command == "inspect":
        inspect_concept(ontology_dir, args.name)
    elif args.command == "search":
        search_concepts(ontology_dir, args.query, args.top_k)
    elif args.command == "validate":
        validate_ontology(ontology_dir)
    elif args.command == "rebuild-faiss":
        rebuild_faiss(ontology_dir)
    else:
        parser.print_help()

def show_stats(ontology_dir):
    report = {}
    for filename in os.listdir(ontology_dir):
        if filename.endswith(".json"):
            path = os.path.join(ontology_dir, filename)
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
                category = filename.replace(".json", "")
                report[category] = {
                    "count": len(data),
                    "nsfw": sum(1 for item in data if item.get("nsfw")),
                    "synthetic": sum(1 for item in data if item.get("synthetic")),
                    "metadata_coverage": sum(1 for item in data if item.get("metadata")) / len(data) if data else 0
                }
    print(json.dumps(report, indent=2))

def inspect_concept(ontology_dir, name):
    matcher = ConceptMatcher(ontology_dir)
    matches = matcher.match(name)
    if matches:
        print(json.dumps([m.model_dump() for m in matches], indent=2))
    else:
        print(f"Concept '{name}' not found.")

def search_concepts(ontology_dir, query, top_k):
    engine = EmbeddingEngine(index_dir="data/faiss_indices")
    try:
        engine.load_index()
    except FileNotFoundError:
        print("FAISS index not found. Please run 'rebuild-faiss' first.")
        return
        
    results = engine.search(query, top_k=top_k)
    output = []
    for record, score in results:
        output.append({
            "canonical": record.canonical,
            "category": record.category,
            "score": round(float(score), 4)
        })
    print(json.dumps(output, indent=2))

def validate_ontology(ontology_dir):
    errors = []
    for filename in os.listdir(ontology_dir):
        if filename.endswith(".json"):
            path = os.path.join(ontology_dir, filename)
            try:
                with open(path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    for i, item in enumerate(data):
                        try:
                            OntologyRecord(**item)
                        except Exception as e:
                            errors.append(f"Error in {filename} at index {i}: {str(e)}")
            except Exception as e:
                errors.append(f"Error reading {filename}: {str(e)}")
    
    if errors:
        for err in errors:
            print(err)
    else:
        print("Ontology validation passed.")

def rebuild_faiss(ontology_dir):
    engine = EmbeddingEngine(index_dir="data/faiss_indices")
    print("Rebuilding FAISS index...")
    engine.build_index(ontology_dir)
    print("FAISS index rebuilt successfully.")

if __name__ == "__main__":
    main()