core-semantic-search-app 1.0.0__tar.gz

core_semantic_search_app-1.0.0/LICENSE.md

@@ -0,0 +1,35 @@
+# NIST Software Licensing Statement
+
+NIST-developed software is provided by NIST as a public service.
+You may use, copy, and distribute copies of the software in any
+medium, provided that you keep intact this entire notice. You may
+improve, modify, and create derivative works of the software or
+any portion of the software, and you may copy and distribute such
+modifications or works. Modified works should carry a notice
+stating that you changed the software and should note the date
+and nature of any such change. Please explicitly acknowledge the
+National Institute of Standards and Technology as the source of
+the software.
+
+NIST-developed software is expressly provided "AS IS." NIST MAKES
+NO WARRANTY OF ANY KIND, EXPRESS, IMPLIED, IN FACT, OR ARISING BY
+OPERATION OF LAW, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
+WARRANTY OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
+NON-INFRINGEMENT, AND DATA ACCURACY. NIST NEITHER REPRESENTS NOR
+WARRANTS THAT THE OPERATION OF THE SOFTWARE WILL BE UNINTERRUPTED
+OR ERROR-FREE, OR THAT ANY DEFECTS WILL BE CORRECTED. NIST DOES
+NOT WARRANT OR MAKE ANY REPRESENTATIONS REGARDING THE USE OF THE
+SOFTWARE OR THE RESULTS THEREOF, INCLUDING BUT NOT LIMITED TO THE
+CORRECTNESS, ACCURACY, RELIABILITY, OR USEFULNESS OF THE
+SOFTWARE.
+
+You are solely responsible for determining the appropriateness of
+using and distributing the software and you assume all risks
+associated with its use, including but not limited to the risks
+and costs of program errors, compliance with applicable laws,
+damage to or loss of data, programs or equipment, and the
+unavailability or interruption of operation. This software is not
+intended to be used in any situation where a failure could cause
+risk of injury or damage to property. The software developed by
+NIST employees is not subject to copyright protection within the
+United States.

core_semantic_search_app-1.0.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include README.rst
+include requirements.txt
+include requirements.core.txt
+recursive-include core_semantic_search_app/static *
+recursive-include core_semantic_search_app/templates *

core_semantic_search_app-1.0.0/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.1
+Name: core_semantic_search_app
+Version: 1.0.0
+Summary: Semantic Search utils for the curator core project
+Home-page: https://github.com/usnistgov/core_semantic_search_app
+Author: NIST IT Lab
+Author-email: itl_inquiries@nist.gov
+Provides-Extra: sentence_transformers
+License-File: LICENSE.md
+
+========================
+Core Semantic Search App
+========================
+
+Semantic search for the curator core project.
+
+Quick start
+===========
+
+1. Add "core_semantic_search_app" to your INSTALLED_APPS setting
+----------------------------------------------------------------
+
+.. code:: python
+
+    INSTALLED_APPS = [
+      ...
+      'core_semantic_search_app',
+    ]
+
+
+2. Include the core_semantic_search_app URLconf in your project urls.py
+-----------------------------------------------------------------------
+
+.. code:: python
+
+    re_path(r'^semantic-search/', include('core_semantic_search_app.urls')),

core_semantic_search_app-1.0.0/README.rst

@@ -0,0 +1,26 @@
+========================
+Core Semantic Search App
+========================
+
+Semantic search for the curator core project.
+
+Quick start
+===========
+
+1. Add "core_semantic_search_app" to your INSTALLED_APPS setting
+----------------------------------------------------------------
+
+.. code:: python
+
+    INSTALLED_APPS = [
+      ...
+      'core_semantic_search_app',
+    ]
+
+
+2. Include the core_semantic_search_app URLconf in your project urls.py
+-----------------------------------------------------------------------
+
+.. code:: python
+
+    re_path(r'^semantic-search/', include('core_semantic_search_app.urls')),

core_semantic_search_app-1.0.0/core_semantic_search_app/__init__.py

@@ -0,0 +1,4 @@
+""" Initialize app for core_semantic_search_app
+"""
+
+default_app_config = "core_semantic_search_app.apps.SemanticSearchAppConfig"

core_semantic_search_app-1.0.0/core_semantic_search_app/admin.py

@@ -0,0 +1,15 @@
+""" Url router for the administration site
+"""
+
+from django.contrib import admin
+
+from core_semantic_search_app.components.document.models import Document
+from core_semantic_search_app.components.model_settings.admin_site import (
+    CustomModelSettingsAdmin,
+)
+from core_semantic_search_app.components.model_settings.models import (
+    ModelSettings,
+)
+
+admin.site.register(Document)
+admin.site.register(ModelSettings, CustomModelSettingsAdmin)

core_semantic_search_app-1.0.0/core_semantic_search_app/apps.py

@@ -0,0 +1,54 @@
+""" Apps file for core_semantic_search_app
+"""
+
+import sys
+
+from django.apps import AppConfig
+from django.db.models.signals import post_save, post_delete
+
+from core_main_app.utils.databases.backend import uses_postgresql_backend
+from core_semantic_search_app.commons.exceptions import SemanticSearchError
+
+
+class SemanticSearchAppConfig(AppConfig):
+    """Core application settings"""
+
+    name = "core_semantic_search_app"
+
+    def ready(self):
+        """Run when the app is ready
+
+        Returns:
+
+        """
+        if "migrate" in sys.argv:
+            return
+
+        _check_settings()
+        _init_signals()
+
+
+def _check_settings():
+    """Check settings
+
+    Returns:
+
+    """
+    if not uses_postgresql_backend():
+        raise SemanticSearchError("PostgreSQL with Pgvector is required.")
+
+
+def _init_signals():
+    """Init Signals
+
+    Returns:
+
+    """
+    from core_main_app.components.data.models import Data
+    from core_semantic_search_app.components.data.watch import (
+        post_save_data,
+        post_delete_data,
+    )
+
+    post_save.connect(post_save_data, sender=Data)
+    post_delete.connect(post_delete_data, sender=Data)

core_semantic_search_app-1.0.0/core_semantic_search_app/commons/exceptions.py

@@ -0,0 +1,9 @@
+""" Semantic Search Exceptions
+"""
+
+
+class SemanticSearchError(Exception):
+    """Exception raised by the Semantic Search app."""
+
+    def __init__(self, message):
+        self.message = message

core_semantic_search_app-1.0.0/core_semantic_search_app/components/data/watch.py

@@ -0,0 +1,30 @@
+""" Watchers for the data collection
+"""
+
+from core_semantic_search_app.components.document.api import (
+    index_documents_from_data,
+    delete_documents_with_data_id,
+)
+
+
+def post_save_data(sender, instance, **kwargs):
+    """Method executed after saving a Data object.
+    Args:
+        sender: Class.
+        instance: Data object.
+        **kwargs: Args.
+
+    """
+
+    index_documents_from_data(data=instance)
+
+
+def post_delete_data(sender, instance, **kwargs):
+    """Method executed after deleting a Data object.
+    Args:
+        sender: Class.
+        instance: Data object.
+        **kwargs: Args.
+
+    """
+    delete_documents_with_data_id(instance.id)

core_semantic_search_app-1.0.0/core_semantic_search_app/components/document/api.py

@@ -0,0 +1,279 @@
+""" Document API
+"""
+
+import logging
+import re
+
+from django.db.models import OuterRef, Exists, IntegerField
+from django.db.models.fields.json import KeyTextTransform
+from django.db.models.functions import Cast
+from pgvector.django import CosineDistance, L2Distance, MaxInnerProduct
+
+from core_main_app.access_control.api import has_perm_administration
+from core_main_app.access_control.decorators import access_control
+from core_main_app.commons.exceptions import ApiError
+from core_main_app.components.data import api as data_api
+from core_main_app.components.workspace import api as workspace_api
+from core_semantic_search_app import tasks as semantic_search_tasks
+from core_semantic_search_app.components.document.models import Document
+from core_semantic_search_app.components.model_settings.models import (
+    ModelSettings,
+)
+from core_semantic_search_app.settings import (
+    INSTALLED_APPS,
+)
+from core_semantic_search_app.utils.chunking_utils import chunk_json_dict
+
+if "core_linked_records_app" in INSTALLED_APPS:
+    from core_linked_records_app.system.pid_path import (
+        api as system_pid_path_api,
+    )
+    from core_linked_records_app.utils.dict import (
+        is_dot_notation_in_dictionary,
+        get_value_from_dot_notation,
+    )
+logger = logging.getLogger(__name__)
+
+VECTOR_FUNCTIONS = {
+    "cosine_similarity": CosineDistance,
+    "l2_distance": L2Distance,
+    "max_inner_product": MaxInnerProduct,
+}
+
+
+# NOTE: called from task
+def generate_documents_from_data(data):
+    """Convert data to documents
+
+    Args:
+        data:
+
+    Returns:
+
+    """
+    data_pid = None
+    if "core_linked_records_app" in INSTALLED_APPS:
+        try:
+            pid_path = system_pid_path_api.get_pid_path_by_template(
+                data.template,
+            ).path
+
+            # If the pid_path does not exist in the document, exit early and return None
+            data_dict = data.get_dict_content()
+            if is_dot_notation_in_dictionary(data_dict, pid_path):
+                data_pid = get_value_from_dot_notation(data_dict, pid_path)
+        except Exception as e:
+            logger.error(
+                f"Unable to get data PID during data to document conversion: {str(e)}"
+            )
+    model_settings = ModelSettings.get()
+
+    target_keys = (
+        model_settings.document_index_fields
+        if model_settings.document_index_strategy == "VALUES"
+        else None
+    )
+
+    text_chunks = chunk_json_dict(
+        json_dict=data.get_dict_content(),
+        chunk_size=model_settings.sliding_window_chunk_length,
+        chunk_overlap=model_settings.sliding_window_chunk_overlap,
+        target_keys=target_keys,
+    )
+
+    documents = list()
+    for text_chunk in text_chunks:
+        document = Document(
+            content=text_chunk,
+            meta={
+                "title": data.title,
+                "data_id": data.id,
+                "data_pid": data_pid,
+            },
+        )
+        documents.append(document)
+    return documents
+
+
+# NOTE: called from data watch
+def delete_documents_with_data_id(data_id):
+    """Delete documents with given data id
+
+    Args:
+        data_id:
+
+    Returns:
+
+    """
+    # Get any existing documents for this data
+    documents = _get_documents_by_data_id(data_id=data_id)
+    # Delete documents
+    documents.delete()
+
+
+# NOTE: called from data watch
+def index_documents_from_data(data):
+    """Index documents extracted from a data
+
+    Args:
+        data:
+
+    Returns:
+
+    """
+    model_settings = ModelSettings.get()
+
+    if not model_settings or not model_settings.embedding_models.keys():
+        return
+
+    # Delete previous documents
+    delete_documents_with_data_id(data.id)
+
+    if not _check_template_name(
+        data, model_settings.document_index_template_name_filter
+    ):
+        return
+
+    # only deal with data in a workspace
+    if not data.workspace:
+        return
+
+    # only deal with public data
+    public_workspaces = workspace_api.get_all_public_workspaces().values_list(
+        "id", flat=True
+    )
+    if data.workspace.id not in public_workspaces:
+        return
+
+    # Start indexing of documents
+    semantic_search_tasks.write_documents.apply_async((data.id,))
+
+
+@access_control(has_perm_administration)
+def reindex(user):
+    """Reindex the full knowledge base
+
+    Returns:
+
+    """
+    # Get all currently indexed documents
+    all_docs = Document.objects.all()
+    # Delete them
+    all_docs.delete()
+
+    # Find all public workspaces
+    public_workspaces = workspace_api.get_all_public_workspaces().values_list(
+        "id", flat=True
+    )
+    # Get all public data
+    all_data = data_api.get_all(user).filter(workspace__in=public_workspaces)
+
+    # Get model settings from database
+    model_settings = ModelSettings.get()
+    # Go through all records
+    for data in all_data:
+        # Check if template name matches pattern
+        if not _check_template_name(
+            data, model_settings.document_index_template_name_filter
+        ):
+            continue
+        # Index documents in tasks
+        semantic_search_tasks.write_documents.apply_async((data.id,))
+
+
+def query(
+    query_embedding=None,
+    top_k=10,
+    threshold=0.8,
+    vector_function="cosine_similarity",
+    data_filters_qs=None,
+):
+    """Query the embedding
+
+    Args:
+        query_embedding:
+        top_k:
+        threshold:
+        data_filters_qs:
+        vector_function:
+
+    Returns:
+
+    """
+    if vector_function not in list(VECTOR_FUNCTIONS.keys()):
+        raise ApiError(
+            f"Vector function should be in: {list(VECTOR_FUNCTIONS.keys())}."
+        )
+    # Get all documents
+    queryset = Document.objects.all()
+
+    if data_filters_qs is not None:
+        # Pull data_id from meta field
+        data_id_expr = Cast(
+            KeyTextTransform("data_id", "meta"),
+            output_field=IntegerField(),
+        )
+        # Filter on extracted data_id
+        queryset = queryset.annotate(_data_id=data_id_expr).filter(
+            Exists(data_filters_qs.filter(id=OuterRef("_data_id")))
+        )
+
+    if query_embedding:
+        # Order documents by distance to query
+        # https://github.com/pgvector/pgvector-python?tab=readme-ov-file#django
+        queryset = queryset.order_by(
+            VECTOR_FUNCTIONS[vector_function]("embedding", query_embedding)
+        )
+
+        # Compute the score between query and documents
+        # https://github.com/pgvector/pgvector?tab=readme-ov-file#distances
+        if vector_function == "cosine_similarity":
+            queryset = queryset.annotate(
+                score=1 - CosineDistance("embedding", query_embedding)
+            )
+        elif vector_function == "max_inner_product":
+            queryset = queryset.annotate(
+                score=-1 * MaxInnerProduct("embedding", query_embedding)
+            )
+        elif vector_function == "l2_distance":
+            queryset = queryset.annotate(
+                score=L2Distance("embedding", query_embedding)
+            )
+
+        # filter score with provided threshold
+        if threshold:
+            if vector_function == "l2_distance":
+                queryset = queryset.filter(score__lt=threshold)
+            else:
+                queryset = queryset.filter(score__gt=threshold)
+
+    # Keep the top k results
+    if top_k:
+        queryset = queryset.all()[:top_k]
+
+    return queryset
+
+
+def _check_template_name(data, pattern):
+    """Check template name matches the pattern
+
+    Args:
+        data:
+        pattern:
+
+    Returns:
+
+    """
+    return re.search(pattern, data.template.version_manager.title)
+
+
+def _get_documents_by_data_id(data_id):
+    """Get documents with data id
+
+    Args:
+        data_id:
+
+    Returns:
+
+    """
+    return Document.objects.filter(meta__data_id=int(data_id))

core_semantic_search_app-1.0.0/core_semantic_search_app/components/document/models.py

@@ -0,0 +1,22 @@
+""" Document models
+"""
+
+from django.db import models
+from django.db.models import JSONField
+from pgvector.django import VectorField
+
+
+class Document(models.Model):
+    """Document Model"""
+
+    embedding = VectorField()
+    content = models.CharField(unique=False, max_length=4000)
+    meta = JSONField(default=dict)
+
+    def __str__(self):
+        """
+
+        Returns:
+
+        """
+        return self.meta.get("title", "Untitled")

core_semantic_search_app-1.0.0/core_semantic_search_app/components/model_settings/admin_site.py

@@ -0,0 +1,49 @@
+""" Custom admin site for the Model Settings model
+"""
+
+from django.contrib import admin
+from django.contrib import messages
+
+from core_semantic_search_app.components.document.api import reindex
+from core_semantic_search_app.components.model_settings.forms import (
+    ModelSettingsAdminForm,
+)
+
+
+@admin.action(description="Reindex the knowledge base")
+def reindex_action(model_admin, request, queryset):
+    """
+
+    Args:
+        model_admin:
+        request:
+        queryset:
+
+    Returns:
+
+    """
+    if not request.user.is_superuser:
+        model_admin.message_user(request, "Permission denied.", messages.ERROR)
+        return
+
+    reindex(request.user)
+
+
+class CustomModelSettingsAdmin(admin.ModelAdmin):
+    """CustomModelSettingsAdmin"""
+
+    form = ModelSettingsAdminForm
+    actions = [reindex_action]
+
+    def has_add_permission(self, request):
+        """Has add permission - only if doesn't exist
+
+        Args:
+            request:
+
+        Returns:
+
+        """
+        if self.model.objects.exists():
+            return False
+        return super().has_add_permission(request)

core_semantic_search_app-1.0.0/core_semantic_search_app/components/model_settings/api.py

@@ -0,0 +1,23 @@
+""" ModelSetting api
+"""
+
+import os
+
+
+def get_api_key(model_dict):
+    """Get API Key from dictionary
+
+    Args:
+        model_dict:
+
+    Returns:
+
+    """
+    # Check if api_key_env is set
+    api_key_env = model_dict.get("api_key_env")
+    # Get value of API Key from env
+    api_key = os.getenv(api_key_env) if api_key_env else None
+    # Get value from dict otherwise
+    api_key = model_dict.get("api_key", "no-key") if not api_key else api_key
+    # Return API Key
+    return api_key

core_semantic_search_app-1.0.0/core_semantic_search_app/components/model_settings/forms.py

@@ -0,0 +1,44 @@
+""" Model Settings forms
+"""
+
+import json
+
+from django import forms
+
+from core_semantic_search_app.components.model_settings.models import (
+    ModelSettings,
+)
+
+
+class PrettyJSONEncoder(json.JSONEncoder):
+    def __init__(self, *args, indent, sort_keys, **kwargs):
+        super().__init__(*args, indent=4, sort_keys=False, **kwargs)
+
+
+class ModelSettingsAdminForm(forms.ModelForm):
+
+    class Meta:
+        model = ModelSettings
+        fields = "__all__"
+        labels = {
+            "sliding_window_chunk_length": "Chunk size (characters)",
+            "sliding_window_chunk_overlap": "Chunk overlap (characters)",
+        }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.fields["embedding_models"].encoder = PrettyJSONEncoder
+        self.fields["embedding_models"].widget.attrs.update(
+            {
+                "placeholder": """{"modelName":{
+        "model": "modelName:version",
+        "base_url": "http://localhost:8080/v1",
+        "api_key": "",
+        "api_key_env": "",
+        "ssl_verify": true,
+        "proxies": {}
+    }
+}
+            """
+            }
+        )