Source code for langchain_community.vectorstores.yellowbrick

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_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]