"""Abstractions for documents and their extractions."""
import base64
import json
import logging
import uuid
from datetime import datetime
from enum import Enum
from typing import Optional, Union
from pydantic import BaseModel, Field
logger = logging.getLogger(__name__)
DataType = Union[str, bytes]
class DocumentType(str, Enum):
"""Types of documents that can be stored."""
CSV = "csv"
DOCX = "docx"
HTML = "html"
JSON = "json"
MD = "md"
PDF = "pdf"
PPTX = "pptx"
TXT = "txt"
XLSX = "xlsx"
GIF = "gif"
PNG = "png"
JPG = "jpg"
JPEG = "jpeg"
SVG = "svg"
MP3 = "mp3"
MP4 = "mp4"
class Document(BaseModel):
id: uuid.UUID = Field(default_factory=uuid.uuid4)
type: DocumentType
data: Union[str, bytes]
metadata: dict
def __init__(self, *args, **kwargs):
data = kwargs.get("data")
if data and isinstance(data, str):
try:
# Try to decode if it's already base64 encoded
kwargs["data"] = base64.b64decode(data)
except:
# If it's not base64, encode it to bytes
kwargs["data"] = data.encode("utf-8")
doc_type = kwargs.get("type")
if isinstance(doc_type, str):
kwargs["type"] = DocumentType(doc_type)
# Generate UUID based on the hash of the data
if "id" not in kwargs:
if isinstance(kwargs["data"], bytes):
data_hash = uuid.uuid5(
uuid.NAMESPACE_DNS, kwargs["data"].decode("utf-8")
)
else:
data_hash = uuid.uuid5(uuid.NAMESPACE_DNS, kwargs["data"])
kwargs["id"] = data_hash # Set the id based on the data hash
super().__init__(*args, **kwargs)
class Config:
arbitrary_types_allowed = True
json_encoders = {
uuid.UUID: str,
bytes: lambda v: base64.b64encode(v).decode("utf-8"),
}
class DocumentStatus(str, Enum):
"""Status of document processing."""
PROCESSING = "processing"
# TODO - Extend support for `partial-failure`
# PARTIAL_FAILURE = "partial-failure"
FAILURE = "failure"
SUCCESS = "success"
class DocumentInfo(BaseModel):
"""Base class for document information handling."""
document_id: uuid.UUID
version: str
size_in_bytes: int
metadata: dict
status: DocumentStatus = DocumentStatus.PROCESSING
user_id: Optional[uuid.UUID] = None
title: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
def convert_to_db_entry(self):
"""Prepare the document info for database entry, extracting certain fields from metadata."""
now = datetime.now()
metadata = self.metadata
if "user_id" in metadata:
metadata["user_id"] = str(metadata["user_id"])
metadata["title"] = metadata.get("title", "N/A")
return {
"document_id": str(self.document_id),
"title": metadata.get("title", "N/A"),
"user_id": metadata.get("user_id", None),
"version": self.version,
"size_in_bytes": self.size_in_bytes,
"metadata": json.dumps(self.metadata),
"created_at": self.created_at or now,
"updated_at": self.updated_at or now,
"status": self.status,
}
class ExtractionType(Enum):
"""Types of extractions that can be performed."""
TXT = "txt"
IMG = "img"
MOV = "mov"
class Extraction(BaseModel):
"""An extraction from a document."""
id: uuid.UUID
type: ExtractionType = ExtractionType.TXT
data: DataType
metadata: dict
document_id: uuid.UUID
class FragmentType(Enum):
"""A type of fragment that can be extracted from a document."""
TEXT = "text"
IMAGE = "image"
class Fragment(BaseModel):
"""A fragment extracted from a document."""
id: uuid.UUID
type: FragmentType
data: DataType
metadata: dict
document_id: uuid.UUID
extraction_id: uuid.UUID
class Entity(BaseModel):
"""An entity extracted from a document."""
category: str
subcategory: Optional[str] = None
value: str
def __str__(self):
return (
f"{self.category}:{self.subcategory}:{self.value}"
if self.subcategory
else f"{self.category}:{self.value}"
)
class Triple(BaseModel):
"""A triple extracted from a document."""
subject: str
predicate: str
object: str
def extract_entities(llm_payload: list[str]) -> dict[str, Entity]:
entities = {}
for entry in llm_payload:
try:
if "], " in entry: # Check if the entry is an entity
entry_val = entry.split("], ")[0] + "]"
entry = entry.split("], ")[1]
colon_count = entry.count(":")
if colon_count == 1:
category, value = entry.split(":")
subcategory = None
elif colon_count >= 2:
parts = entry.split(":", 2)
category, subcategory, value = (
parts[0],
parts[1],
parts[2],
)
else:
raise ValueError("Unexpected entry format")
entities[entry_val] = Entity(
category=category, subcategory=subcategory, value=value
)
except Exception as e:
logger.error(f"Error processing entity {entry}: {e}")
continue
return entities
def extract_triples(
llm_payload: list[str], entities: dict[str, Entity]
) -> list[Triple]:
triples = []
for entry in llm_payload:
try:
if "], " not in entry: # Check if the entry is an entity
elements = entry.split(" ")
subject = elements[0]
predicate = elements[1]
object = " ".join(elements[2:])
subject = entities[subject].value # Use entity.value
if "[" in object and "]" in object:
object = entities[object].value # Use entity.value
triples.append(
Triple(subject=subject, predicate=predicate, object=object)
)
except Exception as e:
logger.error(f"Error processing triplet {entry}: {e}")
continue
return triples
class KGExtraction(BaseModel):
"""An extraction from a document that is part of a knowledge graph."""
entities: dict[str, Entity]
triples: list[Triple]
