Caching LLMs Response With MongoDB Atlas and Vector Search

1
pip install fastapi pymongo openai uvicorn

1
from fastapi import FastAPI
2
server = FastAPI()
3
# root route
4
@server.get("/")
5
async def home():
6
  return { "message": "This is home server" }
7
8
# search route
9
@server.get("/ask")
10
async def search(query: str):
11
  return { "message": "The query is: {}".format(query) }

1
uvicorn app:server --reload

1
curl -X GET "http://127.0.0.1:8000/ask?query=hello+this+is+search+query"

1
{ "message": "The query is: hello this is search query" }

1
from openai import OpenAI
2
3
open_api_key = "..."  # OpenAI API Key
4
openai_client = OpenAI(api_key=open_api_key)
5
language_model = "gpt-3.5-turbo"
6
7
# getTextResponse receive text and ask LLM for answer
8
def getTextResponse(text):
9
	chat_completion = openai_client.chat.completions.create(
10
    	messages=[{"role": "user", "content": text}], model=language_model
11
	)
12
	return chat_completion.choices[0].message.content

1
# ... other import dependencies
2
from llm import getTextResponse
3
# ... other routes
4
# search route
5
@server.get("/ask")
6
async def search(query):
7
	llm_response = getTextResponse(query)
8
	return {"message": "Your AI response is: {}".format(llm_response)}

1
curl -X GET "http://127.0.0.1:8000/ask?query=what+is+llm?"
2
Response
3
{
4
  "message": "Your AI response is: LLM stands for Master of Laws, which is an advanced law degree typically pursued by individuals who have already received a law degree (such as a JD) and want to further specialize or advance their knowledge in a specific area of law. LLM programs are typically one year in length and often focus on areas such as international law, human rights, or commercial law."
5
}

1
text = "Large Language Model"
2
# embedding_function(text: str) -> vector<number>([...N])
3
vector = embedding_function(text)
4
5
# vector = [12, 23, 0.11, 22, 85, 43, ..., 90]

1
# ... other code
2
text_embedding_model = "text-embedding-3-small"
3
# getEmbedding receives text and response embedding (vector) of its original data
4
def getEmbedding(text):
5
	embedding = openai_client.embeddings.create(input=text, model=text_embedding_model)
6
	return embedding.data[0].embedding

1
# ... other import dependencies
2
from llm import getTextResponse, getEmbedding
3
4
# ... other routes
5
6
# search route
7
@server.get("/ask")
8
async def search(query):
9
	llm_response = getTextResponse(query)
10
	query_vector = getEmbedding(query)
11
	print("embedding : ", query_vector)
12
	return {"message": "Your AI response is: {}".format(llm_response)}

1
embedding :  [-0.02254991, 0.031336114, 0.019013261, 0.00017081834, -0.0202526, -0.0020466715, 0.0111036645, -0.0111036645, 0.036172554, 0.04038429, -0.027043771, 0.0046273666, -0.039820038, -0.011456322, 0.0048339227, 0.021824444, 0.0048666694, -0.017501874, 0.03915503, 0.03895351, 0.041311275, ..., 0.046349235]

1
import pymongo
2
3
MONGO_URI = ""  # MongoDB connection string
4
mongo_client = pymongo.MongoClient(MONGO_URI)
5
6
db = mongo_client.get_database("logging")  # database name in mongodb
7
collection = db.get_collection("test")  # collection name
8
9
# document {
10
#   query: string,
11
#   response: string,
12
#   embeddings: vector<number<1408>>
13
# }
14
15
def save_cache(document):
16
	collection.insert_one(document)

1
# ... other import dependencies
2
from db import save_cache
3
4
# ... other code
5
6
# search route
7
@server.get("/ask")
8
async def search(query):
9
	llm_response = getTextResponse(query)
10
	query_vector = getEmbedding(query)
11
12
	document = {
13
  	"response": llm_response,
14
  	"embeddings": query_vector,
15
  	"query": query
16
  	}
17
	save_cache(document)
18
19
	return {"message": "Your AI response is: {}".format(llm_response)}

1
{
2
  "type": "vectorSearch", # vector search identity
3
  "fields": [
4
	{
5
  	# type of data in the asking the system again with a new question (but similar meaning): Howfield (vector information)
6
  	"type": "vector",
7
  	# field of document that store vector information
8
  	"path": "embeddings",
9
  	# dimension of vector, receiving from language model specification.
10
  	# for `text-embedding-3-small` of OpenAI, it's 1408 dimensions of vector.
11
  	"numDimensions": 1408,
12
  	# vector search find similarity of searching document with other. So, closest documents would get high score.
13
  	"similarity": "cosine"
14
	}
15
  ]
16
}

1
# Step 1: Perform vector search
2
{
3
  # $vectorSearch operation
4
    "$vectorSearch": {
5
        # Vector Search index in Atlas
6
        "index": "vector_index",
7
        # vector information in document
8
        "path": "embeddings",
9
        # embedding information of client's query
10
        "queryVector": numerical_embedding,
11
       # number of nearest neighbors to find the closest group. This value is used for 'searchScore' ranking
12
        "numCandidates": 20,
13
           # number of document to return ranked by semantic meaning / searchScore
14
        "limit": 5,
15
    },
16
},
17
# add 'score' field to show vector search score.
18
{
19
  "$addFields": {
20
  "score": {
21
    # vectorSearchScore is attached in $meta for pipeline after vectorSearch operation
22
    "$meta": "vectorSearchScore"
23
    }
24
  }
25
},
26
# return document that score greater than specific ratio
27
# in this case, only top one that has ration > 70% is going to return.
28
{
29
  "$match": {
30
    "score": { "$gte": 0.70 }
31
  }
32
},
33
# remove `embeddings` field. Because it reduce load to retrieve data.
34
{
35
  "$unset": ["embeddings"]
36
}
37
]

1
# ... other code
2
def perform_search_cache(query):
3
	pipeline = [
4
    	{
5
        	"$vectorSearch": {
6
            	"index": "vector_index",
7
            	"path": "embeddings",
8
            	"queryVector": query,
9
            	"numCandidates": 20,
10
            	"limit": 1,
11
        	},
12
    	},
13
    	{"$addFields": {"score": {"$meta": "vectorSearchScore"}}},
14
    	{"$match": {"score": {"$gte": 0.70}}},
15
    	{"$unset": ["embeddings"]},
16
	]
17
	result = collection.aggregate(pipeline)
18
	return result

1
# ... other import dependencies
2
from db import save_cache, perform_search_cache
3
4
# ... other code
5
# search route
6
@server.get("/ask")
7
async def search(query):
8
	query_vector = getEmbedding(query)
9
	cache_response = list(perform_search_cache(query_vector))
10
	response = ""
11
	if len(cache_response) < 1:
12
    	# in case no cache hit
13
    	# perform search and save the cache
14
    	llm_response = getTextResponse(query)
15
    	document = {"response": llm_response, "embeddings": query_vector, "query": query}
16
    	save_cache(document)
17
    	response = llm_response
18
	else:
19
    	# if cache hit, return it
20
    	response = cache_response[0]["response"]
21
    	print("cache hit")
22
    	print(cache_response)
23
	return {"message": "Your AI response is: {}".format(response)}

1
curl -X GET "http://127.0.0.1:8000/ask?query=how+are+things+with+you?"

1
{
2
  "message": "Your AI response is: I'm just a computer program, so I don't have feelings or emotions. But I'm here to help you with anything you need! How can I assist you today?"
3
}

1
curl -X GET "http://127.0.0.1:8000/ask?query=how+are+you+today?"

1
{
2
  "message": "Your AI response is: I'm just a computer program, so I don't have feelings or emotions. But I'm here to help you with anything you need! How can I assist you today?"
3
}

1
cache hit
2
[
3
  {
4
	'_id': ObjectId('6671440cb2bf0b0eb12b75b3'),
5
	'response': "I'm just a computer program, so I don't have feelings or emotions. But I'm here to help you with anything you need! How can I assist you today?",
6
	'query': 'how are thing with you?',
7
	'score': 0.8066450357437134
8
  }
9
]