From 43948d6f17e482f84bc3080b10deca08663466e4 Mon Sep 17 00:00:00 2001 From: yangdx Date: Tue, 13 May 2025 18:27:55 +0800 Subject: [PATCH] Update openai demo --- examples/lightrag_openai_demo.py | 36 ++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/examples/lightrag_openai_demo.py b/examples/lightrag_openai_demo.py index 24e0257c..fa0b37f1 100644 --- a/examples/lightrag_openai_demo.py +++ b/examples/lightrag_openai_demo.py @@ -91,10 +91,46 @@ async def initialize_rag(): async def main(): + # Check if OPENAI_API_KEY environment variable exists + if not os.getenv("OPENAI_API_KEY"): + print( + "Error: OPENAI_API_KEY environment variable is not set. Please set this variable before running the program." + ) + print("You can set the environment variable by running:") + print(" export OPENAI_API_KEY='your-openai-api-key'") + return # Exit the async function + try: + # Clear old data files + files_to_delete = [ + "graph_chunk_entity_relation.graphml", + "kv_store_doc_status.json", + "kv_store_full_docs.json", + "kv_store_text_chunks.json", + "vdb_chunks.json", + "vdb_entities.json", + "vdb_relationships.json", + ] + + for file in files_to_delete: + file_path = os.path.join(WORKING_DIR, file) + if os.path.exists(file_path): + os.remove(file_path) + print(f"Deleting old file:: {file_path}") + # Initialize RAG instance rag = await initialize_rag() + # Test embedding function + test_text = ["This is a test string for embedding."] + embedding = await rag.embedding_func(test_text) + embedding_dim = embedding.shape[1] + print("\n=======================") + print("Test embedding function") + print("========================") + print(f"Test dict: {test_text}") + print(f"Detected embedding dimension: {embedding_dim}\n\n") + with open("./book.txt", "r", encoding="utf-8") as f: await rag.ainsert(f.read())