@ -1,32 +1,17 @@
from __future__ import annotations
import uuid
from typing import Any , , List , Literal , Optional , Tuple , Type
from typing import Any , Dict, Iterable, List , Literal , Optional , Tuple , Union
import numpy as np
import sqlalchemy
from langchain_core . documents import Document
from langchain_core . embeddings import Embeddings
from langchain_core . vectorstores import VectorStore
from sqlalchemy import insert , select
from sqlalchemy . dialects import postgresql
from sqlalchemy . orm import DeclarativeBase , Mapped , mapped_column
from sqlalchemy . orm . session import Session
class _ORMBase ( DeclarativeBase ) :
__tablename__ : str
id : Mapped [ uuid . UUID ]
text : Mapped [ str ]
meta : Mapped [ dict ]
embedding : Mapped [ np . ndarray ]
class PGVecto_rs ( VectorStore ) :
""" VectorStore backed by pgvecto_rs. """
_engine : sqlalchemy . engine . Engine
_table : Type [ _ORMBase ]
_store = None
_embedding : Embeddings
def __init__ (
@ -45,28 +30,22 @@ class PGVecto_rs(VectorStore):
db_url : Database URL .
collection_name : Name of the collection .
new_table : Whether to create a new table or connect to an existing one .
Defaults to False .
If true , the table will be dropped if exists , then recreated .
Defaults to False .
"""
try :
from pgvecto_rs . s qlalchemy import Vector
from pgvecto_rs . s dk import PGVectoRs
except ImportError as e :
raise ImportError (
" Unable to import pgvector_rs , please install with "
" `pip install pgvector_rs`. "
" Unable to import pgvector_rs .sdk , please install with "
' `pip install " pgvector_rs[sdk] " `. '
) from e
class _Table ( _ORMBase ) :
__tablename__ = f " collection_ { collection_name } "
id : Mapped [ uuid . UUID ] = mapped_column (
postgresql . UUID ( as_uuid = True ) , primary_key = True , default = uuid . uuid4
)
text : Mapped [ str ] = mapped_column ( sqlalchemy . String )
meta : Mapped [ dict ] = mapped_column ( postgresql . JSONB )
embedding : Mapped [ np . ndarray ] = mapped_column ( Vector ( dimension ) )
self . _engine = sqlalchemy . create_engine ( db_url )
self . _table = _Table
self . _table . __table__ . create ( self . _engine , checkfirst = not new_table ) # type: ignore
self . _store = PGVectoRs (
db_url = db_url ,
collection_name = collection_name ,
dimension = dimension ,
recreate = new_table ,
)
self . _embedding = embedding
# ================ Create interface =================
@ -90,7 +69,6 @@ class PGVecto_rs(VectorStore):
dimension = dimension ,
db_url = db_url ,
collection_name = collection_name ,
new_table = True ,
)
_self . add_texts ( texts , metadatas , * * kwargs )
return _self
@ -148,19 +126,15 @@ class PGVecto_rs(VectorStore):
List of ids of the added texts .
"""
from pgvecto_rs . sdk import Record
embeddings = self . _embedding . embed_documents ( list ( texts ) )
with Session ( self . _engine ) as _session :
results : List [ str ] = [ ]
for text , embedding , metadata in zip (
texts , embeddings , metadatas or [ dict ( ) ] * len ( list ( texts ) )
) :
t = insert ( self . _table ) . values (
text = text , meta = metadata , embedding = embedding
)
id = _session . execute ( t ) . inserted_primary_key [ 0 ] # type: ignore
results . append ( str ( id ) )
_session . commit ( )
return results
records = [
Record . from_text ( text , embedding , meta )
for text , embedding , meta in zip ( texts , embeddings , metadatas or [ ] )
]
self . _store . insert ( records )
return [ str ( record . id ) for record in records ]
def add_documents ( self , documents : List [ Document ] , * * kwargs : Any ) - > List [ str ] :
""" Run more documents through the embeddings and add to the vectorstore.
@ -185,31 +159,41 @@ class PGVecto_rs(VectorStore):
distance_func : Literal [
" sqrt_euclid " , " neg_dot_prod " , " ned_cos "
] = " sqrt_euclid " ,
filter : Union [ None , Dict [ str , Any ] , Any ] = None ,
* * kwargs : Any ,
) - > List [ Tuple [ Document , float ] ] :
""" Return docs most similar to query vector, with its score. """
with Session ( self . _engine ) as _session :
real_distance_func = (
self . _table . embedding . squared_euclidean_distance
if distance_func == " sqrt_euclid "
else self . _table . embedding . negative_dot_product_distance
if distance_func == " neg_dot_prod "
else self . _table . embedding . negative_cosine_distance
if distance_func == " ned_cos "
else None
)
if real_distance_func is None :
raise ValueError ( " Invalid distance function " )
t = (
select ( self . _table , real_distance_func ( query_vector ) . label ( " score " ) )
. order_by ( " score " )
. limit ( k ) # type: ignore
from pgvecto_rs . sdk . filters import meta_contains
distance_func_map = {
" sqrt_euclid " : " <-> " ,
" neg_dot_prod " : " <#> " ,
" ned_cos " : " <=> " ,
}
if filter is None :
real_filter = None
elif isinstance ( filter , dict ) :
real_filter = meta_contains ( filter )
else :
real_filter = filter
results = self . _store . search (
query_vector ,
distance_func_map [ distance_func ] ,
k ,
filter = real_filter ,
)
return [
(
Document (
page_content = res [ 0 ] . text ,
metadata = res [ 0 ] . meta ,
) ,
res [ 1 ] ,
)
return [
( Document ( page_content = row [ 0 ] . text , metadata = row [ 0 ] . meta ) , row [ 1 ] )
for row in _session . execute ( t )
]
for res in results
]
def similarity_search_by_vector (
self ,
@ -218,11 +202,12 @@ class PGVecto_rs(VectorStore):
distance_func : Literal [
" sqrt_euclid " , " neg_dot_prod " , " ned_cos "
] = " sqrt_euclid " ,
filter : Optional [ Any ] = None ,
* * kwargs : Any ,
) - > List [ Document ] :
return [
doc
for doc , in self . similarity_search_with_score_by_vector (
for doc , _ score in self . similarity_search_with_score_by_vector (
embedding , k , distance_func , * * kwargs
)
]
@ -254,7 +239,7 @@ class PGVecto_rs(VectorStore):
query_vector = self . _embedding . embed_query ( query )
return [
doc
for doc , in self . similarity_search_with_score_by_vector (
for doc , _ score in self . similarity_search_with_score_by_vector (
query_vector , k , distance_func , * * kwargs
)
]
盐粒 Yanli
6 months ago
committed by
GitHub