diff --git a/docs/extras/modules/data_connection/vectorstores/integrations/pgvector.ipynb b/docs/extras/modules/data_connection/vectorstores/integrations/pgvector.ipynb index 381de0ee9f..c36bb47299 100644 --- a/docs/extras/modules/data_connection/vectorstores/integrations/pgvector.ipynb +++ b/docs/extras/modules/data_connection/vectorstores/integrations/pgvector.ipynb @@ -123,7 +123,7 @@ }, { "cell_type": "code", - "execution_count": 62, + "execution_count": 1, "metadata": { "tags": [] }, @@ -138,7 +138,7 @@ }, { "cell_type": "code", - "execution_count": 63, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -152,49 +152,25 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 3, "metadata": {}, "outputs": [], "source": [ - "## PGVector needs the connection string to the database.\n", - "## We will load it from the environment variables.\n", - "import os\n", - "\n", - "CONNECTION_STRING = PGVector.connection_string_from_db_params(\n", - " driver=os.environ.get(\"PGVECTOR_DRIVER\", \"psycopg2\"),\n", - " host=os.environ.get(\"PGVECTOR_HOST\", \"localhost\"),\n", - " port=int(os.environ.get(\"PGVECTOR_PORT\", \"5432\")),\n", - " database=os.environ.get(\"PGVECTOR_DATABASE\", \"postgres\"),\n", - " user=os.environ.get(\"PGVECTOR_USER\", \"postgres\"),\n", - " password=os.environ.get(\"PGVECTOR_PASSWORD\", \"postgres\"),\n", - ")\n", + "# PGVector needs the connection string to the database.\n", + "CONNECTION_STRING = \"postgresql+psycopg2://harrisonchase@localhost:5432/test3\"\n", "\n", - "\n", - "## Example\n", - "# postgresql+psycopg2://username:password@localhost:5432/database_name" - ] - }, - { - "cell_type": "code", - "execution_count": 64, - "metadata": {}, - "outputs": [], - "source": [ - "# ## PGVector needs the connection string to the database.\n", - "# ## We will load it from the environment variables.\n", + "# # Alternatively, you can create it from enviornment variables.\n", "# import os\n", + "\n", "# CONNECTION_STRING = PGVector.connection_string_from_db_params(\n", "# driver=os.environ.get(\"PGVECTOR_DRIVER\", \"psycopg2\"),\n", "# host=os.environ.get(\"PGVECTOR_HOST\", \"localhost\"),\n", "# port=int(os.environ.get(\"PGVECTOR_PORT\", \"5432\")),\n", - "# database=os.environ.get(\"PGVECTOR_DATABASE\", \"rd-embeddings\"),\n", - "# user=os.environ.get(\"PGVECTOR_USER\", \"admin\"),\n", - "# password=os.environ.get(\"PGVECTOR_PASSWORD\", \"password\"),\n", + "# database=os.environ.get(\"PGVECTOR_DATABASE\", \"postgres\"),\n", + "# user=os.environ.get(\"PGVECTOR_USER\", \"postgres\"),\n", + "# password=os.environ.get(\"PGVECTOR_PASSWORD\", \"postgres\"),\n", "# )\n", - "\n", - "\n", - "# ## Example\n", - "# # postgresql+psycopg2://username:password@localhost:5432/database_name" + "\n" ] }, { @@ -206,27 +182,36 @@ }, { "cell_type": "code", - "execution_count": 69, + "execution_count": 16, "metadata": {}, "outputs": [], "source": [ - "# The PGVector Module will try to create a table with the name of the collection. So, make sure that the collection name is unique and the user has the\n", - "# permission to create a table.\n", + "# The PGVector Module will try to create a table with the name of the collection. \n", + "# So, make sure that the collection name is unique and the user has the permission to create a table.\n", + "\n", + "COLLECTION_NAME = \"state_of_the_union_test\"\n", "\n", "db = PGVector.from_documents(\n", " embedding=embeddings,\n", " documents=docs,\n", - " collection_name=\"state_of_the_union\",\n", + " collection_name=COLLECTION_NAME,\n", " connection_string=CONNECTION_STRING,\n", - ")\n", - "\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [], + "source": [ "query = \"What did the president say about Ketanji Brown Jackson\"\n", - "docs_with_score: List[Tuple[Document, float]] = db.similarity_search_with_score(query)" + "docs_with_score = db.similarity_search_with_score(query)" ] }, { "cell_type": "code", - "execution_count": 70, + "execution_count": 18, "metadata": {}, "outputs": [ { @@ -234,7 +219,7 @@ "output_type": "stream", "text": [ "--------------------------------------------------------------------------------\n", - "Score: 0.6076804864602984\n", + "Score: 0.18460171628856903\n", "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", @@ -244,7 +229,7 @@ "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", "--------------------------------------------------------------------------------\n", "--------------------------------------------------------------------------------\n", - "Score: 0.6076804864602984\n", + "Score: 0.18460171628856903\n", "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", @@ -254,21 +239,17 @@ "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", "--------------------------------------------------------------------------------\n", "--------------------------------------------------------------------------------\n", - "Score: 0.659062774389974\n", - "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n", - "\n", - "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n", - "\n", - "We can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \n", + "Score: 0.18470284560586236\n", + "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", - "We’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \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", - "We’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n", + "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n", "\n", - "We’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.\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", "--------------------------------------------------------------------------------\n", "--------------------------------------------------------------------------------\n", - "Score: 0.659062774389974\n", + "Score: 0.21730864082247825\n", "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n", "\n", "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n", @@ -296,183 +277,189 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Working with vectorstore" + "## Working with vectorstore\n", + "\n", + "Above, we created a vectorstore from scratch. However, often times we want to work with an existing vectorstore.\n", + "In order to do that, we can initialize it directly." ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": 8, "metadata": {}, + "outputs": [], "source": [ - "### Uploading a vectorstore" + "store = PGVector(\n", + " collection_name=COLLECTION_NAME,\n", + " connection_string=CONNECTION_STRING,\n", + " embedding_function=embeddings,\n", + ")\n" ] }, { - "cell_type": "code", - "execution_count": 55, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "data = docs\n", - "api_key = os.environ[\"OPENAI_API_KEY\"]\n", - "db = PGVector.from_documents(\n", - " documents=docs,\n", - " embedding=embeddings,\n", - " collection_name=collection_name,\n", - " connection_string=connection_string,\n", - " distance_strategy=DistanceStrategy.COSINE,\n", - " openai_api_key=api_key,\n", - " pre_delete_collection=False,\n", - ")" + "### Add documents\n", + "We can add documents to the existing vectorstore." ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": 19, "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "['048c2e14-1cf3-11ee-8777-e65801318980']" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ - "### Retrieving a vectorstore" + "store.add_documents([Document(page_content=\"foo\")])" ] }, { "cell_type": "code", - "execution_count": 56, + "execution_count": 20, "metadata": {}, "outputs": [], "source": [ - "connection_string = CONNECTION_STRING\n", - "embedding = embeddings\n", - "collection_name = \"state_of_the_union\"\n", - "from langchain.vectorstores.pgvector import DistanceStrategy\n", - "\n", - "store = PGVector(\n", - " connection_string=connection_string,\n", - " embedding_function=embedding,\n", - " collection_name=collection_name,\n", - " distance_strategy=DistanceStrategy.COSINE,\n", - ")\n", - "\n", - "retriever = store.as_retriever()" + "docs_with_score = db.similarity_search_with_score(\"foo\")" ] }, { "cell_type": "code", - "execution_count": 57, + "execution_count": 21, "metadata": {}, "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "vectorstore= search_type='similarity' search_kwargs={}\n" - ] + "data": { + "text/plain": [ + "(Document(page_content='foo', metadata={}), 3.3203430005457335e-09)" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ - "print(retriever)" + "docs_with_score[0]" ] }, { "cell_type": "code", - "execution_count": 83, + "execution_count": 22, "metadata": {}, "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "[(Document(page_content='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\\nTonight, 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\\nOne of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \\n\\nAnd 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.', metadata={'source': '../../../state_of_the_union.txt'}), 0.6075870262188066), (Document(page_content='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\\nTonight, 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\\nOne of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \\n\\nAnd 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.', metadata={'source': '../../../state_of_the_union.txt'}), 0.6075870262188066), (Document(page_content='A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \\n\\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \\n\\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \\n\\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \\n\\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \\n\\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.', metadata={'source': '../../../state_of_the_union.txt'}), 0.6589478388546668), (Document(page_content='A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \\n\\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \\n\\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \\n\\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \\n\\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \\n\\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.', metadata={'source': '../../../state_of_the_union.txt'}), 0.6589478388546668)]\n" - ] + "data": { + "text/plain": [ + "(Document(page_content='A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \\n\\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \\n\\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \\n\\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \\n\\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \\n\\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.', metadata={'source': '../../../state_of_the_union.txt'}),\n", + " 0.2404395365581814)" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ - "# When we have an existing PG VEctor\n", - "DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.EUCLIDEAN\n", - "db1 = PGVector.from_existing_index(\n", + "docs_with_score[1]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Overriding a vectorstore\n", + "\n", + "If you have an existing collection, you override it by doing `from_documents` and setting `pre_delete_collection` = True" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [], + "source": [ + "db = PGVector.from_documents(\n", + " documents=docs,\n", " embedding=embeddings,\n", - " collection_name=\"state_of_the_union\",\n", - " distance_strategy=DEFAULT_DISTANCE_STRATEGY,\n", - " pre_delete_collection=False,\n", + " collection_name=COLLECTION_NAME,\n", " connection_string=CONNECTION_STRING,\n", - ")\n", - "\n", - "query = \"What did the president say about Ketanji Brown Jackson\"\n", - "docs_with_score: List[Tuple[Document, float]] = db1.similarity_search_with_score(query)\n", - "print(docs_with_score)" + " pre_delete_collection=True,\n", + ")" ] }, { "cell_type": "code", - "execution_count": 81, + "execution_count": 24, + "metadata": {}, + "outputs": [], + "source": [ + "docs_with_score = db.similarity_search_with_score(\"foo\")" + ] + }, + { + "cell_type": "code", + "execution_count": 25, "metadata": {}, "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "--------------------------------------------------------------------------------\n", - "Score: 0.6075870262188066\n", - "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", - "--------------------------------------------------------------------------------\n", - "--------------------------------------------------------------------------------\n", - "Score: 0.6075870262188066\n", - "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", - "--------------------------------------------------------------------------------\n", - "--------------------------------------------------------------------------------\n", - "Score: 0.6589478388546668\n", - "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n", - "\n", - "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n", - "\n", - "We can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \n", - "\n", - "We’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \n", - "\n", - "We’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n", - "\n", - "We’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.\n", - "--------------------------------------------------------------------------------\n", - "--------------------------------------------------------------------------------\n", - "Score: 0.6589478388546668\n", - "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n", - "\n", - "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n", - "\n", - "We can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \n", - "\n", - "We’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \n", - "\n", - "We’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \n", - "\n", - "We’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.\n", - "--------------------------------------------------------------------------------\n" - ] + "data": { + "text/plain": [ + "(Document(page_content='A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \\n\\nAnd if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \\n\\nWe can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. \\n\\nWe’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. \\n\\nWe’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. \\n\\nWe’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders.', metadata={'source': '../../../state_of_the_union.txt'}),\n", + " 0.2404115088144465)" + ] + }, + "execution_count": 25, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ - "for doc, score in docs_with_score:\n", - " print(\"-\" * 80)\n", - " print(\"Score: \", score)\n", - " print(doc.page_content)\n", - " print(\"-\" * 80)" + "docs_with_score[0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Using a VectorStore as a Retriever" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 26, "metadata": {}, "outputs": [], - "source": [] + "source": [ + "retriever = store.as_retriever()" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "tags=None metadata=None vectorstore= search_type='similarity' search_kwargs={}\n" + ] + } + ], + "source": [ + "print(retriever)" + ] } ], "metadata": { @@ -491,7 +478,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.6" + "version": "3.9.1" } }, "nbformat": 4, diff --git a/langchain/vectorstores/__init__.py b/langchain/vectorstores/__init__.py index bc98f88cdf..dc524f5117 100644 --- a/langchain/vectorstores/__init__.py +++ b/langchain/vectorstores/__init__.py @@ -29,6 +29,7 @@ from langchain.vectorstores.mongodb_atlas import MongoDBAtlasVectorSearch from langchain.vectorstores.myscale import MyScale, MyScaleSettings from langchain.vectorstores.opensearch_vector_search import OpenSearchVectorSearch from langchain.vectorstores.pgembedding import PGEmbedding +from langchain.vectorstores.pgvector import PGVector from langchain.vectorstores.pinecone import Pinecone from langchain.vectorstores.qdrant import Qdrant from langchain.vectorstores.redis import Redis @@ -95,4 +96,5 @@ __all__ = [ "VectorStore", "Weaviate", "Zilliz", + "PGVector", ] diff --git a/langchain/vectorstores/_pgvector_data_models.py b/langchain/vectorstores/_pgvector_data_models.py new file mode 100644 index 0000000000..1be27e5533 --- /dev/null +++ b/langchain/vectorstores/_pgvector_data_models.py @@ -0,0 +1,26 @@ +import sqlalchemy +from pgvector.sqlalchemy import Vector +from sqlalchemy.dialects.postgresql import JSON, UUID +from sqlalchemy.orm import relationship + +from langchain.vectorstores.pgvector import BaseModel, CollectionStore + + +class EmbeddingStore(BaseModel): + __tablename__ = "langchain_pg_embedding" + + collection_id = sqlalchemy.Column( + UUID(as_uuid=True), + sqlalchemy.ForeignKey( + f"{CollectionStore.__tablename__}.uuid", + ondelete="CASCADE", + ), + ) + collection = relationship(CollectionStore, back_populates="embeddings") + + embedding: Vector = sqlalchemy.Column(Vector(None)) + document = sqlalchemy.Column(sqlalchemy.String, nullable=True) + cmetadata = sqlalchemy.Column(JSON, nullable=True) + + # custom_id : any user defined id + custom_id = sqlalchemy.Column(sqlalchemy.String, nullable=True) diff --git a/langchain/vectorstores/pgvector.py b/langchain/vectorstores/pgvector.py index 2825fc16c8..5654db5f71 100644 --- a/langchain/vectorstores/pgvector.py +++ b/langchain/vectorstores/pgvector.py @@ -7,7 +7,6 @@ import uuid from typing import Any, Dict, Iterable, List, Optional, Tuple, Type import sqlalchemy -from pgvector.sqlalchemy import Vector from sqlalchemy.dialects.postgresql import JSON, UUID from sqlalchemy.orm import Session, declarative_base, relationship @@ -16,6 +15,17 @@ from langchain.embeddings.base import Embeddings from langchain.utils import get_from_dict_or_env from langchain.vectorstores.base import VectorStore + +class DistanceStrategy(str, enum.Enum): + """Enumerator of the Distance strategies.""" + + EUCLIDEAN = "l2" + COSINE = "cosine" + MAX_INNER_PRODUCT = "inner" + + +DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.COSINE + Base = declarative_base() # type: Any @@ -66,58 +76,40 @@ class CollectionStore(BaseModel): return collection, created -class EmbeddingStore(BaseModel): - __tablename__ = "langchain_pg_embedding" - - collection_id = sqlalchemy.Column( - UUID(as_uuid=True), - sqlalchemy.ForeignKey( - f"{CollectionStore.__tablename__}.uuid", - ondelete="CASCADE", - ), - ) - collection = relationship(CollectionStore, back_populates="embeddings") - - embedding: Vector = sqlalchemy.Column(Vector(None)) - document = sqlalchemy.Column(sqlalchemy.String, nullable=True) - cmetadata = sqlalchemy.Column(JSON, nullable=True) - - # custom_id : any user defined id - custom_id = sqlalchemy.Column(sqlalchemy.String, nullable=True) - - -class QueryResult: - EmbeddingStore: EmbeddingStore - distance: float - - -class DistanceStrategy(str, enum.Enum): - """Enumerator of the Distance strategies.""" - - EUCLIDEAN = EmbeddingStore.embedding.l2_distance - COSINE = EmbeddingStore.embedding.cosine_distance - MAX_INNER_PRODUCT = EmbeddingStore.embedding.max_inner_product - - -DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.EUCLIDEAN +class PGVector(VectorStore): + """VectorStore implementation using Postgres and pgvector. + To use, you should have the ``pgvector`` python package installed. -class PGVector(VectorStore): - """ - VectorStore implementation using Postgres and pgvector. - - `connection_string` is a postgres connection string. - - `embedding_function` any embedding function implementing - `langchain.embeddings.base.Embeddings` interface. - - `collection_name` is the name of the collection to use. (default: langchain) - - NOTE: This is not the name of the table, but the name of the collection. + Args: + connection_string: Postgres connection string. + embedding_function: Any embedding function implementing + `langchain.embeddings.base.Embeddings` interface. + collection_name: The name of the collection to use. (default: langchain) + NOTE: This is not the name of the table, but the name of the collection. The tables will be created when initializing the store (if not exists) So, make sure the user has the right permissions to create tables. - - `distance_strategy` is the distance strategy to use. (default: EUCLIDEAN) - - `EUCLIDEAN` is the euclidean distance. - - `COSINE` is the cosine distance. - - `pre_delete_collection` if True, will delete the collection if it exists. - (default: False) - - Useful for testing. + distance_strategy: The distance strategy to use. (default: COSINE) + pre_delete_collection: If True, will delete the collection if it exists. + (default: False). Useful for testing. + + Example: + .. code-block:: python + + from langchain.vectorstores import PGVector + from langchain.embeddings.openai import OpenAIEmbeddings + + CONNECTION_STRING = "postgresql+psycopg2://hwc@localhost:5432/test3" + COLLECTION_NAME = "state_of_the_union_test" + embeddings = OpenAIEmbeddings() + vectorestore = PGVector.from_documents( + embedding=embeddings, + documents=docs, + collection_name=COLLECTION_NAME, + connection_string=CONNECTION_STRING, + ) + + """ def __init__( @@ -134,7 +126,7 @@ class PGVector(VectorStore): self.embedding_function = embedding_function self.collection_name = collection_name self.collection_metadata = collection_metadata - self.distance_strategy = distance_strategy + self._distance_strategy = distance_strategy self.pre_delete_collection = pre_delete_collection self.logger = logger or logging.getLogger(__name__) self.__post_init__() @@ -147,6 +139,9 @@ class PGVector(VectorStore): """ self._conn = self.connect() # self.create_vector_extension() + from langchain.vectorstores._pgvector_data_models import EmbeddingStore + + self.EmbeddingStore = EmbeddingStore self.create_tables_if_not_exists() self.create_collection() @@ -202,7 +197,7 @@ class PGVector(VectorStore): metadatas: Optional[List[dict]] = None, ids: Optional[List[str]] = None, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, - distance_strategy: DistanceStrategy = DistanceStrategy.COSINE, + distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, pre_delete_collection: bool = False, **kwargs: Any, ) -> PGVector: @@ -249,7 +244,7 @@ class PGVector(VectorStore): if not collection: raise ValueError("Collection not found") for text, metadata, embedding, id in zip(texts, metadatas, embeddings, ids): - embedding_store = EmbeddingStore( + embedding_store = self.EmbeddingStore( embedding=embedding, document=text, cmetadata=metadata, @@ -329,6 +324,20 @@ class PGVector(VectorStore): ) return docs + @property + def distance_strategy(self) -> Any: + if self._distance_strategy == "l2": + return self.EmbeddingStore.embedding.l2_distance + elif self._distance_strategy == "cosine": + return self.EmbeddingStore.embedding.cosine_distance + elif self._distance_strategy == "inner": + return self.EmbeddingStore.embedding.max_inner_product + else: + raise ValueError( + f"Got unexpected value for distance: {self._distance_strategy}. " + f"Should be one of `l2`, `cosine`, `inner`." + ) + def similarity_search_with_score_by_vector( self, embedding: List[float], @@ -340,7 +349,7 @@ class PGVector(VectorStore): if not collection: raise ValueError("Collection not found") - filter_by = EmbeddingStore.collection_id == collection.uuid + filter_by = self.EmbeddingStore.collection_id == collection.uuid if filter is not None: filter_clauses = [] @@ -350,28 +359,30 @@ class PGVector(VectorStore): value_case_insensitive = { k.lower(): v for k, v in value.items() } - filter_by_metadata = EmbeddingStore.cmetadata[key].astext.in_( - value_case_insensitive[IN] - ) + filter_by_metadata = self.EmbeddingStore.cmetadata[ + key + ].astext.in_(value_case_insensitive[IN]) filter_clauses.append(filter_by_metadata) else: - filter_by_metadata = EmbeddingStore.cmetadata[ + filter_by_metadata = self.EmbeddingStore.cmetadata[ key ].astext == str(value) filter_clauses.append(filter_by_metadata) filter_by = sqlalchemy.and_(filter_by, *filter_clauses) - results: List[QueryResult] = ( + _type = self.EmbeddingStore + + results: List[Any] = ( session.query( - EmbeddingStore, + self.EmbeddingStore, self.distance_strategy(embedding).label("distance"), # type: ignore ) .filter(filter_by) .order_by(sqlalchemy.asc("distance")) .join( CollectionStore, - EmbeddingStore.collection_id == CollectionStore.uuid, + self.EmbeddingStore.collection_id == CollectionStore.uuid, ) .limit(k) .all() @@ -418,7 +429,7 @@ class PGVector(VectorStore): embedding: Embeddings, metadatas: Optional[List[dict]] = None, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, - distance_strategy: DistanceStrategy = DistanceStrategy.COSINE, + distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, ids: Optional[List[str]] = None, pre_delete_collection: bool = False, **kwargs: Any, @@ -450,7 +461,7 @@ class PGVector(VectorStore): embedding: Embeddings, metadatas: Optional[List[dict]] = None, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, - distance_strategy: DistanceStrategy = DistanceStrategy.COSINE, + distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, ids: Optional[List[str]] = None, pre_delete_collection: bool = False, **kwargs: Any, @@ -493,7 +504,7 @@ class PGVector(VectorStore): cls: Type[PGVector], embedding: Embeddings, collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME, - distance_strategy: DistanceStrategy = DistanceStrategy.COSINE, + distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY, pre_delete_collection: bool = False, **kwargs: Any, ) -> PGVector: