Merge pull request #373 from MITLibraries/USE-181-read-embeddings-with-tda

jonavellecuerdo · web-flow · commit b5202f88204e · 2025-12-15T10:07:54.000-05:00
Use 181 read embeddings with tda
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -0,0 +1,18 @@
+# CODEOWNERS file (from GitHub template at
+# https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners)
+# Each line is a file pattern followed by one or more owners.
+
+################################################################################
+# These owners will be the default owners for everything in the repo. This is commented
+# out in favor of using a team as the default (see below). It is left here as a comment
+# to indicate the primary expert for this code.
+# * @adamshire123
+
+# Teams can be specified as code owners as well. Teams should be identified in
+# the format @org/team-name. Teams must have explicit write access to the
+# repository.
+* @mitlibraries/dataeng
+
+# We set the senior engineer in the team as the owner of the CODEOWNERS file as
+# a layer of protection for unauthorized changes.
+/.github/CODEOWNERS @ghukill
diff --git a/.github/pull-request-template.md b/.github/pull-request-template.md
@@ -15,15 +15,5 @@ YES | NO
 ### What are the relevant tickets?
 - Include links to Jira Software and/or Jira Service Management tickets here.
 
-### Developer
-- [ ] All new ENV is documented in README
-- [ ] All new ENV has been added to staging and production environments
-- [ ] All related Jira tickets are linked in commit message(s)
-- [ ] Stakeholder approval has been confirmed (or is not needed)
-
-### Code Reviewer(s)
-- [ ] The commit message is clear and follows our guidelines (not just this PR message)
-- [ ] There are appropriate tests covering any new functionality
-- [ ] The provided documentation is sufficient for understanding any new functionality introduced
-- [ ] Any manual tests have been performed **or** provided examples have been verified
-- [ ] New dependencies are appropriate or there were no changes
+### Code review
+* Code review best practices are documented [here](https://mitlibraries.github.io/guides/collaboration/code_review.html) and you are encouraged to have a constructive dialogue with your reviewers about their preferences and expectations.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -24,6 +24,6 @@ repos:
         types: ["python"]
       - id: pip-audit
         name: pip-audit
-        entry: pipenv run pip-audit --ignore-vuln GHSA-4xh5-x5gv-qwph
+        entry: pipenv run pip-audit
         language: system
         pass_filenames: false
diff --git a/Makefile b/Makefile
@@ -1,4 +1,5 @@
 ### This is the Terraform-generated header for timdex-index-manager-dev ###
+.PHONY: test
 ECR_NAME_DEV:=timdex-index-manager-dev
 ECR_URL_DEV:=222053980223.dkr.ecr.us-east-1.amazonaws.com/timdex-index-manager-dev
 ### End of Terraform-generated header ###
@@ -28,7 +29,7 @@ update: install # Update Python dependencies
 ######################
 
 test: # Run tests and print a coverage report
-	pipenv run coverage run --source=tim -m pytest -vv
+	pipenv run coverage run --source=tim -m pytest -vv 
 	pipenv run coverage report -m
 
 coveralls: test # Write coverage data to an LCOV report
@@ -50,7 +51,7 @@ ruff: # Run 'ruff' linter and print a preview of errors
 	pipenv run ruff check .
 
 safety: # Check for security vulnerabilities and verify Pipfile.lock is up-to-date
-	pipenv run pip-audit --ignore-vuln GHSA-4xh5-x5gv-qwph
+	pipenv run pip-audit
 	pipenv verify
 
 lint-apply: black-apply ruff-apply # Apply changes with 'black' and resolve 'fixable errors' with 'ruff'
diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/config/opensearch_mappings.json b/config/opensearch_mappings.json
@@ -132,6 +132,9 @@
       "edition": {
         "type": "text"
       },
+      "embedding_full_record": {
+        "type": "rank_features"
+      },
       "file_formats": {
         "type": "keyword",
         "normalizer": "lowercase"
diff --git a/tests/fixtures/cassettes/bulk_update_embeddings_success.yaml b/tests/fixtures/cassettes/bulk_update_embeddings_success.yaml
@@ -0,0 +1,22 @@
+interactions:
+- request:
+    body: null
+    headers:
+      content-type:
+      - application/json
+      user-agent:
+      - opensearch-py/2.8.0 (Python 3.12.11)
+    method: GET
+    uri: http://localhost:9200/_cat/aliases?format=json
+  response:
+    body:
+      string: '[{"alias":"all-current","index":"libguides-2025-12-11t16-36-09","filter":"-","routing.index":"-","routing.search":"-","is_write_index":"-"},{"alias":"libguides","index":"libguides-2025-12-11t16-36-09","filter":"-","routing.index":"-","routing.search":"-","is_write_index":"-"},{"alias":"all-current","index":"test-index-2025-12-11t16-58-08","filter":"-","routing.index":"-","routing.search":"-","is_write_index":"-"},{"alias":"test-index","index":"test-index-2025-12-11t16-58-08","filter":"-","routing.index":"-","routing.search":"-","is_write_index":"-"},{"alias":".kibana","index":".kibana_1","filter":"-","routing.index":"-","routing.search":"-","is_write_index":"-"}]'
+    headers:
+      content-length:
+      - '671'
+      content-type:
+      - application/json; charset=UTF-8
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/fixtures/cassettes/opensearch/bulk_update_raises_bulk_operation_error_if_record_not_found.yaml b/tests/fixtures/cassettes/opensearch/bulk_update_raises_bulk_operation_error_if_record_not_found.yaml
@@ -0,0 +1,30 @@
+interactions:
+- request:
+    body: '{"update":{"_id":"i-am-not-found","_index":"test-index"}}
+
+      {"doc":{"timdex_record_id":"i-am-not-found","title":"Materials Science & Engineering
+      (UPDATED)"}}
+
+      '
+    headers:
+      Content-Length:
+      - '156'
+      content-type:
+      - application/json
+      user-agent:
+      - opensearch-py/2.8.0 (Python 3.12.11)
+    method: POST
+    uri: http://localhost:9200/_bulk
+  response:
+    body:
+      string: '{"took":9,"errors":true,"items":[{"update":{"_index":"test-index-2025-12-11t16-58-08","_id":"i-am-not-found","status":404,"error":{"type":"document_missing_exception","reason":"[i-am-not-found]:
+        document missing","index":"test-index-2025-12-11t16-58-08","shard":"0","index_uuid":"in04_JvQS5qqCvUXeZta_g"}}}]}'
+    headers:
+      content-length:
+      - '308'
+      content-type:
+      - application/json; charset=UTF-8
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/fixtures/cassettes/opensearch/bulk_update_updates_records.yaml b/tests/fixtures/cassettes/opensearch/bulk_update_updates_records.yaml
@@ -0,0 +1,51 @@
+interactions:
+- request:
+    body: '{"update":{"_id":"libguides:guides-175846","_index":"test-index"}}
+
+      {"doc":{"timdex_record_id":"libguides:guides-175846","title":"Materials Science
+      & Engineering (UPDATED)"}}
+
+      '
+    headers:
+      Content-Length:
+      - '174'
+      content-type:
+      - application/json
+      user-agent:
+      - opensearch-py/2.8.0 (Python 3.12.11)
+    method: POST
+    uri: http://localhost:9200/_bulk
+  response:
+    body:
+      string: '{"took":7,"errors":false,"items":[{"update":{"_index":"test-index-2025-12-11t16-58-08","_id":"libguides:guides-175846","_version":4,"result":"updated","_shards":{"total":2,"successful":1,"failed":0},"_seq_no":7,"_primary_term":1,"status":200}}]}'
+    headers:
+      content-length:
+      - '245'
+      content-type:
+      - application/json; charset=UTF-8
+    status:
+      code: 200
+      message: OK
+- request:
+    body: null
+    headers:
+      Content-Length:
+      - '0'
+      content-type:
+      - application/json
+      user-agent:
+      - opensearch-py/2.8.0 (Python 3.12.2)
+    method: POST
+    uri: http://localhost:9200/test-index/_refresh
+  response:
+    body:
+      string: '{"_shards":{"total":2,"successful":1,"failed":0}}'
+    headers:
+      content-length:
+      - '49'
+      content-type:
+      - application/json; charset=UTF-8
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/fixtures/dataset/data/embeddings/year=2025/month=12/day=10/629d15f4-84e4-4b32-92c1-1b1debd377fb-0.parquet b/tests/fixtures/dataset/data/embeddings/year=2025/month=12/day=10/629d15f4-84e4-4b32-92c1-1b1debd377fb-0.parquet
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -2,10 +2,11 @@
 import re
 from unittest.mock import patch
 
+import pytest
 from freezegun import freeze_time
 
 from tim.cli import main
-from tim.errors import BulkIndexingError
+from tim.errors import BulkIndexingError, BulkOperationError
 
 from .conftest import EXIT_CODES, my_vcr
 
@@ -274,6 +275,72 @@ def test_bulk_update_with_source_raise_bulk_indexing_error(
     )
 
 
+@pytest.mark.parametrize(
+    "bulk_update_return",
+    [
+        {"updated": 1, "errors": 0, "total": 1},
+        {"updated": 0, "errors": 1, "total": 1},
+    ],
+)
+@patch("tim.helpers.validate_bulk_cli_options")
+@patch("tim.opensearch.bulk_update")
+def test_bulk_update_embeddings_logs_complete(
+    mock_bulk_update,
+    mock_validate_bulk_cli_options,
+    bulk_update_return,
+    caplog,
+    monkeypatch,
+    runner,
+):
+    monkeypatch.delenv("TIMDEX_OPENSEARCH_ENDPOINT", raising=False)
+    mock_bulk_update.return_value = bulk_update_return
+    mock_validate_bulk_cli_options.return_value = "libguides"
+
+    result = runner.invoke(
+        main,
+        [
+            "bulk-update-embeddings",
+            "--source",
+            "libguides",
+            "--run-id",
+            "85cfe316-089c-4639-a5af-c861a7321493",
+            "tests/fixtures/dataset",
+        ],
+    )
+
+    assert result.exit_code == EXIT_CODES["success"]
+    assert (
+        f"Bulk update with embeddings complete: {json.dumps(mock_bulk_update())}"
+        in caplog.text
+    )
+
+
+@patch("tim.helpers.validate_bulk_cli_options")
+@patch("tim.opensearch.bulk_update")
+def test_bulk_update_embeddings_exit_bulk_operation_error(
+    mock_bulk_update, mock_validate_bulk_cli_options, caplog, monkeypatch, runner
+):
+    monkeypatch.delenv("TIMDEX_OPENSEARCH_ENDPOINT", raising=False)
+    mock_bulk_update.side_effect = BulkOperationError(
+        action="update", record="alma:0", index="index", error="exception"
+    )
+    mock_validate_bulk_cli_options.return_value = "libguides"
+
+    result = runner.invoke(
+        main,
+        [
+            "bulk-update-embeddings",
+            "--source",
+            "libguides",
+            "--run-id",
+            "85cfe316-089c-4639-a5af-c861a7321493",
+            "tests/fixtures/dataset",
+        ],
+    )
+    assert result.exit_code == EXIT_CODES["error"]
+    assert "Bulk update with embeddings failed" in caplog.text
+
+
 @patch("tim.opensearch.create_index")
 @patch("tim.opensearch.promote_index")
 @patch("tim.opensearch.get_index_aliases")
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -51,6 +51,17 @@ def test_generate_bulk_actions_delete():
     }
 
 
+def test_generate_bulk_actions_update():
+    records = [{"timdex_record_id": "12345"}]
+    actions = helpers.generate_bulk_actions("test-index", records, "update")
+    assert next(actions) == {
+        "_op_type": "update",
+        "_index": "test-index",
+        "_id": "12345",
+        "doc": {"timdex_record_id": "12345"},
+    }
+
+
 def test_generate_bulk_actions_invalid_action_raises_error():
     records = [{"timdex_record_id": "12345", "other_fields": "some_data"}]
     actions = helpers.generate_bulk_actions("test-index", records, "wrong")
diff --git a/tests/test_opensearch.py b/tests/test_opensearch.py
@@ -9,6 +9,7 @@
 from tim.config import PRIMARY_ALIAS
 from tim.errors import (
     AliasNotFoundError,
+    BulkOperationError,
     IndexExistsError,
     IndexNotFoundError,
 )
@@ -532,3 +533,34 @@ def test_bulk_delete_logs_error_if_record_not_found(
         "Record to delete 'i-am-not-found' was not found in index 'test-index'."
         in caplog.text
     )
+
+
+@my_vcr.use_cassette("opensearch/bulk_update_updates_records.yaml")
+def test_bulk_update_updates_records(test_opensearch_client):
+    updates = [
+        {
+            "timdex_record_id": "libguides:guides-175846",
+            "title": "Materials Science & Engineering (UPDATED)",
+        }
+    ]
+    assert tim_os.bulk_update(test_opensearch_client, "test-index", iter(updates)) == {
+        "updated": 1,
+        "errors": 0,
+        "total": 1,
+    }
+
+
+@my_vcr.use_cassette(
+    "opensearch/bulk_update_raises_bulk_operation_error_if_record_not_found.yaml"
+)
+def test_bulk_update_raises_bulk_operation_error_if_record_not_found(
+    test_opensearch_client,
+):
+    updates = [
+        {
+            "timdex_record_id": "i-am-not-found",
+            "title": "Materials Science & Engineering (UPDATED)",
+        }
+    ]
+    with pytest.raises(BulkOperationError):
+        tim_os.bulk_update(test_opensearch_client, "test-index", iter(updates))
diff --git a/tim/cli.py b/tim/cli.py
diff --git a/tim/helpers.py b/tim/helpers.py
diff --git a/tim/opensearch.py b/tim/opensearch.py
