Refactor logic to support local mode with BedrockChat and improve code quality

Author: kris-szlapa

README.md

@@ -72,6 +72,19 @@ You will now be able to use AWS and SAM CLI commands to access the dev account.
 
 When the token expires, you may need to reauthorise using `make aws-login`
 
+### Running locally
+
+If you want to test the EPS query functionality without starting the Slack bot, you can run the application in local mode. This is useful for development, debugging, or when Slack credentials are not available.
+
+To do this, set the `LOCAL_MODE` environment variable to `1`:
+
+```bash
+export LOCAL_MODE=1
+poetry run python packages/slackbot/app.py
+```
+
+For running the query tool only, see the [querytool instructions](packages/querytool/README.md).
+
 ### CI Setup
 
 The GitHub Actions require a secret to exist on the repo called "SONAR_TOKEN".

packages/querytool/README.md

@@ -1,41 +1,92 @@
-# Query tool
+# Query Tool
 
-This is the tool that augments incoming user queries with EPS specific data.
+This module powers EPS Assist's ability to interpret technical queries by leveraging EPS documentation and returning answers generated from retrieved context.
+
+It uses LangChain, semantic embeddings, and DuckDB to provide retrieval-augmented generation (RAG) for queries related to EPS APIs, SCAL, and other NHS documentation.
 
 ## Prerequisites
 
-- Python 3.12
-- Poetry
-- .env file 
+All runtime and development dependencies (Python, Node, Poetry, Widdershins, etc.) are installed automatically when you open the project in the devcontainer.
+
+Environment variables required for authentication are managed via `.envrc`.
+
+Ensure your .envrc file includes the following:
+
+```bash
+export TOKENIZERS_PARALLELISM=false
+export BEDROCK_MODEL_ID=<your_bedrock_model_id>
+export AWS_BEARER_TOKEN_BEDROCK=<your_bearer_token>
+```
+
+If you're working outside the devcontainer, you can install all dependencies manually by running:
+
+```bash
+make install
+```
+
+## Preparing the Document Corpus
+
+Before the assistant can answer queries, the EPS documentation must be parsed and embedded into a searchable vector store (DuckDB).
+
+### 1. Convert SCAL CSVs to Markdown
+
+```bash
+poetry run python packages/querytool/eps_assist/preprocessors/prepare_scal.py
+```
+
+### 2. Convert OAS JSON to Markdown
+
+This pulls the latest NHS OpenAPI specification and converts it to Markdown using Widdershins.
 
-Widdershins:
+> **Note**: To suppress noisy Node.js warnings during the conversion, use the `NODE_NO_WARNINGS=1` environment variable as shown below.
 
-    npm install -g widdershins
+```bash
+NODE_NO_WARNINGS=1 poetry run python packages/querytool/eps_assist/preprocessors/prepare_oas.py
+```
 
-Load environment variables:
+### 3. Build or Rebuild the Vector Store
 
-    source .env
+This loads all Markdown files into DuckDB with semantic chunking and embeddings.
 
-To set up run:
+```bash
+poetry run python packages/querytool/eps_assist/transform.py
+```
 
-    poetry install
+A new `eps_corpus.db` file will be created in the same directory.
 
-## Updating corpus
+## Running Queries Locally
 
-To prepare the SCAL files for processing, run:
+You can run sample questions against the vector store directly:
 
-    poetry run python querytool/eps_assist/preprocessors/prepare_scal.py
+```bash
+poetry run python packages/querytool/eps_assist/query.py
+```
 
-To prepare the OAS file for processing, run:
+This script:
+- Connects to `eps_corpus.db`
+- Retrieves relevant document chunks
+- Sends the prompt to Claude 3 via Amazon Bedrock
+- Outputs the model's answer in your terminal
 
-    poetry run python querytool/eps_assist/preprocessors/prepare_oas.py
+## Notes
 
-To run the ingestion and transformation of documents into the vector store, run:
+- Ensure that `eps_corpus.db` exists before querying. If in doubt, re-run the transformation step.
+- Claude 3 is accessed using the AWS Bedrock API via `boto3` and LangChain.
+- Vector storage is file-based (DuckDB), so no external database or service is required.
+- Environment variables (e.g., AWS credentials) are expected to be managed via `.envrc` in the project root.
 
-    poetry run python querytool/eps_assist/transform.py
+## File Structure Overview
 
-## Running samples queries
+```
+eps_assist/
+├── docs/               # Source documentation (.md) for SCAL, OAS, etc.
+├── preprocessors/      # Scripts for cleaning and converting raw files
+├── query.py            # Executes a full question-answering example
+├── transform.py        # Converts docs to vector store (DuckDB)
 
-To run a query, run:
+preprocessors/
+├── prepare_scal.py     # Converts SCAL CSV to Markdown
+├── prepare_oas.py      # Fetches & converts OpenAPI to Markdown
+```
 
-    poetry run python querytool/eps_assist/query.py
+This module is a self-contained tool that can also be used outside of Slack for testing or integration in other EPS-related projects.

packages/querytool/eps_assist/preprocessors/prepare_oas.py

@@ -4,17 +4,17 @@
 oas_url = "https://digital.nhs.uk/restapi/oas/324177"
 oas_content = requests.get(oas_url)
 
-with open("./querytool/eps_assist/preprocessors/.eps_oas.json", "w") as f:
+with open("packages/querytool/eps_assist/preprocessors/.eps_oas.json", "w") as f:
     f.write(oas_content.text)
 
 oas_version = oas_content.json()["info"]["version"]
 
-with open("./querytool/eps_assist/docs/eps_oas.version", "w") as f:
+with open("packages/querytool/eps_assist/docs/eps_oas.version", "w") as f:
     f.write(oas_version)
 
 subprocess.check_output(
     ["widdershins",
      "--expandBody",
      "true",
-     "querytool/eps_assist/preprocessors/.eps_oas.json",
-     "querytool/eps_assist/docs/eps_output.md"])
+     "packages/querytool/eps_assist/preprocessors/.eps_oas.json",
+     "packages/querytool/eps_assist/docs/eps_output.md"])

packages/querytool/eps_assist/preprocessors/prepare_scal.py

@@ -1,10 +1,7 @@
 import csv
 
-file_paths = ["./querytool/eps_assist/preprocessors/prescribing_scal.csv",
-              "./querytool/eps_assist/preprocessors/dispensing_scal.csv"]
-
-
-file_paths = ["./querytool/eps_assist/preprocessors/prescribing_scal.csv"]
+file_paths = ["packages/querytool/eps_assist/preprocessors/prescribing_scal.csv",
+              "packages/querytool/eps_assist/preprocessors/dispensing_scal.csv"]
 
 
 def clean_texts(texts: list[str]) -> str:
@@ -50,7 +47,17 @@ def process_file(path: str) -> str:
 
                     doc.append(
                         clean_texts(
-                            [section_id, item, detail, ". Related docs: ", helpful_docs, ". Risk Logs: ", risk_logs, ". Requirement assessed by: ", assessment_type]
+                            [
+                                section_id,
+                                item,
+                                detail,
+                                ". Related docs: ",
+                                helpful_docs,
+                                ". Risk Logs: ",
+                                risk_logs,
+                                ". Requirement assessed by: ",
+                                assessment_type,
+                            ]
                         )
                     )
                 continue
@@ -74,11 +81,13 @@ def process_file(path: str) -> str:
             else:
                 if len(related_desc) > 0:
                     doc.append(
-                        f"SCAL Requirement {section_id} {clean_text(requirement_section)}: {clean_text(requirement_or_section)} related info: {bullet_points_to_sentences(related_desc)}"
+                        f"SCAL Requirement {section_id} {clean_text(requirement_section)}: "
+                        f"{clean_text(requirement_or_section)} related info: {clean_text(related_desc)}"
                     )
                 else:
                     doc.append(
-                        f"SCAL Requirement {section_id} {clean_text(requirement_section)}: {clean_text(requirement_or_section)}"
+                        f"SCAL Requirement {section_id} {clean_text(requirement_section)}: "
+                        f"{clean_text(requirement_or_section)}"
                     )
 
     return "\n\n".join(doc)

packages/querytool/eps_assist/query.py

@@ -1,62 +1,65 @@
-# LLM
+import duckdb
+import os
+from pathlib import Path
+
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 from langchain_community.vectorstores import DuckDB
-
-# LLM
-from langchain_openai import AzureChatOpenAI
+from langchain_community.chat_models import BedrockChat
 from langchain.callbacks.manager import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-
-# QA chain
 from langchain.chains import RetrievalQA
-from langchain import hub
 from langchain.prompts import PromptTemplate
 
-import duckdb
-import os
-
+# Load the model ID for Amazon Bedrock from
+model_id = os.getenv("BEDROCK_MODEL_ID")
+if not model_id:
+    raise EnvironmentError("BEDROCK_MODEL_ID environment variable is not set.")
 
+# Set up embeddings
 embedding_function = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
 
-DB_PATH = "./eps_corpus.db"
-
+# Connect to the local DuckDB vector store
+DB_PATH = Path(__file__).resolve().parent / "eps_corpus.db"
 if os.path.exists(DB_PATH):
-    print(f"connecting to existing db ({DB_PATH})")
+    print(f"Connecting to existing db ({DB_PATH})")
     conn = duckdb.connect(DB_PATH)
     vector_store = DuckDB(connection=conn, embedding=embedding_function)
 
 else:
-    # load into database
-    raise Exception(f"db was not found ({DB_PATH})")
-
+    # Load into database
+    raise FileNotFoundError(f"DB not found at ({DB_PATH})")
 
-prompt = hub.pull("rlm/rag-prompt")
-
-llm = AzureChatOpenAI(
-    openai_api_version="2023-08-01-preview",
-    azure_deployment="eps-assistant-model",
-    callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
-    model="gpt-4",
+# Load Claude via Amazon Bedrock
+llm = BedrockChat(
+    model_id=model_id,
+    callback_manager=CallbackManager([StreamingStdOutCallbackHandler()])
 )
 
-
-question = """Are Rate Limits applied to the EPS for Dispenser API or any part thereof? If so, what are they and where do they apply, please?"""
-
+# Build the prompt
 prompt = PromptTemplate.from_template(
-    """[INST]<<SYS>> You are an assistant for question-answering tasks relating to the Electronic Prescribing Services EPS API.
-    Use the following pieces of retrieved context to answer the question.
-    Rate limits apply to all APIs (e.g. EPS dispensing, EPS prescribing, PDS)
-    If you don't know the answer, just say that you don't know. keep the answer concise and include technical references where possible
-    code samples should be used to support the answer where appropriate.<</SYS>> \nQuestion: {question} \nContext: {context} \n  Answer: [/INST]"""
+    """[INST]<<SYS>> You are an assistant for question-answering tasks relating to the Electronic Prescribing
+    Services EPS API. Use the following pieces of retrieved context to answer the question. Rate limits apply
+    to all APIs (e.g. EPS dispensing, EPS prescribing, PDS). If you don't know the answer, just say that you
+    don't know. Keep the answer concise and include technical references where possible. Code samples should
+    be used to support the answer where appropriate.<</SYS>> \nQuestion: {question} \nContext: {context} \n
+    Answer: [/INST]"""
 )
 
+# Create the Retrieval QA chain
 qa_chain = RetrievalQA.from_chain_type(
-    llm,
+    llm=llm,
     retriever=vector_store.as_retriever(),
-    chain_type_kwargs={"prompt": prompt},
+    chain_type_kwargs={"prompt": prompt}
 )
 
-result = qa_chain.invoke(question)
+# Define the question
+question = (
+    "Are Rate Limits applied to the EPS for Dispenser API or any part thereof? If so, what are they and "
+    "where do they apply, please?"
+)
 
+# Run the chain
+result = qa_chain.invoke(question)
 
+# Output the result
 print(result["result"])

packages/querytool/eps_assist/transform.py

@@ -1,19 +1,27 @@
 import duckdb
 import os
+from pathlib import Path
 
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 from langchain.docstore.document import Document
 from langchain_community.vectorstores import DuckDB
 from semantic_text_splitter import MarkdownSplitter
 from tokenizers import Tokenizer
 
-tokenizer = Tokenizer.from_pretrained("bert-base-uncased")
-splitter = MarkdownSplitter.from_huggingface_tokenizer(tokenizer)
+# Set up the base directory and paths
+BASE_DIR = Path(__file__).resolve().parent
+DB_PATH = BASE_DIR / "eps_corpus.db"
+CORPUS_PATH = BASE_DIR / "docs"
 
+# Toggle this flag to control DB rebuilding
+REBUILD_DB = True
+
+# Set up embeddings
 embedding_function = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
 
-DB_PATH = "./eps_corpus.db"
-CORPUS_PATH = "./querytool/eps_assist/docs/"
+# Set up splitter
+tokenizer = Tokenizer.from_pretrained("bert-base-uncased")
+splitter = MarkdownSplitter.from_huggingface_tokenizer(tokenizer)
 
 
 def connect_to_existing_vector_store():
@@ -30,8 +38,7 @@ def create_vector_store_file():
     vector_store = DuckDB(connection=conn, embedding=embedding_function)
 
     for file in os.listdir(CORPUS_PATH):
-
-        file_path = f"{CORPUS_PATH}{file}"
+        file_path = CORPUS_PATH / file
 
         with open(file_path) as doc:
             doc_text = doc.read()
@@ -42,7 +49,6 @@ def create_vector_store_file():
             #    chunks = doc_text.split("SCAL requirement")
             # else:
             chunks = splitter.chunks(doc_text, chunk_capacity=(200, 1000))
-
             docs = [Document(page_content=chunk) for chunk in chunks]
 
             print(f"adding {len(docs)} documents to vector store...")
@@ -51,8 +57,8 @@ def create_vector_store_file():
 
 if __name__ == "__main__":
 
-    # normally we just want to recreate the file, remove this if you want to test queries
-    if True and os.path.exists(DB_PATH):
+    if REBUILD_DB and os.path.exists(DB_PATH):
+        print(f"removing existing db file ({DB_PATH})")
         os.remove(DB_PATH)
 
     if os.path.exists(DB_PATH):
@@ -61,13 +67,17 @@ def create_vector_store_file():
         create_vector_store_file()
         vector_store = connect_to_existing_vector_store()
 
-    results = vector_store.similarity_search(
-        """1.6.5 “For eRD prescriptions, the dispenser must see:
-• the current issue
-• the total number of authorised issues for both the prescription and line items on the prescription.
-There is nothing we can see at prescription level that defines the current issue or the number of authorised issues (we can only see it at MedicationRequest level) Eg: Prescription number: 8F4A22-C81007-000012"""
+    test_query = (
+        "1.6.5 “For eRD prescriptions, the dispenser must see:\n"
+        "• the current issue\n"
+        "• the total number of authorised issues for both the prescription and line items on the prescription.\n"
+        "There is nothing we can see at prescription level that defines the current issue or the number of "
+        "authorised issues (we can only see it at MedicationRequest level)\n"
+        "Eg: Prescription number: 8F4A22-C81007-000012"
     )
 
+    results = vector_store.similarity_search(test_query)
+
     for result in results:
         print("*" * 100)
         print(result)