wordlift-sdk 2.7.1__py3-none-any.whl

wordlift_sdk/kg/entity_store.py

@@ -0,0 +1,94 @@
+import logging
+from typing import List, AsyncGenerator
+
+import pandas as pd
+from gql import Client, gql
+from pandas.core.interchange.dataframe_protocol import DataFrame
+from tenacity import AsyncRetrying, stop_after_attempt, wait_fixed, RetryError
+
+from wordlift_sdk.kg import Entity
+
+logger = logging.getLogger(__name__)
+
+
+class EntityStore:
+    _gql_client: Client
+
+    def __init__(self, gql_client: Client):
+        self._gql_client = gql_client
+
+    async def url_id(self, url_list: List[str] = None) -> AsyncGenerator[Entity, None]:
+        # the query.
+        query = gql(
+            """
+            query entities_url_id($urls: [String]!) {
+                entities(query: { urlConstraint: { in: $urls } } ) {
+                    url: string(name:"schema:url")
+                    id: iri
+                }
+            }
+            """
+        )
+
+        # the variables.
+        values = {"urls": url_list}
+
+        try:
+            async for attempt in AsyncRetrying(stop=stop_after_attempt(3), wait=wait_fixed(2)):
+                with attempt:
+                    logger.debug(
+                        'Loading data from GraphQL with attempt %d', attempt.retry_state.attempt_number)
+                    response = await self._gql_client.execute_async(query, variable_values=values)
+                    for item in response['entities']:
+                        yield Entity(item)
+        except RetryError as e:
+            logger.error('Error loading data from GraphQL', exc_info=True)
+
+    async def url_iri(self, url_list: List[str] = None) -> AsyncGenerator[Entity, None]:
+        # the query.
+        query = gql(
+            """
+            query entities_url_iri($urls: [String]!) {
+                entities(query: { urlConstraint: { in: $urls } } ) {
+                    url: string(name:"schema:url")
+                    iri
+                }
+            }
+            """
+        )
+
+        # the variables.
+        values = {"urls": url_list}
+
+        try:
+            async for attempt in AsyncRetrying(stop=stop_after_attempt(3), wait=wait_fixed(2)):
+                with attempt:
+                    logger.debug(
+                        'Loading data from GraphQL with attempt %d', attempt.retry_state.attempt_number)
+                    response = await self._gql_client.execute_async(query, variable_values=values)
+                    for item in response['entities']:
+                        yield Entity(item)
+        except RetryError as e:
+            logger.error('Error loading data from GraphQL', exc_info=True)
+
+    async def url_id_as_dataframe(self, url_list: List[str] = None) -> DataFrame:
+        """
+        Get the Entity URL ID maps as a Pandas dataframe.
+
+        :param url_list:
+        :return:
+        """
+        return pd.DataFrame.from_records(
+            data=[(entity.url, entity.iri) async for entity in self.url_id(url_list)],
+            columns=("url", "id"))
+
+    async def url_iri_as_dataframe(self, url_list: List[str] = None) -> DataFrame:
+        """
+        Get the Entity URL ID maps as a Pandas dataframe.
+
+        :param url_list:
+        :return:
+        """
+        return pd.DataFrame.from_records(
+            data=[(entity.url, entity.iri) async for entity in self.url_iri(url_list)],
+            columns=("url", "iri"))

wordlift_sdk/kg/entity_store_factory.py

@@ -0,0 +1,13 @@
+from gql import Client
+
+from wordlift_sdk.kg import EntityStore
+
+
+class EntityStoreFactory:
+    _gql_client: Client
+
+    def __init__(self, gql_client: Client):
+        self._gql_client = gql_client
+
+    def create(self):
+        return EntityStore(self._gql_client)

wordlift_sdk/kg/relation/relation_service.py

@@ -0,0 +1,78 @@
+import wordlift_client
+from wordlift_client import Configuration, VectorSearchQueriesApi, VectorSearchQueryRequest
+
+from wordlift_sdk.kg import Entity
+
+
+class RelationService:
+    _configuration: Configuration
+
+    def __init__(self, configuration: Configuration):
+        self._configuration = configuration
+
+    def get_relations(self, entity: Entity):
+        async with wordlift_client.ApiClient(self._configuration) as api_client:
+            # Search for related pages
+            search_api = VectorSearchQueriesApi(api_client)
+
+            related_request = VectorSearchQueryRequest(
+                query_url=entity_url,
+                similarity_top_k=100,
+                fields=["schema:url", "schema:headline", "ex-private:category", "ex-private:subCategory",
+                        "ex-private:subSubCategory", "ex-private:subSubSubCategory", "ex-private:location",
+                        "ex-private:pagesSection"]
+            )
+
+            try:
+                related_page = await search_api.create_query(vector_search_query_request=related_request)
+                print(f"Number of related items found: {len(related_page.items)}")
+            except Exception as e:
+                logger.error(f"Error during vector search: {e}")
+                return False
+
+            # Filter and re-rank results
+            filtered_results = []
+            for item in related_page.items:
+                item_url = safe_get_field(item, "schema:url")
+                if item_url == entity_url:
+                    print(f"Skipping original entity: {item_url}")
+                    continue
+
+                item_category = safe_get_field(item, "ex-private:category")
+                item_sub_category = safe_get_field(item, "ex-private:subCategory")
+                item_sub_sub_category = safe_get_field(item, "ex-private:subSubCategory")
+                item_sub_sub_sub_category = safe_get_field(item, "ex-private:subSubSubCategory")
+
+                print(f"Processing item: {item_url}")
+                print(
+                    f"Item categories: {item_category} > {item_sub_category} > {item_sub_sub_category} > {item_sub_sub_sub_category}")
+
+                if (item_category == main_category or
+                        item_sub_category == sub_category or
+                        item_sub_sub_category == sub_sub_category or
+                        item_sub_sub_sub_category == sub_sub_sub_category):
+
+                    score = item.score
+                    item_location = safe_get_field(item, "ex-private:location")
+                    item_pages_section = safe_get_field(item, "ex-private:pagesSection")
+
+                    print(f"Item location: {item_location}")
+                    print(f"Item pages section: {item_pages_section}")
+
+                    if item_location == location:
+                        score += 0.2
+                        print("Location match, score boosted")
+                    if item_pages_section == "Top":
+                        score += 0.1
+                        print("Top page, score boosted")
+
+                    filtered_results.append({
+                        "url": item_url,
+                        "headline": safe_get_field(item, "schema:headline"),
+                        "score": score
+                    })
+                    print(f"Item added to filtered results. Score: {score}")
+                else:
+                    print("Item categories don't match any of our category levels, skipped")
+
+                print("---")

wordlift_sdk/main.py

@@ -0,0 +1,7 @@
+from .container.application_container import ApplicationContainer
+
+
+async def run_kg_import_workflow():
+    application_container = ApplicationContainer()
+    workflow = await application_container.create_kg_import_workflow()
+    await workflow.run()