Neo4j Agent Memory 시리즈 - 1편

import asyncio
import logging
import os
import time
from pathlib import Path
from dotenv import load_dotenv
from neo4j_agent_memory import MemoryClient, MemorySettings
from neo4j_agent_memory.extraction.pipeline import ExtractionPipeline
from neo4j_agent_memory.extraction.gliner_extractor import GLiNEREntityExtractor
from neo4j_agent_memory.extraction.llm_extractor import LLMEntityExtractor
from neo4j_agent_memory.extraction.spacy_extractor import SpacyEntityExtractor

load_dotenv()

LOG_DIR = Path("logs")
LOG_FILE = LOG_DIR / "three_extractor_time.log"
logger = logging.getLogger(__name__)


def get_torch_device() -> str:
    """Prefer GPU acceleration for GLiNER when PyTorch can use it."""
    try:
        import torch
    except ImportError:
        return "cpu"

    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def setup_logging():
    """Configure logs to write to both terminal and the script-specific log file."""
    LOG_DIR.mkdir(exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler(LOG_FILE, encoding="utf-8"),
        ],
    )


async def extract_and_log(extractor, content: str):
    """Run extraction once, log per-stage timing and entity results, return entities."""
    if extractor is None:
        logger.info("  No extractor configured")
        return []

    t0 = time.perf_counter()

    if isinstance(extractor, ExtractionPipeline):
        pipeline_result = await extractor.extract_with_details(content)
        total_ms = (time.perf_counter() - t0) * 1000

        logger.info("  Pipeline total: %.1fms", total_ms)
        for stage in pipeline_result.stage_results:
            if stage.success:
                entity_names = [e.name for e in stage.result.entities]
                logger.info(
                    "    ✓ %s: %d entities in %.1fms%s",
                    stage.stage_name,
                    len(entity_names),
                    stage.duration_ms,
                    f" → {entity_names}" if entity_names else "",
                )
            else:
                logger.info(
                    "    ✗ %s: failed in %.1fms",
                    stage.stage_name,
                    stage.duration_ms,
                )

        return pipeline_result.final_result.entities
    else:
        result = await extractor.extract(content)
        elapsed_ms = (time.perf_counter() - t0) * 1000
        logger.info(
            "  %s: %d entities in %.1fms",
            extractor.__class__.__name__,
            result.entity_count,
            elapsed_ms,
        )
        return result.entities


async def prewarm_extractor(extractor, content: str) -> bool:
    """Warm one extractor before measured extraction when it has lazy startup work."""
    if extractor is None:
        return False

    try:
        if isinstance(extractor, GLiNEREntityExtractor):
            t0 = time.perf_counter()
            await extractor.extract(content)
            logger.info(
                "Pre-warmed GLiNER extractor in %.1fms",
                (time.perf_counter() - t0) * 1000,
            )
            return True

        if isinstance(extractor, SpacyEntityExtractor):
            t0 = time.perf_counter()
            await extractor.extract(content)
            logger.info(
                "Pre-warmed spaCy extractor in %.1fms",
                (time.perf_counter() - t0) * 1000,
            )
            return True

        if isinstance(extractor, LLMEntityExtractor):
            t0 = time.perf_counter()
            extractor._ensure_client()
            logger.info(
                "Pre-warmed LLM extractor client in %.1fms",
                (time.perf_counter() - t0) * 1000,
            )
            return True
    except Exception as exc:
        logger.warning(
            "Pre-warm for %s failed: %s",
            extractor.__class__.__name__,
            exc,
        )
        return True

    return False


async def prewarm_extractors(extractor, content: str):
    """Warm configured extractors before measured extraction."""
    if extractor is None:
        return

    if await prewarm_extractor(extractor, content):
        return

    if isinstance(extractor, ExtractionPipeline):
        for stage in extractor.stages:
            stage_extractor = getattr(stage, "_extractor", None)
            await prewarm_extractor(stage_extractor, content)


async def main():
    setup_logging()
    gliner_device = get_torch_device()
    logger.info("Using GLiNER device: %s", gliner_device)

    settings = MemorySettings(
        neo4j={"uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), "username": os.getenv("NEO4J_USER", "neo4j"), "password": os.getenv("NEO4J_PASSWORD", "password123")},
        extraction={"gliner_device": gliner_device},
    )

    async with MemoryClient(settings) as client:
        logger.info("✓ Connected to Neo4j!")

        content = "Hello! I'm looking for running shoes. I really like Nike products."

        await prewarm_extractors(client._extractor, content)

        logger.info("\nExtraction pipeline:")
        entities = await extract_and_log(client._extractor, content)

        t0 = time.perf_counter()
        message = await client.short_term.add_message(
            session_id="test-session-001",
            role="user",
            content=content,
            extraction_mode="explicit",
            explicit_mentions=entities,
        )
        elapsed_ms = (time.perf_counter() - t0) * 1000
        logger.info("\n✓ Stored message: %s (%.1fms)", message.id, elapsed_ms)

        conversation = await client.short_term.get_conversation(
            session_id="test-session-001"
        )
        logger.info("✓ Retrieved %d message(s)", len(conversation.messages))
        for msg in conversation.messages:
            logger.info("  [%s]: %s", msg.role, msg.content)

    logger.info("✓ Connection closed!")


if __name__ == "__main__":
    asyncio.run(main())