File size: 3,739 Bytes
7699638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | #!/usr/bin/env python3
"""
Script: Ingest documents into both Basic RAG (FAISS) and GraphRAG (TigerGraph).
Usage:
python scripts/ingest_documents.py --data data/sample_docs/
python scripts/ingest_documents.py --data /path/to/your/corpus/
"""
import sys
import json
import pickle
import logging
import argparse
from pathlib import Path
from dotenv import load_dotenv
# Load .env before importing any modules that use environment variables
load_dotenv(Path(__file__).parent.parent / ".env")
sys.path.insert(0, str(Path(__file__).parent.parent))
from backend.graph.tigergraph_client import TigerGraphClient
from backend.graph.ingestion import DocumentIngestionPipeline, chunk_text
from backend.rag.basic_rag import BasicRAG
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)
def ingest_basic_rag(data_dir: str, faiss_path: str = "data/faiss_index.pkl"):
"""Build FAISS index for Basic RAG from raw documents."""
logger.info("📦 Building FAISS index for Basic RAG...")
rag = BasicRAG()
path = Path(data_dir)
files = list(path.glob("*.txt")) + list(path.glob("*.md"))
all_chunks = []
all_metadata = []
for f in files:
with open(f, "r", encoding="utf-8") as fh:
content = fh.read()
chunks = chunk_text(content)
all_chunks.extend(chunks)
all_metadata.extend([{"source": str(f), "file": f.name}] * len(chunks))
logger.info(f" {f.name}: {len(chunks)} chunks")
rag.add_documents(all_chunks, all_metadata)
# Save FAISS index for reuse
with open(faiss_path, "wb") as f:
pickle.dump({"index": rag.index, "chunks": rag.chunks, "metadata": rag.chunk_metadata}, f)
logger.info(f"✅ FAISS index saved: {faiss_path} ({len(all_chunks)} total chunks)")
return rag
def ingest_graph_rag(data_dir: str, tg_client: TigerGraphClient):
"""Ingest documents into TigerGraph knowledge graph."""
logger.info("🕸️ Ingesting documents into TigerGraph...")
pipeline = DocumentIngestionPipeline(tg_client)
pipeline.ingest_directory(data_dir)
stats = pipeline.get_stats()
logger.info(f"\n✅ TigerGraph ingestion complete:")
logger.info(f" Docs processed: {stats['docs_processed']}")
logger.info(f" Chunks created: {stats['chunks_created']}")
logger.info(f" Entities extracted: {stats['entities_extracted']}")
logger.info(f" Relationships created: {stats['relationships_extracted']}")
return stats
def main():
parser = argparse.ArgumentParser(description="Ingest documents into RAG pipelines")
parser.add_argument("--data", default="data/sample_docs/", help="Directory with .txt/.md files")
parser.add_argument("--skip-faiss", action="store_true", help="Skip FAISS index build")
parser.add_argument("--skip-graph", action="store_true", help="Skip TigerGraph ingestion")
parser.add_argument("--faiss-output", default="data/faiss_index.pkl", help="FAISS index output path")
args = parser.parse_args()
Path("data").mkdir(exist_ok=True)
if not args.skip_faiss:
ingest_basic_rag(args.data, args.faiss_output)
if not args.skip_graph:
try:
tg = TigerGraphClient()
tg.connect()
ingest_graph_rag(args.data, tg)
except Exception as e:
logger.error(f"TigerGraph ingestion failed: {e}")
logger.info("Hint: Run scripts/setup_tigergraph.py first, or check your .env credentials.")
logger.info("\n🎉 Ingestion complete! Now run the server:")
logger.info(" cd backend && uvicorn api.server:app --reload --port 8000")
if __name__ == "__main__":
main()
|