Process Documents in Batch

from neo4j_agent_memory.extraction import GLiNEREntityExtractor

extractor = GLiNEREntityExtractor.for_schema("ecommerce")

# List of documents to process
documents = [
    "Customer John ordered Nike Air Max shoes from our Manhattan store.",
    "Jane returned the Adidas Ultraboost due to sizing issues.",
    "Order #12345 shipped via FedEx to Brooklyn, NY.",
    # ... hundreds more
]

# Batch extraction
result = await extractor.extract_batch(
    texts=documents,
    batch_size=10,
    max_concurrency=5,
)

print(f"Processed: {result.successful_items}/{result.total_items}")
print(f"Total entities: {result.total_entities}")
print(f"Total relations: {result.total_relations}")

def on_progress(completed: int, total: int):
    percent = (completed / total) * 100
    print(f"Progress: {completed}/{total} ({percent:.1f}%)")

result = await extractor.extract_batch(
    texts=documents,
    batch_size=10,
    max_concurrency=5,
    on_progress=on_progress,
)

result = await extractor.extract_batch(texts=documents)

for i, item_result in enumerate(result.items):
    if item_result.success:
        print(f"Doc {i}: {item_result.entity_count} entities")
        for entity in item_result.entities:
            print(f"  - {entity.name} ({entity.type})")
    else:
        print(f"Doc {i}: Failed - {item_result.error}")

from neo4j_agent_memory.extraction import StreamingExtractor

# Wrap extractor for streaming
streamer = StreamingExtractor(
    extractor=extractor,
    chunk_size=4000,    # Characters per chunk
    overlap=200,        # Overlap to avoid splitting entities
)

# Process a long document
long_document = open("annual_report.txt").read()  # 100K+ characters

# Stream results as they're extracted
async for chunk_result in streamer.extract_streaming(long_document):
    print(f"Chunk {chunk_result.chunk.index}: {chunk_result.entity_count} entities")

    # Store entities as they come
    for entity in chunk_result.entities:
        await client.long_term.add_entity(
            name=entity.name,
            entity_type=entity.type,
        )

# Extract with automatic cross-chunk deduplication
result = await streamer.extract(long_document, deduplicate=True)

print(f"Total chunks: {result.stats.total_chunks}")
print(f"Raw entities: {result.stats.raw_entity_count}")
print(f"Unique entities: {result.stats.deduplicated_entities}")

for entity in result.entities:
    print(f"  {entity.name} ({entity.type})")

from neo4j_agent_memory import MemoryClient

client = MemoryClient(
    neo4j_uri="bolt://localhost:7687",
    neo4j_user="neo4j",
    neo4j_password="password",
)

# Extract batch
result = await extractor.extract_batch(texts=documents)

# Store all entities
for item in result.items:
    if not item.success:
        continue

    for entity in item.entities:
        await client.long_term.add_entity(
            name=entity.name,
            entity_type=entity.type,
            properties={
                "confidence": entity.confidence,
                "source_index": item.index,
            },
        )

from neo4j_agent_memory.memory import DeduplicationConfig

dedup_config = DeduplicationConfig(
    auto_merge_threshold=0.95,
    flag_threshold=0.85,
)

stored_count = 0
merged_count = 0

for item in result.items:
    if not item.success:
        continue

    for entity in item.entities:
        stored, dedup_result = await client.long_term.add_entity(
            name=entity.name,
            entity_type=entity.type,
            deduplication=dedup_config,
        )

        if dedup_result.action == "created":
            stored_count += 1
        elif dedup_result.action == "merged":
            merged_count += 1

print(f"New entities: {stored_count}")
print(f"Merged with existing: {merged_count}")

# Register the extractor for provenance
await client.long_term.register_extractor(
    name="GLiNEREntityExtractor",
    version="1.0",
    config={"schema": "ecommerce"},
)

# Store with provenance
for doc_index, item in enumerate(result.items):
    if not item.success:
        continue

    # Store document as a message/source
    message = await client.short_term.add_message(
        role="system",
        content=documents[doc_index],
        session_id="batch-import",
        metadata={"doc_index": doc_index},
    )

    for entity in item.entities:
        stored = await client.long_term.add_entity(
            name=entity.name,
            entity_type=entity.type,
        )

        # Link to source document
        await client.long_term.link_entity_to_message(
            entity=stored,
            message_id=message.id,
            confidence=entity.confidence,
            start_pos=entity.start,
            end_pos=entity.end,
        )

        # Link to extractor
        await client.long_term.link_entity_to_extractor(
            entity=stored,
            extractor_name="GLiNEREntityExtractor",
            confidence=entity.confidence,
        )

# Smaller batches for memory-constrained environments
result = await extractor.extract_batch(
    texts=documents,
    batch_size=5,  # Smaller batches
    max_concurrency=2,
)

# Larger batches for high-memory environments with GPU
result = await extractor.extract_batch(
    texts=documents,
    batch_size=50,  # Larger batches
    max_concurrency=10,
)

# High-volume, lower precision
fast_extractor = SpacyEntityExtractor()  # Fastest

# Balanced
balanced_extractor = GLiNEREntityExtractor.for_schema("ecommerce")

# High precision, slower
accurate_pipeline = ExtractionPipeline(
    stages=[
        GLiNEREntityExtractor.for_schema("ecommerce"),
        LLMEntityExtractor(model="gpt-4o-mini"),
    ],
)

result = await extractor.extract_batch(
    texts=documents,
    fail_fast=False,  # Continue on errors (default)
)

# Check for failures
failed = [item for item in result.items if not item.success]
print(f"Failed: {len(failed)}/{result.total_items}")

for item in failed:
    print(f"  Index {item.index}: {item.error}")

# Filter out invalid documents
valid_docs = [
    doc for doc in documents
    if doc.get("content") and len(doc["content"]) > 10
]

print(f"Processing {len(valid_docs)} of {len(documents)} documents")
result = await extractor.extract_batch(texts=[d["content"] for d in valid_docs])

import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def on_progress(completed, total):
    logger.info(f"Progress: {completed}/{total}")

result = await extractor.extract_batch(
    texts=documents,
    on_progress=on_progress,
)

for item in result.items:
    if not item.success:
        logger.error(f"Failed doc {item.index}: {item.error}")

# For very large document sets, process in chunks
CHUNK_SIZE = 1000

for i in range(0, len(all_documents), CHUNK_SIZE):
    chunk = all_documents[i:i + CHUNK_SIZE]
    print(f"Processing chunk {i//CHUNK_SIZE + 1}")

    result = await extractor.extract_batch(texts=chunk)

    # Store results before next chunk
    await store_results(result)

# For known-clean data, skip deduplication
await client.long_term.add_entity(
    name=entity.name,
    entity_type=entity.type,
    deduplicate=False,  # Faster for trusted sources
)

# For user-generated content, enable deduplication
await client.long_term.add_entity(
    name=entity.name,
    entity_type=entity.type,
    deduplication=dedup_config,
)