#!/usr/bin/env python

# index.py - create a set of word embeddings and cache them for later use
# see: https://github.com/sunny2309/llamaindex_tutorials/

# Eric Lease Morgan <emorgan@nd.edu>
# (c) University of Notre Dame; distributed under a GNU Public License

# March 22, 2024 - got it working using LlamaIndex and an open source model
# March 23, 2024 - migrating to study carrels
# April 26, 2025 - implemented Ollama embeddings, thus removing network necessity


# configure
EMBEDDINGS    = 'all-minilm'
CONFIGURATION = 'localLibrary'
LLM           = 'llm'
CHUNKSIZE     = 1024
CHUNKOVERLAP  = 20
EXTENSION     = '.txt'

# require
from json                               import loads
from llama_index.core                   import Document, Settings, VectorStoreIndex
from llama_index.core.node_parser       import SimpleNodeParser, SentenceSplitter
from llama_index.embeddings.ollama      import OllamaEmbedding
from pathlib                            import Path
from rdr                                import bibliography, configuration, TXT, ETC
from sys                                import stderr, exit, argv

# get input
if len( argv ) != 2 : exit( "Usage: " + argv[ 0 ] + " <carrel>" )
carrel = argv[ 1 ]

# initialize
library              = Path( configuration( CONFIGURATION ) )
txt                  = library/carrel/TXT
llm                  = library/carrel/ETC/LLM
parser               = SentenceSplitter( chunk_size=CHUNKSIZE, chunk_overlap=CHUNKOVERLAP )
Settings.embed_model = OllamaEmbedding( model_name=EMBEDDINGS )

# make sane
llm.mkdir( exist_ok=True ) 

# process each bibliographic item; create a list of documents
stderr.write( "Step #1 of 4: Reading documents\n" )
bibliographics = loads( bibliography( carrel, format='json' ) )
documents      = []
for bibliographic in bibliographics :

	# parse
	identifier = bibliographic[ 'id' ]
	author     = bibliographic[ 'author' ]
	title      = bibliographic[ 'title' ]
	date       = bibliographic[ 'date' ]
	extension  = bibliographic[ 'extension' ]
	
	# parse some more; be careful!
	#author = author.split( ',' )[ 0 ]
	
	# create and update a document
	with open( txt/( identifier + EXTENSION ) ) as handle : document = Document( text=handle.read() )
	document.doc_id   = identifier
	document.metadata = { 'identifier' : identifier,
	                      'author'     : author, 
	                      'title'      : title,
	                      'date'       : date,
	                      'file'       : identifier + EXTENSION }

	# update the list of documents
	documents.append( document )

# parse the documents into nodes
stderr.write( "Step #2 of 4: Parsing documents\n" )
nodes = parser.get_nodes_from_documents( documents )

# index the nodes
stderr.write( "Step #3 of 4: Indexing\n" )
index = VectorStoreIndex( nodes )

# save and done
stderr.write( "Step #4 of 4: Saving\n" )
index.storage_context.persist( persist_dir=llm )
exit()