110 lines
4.6 KiB
Python
110 lines
4.6 KiB
Python
from common.extensions import db
|
|
from .user import User, Tenant
|
|
from pgvector.sqlalchemy import Vector
|
|
|
|
|
|
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=False)
|
|
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
|
|
languages = db.relationship('DocumentLanguage', backref='document', lazy=True)
|
|
|
|
def __repr__(self):
|
|
return f"<Document {self.id}: {self.name}>"
|
|
|
|
|
|
class DocumentLanguage(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
document_id = db.Column(db.Integer, db.ForeignKey(Document.id), nullable=False)
|
|
language = db.Column(db.String(2), nullable=False)
|
|
user_context = db.Column(db.Text, nullable=True)
|
|
system_context = db.Column(db.Text, nullable=True)
|
|
latest_version_id = db.Column(db.Integer, db.ForeignKey('document_version.id'), 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=False)
|
|
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_language',
|
|
lazy='joined',
|
|
foreign_keys='DocumentVersion.doc_lang_id'
|
|
)
|
|
latest_version = db.relationship(
|
|
'DocumentVersion',
|
|
uselist=False,
|
|
foreign_keys=[latest_version_id]
|
|
)
|
|
|
|
def __repr__(self):
|
|
return f"<DocumentLanguage {self.document_id}.{self.language}>"
|
|
|
|
|
|
class DocumentVersion(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
doc_lang_id = db.Column(db.Integer, db.ForeignKey(DocumentLanguage.id), nullable=False)
|
|
url = db.Column(db.String(200), nullable=True)
|
|
file_location = db.Column(db.String(255), nullable=True)
|
|
file_name = db.Column(db.String(200), nullable=True)
|
|
file_type = db.Column(db.String(20), 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('EmbeddingMistral', 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_language.document.tenant_id}/{self.document_language.document.id}/{self.document_language.language}"
|
|
|
|
def calc_file_name(self):
|
|
return f"{self.id}.{self.file_type}"
|
|
|
|
|
|
class EmbeddingMistral(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
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)
|
|
|
|
# 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)
|
|
|
|
|
|
class EmbeddingSmallOpenAI(db.Model):
|
|
id = db.Column(db.Integer, primary_key=True)
|
|
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)
|
|
|
|
# 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)
|