mirror of
https://github.com/run-llama/create-llama.git
synced 2026-07-09 15:26:47 -04:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 02f360e8b4 |
@@ -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",
|
||||
|
||||
@@ -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
|
||||
+4
-2
@@ -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
|
||||
Reference in New Issue
Block a user