Update openai compatible demo
This commit is contained in:
@@ -7,9 +7,12 @@ from lightrag import LightRAG, QueryParam
|
|||||||
from lightrag.llm.openai import openai_complete_if_cache
|
from lightrag.llm.openai import openai_complete_if_cache
|
||||||
from lightrag.llm.ollama import ollama_embed
|
from lightrag.llm.ollama import ollama_embed
|
||||||
from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
|
from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
|
||||||
import numpy as np
|
|
||||||
from lightrag.kg.shared_storage import initialize_pipeline_status
|
from lightrag.kg.shared_storage import initialize_pipeline_status
|
||||||
|
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv(dotenv_path=".env", override=False)
|
||||||
|
|
||||||
WORKING_DIR = "./dickens"
|
WORKING_DIR = "./dickens"
|
||||||
|
|
||||||
|
|
||||||
@@ -86,43 +89,16 @@ async def llm_model_func(
|
|||||||
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
|
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
|
||||||
) -> str:
|
) -> str:
|
||||||
return await openai_complete_if_cache(
|
return await openai_complete_if_cache(
|
||||||
"deepseek-chat",
|
os.getenv("LLM_MODEL", "deepseek-chat"),
|
||||||
prompt,
|
prompt,
|
||||||
system_prompt=system_prompt,
|
system_prompt=system_prompt,
|
||||||
history_messages=history_messages,
|
history_messages=history_messages,
|
||||||
api_key=os.getenv("OPENAI_API_KEY"),
|
api_key=os.getenv("LLM_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY"),
|
||||||
base_url="https://api.deepseek.com",
|
base_url=os.getenv("LLM_BINDING_HOST", "https://api.deepseek.com"),
|
||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
async def embedding_func(texts: list[str]) -> np.ndarray:
|
|
||||||
return await ollama_embed(
|
|
||||||
texts=texts,
|
|
||||||
embed_model="bge-m3:latest",
|
|
||||||
host="http://localhost:11434",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
async def get_embedding_dim():
|
|
||||||
test_text = ["This is a test sentence."]
|
|
||||||
embedding = await embedding_func(test_text)
|
|
||||||
embedding_dim = embedding.shape[1]
|
|
||||||
return embedding_dim
|
|
||||||
|
|
||||||
|
|
||||||
# function test
|
|
||||||
async def test_funcs():
|
|
||||||
result = await llm_model_func("How are you?")
|
|
||||||
print("llm_model_func: ", result)
|
|
||||||
|
|
||||||
result = await embedding_func(["How are you?"])
|
|
||||||
print("embedding_func: ", result)
|
|
||||||
|
|
||||||
|
|
||||||
# asyncio.run(test_funcs())
|
|
||||||
|
|
||||||
|
|
||||||
async def print_stream(stream):
|
async def print_stream(stream):
|
||||||
async for chunk in stream:
|
async for chunk in stream:
|
||||||
if chunk:
|
if chunk:
|
||||||
@@ -130,16 +106,17 @@ async def print_stream(stream):
|
|||||||
|
|
||||||
|
|
||||||
async def initialize_rag():
|
async def initialize_rag():
|
||||||
embedding_dimension = await get_embedding_dim()
|
|
||||||
print(f"Detected embedding dimension: {embedding_dimension}")
|
|
||||||
|
|
||||||
rag = LightRAG(
|
rag = LightRAG(
|
||||||
working_dir=WORKING_DIR,
|
working_dir=WORKING_DIR,
|
||||||
llm_model_func=llm_model_func,
|
llm_model_func=llm_model_func,
|
||||||
embedding_func=EmbeddingFunc(
|
embedding_func=EmbeddingFunc(
|
||||||
embedding_dim=embedding_dimension,
|
embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
|
||||||
max_token_size=8192,
|
max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
|
||||||
func=embedding_func,
|
func=lambda texts: ollama_embed(
|
||||||
|
texts,
|
||||||
|
embed_model=os.getenv("EMBEDDING_MODEL", "bge-m3:latest"),
|
||||||
|
host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
|
||||||
|
),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -151,9 +128,36 @@ async def initialize_rag():
|
|||||||
|
|
||||||
async def main():
|
async def main():
|
||||||
try:
|
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
|
# Initialize RAG instance
|
||||||
rag = await initialize_rag()
|
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:
|
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||||
await rag.ainsert(f.read())
|
await rag.ainsert(f.read())
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user