Full support for different Embedding models in the database. There was an error.

common/models/document.py

@@ -75,7 +75,7 @@ class DocumentVersion(db.Model):
     processing_error = db.Column(db.String(255), nullable=True)
 
     # Relations
-    embeddings = db.relationship('EmbeddingMistral', backref='document_version', lazy=True)
+    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}>"
@@ -87,23 +87,41 @@ class DocumentVersion(db.Model):
         return f"{self.id}.{self.file_type}"
 
 
-class EmbeddingMistral(db.Model):
+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(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)
+
+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',
+    }