Skip to content
This repository has been archived by the owner on Oct 15, 2024. It is now read-only.

Commit

Permalink
Adds embeddings filter option to return scores in state (langchain-ai…
Browse files Browse the repository at this point in the history
  • Loading branch information
jacoblee93 authored Nov 9, 2023
1 parent 18601bd commit 76283e9
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -67,4 +67,6 @@ def compress_documents(
similarity[included_idxs] > self.similarity_threshold
)
included_idxs = included_idxs[similar_enough]
for i in included_idxs:
stateful_documents[i].state["query_similarity_score"] = similarity[i]
return [stateful_documents[i] for i in included_idxs]
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@ def test_embeddings_filter_with_state() -> None:
state = {"embedded_doc": np.zeros(len(embedded_query))}
docs = [_DocumentWithState(page_content=t, state=state) for t in texts]
docs[-1].state = {"embedded_doc": embedded_query}
relevant_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.75)
relevant_filter = EmbeddingsFilter(
embeddings=embeddings, similarity_threshold=0.75, return_similarity_scores=True
)
actual = relevant_filter.compress_documents(docs, query)
assert len(actual) == 1
assert texts[-1] == actual[0].page_content

0 comments on commit 76283e9

Please sign in to comment.