from __future__ import annotations
import json
import logging
import uuid
import warnings
from itertools import repeat
from typing import (
Any,
Iterable,
List,
Optional,
Tuple,
Type,
)
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore
from langchain_community.docstore.document import Document
logger = logging.getLogger(__name__)
[docs]class Yellowbrick(VectorStore):
"""Wrapper around Yellowbrick as a vector database.
Example:
.. code-block:: python
from langchain_community.vectorstores import Yellowbrick
from langchain_community.embeddings.openai import OpenAIEmbeddings
...
"""
[docs] def __init__(
self,
embedding: Embeddings,
connection_string: str,
table: str,
) -> None:
"""Initialize with yellowbrick client.
Args:
embedding: Embedding operator
connection_string: Format 'postgres://username:password@host:port/database'
table: Table used to store / retrieve embeddings from
"""
import psycopg2
if not isinstance(embedding, Embeddings):
warnings.warn("embeddings input must be Embeddings object.")
self.connection_string = connection_string
self._table = table
self._embedding = embedding
self._connection = psycopg2.connect(connection_string)
self.__post_init__()
def __post_init__(
self,
) -> None:
"""Initialize the store."""
self.check_database_utf8()
self.create_table_if_not_exists()
def __del__(self) -> None:
if self._connection:
self._connection.close()
[docs] def create_table_if_not_exists(self) -> None:
"""
Helper function: create table if not exists
"""
from psycopg2 import sql
cursor = self._connection.cursor()
cursor.execute(
sql.SQL(
"CREATE TABLE IF NOT EXISTS {} ( \
id UUID, \
embedding_id INTEGER, \
text VARCHAR(60000), \
metadata VARCHAR(1024), \
embedding FLOAT)"
).format(sql.Identifier(self._table))
)
self._connection.commit()
cursor.close()
[docs] def drop(self, table: str) -> None:
"""
Helper function: Drop data
"""
from psycopg2 import sql
cursor = self._connection.cursor()
cursor.execute(sql.SQL("DROP TABLE IF EXISTS {}").format(sql.Identifier(table)))
self._connection.commit()
cursor.close()
[docs] def check_database_utf8(self) -> bool:
"""
Helper function: Test the database is UTF-8 encoded
"""
cursor = self._connection.cursor()
query = "SELECT pg_encoding_to_char(encoding) \
FROM pg_database \
WHERE datname = current_database();"
cursor.execute(query)
encoding = cursor.fetchone()[0]
cursor.close()
if encoding.lower() == "utf8" or encoding.lower() == "utf-8":
return True
else:
raise Exception(
f"Database \
'{self.connection_string.split('/')[-1]}' encoding is not UTF-8"
)
[docs] def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> List[str]:
"""Add more texts to the vectorstore index.
Args:
texts: Iterable of strings to add to the vectorstore.
metadatas: Optional list of metadatas associated with the texts.
kwargs: vectorstore specific parameters
"""
from psycopg2 import sql
texts = list(texts)
cursor = self._connection.cursor()
embeddings = self._embedding.embed_documents(list(texts))
results = []
if not metadatas:
metadatas = [{} for _ in texts]
for id in range(len(embeddings)):
doc_uuid = uuid.uuid4()
results.append(str(doc_uuid))
data_input = [
(str(id), embedding_id, text, json.dumps(metadata), embedding)
for id, embedding_id, text, metadata, embedding in zip(
repeat(doc_uuid),
range(len(embeddings[id])),
repeat(texts[id]),
repeat(metadatas[id]),
embeddings[id],
)
]
flattened_input = [val for sublist in data_input for val in sublist]
insert_query = sql.SQL(
"INSERT INTO {t} \
(id, embedding_id, text, metadata, embedding) VALUES {v}"
).format(
t=sql.Identifier(self._table),
v=(
sql.SQL(",").join(
[
sql.SQL("(%s,%s,%s,%s,%s)")
for _ in range(len(embeddings[id]))
]
)
),
)
cursor.execute(insert_query, flattened_input)
self._connection.commit()
return results
[docs] @classmethod
def from_texts(
cls: Type[Yellowbrick],
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
connection_string: str = "",
table: str = "langchain",
**kwargs: Any,
) -> Yellowbrick:
"""Add texts to the vectorstore index.
Args:
texts: Iterable of strings to add to the vectorstore.
metadatas: Optional list of metadatas associated with the texts.
connection_string: URI to Yellowbrick instance
embedding: Embedding function
table: table to store embeddings
kwargs: vectorstore specific parameters
"""
if connection_string is None:
raise ValueError("connection_string must be provided")
vss = cls(
embedding=embedding,
connection_string=connection_string,
table=table,
)
vss.add_texts(texts=texts, metadatas=metadatas)
return vss
[docs] def similarity_search_with_score_by_vector(
self, embedding: List[float], k: int = 4, **kwargs: Any
) -> List[Tuple[Document, float]]:
"""Perform a similarity search with Yellowbrick with vector
Args:
embedding (List[float]): query embedding
k (int, optional): Top K neighbors to retrieve. Defaults to 4.
NOTE: Please do not let end-user fill this and always be aware
of SQL injection.
Returns:
List[Document, float]: List of Documents and scores
"""
from psycopg2 import sql
cursor = self._connection.cursor()
tmp_table = "tmp_" + self._table
cursor.execute(
sql.SQL(
"CREATE TEMPORARY TABLE {} ( \
embedding_id INTEGER, embedding FLOAT)"
).format(sql.Identifier(tmp_table))
)
self._connection.commit()
data_input = [
(embedding_id, embedding)
for embedding_id, embedding in zip(range(len(embedding)), embedding)
]
flattened_input = [val for sublist in data_input for val in sublist]
insert_query = sql.SQL(
"INSERT INTO {t} \
(embedding_id, embedding) VALUES {v}"
).format(
t=sql.Identifier(tmp_table),
v=sql.SQL(",").join([sql.SQL("(%s,%s)") for _ in range(len(embedding))]),
)
cursor.execute(insert_query, flattened_input)
self._connection.commit()
sql_query = sql.SQL(
"SELECT text, \
metadata, \
sum(v1.embedding * v2.embedding) / \
( sqrt(sum(v1.embedding * v1.embedding)) * \
sqrt(sum(v2.embedding * v2.embedding))) AS score \
FROM {v1} v1 INNER JOIN {v2} v2 \
ON v1.embedding_id = v2.embedding_id \
GROUP BY v2.id, v2.text, v2.metadata \
ORDER BY score DESC \
LIMIT %s"
).format(v1=sql.Identifier(tmp_table), v2=sql.Identifier(self._table))
cursor.execute(sql_query, (k,))
results = cursor.fetchall()
self.drop(tmp_table)
documents: List[Tuple[Document, float]] = []
for result in results:
metadata = json.loads(result[1]) or {}
doc = Document(page_content=result[0], metadata=metadata)
documents.append((doc, result[2]))
cursor.close()
return documents
[docs] def similarity_search(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Document]:
"""Perform a similarity search with Yellowbrick
Args:
query (str): query string
k (int, optional): Top K neighbors to retrieve. Defaults to 4.
NOTE: Please do not let end-user fill this and always be aware
of SQL injection.
Returns:
List[Document]: List of Documents
"""
embedding = self._embedding.embed_query(query)
documents = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k
)
return [doc for doc, _ in documents]
[docs] def similarity_search_with_score(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Tuple[Document, float]]:
"""Perform a similarity search with Yellowbrick
Args:
query (str): query string
k (int, optional): Top K neighbors to retrieve. Defaults to 4.
NOTE: Please do not let end-user fill this and always be aware
of SQL injection.
Returns:
List[Document]: List of (Document, similarity)
"""
embedding = self._embedding.embed_query(query)
documents = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k
)
return documents
[docs] def similarity_search_by_vector(
self, embedding: List[float], k: int = 4, **kwargs: Any
) -> List[Document]:
"""Perform a similarity search with Yellowbrick by vectors
Args:
embedding (List[float]): query embedding
k (int, optional): Top K neighbors to retrieve. Defaults to 4.
NOTE: Please do not let end-user fill this and always be aware
of SQL injection.
Returns:
List[Document]: List of documents
"""
documents = self.similarity_search_with_score_by_vector(
embedding=embedding, k=k
)
return [doc for doc, _ in documents]