9782e31ae5
- Improved DocumentVersion storage attributes to reflect Minio settings - Added size to DocumentVersions to easily calculate usage - License / LicenseTier forms and views added
117 lines
4.5 KiB
Python
117 lines
4.5 KiB
Python
from common.extensions import db
|
|
from .user import User, Tenant
|
|
from pgvector.sqlalchemy import Vector
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
|
|
|
|
class Document(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
name = db.Column(db.String(100), nullable=False)
|
|
tenant_id = db.Column(db.Integer, db.ForeignKey(Tenant.id), nullable=False)
|
|
valid_from = db.Column(db.DateTime, nullable=True)
|
|
valid_to = db.Column(db.DateTime, nullable=True)
|
|
|
|
# Versioning Information
|
|
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
|
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
|
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
|
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
|
|
|
# Relations
|
|
versions = db.relationship('DocumentVersion', backref='document', lazy=True)
|
|
|
|
def __repr__(self):
|
|
return f"<Document {self.id}: {self.name}>"
|
|
|
|
|
|
class DocumentVersion(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
doc_id = db.Column(db.Integer, db.ForeignKey(Document.id), nullable=False)
|
|
url = db.Column(db.String(200), nullable=True)
|
|
bucket_name = db.Column(db.String(255), nullable=True)
|
|
object_name = db.Column(db.String(200), nullable=True)
|
|
file_type = db.Column(db.String(20), nullable=True)
|
|
file_size = db.Column(db.Float, nullable=True)
|
|
language = db.Column(db.String(2), nullable=False)
|
|
user_context = db.Column(db.Text, nullable=True)
|
|
system_context = db.Column(db.Text, nullable=True)
|
|
user_metadata = db.Column(JSONB, nullable=True)
|
|
system_metadata = db.Column(JSONB, nullable=True)
|
|
|
|
# Versioning Information
|
|
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
|
created_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
|
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
|
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
|
|
|
# Processing Information
|
|
processing = db.Column(db.Boolean, nullable=False, default=False)
|
|
processing_started_at = db.Column(db.DateTime, nullable=True)
|
|
processing_finished_at = db.Column(db.DateTime, nullable=True)
|
|
processing_error = db.Column(db.String(255), nullable=True)
|
|
|
|
# Relations
|
|
embeddings = db.relationship('Embedding', backref='document_version', lazy=True)
|
|
|
|
def __repr__(self):
|
|
return f"<DocumentVersion {self.document_language.document_id}.{self.document_language.language}>.{self.id}>"
|
|
|
|
def calc_file_location(self):
|
|
return f"{self.document.tenant_id}/{self.document.id}/{self.language}"
|
|
|
|
def calc_file_name(self):
|
|
return f"{self.id}.{self.file_type}"
|
|
|
|
|
|
class Embedding(db.Model):
|
|
__tablename__ = 'embeddings'
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
type = db.Column(db.String(30), nullable=False)
|
|
doc_vers_id = db.Column(db.Integer, db.ForeignKey(DocumentVersion.id), nullable=False)
|
|
active = db.Column(db.Boolean, nullable=False, default=True)
|
|
chunk = db.Column(db.Text, nullable=False)
|
|
|
|
__mapper_args__ = {
|
|
'polymorphic_identity': 'embedding',
|
|
'polymorphic_on': type
|
|
}
|
|
|
|
|
|
class EmbeddingMistral(Embedding):
|
|
__tablename__ = 'embedding_mistral'
|
|
id = db.Column(db.Integer, db.ForeignKey('embeddings.id'), primary_key=True)
|
|
|
|
# 1024 is the MISTRAL Embedding dimension.
|
|
# If another embedding model is chosen, this dimension may need to be changed.
|
|
embedding = db.Column(Vector(1024), nullable=False)
|
|
|
|
__mapper_args__ = {
|
|
'polymorphic_identity': 'embedding_mistral',
|
|
}
|
|
|
|
|
|
class EmbeddingSmallOpenAI(Embedding):
|
|
__tablename__ = 'embedding_small_openai'
|
|
id = db.Column(db.Integer, db.ForeignKey('embeddings.id'), primary_key=True)
|
|
|
|
# 1536 is the OpenAI Small Embedding dimension.
|
|
# If another embedding model is chosen, this dimension may need to be changed.
|
|
embedding = db.Column(Vector(1536), nullable=False)
|
|
|
|
__mapper_args__ = {
|
|
'polymorphic_identity': 'embedding_small_openai',
|
|
}
|
|
|
|
|
|
class EmbeddingLargeOpenAI(Embedding):
|
|
__tablename__ = 'embedding_large_openai'
|
|
id = db.Column(db.Integer, db.ForeignKey('embeddings.id'), primary_key=True)
|
|
|
|
# 3072 is the OpenAI Large Embedding dimension.
|
|
# If another embedding model is chosen, this dimension may need to be changed.
|
|
embedding = db.Column(Vector(3072), nullable=False)
|
|
|
|
__mapper_args__ = {
|
|
'polymorphic_identity': 'embedding_large_openai',
|
|
}
|