github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-24 15:43:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

community: Add partition parameter to DashVector (#19023)

Browse Source 
**Description**: DashVector Add partition parameter
**Twitter handle**: @CailinWang_

---------

Co-authored-by: root <root@Bluedot-AI>
This commit is contained in:
Cailin Wang 2024-03-17 06:20:30 +08:00 committed by GitHub
parent e64cf1aba4
commit 7cd87d2f6a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 38 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
libs/community/langchain_community/vectorstores/dashvector.py
View File

@ -66,16 +66,23 @@ class DashVector(VectorStore):
self._embedding = embedding self._embedding = embedding
self._text_field = text_field self._text_field = text_field
def _create_partition_if_not_exists(self, partition: str) -> None:
"""Create a Partition in current Collection."""
self._collection.create_partition(partition)
def _similarity_search_with_score_by_vector( def _similarity_search_with_score_by_vector(
self, self,
embedding: List[float], embedding: List[float],
k: int = 4, k: int = 4,
filter: Optional[str] = None, filter: Optional[str] = None,
partition: str = "default",
) -> List[Tuple[Document, float]]: ) -> List[Tuple[Document, float]]:
"""Return docs most similar to query vector, along with scores""" """Return docs most similar to query vector, along with scores"""
# query by vector # query by vector
ret = self._collection.query(embedding, topk=k, filter=filter) ret = self._collection.query(
embedding, topk=k, filter=filter, partition=partition
)
if not ret: if not ret:
raise ValueError( raise ValueError(
f"Fail to query docs by vector, error {self._collection.message}" f"Fail to query docs by vector, error {self._collection.message}"
@ -95,6 +102,7 @@ class DashVector(VectorStore):
metadatas: Optional[List[dict]] = None, metadatas: Optional[List[dict]] = None,
ids: Optional[List[str]] = None, ids: Optional[List[str]] = None,
batch_size: int = 25, batch_size: int = 25,
partition: str = "default",
**kwargs: Any, **kwargs: Any,
) -> List[str]: ) -> List[str]:
"""Run more texts through the embeddings and add to the vectorstore. """Run more texts through the embeddings and add to the vectorstore.
@ -104,11 +112,13 @@ class DashVector(VectorStore):
metadatas: Optional list of metadatas associated with the texts. metadatas: Optional list of metadatas associated with the texts.
ids: Optional list of ids associated with the texts. ids: Optional list of ids associated with the texts.
batch_size: Optional batch size to upsert docs. batch_size: Optional batch size to upsert docs.
partition: a partition name in collection. [optional].
kwargs: vectorstore specific parameters kwargs: vectorstore specific parameters
Returns: Returns:
List of ids from adding the texts into the vectorstore. List of ids from adding the texts into the vectorstore.
""" """
self._create_partition_if_not_exists(partition)
ids = ids or [str(uuid.uuid4().hex) for _ in texts] ids = ids or [str(uuid.uuid4().hex) for _ in texts]
text_list = list(texts) text_list = list(texts)
for i in range(0, len(text_list), batch_size): for i in range(0, len(text_list), batch_size):
@ -129,7 +139,7 @@ class DashVector(VectorStore):
# batch upsert to collection # batch upsert to collection
docs = list(zip(batch_ids, batch_embeddings, batch_metadatas)) docs = list(zip(batch_ids, batch_embeddings, batch_metadatas))
ret = self._collection.upsert(docs) ret = self._collection.upsert(docs, partition=partition)
if not ret: if not ret:
raise ValueError( raise ValueError(
f"Fail to upsert docs to dashvector vector database," f"Fail to upsert docs to dashvector vector database,"
@ -137,23 +147,27 @@ class DashVector(VectorStore):
) )
return ids return ids
def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> bool: def delete(
self, ids: Optional[List[str]] = None, partition: str = "default", **kwargs: Any
) -> bool:
"""Delete by vector ID. """Delete by vector ID.
Args: Args:
ids: List of ids to delete. ids: List of ids to delete.
partition: a partition name in collection. [optional].
Returns: Returns:
True if deletion is successful, True if deletion is successful,
False otherwise. False otherwise.
""" """
return bool(self._collection.delete(ids)) return bool(self._collection.delete(ids, partition=partition))
def similarity_search( def similarity_search(
self, self,
query: str, query: str,
k: int = 4, k: int = 4,
filter: Optional[str] = None, filter: Optional[str] = None,
partition: str = "default",
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Return docs most similar to query. """Return docs most similar to query.
@ -163,12 +177,15 @@ class DashVector(VectorStore):
k: Number of documents to return. Default to 4. k: Number of documents to return. Default to 4.
filter: Doc fields filter conditions that meet the SQL where clause filter: Doc fields filter conditions that meet the SQL where clause
specification. specification.
partition: a partition name in collection. [optional].
Returns: Returns:
List of Documents most similar to the query text. List of Documents most similar to the query text.
""" """
docs_and_scores = self.similarity_search_with_relevance_scores(query, k, filter) docs_and_scores = self.similarity_search_with_relevance_scores(
query, k, filter, partition
)
return [doc for doc, _ in docs_and_scores] return [doc for doc, _ in docs_and_scores]
def similarity_search_with_relevance_scores( def similarity_search_with_relevance_scores(
@ -176,6 +193,7 @@ class DashVector(VectorStore):
query: str, query: str,
k: int = 4, k: int = 4,
filter: Optional[str] = None, filter: Optional[str] = None,
partition: str = "default",
**kwargs: Any, **kwargs: Any,
) -> List[Tuple[Document, float]]: ) -> List[Tuple[Document, float]]:
"""Return docs most similar to query text , alone with relevance scores. """Return docs most similar to query text , alone with relevance scores.
@ -187,6 +205,7 @@ class DashVector(VectorStore):
k: Number of Documents to return. Defaults to 4. k: Number of Documents to return. Defaults to 4.
filter: Doc fields filter conditions that meet the SQL where clause filter: Doc fields filter conditions that meet the SQL where clause
specification. specification.
partition: a partition name in collection. [optional].
Returns: Returns:
List of Tuples of (doc, similarity_score) List of Tuples of (doc, similarity_score)
@ -194,7 +213,7 @@ class DashVector(VectorStore):
embedding = self._embedding.embed_query(query) embedding = self._embedding.embed_query(query)
return self._similarity_search_with_score_by_vector( return self._similarity_search_with_score_by_vector(
embedding, k=k, filter=filter embedding, k=k, filter=filter, partition=partition
) )
def similarity_search_by_vector( def similarity_search_by_vector(
@ -202,6 +221,7 @@ class DashVector(VectorStore):
embedding: List[float], embedding: List[float],
k: int = 4, k: int = 4,
filter: Optional[str] = None, filter: Optional[str] = None,
partition: str = "default",
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Return docs most similar to embedding vector. """Return docs most similar to embedding vector.
@ -211,12 +231,13 @@ class DashVector(VectorStore):
k: Number of Documents to return. Defaults to 4. k: Number of Documents to return. Defaults to 4.
filter: Doc fields filter conditions that meet the SQL where clause filter: Doc fields filter conditions that meet the SQL where clause
specification. specification.
partition: a partition name in collection. [optional].
Returns: Returns:
List of Documents most similar to the query vector. List of Documents most similar to the query vector.
""" """
docs_and_scores = self._similarity_search_with_score_by_vector( docs_and_scores = self._similarity_search_with_score_by_vector(
embedding, k, filter embedding, k, filter, partition
) )
return [doc for doc, _ in docs_and_scores] return [doc for doc, _ in docs_and_scores]
@ -227,6 +248,7 @@ class DashVector(VectorStore):
fetch_k: int = 20, fetch_k: int = 20,
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
filter: Optional[dict] = None, filter: Optional[dict] = None,
partition: str = "default",
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Return docs selected using the maximal marginal relevance. """Return docs selected using the maximal marginal relevance.
@ -244,13 +266,14 @@ class DashVector(VectorStore):
Defaults to 0.5. Defaults to 0.5.
filter: Doc fields filter conditions that meet the SQL where clause filter: Doc fields filter conditions that meet the SQL where clause
specification. specification.
partition: a partition name in collection. [optional].
Returns: Returns:
List of Documents selected by maximal marginal relevance. List of Documents selected by maximal marginal relevance.
""" """
embedding = self._embedding.embed_query(query) embedding = self._embedding.embed_query(query)
return self.max_marginal_relevance_search_by_vector( return self.max_marginal_relevance_search_by_vector(
embedding, k, fetch_k, lambda_mult, filter embedding, k, fetch_k, lambda_mult, filter, partition
) )
def max_marginal_relevance_search_by_vector( def max_marginal_relevance_search_by_vector(
@ -260,6 +283,7 @@ class DashVector(VectorStore):
fetch_k: int = 20, fetch_k: int = 20,
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
filter: Optional[dict] = None, filter: Optional[dict] = None,
partition: str = "default",
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
"""Return docs selected using the maximal marginal relevance. """Return docs selected using the maximal marginal relevance.
@ -277,6 +301,7 @@ class DashVector(VectorStore):
Defaults to 0.5. Defaults to 0.5.
filter: Doc fields filter conditions that meet the SQL where clause filter: Doc fields filter conditions that meet the SQL where clause
specification. specification.
partition: a partition name in collection. [optional].
Returns: Returns:
List of Documents selected by maximal marginal relevance. List of Documents selected by maximal marginal relevance.
@ -284,7 +309,11 @@ class DashVector(VectorStore):
# query by vector # query by vector
ret = self._collection.query( ret = self._collection.query(
embedding, topk=fetch_k, filter=filter, include_vector=True embedding,
topk=fetch_k,
filter=filter,
partition=partition,
include_vector=True,
) )
if not ret: if not ret:
raise ValueError( raise ValueError(