morphik-core/core/models/chunk.py

from typing import Any, Dict, List

from pydantic import BaseModel, Field
import numpy as np

Embedding = List[float] | List[List[float]] | np.ndarray


class DocumentChunk(BaseModel):
    """Represents a chunk stored in VectorStore"""

    document_id: str  # external_id of parent document
    content: str
    embedding: Embedding
    chunk_number: int
    # chunk-specific metadata
    metadata: Dict[str, Any] = Field(default_factory=dict)
    score: float = 0.0

    model_config = {"arbitrary_types_allowed": True}


class Chunk(BaseModel):
    """Represents a chunk containing content and metadata"""

    content: str
    metadata: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"arbitrary_types_allowed": True}

    def to_document_chunk(
        self, document_id: str, chunk_number: int, embedding: Embedding
    ) -> DocumentChunk:
        return DocumentChunk(
            document_id=document_id,
            content=self.content,
            embedding=embedding,
            chunk_number=chunk_number,
            metadata=self.metadata,
        )