Qdrant filter integrations (#142)

* Filtered attributes from config request when creating an embedding

* Filter params for neural search

* Query for getting unique values by attribute id

* Submodules change

* Update embedding payloads mutation

* Endpoint for unique values changed

* Removed unused code

* Removed unsed code

* PR comments

* Import export field for filter attributes

* Submodules updated

* PR comments

* PR comments

Author: lumburovskalina

controller/embedding/connector.py

@@ -10,9 +10,8 @@ def request_listing_recommended_encoders() -> Any:
     url = f"{BASE_URI}/classification/recommend/TEXT"  # TODO does here have to be a data type?
     return service_requests.get_call_or_raise(url)
 
-def request_embedding(
-    project_id: str, embedding_id: str
-) -> Any:
+
+def request_embedding(project_id: str, embedding_id: str) -> Any:
     url = f"{BASE_URI}/embed"
     data = {
         "project_id": str(project_id),

controller/embedding/manager.py

@@ -185,6 +185,7 @@ def __recreate_embedding(project_id: str, embedding_id: str) -> Embedding:
         model=old_embedding_item.model,
         platform=old_embedding_item.platform,
         api_token=old_embedding_item.api_token,
+        filter_attributes=old_embedding_item.filter_attributes,
         additional_data=old_embedding_item.additional_data,
         with_commit=False,
     )
@@ -212,3 +213,17 @@ def __recreate_embedding(project_id: str, embedding_id: str) -> Embedding:
     connector.request_deleting_embedding(project_id, old_id)
     daemon.run(connector.request_embedding, project_id, new_embedding_item.id)
     return new_embedding_item
+
+
+def update_embedding_payload(
+    project_id: str, embedding_id: str, filter_attributes: List[str]
+) -> None:
+    notification.send_organization_update(
+        project_id=project_id,
+        message=f"upload_embedding_payload:{str(embedding_id)}:start",
+    )
+    embedding.update_embedding_filter_attributes(
+        project_id, embedding_id, filter_attributes, with_commit=True
+    )
+    connector.request_deleting_embedding(project_id, embedding_id)
+    connector.request_tensor_upload(project_id, embedding_id)

controller/record/manager.py

@@ -1,8 +1,13 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 
 from graphql_api.types import ExtendedSearch
 from submodules.model import Record, Attribute
-from submodules.model.business_objects import general, record, user_session, embedding
+from submodules.model.business_objects import (
+    record,
+    user_session,
+    embedding,
+    attribute,
+)
 from service.search import search
 
 from controller.record import neural_search_connector
@@ -19,9 +24,10 @@ def get_records_by_similarity_search(
     user_id: str,
     embedding_id: str,
     record_id: str,
+    att_filter: Optional[List[Dict[str, Any]]] = None,
 ) -> ExtendedSearch:
     record_ids = neural_search_connector.request_most_similar_record_ids(
-        project_id, embedding_id, record_id, 100
+        project_id, embedding_id, record_id, 100, att_filter
     )
     if not len(record_ids):
         record_ids = [record_id]
@@ -102,3 +108,7 @@ def __reupload_embeddings(project_id: str) -> None:
     embeddings = embedding.get_finished_embeddings(project_id)
     for e in embeddings:
         embedding_manager.request_tensor_upload(project_id, str(e.id))
+
+
+def get_unique_values_by_attributes(project_id: str) -> Dict[str, List[str]]:
+    return attribute.get_unique_values_by_attributes(project_id)

controller/record/neural_search_connector.py

@@ -1,16 +1,19 @@
 import os
-from typing import List
+from typing import List, Optional, Dict, Any
 
 from util import service_requests
 
 BASE_URI = os.getenv("NEURAL_SEARCH")
 
 
 def request_most_similar_record_ids(
-    project_id: str, embedding_id: str, record_id: str, limit: int
+    project_id: str,
+    embedding_id: str,
+    record_id: str,
+    limit: int,
+    att_filter: Optional[List[Dict[str, Any]]] = None,
 ) -> List[str]:
     url = f"{BASE_URI}/most_similar?project_id={project_id}&embedding_id={embedding_id}&record_id={record_id}&limit={limit}"
 
-    # changed from get to post so we can send the filter -> however currently filter isn't part of the prototype so None
-    result = service_requests.post_call_or_raise(url, None)
+    result = service_requests.post_call_or_raise(url, att_filter)
     return result

controller/transfer/project_transfer_manager.py

@@ -397,6 +397,7 @@ def __transform_embedding_by_name(embedding_name: str):
                 model=embedding_item.get(
                     "model",
                 ),
+                filter_attributes=embedding_item.get("filter_attributes"),
                 additional_data=embedding_item.get(
                     "additional_data",
                 ),
@@ -929,7 +930,6 @@ def __post_processing_import_threaded(
                 f"Tokenization finished, continue with embedding handling of project {project_id}"
             )
             break
-
     if not data.get(
         "embedding_tensors_data",
     ):
@@ -1177,6 +1177,7 @@ def get_project_export_dump(
             "finished_at": embedding_item.finished_at,
             "platform": embedding_item.platform,
             "model": embedding_item.model,
+            "filter_attributes": embedding_item.filter_attributes,
             "additional_data": embedding_item.additional_data,
         }
         for embedding_item in embeddings

graphql_api/mutation/embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, List, Optional
 from controller.auth import manager as auth
 from controller.embedding import manager
 from controller.auth.manager import get_user_by_info
@@ -31,6 +31,8 @@ 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")
+        filter_attributes = config.get("filterAttributes")
+
         additional_data = None
         if config.get("base") is not None:
             additional_data = {
@@ -39,10 +41,6 @@ def mutate(self, info, project_id: str, attribute_id: str, config: Dict[str, Any
                 "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,
@@ -58,7 +56,7 @@ def mutate(self, info, project_id: str, attribute_id: str, config: Dict[str, Any
                 "api_token": api_token,
                 "terms_text": terms_text,
                 "terms_accepted": terms_accepted,
-                "filter_attributes": relevant_attribute_list,
+                "filter_attributes": filter_attributes,
                 "additional_data": additional_data,
             },
         )
@@ -85,6 +83,31 @@ def mutate(self, info, project_id: str, embedding_id: str):
         return DeleteEmbedding(ok=True)
 
 
+class UpdateEmbeddingPayload(graphene.Mutation):
+    class Arguments:
+        project_id = graphene.ID(required=True)
+        embedding_id = graphene.ID(required=True)
+        filter_attributes = graphene.JSONString(required=False)
+
+    ok = graphene.Boolean()
+
+    def mutate(
+        self,
+        info,
+        project_id: str,
+        embedding_id: str,
+        filter_attributes: Optional[List[str]] = None,
+    ):
+        auth.check_demo_access(info)
+        auth.check_project_access(info, project_id)
+        manager.update_embedding_payload(project_id, embedding_id, filter_attributes)
+        notification.send_organization_update(
+            project_id, f"embedding_updated:{embedding_id}"
+        )
+        return UpdateEmbeddingPayload(ok=True)
+
+
 class EmbeddingMutation(graphene.ObjectType):
     create_embedding = CreateEmbedding.Field()
     delete_embedding = DeleteEmbedding.Field()
+    update_embedding_payload = UpdateEmbeddingPayload.Field()

graphql_api/query/record.py

@@ -45,6 +45,7 @@ class RecordQuery(graphene.ObjectType):
         project_id=graphene.ID(required=True),
         embedding_id=graphene.ID(required=True),
         record_id=graphene.ID(required=True),
+        att_filter=graphene.JSONString(required=False),
     )
 
     tokenize_record = graphene.Field(
@@ -59,6 +60,10 @@ class RecordQuery(graphene.ObjectType):
         record_ids=graphene.List(graphene.ID, required=True),
     )
 
+    unique_values_by_attributes = graphene.Field(
+        graphene.JSONString, project_id=graphene.ID(required=True)
+    )
+
     def resolve_all_records(self, info, project_id: str) -> List[Record]:
         auth.check_project_access(info, project_id)
         return manager.get_all_records(project_id)
@@ -101,13 +106,18 @@ def resolve_search_records_extended(
         )
 
     def resolve_search_records_by_similarity(
-        self, info, project_id: str, embedding_id: str, record_id: str
+        self,
+        info,
+        project_id: str,
+        embedding_id: str,
+        record_id: str,
+        att_filter: Optional[List[Dict[str, Any]]] = None,
     ) -> ExtendedSearch:
         auth.check_demo_access(info)
         auth.check_project_access(info, project_id)
         user_id = auth.get_user_by_info(info).id
         return manager.get_records_by_similarity_search(
-            project_id, user_id, embedding_id, record_id
+            project_id, user_id, embedding_id, record_id, att_filter
         )
 
     def resolve_tokenize_record(self, info, record_id: str) -> TokenizedRecord:
@@ -127,3 +137,8 @@ def resolve_record_comments(
         auth.check_project_access(info, project_id)
         user_id = auth.get_user_id_by_info(info)
         return comment_manager.get_record_comments(project_id, user_id, record_ids)
+
+    def resolve_unique_values_by_attributes(self, info, project_id: str) -> str:
+        auth.check_demo_access(info)
+        auth.check_project_access(info, project_id)
+        return manager.get_unique_values_by_attributes(project_id)