refactor (embedding utils): refactor add_embeddings_to_source (#36)

* refactor (embedding utils): refactor add_embeddings_to_source

* refactor (embedding utils): use list comprehension

* refactor (transform): simplify structure of named tuple

* refactor (embedding utils): use named tuple instead of tuple

* refactor (embedding utils): use list comprehension instead of map

* refactor (embedding utils): use list comprehension instead of map

* test (embedding utils): add test for preprocess_batch

* refactor (embedding utils): simplify code

* refactor (embedding utils): split list comprehensions into multiple lines

Author: tobiasschweizer

src/tasks.py

@@ -13,7 +13,7 @@
 from config.opensearch_config import OpenSearchConfig
 from utils.queue_utils import HarvestEventQueue
 from utils.embedding_utils import preprocess_batch, add_embeddings_to_source, SourceWithEmbeddingText, \
-    get_embedding_text_from_fields
+    get_embedding_text_from_fields, OpenSearchSourceWithEmbedding
 from utils import normalize_datacite_json
 from typing import Any
 from celery.utils.log import get_task_logger
@@ -164,7 +164,6 @@ def transform_batch(self: Any, batch: list[HarvestEventQueue], index_name: str)
                 validate(instance=normalized_record, schema=self.schema)
                 normalized.append(SourceWithEmbeddingText(src=normalized_record,
                                                           textToEmbed=get_embedding_text_from_fields(normalized_record),
-                                                          file=Path(''),
                                                           event=harvest_event
                                                           ))
 
@@ -182,10 +181,10 @@ def transform_batch(self: Any, batch: list[HarvestEventQueue], index_name: str)
 
         try:
             logger.info(f'About to Calculate embeddings for {len(normalized)}')
-            src_with_emb: list[tuple[dict[str, Any], SourceWithEmbeddingText]] = add_embeddings_to_source(normalized,
+            src_with_emb: list[OpenSearchSourceWithEmbedding] = add_embeddings_to_source(normalized,
                                                                                        self.embedding_transformer)
             logger.info(f'Calculated embeddings for {len(src_with_emb)}')
-            preprocessed = preprocess_batch(list(map(lambda el: el[0], src_with_emb)), index_name)
+            preprocessed = preprocess_batch([src_with_emb_ele.src for src_with_emb_ele in src_with_emb], index_name)
         except Exception as e:
             logger.error(f'Could not calculate embeddings: {e}')
             raise e
@@ -201,25 +200,20 @@ def transform_batch(self: Any, batch: list[HarvestEventQueue], index_name: str)
             for rec in src_with_emb:
                 # write to records table
 
-                if rec[1].event is None:
-                    raise ValueError(f'Original HarvestEvent not found')
-
-                #logger.info(rec[1].event.record_identifier)
-
-                record_identifier = rec[1].event.record_identifier
-                datestamp = rec[1].event.datestamp
-                repository_id = rec[1].event.repository_id
-                endpoint_id = rec[1].event.endpoint_id
+                record_identifier = rec.harvest_event.record_identifier
+                datestamp = rec.harvest_event.datestamp
+                repository_id = rec.harvest_event.repository_id
+                endpoint_id = rec.harvest_event.endpoint_id
                 resource_type = 'Dataset' # TODO: get this information from record
-                title = rec[0]['titles'][0]['title']
-                xml = rec[1].event.xml
+                title = rec.src['titles'][0]['title']
+                xml = rec.harvest_event.xml
                 protocol = 'OAI-PMH'
-                doi = rec[0].get('doi')
-                url = rec[0].get('url')
-                embeddings = rec[0]['emb']
-                datacite_json = json.dumps({**rec[0], 'emb': None})
+                doi = rec.src.get('doi')
+                url = rec.src.get('url')
+                embeddings = rec.src['emb']
+                datacite_json = json.dumps({**rec.src, 'emb': None})
                 opensearch_synced = True
-                additional_metadata = rec[1].event.additional_metadata
+                additional_metadata = rec.harvest_event.additional_metadata
 
                 # https://neon.com/postgresql/postgresql-tutorial/postgresql-upsert
                 cur.execute("""
@@ -282,7 +276,7 @@ def transform_batch(self: Any, batch: list[HarvestEventQueue], index_name: str)
                     UPDATE harvest_events 
                     SET error_message = NULL
                     WHERE id = %s  
-                    """, [rec[1].event.id]
+                    """, [rec.harvest_event.id]
                 )
 
         except BulkIndexError as e:

src/transform.py

@@ -62,11 +62,13 @@ class AdditionalMetadataParams(BaseModel):
     endpoint: str
     protocol: str
 
+
 class HarvestParams(BaseModel):
     metadata_prefix: str
     set: Optional[list[str]]
     additional_metadata_params: Optional[AdditionalMetadataParams]
 
+
 class EndpointConfig(BaseModel):
     name: str
     harvest_url: str
@@ -82,8 +84,8 @@ class Config(BaseModel):
 class HarvestEventCreateRequest(BaseModel):
     record_identifier: str
     datestamp: datetime
-    raw_metadata: str # XML
-    additional_metadata: Optional[str] = None # XML or JSON (stringified)
+    raw_metadata: str  # XML
+    additional_metadata: Optional[str] = None  # XML or JSON (stringified)
     harvest_url: str
     repo_code: str
     harvest_run_id: str
@@ -97,6 +99,7 @@ class HarvestEventCreateResponse(BaseModel):
 class HarvestRunCreateRequest(BaseModel):
     harvest_url: str
 
+
 class HarvestRunGetResponse(BaseModel):
     id: Optional[str] = Field(None, description='ID of the harvest run')
     status: Optional[str] = Field(None, description='Status of the harvest run: open|closed|failed')
@@ -121,7 +124,8 @@ class HarvestRunCloseResponse(BaseModel):
 
 
 def get_latest_harvest_run_in_db(harvest_url: str) -> HarvestRunGetResponse:
-    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address, password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
+    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address,
+                         password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
 
         cur = conn.cursor()
 
@@ -141,15 +145,16 @@ def get_latest_harvest_run_in_db(harvest_url: str) -> HarvestRunGetResponse:
         else:
             return HarvestRunGetResponse(id=None, status=None)
 
+
 def create_harvest_run_in_db(harvest_url: str) -> HarvestRunCreateResponse:
     """
     Creates a new entry in harvest_runs and returns its id.
 
     :param harvest_url: The new entry to be created.
     """
 
-    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address, password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
-
+    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address,
+                         password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
         cur = conn.cursor()
         # TODO: only allow one open harvest run per endpoint
         # TODO check (in one transaction):
@@ -192,15 +197,19 @@ def create_harvest_run_in_db(harvest_url: str) -> HarvestRunCreateResponse:
             id=str(new_harvest_run['id']),
             from_date=new_harvest_run['from_date'],
             until_date=new_harvest_run['until_date'],
-            endpoint_config=EndpointConfig(name=new_harvest_run['name'], harvest_url=new_harvest_run['harvest_url'], code=new_harvest_run['code'], protocol=new_harvest_run['protocol'],
-                                   harvest_params=HarvestParams(metadata_prefix=new_harvest_run['harvest_params'].get('metadata_prefix'), set=new_harvest_run['harvest_params'].get('set'), additional_metadata_params=new_harvest_run['harvest_params'].get('additional_metadata_params')))
+            endpoint_config=EndpointConfig(name=new_harvest_run['name'], harvest_url=new_harvest_run['harvest_url'],
+                                           code=new_harvest_run['code'], protocol=new_harvest_run['protocol'],
+                                           harvest_params=HarvestParams(
+                                               metadata_prefix=new_harvest_run['harvest_params'].get('metadata_prefix'),
+                                               set=new_harvest_run['harvest_params'].get('set'),
+                                               additional_metadata_params=new_harvest_run['harvest_params'].get(
+                                                   'additional_metadata_params')))
         )
 
 
 def close_harvest_run_in_db(harvest_run: HarvestRunCloseRequest) -> HarvestRunCloseResponse:
     with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address,
                          password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
-
         cur = conn.cursor()
 
         state = 'closed' if harvest_run.success else 'failed'
@@ -224,15 +233,16 @@ def close_harvest_run_in_db(harvest_run: HarvestRunCloseRequest) -> HarvestRunCl
 
         return HarvestRunCloseResponse(id=harvest_run.id)
 
+
 def create_harvest_event_in_db(harvest_event: HarvestEventCreateRequest) -> HarvestEventCreateResponse:
     """
     Creates a record in table harvest_events
 
     :param harvest_event: The new record to be created.
     """
 
-    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address, password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
-
+    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address,
+                         password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
         cur = conn.cursor()
 
         cur.execute("""
@@ -260,8 +270,9 @@ def create_harvest_event_in_db(harvest_event: HarvestEventCreateRequest) -> Harv
                             (SELECT id FROM harvest_runs WHERE id = %s and status = 'open'),
                             %s
                             );
-                        """, (harvest_event.record_identifier, harvest_event.datestamp, harvest_event.raw_metadata, harvest_event.additional_metadata, harvest_event.repo_code, harvest_event.harvest_url, 'OAI-PMH', 'XML', harvest_event.harvest_run_id, harvest_event.is_deleted))
-
+                        """, (harvest_event.record_identifier, harvest_event.datestamp, harvest_event.raw_metadata,
+                              harvest_event.additional_metadata, harvest_event.repo_code, harvest_event.harvest_url,
+                              'OAI-PMH', 'XML', harvest_event.harvest_run_id, harvest_event.is_deleted))
 
         cur.execute("""
         SELECT id 
@@ -301,10 +312,14 @@ def get_config_from_db() -> list[EndpointConfig]:
 JOIN repositories r ON e.repository_id = r.id
             """)
             for doc in cur.fetchall():
-
                 endpoints.append(
-                    EndpointConfig(name=doc['name'], harvest_url=doc['harvest_url'], code=doc['code'], protocol=doc['protocol'],
-                                   harvest_params=HarvestParams(metadata_prefix=doc['harvest_params'].get('metadata_prefix'), set=doc['harvest_params'].get('set'), additional_metadata_params=doc['harvest_params'].get('additional_metadata_params'))))
+                    EndpointConfig(name=doc['name'], harvest_url=doc['harvest_url'], code=doc['code'],
+                                   protocol=doc['protocol'],
+                                   harvest_params=HarvestParams(
+                                       metadata_prefix=doc['harvest_params'].get('metadata_prefix'),
+                                       set=doc['harvest_params'].get('set'),
+                                       additional_metadata_params=doc['harvest_params'].get(
+                                           'additional_metadata_params'))))
 
         return endpoints
     except JSONDecodeError as e:
@@ -331,7 +346,8 @@ def create_jobs_in_queue(harvest_run_id: str) -> int:
 
     logger.info(f'Preparing jobs')
 
-    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address, password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
+    with psycopg.connect(dbname=postgres_config.user, user=postgres_config.user, host=postgres_config.address,
+                         password=postgres_config.password, port=postgres_config.port, row_factory=dict_row) as conn:
 
         cur = conn.cursor()
 
@@ -363,13 +379,15 @@ def create_jobs_in_queue(harvest_run_id: str) -> int:
             """, (harvest_run_id, limit, offset))
 
             for doc in cur.fetchall():
-
                 # https://www.psycopg.org/psycopg3/docs/basic/adapt.html#uuid-adaptation
                 # https://docs.python.org/3/library/uuid.html#uuid.UUID
                 # str(uuid) returns a string in the form 12345678-1234-5678-1234-567812345678 where the 32 hexadecimal digits represent the UUID.
                 batch.append(
                     HarvestEventQueue(id=str(doc['id']), xml=doc['record'], repository_id=str(doc['repository_id']),
-                                 endpoint_id=str(doc['endpoint_id']), record_identifier=doc['record_identifier'], code=doc['code'], harvest_url=doc['harvest_url'], additional_metadata=doc['additional_metadata'], is_deleted=doc['is_deleted'], datestamp=doc['datestamp'].strftime('%Y-%m-%d %H:%M:%S.%f%z'))
+                                      endpoint_id=str(doc['endpoint_id']), record_identifier=doc['record_identifier'],
+                                      code=doc['code'], harvest_url=doc['harvest_url'],
+                                      additional_metadata=doc['additional_metadata'], is_deleted=doc['is_deleted'],
+                                      datestamp=doc['datestamp'].strftime('%Y-%m-%d %H:%M:%S.%f%z'))
                 )
 
             if len(batch) == 0:
@@ -385,13 +403,15 @@ def create_jobs_in_queue(harvest_run_id: str) -> int:
             offset += limit
             # will be false if query returned fewer results than limit
             fetch = len(batch) == limit
-            #fetch = False
+            # fetch = False
             batch = []
 
     return tasks
 
+
 @app.get('/index', tags=['index'])
-def init_index(harvest_run_id: str = Query(default=None, description='Id of the harvest run to be indexed')) -> IndexGetResponse:
+def init_index(
+    harvest_run_id: str = Query(default=None, description='Id of the harvest run to be indexed')) -> IndexGetResponse:
     # this long-running method is synchronous and runs in an external threadpool, see https://fastapi.tiangolo.com/async/#path-operation-functions
     # this way, it does not block the server
     try:
@@ -422,24 +442,31 @@ def get_config() -> Config:
 @app.post('/harvest_event', tags=['harvest_event'], summary='Register a new harvest event')
 def create_harvest_event(harvest_event: HarvestEventCreateRequest) -> HarvestEventCreateResponse:
     try:
-        #logger.debug(harvest_event)
+        # logger.debug(harvest_event)
         return create_harvest_event_in_db(harvest_event)
     except psycopg_errors.UniqueViolation as e:
         logger.exception(f'Harvest event could not be created for given harvest run')
-        raise HTTPException(status_code=400, detail='Harvest event could not be created for the given harvest run because the record identifier already exists.')
+        raise HTTPException(status_code=400,
+                            detail='Harvest event could not be created for the given harvest run because the record identifier already exists.')
     except Exception as e:
         logger.exception(f'An error occurred when creating harvest event: {e}')
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get('/harvest_run', tags=['harvest_run'], summary='Get id and status of the latest harvest run for a given endpoint.', description='If no harvest run exists for the given endpoint, id and status will be null in the response.')
-def get_harvest_run(harvest_url: str = Query(default=None, description='harvest url of the endpoint')) -> HarvestRunGetResponse:
+
+@app.get('/harvest_run', tags=['harvest_run'],
+         summary='Get id and status of the latest harvest run for a given endpoint.',
+         description='If no harvest run exists for the given endpoint, id and status will be null in the response.')
+def get_harvest_run(
+    harvest_url: str = Query(default=None, description='harvest url of the endpoint')) -> HarvestRunGetResponse:
     try:
         return get_latest_harvest_run_in_db(harvest_url)
     except Exception as e:
         logger.exception(f'An error occurred when getting harvest run: {e}')
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post('/harvest_run', tags=['harvest_run'], summary='Create a new havest run for a given endpoint.', description='A new harvest run can only be created if no other open harvest run exists for the same endpoint.')
+
+@app.post('/harvest_run', tags=['harvest_run'], summary='Create a new havest run for a given endpoint.',
+          description='A new harvest run can only be created if no other open harvest run exists for the same endpoint.')
 def create_harvest_run(harvest_run: HarvestRunCreateRequest) -> HarvestRunCreateResponse:
     try:
         logger.debug(harvest_run)
@@ -451,11 +478,11 @@ def create_harvest_run(harvest_run: HarvestRunCreateRequest) -> HarvestRunCreate
         logger.exception(f'An error occurred when creating harvest event: {e}')
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @app.put('/harvest_run', tags=['harvest_run'], summary='Close an open harvest run for a given endpoint.')
 def close_harvest_run(harvest_run: HarvestRunCloseRequest) -> HarvestRunCloseResponse:
     try:
         return close_harvest_run_in_db(harvest_run)
     except Exception as e:
         logger.exception(f'An error occurred when closing harvest event: {e}')
         raise HTTPException(status_code=500, detail=str(e))
-