gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

集中处理环境变量

Browse Source
This commit is contained in:
90houlaoheshang
2024-11-06 11:13:37 +08:00
parent deca6305a9
commit 846937195e
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
examples/lightrag_api_openai_compatible_demo.py
View File

@@ -18,6 +18,13 @@ app = FastAPI(title="LightRAG API", description="API for RAG operations")
# Configure working directory # Configure working directory
WORKING_DIR = os.environ.get("RAG_DIR", f"{DEFAULT_RAG_DIR}") WORKING_DIR = os.environ.get("RAG_DIR", f"{DEFAULT_RAG_DIR}")
print(f"WORKING_DIR: {WORKING_DIR}") print(f"WORKING_DIR: {WORKING_DIR}")
LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4o-mini")
print(f"LLM_MODEL: {LLM_MODEL}")
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large")
print(f"EMBEDDING_MODEL: {EMBEDDING_MODEL}")
EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192))
print(f"EMBEDDING_MAX_TOKEN_SIZE: {EMBEDDING_MAX_TOKEN_SIZE}")
if not os.path.exists(WORKING_DIR): if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR) os.mkdir(WORKING_DIR)
@@ -29,7 +36,7 @@ async def llm_model_func(
prompt, system_prompt=None, history_messages=[], **kwargs prompt, system_prompt=None, history_messages=[], **kwargs
) -> str: ) -> str:
return await openai_complete_if_cache( return await openai_complete_if_cache(
os.environ.get("LLM_MODEL", "gpt-4o-mini"), LLM_MODEL,
prompt, prompt,
system_prompt=system_prompt, system_prompt=system_prompt,
history_messages=history_messages, history_messages=history_messages,
@@ -43,7 +50,7 @@ async def llm_model_func(
async def embedding_func(texts: list[str]) -> np.ndarray: async def embedding_func(texts: list[str]) -> np.ndarray:
return await openai_embedding( return await openai_embedding(
texts, texts,
model=os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large"), model=EMBEDDING_MODEL,
) )
@@ -60,7 +67,7 @@ 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_dim=asyncio.run(get_embedding_dim()), embedding_func=EmbeddingFunc(embedding_dim=asyncio.run(get_embedding_dim()),
max_token_size=os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192), max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
func=embedding_func), func=embedding_func),
) )