diff --git a/examples/lightrag_yangdx.py b/examples/lightrag_yangdx.py index 162900c4..72f42cc4 100644 --- a/examples/lightrag_yangdx.py +++ b/examples/lightrag_yangdx.py @@ -2,34 +2,46 @@ import asyncio import os import inspect import logging +from dotenv import load_dotenv from lightrag import LightRAG, QueryParam -from lightrag.llm import ollama_model_complete, ollama_embedding +from lightrag.llm import openai_complete_if_cache, ollama_embedding from lightrag.utils import EmbeddingFunc -WORKING_DIR = "./dickens" +load_dotenv() + +WORKING_DIR = "./examples/input" logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO) +async def llm_model_func( + prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs +) -> str: + return await openai_complete_if_cache( + "deepseek-chat", + prompt, + system_prompt=system_prompt, + history_messages=history_messages, + api_key=os.getenv("DEEPSEEK_API_KEY"), + base_url=os.getenv("DEEPSEEK__ENDPOINT"), + **kwargs, + ) + if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) rag = LightRAG( working_dir=WORKING_DIR, - llm_model_func=ollama_model_complete, - llm_model_name="gemma2:2b", - llm_model_max_async=4, - llm_model_max_token_size=32768, - llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}}, + llm_model_func=llm_model_func, embedding_func=EmbeddingFunc( embedding_dim=768, max_token_size=8192, func=lambda texts: ollama_embedding( - texts, embed_model="nomic-embed-text", host="http://localhost:11434" + texts, embed_model="nomic-embed-text", host="http://m4.lan.znipower.com:11434" ), ), ) -with open("./book.txt", "r", encoding="utf-8") as f: +with open("./input/book.txt", "r", encoding="utf-8") as f: rag.insert(f.read()) # Perform naive search