Vector database agnostic integration with LangChain

This notebook teaches you how to:

1. Apply a RAG framework by connecting a vector database to one of the watsonx foundation models and utility functions from the Watson Machine Learning service within watsonx.ai and Langchain,
2. Build up a knowledge base,
3. Create an embedding function to generate a Q&A resource for users

Set up the environment

Install and import the dependencies

!pip install "langchain==0.1.7" | tail -n 1
!pip install langchain_ibm | tail -n 1
!pip install wget | tail -n 1
!pip install sentence-transformers | tail -n 1
!pip install "chromadb==0.3.26" | tail -n 1
!pip install "ibm-watsonx-ai>=0.1.7" | tail -n 1
!pip install "pydantic==1.10.0" | tail -n 1
!pip install "sqlalchemy==2.0.1" | tail -n 1

Install vector database's Python library

You must choose your own vector database and update "vector store".

!pip install ["vector store"] | tail -n 1

Foundation models on watsonx.ai

Import IBM Foundation Models and utility functions from WML and Langchain, build up knowledge base, create embedding fuction

from langchain.embeddings import HuggingFaceEmbeddings

from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter

from langchain.chains import RetrievalQA

import os 

import getpass

import wget

Import vector database

You must choose your own vector database and update "VECTOR STORE".

from langchain.vectorstores import ["VECTOR STORE"]

watsonx API connection (insert WML api key)

This cell defines the credentials required to work with watsonx API for Foundation Model inferencing.

Action: Provide the IBM Cloud user API key. For details, see documentation: https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui

The API requires project id that provides the context for the call. We will obtain the id from the project in which this notebook runs. Otherwise, please provide the project id.

credentials = {
    "url": "https://us-south.ml.cloud.ibm.com",
    "apikey": getpass.getpass("Please enter your WML api key (hit enter): ")
}


try:
    project_id = os.environ["PROJECT_ID"] 
except KeyError:
    project_id = input("Please enter your project_id (hit enter): ")

Define model, define model parameteres, LangChain CustomLLM wrapper for watsonx model

from ibm_watsonx_ai.foundation_models.utils.enums import ModelTypes

model_id = ModelTypes.GRANITE_13B_CHAT_V2

### Define model behavior parameters

from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
from ibm_watsonx_ai.foundation_models.utils.enums import DecodingMethods

parameters = {
    GenParams.DECODING_METHOD: DecodingMethods.GREEDY,
    GenParams.MIN_NEW_TOKENS: 1,
    GenParams.MAX_NEW_TOKENS: 100,
    GenParams.STOP_SEQUENCES: ["<|endoftext|>"]
}

### Create an instance of your watsonx LLM
from langchain_ibm import WatsonxLLM

watsonx_granite = WatsonxLLM(
    model_id=model_id.value,
    url=credentials.get("url"),
    apikey=credentials.get("apikey"),
    project_id=project_id,
    params=parameters
)

Document data loading

filename = 'state_of_the_union.txt'
url = 'https://raw.github.com/IBM/watson-machine-learning-samples/master/cloud/data/foundation_models/state_of_the_union.txt'
if not os.path.isfile(filename):
    wget.download(url, out=filename)
loader = TextLoader(filename)
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)


embeddings = HuggingFaceEmbeddings()

Connecting to vector database from Langchain

Check Langchain documentation for each vector db (https://python.langchain.com/docs/integrations/vectorstores) for the specific formatting. Check whether it is an in memory db or cloud based db. If cloud based, user will need to create an account and obtain api key. See example below:

url = ["URL"]
api_key = ["API KEY"]
docsearch = "VECTOR STORE".from_documents(
    docs,
    embeddings,
    url=url,
    prefer_grpc=True,
    api_key=api_key,
    collection_name=["COLLECTION NAME"],
)

Generate a retrieval-augmented response to a question

Create RAG chain

qa = RetrievalQA.from_chain_type(llm=watsonx_granite, chain_type="stuff", retriever=docsearch.as_retriever())

Select questions

query = "What did the president say about Ketanji Brown Jackson"
qa.invoke(query)

Copyright © 2024 IBM. This notebook and its source code are released under the terms of the MIT License.