refactor to use langgraph + dagster

Containerfile

@@ -0,0 +1,12 @@
+FROM python:3.12
+
+ENV VIRTUAL_ENV=/usr/local
+
+ENV DAGSTER_HOME=/opt/dagster/dagster_home/
+RUN mkdir -p $DAGSTER_HOME /opt/dagster/app
+COPY dagster.yaml $DAGSTER_HOME
+
+WORKDIR /opt/dagster/app
+
+COPY requirements.lock pyproject.toml .env README.md ./
+RUN pip install --no-cache-dir -r requirements.lock

compose.yml

@@ -0,0 +1,99 @@
+services:
+  datastore-postgresql:
+    image: postgres:11
+    container_name: datastore-postgresql
+    environment:
+      POSTGRES_USER: "postgres_user"
+      POSTGRES_PASSWORD: "postgres_password"
+      POSTGRES_DB: "postgres_db"
+    networks:
+      - datastore-network
+
+  datastore-code:
+    container_name: datastore-code
+    restart: always
+    build:
+      context: .
+      dockerfile: Containerfile
+    image: datastore
+    security_opt:
+      - "label:disable"
+    entrypoint:
+      - dagster
+      - api
+      - grpc
+      - -h
+      - "0.0.0.0"
+      - -p
+      - "4000"
+      - -m
+      - src.datastore
+    environment:
+      DAGSTER_POSTGRES_USER: "postgres_user"
+      DAGSTER_POSTGRES_PASSWORD: "postgres_password"
+      DAGSTER_POSTGRES_DB: "postgres_db"
+      DAGSTER_CURRENT_IMAGE: "datastore"
+    volumes:
+      - ./src:/opt/dagster/app/src
+      - ./config:/opt/dagster/app/config
+      - ./data:/opt/dagster/app/data
+    networks:
+      - datastore-network
+
+  datastore-server:
+    container_name: datastore-server
+    build:
+      context: .
+      dockerfile: Containerfile
+    entrypoint:
+      - dagster-webserver
+      - -h
+      - "0.0.0.0"
+      - -p
+      - "3000"
+    expose:
+      - "3000"
+    ports:
+      - 3000:3000
+    security_opt:
+      - "label:disable"
+    environment:
+      DAGSTER_POSTGRES_USER: "postgres_user"
+      DAGSTER_POSTGRES_PASSWORD: "postgres_password"
+      DAGSTER_POSTGRES_DB: "postgres_db"
+    volumes:
+      - ./src:/opt/dagster/app/src
+      - ./config:/opt/dagster/app/config
+      - ./data:/opt/dagster/app/data
+    networks:
+      - datastore-network
+
+  datastore-daemon:
+    container_name: datastore-daemon
+    restart: on-failure
+    build:
+      context: .
+      dockerfile: Containerfile
+    entrypoint:
+      - dagster-daemon
+      - run
+    security_opt:
+      - "label:disable"
+    environment:
+      DAGSTER_POSTGRES_USER: "postgres_user"
+      DAGSTER_POSTGRES_PASSWORD: "postgres_password"
+      DAGSTER_POSTGRES_DB: "postgres_db"
+    volumes:
+      - ./src:/opt/dagster/app/src
+      - ./config:/opt/dagster/app/config
+      - ./data:/opt/dagster/app/data
+    networks:
+      - datastore-network
+    depends_on:
+      - datastore-postgresql
+      - datastore-code
+
+networks:
+  datastore-network:
+    driver: bridge
+    name: datastore-network

config/config.toml

@@ -0,0 +1,10 @@
+[datastore]
+index_name = "data-catalogue"
+embed_model = "text-embedding-3-large"
+embed_dim = 3072
+chunk_size = 256
+chunk_overlap = 32
+
+[model]
+llm = "gpt-3.5-turbo"
+top_k = 30

dagster.yaml

@@ -0,0 +1,52 @@
+scheduler:
+  module: dagster.core.scheduler
+  class: DagsterDaemonScheduler
+
+run_retries:
+  enabled: true
+  max_retries: 3
+
+run_storage:
+  module: dagster_postgres.run_storage
+  class: PostgresRunStorage
+  config:
+    postgres_db:
+      hostname: datastore-postgresql
+      username:
+        env: DAGSTER_POSTGRES_USER
+      password:
+        env: DAGSTER_POSTGRES_PASSWORD
+      db_name:
+        env: DAGSTER_POSTGRES_DB
+      port: 5432
+
+schedule_storage:
+  module: dagster_postgres.schedule_storage
+  class: PostgresScheduleStorage
+  config:
+    postgres_db:
+      hostname: datastore-postgresql
+      username:
+        env: DAGSTER_POSTGRES_USER
+      password:
+        env: DAGSTER_POSTGRES_PASSWORD
+      db_name:
+        env: DAGSTER_POSTGRES_DB
+      port: 5432
+
+event_log_storage:
+  module: dagster_postgres.event_log
+  class: PostgresEventLogStorage
+  config:
+    postgres_db:
+      hostname: datastore-postgresql
+      username:
+        env: DAGSTER_POSTGRES_USER
+      password:
+        env: DAGSTER_POSTGRES_PASSWORD
+      db_name:
+        env: DAGSTER_POSTGRES_DB
+      port: 5432
+
+telemetry:
+  enabled: false

pyproject.toml

@@ -6,18 +6,26 @@ 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-postgres>=0.23.15",
     "dagster-openai>=0.23.9",
+    "requests>=2.31.0",
+    "polars>=0.20.25",
+    "python-dotenv>=1.0.1",
+    "sickle>=0.7.0",
+    "lxml>=5.2.2",
+    "pydantic-settings>=2.3.4",
+    "pdfminer-six>=20240706",
+    "dateparser>=1.2.0",
+    "langchain>=0.2.12",
+    "langchain-community>=0.2.10",
+    "langchain-pinecone>=0.1.3",
+    "langchain-openai>=0.1.20",
+    "langchain-experimental>=0.0.63",
+    "langgraph>=0.1.19",
+    "langchainhub>=0.1.20",
+    "fastapi[standard]>=0.112.0",
 ]
 readme = "README.md"
 requires-python = ">= 3.8"
@@ -41,4 +49,4 @@ allow-direct-references = true
 packages = ["src/semantic_catalogue"]
 
 [tool.dagster]
-module_name = "src"
+module_name = "src.datastore"

reports/DOCS.md

@@ -0,0 +1,72 @@
+# Introduction
+
+* Unify search across catalogues
+* Uses semantic search with RAG for results explainability
+* Llama Index framework
+
+# Methodology
+
+## Pre-processing
+
+For each catalogue their respective API was used to return dataset metadata. Each returned result contained descriptive information regarding datasets, which form the bulk of text data used by the semantic search system to return results. For the CDRC catalogue, PDFs were also processed to extract text. Other metadata was also returned which may be used by the final system; for example, data creation date.
+
+## Datastore
+
+The description of each dataset were then saved into individual text files, identifiable by a unique ID. These files were then embedded using OpenAI embeddings, and uploaded to the Pinecone database, alongside any metadata. Descriptions were 'chunked' into individual segments 1024 tokens in length. For each chunk, the dataset title as embedded at the start.
+
+## RAG Model
+
+A RAG system was then built which embeds a user query using the same embedding model, and returns the top 'k' results ranked by cosine similarity from the Pinecone database. To ensure that results are ranked by dataset, a custom document grouping postprocessor was defined, which grouped all document chunks relating to the same dataset. The highest score from any chunk is used to rank grouped documents.
+
+An adjustable 'alpha' value was used to allow for a mixture of traditional 'sparse vector' search (e.g. BM25: keyword search), and the 'dense vector' search, using the LLM embeddings.
+
+For each unique document returned, an explainable 'Ask AI' option was added, which feeds the grouped document into a GPT LLM with the following prompt:
+
+```python
+prompt = """
+Below is a dataset description that is relevant to a researchers query.
+
+Explain the relevance of this dataset to the query in under 50 words. Use your own knowledge or the data profile. Do not say it is unrelated; attempt to find a relevant connection.
+
+---------------------
+Query: "{query_str}"
+
+Dataset description:
+
+{context_str}
+---------------------
+"""
+```
+This approach ensures that users receive not only relevant search results but also understandable explanations regarding the relevance of each dataset to their query.
+
+# System architecture
+
+## Overview
+
+![]('./figs/system.png')
+
+## Data flow
+
+(Describe the flow of data from the catalogues to the end-user.)
+
+## Implementation details
+
+* Tools and Libraries: OpenAI API, Pinecone, Llama Index
+* Challenges: (Detail any challenges and solutions.)
+
+# Evaluation and results
+
+* Performance Metrics: Search accuracy, response time, user feedback
+* Comparison: Effectiveness of keyword search vs. dense vector search
+
+# Future work and improvements
+
+* Potential improvements and future enhancements
+* Discuss limitations of the current implementation
+
+# Conclusion
+
+Summarise the key points and the impact of the unified search system.
+References
+
+(List any academic papers, tools, or libraries referenced.)

requirements-dev.lock

@@ -7,14 +7,16 @@
 #   all-features: false
 #   with-sources: false
 #   generate-hashes: false
+#   universal: false
 
 -e file:.
 aiohttp==3.9.5
     # via langchain
     # via langchain-community
+    # via langchain-pinecone
 aiosignal==1.3.1
     # via aiohttp
-alembic==1.13.1
+alembic==1.13.2
     # via dagster
 aniso8601==9.0.1
     # via graphene
@@ -32,51 +34,70 @@ attrs==23.2.0
     # via aiohttp
 backoff==2.2.1
     # via gql
-certifi==2024.6.2
+certifi==2024.7.4
     # via httpcore
     # via httpx
     # via pinecone-client
     # via requests
+cffi==1.16.0
+    # via cryptography
 charset-normalizer==3.3.2
+    # via pdfminer-six
     # via requests
 click==8.1.7
     # via dagster
     # via dagster-webserver
+    # via typer
     # via uvicorn
 coloredlogs==14.0
     # via dagster
-croniter==2.0.5
+croniter==3.0.3
     # via dagster
-dagster==1.7.9
+cryptography==43.0.0
+    # via pdfminer-six
+dagster==1.7.15
     # via dagster-graphql
     # via dagster-openai
+    # via dagster-postgres
     # via dagster-webserver
     # via semantic-catalogue
-dagster-graphql==1.7.9
+dagster-graphql==1.7.15
     # via dagster-webserver
-dagster-openai==0.23.9
+dagster-openai==0.23.15
     # via semantic-catalogue
-dagster-pipes==1.7.9
+dagster-pipes==1.7.15
     # via dagster
-dagster-webserver==1.7.9
+dagster-postgres==0.23.15
+    # via semantic-catalogue
+dagster-webserver==1.7.15
     # via semantic-catalogue
-dataclasses-json==0.6.6
+dataclasses-json==0.6.7
     # via langchain-community
+dateparser==1.2.0
+    # via semantic-catalogue
 decorator==5.1.1
     # via ipdb
     # via ipython
 distro==1.9.0
     # via openai
+dnspython==2.6.1
+    # via email-validator
 docstring-parser==0.16
     # via dagster
+email-validator==2.2.0
+    # via fastapi
 executing==2.0.1
     # via stack-data
-filelock==3.14.0
+fastapi==0.112.0
+    # via semantic-catalogue
+fastapi-cli==0.0.5
+    # via fastapi
+filelock==3.15.4
     # via dagster
 frozenlist==1.4.1
     # via aiohttp
     # via aiosignal
-fsspec==2024.5.0
+fsspec==2024.6.1
     # via universal-pathlib
 gql==3.5.0
     # via dagster-graphql
@@ -90,7 +111,7 @@ graphql-relay==3.2.0
     # via graphene
 greenlet==3.0.3
     # via sqlalchemy
-grpcio==1.64.0
+grpcio==1.64.1
     # via dagster
     # via grpcio-health-checking
 grpcio-health-checking==1.62.2
@@ -103,46 +124,61 @@ httpcore==1.0.5
 httptools==0.6.1
     # via uvicorn
 httpx==0.27.0
+    # via fastapi
     # via openai
 humanfriendly==10.0
     # via coloredlogs
 idna==3.7
     # via anyio
+    # via email-validator
     # via httpx
     # via requests
     # via yarl
 ipdb==0.13.13
-ipython==8.25.0
+ipython==8.26.0
     # via ipdb
 jedi==0.19.1
     # via ipython
 jinja2==3.1.4
     # via dagster
+    # via fastapi
 jsonpatch==1.33
     # via langchain-core
-jsonpointer==2.4
+jsonpointer==3.0.0
     # via jsonpatch
-langchain==0.2.1
+langchain==0.2.12
     # via langchain-community
     # via semantic-catalogue
-langchain-community==0.2.1
+langchain-community==0.2.10
+    # via langchain-experimental
     # via semantic-catalogue
-langchain-core==0.2.3
+langchain-core==0.2.27
     # via langchain
     # via langchain-community
+    # via langchain-experimental
     # via langchain-openai
     # via langchain-pinecone
     # via langchain-text-splitters
-langchain-openai==0.1.8
+    # via langgraph
+langchain-experimental==0.0.63
     # via semantic-catalogue
-langchain-pinecone==0.1.1
+langchain-openai==0.1.20
     # via semantic-catalogue
-langchain-text-splitters==0.2.0
+langchain-pinecone==0.1.3
+    # via semantic-catalogue
+langchain-text-splitters==0.2.2
     # via langchain
-langsmith==0.1.67
+langchainhub==0.1.20
+    # via semantic-catalogue
+langgraph==0.1.19
+    # via semantic-catalogue
+langsmith==0.1.96
     # via langchain
     # via langchain-community
     # via langchain-core
+lxml==5.2.2
+    # via semantic-catalogue
+    # via sickle
 mako==1.3.5
     # via alembic
 markdown-it-py==3.0.0
@@ -150,7 +186,7 @@ markdown-it-py==3.0.0
 markupsafe==2.1.5
     # via jinja2
     # via mako
-marshmallow==3.21.2
+marshmallow==3.21.3
     # via dataclasses-json
 matplotlib-inline==0.1.7
     # via ipython
@@ -165,64 +201,85 @@ numpy==1.26.4
     # via langchain
     # via langchain-community
     # via langchain-pinecone
-openai==1.30.5
+openai==1.37.1
     # via dagster-openai
     # via langchain-openai
-orjson==3.10.3
+orjson==3.10.6
     # via langsmith
-packaging==23.2
+packaging==24.1
     # via dagster
     # via langchain-core
+    # via langchainhub
     # via marshmallow
 parso==0.8.4
     # via jedi
+pdfminer-six==20240706
+    # via semantic-catalogue
 pendulum==3.0.0
     # via dagster
 pexpect==4.9.0
     # via ipython
-pinecone-client==3.2.2
+pinecone-client==5.0.1
     # via langchain-pinecone
-polars==0.20.31
+pinecone-plugin-inference==1.0.3
+    # via pinecone-client
+pinecone-plugin-interface==0.0.7
+    # via pinecone-client
+    # via pinecone-plugin-inference
+polars==1.3.0
     # via semantic-catalogue
-prompt-toolkit==3.0.45
+prompt-toolkit==3.0.47
     # via ipython
-protobuf==4.25.3
+protobuf==4.25.4
     # via dagster
     # via grpcio-health-checking
+psycopg2-binary==2.9.9
+    # via dagster-postgres
 ptyprocess==0.7.0
     # via pexpect
-pure-eval==0.2.2
+pure-eval==0.2.3
     # via stack-data
-pydantic==2.7.2
+pycparser==2.22
+    # via cffi
+pydantic==2.8.2
     # via dagster
+    # via fastapi
     # via langchain
     # via langchain-core
     # via langsmith
     # via openai
-pydantic-core==2.18.3
+    # via pydantic-settings
+pydantic-core==2.20.1
     # via pydantic
+pydantic-settings==2.3.4
+    # via semantic-catalogue
 pygments==2.18.0
     # via ipython
     # via rich
 python-dateutil==2.9.0.post0
     # via croniter
-    # via dagster
+    # via dateparser
     # via pendulum
     # via time-machine
 python-dotenv==1.0.1
     # via dagster
+    # via pydantic-settings
     # via semantic-catalogue
     # via uvicorn
+python-multipart==0.0.9
+    # via fastapi
 pytz==2024.1
     # via croniter
     # via dagster
+    # via dateparser
 pyyaml==6.0.1
     # via dagster
     # via langchain
     # via langchain-community
     # via langchain-core
     # via uvicorn
-regex==2024.5.15
+regex==2024.7.24
+    # via dateparser
     # via tiktoken
 requests==2.32.3
     # via dagster
@@ -230,16 +287,23 @@ requests==2.32.3
     # via gql
     # via langchain
     # via langchain-community
+    # via langchainhub
     # via langsmith
     # via requests-toolbelt
     # via semantic-catalogue
+    # via sickle
     # via tiktoken
 requests-toolbelt==1.0.0
     # via gql
 rich==13.7.1
     # via dagster
-setuptools==70.0.0
+    # via typer
+setuptools==72.1.0
     # via dagster
+shellingham==1.5.4
+    # via typer
+sickle==0.7.0
+    # via semantic-catalogue
 six==1.16.0
     # via asttokens
     # via python-dateutil
@@ -248,7 +312,7 @@ sniffio==1.3.1
     # via httpx
     # via openai
 sourcery-cli==1.18.0
-sqlalchemy==2.0.30
+sqlalchemy==2.0.31
     # via alembic
     # via dagster
     # via langchain
@@ -258,17 +322,18 @@ stack-data==0.6.3
 starlette==0.37.2
     # via dagster-graphql
     # via dagster-webserver
-structlog==24.2.0
+    # via fastapi
+structlog==24.4.0
     # via dagster
 tabulate==0.9.0
     # via dagster
-tenacity==8.3.0
+tenacity==8.5.0
     # via langchain
     # via langchain-community
     # via langchain-core
 tiktoken==0.7.0
     # via langchain-openai
-time-machine==2.14.1
+time-machine==2.14.2
     # via pendulum
 tomli==2.0.1
     # via dagster
@@ -281,26 +346,38 @@ tqdm==4.66.4
 traitlets==5.14.3
     # via ipython
     # via matplotlib-inline
-typing-extensions==4.12.1
+typer==0.12.3
+    # via fastapi-cli
+types-requests==2.32.0.20240712
+    # via langchainhub
+typing-extensions==4.12.2
     # via alembic
     # via dagster
+    # via fastapi
+    # via langchain-core
     # via openai
     # via pinecone-client
     # via pydantic
     # via pydantic-core
     # via sqlalchemy
+    # via typer
     # via typing-inspect
 typing-inspect==0.9.0
     # via dataclasses-json
 tzdata==2024.1
     # via pendulum
+tzlocal==5.2
+    # via dateparser
 universal-pathlib==0.2.2
     # via dagster
-urllib3==2.2.1
+urllib3==2.2.2
     # via pinecone-client
     # via requests
-uvicorn==0.30.1
+    # via types-requests
+uvicorn==0.30.3
     # via dagster-webserver
+    # via fastapi
+    # via fastapi-cli
 uvloop==0.19.0
     # via uvicorn
 watchdog==4.0.1

requirements.lock

@@ -7,14 +7,16 @@
 #   all-features: false
 #   with-sources: false
 #   generate-hashes: false
+#   universal: false
 
 -e file:.
 aiohttp==3.9.5
     # via langchain
     # via langchain-community
+    # via langchain-pinecone
 aiosignal==1.3.1
     # via aiohttp
-alembic==1.13.1
+alembic==1.13.2
     # via dagster
 aniso8601==9.0.1
     # via graphene
@@ -30,46 +32,65 @@ attrs==23.2.0
     # via aiohttp
 backoff==2.2.1
     # via gql
-certifi==2024.6.2
+certifi==2024.7.4
     # via httpcore
     # via httpx
     # via pinecone-client
     # via requests
+cffi==1.16.0
+    # via cryptography
 charset-normalizer==3.3.2
+    # via pdfminer-six
     # via requests
 click==8.1.7
     # via dagster
     # via dagster-webserver
+    # via typer
     # via uvicorn
 coloredlogs==14.0
     # via dagster
-croniter==2.0.5
+croniter==3.0.3
     # via dagster
-dagster==1.7.9
+cryptography==43.0.0
+    # via pdfminer-six
+dagster==1.7.15
     # via dagster-graphql
     # via dagster-openai
+    # via dagster-postgres
     # via dagster-webserver
     # via semantic-catalogue
-dagster-graphql==1.7.9
+dagster-graphql==1.7.15
     # via dagster-webserver
-dagster-openai==0.23.9
+dagster-openai==0.23.15
     # via semantic-catalogue
-dagster-pipes==1.7.9
+dagster-pipes==1.7.15
     # via dagster
-dagster-webserver==1.7.9
+dagster-postgres==0.23.15
+    # via semantic-catalogue
+dagster-webserver==1.7.15
     # via semantic-catalogue
-dataclasses-json==0.6.6
+dataclasses-json==0.6.7
     # via langchain-community
+dateparser==1.2.0
+    # via semantic-catalogue
 distro==1.9.0
     # via openai
+dnspython==2.6.1
+    # via email-validator
 docstring-parser==0.16
     # via dagster
-filelock==3.14.0
+email-validator==2.2.0
+    # via fastapi
+fastapi==0.112.0
+    # via semantic-catalogue
+fastapi-cli==0.0.5
+    # via fastapi
+filelock==3.15.4
     # via dagster
 frozenlist==1.4.1
     # via aiohttp
     # via aiosignal
-fsspec==2024.5.0
+fsspec==2024.6.1
     # via universal-pathlib
 gql==3.5.0
     # via dagster-graphql
@@ -83,7 +104,7 @@ graphql-relay==3.2.0
     # via graphene
 greenlet==3.0.3
     # via sqlalchemy
-grpcio==1.64.0
+grpcio==1.64.1
     # via dagster
     # via grpcio-health-checking
 grpcio-health-checking==1.62.2
@@ -96,41 +117,56 @@ httpcore==1.0.5
 httptools==0.6.1
     # via uvicorn
 httpx==0.27.0
+    # via fastapi
     # via openai
 humanfriendly==10.0
     # via coloredlogs
 idna==3.7
     # via anyio
+    # via email-validator
     # via httpx
     # via requests
     # via yarl
 jinja2==3.1.4
     # via dagster
+    # via fastapi
 jsonpatch==1.33
     # via langchain-core
-jsonpointer==2.4
+jsonpointer==3.0.0
     # via jsonpatch
-langchain==0.2.1
+langchain==0.2.12
     # via langchain-community
     # via semantic-catalogue
-langchain-community==0.2.1
+langchain-community==0.2.10
+    # via langchain-experimental
     # via semantic-catalogue
-langchain-core==0.2.3
+langchain-core==0.2.27
     # via langchain
     # via langchain-community
+    # via langchain-experimental
     # via langchain-openai
     # via langchain-pinecone
     # via langchain-text-splitters
-langchain-openai==0.1.8
+    # via langgraph
+langchain-experimental==0.0.63
     # via semantic-catalogue
-langchain-pinecone==0.1.1
+langchain-openai==0.1.20
     # via semantic-catalogue
-langchain-text-splitters==0.2.0
+langchain-pinecone==0.1.3
+    # via semantic-catalogue
+langchain-text-splitters==0.2.2
     # via langchain
-langsmith==0.1.67
+langchainhub==0.1.20
+    # via semantic-catalogue
+langgraph==0.1.19
+    # via semantic-catalogue
+langsmith==0.1.96
     # via langchain
     # via langchain-community
     # via langchain-core
+lxml==5.2.2
+    # via semantic-catalogue
+    # via sickle
 mako==1.3.5
     # via alembic
 markdown-it-py==3.0.0
@@ -138,7 +174,7 @@ markdown-it-py==3.0.0
 markupsafe==2.1.5
     # via jinja2
     # via mako
-marshmallow==3.21.2
+marshmallow==3.21.3
     # via dataclasses-json
 mdurl==0.1.2
     # via markdown-it-py
@@ -151,53 +187,74 @@ numpy==1.26.4
     # via langchain
     # via langchain-community
     # via langchain-pinecone
-openai==1.30.5
+openai==1.37.1
     # via dagster-openai
     # via langchain-openai
-orjson==3.10.3
+orjson==3.10.6
     # via langsmith
-packaging==23.2
+packaging==24.1
     # via dagster
     # via langchain-core
+    # via langchainhub
     # via marshmallow
+pdfminer-six==20240706
+    # via semantic-catalogue
 pendulum==3.0.0
     # via dagster
-pinecone-client==3.2.2
+pinecone-client==5.0.1
     # via langchain-pinecone
-polars==0.20.31
+pinecone-plugin-inference==1.0.3
+    # via pinecone-client
+pinecone-plugin-interface==0.0.7
+    # via pinecone-client
+    # via pinecone-plugin-inference
+polars==1.3.0
     # via semantic-catalogue
-protobuf==4.25.3
+protobuf==4.25.4
     # via dagster
     # via grpcio-health-checking
-pydantic==2.7.2
+psycopg2-binary==2.9.9
+    # via dagster-postgres
+pycparser==2.22
+    # via cffi
+pydantic==2.8.2
     # via dagster
+    # via fastapi
     # via langchain
     # via langchain-core
     # via langsmith
     # via openai
-pydantic-core==2.18.3
+    # via pydantic-settings
+pydantic-core==2.20.1
     # via pydantic
+pydantic-settings==2.3.4
+    # via semantic-catalogue
 pygments==2.18.0
     # via rich
 python-dateutil==2.9.0.post0
     # via croniter
-    # via dagster
+    # via dateparser
     # via pendulum
     # via time-machine
 python-dotenv==1.0.1
     # via dagster
+    # via pydantic-settings
     # via semantic-catalogue
     # via uvicorn
+python-multipart==0.0.9
+    # via fastapi
 pytz==2024.1
     # via croniter
     # via dagster
+    # via dateparser
 pyyaml==6.0.1
     # via dagster
     # via langchain
     # via langchain-community
     # via langchain-core
     # via uvicorn
-regex==2024.5.15
+regex==2024.7.24
+    # via dateparser
     # via tiktoken
 requests==2.32.3
     # via dagster
@@ -205,23 +262,30 @@ requests==2.32.3
     # via gql
     # via langchain
     # via langchain-community
+    # via langchainhub
     # via langsmith
     # via requests-toolbelt
     # via semantic-catalogue
+    # via sickle
     # via tiktoken
 requests-toolbelt==1.0.0
     # via gql
 rich==13.7.1
     # via dagster
-setuptools==70.0.0
+    # via typer
+setuptools==72.1.0
     # via dagster
+shellingham==1.5.4
+    # via typer
+sickle==0.7.0
+    # via semantic-catalogue
 six==1.16.0
     # via python-dateutil
 sniffio==1.3.1
     # via anyio
     # via httpx
     # via openai
-sqlalchemy==2.0.30
+sqlalchemy==2.0.31
     # via alembic
     # via dagster
     # via langchain
@@ -229,17 +293,18 @@ sqlalchemy==2.0.30
 starlette==0.37.2
     # via dagster-graphql
     # via dagster-webserver
-structlog==24.2.0
+    # via fastapi
+structlog==24.4.0
     # via dagster
 tabulate==0.9.0
     # via dagster
-tenacity==8.3.0
+tenacity==8.5.0
     # via langchain
     # via langchain-community
     # via langchain-core
 tiktoken==0.7.0
     # via langchain-openai
-time-machine==2.14.1
+time-machine==2.14.2
     # via pendulum
 tomli==2.0.1
     # via dagster
@@ -249,26 +314,38 @@ tqdm==4.66.4
     # via dagster
     # via openai
     # via pinecone-client
-typing-extensions==4.12.1
+typer==0.12.3
+    # via fastapi-cli
+types-requests==2.32.0.20240712
+    # via langchainhub
+typing-extensions==4.12.2
     # via alembic
     # via dagster
+    # via fastapi
+    # via langchain-core
     # via openai
     # via pinecone-client
     # via pydantic
     # via pydantic-core
     # via sqlalchemy
+    # via typer
     # via typing-inspect
 typing-inspect==0.9.0
     # via dataclasses-json
 tzdata==2024.1
     # via pendulum
+tzlocal==5.2
+    # via dateparser
 universal-pathlib==0.2.2
     # via dagster
-urllib3==2.2.1
+urllib3==2.2.2
     # via pinecone-client
     # via requests
-uvicorn==0.30.1
+    # via types-requests
+uvicorn==0.30.3
     # via dagster-webserver
+    # via fastapi
+    # via fastapi-cli
 uvloop==0.19.0
     # via uvicorn
 watchdog==4.0.1

src/__init__.py

@@ -1,14 +0,0 @@
-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/datastore.py

@@ -1,48 +0,0 @@
-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/logging.py

@@ -0,0 +1,3 @@
+import logging.config
+
+logger = logging.getLogger("data-catalogue")

src/common/settings.py

@@ -0,0 +1,28 @@
+import tomllib
+
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+with open("./config/config.toml", "rb") as f:
+    Config = tomllib.load(f)
+
+
+class DataStoreSettings(BaseSettings):
+    index_name: str = Field(min_length=1)
+    embed_model: str = Field(min_length=1)
+    embed_dim: int = Field(gt=0, le=10_000)
+    chunk_size: int = Field(gt=0, le=10_000)
+    chunk_overlap: int = Field(ge=0, le=10_000)
+
+
+class ModelSettings(BaseSettings):
+    llm: str = Field(min_length=1)
+    top_k: int = Field(gt=0, le=100)
+
+
+class Settings(BaseSettings):
+    model: ModelSettings = ModelSettings.model_validate(Config["model"])
+    datastore: DataStoreSettings = DataStoreSettings.model_validate(Config["datastore"])
+
+
+cfg = Settings()

src/common/utils.py

@@ -1,12 +1,21 @@
 from pathlib import Path
 
+from dotenv import load_dotenv
+
+load_dotenv()
+
 
 class Paths:
-    DATA = Path("data")
+    DATA: Path = Path("data")
     ADR = DATA / "adr"
+    UKDS = DATA / "ukds"
+    CDRC = DATA / "cdrc"
+
+    @classmethod
+    def ensure_directories_exist(cls):
+        cls.ADR.mkdir(parents=True, exist_ok=True)
+        cls.UKDS.mkdir(parents=True, exist_ok=True)
+        cls.CDRC.mkdir(parents=True, exist_ok=True)
 
 
-class Consts:
-    INDEX_NAME = "data-catalogue"
-    EMBEDDING_MODEL = "text-embedding-3-large"
-    EMBEDDING_DIM = 3072
+Paths.ensure_directories_exist()

src/datastore/__init__.py

@@ -0,0 +1,18 @@
+from dagster import Definitions, load_assets_from_modules
+
+from src.datastore.assets import adr, cdrc, datastore, ukds
+from src.datastore.jobs import adr_job, cdrc_job, ukds_job
+from src.datastore.resources import openai_resource
+from src.datastore.schedules import adr_schedule, cdrc_schedule, ukds_schedule
+
+adr_assets = load_assets_from_modules(modules=[adr], group_name="adr_assets")
+ukds_assets = load_assets_from_modules(modules=[ukds], group_name="ukds_assets")
+cdrc_assets = load_assets_from_modules(modules=[cdrc], group_name="cdrc_assets")
+datastore_assets = load_assets_from_modules(modules=[datastore], group_name="datastore")
+
+defs = Definitions(
+    assets=[*ukds_assets, *adr_assets, *cdrc_assets, *datastore_assets],
+    jobs=[adr_job, ukds_job, cdrc_job],
+    schedules=[adr_schedule, ukds_schedule, cdrc_schedule],
+    resources={"openai": openai_resource},
+)

src/{assets → datastore/assets}/adr.py

@@ -1,10 +1,10 @@
 import itertools
 import json
-import logging
 
 import polars as pl
 import requests
 from dagster import AssetExecutionContext, asset
+from tenacity import retry, stop_after_attempt, wait_exponential
 from tqdm import tqdm
 
 from src.common.utils import Paths
@@ -15,7 +15,7 @@ BASE_URL = "https://api-datacatalogue.adruk.org/api"
 
 
 @asset
-def adr_session():
+def adr_session() -> requests.Session:
     session = requests.Session()
     session.headers.update({"X-API-Version": API_VERSION})
     return session
@@ -28,16 +28,26 @@ def adr_datasets_id(
     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:
+        datasets_page = _fetch_datasets_page(context, adr_session, page_number)
+        if "end" in datasets_page:
+            context.log.info(f"End of pages reached at {datasets_page['end']}")
             break
         datasets.extend(datasets_page)
-    df = pl.DataFrame(datasets)
+    df = (
+        pl.DataFrame(datasets)
+        .select(["origin", "id", "searchResultType", "title"])
+        .filter(pl.col("searchResultType") == "PHYSICAL")
+        .with_columns(pl.col("origin").struct[0].alias("origin_id"))
+    )
     df.write_parquet(Paths.ADR / "adr_datasets_id.parquet")
+
     return df
 
 
-def _fetch_datasets_page(adr_session: requests.Session, page_number: int) -> dict:
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
+def _fetch_datasets_page(
+    context: AssetExecutionContext, adr_session: requests.Session, page_number: int
+) -> dict:
     params = {
         "pageSize": PAGE_SIZE,
         "pageNumber": page_number,
@@ -50,11 +60,16 @@ def _fetch_datasets_page(adr_session: requests.Session, page_number: int) -> dic
     try:
         response = adr_session.get(f"{BASE_URL}/{{sql}}/dataset", params=params)
         response.raise_for_status()
-        return json.loads(response.content)["content"]
+        content = json.loads(response.content)["content"]
+        if not content:
+            return {"end": page_number}
+        return content
     except requests.HTTPError as http_err:
-        logging.error(f"HTTP error occurred: {http_err}")
+        context.log.error(f"HTTP error occurred: {http_err}")
+        return {}
     except Exception as err:
-        logging.error(f"Other error occurred: {err}")
+        context.log.error(f"Other error occurred: {err}")
+        return {}
 
 
 @asset
@@ -63,19 +78,26 @@ def adr_datasets(
     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)):
+    for row in tqdm(adr_datasets_id.rows(named=True), total=len(adr_datasets_id)):
         dataset = _fetch_dataset_info(context, adr_session, row)
-        datasets_list.append(dataset)
+        if dataset:
+            datasets_list.append(dataset)
     df = pl.DataFrame(datasets_list)
-    df.write_parquet(Paths.ADR / "adr_datasets.parquet")
+    (
+        df.with_columns(
+            pl.col("coverage").struct[0].alias("coverage_0"),
+            pl.col("coverage").struct[1].alias("coverage_1"),
+        )
+        .drop("coverage")
+        .write_parquet(Paths.ADR / "adr_datasets.parquet")
+    )
+
     return df
 
 
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
 def _fetch_dataset_info(
     context: AssetExecutionContext, adr_session: requests.Session, row: dict
 ) -> dict:
@@ -85,8 +107,10 @@ def _fetch_dataset_info(
         response.raise_for_status()
     except requests.HTTPError as http_err:
         context.log.error(f"HTTP error occurred: {http_err}")
+        return {}
     except Exception as err:
         context.log.error(f"Other error occurred: {err}")
+        return {}
     content = json.loads(response.content)
 
     return {
@@ -105,10 +129,12 @@ def _fetch_dataset_info(
 
 @asset
 def adr_descriptions(adr_datasets: pl.DataFrame) -> None:
+    outdir = Paths.ADR / "txt"
+    outdir.mkdir(parents=True, exist_ok=True)
+
     for item in adr_datasets.rows(named=True):
         with open(
-            Paths.ADR
-            / f"descriptions/{item['id']}-{item['origin_id']}-description.txt",
+            outdir / f"{item['id']}-{item['origin_id']}-description.txt",
             "w",
         ) as f:
             f.write(

src/datastore/assets/cdrc.py

@@ -0,0 +1,113 @@
+import itertools
+import os
+import re
+
+import polars as pl
+import requests
+from dagster import AssetExecutionContext, asset
+from dotenv import load_dotenv
+from tqdm import tqdm
+
+from src.common.utils import Paths
+
+load_dotenv()
+
+METADATA_URL = (
+    "https://data.cdrc.ac.uk/api/3/action/current_package_list_with_resources"
+)
+LOGIN_URL = "https://data.cdrc.ac.uk/user/login"
+
+
+@asset
+def cdrc_metadata(context: AssetExecutionContext) -> list[dict]:
+    try:
+        r = requests.get(METADATA_URL)
+        r.raise_for_status()
+    except requests.HTTPError as http_err:
+        context.log.error(f"HTTP error occurred: {http_err}")
+        raise
+    except Exception as err:
+        context.log.error(f"Other error occurred: {err}")
+        raise
+    catalogue_metadata = r.json()["result"][0]
+    return catalogue_metadata
+
+
+@asset
+def cdrc_notes(cdrc_metadata: list[dict]):
+    outdir = Paths.CDRC / "txt"
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    df = pl.DataFrame(cdrc_metadata).drop(["resources", "tags", "extras"])
+    df.write_parquet(Paths.CDRC / "cdrc_metadata.parquet")
+
+    for item in df.rows(named=True):
+        with open(Paths.CDRC / "txt" / f"{item['id']}-notes.txt", "w") as f:
+            f.write(
+                f"Dataset Title: {item['title']}"
+                "\n\nDescription: \n\n"
+                f"{re.sub('<[^<]+?>','', item['notes'])}"
+            )
+
+
+@asset
+def cdrc_resources(cdrc_metadata: list[dict]) -> pl.DataFrame:
+    resources = list(
+        itertools.chain.from_iterable(
+            [item["resources"] if "resources" in item else [] for item in cdrc_metadata]
+        )
+    )
+    df = pl.concat(
+        [
+            pl.DataFrame(cdrc_metadata).explode("resources").drop("resources"),
+            pl.DataFrame(resources)
+            .rename(
+                {"id": "resource_id", "url": "resource_url", "name": "resource_name"},
+            )
+            .drop(["state", "revision_timestamp"]),
+        ],
+        how="horizontal",
+    ).filter((pl.col("format") == "pdf") & (pl.col("resource_url") != ""))
+    df.write_parquet(Paths.CDRC / "cdrc_resource_metadata.parquet")
+    return df
+
+
+@asset
+def cdrc_session() -> requests.Session:
+    session = requests.Session()
+    session.post(
+        LOGIN_URL,
+        data={
+            "name": os.getenv("CDRC_USERNAME"),
+            "pass": os.getenv("CDRC_PASSWORD"),
+            "form_build_id": os.getenv("CDRC_FORM_BUILD_ID"),
+            "form_id": "user_login",
+            "op": "Log in",
+        },
+    )
+    return session
+
+
+@asset
+def cdrc_pdfs(
+    context: AssetExecutionContext,
+    cdrc_session: requests.Session,
+    cdrc_resources: pl.DataFrame,
+):
+    outdir = Paths.CDRC / "pdf"
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    for item in tqdm(cdrc_resources.rows(named=True)):
+        context.log.info(f"Processing {item['resource_url']}...")
+        filepath = outdir / f"{item['id']}-{item['resource_id']}.pdf"
+        try:
+            file = cdrc_session.get(item["resource_url"])
+            file.raise_for_status()
+        except requests.HTTPError as http_err:
+            context.log.error(f"HTTP error occurred: {http_err}")
+            continue
+        except Exception as err:
+            context.log.error(f"Other error occurred: {err}")
+            continue
+        with open(filepath, "wb") as f:
+            f.write(file.content)

src/datastore/assets/datastore.py

@@ -0,0 +1,99 @@
+import time
+from pathlib import Path
+
+from dagster import (
+    AssetExecutionContext,
+    AutoMaterializePolicy,
+    AutoMaterializeRule,
+    asset,
+)
+from dagster_openai import OpenAIResource
+from langchain_community.document_loaders import DirectoryLoader
+from langchain_experimental.text_splitter import SemanticChunker
+from langchain_openai import OpenAIEmbeddings
+from langchain_pinecone import PineconeVectorStore
+from pinecone import Pinecone, ServerlessSpec
+
+from src.common.settings import cfg
+from src.common.utils import Paths
+from src.datastore.loaders import ADRLoader, CDRCLoader, UKDSLoader
+
+wait_on_all_parents_policy = AutoMaterializePolicy.eager().with_rules(
+    AutoMaterializeRule.skip_on_not_all_parents_updated()
+)
+
+
+def _process_documents(
+    context: AssetExecutionContext,
+    openai: OpenAIResource,
+    paths: list[Path],
+    glob_patterns: list[str],
+    loader_classes: list[type],
+):
+    documents = []
+    for path, glob_pattern, loader_cls in zip(paths, glob_patterns, loader_classes):
+        loader = DirectoryLoader(
+            str(path),
+            glob=glob_pattern,
+            loader_cls=loader_cls,
+            use_multithreading=True,
+            show_progress=True,
+        )
+        documents.extend(loader.load())
+
+    with openai.get_client(context) as client:
+        embeddings = OpenAIEmbeddings(
+            client=client.embeddings, model=cfg.datastore.embed_model
+        )
+
+    text_splitter = SemanticChunker(
+        embeddings=embeddings, breakpoint_threshold_type="percentile"
+    )
+    docs = text_splitter.split_documents(documents)
+
+    vectorstore = PineconeVectorStore(
+        index_name=cfg.datastore.index_name, embedding=embeddings
+    )
+    vectorstore.add_documents(documents=docs)
+
+
+@asset(
+    compute_kind="Pinecone",
+    deps=["adr_descriptions", "ukds_abstracts", "cdrc_notes", "cdrc_pdfs"],
+    auto_materialize_policy=wait_on_all_parents_policy,
+)
+def pinecone_index(context: AssetExecutionContext, openai: OpenAIResource):
+    pc = Pinecone()
+    if cfg.datastore.index_name in [index["name"] for index in pc.list_indexes()]:
+        pc.delete_index(cfg.datastore.index_name)
+
+    pc.create_index(
+        name=cfg.datastore.index_name,
+        dimension=cfg.datastore.embed_dim,
+        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
+        metric="cosine",
+    )
+    while not pc.describe_index(cfg.datastore.index_name).status["ready"]:
+        time.sleep(1)
+
+    _process_documents(
+        context,
+        openai,
+        paths=[Paths.ADR / "txt"],
+        glob_patterns=["*.txt"],
+        loader_classes=[ADRLoader],
+    )
+    _process_documents(
+        context,
+        openai,
+        paths=[Paths.CDRC / "txt", Paths.CDRC / "pdf"],
+        glob_patterns=["*.txt", "*.pdf"],
+        loader_classes=[CDRCLoader, CDRCLoader],
+    )
+    _process_documents(
+        context,
+        openai,
+        paths=[Paths.UKDS / "txt"],
+        loader_classes=[UKDSLoader],
+        glob_patterns=["*.txt"],
+    )

src/datastore/assets/ukds.py

@@ -0,0 +1,122 @@
+import re
+import xml.etree.ElementTree as ET
+
+import polars as pl
+import requests
+from dagster import AssetExecutionContext, asset
+from tenacity import retry, stop_after_attempt, wait_exponential
+from tqdm import tqdm
+
+from src.common.utils import Paths
+
+BASE_URL = "https://oai.ukdataservice.ac.uk:8443/oai/provider"
+PARAMS = {"verb": "ListIdentifiers", "metadataPrefix": "ddi", "set": "DataCollections"}
+NAMESPACES = {"oai": "http://www.openarchives.org/OAI/2.0/", "ns2": "ddi:codebook:2_5"}
+
+
+@asset
+def ukds_identifiers() -> list[str]:
+    params = PARAMS.copy()
+
+    identifiers = []
+    while True:
+        response = requests.get(BASE_URL, params=params)
+        if response.status_code != 200:
+            print(f"Failed to fetch data. Status code: {response.status_code}")
+            break
+
+        root = ET.fromstring(response.content)
+
+        headers = root.findall(".//oai:header", NAMESPACES)
+        for header in headers:
+            if header.attrib.get("status") != "deleted":
+                identifier = header.find(".//oai:identifier", NAMESPACES)
+                if identifier is not None:
+                    identifiers.append(identifier.text)
+
+        token_elem = root.find(".//oai:resumptionToken", NAMESPACES)
+        if token_elem is None or token_elem.text is None:
+            break
+
+        params = {"verb": "ListIdentifiers", "resumptionToken": token_elem.text}
+    return identifiers
+
+
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
+def _fetch_metadata(context: AssetExecutionContext, identifier: int):
+    metadata_url = (
+        f"{BASE_URL}?verb=GetRecord&identifier={identifier}&metadataPrefix=ddi"
+    )
+    try:
+        response = requests.get(metadata_url)
+        response.raise_for_status()
+    except requests.HTTPError as http_err:
+        context.log.error(f"HTTP error occurred: {http_err}")
+        raise
+    except Exception as err:
+        context.log.error(f"Other error occurred: {err}")
+        raise
+
+    root = ET.fromstring(response.content)
+    return root
+
+
+@asset
+def ukds_datasets(
+    context: AssetExecutionContext, ukds_identifiers: list[str]
+) -> pl.DataFrame:
+    data = []
+    for identifier in tqdm(ukds_identifiers):
+        context.log.info(f"Fetching identifier {identifier}")
+        metadata = _fetch_metadata(context, identifier).find(
+            ".//ns2:stdyDscr", NAMESPACES
+        )
+
+        if metadata is None:
+            continue
+
+        abstract = "\n".join(
+            [
+                re.sub("<[^<]+?>", "", m.text)
+                for m in metadata.findall(".//ns2:abstract", NAMESPACES)
+                if m.text is not None
+            ]
+        )
+        date = metadata.find(".//ns2:depDate", NAMESPACES)
+        title = metadata.find(".//ns2:titl", NAMESPACES)
+        keywords = [
+            m.text
+            for m in metadata.findall(".//ns2:keyword", NAMESPACES)
+            if m.text is not None
+        ]
+        doi = metadata.find(".//ns2:holdings", NAMESPACES)
+        url = f"https://beta.ukdataservice.ac.uk/datacatalogue/studies/study?id={identifier}"
+        data.append(
+            {
+                "title": title.text if title is not None else None,
+                "abstract": abstract,
+                "date": date.text if date is not None else None,
+                "keywords": keywords,
+                "doi": doi.get("URI") if doi is not None else None,
+                "url": url,
+            }
+        )
+
+    df = pl.DataFrame(data)
+    df.write_parquet(Paths.UKDS / "ukds.parquet")
+    return df
+
+
+@asset
+def ukds_abstracts(ukds_datasets: pl.DataFrame):
+    outdir = Paths.UKDS / "txt"
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    for row in ukds_datasets.rows(named=True):
+        id = row["url"].split("=")[-1]
+        abstract = row["abstract"].replace(
+            "Abstract copyright UK Data Service and data collection copyright owner.",
+            "",
+        )
+        with open(outdir / f"{id}-abstract.txt", "w") as f:
+            f.write(f"Dataset Title: {row['title']}" f"\n\nDescription: \n\n{abstract}")

src/{jobs.py → datastore/jobs.py}

@@ -4,4 +4,17 @@ 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"])
+ukds_job = define_asset_job(
+    "ukds",
+    selection=["ukds_identifiers", "ukds_datasets", "ukds_abstracts"],
+)
+cdrc_job = define_asset_job(
+    "cdrc",
+    selection=[
+        "cdrc_session",
+        "cdrc_metadata",
+        "cdrc_notes",
+        "cdrc_resources",
+        "cdrc_pdfs",
+    ],
+)

src/datastore/loaders.py

@@ -0,0 +1,154 @@
+from pathlib import Path
+from typing import Iterator
+
+import dateparser
+import polars as pl
+from langchain_community.document_loaders import PDFMinerLoader
+from langchain_core.document_loaders import BaseLoader
+from langchain_core.documents import Document
+
+from src.common.utils import Paths
+
+
+class CDRCLoader(BaseLoader):
+    def __init__(self, file_path: str) -> None:
+        self.file_path = file_path
+
+    def lazy_load(self) -> Iterator[Document]:
+        if self.file_path.endswith(".pdf"):
+            document = PDFMinerLoader(self.file_path).load()
+            metadata = self._add_cdrc_pdf_metadata(self.file_path)
+            document[0].metadata |= metadata
+            yield document[0]
+        elif self.file_path.endswith(".txt"):
+            with open(self.file_path, encoding="utf-8") as f:
+                yield Document(
+                    page_content=f.read(),
+                    metadata={"source": self.file_path}
+                    | self._add_cdrc_txt_metadata(self.file_path),
+                )
+
+    @staticmethod
+    def _add_cdrc_txt_metadata(file_path: str) -> dict[str, str]:
+        id = Path(file_path).stem.rsplit("-", maxsplit=1)[0]
+        cdrc_meta = pl.read_parquet(Paths.CDRC / "cdrc_metadata.parquet")
+
+        metadata = cdrc_meta.filter(pl.col("id") == id)
+        iso_date = dateparser.parse(metadata["metadata_created"][0]).isoformat()  # type: ignore
+        return {
+            "title": metadata["title"][0],
+            "id": metadata["id"][0],
+            "url": metadata["url"][0],
+            "date_created": iso_date,
+            "source": "CDRC",
+        }
+
+    @staticmethod
+    def _add_cdrc_pdf_metadata(file_path: str) -> dict[str, str]:
+        id = Path(file_path).stem
+        main_id = "-".join(id.split("-")[:5])
+        resource_id = "-".join(id.split("-")[5:])
+
+        cdrc_meta = pl.read_parquet(Paths.CDRC / "cdrc_metadata.parquet")
+        cdrc_pdf_meta = pl.read_parquet(Paths.CDRC / "cdrc_resource_metadata.parquet")
+
+        resource = cdrc_pdf_meta.filter(pl.col("resource_id") == resource_id)
+        metadata = cdrc_meta.filter(pl.col("id") == main_id)
+
+        iso_date = dateparser.parse(resource["created"][0]).isoformat()  # type: ignore
+        return {
+            "title": metadata["title"][0],
+            "id": metadata["id"][0],
+            "url": metadata["url"][0],
+            "date_created": iso_date,
+            "source": "CDRC",
+        }
+
+
+class ADRLoader(BaseLoader):
+    def __init__(self, file_path: str) -> None:
+        self.file_path = file_path
+
+    def lazy_load(self) -> Iterator[Document]:
+        with open(self.file_path, encoding="utf-8") as f:
+            yield Document(
+                page_content=f.read(),
+                metadata={"source": self.file_path}
+                | self._add_adr_metadata(self.file_path),
+            )
+
+    @staticmethod
+    def _add_adr_metadata(file_path: str) -> dict[str, str]:
+        doc_id, origin_id, _ = Path(file_path).stem.split("-")
+        metadata = (
+            pl.scan_parquet(Paths.ADR / "adr_datasets.parquet")
+            .filter((pl.col("id") == doc_id) & (pl.col("origin_id") == origin_id))
+            .collect()[0]
+            .to_dict(as_series=False)
+        )
+        if len(metadata["id"]) == 0:
+            return {
+                "title": "",
+                "id": f"{doc_id}-{origin_id}",
+                "url": "",
+                "date_created": "",
+                "source": "ADR",
+            }
+
+        date_created = metadata["coverage_1"][0]["distributionReleaseDate"]
+        date_created = (
+            dateparser.parse(date_created).isoformat()  # type: ignore
+            if isinstance(date_created, str)
+            else ""
+        )
+        return {
+            "title": metadata["name"][0],
+            "id": f"{doc_id}-{origin_id}",
+            "url": metadata["url"][0],
+            "date_created": date_created,
+            "source": "ADR",
+        }
+
+
+class UKDSLoader(BaseLoader):
+    def __init__(self, file_path: str) -> None:
+        self.file_path = file_path
+
+    def lazy_load(self) -> Iterator[Document]:
+        with open(self.file_path, encoding="utf-8") as f:
+            yield Document(
+                page_content=f.read(),
+                metadata={"source": self.file_path}
+                | self._add_ukds_metadata(self.file_path),
+            )
+
+    @staticmethod
+    def _add_ukds_metadata(file_path: str) -> dict[str, str]:
+        doc_id = Path(file_path).stem.split("-")[0]
+        metadata = (
+            pl.scan_parquet(Paths.UKDS / "ukds.parquet")
+            .with_columns(pl.col("url").str.split("=").list[1].alias("id"))
+            .filter(pl.col("id") == doc_id)
+            .collect()
+            .to_dict(as_series=False)
+        )
+        if len(metadata["id"]) == 0:
+            return {
+                "title": "",
+                "id": doc_id,
+                "url": "",
+                "date_created": "",
+                "source": "UKDS",
+            }
+        date_created = (
+            dateparser.parse(metadata["date"][0]).isoformat()  # type: ignore
+            if isinstance(metadata["date"][0], str)
+            else ""
+        )
+        return {
+            "title": metadata["title"][0],
+            "id": doc_id,
+            "url": metadata["url"][0],
+            "date_created": date_created,
+            "source": "UKDS",
+        }

src/datastore/schedules.py

@@ -0,0 +1,22 @@
+from dagster import DefaultScheduleStatus, ScheduleDefinition
+
+from src.datastore.jobs import adr_job, cdrc_job, ukds_job
+
+adr_schedule = ScheduleDefinition(
+    job=adr_job,
+    cron_schedule="0 0 * * 0",
+    name="adr_weekly_schedule",
+    default_status=DefaultScheduleStatus.RUNNING,
+)
+cdrc_schedule = ScheduleDefinition(
+    job=cdrc_job,
+    cron_schedule="0 0 * * 0",
+    name="cdrc_weekly_schedule",
+    default_status=DefaultScheduleStatus.RUNNING,
+)
+ukds_schedule = ScheduleDefinition(
+    job=ukds_job,
+    cron_schedule="0 0 * * 0",
+    name="ukds_weekly_schedule",
+    default_status=DefaultScheduleStatus.RUNNING,
+)

src/model/answer.py

@@ -0,0 +1,32 @@
+from dotenv import load_dotenv
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_openai import ChatOpenAI
+
+from src.common.settings import cfg
+from src.model.grader import structured_llm_grader
+
+_ = load_dotenv()
+
+
+class GradeAnswer(BaseModel):
+    """Binary score to assess answer addresses question."""
+
+    binary_score: str = Field(
+        description="Answer addresses the question, 'yes' or 'no'"
+    )
+
+
+llm = ChatOpenAI(model=cfg.model.llm, temperature=0)
+structured_llm_grader = llm.with_structured_output(GradeAnswer)
+
+# Prompt
+system = """You are a grader assessing whether an answer addresses / resolves a question \n 
+     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question."""
+answer_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", system),
+        ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
+    ]
+)
+answer_grader = answer_prompt | structured_llm_grader

src/model/dag.py

@@ -0,0 +1,79 @@
+@asset(compute_kind="OpenAI", deps=["pinecone_index"])
+def search(context: AssetExecutionContext, openai: OpenAIResource):
+    with openai.get_client(context) as client:
+        embeddings = OpenAIEmbeddings(
+            client=client.embeddings, model=cfg.datastore.embed_model
+        )
+        vectorstore = PineconeVectorStore(
+            index_name=cfg.datastore.index_name, embedding=embeddings
+        )
+
+    # retriever = vectorstore.as_retriever()
+
+
+from typing import Any, Dict
+
+import requests
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+app = FastAPI()
+
+
+class RunConfig(BaseModel):
+    run_config: Dict[str, Any]
+
+
+def launch_dagster_run(pipeline_name: str, run_config: dict, mode: str = "default"):
+    url = "http://<your-dagster-instance>/graphql"
+    headers = {
+        "Content-Type": "application/json",
+    }
+    query = """
+    mutation($pipelineName: String!, $runConfigData: RunConfigData, $mode: String!) {
+        launchPipelineExecution(
+            executionParams: {
+                selector: {
+                    pipelineName: $pipelineName
+                },
+                runConfigData: $runConfigData,
+                mode: $mode
+            }
+        ) {
+            __typename
+            ... on LaunchPipelineRunSuccess {
+                run {
+                    runId
+                }
+            }
+            ... on PythonError {
+                message
+                stack
+            }
+        }
+    }
+    """
+    variables = {
+        "pipelineName": pipeline_name,
+        "runConfigData": run_config,
+        "mode": mode,
+    }
+    response = requests.post(
+        url, json={"query": query, "variables": variables}, headers=headers
+    )
+    return response.json()
+
+
+@app.post("/trigger-asset")
+def trigger_asset(pipeline_name: str, run_config: RunConfig, mode: str = "default"):
+    try:
+        result = launch_dagster_run(pipeline_name, run_config.run_config, mode)
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

src/model/grader.py

@@ -0,0 +1,34 @@
+from dotenv import load_dotenv
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_openai import ChatOpenAI
+
+from src.common.settings import cfg
+
+_ = load_dotenv()
+
+
+class GradeDocuments(BaseModel):
+    """Binary score for relevance check on retrieved documents."""
+
+    binary_score: str = Field(
+        description="Documents are relevant to the query, 'yes' or 'no'"
+    )
+
+
+llm = ChatOpenAI(model=cfg.model.llm, temperature=0)
+structured_llm_grader = llm.with_structured_output(GradeDocuments)
+system = """
+You are a grader assessing relevance of a retrieved document to a user query. \n 
+It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
+If the document contains keyword(s) or semantic meaning related to the user query, grade it as relevant. \n
+Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the query.
+"""
+grade_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", system),
+        ("human", "Retrieved document: \n\n {document} \n\n User query: {query}"),
+    ]
+)
+
+retrieval_grader = grade_prompt | structured_llm_grader

src/model/hallucination.py

@@ -0,0 +1,32 @@
+from dotenv import load_dotenv
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_openai import ChatOpenAI
+
+from src.common.settings import cfg
+from src.model.grader import structured_llm_grader
+
+_ = load_dotenv()
+
+
+class GradeHallucinations(BaseModel):
+    """Binary score for hallucination present in generation answer."""
+
+    binary_score: str = Field(
+        description="Answer is grounded in the facts, 'yes' or 'no'"
+    )
+
+
+llm = ChatOpenAI(model=cfg.model.llm, temperature=0)
+structured_llm_grader = llm.with_structured_output(GradeHallucinations)
+
+system = """You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \n 
+     Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts."""
+hallucination_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", system),
+        ("human", "Set of facts: \n\n {document} \n\n LLM generation: {generation}"),
+    ]
+)
+
+hallucination_grader = hallucination_prompt | structured_llm_grader

src/model/model.py

@@ -0,0 +1,132 @@
+from typing import TypedDict
+
+from dotenv import load_dotenv
+from langchain_core.documents import Document
+from langchain_openai import OpenAIEmbeddings
+from langchain_pinecone import PineconeVectorStore
+from langgraph.graph import END, START, StateGraph
+
+from src.common.settings import cfg
+from src.model.grader import retrieval_grader
+from src.model.hallucination import hallucination_grader
+from src.model.rag import rag_chain
+
+_ = load_dotenv()
+
+
+class SearchState(TypedDict):
+    query: str
+    documents: list[str]
+
+
+class GenerationState(TypedDict):
+    query: str
+    document: str
+    generation: str
+
+
+def _group_by_document(documents):
+    grouped_id: dict[str, list[tuple[Document, float]]] = {}
+
+    for node, score in documents:
+        id = node.metadata["id"]
+        if id not in grouped_id:
+            grouped_id[id] = []
+        grouped_id[id].append((node, score))
+
+    out_nodes = []
+    for group in grouped_id.values():
+        nodes = [n[0] for n in group]
+        scores = [n[1] for n in group]
+        content = "\n--------------------\n".join([n.page_content for n in nodes])
+        document = Document(
+            page_content=content, metadata=nodes[0].metadata | {"score": max(scores)}
+        )
+        out_nodes.append(document)
+    return out_nodes
+
+
+def retrieve(state):
+    print("---RETRIEVE---")
+    embeddings = OpenAIEmbeddings(model=cfg.datastore.embed_model)
+    vectorstore = PineconeVectorStore(
+        index_name=cfg.datastore.index_name,
+        embedding=embeddings,
+    )
+    query = state["query"]
+    documents = vectorstore.similarity_search_with_score(query=query, k=cfg.model.top_k)
+    documents = _group_by_document(documents)
+    return {"documents": documents, "query": query}
+
+
+def grade_documents(state):
+    print("---CHECK DOCUMENT RELEVANCE TO QUERY---")
+    query = state["query"]
+    documents = state["documents"]
+
+    filtered_docs = []
+    for d in documents:
+        score = retrieval_grader.invoke({"query": query, "document": d.page_content})
+        grade = score.binary_score  # type: ignore
+        if grade == "yes":
+            print("---GRADE: DOCUMENT RELEVANT---")
+            filtered_docs.append(d)
+        else:
+            print("---GRADE: DOCUMENT NOT RELEVANT---")
+            continue
+    return {"documents": filtered_docs, "query": query}
+
+
+def generation(state):
+    query = state["query"]
+    document = state["document"]
+
+    generation = rag_chain.invoke({"query": query, "context": document})
+    return {"query": query, "document": document, "generation": generation}
+
+
+def grade_generation(state):
+    query = state["query"]
+    document = state["document"]
+    generation = state["generation"]
+
+    score = hallucination_grader.invoke(
+        {"document": document, "generation": generation}
+    )
+    grade = score.binary_score
+
+    if grade == "yes":
+        return {"query": query, "document": document, "generation": generation}
+    else:
+        return {"query": query, "document": document, "generation": "Hallucination"}
+
+
+def search_graph():
+    workflow = StateGraph(SearchState)
+    workflow.add_node("retrieve", retrieve)
+    # workflow.add_node("grade_documents", grade_documents)
+
+    workflow.add_edge(START, "retrieve")
+    workflow.add_edge("retrieve", END)
+    # workflow.add_edge("retrieve", "grade_documents")
+    # workflow.add_edge("grade_documents", END)
+    return workflow.compile()
+
+
+def generation_graph():
+    workflow = StateGraph(GenerationState)
+    workflow.add_node("gen", generation)
+    workflow.add_node("grade_generation", grade_generation)
+
+    workflow.add_edge(START, "gen")
+    workflow.add_edge("gen", "grade_generation")
+    workflow.add_edge("grade_generation", END)
+    return workflow.compile()
+
+
+search_app = search_graph()
+thread_id = 42
+q = "What is the capital of France?"
+
+out = search_app.invoke({"query": q}, config={"configurable": {"thread_id": thread_id}})
+[d.dict() for d in out["documents"]]

src/model/rag.py

@@ -0,0 +1,23 @@
+from dotenv import load_dotenv
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from src.common.settings import cfg
+
+_ = load_dotenv()
+human = """
+A user has queried a data catalogue, which has returned a relevant dataset.
+
+Explain the relevance of this dataset to the query in under three sentences. Use your own knowledge or the data profile. Do not say it is unrelated; attempt to find a relevant connection.
+
+Query: "{query}"
+
+Dataset description:
+
+{context}
+"""
+
+gen_prompt = ChatPromptTemplate.from_messages([("human", human)])
+llm = ChatOpenAI(model=cfg.model.llm, temperature=0)
+rag_chain = gen_prompt | llm | StrOutputParser()

src/search_api/api.py

@@ -0,0 +1,55 @@
+from uuid import UUID, uuid4
+
+from fastapi import FastAPI
+from langchain_core.documents import Document
+
+from src.model.model import generation_graph, search_graph
+
+app = FastAPI()
+document_store = {}
+query_mapping = {}
+
+
+search_app = search_graph()
+gen_app = generation_graph()
+
+
+@app.get("/")
+def index():
+    return {"message": "Make a post request to /query."}
+
+
+@app.post("/query")
+async def query(q: str) -> dict:
+    thread_id = uuid4()
+    out = search_app.invoke(
+        {"query": q}, config={"configurable": {"thread_id": thread_id}}
+    )
+    docs_dict = [d.dict() for d in out["documents"]]
+    document_store[thread_id] = docs_dict
+    query_mapping[thread_id] = q
+    return {"thread_id": thread_id, "query": q, "documents": docs_dict}
+
+
+@app.get("/explain/{thread_id}")
+async def explain(thread_id: UUID, docid: int) -> dict:
+    doc_dict = document_store[thread_id][docid]
+    document = Document(
+        page_content=doc_dict["page_content"],
+        metadata=doc_dict["metadata"],
+    )
+    query = query_mapping[thread_id]
+    generation_state = gen_app.invoke(
+        {"query": query, "document": document},
+        config={"configurable": {"thread_id": thread_id}},
+    )
+    generation = generation_state["generation"]
+
+    return {
+        "generation": generation,
+        "metadata": {
+            "thread_id": thread_id,
+            "query": query,
+            "related_dataset": doc_dict,
+        },
+    }

src/search_api/streamlit_app.py

@@ -0,0 +1,66 @@
+from subprocess import Popen
+from time import sleep
+
+import requests
+import streamlit as st
+
+
+def main():
+    st.title("CDRC Semantic Search App")
+
+    with st.spinner("Loading..."):
+        while True:
+            try:
+                r = requests.get("http://localhost:8000/")
+                if r.status_code == 200:
+                    break
+            except requests.exceptions.ConnectionError:
+                Popen(["uvicorn", "search_service.api:app", "--port", "8000"])
+                sleep(10)
+
+    # use_llm = st.toggle("Activate LLM")
+    use_llm = False
+    text = st.text_input("Query")
+    if text == "":
+        return None
+
+    r = requests.get(
+        "http://localhost:8000/query", params={"q": text, "use_llm": use_llm}
+    )
+    if r.status_code != 200:
+        st.error("No results :(")
+        return None
+
+    if use_llm:
+        response, metadata = r.json()
+        responses = []
+        for r in response["response"].split("---------------------"):
+            if all(x in r for x in ["Summary: ", "Relevance: "]):
+                responses.append(r)
+            else:
+                responses.append(None)
+
+        for res, meta in zip(responses, metadata.values(), strict=False):
+            st.subheader(meta["title"])
+            if res:
+                summary, relevance = res.split("Summary: ")[1].split("Relevance: ")
+                st.caption(summary)
+                st.caption(f":red[{relevance}]")
+            # st.caption(f"Score: :red[{meta['score']:.3f}]")
+            else:
+                st.caption("LLM did not return a response.")
+            if meta["url"] != "None":
+                st.write(meta["url"])
+            st.divider()
+    else:
+        metadata = r.json()
+        for meta in metadata:
+            st.subheader(meta["title"])
+            st.caption(f"Score: :red[{meta['score']:.3f}]")
+            if meta["url"] != "None":
+                st.write(meta["url"])
+            st.divider()
+
+
+if __name__ == "__main__":
+    main()

src/sensors.py

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