from __future__ import annotations
import logging
import uuid
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.utils import get_from_dict_or_env
from langchain_core.vectorstores import VectorStore
ADA_TOKEN_COUNT = 1536
_LANGCHAIN_DEFAULT_TABLE_NAME = "langchain_pg_embedding"
[docs]
class Hologres(VectorStore):
"""`Hologres API` vector store.
- `connection_string` is a hologres connection string.
- `embedding_function` any embedding function implementing
`langchain.embeddings.base.Embeddings` interface.
- `ndims` is the number of dimensions of the embedding output.
- `table_name` is the name of the table to store embeddings and data.
(default: langchain_pg_embedding)
- NOTE: The table will be created when initializing the store (if not exists)
So, make sure the user has the right permissions to create tables.
- `pre_delete_table` if True, will delete the table if it exists.
(default: False)
- Useful for testing.
"""
[docs]
def __init__(
self,
connection_string: str,
embedding_function: Embeddings,
ndims: int = ADA_TOKEN_COUNT,
table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,
pre_delete_table: bool = False,
logger: Optional[logging.Logger] = None,
) -> None:
self.connection_string = connection_string
self.ndims = ndims
self.table_name = table_name
self.embedding_function = embedding_function
self.pre_delete_table = pre_delete_table
self.logger = logger or logging.getLogger(__name__)
self.__post_init__()
def __post_init__(
self,
) -> None:
"""
Initialize the store.
"""
from hologres_vector import HologresVector
self.storage = HologresVector(
self.connection_string,
ndims=self.ndims,
table_name=self.table_name,
table_schema={"document": "text"},
pre_delete_table=self.pre_delete_table,
)
@property
def embeddings(self) -> Embeddings:
return self.embedding_function
@classmethod
def __from(
cls,
texts: List[str],
embeddings: List[List[float]],
embedding_function: Embeddings,
metadatas: Optional[List[dict]] = None,
ids: Optional[List[str]] = None,
ndims: int = ADA_TOKEN_COUNT,
table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,
pre_delete_table: bool = False,
**kwargs: Any,
) -> Hologres:
if ids is None:
ids = [str(uuid.uuid4()) for _ in texts]
if not metadatas:
metadatas = [{} for _ in texts]
connection_string = cls.get_connection_string(kwargs)
store = cls(
connection_string=connection_string,
embedding_function=embedding_function,
ndims=ndims,
table_name=table_name,
pre_delete_table=pre_delete_table,
)
store.add_embeddings(
texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs
)
return store
[docs]
def add_embeddings(
self,
texts: Iterable[str],
embeddings: List[List[float]],
metadatas: List[dict],
ids: List[str],
**kwargs: Any,
) -> None:
"""Add embeddings to the vectorstore.
Args:
texts: Iterable of strings to add to the vectorstore.
embeddings: List of list of embedding vectors.
metadatas: List of metadatas associated with the texts.
kwargs: vectorstore specific parameters
"""
try:
schema_datas = [{"document": t} for t in texts]
self.storage.upsert_vectors(embeddings, ids, metadatas, schema_datas)
except Exception as e:
self.logger.exception(e)
[docs]
def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
ids: Optional[List[str]] = None,
**kwargs: Any,
) -> List[str]:
"""Run more texts through the embeddings and add to the vectorstore.
Args:
texts: Iterable of strings to add to the vectorstore.
metadatas: Optional list of metadatas associated with the texts.
kwargs: vectorstore specific parameters
Returns:
List of ids from adding the texts into the vectorstore.
"""
if ids is None:
ids = [str(uuid.uuid4()) for _ in texts]
embeddings = self.embedding_function.embed_documents(list(texts))
if not metadatas:
metadatas = [{} for _ in texts]
self.add_embeddings(texts, embeddings, metadatas, ids, **kwargs)
return ids
[docs]
def similarity_search(
self,
query: str,
k: int = 4,
filter: Optional[dict] = None,
**kwargs: Any,
) -> List[Document]:
"""Run similarity search with Hologres with distance.
Args:
query (str): Query text to search for.
k (int): Number of results to return. Defaults to 4.
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
Returns:
List of Documents most similar to the query.
"""
embedding = self.embedding_function.embed_query(text=query)
return self.similarity_search_by_vector(
embedding=embedding,
k=k,
filter=filter,
)
[docs]
def similarity_search_by_vector(
self,
embedding: List[float],
k: int = 4,
filter: Optional[dict] = None,
**kwargs: Any,
) -> List[Document]:
"""Return docs most similar to embedding vector.
Args:
embedding: Embedding to look up documents similar to.
k: Number of Documents to return. Defaults to 4.
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
Returns:
List of Documents most similar to the query vector.
"""
docs_and_scores = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k, filter=filter
)
return [doc for doc, _ in docs_and_scores]
[docs]
def similarity_search_with_score(
self,
query: str,
k: int = 4,
filter: Optional[dict] = None,
) -> List[Tuple[Document, float]]:
"""Return docs most similar to query.
Args:
query: Text to look up documents similar to.
k: Number of Documents to return. Defaults to 4.
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
Returns:
List of Documents most similar to the query and score for each
"""
embedding = self.embedding_function.embed_query(query)
docs = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k, filter=filter
)
return docs
[docs]
def similarity_search_with_score_by_vector(
self,
embedding: List[float],
k: int = 4,
filter: Optional[dict] = None,
) -> List[Tuple[Document, float]]:
results: List[dict[str, Any]] = self.storage.search(
embedding, k=k, select_columns=["document"], metadata_filters=filter
)
docs = [
(
Document(
page_content=result["document"],
metadata=result["metadata"],
),
result["distance"],
)
for result in results
]
return docs
[docs]
@classmethod
def from_texts(
cls: Type[Hologres],
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
ndims: int = ADA_TOKEN_COUNT,
table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,
ids: Optional[List[str]] = None,
pre_delete_table: bool = False,
**kwargs: Any,
) -> Hologres:
"""
Return VectorStore initialized from texts and embeddings.
Hologres connection string is required
"Either pass it as a parameter
or set the HOLOGRES_CONNECTION_STRING environment variable.
Create the connection string by calling
HologresVector.connection_string_from_db_params
"""
embeddings = embedding.embed_documents(list(texts))
return cls.__from(
texts,
embeddings,
embedding,
metadatas=metadatas,
ids=ids,
ndims=ndims,
table_name=table_name,
pre_delete_table=pre_delete_table,
**kwargs,
)
[docs]
@classmethod
def from_embeddings(
cls,
text_embeddings: List[Tuple[str, List[float]]],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
ndims: int = ADA_TOKEN_COUNT,
table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,
ids: Optional[List[str]] = None,
pre_delete_table: bool = False,
**kwargs: Any,
) -> Hologres:
"""Construct Hologres wrapper from raw documents and pre-
generated embeddings.
Return VectorStore initialized from documents and embeddings.
Hologres connection string is required
"Either pass it as a parameter
or set the HOLOGRES_CONNECTION_STRING environment variable.
Create the connection string by calling
HologresVector.connection_string_from_db_params
Example:
.. code-block:: python
from langchain_community.vectorstores import Hologres
from langchain_community.embeddings import OpenAIEmbeddings
embeddings = OpenAIEmbeddings()
text_embeddings = embeddings.embed_documents(texts)
text_embedding_pairs = list(zip(texts, text_embeddings))
faiss = Hologres.from_embeddings(text_embedding_pairs, embeddings)
"""
texts = [t[0] for t in text_embeddings]
embeddings = [t[1] for t in text_embeddings]
return cls.__from(
texts,
embeddings,
embedding,
metadatas=metadatas,
ids=ids,
ndims=ndims,
table_name=table_name,
pre_delete_table=pre_delete_table,
**kwargs,
)
[docs]
@classmethod
def from_existing_index(
cls: Type[Hologres],
embedding: Embeddings,
ndims: int = ADA_TOKEN_COUNT,
table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,
pre_delete_table: bool = False,
**kwargs: Any,
) -> Hologres:
"""
Get instance of an existing Hologres store.This method will
return the instance of the store without inserting any new
embeddings
"""
connection_string = cls.get_connection_string(kwargs)
store = cls(
connection_string=connection_string,
ndims=ndims,
table_name=table_name,
embedding_function=embedding,
pre_delete_table=pre_delete_table,
)
return store
[docs]
@classmethod
def get_connection_string(cls, kwargs: Dict[str, Any]) -> str:
connection_string: str = get_from_dict_or_env(
data=kwargs,
key="connection_string",
env_key="HOLOGRES_CONNECTION_STRING",
)
if not connection_string:
raise ValueError(
"Hologres connection string is required"
"Either pass it as a parameter"
"or set the HOLOGRES_CONNECTION_STRING environment variable."
"Create the connection string by calling"
"HologresVector.connection_string_from_db_params"
)
return connection_string
[docs]
@classmethod
def from_documents(
cls: Type[Hologres],
documents: List[Document],
embedding: Embeddings,
ndims: int = ADA_TOKEN_COUNT,
table_name: str = _LANGCHAIN_DEFAULT_TABLE_NAME,
ids: Optional[List[str]] = None,
pre_delete_collection: bool = False,
**kwargs: Any,
) -> Hologres:
"""
Return VectorStore initialized from documents and embeddings.
Hologres connection string is required
"Either pass it as a parameter
or set the HOLOGRES_CONNECTION_STRING environment variable.
Create the connection string by calling
HologresVector.connection_string_from_db_params
"""
texts = [d.page_content for d in documents]
metadatas = [d.metadata for d in documents]
connection_string = cls.get_connection_string(kwargs)
kwargs["connection_string"] = connection_string
return cls.from_texts(
texts=texts,
pre_delete_collection=pre_delete_collection,
embedding=embedding,
metadatas=metadatas,
ids=ids,
ndims=ndims,
table_name=table_name,
**kwargs,
)
[docs]
@classmethod
def connection_string_from_db_params(
cls,
host: str,
port: int,
database: str,
user: str,
password: str,
) -> str:
"""Return connection string from database parameters."""
return (
f"dbname={database} user={user} password={password} host={host} port={port}"
)
