Update openai demo
Browse files
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 |
|