Skip to content

feat: retrieve embeddings from database only when necessary #119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: retrieve embeddings from database only when necessary #119

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 31 additions & 9 deletions langchain_postgres/vectorstores.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1060,9 +1060,9 @@ def _results_to_docs_and_scores(self, results: Any) -> List[Tuple[Document, floa
docs = [
(
Document(
id=str(result.EmbeddingStore.id),
page_content=result.EmbeddingStore.document,
metadata=result.EmbeddingStore.cmetadata,
id=str(result.id),
page_content=result.document,
metadata=result.cmetadata,
),
result.distance if self.embeddings is not None else None,
)
Expand Down Expand Up @@ -1395,8 +1395,16 @@ def __query_collection(
embedding: List[float],
k: int = 4,
filter: Optional[Dict[str, str]] = None,
retrieve_embeddings: bool = False,
) -> Sequence[Any]:
"""Query the collection."""
columns_to_select = [
self.EmbeddingStore.id,
self.EmbeddingStore.document,
self.EmbeddingStore.cmetadata,
]
if retrieve_embeddings:
columns_to_select.append(self.EmbeddingStore.embedding)
with self._make_sync_session() as session: # type: ignore[arg-type]
collection = self.get_collection(session)
if not collection:
Expand All @@ -1417,7 +1425,7 @@ def __query_collection(

results: List[Any] = (
session.query(
self.EmbeddingStore,
*columns_to_select,
self.distance_strategy(embedding).label("distance"),
)
.filter(*filter_by)
Expand All @@ -1438,8 +1446,16 @@ async def __aquery_collection(
embedding: List[float],
k: int = 4,
filter: Optional[Dict[str, str]] = None,
retrieve_embeddings: bool = False,
) -> Sequence[Any]:
"""Query the collection."""
columns_to_select = [
self.EmbeddingStore.id,
self.EmbeddingStore.document,
self.EmbeddingStore.cmetadata,
]
if retrieve_embeddings:
columns_to_select.append(self.EmbeddingStore.embedding)
async with self._make_async_session() as session: # type: ignore[arg-type]
collection = await self.aget_collection(session)
if not collection:
Expand All @@ -1460,7 +1476,7 @@ async def __aquery_collection(

stmt = (
select(
self.EmbeddingStore,
*columns_to_select,
self.distance_strategy(embedding).label("distance"),
)
.filter(*filter_by)
Expand Down Expand Up @@ -1899,9 +1915,11 @@ def max_marginal_relevance_search_with_score_by_vector(
relevance to the query and score for each.
"""
assert not self._async_engine, "This method must be called without async_mode"
results = self.__query_collection(embedding=embedding, k=fetch_k, filter=filter)
results = self.__query_collection(
embedding=embedding, k=fetch_k, filter=filter, retrieve_embeddings=True
)

embedding_list = [result.EmbeddingStore.embedding for result in results]
embedding_list = [result.embedding for result in results]

mmr_selected = maximal_marginal_relevance(
np.array(embedding, dtype=np.float32),
Expand Down Expand Up @@ -1947,10 +1965,14 @@ async def amax_marginal_relevance_search_with_score_by_vector(
await self.__apost_init__() # Lazy async init
async with self._make_async_session() as session:
results = await self.__aquery_collection(
session=session, embedding=embedding, k=fetch_k, filter=filter
session=session,
embedding=embedding,
k=fetch_k,
filter=filter,
retrieve_embeddings=True,
)

embedding_list = [result.EmbeddingStore.embedding for result in results]
embedding_list = [result.embedding for result in results]

mmr_selected = maximal_marginal_relevance(
np.array(embedding, dtype=np.float32),
Expand Down