gzdaniel commited on
Commit
15ee612
·
1 Parent(s): 979c6d6

Update openai compatible demo

Browse files
examples/lightrag_openai_compatible_demo.py CHANGED
@@ -7,9 +7,12 @@ from lightrag import LightRAG, QueryParam
7
  from lightrag.llm.openai import openai_complete_if_cache
8
  from lightrag.llm.ollama import ollama_embed
9
  from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
10
- import numpy as np
11
  from lightrag.kg.shared_storage import initialize_pipeline_status
12
 
 
 
 
 
13
  WORKING_DIR = "./dickens"
14
 
15
 
@@ -86,43 +89,16 @@ async def llm_model_func(
86
  prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
87
  ) -> str:
88
  return await openai_complete_if_cache(
89
- "deepseek-chat",
90
  prompt,
91
  system_prompt=system_prompt,
92
  history_messages=history_messages,
93
- api_key=os.getenv("OPENAI_API_KEY"),
94
- base_url="https://api.deepseek.com",
95
  **kwargs,
96
  )
97
 
98
 
99
- async def embedding_func(texts: list[str]) -> np.ndarray:
100
- return await ollama_embed(
101
- texts=texts,
102
- embed_model="bge-m3:latest",
103
- host="http://localhost:11434",
104
- )
105
-
106
-
107
- async def get_embedding_dim():
108
- test_text = ["This is a test sentence."]
109
- embedding = await embedding_func(test_text)
110
- embedding_dim = embedding.shape[1]
111
- return embedding_dim
112
-
113
-
114
- # function test
115
- async def test_funcs():
116
- result = await llm_model_func("How are you?")
117
- print("llm_model_func: ", result)
118
-
119
- result = await embedding_func(["How are you?"])
120
- print("embedding_func: ", result)
121
-
122
-
123
- # asyncio.run(test_funcs())
124
-
125
-
126
  async def print_stream(stream):
127
  async for chunk in stream:
128
  if chunk:
@@ -130,16 +106,17 @@ async def print_stream(stream):
130
 
131
 
132
  async def initialize_rag():
133
- embedding_dimension = await get_embedding_dim()
134
- print(f"Detected embedding dimension: {embedding_dimension}")
135
-
136
  rag = LightRAG(
137
  working_dir=WORKING_DIR,
138
  llm_model_func=llm_model_func,
139
  embedding_func=EmbeddingFunc(
140
- embedding_dim=embedding_dimension,
141
- max_token_size=8192,
142
- func=embedding_func,
 
 
 
 
143
  ),
144
  )
145
 
@@ -151,9 +128,36 @@ async def initialize_rag():
151
 
152
  async def main():
153
  try:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
154
  # Initialize RAG instance
155
  rag = await initialize_rag()
156
 
 
 
 
 
 
 
 
 
 
 
157
  with open("./book.txt", "r", encoding="utf-8") as f:
158
  await rag.ainsert(f.read())
159
 
 
7
  from lightrag.llm.openai import openai_complete_if_cache
8
  from lightrag.llm.ollama import ollama_embed
9
  from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
 
10
  from lightrag.kg.shared_storage import initialize_pipeline_status
11
 
12
+ from dotenv import load_dotenv
13
+
14
+ load_dotenv(dotenv_path=".env", override=False)
15
+
16
  WORKING_DIR = "./dickens"
17
 
18
 
 
89
  prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
90
  ) -> str:
91
  return await openai_complete_if_cache(
92
+ os.getenv("LLM_MODEL", "deepseek-chat"),
93
  prompt,
94
  system_prompt=system_prompt,
95
  history_messages=history_messages,
96
+ api_key=os.getenv("LLM_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY"),
97
+ base_url=os.getenv("LLM_BINDING_HOST", "https://api.deepseek.com"),
98
  **kwargs,
99
  )
100
 
101
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
102
  async def print_stream(stream):
103
  async for chunk in stream:
104
  if chunk:
 
106
 
107
 
108
  async def initialize_rag():
 
 
 
109
  rag = LightRAG(
110
  working_dir=WORKING_DIR,
111
  llm_model_func=llm_model_func,
112
  embedding_func=EmbeddingFunc(
113
+ embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
114
+ max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
115
+ func=lambda texts: ollama_embed(
116
+ texts,
117
+ embed_model=os.getenv("EMBEDDING_MODEL", "bge-m3:latest"),
118
+ host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
119
+ ),
120
  ),
121
  )
122
 
 
128
 
129
  async def main():
130
  try:
131
+ # Clear old data files
132
+ files_to_delete = [
133
+ "graph_chunk_entity_relation.graphml",
134
+ "kv_store_doc_status.json",
135
+ "kv_store_full_docs.json",
136
+ "kv_store_text_chunks.json",
137
+ "vdb_chunks.json",
138
+ "vdb_entities.json",
139
+ "vdb_relationships.json",
140
+ ]
141
+
142
+ for file in files_to_delete:
143
+ file_path = os.path.join(WORKING_DIR, file)
144
+ if os.path.exists(file_path):
145
+ os.remove(file_path)
146
+ print(f"Deleting old file:: {file_path}")
147
+
148
  # Initialize RAG instance
149
  rag = await initialize_rag()
150
 
151
+ # Test embedding function
152
+ test_text = ["This is a test string for embedding."]
153
+ embedding = await rag.embedding_func(test_text)
154
+ embedding_dim = embedding.shape[1]
155
+ print("\n=======================")
156
+ print("Test embedding function")
157
+ print("========================")
158
+ print(f"Test dict: {test_text}")
159
+ print(f"Detected embedding dimension: {embedding_dim}\n\n")
160
+
161
  with open("./book.txt", "r", encoding="utf-8") as f:
162
  await rag.ainsert(f.read())
163