chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,114 @@
|
||||
"""
|
||||
Sometimes you need to switch a storage solution, but you want to save LLM token and time.
|
||||
This handy script helps you to copy the LLM caches from one storage solution to another.
|
||||
(Not all the storage impl are supported)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from lightrag.kg.postgres_impl import PostgreSQLDB, PGKVStorage
|
||||
from lightrag.kg.json_kv_impl import JsonKVStorage
|
||||
from lightrag.namespace import NameSpace
|
||||
|
||||
load_dotenv()
|
||||
ROOT_DIR = os.environ.get("ROOT_DIR")
|
||||
WORKING_DIR = f"{ROOT_DIR}/dickens"
|
||||
|
||||
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
# AGE
|
||||
os.environ["AGE_GRAPH_NAME"] = "chinese"
|
||||
|
||||
postgres_db = PostgreSQLDB(
|
||||
config={
|
||||
"host": "localhost",
|
||||
"port": 15432,
|
||||
"user": "rag",
|
||||
"password": "rag",
|
||||
"database": "r2",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def copy_from_postgres_to_json():
|
||||
await postgres_db.initdb()
|
||||
|
||||
from_llm_response_cache = PGKVStorage(
|
||||
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
||||
global_config={"embedding_batch_num": 6},
|
||||
embedding_func=None,
|
||||
db=postgres_db,
|
||||
)
|
||||
|
||||
to_llm_response_cache = JsonKVStorage(
|
||||
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
||||
global_config={"working_dir": WORKING_DIR},
|
||||
embedding_func=None,
|
||||
)
|
||||
|
||||
# Get all cache data using the new flattened structure
|
||||
all_data = await from_llm_response_cache.get_all()
|
||||
|
||||
# Convert flattened data to hierarchical structure for JsonKVStorage
|
||||
kv = {}
|
||||
for flattened_key, cache_entry in all_data.items():
|
||||
# Parse flattened key: {mode}:{cache_type}:{hash}
|
||||
parts = flattened_key.split(":", 2)
|
||||
if len(parts) == 3:
|
||||
mode, cache_type, hash_value = parts
|
||||
if mode not in kv:
|
||||
kv[mode] = {}
|
||||
kv[mode][hash_value] = cache_entry
|
||||
print(f"Copying {flattened_key} -> {mode}[{hash_value}]")
|
||||
else:
|
||||
print(f"Skipping invalid key format: {flattened_key}")
|
||||
|
||||
await to_llm_response_cache.upsert(kv)
|
||||
await to_llm_response_cache.index_done_callback()
|
||||
print("Mission accomplished!")
|
||||
|
||||
|
||||
async def copy_from_json_to_postgres():
|
||||
await postgres_db.initdb()
|
||||
|
||||
from_llm_response_cache = JsonKVStorage(
|
||||
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
||||
global_config={"working_dir": WORKING_DIR},
|
||||
embedding_func=None,
|
||||
)
|
||||
|
||||
to_llm_response_cache = PGKVStorage(
|
||||
namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
|
||||
global_config={"embedding_batch_num": 6},
|
||||
embedding_func=None,
|
||||
db=postgres_db,
|
||||
)
|
||||
|
||||
# Get all cache data from JsonKVStorage (hierarchical structure)
|
||||
all_data = await from_llm_response_cache.get_all()
|
||||
|
||||
# Convert hierarchical data to flattened structure for PGKVStorage
|
||||
flattened_data = {}
|
||||
for mode, mode_data in all_data.items():
|
||||
print(f"Processing mode: {mode}")
|
||||
for hash_value, cache_entry in mode_data.items():
|
||||
# Determine cache_type from cache entry or use default
|
||||
cache_type = cache_entry.get("cache_type", "extract")
|
||||
# Create flattened key: {mode}:{cache_type}:{hash}
|
||||
flattened_key = f"{mode}:{cache_type}:{hash_value}"
|
||||
flattened_data[flattened_key] = cache_entry
|
||||
print(f"\tConverting {mode}[{hash_value}] -> {flattened_key}")
|
||||
|
||||
# Upsert the flattened data
|
||||
await to_llm_response_cache.upsert(flattened_data)
|
||||
print("Mission accomplished!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(copy_from_json_to_postgres())
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
LightRAG meets Amazon Bedrock ⛰️
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.bedrock import bedrock_complete, bedrock_embed
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
logging.getLogger("aiobotocore").setLevel(logging.WARNING)
|
||||
|
||||
WORKING_DIR = "./dickens"
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=bedrock_complete,
|
||||
llm_model_name="Anthropic Claude 3 Haiku // Amazon Bedrock",
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=1024, max_token_size=8192, func=bedrock_embed
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
for mode in ["naive", "local", "global", "hybrid"]:
|
||||
print("\n+-" + "-" * len(mode) + "-+")
|
||||
print(f"| {mode.capitalize()} |")
|
||||
print("+-" + "-" * len(mode) + "-+\n")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode=mode)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,354 @@
|
||||
import asyncio
|
||||
import os
|
||||
import inspect
|
||||
import logging
|
||||
import logging.config
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
|
||||
|
||||
import requests
|
||||
import numpy as np
|
||||
from dotenv import load_dotenv
|
||||
|
||||
"""This code is a modified version of lightrag_openai_demo.py"""
|
||||
|
||||
# ideally, as always, env!
|
||||
load_dotenv(dotenv_path=".env", override=False)
|
||||
|
||||
|
||||
""" ----========= IMPORTANT CHANGE THIS! =========---- """
|
||||
cloudflare_api_key = "YOUR_API_KEY"
|
||||
account_id = "YOUR_ACCOUNT ID" # This is unique to your Cloudflare account
|
||||
|
||||
# Authomatically changes
|
||||
api_base_url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/"
|
||||
|
||||
|
||||
# choose an embedding model
|
||||
EMBEDDING_MODEL = "@cf/baai/bge-m3"
|
||||
# choose a generative model
|
||||
LLM_MODEL = "@cf/meta/llama-3.2-3b-instruct"
|
||||
|
||||
WORKING_DIR = "../dickens" # you can change output as desired
|
||||
|
||||
|
||||
# Cloudflare init
|
||||
class CloudflareWorker:
|
||||
def __init__(
|
||||
self,
|
||||
cloudflare_api_key: str,
|
||||
api_base_url: str,
|
||||
llm_model_name: str,
|
||||
embedding_model_name: str,
|
||||
max_tokens: int = 4080,
|
||||
max_response_tokens: int = 4080,
|
||||
):
|
||||
self.cloudflare_api_key = cloudflare_api_key
|
||||
self.api_base_url = api_base_url
|
||||
self.llm_model_name = llm_model_name
|
||||
self.embedding_model_name = embedding_model_name
|
||||
self.max_tokens = max_tokens
|
||||
self.max_response_tokens = max_response_tokens
|
||||
|
||||
async def _send_request(self, model_name: str, input_: dict, debug_log: str):
|
||||
headers = {"Authorization": f"Bearer {self.cloudflare_api_key}"}
|
||||
|
||||
print(f"""
|
||||
data sent to Cloudflare
|
||||
~~~~~~~~~~~
|
||||
{debug_log}
|
||||
""")
|
||||
|
||||
try:
|
||||
response_raw = requests.post(
|
||||
f"{self.api_base_url}{model_name}", headers=headers, json=input_
|
||||
).json()
|
||||
print(f"""
|
||||
Cloudflare worker responded with:
|
||||
~~~~~~~~~~~
|
||||
{str(response_raw)}
|
||||
""")
|
||||
result = response_raw.get("result", {})
|
||||
|
||||
if "data" in result: # Embedding case
|
||||
return np.array(result["data"])
|
||||
|
||||
if "response" in result: # LLM response
|
||||
return result["response"]
|
||||
|
||||
raise ValueError("Unexpected Cloudflare response format")
|
||||
|
||||
except Exception as e:
|
||||
print(f"""
|
||||
Cloudflare API returned:
|
||||
~~~~~~~~~
|
||||
Error: {e}
|
||||
""")
|
||||
input("Press Enter to continue...")
|
||||
return None
|
||||
|
||||
async def query(self, prompt, system_prompt: str = "", **kwargs) -> str:
|
||||
# since no caching is used and we don't want to mess with everything lightrag, pop the kwarg it is
|
||||
kwargs.pop("hashing_kv", None)
|
||||
|
||||
message = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
|
||||
input_ = {
|
||||
"messages": message,
|
||||
"max_tokens": self.max_tokens,
|
||||
"response_token_limit": self.max_response_tokens,
|
||||
}
|
||||
|
||||
return await self._send_request(
|
||||
self.llm_model_name,
|
||||
input_,
|
||||
debug_log=f"\n- model used {self.llm_model_name}\n- system prompt: {system_prompt}\n- query: {prompt}",
|
||||
)
|
||||
|
||||
async def embedding_chunk(self, texts: list[str]) -> np.ndarray:
|
||||
print(f"""
|
||||
TEXT inputted
|
||||
~~~~~
|
||||
{texts}
|
||||
""")
|
||||
|
||||
input_ = {
|
||||
"text": texts,
|
||||
"max_tokens": self.max_tokens,
|
||||
"response_token_limit": self.max_response_tokens,
|
||||
}
|
||||
|
||||
return await self._send_request(
|
||||
self.embedding_model_name,
|
||||
input_,
|
||||
debug_log=f"\n-llm model name {self.embedding_model_name}\n- texts: {texts}",
|
||||
)
|
||||
|
||||
|
||||
def configure_logging():
|
||||
"""Configure logging for the application"""
|
||||
|
||||
# Reset any existing handlers to ensure clean configuration
|
||||
for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
|
||||
logger_instance = logging.getLogger(logger_name)
|
||||
logger_instance.handlers = []
|
||||
logger_instance.filters = []
|
||||
|
||||
# Get log directory path from environment variable or use current directory
|
||||
log_dir = os.getenv("LOG_DIR", os.getcwd())
|
||||
log_file_path = os.path.abspath(
|
||||
os.path.join(log_dir, "lightrag_cloudflare_worker_demo.log")
|
||||
)
|
||||
|
||||
print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
|
||||
os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
|
||||
|
||||
# Get log file max size and backup count from environment variables
|
||||
log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760)) # Default 10MB
|
||||
log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5)) # Default 5 backups
|
||||
|
||||
logging.config.dictConfig(
|
||||
{
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"default": {
|
||||
"format": "%(levelname)s: %(message)s",
|
||||
},
|
||||
"detailed": {
|
||||
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"formatter": "default",
|
||||
"class": "logging.StreamHandler",
|
||||
"stream": "ext://sys.stderr",
|
||||
},
|
||||
"file": {
|
||||
"formatter": "detailed",
|
||||
"class": "logging.handlers.RotatingFileHandler",
|
||||
"filename": log_file_path,
|
||||
"maxBytes": log_max_bytes,
|
||||
"backupCount": log_backup_count,
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
},
|
||||
"loggers": {
|
||||
"lightrag": {
|
||||
"handlers": ["console", "file"],
|
||||
"level": "INFO",
|
||||
"propagate": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Set the logger level to INFO
|
||||
logger.setLevel(logging.INFO)
|
||||
# Enable verbose debug if needed
|
||||
set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
|
||||
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
cloudflare_worker = CloudflareWorker(
|
||||
cloudflare_api_key=cloudflare_api_key,
|
||||
api_base_url=api_base_url,
|
||||
embedding_model_name=EMBEDDING_MODEL,
|
||||
llm_model_name=LLM_MODEL,
|
||||
)
|
||||
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
max_parallel_insert=2,
|
||||
llm_model_func=cloudflare_worker.query,
|
||||
llm_model_name=os.getenv("LLM_MODEL", LLM_MODEL),
|
||||
summary_max_tokens=4080,
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
|
||||
max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "2048")),
|
||||
func=lambda texts: cloudflare_worker.embedding_chunk(
|
||||
texts,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
async def print_stream(stream):
|
||||
async for chunk in stream:
|
||||
print(chunk, end="", flush=True)
|
||||
|
||||
|
||||
async def main():
|
||||
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
|
||||
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")
|
||||
|
||||
# Locate the location of what is needed to be added to the knowledge
|
||||
# Can add several simultaneously by modifying code
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
await rag.ainsert(f.read())
|
||||
|
||||
# Perform naive search
|
||||
print("\n=====================")
|
||||
print("Query mode: naive")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="naive", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
# Perform local search
|
||||
print("\n=====================")
|
||||
print("Query mode: local")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="local", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
# Perform global search
|
||||
print("\n=====================")
|
||||
print("Query mode: global")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="global", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
# Perform hybrid search
|
||||
print("\n=====================")
|
||||
print("Query mode: hybrid")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="hybrid", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
""" FOR TESTING (if you want to test straight away, after building. Uncomment this part"""
|
||||
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("AI ASSISTANT READY!")
|
||||
print("Ask questions about (your uploaded) regulations")
|
||||
print("Type 'quit' to exit")
|
||||
print("=" * 60)
|
||||
|
||||
while True:
|
||||
question = input("\n🔥 Your question: ")
|
||||
|
||||
if question.lower() in ['quit', 'exit', 'bye']:
|
||||
break
|
||||
|
||||
print("\nThinking...")
|
||||
response = await rag.aquery(question, param=QueryParam(mode="hybrid"))
|
||||
print(f"\nAnswer: {response}")
|
||||
|
||||
"""
|
||||
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
finally:
|
||||
if rag:
|
||||
await rag.llm_response_cache.index_done_callback()
|
||||
await rag.finalize_storages()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Configure logging before running the main function
|
||||
configure_logging()
|
||||
asyncio.run(main())
|
||||
print("\nDone!")
|
||||
@@ -0,0 +1,235 @@
|
||||
import os
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
import logging.config
|
||||
from functools import partial
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.openai import openai_complete_if_cache
|
||||
from lightrag.llm.ollama import ollama_embed
|
||||
from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv(dotenv_path=".env", override=False)
|
||||
|
||||
WORKING_DIR = "./dickens"
|
||||
|
||||
|
||||
def configure_logging():
|
||||
"""Configure logging for the application"""
|
||||
|
||||
# Reset any existing handlers to ensure clean configuration
|
||||
for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
|
||||
logger_instance = logging.getLogger(logger_name)
|
||||
logger_instance.handlers = []
|
||||
logger_instance.filters = []
|
||||
|
||||
# Get log directory path from environment variable or use current directory
|
||||
log_dir = os.getenv("LOG_DIR", os.getcwd())
|
||||
log_file_path = os.path.abspath(
|
||||
os.path.join(log_dir, "lightrag_compatible_demo.log")
|
||||
)
|
||||
|
||||
print(f"\nLightRAG compatible demo log file: {log_file_path}\n")
|
||||
os.makedirs(os.path.dirname(log_file_path), exist_ok=True)
|
||||
|
||||
# Get log file max size and backup count from environment variables
|
||||
log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760)) # Default 10MB
|
||||
log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5)) # Default 5 backups
|
||||
|
||||
logging.config.dictConfig(
|
||||
{
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"default": {
|
||||
"format": "%(levelname)s: %(message)s",
|
||||
},
|
||||
"detailed": {
|
||||
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"formatter": "default",
|
||||
"class": "logging.StreamHandler",
|
||||
"stream": "ext://sys.stderr",
|
||||
},
|
||||
"file": {
|
||||
"formatter": "detailed",
|
||||
"class": "logging.handlers.RotatingFileHandler",
|
||||
"filename": log_file_path,
|
||||
"maxBytes": log_max_bytes,
|
||||
"backupCount": log_backup_count,
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
},
|
||||
"loggers": {
|
||||
"lightrag": {
|
||||
"handlers": ["console", "file"],
|
||||
"level": "INFO",
|
||||
"propagate": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Set the logger level to INFO
|
||||
logger.setLevel(logging.INFO)
|
||||
# Enable verbose debug if needed
|
||||
set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
|
||||
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
async def llm_model_func(
|
||||
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
|
||||
) -> str:
|
||||
return await openai_complete_if_cache(
|
||||
os.getenv("LLM_MODEL", "deepseek-chat"),
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
api_key=os.getenv("LLM_BINDING_API_KEY") or os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("LLM_BINDING_HOST", "https://api.deepseek.com"),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def print_stream(stream):
|
||||
async for chunk in stream:
|
||||
if chunk:
|
||||
print(chunk, end="", flush=True)
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
# Note: ollama_embed is decorated with @wrap_embedding_func_with_attrs,
|
||||
# which wraps it in an EmbeddingFunc. Using .func accesses the original
|
||||
# unwrapped function to avoid double wrapping when we create our own
|
||||
# EmbeddingFunc with custom configuration (embedding_dim, max_token_size and prefixes).
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=int(os.getenv("EMBEDDING_DIM", "1024")),
|
||||
max_token_size=int(os.getenv("MAX_EMBED_TOKENS", "8192")),
|
||||
supports_asymmetric=True,
|
||||
func=partial(
|
||||
ollama_embed.func, # Access the unwrapped function to avoid double EmbeddingFunc wrapping
|
||||
embed_model=os.getenv("EMBEDDING_MODEL", "FRIDA:latest"),
|
||||
host=os.getenv("EMBEDDING_BINDING_HOST", "http://localhost:11434"),
|
||||
query_prefix=os.getenv("EMBEDDING_QUERY_PREFIX", "search_query: "),
|
||||
document_prefix=os.getenv(
|
||||
"EMBEDDING_DOCUMENT_PREFIX", "search_document: "
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
async def main():
|
||||
rag = None
|
||||
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
|
||||
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:
|
||||
await rag.ainsert(f.read())
|
||||
|
||||
# Perform naive search
|
||||
print("\n=====================")
|
||||
print("Query mode: naive")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="naive", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
# Perform local search
|
||||
print("\n=====================")
|
||||
print("Query mode: local")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="local", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
# Perform global search
|
||||
print("\n=====================")
|
||||
print("Query mode: global")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="global", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
# Perform hybrid search
|
||||
print("\n=====================")
|
||||
print("Query mode: hybrid")
|
||||
print("=====================")
|
||||
resp = await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="hybrid", stream=True),
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
await print_stream(resp)
|
||||
else:
|
||||
print(resp)
|
||||
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
finally:
|
||||
if rag:
|
||||
await rag.finalize_storages()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Configure logging before running the main function
|
||||
configure_logging()
|
||||
asyncio.run(main())
|
||||
print("\nDone!")
|
||||
@@ -0,0 +1,79 @@
|
||||
import os
|
||||
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.hf import hf_model_complete, hf_embed
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
WORKING_DIR = "./dickens"
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=hf_model_complete,
|
||||
llm_model_name="meta-llama/Llama-3.1-8B-Instruct",
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=384,
|
||||
max_token_size=5000,
|
||||
func=lambda texts: hf_embed(
|
||||
texts,
|
||||
tokenizer=AutoTokenizer.from_pretrained(
|
||||
"sentence-transformers/all-MiniLM-L6-v2"
|
||||
),
|
||||
embed_model=AutoModel.from_pretrained(
|
||||
"sentence-transformers/all-MiniLM-L6-v2"
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Perform naive search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform local search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform global search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="global")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform hybrid search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,139 @@
|
||||
import os
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.llama_index_impl import (
|
||||
llama_index_complete_if_cache,
|
||||
llama_index_embed,
|
||||
)
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from llama_index.llms.openai import OpenAI
|
||||
from llama_index.embeddings.openai import OpenAIEmbedding
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
|
||||
# Configure working directory
|
||||
WORKING_DIR = "./index_default"
|
||||
print(f"WORKING_DIR: {WORKING_DIR}")
|
||||
|
||||
# Model configuration
|
||||
LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4")
|
||||
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}")
|
||||
|
||||
# OpenAI configuration
|
||||
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "your-api-key-here")
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
print(f"Creating working directory: {WORKING_DIR}")
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
# Initialize LLM function
|
||||
async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
|
||||
try:
|
||||
# Initialize OpenAI if not in kwargs
|
||||
if "llm_instance" not in kwargs:
|
||||
llm_instance = OpenAI(
|
||||
model=LLM_MODEL,
|
||||
api_key=OPENAI_API_KEY,
|
||||
temperature=0.7,
|
||||
)
|
||||
kwargs["llm_instance"] = llm_instance
|
||||
|
||||
response = await llama_index_complete_if_cache(
|
||||
kwargs["llm_instance"],
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
**kwargs,
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
print(f"LLM request failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Initialize embedding function
|
||||
async def embedding_func(texts):
|
||||
try:
|
||||
embed_model = OpenAIEmbedding(
|
||||
model=EMBEDDING_MODEL,
|
||||
api_key=OPENAI_API_KEY,
|
||||
)
|
||||
return await llama_index_embed(texts, embed_model=embed_model)
|
||||
except Exception as e:
|
||||
print(f"Embedding failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Get embedding dimension
|
||||
async def get_embedding_dim():
|
||||
test_text = ["This is a test sentence."]
|
||||
embedding = await embedding_func(test_text)
|
||||
embedding_dim = embedding.shape[1]
|
||||
print(f"embedding_dim={embedding_dim}")
|
||||
return embedding_dim
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
embedding_dimension = await get_embedding_dim()
|
||||
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=embedding_dimension,
|
||||
max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
|
||||
func=embedding_func,
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize RAG instance
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
# Insert example text
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Test different query modes
|
||||
print("\nNaive Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nLocal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nGlobal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="global")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nHybrid Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,141 @@
|
||||
import os
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.llama_index_impl import (
|
||||
llama_index_complete_if_cache,
|
||||
llama_index_embed,
|
||||
)
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from llama_index.llms.litellm import LiteLLM
|
||||
from llama_index.embeddings.litellm import LiteLLMEmbedding
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
|
||||
# Configure working directory
|
||||
WORKING_DIR = "./index_default"
|
||||
print(f"WORKING_DIR: {WORKING_DIR}")
|
||||
|
||||
# Model configuration
|
||||
LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-4")
|
||||
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}")
|
||||
|
||||
# LiteLLM configuration
|
||||
LITELLM_URL = os.environ.get("LITELLM_URL", "http://localhost:4000")
|
||||
print(f"LITELLM_URL: {LITELLM_URL}")
|
||||
LITELLM_KEY = os.environ.get("LITELLM_KEY", "sk-1234")
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
# Initialize LLM function
|
||||
async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
|
||||
try:
|
||||
# Initialize LiteLLM if not in kwargs
|
||||
if "llm_instance" not in kwargs:
|
||||
llm_instance = LiteLLM(
|
||||
model=f"openai/{LLM_MODEL}", # Format: "provider/model_name"
|
||||
api_base=LITELLM_URL,
|
||||
api_key=LITELLM_KEY,
|
||||
temperature=0.7,
|
||||
)
|
||||
kwargs["llm_instance"] = llm_instance
|
||||
|
||||
response = await llama_index_complete_if_cache(
|
||||
kwargs["llm_instance"],
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
print(f"LLM request failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Initialize embedding function
|
||||
async def embedding_func(texts):
|
||||
try:
|
||||
embed_model = LiteLLMEmbedding(
|
||||
model_name=f"openai/{EMBEDDING_MODEL}",
|
||||
api_base=LITELLM_URL,
|
||||
api_key=LITELLM_KEY,
|
||||
)
|
||||
return await llama_index_embed(texts, embed_model=embed_model)
|
||||
except Exception as e:
|
||||
print(f"Embedding failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Get embedding dimension
|
||||
async def get_embedding_dim():
|
||||
test_text = ["This is a test sentence."]
|
||||
embedding = await embedding_func(test_text)
|
||||
embedding_dim = embedding.shape[1]
|
||||
print(f"embedding_dim={embedding_dim}")
|
||||
return embedding_dim
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
embedding_dimension = await get_embedding_dim()
|
||||
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=embedding_dimension,
|
||||
max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
|
||||
func=embedding_func,
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize RAG instance
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
# Insert example text
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Test different query modes
|
||||
print("\nNaive Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nLocal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nGlobal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="global")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nHybrid Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,152 @@
|
||||
import os
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.llama_index_impl import (
|
||||
llama_index_complete_if_cache,
|
||||
llama_index_embed,
|
||||
)
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from llama_index.llms.litellm import LiteLLM
|
||||
from llama_index.embeddings.litellm import LiteLLMEmbedding
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
|
||||
# Configure working directory
|
||||
WORKING_DIR = "./index_default"
|
||||
print(f"WORKING_DIR: {WORKING_DIR}")
|
||||
|
||||
# Model configuration
|
||||
LLM_MODEL = os.environ.get("LLM_MODEL", "gemma-3-4b")
|
||||
print(f"LLM_MODEL: {LLM_MODEL}")
|
||||
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "arctic-embed")
|
||||
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}")
|
||||
|
||||
# LiteLLM configuration
|
||||
LITELLM_URL = os.environ.get("LITELLM_URL", "http://localhost:4000")
|
||||
print(f"LITELLM_URL: {LITELLM_URL}")
|
||||
LITELLM_KEY = os.environ.get("LITELLM_KEY", "sk-4JdvGFKqSA3S0k_5p0xufw")
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
# Initialize LLM function
|
||||
async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
|
||||
try:
|
||||
# Initialize LiteLLM if not in kwargs
|
||||
if "llm_instance" not in kwargs:
|
||||
llm_instance = LiteLLM(
|
||||
model=f"openai/{LLM_MODEL}", # Format: "provider/model_name"
|
||||
api_base=LITELLM_URL,
|
||||
api_key=LITELLM_KEY,
|
||||
temperature=0.7,
|
||||
)
|
||||
kwargs["llm_instance"] = llm_instance
|
||||
|
||||
chat_kwargs = {}
|
||||
chat_kwargs["litellm_params"] = {
|
||||
"metadata": {
|
||||
"opik": {
|
||||
"project_name": "lightrag_llamaindex_litellm_opik_demo",
|
||||
"tags": ["lightrag", "litellm"],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
response = await llama_index_complete_if_cache(
|
||||
kwargs["llm_instance"],
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
chat_kwargs=chat_kwargs,
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
print(f"LLM request failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Initialize embedding function
|
||||
async def embedding_func(texts):
|
||||
try:
|
||||
embed_model = LiteLLMEmbedding(
|
||||
model_name=f"openai/{EMBEDDING_MODEL}",
|
||||
api_base=LITELLM_URL,
|
||||
api_key=LITELLM_KEY,
|
||||
)
|
||||
return await llama_index_embed(texts, embed_model=embed_model)
|
||||
except Exception as e:
|
||||
print(f"Embedding failed: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
# Get embedding dimension
|
||||
async def get_embedding_dim():
|
||||
test_text = ["This is a test sentence."]
|
||||
embedding = await embedding_func(test_text)
|
||||
embedding_dim = embedding.shape[1]
|
||||
print(f"embedding_dim={embedding_dim}")
|
||||
return embedding_dim
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
embedding_dimension = await get_embedding_dim()
|
||||
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=embedding_dimension,
|
||||
max_token_size=EMBEDDING_MAX_TOKEN_SIZE,
|
||||
func=embedding_func,
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize RAG instance
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
# Insert example text
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Test different query modes
|
||||
print("\nNaive Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nLocal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nGlobal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="global")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nHybrid Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,107 @@
|
||||
import os
|
||||
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.lmdeploy import lmdeploy_model_if_cache
|
||||
from lightrag.llm.hf import hf_embed
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
WORKING_DIR = "./dickens"
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
|
||||
async def lmdeploy_model_complete(
|
||||
prompt=None,
|
||||
system_prompt=None,
|
||||
history_messages=[],
|
||||
keyword_extraction=False,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
|
||||
return await lmdeploy_model_if_cache(
|
||||
model_name,
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
## please specify chat_template if your local path does not follow original HF file name,
|
||||
## or model_name is a pytorch model on huggingface.co,
|
||||
## you can refer to https://github.com/InternLM/lmdeploy/blob/main/lmdeploy/model.py
|
||||
## for a list of chat_template available in lmdeploy.
|
||||
chat_template="llama3",
|
||||
# model_format ='awq', # if you are using awq quantization model.
|
||||
# quant_policy=8, # if you want to use online kv cache, 4=kv int4, 8=kv int8.
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=lmdeploy_model_complete,
|
||||
llm_model_name="meta-llama/Llama-3.1-8B-Instruct", # please use definite path for local model
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=384,
|
||||
max_token_size=5000,
|
||||
func=lambda texts: hf_embed(
|
||||
texts,
|
||||
tokenizer=AutoTokenizer.from_pretrained(
|
||||
"sentence-transformers/all-MiniLM-L6-v2"
|
||||
),
|
||||
embed_model=AutoModel.from_pretrained(
|
||||
"sentence-transformers/all-MiniLM-L6-v2"
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize RAG instance
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
# Insert example text
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Test different query modes
|
||||
print("\nNaive Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nLocal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nGlobal Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="global")
|
||||
)
|
||||
)
|
||||
|
||||
print("\nHybrid Search:")
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,168 @@
|
||||
import os
|
||||
import asyncio
|
||||
import nest_asyncio
|
||||
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm import (
|
||||
openai_complete_if_cache,
|
||||
nvidia_openai_embed,
|
||||
)
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
import numpy as np
|
||||
|
||||
# for custom llm_model_func
|
||||
from lightrag.utils import locate_json_string_body_from_string
|
||||
|
||||
nest_asyncio.apply()
|
||||
|
||||
WORKING_DIR = "./dickens"
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
# some method to use your API key (choose one)
|
||||
# NVIDIA_OPENAI_API_KEY = os.getenv("NVIDIA_OPENAI_API_KEY")
|
||||
NVIDIA_OPENAI_API_KEY = "nvapi-xxxx" # your api key
|
||||
|
||||
# using pre-defined function for nvidia LLM API. OpenAI compatible
|
||||
# llm_model_func = nvidia_openai_complete
|
||||
|
||||
|
||||
# If you trying to make custom llm_model_func to use llm model on NVIDIA API like other example:
|
||||
async def llm_model_func(
|
||||
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
|
||||
) -> str:
|
||||
result = await openai_complete_if_cache(
|
||||
"nvidia/llama-3.1-nemotron-70b-instruct",
|
||||
prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages,
|
||||
api_key=NVIDIA_OPENAI_API_KEY,
|
||||
base_url="https://integrate.api.nvidia.com/v1",
|
||||
**kwargs,
|
||||
)
|
||||
if keyword_extraction:
|
||||
return locate_json_string_body_from_string(result)
|
||||
return result
|
||||
|
||||
|
||||
# custom embedding
|
||||
nvidia_embed_model = "nvidia/nv-embedqa-e5-v5"
|
||||
|
||||
|
||||
async def indexing_embedding_func(texts: list[str]) -> np.ndarray:
|
||||
return await nvidia_openai_embed(
|
||||
texts,
|
||||
model=nvidia_embed_model, # maximum 512 token
|
||||
# model="nvidia/llama-3.2-nv-embedqa-1b-v1",
|
||||
api_key=NVIDIA_OPENAI_API_KEY,
|
||||
base_url="https://integrate.api.nvidia.com/v1",
|
||||
input_type="passage",
|
||||
trunc="END", # handling on server side if input token is longer than maximum token
|
||||
encode="float",
|
||||
)
|
||||
|
||||
|
||||
async def query_embedding_func(texts: list[str]) -> np.ndarray:
|
||||
return await nvidia_openai_embed(
|
||||
texts,
|
||||
model=nvidia_embed_model, # maximum 512 token
|
||||
# model="nvidia/llama-3.2-nv-embedqa-1b-v1",
|
||||
api_key=NVIDIA_OPENAI_API_KEY,
|
||||
base_url="https://integrate.api.nvidia.com/v1",
|
||||
input_type="query",
|
||||
trunc="END", # handling on server side if input token is longer than maximum token
|
||||
encode="float",
|
||||
)
|
||||
|
||||
|
||||
# dimension are same
|
||||
async def get_embedding_dim():
|
||||
test_text = ["This is a test sentence."]
|
||||
embedding = await indexing_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 indexing_embedding_func(["How are you?"])
|
||||
print("embedding_func: ", result)
|
||||
|
||||
|
||||
# asyncio.run(test_funcs())
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
embedding_dimension = await get_embedding_dim()
|
||||
print(f"Detected embedding dimension: {embedding_dimension}")
|
||||
|
||||
# lightRAG class during indexing
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
# llm_model_name="meta/llama3-70b-instruct", #un comment if
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=embedding_dimension,
|
||||
max_token_size=512, # maximum token size, somehow it's still exceed maximum number of token
|
||||
# so truncate (trunc) parameter on embedding_func will handle it and try to examine the tokenizer used in LightRAG
|
||||
# so you can adjust to be able to fit the NVIDIA model (future work)
|
||||
func=indexing_embedding_func,
|
||||
),
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
async def main():
|
||||
try:
|
||||
# Initialize RAG instance
|
||||
rag = await initialize_rag()
|
||||
|
||||
# reading file
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
await rag.ainsert(f.read())
|
||||
|
||||
# Perform naive search
|
||||
print("==============Naive===============")
|
||||
print(
|
||||
await rag.aquery(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform local search
|
||||
print("==============local===============")
|
||||
print(
|
||||
await rag.aquery(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform global search
|
||||
print("==============global===============")
|
||||
print(
|
||||
await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="global"),
|
||||
)
|
||||
)
|
||||
|
||||
# Perform hybrid search
|
||||
print("==============hybrid===============")
|
||||
print(
|
||||
await rag.aquery(
|
||||
"What are the top themes in this story?",
|
||||
param=QueryParam(mode="hybrid"),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"An error occurred: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,109 @@
|
||||
import os
|
||||
import asyncio
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.ollama import ollama_embed, openai_complete_if_cache
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
|
||||
# WorkingDir
|
||||
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
WORKING_DIR = os.path.join(ROOT_DIR, "myKG")
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
print(f"WorkingDir: {WORKING_DIR}")
|
||||
|
||||
# redis
|
||||
os.environ["REDIS_URI"] = "redis://localhost:6379"
|
||||
|
||||
# neo4j
|
||||
BATCH_SIZE_NODES = 500
|
||||
BATCH_SIZE_EDGES = 100
|
||||
os.environ["NEO4J_URI"] = "neo4j://localhost:7687"
|
||||
os.environ["NEO4J_USERNAME"] = "neo4j"
|
||||
os.environ["NEO4J_PASSWORD"] = "12345678"
|
||||
|
||||
# milvus
|
||||
os.environ["MILVUS_URI"] = "http://localhost:19530"
|
||||
os.environ["MILVUS_USER"] = "root"
|
||||
os.environ["MILVUS_PASSWORD"] = "Milvus"
|
||||
os.environ["MILVUS_DB_NAME"] = "lightrag"
|
||||
|
||||
|
||||
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="",
|
||||
base_url="",
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
embedding_func = EmbeddingFunc(
|
||||
embedding_dim=768,
|
||||
max_token_size=512,
|
||||
func=lambda texts: ollama_embed(
|
||||
texts, embed_model="shaw/dmeta-embedding-zh", host="http://117.50.173.35:11434"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def initialize_rag():
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=llm_model_func,
|
||||
summary_max_tokens=10000,
|
||||
embedding_func=embedding_func,
|
||||
chunk_token_size=512,
|
||||
chunk_overlap_token_size=256,
|
||||
kv_storage="RedisKVStorage",
|
||||
graph_storage="Neo4JStorage",
|
||||
vector_storage="MilvusVectorDBStorage",
|
||||
doc_status_storage="RedisKVStorage",
|
||||
)
|
||||
|
||||
await rag.initialize_storages() # Auto-initializes pipeline_status
|
||||
return rag
|
||||
|
||||
|
||||
def main():
|
||||
# Initialize RAG instance
|
||||
rag = asyncio.run(initialize_rag())
|
||||
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Perform naive search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="naive")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform local search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="local")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform global search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="global")
|
||||
)
|
||||
)
|
||||
|
||||
# Perform hybrid search
|
||||
print(
|
||||
rag.query(
|
||||
"What are the top themes in this story?", param=QueryParam(mode="hybrid")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user