266 lines
8.8 KiB
Python
266 lines
8.8 KiB
Python
|
import hashlib
|
||
|
|
from abc import ABC, abstractmethod
|
||
|
|
from collections import namedtuple
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from typing import Any, Dict, List, Optional
|
||
|
|
|
||
|
|
from pydantic import Field
|
||
|
|
|
||
|
|
from guardrails.vectordb import VectorDBBase
|
||
|
|
|
||
|
|
try:
|
||
|
|
from sqlalchemy.exc import IntegrityError
|
||
|
|
except ImportError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Document:
|
||
|
|
"""Document holds text and metadata of a document.
|
||
|
|
|
||
|
|
Examples of documents are PDFs, Word documents, etc. A collection of
|
||
|
|
related text in an NLP application can be thought of a document as
|
||
|
|
well.
|
||
|
|
"""
|
||
|
|
|
||
|
|
id: str
|
||
|
|
pages: Dict[int, str]
|
||
|
|
metadata: Dict[Any, Any] = Field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
# PageCoordinates is a datastructure that points to the location
|
||
|
|
# of a page in a document.
|
||
|
|
PageCoordinates = namedtuple("PageCoordinates", ["doc_id", "page_num"])
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Page:
|
||
|
|
"""Page holds text and metadata of a page in a document.
|
||
|
|
|
||
|
|
It also containts the coordinates of the page in the document.
|
||
|
|
"""
|
||
|
|
|
||
|
|
cordinates: PageCoordinates
|
||
|
|
text: str
|
||
|
|
metadata: Dict[Any, Any]
|
||
|
|
|
||
|
|
|
||
|
|
class DocumentStoreBase(ABC):
|
||
|
|
"""Abstract class for a store that can store text, and metadata from
|
||
|
|
documents.
|
||
|
|
|
||
|
|
The store can be queried by text for similar documents.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, vector_db: VectorDBBase, path: Optional[str] = None): ...
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def add_document(self, document: Document) -> None:
|
||
|
|
"""Adds a document to the store.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
document: Document object to be added
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
None if the document was added successfully
|
||
|
|
"""
|
||
|
|
...
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def search(self, query: str, k: int = 4) -> List[Page]:
|
||
|
|
"""Searches for pages which contain the text similar to the query.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
query: Text to search for.
|
||
|
|
k: Number of similar pages to return.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List[Pages] List of pages which contains similar texts
|
||
|
|
"""
|
||
|
|
...
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def add_text(self, text: str, meta: Dict[Any, Any]) -> str:
|
||
|
|
"""Adds a text to the store.
|
||
|
|
Args:
|
||
|
|
text: Text to add.
|
||
|
|
meta: Metadata to associate with the text.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The id of the text.
|
||
|
|
"""
|
||
|
|
...
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def add_texts(self, texts: Dict[str, Dict[Any, Any]]) -> List[str]:
|
||
|
|
"""Adds a list of texts to the store.
|
||
|
|
Args:
|
||
|
|
texts: List of texts to add, and their associalted metadata.
|
||
|
|
example:
|
||
|
|
``` json
|
||
|
|
[{"I am feeling good", {"sentiment": "postive"}}]
|
||
|
|
```
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of ids of the texts."""
|
||
|
|
...
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def flush():
|
||
|
|
"""Flushes the store to disk."""
|
||
|
|
...
|
||
|
|
|
||
|
|
|
||
|
|
try:
|
||
|
|
import sqlalchemy
|
||
|
|
from sqlalchemy.orm import Mapped, Session, declarative_base, mapped_column
|
||
|
|
|
||
|
|
class RealEphemeralDocumentStore(DocumentStoreBase):
|
||
|
|
"""EphemeralDocumentStore is a document store that stores the documents
|
||
|
|
on local disk and use a ephemeral vector store like Faiss."""
|
||
|
|
|
||
|
|
def __init__(self, vector_db: VectorDBBase, path: Optional[str] = None):
|
||
|
|
"""Creates a new EphemeralDocumentStore.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
vector_db: VectorDBBase instance to use for storing the vectors.
|
||
|
|
path: Path to the database file store metadata.
|
||
|
|
"""
|
||
|
|
self._vector_db = vector_db
|
||
|
|
self._storage = RealSQLMetadataStore(path=path)
|
||
|
|
|
||
|
|
def add_document(self, document: Document):
|
||
|
|
# Add the document, in case the document is already there it
|
||
|
|
# would raise an exception and we assume the document and
|
||
|
|
# vectors are present.
|
||
|
|
try:
|
||
|
|
self._storage.add_docs(
|
||
|
|
[document], vdb_last_index=self._vector_db.last_index()
|
||
|
|
)
|
||
|
|
except IntegrityError:
|
||
|
|
return
|
||
|
|
self._vector_db.add_texts(list(document.pages.values()))
|
||
|
|
|
||
|
|
def add_text(self, text: str, meta: Dict[Any, Any]) -> str:
|
||
|
|
hash = hashlib.md5()
|
||
|
|
hash.update(text.encode("utf-8"))
|
||
|
|
hash.update(str(meta).encode("utf-8"))
|
||
|
|
id = hash.hexdigest()
|
||
|
|
|
||
|
|
doc = Document(id, {0: text}, meta)
|
||
|
|
self.add_document(doc)
|
||
|
|
return doc.id
|
||
|
|
|
||
|
|
def add_texts(self, texts: Dict[str, Dict[Any, Any]]) -> List[str]:
|
||
|
|
doc_ids = []
|
||
|
|
for text, meta in texts.items():
|
||
|
|
doc_id = self.add_text(text, meta)
|
||
|
|
doc_ids.append(doc_id)
|
||
|
|
return doc_ids
|
||
|
|
|
||
|
|
def search(self, query: str, k: int = 4) -> List[Page]:
|
||
|
|
vector_db_indexes = self._vector_db.similarity_search(query, k)
|
||
|
|
filtered_ids = list(filter(lambda x: x != -1, vector_db_indexes))
|
||
|
|
return self._storage.get_pages_for_for_indexes(filtered_ids)
|
||
|
|
|
||
|
|
def search_with_threshold(
|
||
|
|
self, query: str, threshold: float, k: int = 4
|
||
|
|
) -> List[Page]:
|
||
|
|
vector_db_indexes = self._vector_db.similarity_search_with_threshold(
|
||
|
|
query, k, threshold
|
||
|
|
)
|
||
|
|
filtered_ids = list(filter(lambda x: x != -1, vector_db_indexes))
|
||
|
|
return self._storage.get_pages_for_for_indexes(filtered_ids)
|
||
|
|
|
||
|
|
def flush(self, path: Optional[str] = None):
|
||
|
|
self._vector_db.save(path)
|
||
|
|
|
||
|
|
Base = declarative_base()
|
||
|
|
|
||
|
|
class RealSqlDocument(Base):
|
||
|
|
__tablename__ = "documents"
|
||
|
|
|
||
|
|
id: Mapped[int] = mapped_column(primary_key=True) # type: ignore
|
||
|
|
page_num: Mapped[int] = mapped_column(sqlalchemy.Integer, primary_key=True) # type: ignore
|
||
|
|
text: Mapped[str] = mapped_column(sqlalchemy.String) # type: ignore
|
||
|
|
meta: Mapped[dict] = mapped_column(sqlalchemy.PickleType) # type: ignore
|
||
|
|
vector_index: Mapped[int] = mapped_column(sqlalchemy.Integer) # type: ignore
|
||
|
|
|
||
|
|
class RealSQLMetadataStore:
|
||
|
|
def __init__(self, path: Optional[str] = None):
|
||
|
|
conn = f"sqlite:///{path}" if path is not None else "sqlite://"
|
||
|
|
self._engine = sqlalchemy.create_engine(conn) # type: ignore
|
||
|
|
RealSqlDocument.metadata.create_all(self._engine, checkfirst=True)
|
||
|
|
|
||
|
|
def add_docs(self, docs: List[Document], vdb_last_index: int):
|
||
|
|
vector_id = vdb_last_index
|
||
|
|
with Session(self._engine) as session:
|
||
|
|
for doc in docs:
|
||
|
|
for page_num, text in doc.pages.items():
|
||
|
|
session.add(
|
||
|
|
RealSqlDocument(
|
||
|
|
id=doc.id,
|
||
|
|
page_num=page_num,
|
||
|
|
text=text,
|
||
|
|
meta=doc.metadata,
|
||
|
|
vector_index=vector_id,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
vector_id += 1
|
||
|
|
|
||
|
|
session.commit()
|
||
|
|
|
||
|
|
def get_pages_for_for_indexes(self, indexes: List[int]) -> List[Page]:
|
||
|
|
pages: List[Page] = []
|
||
|
|
with Session(self._engine) as session:
|
||
|
|
for index in indexes:
|
||
|
|
query = sqlalchemy.select(RealSqlDocument).where(
|
||
|
|
RealSqlDocument.vector_index == index
|
||
|
|
)
|
||
|
|
sql_docs = session.execute(query)
|
||
|
|
sql_doc = sql_docs.first()
|
||
|
|
if sql_doc is None:
|
||
|
|
continue
|
||
|
|
sql_doc = sql_doc[0]
|
||
|
|
pages.append(
|
||
|
|
Page(
|
||
|
|
PageCoordinates(sql_doc.id, sql_doc.page_num),
|
||
|
|
sql_doc.text,
|
||
|
|
sql_doc.meta,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
return pages
|
||
|
|
|
||
|
|
EphemeralDocumentStore = RealEphemeralDocumentStore
|
||
|
|
SQLDocument = RealSqlDocument
|
||
|
|
SQLMetadataStore = RealSQLMetadataStore
|
||
|
|
except ImportError:
|
||
|
|
|
||
|
|
class FallbackEphemeralDocumentStore:
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
# Why don't we just raise this when the import
|
||
|
|
# error occurs instead of at runtime?
|
||
|
|
raise ImportError(
|
||
|
|
"SQLAlchemy is required for EphemeralDocumentStore"
|
||
|
|
"Please install it using `poetry add SqlAlchemy`"
|
||
|
|
)
|
||
|
|
|
||
|
|
class FallbackSQLDocument:
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
raise ImportError(
|
||
|
|
"SQLAlchemy is required for SQLDocument"
|
||
|
|
"Please install it using `poetry add SqlAlchemy`"
|
||
|
|
)
|
||
|
|
|
||
|
|
class FallbackSQLMetadataStore:
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
raise ImportError(
|
||
|
|
"SQLAlchemy is required for SQLMetadataStore"
|
||
|
|
"Please install it using `poetry add SqlAlchemy`"
|
||
|
|
)
|
||
|
|
|
||
|
|
EphemeralDocumentStore = FallbackEphemeralDocumentStore
|
||
|
|
SQLDocument = FallbackSQLDocument
|
||
|
|
SQLMetadataStore = FallbackSQLMetadataStore
|