wehub-resource-sync
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
|
|
"""This example demonstrates working with MongoDB Atlas. You can use Mongo Atlas in 2 ways:
|
|
|
|
1. Only as a vector DB (and use local or another mongo for data storage)
|
|
In this case, you can set the following environment variable:
|
|
MONGO_ATLAS_CONNECTION_URI=mongo+srv://<username>:<password>@<cluster-domain>...
|
|
|
|
2. For both vector and data storage
|
|
In this case, you can set the following environment configuration:
|
|
LLMWareConfig().set_config("collection_db_uri", "mongo+srv://<username>:<password>@<cluster-domain>..."
|
|
|
|
This example demonstrates the 2nd approach.
|
|
|
|
"""
|
|
|
|
import os
|
|
from llmware.configs import LLMWareConfig
|
|
from llmware.library import Library
|
|
from llmware.retrieval import Query
|
|
from llmware.setup import Setup
|
|
|
|
|
|
# Use MongoDB Atlas for both data storage and vector embedding storage/search
|
|
|
|
|
|
def using_mongo_atlas(mongo_atlas_connection_string):
|
|
|
|
LLMWareConfig().set_config("collection_db_uri", mongo_atlas_connection_string)
|
|
|
|
# Create a library and populate it with some sample documents
|
|
|
|
library_name = "test_mongo_atlas"
|
|
print(f"\n > Creating library {library_name}...")
|
|
|
|
library = Library().create_new_library(library_name)
|
|
print(f"\n > Loading the llmware Sample Files...")
|
|
|
|
sample_files_path = Setup().load_sample_files()
|
|
print(f"\n > Adding some files to the library...")
|
|
|
|
library.add_files(input_folder_path=os.path.join(sample_files_path, "Agreements"))
|
|
|
|
# Create vector embeddings using Mongo Atlas
|
|
print(f"\n > Generating embedding vectors (using the 'mini-lm-sbert' model) and storing them (in 'Mongo Atlas')...")
|
|
library.install_new_embedding(embedding_model_name="mini-lm-sbert", vector_db="mongo_atlas")
|
|
|
|
# Do a semantic search in the library
|
|
print(f"\n > Running a query for 'Salary'...")
|
|
query_results = Query(library).semantic_query(query="salary", result_count=10, results_only=True)
|
|
|
|
print(f"\n\nResults for 'Salary' in {library_name}:\n")
|
|
for query_result in query_results:
|
|
print(
|
|
"File: " + query_result["file_source"] + " (Page " + str(query_result["page_num"]) + "):\n" + query_result[
|
|
"text"] + "\n")
|
|
|
|
return query_results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
# Set this env var appropriately or just paste in your Mongo Atlas connection string:
|
|
mongo_config = os.environ["MONGO_ATLAS_CONNECTION_URI"]
|
|
|
|
output = using_mongo_atlas(mongo_config)
|
|
|
|
|
|
|
|
|
|
|
|
|