scripts/seed_rag_data.py

#!/usr/bin/env python3
"""Seed Pinecone namespaces with role-specific corpus documents."""

import argparse
import asyncio
import os
from pathlib import Path

from dotenv import load_dotenv
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

from app.core.rag import get_rag_service
from app.core.schemas import TeamRole

load_dotenv()

ROLE_DIRECTORIES = {
    "product_owner": "product_owner",
    "business_analyst": "business_analyst",
    "solution_architect": "solution_architect",
    "data_architect": "data_architect",
    "security_analyst": "security_analyst",
    "ux_designer": "ux_designer",
    "api_designer": "api_designer",
    "qa_strategist": "qa_strategist",
    "devops_architect": "devops_architect",
    "technical_writer": "technical_writer",
}

CORPUS_DIR = Path(__file__).parent.parent / "corpus_rag"


def _build_splitter() -> RecursiveCharacterTextSplitter:
    return RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        add_start_index=True,
        separators=["\n\n", "\n", ". ", " ", ""],
    )


async def seed_collection(
    role_name: str,
    directory_name: str,
    dry_run: bool = False,
) -> dict:
    result = {
        "role": role_name,
        "files_found": 0,
        "chunks_found": 0,
        "chunks_inserted": 0,
        "errors": [],
    }

    try:
        role = TeamRole(role_name)
    except ValueError:
        result["errors"].append(f"Invalid role: {role_name}")
        return result

    dir_path = CORPUS_DIR / directory_name
    if not dir_path.exists():
        result["errors"].append(f"Directory not found: {dir_path}")
        return result

    splitter = _build_splitter()
    documents: list[Document] = []
    supported_extensions = {".md", ".txt", ".yaml", ".yml"}

    for file_path in dir_path.glob("**/*"):
        if file_path.is_dir() or file_path.suffix.lower() not in supported_extensions:
            continue

        result["files_found"] += 1

        try:
            content = file_path.read_text(encoding="utf-8")
        except Exception as exc:
            result["errors"].append(f"Error reading {file_path.name}: {exc}")
            continue

        if not content.strip():
            continue

        chunks = splitter.split_text(content)
        for idx, chunk in enumerate(chunks):
            documents.append(
                Document(
                    page_content=chunk,
                    metadata={
                        "source": file_path.name,
                        "chunk_index": idx,
                        "total_chunks": len(chunks),
                        "role": role_name,
                        "file_type": file_path.suffix,
                        "file_path": str(file_path.relative_to(CORPUS_DIR)),
                    },
                )
            )

    result["chunks_found"] = len(documents)
    if not documents:
        result["errors"].append("No documents found to seed")
        return result

    if dry_run:
        print(f"  [DRY RUN] Would insert {len(documents)} chunks")
        return result

    rag_service = get_rag_service()
    if not rag_service.is_pinecone_available():
        result["errors"].append("Pinecone backend not available")
        return result

    try:
        ids = await rag_service.add_documents(documents=documents, role=role)
        result["chunks_inserted"] = len(ids)
    except Exception as exc:
        result["errors"].append(f"Error inserting documents: {exc}")

    return result


async def seed_all(roles: list[str] | None = None, dry_run: bool = False) -> int:
    to_seed = (
        ROLE_DIRECTORIES
        if not roles
        else {
            name: directory
            for name, directory in ROLE_DIRECTORIES.items()
            if name in roles
        }
    )

    if not to_seed:
        print(f"Error: no valid roles in {roles}")
        print(f"Valid roles: {list(ROLE_DIRECTORIES.keys())}")
        return 1

    print(f"Seeding RAG namespaces{' [DRY RUN]' if dry_run else ''}...")
    print(f"Corpus directory: {CORPUS_DIR}")
    print()

    total_chunks = 0
    total_inserted = 0
    total_errors = 0

    for role_name, directory in to_seed.items():
        print(f"Processing: {role_name}")
        result = await seed_collection(role_name, directory, dry_run)

        total_chunks += result["chunks_found"]
        total_inserted += result["chunks_inserted"]

        if result["errors"]:
            total_errors += len(result["errors"])
            for err in result["errors"]:
                print(f"  - {err}")
        elif dry_run:
            print(
                f"  - Found {result['chunks_found']} chunks from {result['files_found']} files"
            )
        else:
            print(
                f"  - Inserted {result['chunks_inserted']} chunks from {result['files_found']} files"
            )

    print("\n" + "=" * 50)
    print("SEEDING COMPLETE")
    print("=" * 50)
    print(f"Total chunks found: {total_chunks}")
    if not dry_run:
        print(f"Total chunks inserted: {total_inserted}")
    if total_errors:
        print(f"Total errors: {total_errors}")

    return 0 if total_errors == 0 else 1


def create_corpus_directories() -> None:
    CORPUS_DIR.mkdir(exist_ok=True)
    for role_name, directory in ROLE_DIRECTORIES.items():
        dir_path = CORPUS_DIR / directory
        dir_path.mkdir(exist_ok=True)
        readme_path = dir_path / "README.md"
        if not readme_path.exists():
            readme_path.write_text(
                f"# {role_name.replace('_', ' ').title()} Examples\n\n"
                f"Place {role_name} examples here.\n",
                encoding="utf-8",
            )


def _validate_env(dry_run: bool) -> bool:
    if dry_run:
        return True

    required = ["PINECONE_API_KEY", "PINECONE_INDEX", "NVIDIA_API_KEY"]
    missing = [key for key in required if not os.getenv(key)]
    if missing:
        print(f"Error: missing required environment variables: {', '.join(missing)}")
        return False
    return True


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Seed Pinecone RAG namespaces with example documents"
    )
    parser.add_argument("--role", type=str, help="Specific role to seed")
    parser.add_argument("--dry-run", action="store_true", help="Preview without upsert")
    parser.add_argument(
        "--create-dirs",
        action="store_true",
        help="Create corpus_rag directory structure",
    )
    parser.add_argument(
        "--list-roles",
        action="store_true",
        help="List available role names",
    )

    args = parser.parse_args(argv)

    if args.list_roles:
        for role in ROLE_DIRECTORIES:
            print(role)
        return 0

    if args.create_dirs:
        create_corpus_directories()
        return 0

    if not _validate_env(args.dry_run):
        return 1

    roles = [args.role] if args.role else None
    return asyncio.run(seed_all(roles=roles, dry_run=args.dry_run))


if __name__ == "__main__":
    raise SystemExit(main())