gzdaniel commited on
Commit
4587a99
·
1 Parent(s): 15ee612

Update openai demo

Browse files
Files changed (1) hide show
  1. examples/lightrag_openai_demo.py +36 -0
examples/lightrag_openai_demo.py CHANGED
@@ -91,10 +91,46 @@ async def initialize_rag():
91
 
92
 
93
  async def main():
 
 
 
 
 
 
 
 
 
94
  try:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
95
  # Initialize RAG instance
96
  rag = await initialize_rag()
97
 
 
 
 
 
 
 
 
 
 
 
98
  with open("./book.txt", "r", encoding="utf-8") as f:
99
  await rag.ainsert(f.read())
100
 
 
91
 
92
 
93
  async def main():
94
+ # Check if OPENAI_API_KEY environment variable exists
95
+ if not os.getenv("OPENAI_API_KEY"):
96
+ print(
97
+ "Error: OPENAI_API_KEY environment variable is not set. Please set this variable before running the program."
98
+ )
99
+ print("You can set the environment variable by running:")
100
+ print(" export OPENAI_API_KEY='your-openai-api-key'")
101
+ return # Exit the async function
102
+
103
  try:
104
+ # Clear old data files
105
+ files_to_delete = [
106
+ "graph_chunk_entity_relation.graphml",
107
+ "kv_store_doc_status.json",
108
+ "kv_store_full_docs.json",
109
+ "kv_store_text_chunks.json",
110
+ "vdb_chunks.json",
111
+ "vdb_entities.json",
112
+ "vdb_relationships.json",
113
+ ]
114
+
115
+ for file in files_to_delete:
116
+ file_path = os.path.join(WORKING_DIR, file)
117
+ if os.path.exists(file_path):
118
+ os.remove(file_path)
119
+ print(f"Deleting old file:: {file_path}")
120
+
121
  # Initialize RAG instance
122
  rag = await initialize_rag()
123
 
124
+ # Test embedding function
125
+ test_text = ["This is a test string for embedding."]
126
+ embedding = await rag.embedding_func(test_text)
127
+ embedding_dim = embedding.shape[1]
128
+ print("\n=======================")
129
+ print("Test embedding function")
130
+ print("========================")
131
+ print(f"Test dict: {test_text}")
132
+ print(f"Detected embedding dimension: {embedding_dim}\n\n")
133
+
134
  with open("./book.txt", "r", encoding="utf-8") as f:
135
  await rag.ainsert(f.read())
136