initial commit

.gitignore

@@ -0,0 +1,15 @@
+data/
+
+# python generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# venv
+.venv
+.old
+.envrc
+.env

.python-version

@@ -0,0 +1 @@
+3.12.2

README.md

@@ -0,0 +1,3 @@
+# semantic-catalogue
+
+Describe your project here.

pyproject.toml

@@ -0,0 +1,44 @@
+[project]
+name = "semantic-catalogue"
+version = "0.1.0"
+description = "Add your description here"
+authors = [
+    { name = "cjber", email = "cjberragan@gmail.com" }
+]
+dependencies = [
+    "requests>=2.31.0",
+    "langchain>=0.2.0",
+    "langchain-community>=0.2.0",
+    "langchain-openai>=0.1.7",
+    "langchain-pinecone>=0.1.1",
+    "polars>=0.20.25",
+    # "duckdb>=0.10.2",
+    # "duckdb-engine>=0.12.0",
+    "python-dotenv>=1.0.1",
+    "dagster>=1.7.8",
+    "dagster-webserver>=1.7.8",
+    "dagster-openai>=0.23.9",
+]
+readme = "README.md"
+requires-python = ">= 3.8"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.rye]
+managed = true
+dev-dependencies = [
+    "ipython>=8.24.0",
+    "sourcery-cli>=1.18.0",
+    "ipdb>=0.13.13",
+]
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/semantic_catalogue"]
+
+[tool.dagster]
+module_name = "src"

requirements-dev.lock

@@ -0,0 +1,316 @@
+# generated by rye
+# use `rye lock` or `rye sync` to update this lockfile
+#
+# last locked with the following flags:
+#   pre: false
+#   features: []
+#   all-features: false
+#   with-sources: false
+#   generate-hashes: false
+
+-e file:.
+aiohttp==3.9.5
+    # via langchain
+    # via langchain-community
+aiosignal==1.3.1
+    # via aiohttp
+alembic==1.13.1
+    # via dagster
+aniso8601==9.0.1
+    # via graphene
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.4.0
+    # via gql
+    # via httpx
+    # via openai
+    # via starlette
+    # via watchfiles
+asttokens==2.4.1
+    # via stack-data
+attrs==23.2.0
+    # via aiohttp
+backoff==2.2.1
+    # via gql
+certifi==2024.6.2
+    # via httpcore
+    # via httpx
+    # via pinecone-client
+    # via requests
+charset-normalizer==3.3.2
+    # via requests
+click==8.1.7
+    # via dagster
+    # via dagster-webserver
+    # via uvicorn
+coloredlogs==14.0
+    # via dagster
+croniter==2.0.5
+    # via dagster
+dagster==1.7.9
+    # via dagster-graphql
+    # via dagster-openai
+    # via dagster-webserver
+    # via semantic-catalogue
+dagster-graphql==1.7.9
+    # via dagster-webserver
+dagster-openai==0.23.9
+    # via semantic-catalogue
+dagster-pipes==1.7.9
+    # via dagster
+dagster-webserver==1.7.9
+    # via semantic-catalogue
+dataclasses-json==0.6.6
+    # via langchain-community
+decorator==5.1.1
+    # via ipdb
+    # via ipython
+distro==1.9.0
+    # via openai
+docstring-parser==0.16
+    # via dagster
+executing==2.0.1
+    # via stack-data
+filelock==3.14.0
+    # via dagster
+frozenlist==1.4.1
+    # via aiohttp
+    # via aiosignal
+fsspec==2024.5.0
+    # via universal-pathlib
+gql==3.5.0
+    # via dagster-graphql
+graphene==3.3
+    # via dagster-graphql
+graphql-core==3.2.3
+    # via gql
+    # via graphene
+    # via graphql-relay
+graphql-relay==3.2.0
+    # via graphene
+greenlet==3.0.3
+    # via sqlalchemy
+grpcio==1.64.0
+    # via dagster
+    # via grpcio-health-checking
+grpcio-health-checking==1.62.2
+    # via dagster
+h11==0.14.0
+    # via httpcore
+    # via uvicorn
+httpcore==1.0.5
+    # via httpx
+httptools==0.6.1
+    # via uvicorn
+httpx==0.27.0
+    # via openai
+humanfriendly==10.0
+    # via coloredlogs
+idna==3.7
+    # via anyio
+    # via httpx
+    # via requests
+    # via yarl
+ipdb==0.13.13
+ipython==8.25.0
+    # via ipdb
+jedi==0.19.1
+    # via ipython
+jinja2==3.1.4
+    # via dagster
+jsonpatch==1.33
+    # via langchain-core
+jsonpointer==2.4
+    # via jsonpatch
+langchain==0.2.1
+    # via langchain-community
+    # via semantic-catalogue
+langchain-community==0.2.1
+    # via semantic-catalogue
+langchain-core==0.2.3
+    # via langchain
+    # via langchain-community
+    # via langchain-openai
+    # via langchain-pinecone
+    # via langchain-text-splitters
+langchain-openai==0.1.8
+    # via semantic-catalogue
+langchain-pinecone==0.1.1
+    # via semantic-catalogue
+langchain-text-splitters==0.2.0
+    # via langchain
+langsmith==0.1.67
+    # via langchain
+    # via langchain-community
+    # via langchain-core
+mako==1.3.5
+    # via alembic
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==2.1.5
+    # via jinja2
+    # via mako
+marshmallow==3.21.2
+    # via dataclasses-json
+matplotlib-inline==0.1.7
+    # via ipython
+mdurl==0.1.2
+    # via markdown-it-py
+multidict==6.0.5
+    # via aiohttp
+    # via yarl
+mypy-extensions==1.0.0
+    # via typing-inspect
+numpy==1.26.4
+    # via langchain
+    # via langchain-community
+    # via langchain-pinecone
+openai==1.30.5
+    # via dagster-openai
+    # via langchain-openai
+orjson==3.10.3
+    # via langsmith
+packaging==23.2
+    # via dagster
+    # via langchain-core
+    # via marshmallow
+parso==0.8.4
+    # via jedi
+pendulum==3.0.0
+    # via dagster
+pexpect==4.9.0
+    # via ipython
+pinecone-client==3.2.2
+    # via langchain-pinecone
+polars==0.20.31
+    # via semantic-catalogue
+prompt-toolkit==3.0.45
+    # via ipython
+protobuf==4.25.3
+    # via dagster
+    # via grpcio-health-checking
+ptyprocess==0.7.0
+    # via pexpect
+pure-eval==0.2.2
+    # via stack-data
+pydantic==2.7.2
+    # via dagster
+    # via langchain
+    # via langchain-core
+    # via langsmith
+    # via openai
+pydantic-core==2.18.3
+    # via pydantic
+pygments==2.18.0
+    # via ipython
+    # via rich
+python-dateutil==2.9.0.post0
+    # via croniter
+    # via dagster
+    # via pendulum
+    # via time-machine
+python-dotenv==1.0.1
+    # via dagster
+    # via semantic-catalogue
+    # via uvicorn
+pytz==2024.1
+    # via croniter
+    # via dagster
+pyyaml==6.0.1
+    # via dagster
+    # via langchain
+    # via langchain-community
+    # via langchain-core
+    # via uvicorn
+regex==2024.5.15
+    # via tiktoken
+requests==2.32.3
+    # via dagster
+    # via dagster-graphql
+    # via gql
+    # via langchain
+    # via langchain-community
+    # via langsmith
+    # via requests-toolbelt
+    # via semantic-catalogue
+    # via tiktoken
+requests-toolbelt==1.0.0
+    # via gql
+rich==13.7.1
+    # via dagster
+setuptools==70.0.0
+    # via dagster
+six==1.16.0
+    # via asttokens
+    # via python-dateutil
+sniffio==1.3.1
+    # via anyio
+    # via httpx
+    # via openai
+sourcery-cli==1.18.0
+sqlalchemy==2.0.30
+    # via alembic
+    # via dagster
+    # via langchain
+    # via langchain-community
+stack-data==0.6.3
+    # via ipython
+starlette==0.37.2
+    # via dagster-graphql
+    # via dagster-webserver
+structlog==24.2.0
+    # via dagster
+tabulate==0.9.0
+    # via dagster
+tenacity==8.3.0
+    # via langchain
+    # via langchain-community
+    # via langchain-core
+tiktoken==0.7.0
+    # via langchain-openai
+time-machine==2.14.1
+    # via pendulum
+tomli==2.0.1
+    # via dagster
+toposort==1.10
+    # via dagster
+tqdm==4.66.4
+    # via dagster
+    # via openai
+    # via pinecone-client
+traitlets==5.14.3
+    # via ipython
+    # via matplotlib-inline
+typing-extensions==4.12.1
+    # via alembic
+    # via dagster
+    # via openai
+    # via pinecone-client
+    # via pydantic
+    # via pydantic-core
+    # via sqlalchemy
+    # via typing-inspect
+typing-inspect==0.9.0
+    # via dataclasses-json
+tzdata==2024.1
+    # via pendulum
+universal-pathlib==0.2.2
+    # via dagster
+urllib3==2.2.1
+    # via pinecone-client
+    # via requests
+uvicorn==0.30.1
+    # via dagster-webserver
+uvloop==0.19.0
+    # via uvicorn
+watchdog==4.0.1
+    # via dagster
+watchfiles==0.22.0
+    # via uvicorn
+wcwidth==0.2.13
+    # via prompt-toolkit
+websockets==12.0
+    # via uvicorn
+yarl==1.9.4
+    # via aiohttp
+    # via gql

requirements.lock

@@ -0,0 +1,282 @@
+# generated by rye
+# use `rye lock` or `rye sync` to update this lockfile
+#
+# last locked with the following flags:
+#   pre: false
+#   features: []
+#   all-features: false
+#   with-sources: false
+#   generate-hashes: false
+
+-e file:.
+aiohttp==3.9.5
+    # via langchain
+    # via langchain-community
+aiosignal==1.3.1
+    # via aiohttp
+alembic==1.13.1
+    # via dagster
+aniso8601==9.0.1
+    # via graphene
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.4.0
+    # via gql
+    # via httpx
+    # via openai
+    # via starlette
+    # via watchfiles
+attrs==23.2.0
+    # via aiohttp
+backoff==2.2.1
+    # via gql
+certifi==2024.6.2
+    # via httpcore
+    # via httpx
+    # via pinecone-client
+    # via requests
+charset-normalizer==3.3.2
+    # via requests
+click==8.1.7
+    # via dagster
+    # via dagster-webserver
+    # via uvicorn
+coloredlogs==14.0
+    # via dagster
+croniter==2.0.5
+    # via dagster
+dagster==1.7.9
+    # via dagster-graphql
+    # via dagster-openai
+    # via dagster-webserver
+    # via semantic-catalogue
+dagster-graphql==1.7.9
+    # via dagster-webserver
+dagster-openai==0.23.9
+    # via semantic-catalogue
+dagster-pipes==1.7.9
+    # via dagster
+dagster-webserver==1.7.9
+    # via semantic-catalogue
+dataclasses-json==0.6.6
+    # via langchain-community
+distro==1.9.0
+    # via openai
+docstring-parser==0.16
+    # via dagster
+filelock==3.14.0
+    # via dagster
+frozenlist==1.4.1
+    # via aiohttp
+    # via aiosignal
+fsspec==2024.5.0
+    # via universal-pathlib
+gql==3.5.0
+    # via dagster-graphql
+graphene==3.3
+    # via dagster-graphql
+graphql-core==3.2.3
+    # via gql
+    # via graphene
+    # via graphql-relay
+graphql-relay==3.2.0
+    # via graphene
+greenlet==3.0.3
+    # via sqlalchemy
+grpcio==1.64.0
+    # via dagster
+    # via grpcio-health-checking
+grpcio-health-checking==1.62.2
+    # via dagster
+h11==0.14.0
+    # via httpcore
+    # via uvicorn
+httpcore==1.0.5
+    # via httpx
+httptools==0.6.1
+    # via uvicorn
+httpx==0.27.0
+    # via openai
+humanfriendly==10.0
+    # via coloredlogs
+idna==3.7
+    # via anyio
+    # via httpx
+    # via requests
+    # via yarl
+jinja2==3.1.4
+    # via dagster
+jsonpatch==1.33
+    # via langchain-core
+jsonpointer==2.4
+    # via jsonpatch
+langchain==0.2.1
+    # via langchain-community
+    # via semantic-catalogue
+langchain-community==0.2.1
+    # via semantic-catalogue
+langchain-core==0.2.3
+    # via langchain
+    # via langchain-community
+    # via langchain-openai
+    # via langchain-pinecone
+    # via langchain-text-splitters
+langchain-openai==0.1.8
+    # via semantic-catalogue
+langchain-pinecone==0.1.1
+    # via semantic-catalogue
+langchain-text-splitters==0.2.0
+    # via langchain
+langsmith==0.1.67
+    # via langchain
+    # via langchain-community
+    # via langchain-core
+mako==1.3.5
+    # via alembic
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==2.1.5
+    # via jinja2
+    # via mako
+marshmallow==3.21.2
+    # via dataclasses-json
+mdurl==0.1.2
+    # via markdown-it-py
+multidict==6.0.5
+    # via aiohttp
+    # via yarl
+mypy-extensions==1.0.0
+    # via typing-inspect
+numpy==1.26.4
+    # via langchain
+    # via langchain-community
+    # via langchain-pinecone
+openai==1.30.5
+    # via dagster-openai
+    # via langchain-openai
+orjson==3.10.3
+    # via langsmith
+packaging==23.2
+    # via dagster
+    # via langchain-core
+    # via marshmallow
+pendulum==3.0.0
+    # via dagster
+pinecone-client==3.2.2
+    # via langchain-pinecone
+polars==0.20.31
+    # via semantic-catalogue
+protobuf==4.25.3
+    # via dagster
+    # via grpcio-health-checking
+pydantic==2.7.2
+    # via dagster
+    # via langchain
+    # via langchain-core
+    # via langsmith
+    # via openai
+pydantic-core==2.18.3
+    # via pydantic
+pygments==2.18.0
+    # via rich
+python-dateutil==2.9.0.post0
+    # via croniter
+    # via dagster
+    # via pendulum
+    # via time-machine
+python-dotenv==1.0.1
+    # via dagster
+    # via semantic-catalogue
+    # via uvicorn
+pytz==2024.1
+    # via croniter
+    # via dagster
+pyyaml==6.0.1
+    # via dagster
+    # via langchain
+    # via langchain-community
+    # via langchain-core
+    # via uvicorn
+regex==2024.5.15
+    # via tiktoken
+requests==2.32.3
+    # via dagster
+    # via dagster-graphql
+    # via gql
+    # via langchain
+    # via langchain-community
+    # via langsmith
+    # via requests-toolbelt
+    # via semantic-catalogue
+    # via tiktoken
+requests-toolbelt==1.0.0
+    # via gql
+rich==13.7.1
+    # via dagster
+setuptools==70.0.0
+    # via dagster
+six==1.16.0
+    # via python-dateutil
+sniffio==1.3.1
+    # via anyio
+    # via httpx
+    # via openai
+sqlalchemy==2.0.30
+    # via alembic
+    # via dagster
+    # via langchain
+    # via langchain-community
+starlette==0.37.2
+    # via dagster-graphql
+    # via dagster-webserver
+structlog==24.2.0
+    # via dagster
+tabulate==0.9.0
+    # via dagster
+tenacity==8.3.0
+    # via langchain
+    # via langchain-community
+    # via langchain-core
+tiktoken==0.7.0
+    # via langchain-openai
+time-machine==2.14.1
+    # via pendulum
+tomli==2.0.1
+    # via dagster
+toposort==1.10
+    # via dagster
+tqdm==4.66.4
+    # via dagster
+    # via openai
+    # via pinecone-client
+typing-extensions==4.12.1
+    # via alembic
+    # via dagster
+    # via openai
+    # via pinecone-client
+    # via pydantic
+    # via pydantic-core
+    # via sqlalchemy
+    # via typing-inspect
+typing-inspect==0.9.0
+    # via dataclasses-json
+tzdata==2024.1
+    # via pendulum
+universal-pathlib==0.2.2
+    # via dagster
+urllib3==2.2.1
+    # via pinecone-client
+    # via requests
+uvicorn==0.30.1
+    # via dagster-webserver
+uvloop==0.19.0
+    # via uvicorn
+watchdog==4.0.1
+    # via dagster
+watchfiles==0.22.0
+    # via uvicorn
+websockets==12.0
+    # via uvicorn
+yarl==1.9.4
+    # via aiohttp
+    # via gql

src/__init__.py

@@ -0,0 +1,14 @@
+from dagster import Definitions, load_assets_from_modules
+
+from src.assets import adr, datastore
+from src.jobs import adr_job
+from src.resources import openai_resource
+
+adr_assets = load_assets_from_modules(modules=[adr], group_name="adr_assets")
+datastore_assets = load_assets_from_modules(modules=[datastore], group_name="datastore")
+
+defs = Definitions(
+    assets=[*adr_assets, *datastore_assets],
+    jobs=[adr_job],
+    resources={"openai": openai_resource},
+)

src/assets/adr.py

@@ -0,0 +1,118 @@
+import itertools
+import json
+import logging
+
+import polars as pl
+import requests
+from dagster import AssetExecutionContext, asset
+from tqdm import tqdm
+
+from src.common.utils import Paths
+
+PAGE_SIZE = 100
+API_VERSION = "2.0"
+BASE_URL = "https://api-datacatalogue.adruk.org/api"
+
+
+@asset
+def adr_session():
+    session = requests.Session()
+    session.headers.update({"X-API-Version": API_VERSION})
+    return session
+
+
+@asset
+def adr_datasets_id(
+    context: AssetExecutionContext, adr_session: requests.Session
+) -> pl.DataFrame:
+    datasets = []
+    for page_number in itertools.count(start=1):
+        context.log.info(f"Fetching page {page_number}")
+        datasets_page = _fetch_datasets_page(adr_session, page_number)
+        if not datasets_page:
+            break
+        datasets.extend(datasets_page)
+    df = pl.DataFrame(datasets)
+    df.write_parquet(Paths.ADR / "adr_datasets_id.parquet")
+    return df
+
+
+def _fetch_datasets_page(adr_session: requests.Session, page_number: int) -> dict:
+    params = {
+        "pageSize": PAGE_SIZE,
+        "pageNumber": page_number,
+        "include": "dataset::dataelement::dataclass",
+        "additionalProp1": "string",
+        "additionalProp2": "string",
+        "additionalProp3": "string",
+        "state": "START",
+    }
+    try:
+        response = adr_session.get(f"{BASE_URL}/{{sql}}/dataset", params=params)
+        response.raise_for_status()
+        return json.loads(response.content)["content"]
+    except requests.HTTPError as http_err:
+        logging.error(f"HTTP error occurred: {http_err}")
+    except Exception as err:
+        logging.error(f"Other error occurred: {err}")
+
+
+@asset
+def adr_datasets(
+    context: AssetExecutionContext,
+    adr_session: requests.Session,
+    adr_datasets_id: pl.DataFrame,
+) -> pl.DataFrame:
+    df = adr_datasets_id.filter(pl.col("searchResultType") == "PHYSICAL").with_columns(
+        pl.col("origin").struct[0].alias("origin_id")
+    )
+
+    datasets_list = []
+    for row in tqdm(df.rows(named=True), total=len(df)):
+        dataset = _fetch_dataset_info(context, adr_session, row)
+        datasets_list.append(dataset)
+    df = pl.DataFrame(datasets_list)
+    df.write_parquet(Paths.ADR / "adr_datasets.parquet")
+    return df
+
+
+def _fetch_dataset_info(
+    context: AssetExecutionContext, adr_session: requests.Session, row: dict
+) -> dict:
+    url = f"{BASE_URL}/dataset/{row['id']}?originId={row['origin_id']}"
+    try:
+        response = adr_session.get(url)
+        response.raise_for_status()
+    except requests.HTTPError as http_err:
+        context.log.error(f"HTTP error occurred: {http_err}")
+    except Exception as err:
+        context.log.error(f"Other error occurred: {err}")
+    content = json.loads(response.content)
+
+    return {
+        "id": row["id"],
+        "name": row["title"],
+        "origin_name": content["origin"]["name"],
+        "origin_id": row["origin_id"],
+        "doi": content["summary"].get("doiName"),
+        "url": content["origin"]["link"],
+        "keywords": content["summary"].get("keywords"),
+        "abstract": content["summary"]["abstract"],
+        "description": content.get("documentation", {}).get("description"),
+        "coverage": content["coverage"],
+    }
+
+
+@asset
+def adr_descriptions(adr_datasets: pl.DataFrame) -> None:
+    for item in adr_datasets.rows(named=True):
+        with open(
+            Paths.ADR
+            / f"descriptions/{item['id']}-{item['origin_id']}-description.txt",
+            "w",
+        ) as f:
+            f.write(
+                f"Dataset Title: {item['name']}"
+                f"\n\nDescription: \n\n{item['description']}"
+                f"\n\nAbstract: \n\n{item['abstract']}"
+            )

src/assets/datastore.py

@@ -0,0 +1,48 @@
+from dagster import AssetExecutionContext, asset
+from dagster_openai import OpenAIResource
+from dotenv import load_dotenv
+from langchain_community.document_loaders import DirectoryLoader, TextLoader
+from langchain_openai import OpenAIEmbeddings
+from langchain_pinecone import PineconeVectorStore
+from langchain_text_splitters import CharacterTextSplitter
+from pinecone import Pinecone, ServerlessSpec
+
+from src.common.utils import Consts, Paths
+
+load_dotenv()
+
+
+@asset(compute_kind="Pinecone")
+def pinecone_index(context: AssetExecutionContext):
+    pc = Pinecone()
+    if Consts.INDEX_NAME in [index["name"] for index in pc.list_indexes()]:
+        pc.delete_index(Consts.INDEX_NAME)
+
+    pc.create_index(
+        name=Consts.INDEX_NAME,
+        dimension=Consts.EMBEDDING_DIM,
+        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
+        metric="cosine",
+    )
+
+
+@asset(compute_kind="OpenAI", deps=["adr_descriptions", "pinecone_index"])
+def adr_pinecone(context: AssetExecutionContext, openai: OpenAIResource):
+    loader = DirectoryLoader(
+        Paths.ADR / "descriptions",
+        glob="*.txt",
+        loader_cls=TextLoader,
+        use_multithreading=True,
+        show_progress=True,
+    )
+    documents = loader.load()
+    text_splitter = CharacterTextSplitter(chunk_size=1024, chunk_overlap=0)
+    docs = text_splitter.split_documents(documents)
+
+    with openai.get_client(context) as client:
+        embeddings = OpenAIEmbeddings(
+            client=client.embeddings,
+            model=Consts.EMBEDDING_MODEL,
+        )
+
+    PineconeVectorStore.from_documents(docs, embeddings, index_name=Consts.INDEX_NAME)

src/common/utils.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+
+class Paths:
+    DATA = Path("data")
+    ADR = DATA / "adr"
+
+
+class Consts:
+    INDEX_NAME = "data-catalogue"
+    EMBEDDING_MODEL = "text-embedding-3-large"
+    EMBEDDING_DIM = 3072

src/jobs.py

@@ -0,0 +1,7 @@
+from dagster import define_asset_job
+
+adr_job = define_asset_job(
+    "adr",
+    selection=["adr_session", "adr_datasets_id", "adr_datasets", "adr_descriptions"],
+)
+pinecone_job = define_asset_job("pinecone", selection=["adr_pinecone"])

src/resources.py

@@ -0,0 +1,4 @@
+from dagster import EnvVar
+from dagster_openai import OpenAIResource
+
+openai_resource = OpenAIResource(api_key=EnvVar("OPENAI_API_KEY"))

src/sensors.py

@@ -0,0 +1,2 @@
+# see https://docs.dagster.io/concepts/partitions-schedules-sensors/sensors
+# cursor to check which files are new