Creating a custom retriever with FAISS vector database that returns similarity scores

[Louva1]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

python

# currently initialising my retriever as follows
retriever = vectorstore.as_retriever(
            search_type=search_type,
            search_kwargs={'k': retrieval_param, 'fetch_k': 10, 'lambda': 0},
            )

# I am in need of something like this, yet this does not work for me since the function is dependent on a vectorstore input that is cannot handle right now. 
@chain 
def retriever(query: str) -> List[Document]:
    docs, scores = zip(*vectorstore.similarity_search_with_score(query))
    for doc, score in zip(docs, scores):
        doc.metadata["score"] = score
    return docs

Description

In my graph I am using a retriever as a tool, I already made a custom retriever tool for this. Currently this tool makes use of a standard retriever (see code; vectorstore.as_retriever). Given that I want to store the similarity scores I want to create a custom retriever as well. I found in the LangChain documentation an example with a @chain decorator. However, this does not work for me since I cannot pass the vectorstore to this function. Can someone help me out on how I can create a custom retriever that has these functionalities and will work in my retriever tool?

System Info

windows - python 3.12.4
langchain == 0.2.7
langchain-community == 0.2.7
langchain-core == 0.2.19
langchain-openai == 0.1.16
langchain-text-splitters == 0.2.2
langgraph == 0.1.8

[dosubot[bot]]

To create a custom retriever with a FAISS vector database that returns similarity scores and works in your retriever tool, follow these steps:

1. Install the necessary packages:

   pip install -qU langchain-community faiss-cpu

2. Initialize the FAISS vector store:

   import faiss
   from langchain_openai import OpenAIEmbeddings
   from langchain_community.docstore.in_memory import InMemoryDocstore
   from langchain_community.vectorstores import FAISS
   
   embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
   index = faiss.IndexFlatL2(len(embeddings.embed_query("hello world")))
   
   vector_store = FAISS(
       embedding_function=embeddings,
       index=index,
       docstore=InMemoryDocstore(),
       index_to_docstore_id={},
   )

3. Add documents to the vector store:

   from uuid import uuid4
   from langchain_core.documents import Document
   
   documents = [
       Document(page_content="I had chocolate chip pancakes and scrambled eggs for breakfast this morning.", metadata={"source": "tweet"}),
       Document(page_content="The weather forecast for tomorrow is cloudy and overcast, with a high of 62 degrees.", metadata={"source": "news"}),
       # Add more documents as needed
   ]
   uuids = [str(uuid4()) for _ in range(len(documents))]
   
   vector_store.add_documents(documents=documents, ids=uuids)

4. Query the vector store with similarity scores:

   results = vector_store.similarity_search_with_score(
       "Will it be hot tomorrow?", k=1, filter={"source": "news"}
   )
   for res, score in results:
       print(f"* [SIM={score:3f}] {res.page_content} [{res.metadata}]")

5. Transform the vector store into a retriever:

   from typing import List
   from langchain_core.runnables import chain
   
   @chain
   def retriever(query: str) -> List[Document]:
       docs, scores = zip(*vector_store.similarity_search_with_score(query))
       for doc, score in zip(docs, scores):
           doc.metadata["score"] = score
   
       return docs
   
   result = retriever.invoke("Stealing from the bank is a crime")
   for doc in result:
       print(f"Content: {doc.page_content}, Score: {doc.metadata['score']}")

This setup will allow you to create a custom retriever with a FAISS vector database that returns similarity scores and integrates with your retriever tool. You can also specify additional search parameters, such as threshold scores and top-k, to fine-tune the retrieval process. The similarity_search_with_score method in the FAISS vector store supports filtering by metadata and setting a score threshold, which can be useful for more refined searches ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[Louva1]

@hwchase17, the bot answer is not addressing the part of my question about passing the vector_store to the retriever function (I do not have it as a global variable). Can you suggest a solution?