Azure platform on embeddings (#145)

* Additional data field on the embeddings table

* Added azure as embedding platform

* Import/export additional data

* Submodule dev change

---------

Co-authored-by: JWittmeyer <[email protected]>

Author: lumburovskalina

alembic/versions/0714589d508e_additional_data_on_the_embeddings_table.py

@@ -0,0 +1,28 @@
+"""Additional data on the embeddings table
+
+Revision ID: 0714589d508e
+Revises: 73798599a917
+Create Date: 2023-08-03 13:21:49.464532
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '0714589d508e'
+down_revision = '73798599a917'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('embedding', sa.Column('additional_data', sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('embedding', 'additional_data')
+    # ### end Alembic commands ###

controller/embedding/manager.py

@@ -43,25 +43,15 @@ def get_recommended_encoders(is_managed: bool) -> List[Any]:
     return recommendations
 
 
+def create_embedding(project_id: str, embedding_id: str) -> None:
+    daemon.run(connector.request_embedding, project_id, embedding_id)
 
-def create_embedding(
-    project_id: str, embedding_id: str
-) -> None:
-    daemon.run(
-        connector.request_embedding,
-        project_id,
-        embedding_id
-    )
 
 def create_embeddings_one_by_one(
     project_id: str,
     embeddings_ids: List[str],
 ) -> None:
-    daemon.run(
-        __embed_one_by_one_helper,
-        project_id,
-        embeddings_ids
-    )
+    daemon.run(__embed_one_by_one_helper, project_id, embeddings_ids)
 
 
 def request_tensor_upload(project_id: str, embedding_id: str) -> Any:
@@ -74,10 +64,7 @@ def delete_embedding(project_id: str, embedding_id: str) -> None:
     connector.request_deleting_embedding(project_id, embedding_id)
 
 
-def __embed_one_by_one_helper(
-    project_id: str,
-    embeddings_ids: List[str]
-) -> None:
+def __embed_one_by_one_helper(project_id: str, embeddings_ids: List[str]) -> None:
     for embedding_id in embeddings_ids:
         connector.request_embedding(project_id, embedding_id)
         time.sleep(5)
@@ -86,7 +73,12 @@ def __embed_one_by_one_helper(
 
 
 def get_embedding_name(
-    project_id: str, attribute_id: str,  platform: str, embedding_type: str, model: Optional[str] = None, apiToken: Optional[str] = None
+    project_id: str,
+    attribute_id: str,
+    platform: str,
+    embedding_type: str,
+    model: Optional[str] = None,
+    api_token: Optional[str] = None,
 ) -> str:
     if embedding_type not in [
         enums.EmbeddingType.ON_ATTRIBUTE.value,
@@ -109,19 +101,21 @@ def get_embedding_name(
     if model:
         name += f"-{model}"
 
-    if apiToken:
-        name += f"-{apiToken[:3]}...{apiToken[-4:]}"
-    
+    if api_token:
+        name += f"-{api_token[:3]}...{api_token[-4:]}"
+
     return name
 
 
-def recreate_embeddings(project_id: str, embedding_ids: Optional[List[str]] = None) -> None:
+def recreate_embeddings(
+    project_id: str, embedding_ids: Optional[List[str]] = None
+) -> None:
     if not embedding_ids:
         embeddings = embedding.get_all_embeddings_by_project_id(project_id)
         if len(embeddings) == 0:
             return
         embedding_ids = [str(embed.id) for embed in embeddings]
-    
+
     set_to_wait = False
     for embedding_id in embedding_ids:
         set_to_wait = True
@@ -130,10 +124,10 @@ def recreate_embeddings(project_id: str, embedding_ids: Optional[List[str]] = No
 
     if set_to_wait:
         notification.send_organization_update(
-                project_id,
-                f"embedding:{None}:state:{enums.EmbeddingState.WAITING.value}",
-            )
-        
+            project_id,
+            f"embedding:{None}:state:{enums.EmbeddingState.WAITING.value}",
+        )
+
     for embedding_id in embedding_ids:
         new_id = None
         try:
@@ -147,13 +141,17 @@ def recreate_embeddings(project_id: str, embedding_ids: Optional[List[str]] = No
                 embedding_item = general.refresh(embedding_item)
                 if not embedding_item:
                     raise Exception("Embedding not found")
-                elif embedding_item.state == enums.EmbeddingState.FAILED.value or embedding_item.state == enums.EmbeddingState.FINISHED.value:
+                elif (
+                    embedding_item.state == enums.EmbeddingState.FAILED.value
+                    or embedding_item.state == enums.EmbeddingState.FINISHED.value
+                ):
                     break
                 else:
                     time.sleep(1)
         except Exception as e:
             print(
-                f"Error while recreating embedding for {project_id} with id {embedding_id} - {e}", flush=True
+                f"Error while recreating embedding for {project_id} with id {embedding_id} - {e}",
+                flush=True,
             )
             notification.send_organization_update(
                 project_id,
@@ -162,23 +160,19 @@ def recreate_embeddings(project_id: str, embedding_ids: Optional[List[str]] = No
             old_embedding_item = embedding.get(project_id, embedding_id)
             if old_embedding_item:
                 old_embedding_item.state = enums.EmbeddingState.FAILED.value
-            
+
             if new_id:
                 new_embedding_item = embedding.get(project_id, new_id)
                 if new_embedding_item:
                     new_embedding_item.state = enums.EmbeddingState.FAILED.value
             general.commit()
 
-
         notification.send_organization_update(
             project_id=project_id, message="embedding:finished:all"
         )
 
 
-
-def __recreate_embedding(
-    project_id: str, embedding_id: str
-) -> Embedding:
+def __recreate_embedding(project_id: str, embedding_id: str) -> Embedding:
     old_embedding_item = embedding.get(project_id, embedding_id)
     old_id = old_embedding_item.id
     new_embedding_item = embedding.create(
@@ -191,25 +185,30 @@ def __recreate_embedding(
         model=old_embedding_item.model,
         platform=old_embedding_item.platform,
         api_token=old_embedding_item.api_token,
-        with_commit=False
+        additional_data=old_embedding_item.additional_data,
+        with_commit=False,
     )
     embedding.delete(project_id, embedding_id, with_commit=False)
     embedding.delete_tensors(embedding_id, with_commit=False)
     general.commit()
 
-    if new_embedding_item.platform == enums.EmbeddingPlatform.OPENAI.value or new_embedding_item.platform == enums.EmbeddingPlatform.COHERE.value:
-        agreement_item = agreement.get_by_xfkey(project_id, old_id, enums.AgreementType.EMBEDDING.value)
+    if (
+        new_embedding_item.platform == enums.EmbeddingPlatform.OPENAI.value
+        or new_embedding_item.platform == enums.EmbeddingPlatform.COHERE.value
+        or new_embedding_item.platform == enums.EmbeddingPlatform.AZURE.value
+    ):
+        agreement_item = agreement.get_by_xfkey(
+            project_id, old_id, enums.AgreementType.EMBEDDING.value
+        )
         if not agreement_item:
             new_embedding_item.state = enums.EmbeddingState.FAILED.value
             general.commit()
-            raise Exception(f"No agreement found for embedding {new_embedding_item.name}")
+            raise Exception(
+                f"No agreement found for embedding {new_embedding_item.name}"
+            )
         agreement_item.xfkey = new_embedding_item.id
         general.commit()
 
     connector.request_deleting_embedding(project_id, old_id)
-    daemon.run(
-        connector.request_embedding,
-        project_id,
-        new_embedding_item.id
-    )
-    return new_embedding_item
+    daemon.run(connector.request_embedding, project_id, new_embedding_item.id)
+    return new_embedding_item

controller/task_queue/handler/embedding.py

@@ -18,7 +18,6 @@ def get_task_functions() -> Tuple[Callable, Callable, int]:
 
 
 def __start_task(task: Dict[str, Any]) -> bool:
-
     # check task still relevant
     task_db_obj = task_queue_db_bo.get(task["id"])
     if task_db_obj is None or task_db_obj.is_active:
@@ -47,6 +46,7 @@ def __start_task(task: Dict[str, Any]) -> bool:
     terms_accepted = task["task_info"]["terms_accepted"]
 
     filter_attributes = task["task_info"]["filter_attributes"]
+    additional_data = task["task_info"]["additional_data"]
     embedding_item = embedding_db_bo.create(
         project_id,
         attribute_id,
@@ -58,6 +58,7 @@ def __start_task(task: Dict[str, Any]) -> bool:
         platform=platform,
         api_token=api_token,
         filter_attributes=filter_attributes,
+        additional_data=additional_data,
     )
     if (
         platform == enums.EmbeddingPlatform.OPENAI.value
@@ -80,7 +81,6 @@ def __start_task(task: Dict[str, Any]) -> bool:
 
 
 def __check_finished(task: Dict[str, Any]) -> bool:
-
     embedding_item = embedding_db_bo.get_embedding_by_name(
         task["project_id"], task["task_info"]["embedding_name"]
     )

controller/transfer/project_transfer_manager.py

@@ -397,6 +397,9 @@ def __transform_embedding_by_name(embedding_name: str):
                 model=embedding_item.get(
                     "model",
                 ),
+                additional_data=embedding_item.get(
+                    "additional_data",
+                ),
             )
             embedding_ids[
                 embedding_item.get(
@@ -1174,6 +1177,7 @@ def get_project_export_dump(
             "finished_at": embedding_item.finished_at,
             "platform": embedding_item.platform,
             "model": embedding_item.model,
+            "additional_data": embedding_item.additional_data,
         }
         for embedding_item in embeddings
     ]

graphql_api/mutation/embedding.py

@@ -31,12 +31,18 @@ def mutate(self, info, project_id: str, attribute_id: str, config: Dict[str, Any
         api_token = config.get("apiToken")
         terms_text = config.get("termsText")
         terms_accepted = config.get("termsAccepted")
+        additional_data = None
+        if config.get("base") is not None:
+            additional_data = {
+                "base": config.get("base"),
+                "type": config.get("type"),
+                "version": config.get("version"),
+            }
 
         # prototyping logic, this will be part of config after ui integration
         relevant_attribute_list = attribute_do.get_all_possible_names_for_qdrant(
             project_id
         )
-
         task_queue_manager.add_task(
             project_id,
             TaskType.EMBEDDING,
@@ -53,6 +59,7 @@ def mutate(self, info, project_id: str, attribute_id: str, config: Dict[str, Any
                 "terms_text": terms_text,
                 "terms_accepted": terms_accepted,
                 "filter_attributes": relevant_attribute_list,
+                "additional_data": additional_data,
             },
         )
         notification.send_organization_update(

graphql_api/query/embedding.py

@@ -4,7 +4,12 @@
 
 from controller.misc import manager as misc
 from controller.auth import manager as auth
-from graphql_api.types import EmbeddingPlatform, Encoder, LanguageModel, RecordTokenizationTask
+from graphql_api.types import (
+    EmbeddingPlatform,
+    Encoder,
+    LanguageModel,
+    RecordTokenizationTask,
+)
 from submodules.model import enums
 from submodules.model.business_objects import tokenization, task_queue
 from util import spacy_util
@@ -57,23 +62,28 @@ def resolve_project_tokenization(
     def resolve_embedding_platforms(self, info) -> List[EmbeddingPlatform]:
         return [
             {
-            "platform": enums.EmbeddingPlatform.HUGGINGFACE.value,
-            "terms": None,
-            "link": None
+                "platform": enums.EmbeddingPlatform.HUGGINGFACE.value,
+                "terms": None,
+                "link": None,
             },
             {
-            "platform": enums.EmbeddingPlatform.COHERE.value,
-            "terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the Cohere API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
-            "link": "https://openai.com/policies/api-data-usage-policies"
-            }, 
+                "platform": enums.EmbeddingPlatform.COHERE.value,
+                "terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the Cohere API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
+                "link": "https://cohere.com/terms-of-use",
+            },
+            {
+                "platform": enums.EmbeddingPlatform.OPENAI.value,
+                "terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the OpenAI API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
+                "link": "https://openai.com/policies/api-data-usage-policies",
+            },
             {
-            "platform": enums.EmbeddingPlatform.OPENAI.value,
-            "terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the OpenAI API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
-            "link": "https://openai.com/policies/api-data-usage-policies"
+                "platform": enums.EmbeddingPlatform.PYTHON.value,
+                "terms": None,
+                "link": None,
             },
             {
-            "platform": enums.EmbeddingPlatform.PYTHON.value,
-            "terms": None,
-            "link": None
+                "platform": enums.EmbeddingPlatform.AZURE.value,
+                "terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the Azure API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
+                "link": "https://www.microsoft.com/en-us/legal/terms-of-use",
             },
-        ]
+        ]