99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
import os
|
|
from llama_index.core import VectorStoreIndex, StorageContext
|
|
|
|
def load_index(save_dir: str = "mail_index"):
|
|
"""
|
|
Load the saved index from disk.
|
|
|
|
Args:
|
|
save_dir: Directory where the index is saved
|
|
|
|
Returns:
|
|
Loaded index or None if loading fails
|
|
"""
|
|
try:
|
|
# Load storage context
|
|
storage_context = StorageContext.from_defaults(persist_dir=save_dir)
|
|
|
|
# Load index
|
|
index = VectorStoreIndex.from_vector_store(
|
|
storage_context.vector_store,
|
|
storage_context=storage_context
|
|
)
|
|
|
|
print(f"Index loaded from {save_dir}")
|
|
return index
|
|
|
|
except Exception as e:
|
|
print(f"Error loading index: {e}")
|
|
return None
|
|
|
|
def query_index(index, query: str):
|
|
"""
|
|
Query the loaded index.
|
|
|
|
Args:
|
|
index: The loaded index
|
|
query: The query string
|
|
"""
|
|
if index is None:
|
|
print("No index available for querying.")
|
|
return
|
|
|
|
query_engine = index.as_query_engine()
|
|
response = query_engine.query(query)
|
|
print(f"\nQuery: {query}")
|
|
print(f"Response: {response}")
|
|
|
|
def main():
|
|
save_dir = "mail_index"
|
|
|
|
# Check if index exists
|
|
if not os.path.exists(save_dir) or not os.path.exists(os.path.join(save_dir, "vector_store.json")):
|
|
print(f"Index not found in {save_dir}")
|
|
print("Please run mail_reader_save_load.py first to create the index.")
|
|
return
|
|
|
|
# Load the index
|
|
index = load_index(save_dir)
|
|
|
|
if not index:
|
|
print("Failed to load index.")
|
|
return
|
|
|
|
print("\n" + "="*60)
|
|
print("Email Query Interface")
|
|
print("="*60)
|
|
print("Type 'quit' to exit")
|
|
print("Type 'help' for example queries")
|
|
print("="*60)
|
|
|
|
# Interactive query loop
|
|
while True:
|
|
try:
|
|
query = input("\nEnter your query: ").strip()
|
|
|
|
if query.lower() == 'quit':
|
|
print("Goodbye!")
|
|
break
|
|
elif query.lower() == 'help':
|
|
print("\nExample queries:")
|
|
print("- Hows Berkeley Graduate Student Instructor")
|
|
print("- What emails mention GSR appointments?")
|
|
print("- Find emails about deadlines")
|
|
print("- Search for emails from specific sender")
|
|
print("- Find emails about meetings")
|
|
continue
|
|
elif not query:
|
|
continue
|
|
|
|
query_index(index, query)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nGoodbye!")
|
|
break
|
|
except Exception as e:
|
|
print(f"Error processing query: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |