#!/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())