Compare commits

...

1 Commits

Author SHA1 Message Date
leehuwuj 02f360e8b4 use ingestion pipeline for vector stores 2024-05-13 14:38:38 +07:00
32 changed files with 228 additions and 362 deletions
+2 -2
View File
@@ -37,9 +37,9 @@ const getVectorDBEnvs = (vectorDb?: TemplateVectorDB): EnvVar[] => {
case "mongo":
return [
{
name: "MONGO_URI",
name: "MONGODB_URI",
description:
"For generating a connection URI, see https://docs.timescale.com/use-timescale/latest/services/create-a-service\nThe MongoDB connection URI.",
"For generating a connection URI, see https://www.mongodb.com/docs/manual/reference/connection-string/ \nThe MongoDB connection URI.",
},
{
name: "MONGODB_DATABASE",
+26 -8
View File
@@ -1,7 +1,10 @@
import os
import logging
from llama_parse import LlamaParse
from pydantic import BaseModel, validator
logger = logging.getLogger(__name__)
class FileLoaderConfig(BaseModel):
data_dir: str = "data"
@@ -27,11 +30,26 @@ def llama_parse_parser():
def get_file_documents(config: FileLoaderConfig):
from llama_index.core.readers import SimpleDirectoryReader
reader = SimpleDirectoryReader(
config.data_dir,
recursive=True,
)
if config.use_llama_parse:
parser = llama_parse_parser()
reader.file_extractor = {".pdf": parser}
return reader.load_data()
try:
reader = SimpleDirectoryReader(
config.data_dir,
recursive=True,
filename_as_id=True,
)
if config.use_llama_parse:
parser = llama_parse_parser()
reader.file_extractor = {".pdf": parser}
return reader.load_data()
except ValueError as e:
# Carefully check is get the empty data dir and return as empty document list
import sys, traceback
_, _, exc_traceback = sys.exc_info()
function_name = traceback.extract_tb(exc_traceback)[-1].name
if function_name == "_add_files":
logger.warning(
f"Failed to load file documents, error message: {e} . Return as empty document list."
)
return []
else:
raise e
@@ -1,37 +0,0 @@
from dotenv import load_dotenv
load_dotenv()
import os
import logging
from llama_index.core.storage import StorageContext
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.astra_db import AstraDBVectorStore
from app.settings import init_settings
from app.engine.loaders import get_documents
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Creating new index")
documents = get_documents()
store = AstraDBVectorStore(
token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
api_endpoint=os.environ["ASTRA_DB_ENDPOINT"],
collection_name=os.environ["ASTRA_DB_COLLECTION"],
embedding_dimension=int(os.environ["EMBEDDING_DIM"]),
)
storage_context = StorageContext.from_defaults(vector_store=store)
VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=True, # this will show you a progress bar as the embeddings are created
)
logger.info(f"Successfully created embeddings in the AstraDB")
if __name__ == "__main__":
generate_datasource()
@@ -1,21 +0,0 @@
import logging
import os
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.astra_db import AstraDBVectorStore
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting to index from AstraDB...")
store = AstraDBVectorStore(
token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
api_endpoint=os.environ["ASTRA_DB_ENDPOINT"],
collection_name=os.environ["ASTRA_DB_COLLECTION"],
embedding_dimension=int(os.environ["EMBEDDING_DIM"]),
)
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished connecting to index from AstraDB.")
return index
@@ -0,0 +1,20 @@
import os
from llama_index.vector_stores.astra_db import AstraDBVectorStore
def get_vector_store():
endpoint = os.getenv("ASTRA_DB_ENDPOINT")
token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")
collection = os.getenv("ASTRA_DB_COLLECTION")
if not endpoint or not token or not collection:
raise ValueError(
"Please config ASTRA_DB_ENDPOINT, ASTRA_DB_APPLICATION_TOKEN and ASTRA_DB_COLLECTION"
" to your environment variables or config them in the .env file"
)
store = AstraDBVectorStore(
token=token,
api_endpoint=endpoint,
collection_name=collection,
embedding_dimension=int(os.getenv("EMBEDDING_DIM")),
)
return store
@@ -1,39 +0,0 @@
from dotenv import load_dotenv
load_dotenv()
import os
import logging
from llama_index.core.storage import StorageContext
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.milvus import MilvusVectorStore
from app.settings import init_settings
from app.engine.loaders import get_documents
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Creating new index")
# load the documents and create the index
documents = get_documents()
store = MilvusVectorStore(
uri=os.environ["MILVUS_ADDRESS"],
user=os.getenv("MILVUS_USERNAME"),
password=os.getenv("MILVUS_PASSWORD"),
collection_name=os.getenv("MILVUS_COLLECTION"),
dim=int(os.getenv("EMBEDDING_DIM")),
)
storage_context = StorageContext.from_defaults(vector_store=store)
VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=True, # this will show you a progress bar as the embeddings are created
)
logger.info(f"Successfully created embeddings in the Milvus")
if __name__ == "__main__":
generate_datasource()
@@ -1,22 +0,0 @@
import logging
import os
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.milvus import MilvusVectorStore
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting to index from Milvus...")
store = MilvusVectorStore(
uri=os.getenv("MILVUS_ADDRESS"),
user=os.getenv("MILVUS_USERNAME"),
password=os.getenv("MILVUS_PASSWORD"),
collection_name=os.getenv("MILVUS_COLLECTION"),
dim=int(os.getenv("EMBEDDING_DIM")),
)
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished connecting to index from Milvus.")
return index
@@ -0,0 +1,20 @@
import os
from llama_index.vector_stores.milvus import MilvusVectorStore
def get_vector_store():
address = os.getenv("MILVUS_ADDRESS")
collection = os.getenv("MILVUS_COLLECTION")
if not address or not collection:
raise ValueError(
"Please set MILVUS_ADDRESS and MILVUS_COLLECTION to your environment variables"
" or config them in the .env file"
)
store = MilvusVectorStore(
uri=address,
user=os.getenv("MILVUS_USERNAME"),
password=os.getenv("MILVUS_PASSWORD"),
collection_name=collection,
dim=int(os.getenv("EMBEDDING_DIM")),
)
return store
@@ -1,43 +0,0 @@
from dotenv import load_dotenv
load_dotenv()
import os
import logging
from llama_index.core.storage import StorageContext
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
from app.settings import init_settings
from app.engine.loaders import get_documents
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Creating new index")
# load the documents and create the index
documents = get_documents()
store = MongoDBAtlasVectorSearch(
db_name=os.environ["MONGODB_DATABASE"],
collection_name=os.environ["MONGODB_VECTORS"],
index_name=os.environ["MONGODB_VECTOR_INDEX"],
)
storage_context = StorageContext.from_defaults(vector_store=store)
VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=True, # this will show you a progress bar as the embeddings are created
)
logger.info(
f"Successfully created embeddings in the MongoDB collection {os.environ['MONGODB_VECTORS']}"
)
logger.info(
"""IMPORTANT: You can't query your index yet because you need to create a vector search index in MongoDB's UI now.
See https://github.com/run-llama/mongodb-demo/tree/main?tab=readme-ov-file#create-a-vector-search-index"""
)
if __name__ == "__main__":
generate_datasource()
@@ -1,20 +0,0 @@
import logging
import os
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting to index from MongoDB...")
store = MongoDBAtlasVectorSearch(
db_name=os.environ["MONGODB_DATABASE"],
collection_name=os.environ["MONGODB_VECTORS"],
index_name=os.environ["MONGODB_VECTOR_INDEX"],
)
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished connecting to index from MongoDB.")
return index
@@ -0,0 +1,20 @@
import os
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
def get_vector_store():
db_uri = os.getenv("MONGODB_URI")
db_name = os.getenv("MONGODB_DATABASE")
collection_name = os.getenv("MONGODB_VECTORS")
index_name = os.getenv("MONGODB_VECTOR_INDEX")
if not db_uri or not db_name or not collection_name or not index_name:
raise ValueError(
"Please set MONGODB_URI, MONGODB_DATABASE, MONGODB_VECTORS, and MONGODB_VECTOR_INDEX"
" to your environment variables or config them in .env file"
)
store = MongoDBAtlasVectorSearch(
db_name=db_name,
collection_name=collection_name,
index_name=index_name,
)
return store
@@ -1,2 +0,0 @@
PGVECTOR_SCHEMA = "public"
PGVECTOR_TABLE = "llamaindex_embedding"
@@ -1,35 +0,0 @@
from dotenv import load_dotenv
load_dotenv()
import logging
from llama_index.core.indices import VectorStoreIndex
from llama_index.core.storage import StorageContext
from app.engine.loaders import get_documents
from app.settings import init_settings
from app.engine.utils import init_pg_vector_store_from_env
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Creating new index")
# load the documents and create the index
documents = get_documents()
store = init_pg_vector_store_from_env()
storage_context = StorageContext.from_defaults(vector_store=store)
VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=True, # this will show you a progress bar as the embeddings are created
)
logger.info(
f"Successfully created embeddings in the PG vector store, schema={store.schema_name} table={store.table_name}"
)
if __name__ == "__main__":
generate_datasource()
@@ -1,13 +0,0 @@
import logging
from llama_index.core.indices.vector_store import VectorStoreIndex
from app.engine.utils import init_pg_vector_store_from_env
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting to index from PGVector...")
store = init_pg_vector_store_from_env()
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished connecting to index from PGVector.")
return index
@@ -1,10 +1,12 @@
import os
from llama_index.vector_stores.postgres import PGVectorStore
from urllib.parse import urlparse
from app.engine.constants import PGVECTOR_SCHEMA, PGVECTOR_TABLE
PGVECTOR_SCHEMA = "public"
PGVECTOR_TABLE = "llamaindex_embedding"
def init_pg_vector_store_from_env():
def get_vector_store():
original_conn_string = os.environ.get("PG_CONNECTION_STRING")
if original_conn_string is None or original_conn_string == "":
raise ValueError("PG_CONNECTION_STRING environment variable is not set.")
@@ -1,39 +0,0 @@
from dotenv import load_dotenv
load_dotenv()
import os
import logging
from llama_index.core.storage import StorageContext
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.pinecone import PineconeVectorStore
from app.settings import init_settings
from app.engine.loaders import get_documents
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Creating new index")
# load the documents and create the index
documents = get_documents()
store = PineconeVectorStore(
api_key=os.environ["PINECONE_API_KEY"],
index_name=os.environ["PINECONE_INDEX_NAME"],
environment=os.environ["PINECONE_ENVIRONMENT"],
)
storage_context = StorageContext.from_defaults(vector_store=store)
VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=True, # this will show you a progress bar as the embeddings are created
)
logger.info(
f"Successfully created embeddings and save to your Pinecone index {os.environ['PINECONE_INDEX_NAME']}"
)
if __name__ == "__main__":
generate_datasource()
@@ -1,20 +0,0 @@
import logging
import os
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.pinecone import PineconeVectorStore
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting to index from Pinecone...")
store = PineconeVectorStore(
api_key=os.environ["PINECONE_API_KEY"],
index_name=os.environ["PINECONE_INDEX_NAME"],
environment=os.environ["PINECONE_ENVIRONMENT"],
)
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished connecting to index from Pinecone.")
return index
@@ -0,0 +1,19 @@
import os
from llama_index.vector_stores.pinecone import PineconeVectorStore
def get_vector_store():
api_key = os.getenv("PINECONE_API_KEY")
index_name = os.getenv("PINECONE_INDEX_NAME")
environment = os.getenv("PINECONE_ENVIRONMENT")
if not api_key or not index_name or not environment:
raise ValueError(
"Please set PINECONE_API_KEY, PINECONE_INDEX_NAME, and PINECONE_ENVIRONMENT"
" to your environment variables or config them in the .env file"
)
store = PineconeVectorStore(
api_key=api_key,
index_name=index_name,
environment=environment,
)
return store
@@ -1,37 +0,0 @@
import logging
import os
from app.engine.loaders import get_documents
from app.settings import init_settings
from dotenv import load_dotenv
from llama_index.core.indices import VectorStoreIndex
from llama_index.core.storage import StorageContext
from llama_index.vector_stores.qdrant import QdrantVectorStore
load_dotenv()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Creating new index with Qdrant")
# load the documents and create the index
documents = get_documents()
store = QdrantVectorStore(
collection_name=os.getenv("QDRANT_COLLECTION"),
url=os.getenv("QDRANT_URL"),
api_key=os.getenv("QDRANT_API_KEY"),
)
storage_context = StorageContext.from_defaults(vector_store=store)
VectorStoreIndex.from_documents(
documents,
storage_context=storage_context,
show_progress=True, # this will show you a progress bar as the embeddings are created
)
logger.info(
f"Successfully uploaded documents to the {os.getenv('QDRANT_COLLECTION')} collection."
)
if __name__ == "__main__":
generate_datasource()
@@ -1,20 +0,0 @@
import logging
import os
from llama_index.core.indices import VectorStoreIndex
from llama_index.vector_stores.qdrant import QdrantVectorStore
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting to Qdrant collection..")
store = QdrantVectorStore(
collection_name=os.getenv("QDRANT_COLLECTION"),
url=os.getenv("QDRANT_URL"),
api_key=os.getenv("QDRANT_API_KEY"),
)
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished connecting to Qdrant collection.")
return index
@@ -0,0 +1,19 @@
import os
from llama_index.vector_stores.qdrant import QdrantVectorStore
def get_vector_store():
collection_name = os.getenv("QDRANT_COLLECTION")
url = os.getenv("QDRANT_URL")
api_key = os.getenv("QDRANT_API_KEY")
if not collection_name or not url:
raise ValueError(
"Please set QDRANT_COLLECTION, QDRANT_URL"
" to your environment variables or config them in the .env file"
)
store = QdrantVectorStore(
collection_name=collection_name,
url=url,
api_key=api_key,
)
return store
@@ -12,7 +12,7 @@ import { checkRequiredEnvVars } from "./shared";
dotenv.config();
const mongoUri = process.env.MONGO_URI!;
const mongoUri = process.env.MONGODB_URI!;
const databaseName = process.env.MONGODB_DATABASE!;
const vectorCollectionName = process.env.MONGODB_VECTORS!;
const indexName = process.env.MONGODB_VECTOR_INDEX;
@@ -1,5 +1,5 @@
const REQUIRED_ENV_VARS = [
"MONGO_URI",
"MONGODB_URI",
"MONGODB_DATABASE",
"MONGODB_VECTORS",
"MONGODB_VECTOR_INDEX",
@@ -0,0 +1 @@
STORAGE_DIR = "storage"
@@ -0,0 +1,79 @@
from dotenv import load_dotenv
load_dotenv()
import os
import logging
import asyncio
from llama_index.core.settings import Settings
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.storage import StorageContext
from app.constants import STORAGE_DIR
from app.settings import init_settings
from app.engine.loaders import get_documents
from app.engine.vectordb import get_vector_store
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def get_doc_store():
if not os.path.exists(STORAGE_DIR):
docstore = SimpleDocumentStore()
return docstore
else:
logger.info(f"Loading existing docstore from {STORAGE_DIR}")
return SimpleDocumentStore.from_persist_dir(STORAGE_DIR)
def run_pipeline(docstore, vector_store, documents):
pipeline = IngestionPipeline(
transformations=[
SentenceSplitter(
chunk_size=Settings.chunk_size,
chunk_overlap=Settings.chunk_overlap,
),
Settings.embed_model,
],
docstore=docstore,
docstore_strategy="upserts_and_delete",
vector_store=vector_store,
)
# Run the ingestion pipeline and store the results
nodes = pipeline.run(show_progress=True, documents=documents)
return nodes
def persist_storage(docstore, vector_store):
storage_context = StorageContext.from_defaults(
docstore=docstore,
vector_store=vector_store,
)
storage_context.persist(STORAGE_DIR)
def generate_datasource():
init_settings()
logger.info("Generate index for the provided data")
# Get the stores and documents or create new ones
documents = get_documents()
docstore = get_doc_store()
vector_store = get_vector_store()
# Run the ingestion pipeline
_ = run_pipeline(docstore, vector_store, documents)
# Build the index and persist storage
persist_storage(docstore, vector_store)
logger.info("Finished generating the index")
if __name__ == "__main__":
generate_datasource()
@@ -0,0 +1,16 @@
import logging
import os
from llama_index.core.indices import VectorStoreIndex
from app.engine.vectordb import get_vector_store
logger = logging.getLogger("uvicorn")
def get_index():
logger.info("Connecting vector store...")
store = get_vector_store()
index = VectorStoreIndex.from_vector_store(store)
logger.info("Finished load index from vector store.")
return index