Amazon Bedrock や Neptune を活用した神経生物学的インスパイア型 RAG「HippoRAG」の紹介

class HotpotQANeptuneImporter:
    """Class to handle importing HotpotQA data into Neptune."""

    def __init__(
        self,
        hotpotqa_file_path: str,
        output_dir: str,
        s3_bucket: str,
        s3_prefix: str,
        neptune_endpoint: str,
        neptune_port: int,
        iam_role_arn: str,
        aws_region: str,
        llm_endpoint: Optional[str] = None,
        embedding_endpoint: Optional[str] = None,
        max_workers: int = 4,
        max_examples: int = 10,
        max_docs_per_example: int = 2,
        use_ssl: bool = False
    ):
        """Initialize the importer with configuration."""
        self.hotpotqa_file_path = hotpotqa_file_path
        self.output_dir = output_dir
        self.s3_bucket = s3_bucket
        self.s3_prefix = s3_prefix
        self.neptune_endpoint = neptune_endpoint
        self.neptune_port = neptune_port
        self.iam_role_arn = iam_role_arn
        self.aws_region = aws_region

        # Initialize AWS clients
        self.s3_client = boto3.client('s3', region_name=aws_region)
        self.session = boto3.Session()

        # Initialize data structures
        self.phrase_dict = {}  # Maps phrase text to node ID
        self.phrase_embeddings = {}  # Maps phrase text to embedding vector

def extract_triples_with_llm(self, text: str) -> List[Tuple[str, str, str]]:
    """
    Use an LLM to extract knowledge graph triples from text.
    In this simplified version, just generate simple triples from the text.
    """
    # Simple triple generation from text
    words = text.split()
    if len(words) < 5:
        return []

    # Generate simple triples from the words in the text
    triples = []

    for i in range(min(3, len(words) - 2)):  # Get at most 3 triples
        subject = words[i]
        relation = "related_to"
        obj = words[i+2]
        triples.append((subject, relation, obj))

    return triples

def process_data_to_csv(self, data: List[Dict]) -> None:
    """
    Process HotpotQA data into CSV files for Neptune import.
    """
    logger.info("Processing HotpotQA data to CSV files")

    # Open all CSV files for writing
    with open(os.path.join(self.output_dir, 'phrase_nodes.csv'), 'w', newline='', encoding='utf-8') as phrase_f, \
         open(os.path.join(self.output_dir, 'passage_nodes.csv'), 'w', newline='', encoding='utf-8') as passage_f, \
         open(os.path.join(self.output_dir, 'relation_edges.csv'), 'w', newline='', encoding='utf-8') as relation_f, \
         open(os.path.join(self.output_dir, 'context_edges.csv'), 'w', newline='', encoding='utf-8') as context_f:

        # Write headers
        phrase_f.write('~id,~label,text\n')
        passage_f.write('~id,~label,title,content\n')
        relation_f.write('~id,~from,~to,~label\n')
        context_f.write('~id,~from,~to,~label\n')

        # Process each example
        for example in tqdm(data, desc="Processing examples"):
            self._process_example(example, passage_f, phrase_f, relation_f, context_f)

def _process_example(self, example: Dict, passage_f, phrase_f, relation_f, context_f) -> None:
    """Process a single HotpotQA example into CSV data."""
    # Process limited documents in the context
    doc_count = 0
    for doc_title, paragraphs in example['context']:
        if doc_count >= self.max_docs_per_example:
            break

        # Process each paragraph as a passage
        for paragraph in paragraphs:
            if not paragraph.strip():
                continue

            passage_id = f"passage_{uuid.uuid4().hex}"

            # Write passage node
            passage_f.write(f"{passage_id},passage,\"{self._escape_csv(doc_title)}\",\"{self._escape_csv(paragraph)}\"\n")

            # Extract triples from the paragraph
            triples = self.extract_triples_with_llm(paragraph)

            # Process each triple
            for subject, relation, obj in triples:
                # Create or get subject node ID
                if subject not in self.phrase_dict:
                    subject_id = f"phrase_{uuid.uuid4().hex}"
                    self.phrase_dict[subject] = subject_id
                    phrase_f.write(f"{subject_id},phrase,\"{self._escape_csv(subject)}\"\n")
                else:
                    subject_id = self.phrase_dict[subject]

                # Create or get object node ID
                if obj not in self.phrase_dict:
                    obj_id = f"phrase_{uuid.uuid4().hex}"
                    self.phrase_dict[obj] = obj_id
                    phrase_f.write(f"{obj_id},phrase,\"{self._escape_csv(obj)}\"\n")
                else:
                    obj_id = self.phrase_dict[obj]

                # Write relation edge
                relation_id = f"rel_{uuid.uuid4().hex}"
                relation_f.write(f"{relation_id},{subject_id},{obj_id},\"{self._escape_csv(relation)}\"\n")

                # Write context edges
                context_f.write(f"ctx_{uuid.uuid4().hex},{passage_id},{subject_id},contains\n")
                context_f.write(f"ctx_{uuid.uuid4().hex},{passage_id},{obj_id},contains\n")

        doc_count += 1

def import_to_neptune(self) -> Dict:
    """Import data into Neptune using the bulk loader API."""
    logger.info(f"Importing data to Neptune endpoint {self.neptune_endpoint}")

    loader_endpoint = f"{self.protocol}://{self.neptune_endpoint}:{self.neptune_port}/loader"

    payload = {
        "source": f"s3://{self.s3_bucket}/{self.s3_prefix}/",
        "format": "csv",
        "formatParams": {
            "delimiter": ",",
            "header": "true"
        },
        "iamRoleArn": self.iam_role_arn,
        "region": self.aws_region,
        "failOnError": "FALSE"
    }

    try:
        # Create AWS4Auth for the request
        credentials = self.session.get_credentials()
        if credentials:
            auth = AWS4Auth(
                credentials.access_key,
                credentials.secret_key,
                self.aws_region,
                'neptune-db',
                session_token=credentials.token
            )

        response = requests.post(
            loader_endpoint,
            data=json.dumps(payload),
            headers={"Content-Type": "application/json"},
            timeout=30,
            auth=auth
        )

        response.raise_for_status()
        result = response.json()
        logger.info(f"Neptune load job submitted: {result}")

        return result
    except Exception as e:
        logger.error(f"Failed to submit Neptune load job: {e}")
        raise

def run_pipeline(self) -> None:
    """Run the full data processing and import pipeline."""
    try:
        # Step 1: Test Neptune connectivity
        if not self.test_neptune_connection():
            logger.warning("Neptune connectivity test failed, but proceeding with the pipeline")

        # Step 2: Load HotpotQA data
        data = self.load_hotpotqa_data()

        # Step 3: Process data to CSV files
        self.process_data_to_csv(data)

        # Step 4: Create empty synonym_edges.csv file
        with open(os.path.join(self.output_dir, 'synonym_edges.csv'), 'w', newline='', encoding='utf-8') as f:
            f.write('~id,~from,~to,~label\n')

        # Step 5: Upload CSV files to S3
        self.upload_to_s3()

        # Step 6: Import data to Neptune
        load_result = self.import_to_neptune()
        load_id = load_result.get("payload", {}).get("loadId")

        if load_id:
            # Step 7: Wait for import to complete
            final_status = self.wait_for_load_job(load_id)
            logger.info(f"Final import status: {final_status}")

        logger.info("Pipeline completed successfully")
    except Exception as e:
        logger.error(f"Pipeline failed: {e}")
        raise

from src.hipporag.utils.config_utils import BaseConfig

config = BaseConfig()
config.force_index_from_scratch = True
config.openie_mode = "online"
config.llm_name = "us.anthropic.claude-3-5-haiku-20241022-v1:0"
config.embedding_model_name = "amazon.titan-embed-text-v2:0"
config.aws_region = "us-east-1"
config.save_dir = "outputs"
config.retrieval_top_k = 3

class NeptuneAnalyticsClient:
    """Client for Neptune Analytics operations including personalized PageRank."""

    def __init__(self, graph_id, region='us-east-1'):
        """Initialize Neptune Analytics client."""
        self.graph_id = graph_id
        self.region = region
        self.client = boto3.client('neptune-analytics', region_name=region)
        logger.info(f"Initialized Neptune Analytics client for graph {graph_id}")

    def run_personalized_pagerank(self, seed_nodes, damping_factor=0.85, max_iterations=20, tolerance=0.0001):
        """
        Run personalized PageRank algorithm on Neptune Analytics.

        Args:
            seed_nodes (list): List of seed node IDs for personalized PageRank
            damping_factor (float): Damping factor (default 0.85)
            max_iterations (int): Maximum number of iterations
            tolerance (float): Convergence tolerance

        Returns:
            list: List of (node_id, score) tuples sorted by score descending
        """
        if not seed_nodes:
            logger.warning("No seed nodes provided for personalized PageRank")
            return []

        # Format seed nodes for the query
        seed_list = ",".join([f"'{node}'" for node in seed_nodes])

        # Neptune Analytics personalized PageRank query using openCypher
        query = f"""
        CALL neptune.algo.pagerank({{
            sourceNodes: [{seed_list}],
            dampingFactor: {damping_factor},
            maxIterations: {max_iterations},
            tolerance: {tolerance},
            personalized: true
        }})
        YIELD nodeId, score
        RETURN nodeId, score
        ORDER BY score DESC
        LIMIT 100
        """

        try:
            result = self.execute_analytics_query(query)

            if result and 'results' in result:
                pagerank_results = [(item['nodeId'], item['score']) for item in result['results']]
                logger.info(f"Personalized PageRank returned {len(pagerank_results)} results")
                return pagerank_results

            return []
        except Exception as e:
            logger.error(f"Failed to run personalized PageRank: {e}")
            return []

def index_from_neptune(self, limit=1000):
    """
    Index documents directly from Neptune database.

    Args:
        limit (int): Maximum number of documents to index
    """
    documents = self.neptune_client.get_documents(limit=limit)
    logger.info(f"Indexing {len(documents)} documents from Neptune database")

    # Use the standard indexing process with documents from Neptune
    super().index(documents)