diff --git a/docs/modules/indexes/vectorstores/examples/hologres.ipynb b/docs/modules/indexes/vectorstores/examples/hologres.ipynb new file mode 100644 index 00000000..1d671cd6 --- /dev/null +++ b/docs/modules/indexes/vectorstores/examples/hologres.ipynb @@ -0,0 +1,157 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Hologres\n", + "\n", + ">[Hologres](https://www.alibabacloud.com/help/en/hologres/latest/introduction) is a unified real-time data warehousing service developed by Alibaba Cloud. You can use Hologres to write, update, process, and analyze large amounts of data in real time. \n", + ">Hologres supports standard SQL syntax, is compatible with PostgreSQL, and supports most PostgreSQL functions. Hologres supports online analytical processing (OLAP) and ad hoc analysis for up to petabytes of data, and provides high-concurrency and low-latency online data services. \n", + "\n", + ">Hologres provides **vector database** functionality by adopting [Proxima](https://www.alibabacloud.com/help/en/hologres/latest/vector-processing).\n", + ">Proxima is a high-performance software library developed by Alibaba DAMO Academy. It allows you to search for the nearest neighbors of vectors. Proxima provides higher stability and performance than similar open source software such as Faiss. Proxima allows you to search for similar text or image embeddings with high throughput and low latency. Hologres is deeply integrated with Proxima to provide a high-performance vector search service.\n", + "\n", + "This notebook shows how to use functionality related to the `Hologres Proxima` vector database.\n", + "Click [here](https://www.alibabacloud.com/zh/product/hologres) to fast deploy a Hologres cloud instance." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from langchain.embeddings.openai import OpenAIEmbeddings\n", + "from langchain.text_splitter import CharacterTextSplitter\n", + "from langchain.vectorstores import Hologres" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Split documents and get embeddings by call OpenAI API" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.document_loaders import TextLoader\n", + "\n", + "loader = TextLoader(\"../../../state_of_the_union.txt\")\n", + "documents = loader.load()\n", + "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n", + "docs = text_splitter.split_documents(documents)\n", + "\n", + "embeddings = OpenAIEmbeddings()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Connect to Hologres by setting related ENVIRONMENTS.\n", + "```\n", + "export PG_HOST={host}\n", + "export PG_PORT={port} # Optional, default is 80\n", + "export PG_DATABASE={db_name} # Optional, default is postgres\n", + "export PG_USER={username}\n", + "export PG_PASSWORD={password}\n", + "```\n", + "\n", + "Then store your embeddings and documents into Hologres" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "connection_string = Hologres.connection_string_from_db_params(\n", + " host=os.environ.get(\"PGHOST\", \"localhost\"),\n", + " port=int(os.environ.get(\"PGPORT\", \"80\")),\n", + " database=os.environ.get(\"PGDATABASE\", \"postgres\"),\n", + " user=os.environ.get(\"PGUSER\", \"postgres\"),\n", + " password=os.environ.get(\"PGPASSWORD\", \"postgres\"),\n", + ")\n", + "\n", + "vector_db = Hologres.from_documents(\n", + " docs,\n", + " embeddings,\n", + " connection_string=connection_string,\n", + " table_name=\"langchain_example_embeddings\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Query and retrieve data" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "query = \"What did the president say about Ketanji Brown Jackson\"\n", + "docs = vector_db.similarity_search(query)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n", + "\n", + "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n", + "\n", + "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n", + "\n", + "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n" + ] + } + ], + "source": [ + "print(docs[0].page_content)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.16" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/langchain/vectorstores/__init__.py b/langchain/vectorstores/__init__.py index 3a932c23..c59b06fd 100644 --- a/langchain/vectorstores/__init__.py +++ b/langchain/vectorstores/__init__.py @@ -10,6 +10,7 @@ from langchain.vectorstores.deeplake import DeepLake from langchain.vectorstores.docarray import DocArrayHnswSearch, DocArrayInMemorySearch from langchain.vectorstores.elastic_vector_search import ElasticVectorSearch from langchain.vectorstores.faiss import FAISS +from langchain.vectorstores.hologres import Hologres from langchain.vectorstores.lancedb import LanceDB from langchain.vectorstores.matching_engine import MatchingEngine from langchain.vectorstores.milvus import Milvus @@ -57,6 +58,7 @@ __all__ = [ "DocArrayHnswSearch", "DocArrayInMemorySearch", "Typesense", + "Hologres", "Clickhouse", "ClickhouseSettings", "Tigris", diff --git a/langchain/vectorstores/hologres.py b/langchain/vectorstores/hologres.py new file mode 100644 index 00000000..b19dbbbb --- /dev/null +++ b/langchain/vectorstores/hologres.py @@ -0,0 +1,506 @@ +"""VectorStore wrapper around a Hologres database.""" +from __future__ import annotations + +import json +import logging +import uuid +from typing import Any, Dict, Iterable, List, Optional, Tuple, Type + +from langchain.docstore.document import Document +from langchain.embeddings.base import Embeddings +from langchain.utils import get_from_dict_or_env +from langchain.vectorstores.base import VectorStore + +ADA_TOKEN_COUNT = 1536 +_LANGCHAIN_DEFAULT_TABLE_NAME = "langchain_pg_embedding" + + +class HologresWrapper: + def __init__(self, connection_string: str, ndims: int, table_name: str) -> None: + import psycopg2 + + self.table_name = table_name + self.conn = psycopg2.connect(connection_string) + self.cursor = self.conn.cursor() + self.conn.autocommit = False + self.ndims = ndims + + def create_vector_extension(self) -> None: + self.cursor.execute("create extension if not exists proxima") + self.conn.commit() + + def create_table(self, drop_if_exist: bool = True) -> None: + if drop_if_exist: + self.cursor.execute(f"drop table if exists {self.table_name}") + self.conn.commit() + + self.cursor.execute( + f"""create table if not exists {self.table_name} ( +id text, +embedding float4[] check(array_ndims(embedding) = 1 and \ +array_length(embedding, 1) = {self.ndims}), +metadata json, +document text);""" + ) + self.cursor.execute( + f"call set_table_property('{self.table_name}'" + + """, 'proxima_vectors', +'{"embedding":{"algorithm":"Graph", +"distance_method":"SquaredEuclidean", +"build_params":{"min_flush_proxima_row_count" : 1, +"min_compaction_proxima_row_count" : 1, +"max_total_size_to_merge_mb" : 2000}}}');""" + ) + self.conn.commit() + + def get_by_id(self, id: str) -> List[Tuple]: + statement = ( + f"select id, embedding, metadata, " + f"document from {self.table_name} where id = %s;" + ) + self.cursor.execute( + statement, + (id), + ) + self.conn.commit() + return self.cursor.fetchall() + + def insert( + self, + embedding: List[float], + metadata: dict, + document: str, + id: Optional[str] = None, + ) -> None: + self.cursor.execute( + f'insert into "{self.table_name}" ' + f"values (%s, array{json.dumps(embedding)}::float4[], %s, %s)", + (id if id is not None else "null", json.dumps(metadata), document), + ) + self.conn.commit() + + def query_nearest_neighbours( + self, embedding: List[float], k: int, filter: Optional[Dict[str, str]] = None + ) -> List[Tuple[str, str, float]]: + params = [] + filter_clause = "" + if filter is not None: + conjuncts = [] + for key, val in filter.items(): + conjuncts.append("metadata->>%s=%s") + params.append(key) + params.append(val) + filter_clause = "where " + " and ".join(conjuncts) + + sql = ( + f"select document, metadata::text, " + f"pm_approx_squared_euclidean_distance(array{json.dumps(embedding)}" + f"::float4[], embedding) as distance from" + f" {self.table_name} {filter_clause} order by distance asc limit {k};" + ) + self.cursor.execute(sql, tuple(params)) + self.conn.commit() + return self.cursor.fetchall() + + +class Hologres(VectorStore): + """ + VectorStore implementation using Hologres. + - `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. + """ + + 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. + """ + self.storage = HologresWrapper( + self.connection_string, self.ndims, self.table_name + ) + self.create_vector_extension() + self.create_table() + + def create_vector_extension(self) -> None: + try: + self.storage.create_vector_extension() + except Exception as e: + self.logger.exception(e) + raise e + + def create_table(self) -> None: + self.storage.create_table(self.pre_delete_table) + + @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.uuid1()) 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 + + 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: + for text, metadata, embedding, id in zip(texts, metadatas, embeddings, ids): + self.storage.insert(embedding, metadata, text, id) + except Exception as e: + self.logger.exception(e) + self.storage.conn.commit() + + 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.uuid1()) 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 + + 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, + ) + + 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] + + 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 + + def similarity_search_with_score_by_vector( + self, + embedding: List[float], + k: int = 4, + filter: Optional[dict] = None, + ) -> List[Tuple[Document, float]]: + results: List[Tuple[str, str, float]] = self.storage.query_nearest_neighbours( + embedding, k, filter + ) + + docs = [ + ( + Document( + page_content=result[0], + metadata=json.loads(result[1]), + ), + result[2], + ) + for result in results + ] + return 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. + Postgres connection string is required + "Either pass it as a parameter + or set the HOLOGRES_CONNECTION_STRING environment variable. + """ + 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, + ) + + @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. + Postgres connection string is required + "Either pass it as a parameter + or set the HOLOGRES_CONNECTION_STRING environment variable. + + Example: + .. code-block:: python + + from langchain import Hologres + from langchain.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, + ) + + @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 intsance 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 + + @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( + "Postgres connection string is required" + "Either pass it as a parameter" + "or set the HOLOGRES_CONNECTION_STRING environment variable." + ) + + return connection_string + + @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. + Postgres connection string is required + "Either pass it as a parameter + or set the HOLOGRES_CONNECTION_STRING environment variable. + """ + + 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, + ) + + @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}" + ) diff --git a/tests/integration_tests/vectorstores/test_hologres.py b/tests/integration_tests/vectorstores/test_hologres.py new file mode 100644 index 00000000..1e115753 --- /dev/null +++ b/tests/integration_tests/vectorstores/test_hologres.py @@ -0,0 +1,142 @@ +"""Test Hologres functionality.""" +import os +from typing import List + +from langchain.docstore.document import Document +from langchain.vectorstores.hologres import Hologres +from tests.integration_tests.vectorstores.fake_embeddings import FakeEmbeddings + +CONNECTION_STRING = Hologres.connection_string_from_db_params( + host=os.environ.get("TEST_HOLOGRES_HOST", "localhost"), + port=int(os.environ.get("TEST_HOLOGRES_PORT", "80")), + database=os.environ.get("TEST_HOLOGRES_DATABASE", "postgres"), + user=os.environ.get("TEST_HOLOGRES_USER", "postgres"), + password=os.environ.get("TEST_HOLOGRES_PASSWORD", "postgres"), +) + + +ADA_TOKEN_COUNT = 1536 + + +class FakeEmbeddingsWithAdaDimension(FakeEmbeddings): + """Fake embeddings functionality for testing.""" + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Return simple embeddings.""" + return [ + [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(i)] for i in range(len(texts)) + ] + + def embed_query(self, text: str) -> List[float]: + """Return simple embeddings.""" + return [float(1.0)] * (ADA_TOKEN_COUNT - 1) + [float(0.0)] + + +def test_hologres() -> None: + """Test end to end construction and search.""" + texts = ["foo", "bar", "baz"] + docsearch = Hologres.from_texts( + texts=texts, + table_name="test_table", + embedding=FakeEmbeddingsWithAdaDimension(), + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search("foo", k=1) + assert output == [Document(page_content="foo")] + + +def test_hologres_embeddings() -> None: + """Test end to end construction with embeddings and search.""" + texts = ["foo", "bar", "baz"] + text_embeddings = FakeEmbeddingsWithAdaDimension().embed_documents(texts) + text_embedding_pairs = list(zip(texts, text_embeddings)) + docsearch = Hologres.from_embeddings( + text_embeddings=text_embedding_pairs, + table_name="test_table", + embedding=FakeEmbeddingsWithAdaDimension(), + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search("foo", k=1) + assert output == [Document(page_content="foo")] + + +def test_hologres_with_metadatas() -> None: + """Test end to end construction and search.""" + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + docsearch = Hologres.from_texts( + texts=texts, + table_name="test_table", + embedding=FakeEmbeddingsWithAdaDimension(), + metadatas=metadatas, + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search("foo", k=1) + assert output == [Document(page_content="foo", metadata={"page": "0"})] + + +def test_hologres_with_metadatas_with_scores() -> None: + """Test end to end construction and search.""" + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + docsearch = Hologres.from_texts( + texts=texts, + table_name="test_table", + embedding=FakeEmbeddingsWithAdaDimension(), + metadatas=metadatas, + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search_with_score("foo", k=1) + assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)] + + +def test_hologres_with_filter_match() -> None: + """Test end to end construction and search.""" + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + docsearch = Hologres.from_texts( + texts=texts, + table_name="test_table_filter", + embedding=FakeEmbeddingsWithAdaDimension(), + metadatas=metadatas, + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "0"}) + assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)] + + +def test_hologres_with_filter_distant_match() -> None: + """Test end to end construction and search.""" + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + docsearch = Hologres.from_texts( + texts=texts, + table_name="test_table_filter", + embedding=FakeEmbeddingsWithAdaDimension(), + metadatas=metadatas, + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "2"}) + assert output == [(Document(page_content="baz", metadata={"page": "2"}), 4.0)] + + +def test_hologres_with_filter_no_match() -> None: + """Test end to end construction and search.""" + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + docsearch = Hologres.from_texts( + texts=texts, + table_name="test_table_filter", + embedding=FakeEmbeddingsWithAdaDimension(), + metadatas=metadatas, + connection_string=CONNECTION_STRING, + pre_delete_table=True, + ) + output = docsearch.similarity_search_with_score("foo", k=1, filter={"page": "5"}) + assert output == []