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

fix pre commit

Browse Source
This commit is contained in:
jin
2024-11-12 13:32:40 +08:00
parent cc45d802a6
commit 41599897fb
7 changed files with 345 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
examples/lightrag_api_oracle_demo..py
View File

@@ -1,10 +1,10 @@
from fastapi import FastAPI, HTTPException, File, UploadFile
from contextlib import asynccontextmanager
from pydantic import BaseModel
from typing import Optional
import sys, os
import sys
import os
from pathlib import Path
import asyncio
@@ -13,7 +13,6 @@ from lightrag import LightRAG, QueryParam
from lightrag.llm import openai_complete_if_cache, openai_embedding
from lightrag.utils import EmbeddingFunc
import numpy as np
from datetime import datetime
from lightrag.kg.oracle_impl import OracleDB
@@ -24,8 +23,6 @@ script_directory = Path(__file__).resolve().parent.parent
sys.path.append(os.path.abspath(script_directory))
# Apply nest_asyncio to solve event loop issues
nest_asyncio.apply()
@@ -51,6 +48,7 @@ print(f"EMBEDDING_MAX_TOKEN_SIZE: {EMBEDDING_MAX_TOKEN_SIZE}")
if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR)
async def llm_model_func(
prompt, system_prompt=None, history_messages=[], **kwargs
) -> str:
@@ -80,8 +78,8 @@ async def get_embedding_dim():
embedding_dim = embedding.shape[1]
return embedding_dim
async def init():
async def init():
# Detect embedding dimension
embedding_dimension = await get_embedding_dim()
print(f"Detected embedding dimension: {embedding_dimension}")
@@ -91,15 +89,15 @@ async def init():
# We storage data in unified tables, so we need to set a `workspace` parameter to specify which docs we want to store and query
# Below is an example of how to connect to Oracle Autonomous Database on Oracle Cloud
oracle_db = OracleDB(config={
oracle_db = OracleDB(
config={
"user": "",
"password": "",
"dsn": "",
"config_dir": "",
"wallet_location": "",
"wallet_password": "",
"workspace":""
"workspace": "",
} # specify which docs you want to store and query
)
@@ -119,7 +117,7 @@ async def init():
),
graph_storage="OracleGraphStorage",
kv_storage="OracleKVStorage",
vector_storage="OracleVectorDBStorage"
vector_storage="OracleVectorDBStorage",
)
# Setthe KV/vector/graph storage's `db` property, so all operation will use same connection pool
@@ -129,6 +127,7 @@ async def init():
return rag
# Data models
@@ -152,6 +151,7 @@ class Response(BaseModel):
rag = None # 定义为全局对象
@asynccontextmanager
async def lifespan(app: FastAPI):
global rag
@@ -160,7 +160,10 @@ async def lifespan(app: FastAPI):
yield
app = FastAPI(title="LightRAG API", description="API for RAG operations",lifespan=lifespan)
app = FastAPI(
title="LightRAG API", description="API for RAG operations", lifespan=lifespan
)
@app.post("/query", response_model=Response)
async def query_endpoint(request: QueryRequest):

21
examples/lightrag_oracle_demo.py
View File

@@ -1,11 +1,11 @@
import sys, os
import sys
import os
from pathlib import Path
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm import openai_complete_if_cache, openai_embedding
from lightrag.utils import EmbeddingFunc
import numpy as np
from datetime import datetime
from lightrag.kg.oracle_impl import OracleDB
print(os.getcwd())
@@ -25,6 +25,7 @@ EMBEDMODEL = "cohere.embed-multilingual-v3.0"
if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR)
async def llm_model_func(
prompt, system_prompt=None, history_messages=[], **kwargs
) -> str:
@@ -66,22 +67,21 @@ async def main():
# More docs here https://python-oracledb.readthedocs.io/en/latest/user_guide/connection_handling.html
# We storage data in unified tables, so we need to set a `workspace` parameter to specify which docs we want to store and query
# Below is an example of how to connect to Oracle Autonomous Database on Oracle Cloud
oracle_db = OracleDB(config={
oracle_db = OracleDB(
config={
"user": "username",
"password": "xxxxxxxxx",
"dsn": "xxxxxxx_medium",
"config_dir": "dir/path/to/oracle/config",
"wallet_location": "dir/path/to/oracle/wallet",
"wallet_password": "xxxxxxxxx",
"workspace":"company" # specify which docs you want to store and query
"workspace": "company", # specify which docs you want to store and query
}
)
# Check if Oracle DB tables exist, if not, tables will be created
await oracle_db.check_tables()
# Initialize LightRAG
# We use Oracle DB as the KV/vector/graph storage
rag = LightRAG(
@@ -96,7 +96,7 @@ async def main():
),
graph_storage="OracleGraphStorage",
kv_storage="OracleKVStorage",
vector_storage="OracleVectorDBStorage"
vector_storage="OracleVectorDBStorage",
)
# Setthe KV/vector/graph storage's `db` property, so all operation will use same connection pool
@@ -112,7 +112,12 @@ async def main():
modes = ["naive", "local", "global", "hybrid"]
for mode in modes:
print("=" * 20, mode, "=" * 20)
print(await rag.aquery("What are the top themes in this story?", param=QueryParam(mode=mode)))
print(
await rag.aquery(
"What are the top themes in this story?",
param=QueryParam(mode=mode),
)
)
print("-" * 100, "\n")
except Exception as e:

2
lightrag/base.py
View File

@@ -60,6 +60,7 @@ class BaseVectorStorage(StorageNameSpace):
@dataclass
class BaseKVStorage(Generic[T], StorageNameSpace):
embedding_func: EmbeddingFunc
async def all_keys(self) -> list[str]:
raise NotImplementedError
@@ -85,6 +86,7 @@ class BaseKVStorage(Generic[T], StorageNameSpace):
@dataclass
class BaseGraphStorage(StorageNameSpace):
embedding_func: EmbeddingFunc = None
async def has_node(self, node_id: str) -> bool:
raise NotImplementedError

229
lightrag/kg/oracle_impl.py
View File

@@ -1,9 +1,9 @@
import asyncio
# import html
# import os
from dataclasses import dataclass
from typing import Any, Union, cast
import networkx as nx
from typing import Union
import numpy as np
import array
@@ -16,6 +16,7 @@ from ..base import (
import oracledb
class OracleDB:
def __init__(self, config, **kwargs):
self.host = config.get("host", None)
@@ -45,7 +46,7 @@ class OracleDB:
wallet_password=self.wallet_password,
min=1,
max=self.max,
increment = self.increment
increment=self.increment,
)
logger.info(f"Connected to Oracle database at {self.dsn}")
except Exception as e:
@@ -95,7 +96,9 @@ class OracleDB:
for k, v in TABLES.items():
try:
if k.lower() == "lightrag_graph":
await self.query("SELECT id FROM GRAPH_TABLE (lightrag_graph MATCH (a) COLUMNS (a.id)) fetch first row only")
await self.query(
"SELECT id FROM GRAPH_TABLE (lightrag_graph MATCH (a) COLUMNS (a.id)) fetch first row only"
)
else:
await self.query("SELECT 1 FROM {k}".format(k=k))
except Exception as e:
@@ -109,8 +112,7 @@ class OracleDB:
logger.error(f"Failed to create table {k} in Oracle database")
logger.error(f"Oracle database error: {e}")
logger.info(f"Finished check all tables in Oracle database")
logger.info("Finished check all tables in Oracle database")
async def query(self, sql: str, multirows: bool = False) -> Union[dict, None]:
async with self.pool.acquire() as connection:
@@ -158,9 +160,9 @@ class OracleDB:
print(data)
raise
@dataclass
class OracleKVStorage(BaseKVStorage):
# should pass db object to self.db
def __post_init__(self):
self._data = {}
@@ -170,7 +172,9 @@ class OracleKVStorage(BaseKVStorage):
async def get_by_id(self, id: str) -> Union[dict, None]:
"""根据 id 获取 doc_full 数据."""
SQL = SQL_TEMPLATES["get_by_id_"+self.namespace].format(workspace=self.db.workspace,id=id)
SQL = SQL_TEMPLATES["get_by_id_" + self.namespace].format(
workspace=self.db.workspace, id=id
)
# print("get_by_id:"+SQL)
res = await self.db.query(SQL)
if res:
@@ -183,8 +187,9 @@ class OracleKVStorage(BaseKVStorage):
# Query by id
async def get_by_ids(self, ids: list[str], fields=None) -> Union[list[dict], None]:
"""根据 id 获取 doc_chunks 数据"""
SQL = SQL_TEMPLATES["get_by_ids_"+self.namespace].format(workspace=self.db.workspace,
ids=",".join([f"'{id}'" for id in ids]))
SQL = SQL_TEMPLATES["get_by_ids_" + self.namespace].format(
workspace=self.db.workspace, ids=",".join([f"'{id}'" for id in ids])
)
# print("get_by_ids:"+SQL)
res = await self.db.query(SQL, multirows=True)
if res:
@@ -196,9 +201,11 @@ class OracleKVStorage(BaseKVStorage):
async def filter_keys(self, keys: list[str]) -> set[str]:
"""过滤掉重复内容"""
SQL = SQL_TEMPLATES["filter_keys"].format(table_name=N_T[self.namespace],
SQL = SQL_TEMPLATES["filter_keys"].format(
table_name=N_T[self.namespace],
workspace=self.db.workspace,
ids=",".join([f"'{k}'" for k in keys]))
ids=",".join([f"'{k}'" for k in keys]),
)
res = await self.db.query(SQL, multirows=True)
data = None
if res:
@@ -209,7 +216,6 @@ class OracleKVStorage(BaseKVStorage):
data = set([s for s in keys if s not in exist_keys])
return data
################ INSERT METHODS ################
async def upsert(self, data: dict[str, dict]):
left_data = {k: v for k, v in data.items() if k not in self._data}
@@ -237,12 +243,17 @@ class OracleKVStorage(BaseKVStorage):
d["__vector__"] = embeddings[i]
# print(list_data)
for item in list_data:
merge_sql = SQL_TEMPLATES["merge_chunk"].format(
check_id=item["__id__"]
)
merge_sql = SQL_TEMPLATES["merge_chunk"].format(check_id=item["__id__"])
values = [item["__id__"], item["content"], self.db.workspace, item["tokens"],
item["chunk_order_index"], item["full_doc_id"], item["__vector__"]]
values = [
item["__id__"],
item["content"],
self.db.workspace,
item["tokens"],
item["chunk_order_index"],
item["full_doc_id"],
item["__vector__"],
]
# print(merge_sql)
await self.db.execute(merge_sql, values)
@@ -257,13 +268,11 @@ class OracleKVStorage(BaseKVStorage):
await self.db.execute(merge_sql, values)
return left_data
async def index_done_callback(self):
if self.namespace in ["full_docs", "text_chunks"]:
logger.info("full doc and chunk data had been saved into oracle db!")
@dataclass
class OracleVectorDBStorage(BaseVectorStorage):
cosine_better_than_threshold: float = 0.2
@@ -278,7 +287,6 @@ class OracleVectorDBStorage(BaseVectorStorage):
async def index_done_callback(self):
pass
#################### query method ###############
async def query(self, query: str, top_k=5) -> Union[dict, list[dict]]:
"""从向量数据库中查询数据"""
@@ -287,7 +295,7 @@ class OracleVectorDBStorage(BaseVectorStorage):
# 转换精度
dtype = str(embedding.dtype).upper()
dimension = embedding.shape[0]
embedding_string = ', '.join(map(str, embedding.tolist()))
embedding_string = ", ".join(map(str, embedding.tolist()))
SQL = SQL_TEMPLATES[self.namespace].format(
embedding_string=embedding_string,
@@ -311,7 +319,6 @@ class OracleGraphStorage(BaseGraphStorage):
"""从graphml文件加载图"""
self._max_batch_size = self.global_config["embedding_batch_num"]
#################### insert method ################
async def upsert_node(self, node_id: str, node_data: dict[str, str]):
@@ -336,7 +343,18 @@ class OracleGraphStorage(BaseGraphStorage):
workspace=self.db.workspace, name=entity_name, source_chunk_id=source_id
)
# print(merge_sql)
await self.db.execute(merge_sql, [self.db.workspace,entity_name,entity_type,description,source_id,content,content_vector])
await self.db.execute(
merge_sql,
[
self.db.workspace,
entity_name,
entity_type,
description,
source_id,
content,
content_vector,
],
)
# self._graph.add_node(node_id, **node_data)
async def upsert_edge(
@@ -362,10 +380,26 @@ class OracleGraphStorage(BaseGraphStorage):
embeddings = np.concatenate(embeddings_list)
content_vector = embeddings[0]
merge_sql = SQL_TEMPLATES["merge_edge"].format(
workspace=self.db.workspace,source_name=source_name, target_name=target_name, source_chunk_id=source_chunk_id
workspace=self.db.workspace,
source_name=source_name,
target_name=target_name,
source_chunk_id=source_chunk_id,
)
# print(merge_sql)
await self.db.execute(merge_sql, [self.db.workspace,source_name,target_name,weight,keywords,description,source_chunk_id,content,content_vector])
await self.db.execute(
merge_sql,
[
self.db.workspace,
source_name,
target_name,
weight,
keywords,
description,
source_chunk_id,
content,
content_vector,
],
)
# self._graph.add_edge(source_node_id, target_node_id, **edge_data)
async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
@@ -386,15 +420,18 @@ class OracleGraphStorage(BaseGraphStorage):
nodes_ids = [self._graph.nodes[node_id]["id"] for node_id in nodes]
return embeddings, nodes_ids
async def index_done_callback(self):
"""写入graphhml图文件"""
logger.info("Node and edge data had been saved into oracle db already, so nothing to do here!")
logger.info(
"Node and edge data had been saved into oracle db already, so nothing to do here!"
)
#################### query method #################
async def has_node(self, node_id: str) -> bool:
"""根据节点id检查节点是否存在"""
SQL = SQL_TEMPLATES["has_node"].format(workspace=self.db.workspace, node_id=node_id)
SQL = SQL_TEMPLATES["has_node"].format(
workspace=self.db.workspace, node_id=node_id
)
# print(SQL)
# print(self.db.workspace, node_id)
res = await self.db.query(SQL)
@@ -407,9 +444,11 @@ class OracleGraphStorage(BaseGraphStorage):
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
"""根据源和目标节点id检查边是否存在"""
SQL = SQL_TEMPLATES["has_edge"].format(workspace=self.db.workspace,
SQL = SQL_TEMPLATES["has_edge"].format(
workspace=self.db.workspace,
source_node_id=source_node_id,
target_node_id=target_node_id)
target_node_id=target_node_id,
)
# print(SQL)
res = await self.db.query(SQL)
if res:
@@ -421,7 +460,9 @@ class OracleGraphStorage(BaseGraphStorage):
async def node_degree(self, node_id: str) -> int:
"""根据节点id获取节点的度"""
SQL = SQL_TEMPLATES["node_degree"].format(workspace=self.db.workspace, node_id=node_id)
SQL = SQL_TEMPLATES["node_degree"].format(
workspace=self.db.workspace, node_id=node_id
)
# print(SQL)
res = await self.db.query(SQL)
if res:
@@ -431,17 +472,17 @@ class OracleGraphStorage(BaseGraphStorage):
# print("Edge not exist!")
return 0
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
"""根据源和目标节点id获取边的度"""
degree = await self.node_degree(src_id) + await self.node_degree(tgt_id)
# print("Edge degree",degree)
return degree
async def get_node(self, node_id: str) -> Union[dict, None]:
"""根据节点id获取节点数据"""
SQL = SQL_TEMPLATES["get_node"].format(workspace=self.db.workspace, node_id=node_id)
SQL = SQL_TEMPLATES["get_node"].format(
workspace=self.db.workspace, node_id=node_id
)
# print(self.db.workspace, node_id)
# print(SQL)
res = await self.db.query(SQL)
@@ -456,9 +497,11 @@ class OracleGraphStorage(BaseGraphStorage):
self, source_node_id: str, target_node_id: str
) -> Union[dict, None]:
"""根据源和目标节点id获取边"""
SQL = SQL_TEMPLATES["get_edge"].format(workspace=self.db.workspace,
SQL = SQL_TEMPLATES["get_edge"].format(
workspace=self.db.workspace,
source_node_id=source_node_id,
target_node_id=target_node_id)
target_node_id=target_node_id,
)
res = await self.db.query(SQL)
if res:
# print("Get edge!",self.db.workspace, source_node_id, target_node_id,res[0])
@@ -470,8 +513,9 @@ class OracleGraphStorage(BaseGraphStorage):
async def get_node_edges(self, source_node_id: str):
"""根据节点id获取节点的所有边"""
if await self.has_node(source_node_id):
SQL = SQL_TEMPLATES["get_node_edges"].format(workspace=self.db.workspace,
source_node_id=source_node_id)
SQL = SQL_TEMPLATES["get_node_edges"].format(
workspace=self.db.workspace, source_node_id=source_node_id
)
res = await self.db.query(sql=SQL, multirows=True)
if res:
data = [(i["source_name"], i["target_name"]) for i in res]
@@ -487,12 +531,12 @@ N_T = {
"text_chunks": "LIGHTRAG_DOC_CHUNKS",
"chunks": "LIGHTRAG_DOC_CHUNKS",
"entities": "LIGHTRAG_GRAPH_NODES",
"relationships": "LIGHTRAG_GRAPH_EDGES"
"relationships": "LIGHTRAG_GRAPH_EDGES",
}
TABLES = {
"LIGHTRAG_DOC_FULL":
{"ddl":"""CREATE TABLE LIGHTRAG_DOC_FULL (
"LIGHTRAG_DOC_FULL": {
"ddl": """CREATE TABLE LIGHTRAG_DOC_FULL (
id varchar(256)PRIMARY KEY,
workspace varchar(1024),
doc_name varchar(1024),
@@ -500,10 +544,10 @@ TABLES = {
meta JSON,
createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updatetime TIMESTAMP DEFAULT NULL
)"""},
"LIGHTRAG_DOC_CHUNKS":
{"ddl":"""CREATE TABLE LIGHTRAG_DOC_CHUNKS (
)"""
},
"LIGHTRAG_DOC_CHUNKS": {
"ddl": """CREATE TABLE LIGHTRAG_DOC_CHUNKS (
id varchar(256) PRIMARY KEY,
workspace varchar(1024),
full_doc_id varchar(256),
@@ -513,10 +557,10 @@ TABLES = {
content_vector VECTOR,
createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updatetime TIMESTAMP DEFAULT NULL
)"""},
"LIGHTRAG_GRAPH_NODES":
{"ddl":"""CREATE TABLE LIGHTRAG_GRAPH_NODES (
)"""
},
"LIGHTRAG_GRAPH_NODES": {
"ddl": """CREATE TABLE LIGHTRAG_GRAPH_NODES (
id NUMBER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
workspace varchar(1024),
name varchar(2048),
@@ -527,9 +571,10 @@ TABLES = {
content_vector VECTOR,
createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updatetime TIMESTAMP DEFAULT NULL
)"""},
"LIGHTRAG_GRAPH_EDGES":
{"ddl":"""CREATE TABLE LIGHTRAG_GRAPH_EDGES (
)"""
},
"LIGHTRAG_GRAPH_EDGES": {
"ddl": """CREATE TABLE LIGHTRAG_GRAPH_EDGES (
id NUMBER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
workspace varchar(1024),
source_name varchar(2048),
@@ -542,19 +587,20 @@ TABLES = {
content_vector VECTOR,
createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updatetime TIMESTAMP DEFAULT NULL
)"""},
"LIGHTRAG_LLM_CACHE":
{"ddl":"""CREATE TABLE LIGHTRAG_LLM_CACHE (
)"""
},
"LIGHTRAG_LLM_CACHE": {
"ddl": """CREATE TABLE LIGHTRAG_LLM_CACHE (
id varchar(256) PRIMARY KEY,
send clob,
return clob,
model varchar(1024),
createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updatetime TIMESTAMP DEFAULT NULL
)"""},
"LIGHTRAG_GRAPH":
{"ddl":"""CREATE OR REPLACE PROPERTY GRAPH lightrag_graph
)"""
},
"LIGHTRAG_GRAPH": {
"ddl": """CREATE OR REPLACE PROPERTY GRAPH lightrag_graph
VERTEX TABLES (
lightrag_graph_nodes KEY (id)
LABEL entity
@@ -566,92 +612,66 @@ TABLES = {
DESTINATION KEY (target_name) REFERENCES lightrag_graph_nodes(name)
LABEL has_relation
PROPERTIES (id,workspace,source_name,target_name) -- ,weight, keywords,description,source_chunk_id)
) OPTIONS(ALLOW MIXED PROPERTY TYPES)"""},
) OPTIONS(ALLOW MIXED PROPERTY TYPES)"""
},
}
SQL_TEMPLATES = {
# SQL for KVStorage
"get_by_id_full_docs":
"select ID,NVL(content,'') as content from LIGHTRAG_DOC_FULL where workspace='{workspace}' and ID='{id}'",
"get_by_id_text_chunks":
"select ID,TOKENS,NVL(content,'') as content,CHUNK_ORDER_INDEX,FULL_DOC_ID from LIGHTRAG_DOC_CHUNKS where workspace='{workspace}' and ID='{id}'",
"get_by_ids_full_docs":
"select ID,NVL(content,'') as content from LIGHTRAG_DOC_FULL where workspace='{workspace}' and ID in ({ids})",
"get_by_ids_text_chunks":
"select ID,TOKENS,NVL(content,'') as content,CHUNK_ORDER_INDEX,FULL_DOC_ID from LIGHTRAG_DOC_CHUNKS where workspace='{workspace}' and ID in ({ids})",
"filter_keys":
"select id from {table_name} where workspace='{workspace}' and id in ({ids})",
"get_by_id_full_docs": "select ID,NVL(content,'') as content from LIGHTRAG_DOC_FULL where workspace='{workspace}' and ID='{id}'",
"get_by_id_text_chunks": "select ID,TOKENS,NVL(content,'') as content,CHUNK_ORDER_INDEX,FULL_DOC_ID from LIGHTRAG_DOC_CHUNKS where workspace='{workspace}' and ID='{id}'",
"get_by_ids_full_docs": "select ID,NVL(content,'') as content from LIGHTRAG_DOC_FULL where workspace='{workspace}' and ID in ({ids})",
"get_by_ids_text_chunks": "select ID,TOKENS,NVL(content,'') as content,CHUNK_ORDER_INDEX,FULL_DOC_ID from LIGHTRAG_DOC_CHUNKS where workspace='{workspace}' and ID in ({ids})",
"filter_keys": "select id from {table_name} where workspace='{workspace}' and id in ({ids})",
"merge_doc_full": """ MERGE INTO LIGHTRAG_DOC_FULL a
USING DUAL
ON (a.id = '{check_id}')
WHEN NOT MATCHED THEN
INSERT(id,content,workspace) values(:1,:2,:3)
""",
"merge_chunk": """MERGE INTO LIGHTRAG_DOC_CHUNKS a
USING DUAL
ON (a.id = '{check_id}')
WHEN NOT MATCHED THEN
INSERT(id,content,workspace,tokens,chunk_order_index,full_doc_id,content_vector)
values (:1,:2,:3,:4,:5,:6,:7) """,
# SQL for VectorStorage
"entities":
"""SELECT name as entity_name FROM
"entities": """SELECT name as entity_name FROM
(SELECT id,name,VECTOR_DISTANCE(content_vector,vector('[{embedding_string}]',{dimension},{dtype}),COSINE) as distance
FROM LIGHTRAG_GRAPH_NODES WHERE workspace='{workspace}')
WHERE distance>{better_than_threshold} ORDER BY distance ASC FETCH FIRST {top_k} ROWS ONLY""",
"relationships":
"""SELECT source_name as src_id, target_name as tgt_id FROM
"relationships": """SELECT source_name as src_id, target_name as tgt_id FROM
(SELECT id,source_name,target_name,VECTOR_DISTANCE(content_vector,vector('[{embedding_string}]',{dimension},{dtype}),COSINE) as distance
FROM LIGHTRAG_GRAPH_EDGES WHERE workspace='{workspace}')
WHERE distance>{better_than_threshold} ORDER BY distance ASC FETCH FIRST {top_k} ROWS ONLY""",
"chunks":
"""SELECT id FROM
"chunks": """SELECT id FROM
(SELECT id,VECTOR_DISTANCE(content_vector,vector('[{embedding_string}]',{dimension},{dtype}),COSINE) as distance
FROM LIGHTRAG_DOC_CHUNKS WHERE workspace='{workspace}')
WHERE distance>{better_than_threshold} ORDER BY distance ASC FETCH FIRST {top_k} ROWS ONLY""",
# SQL for GraphStorage
"has_node":
"""SELECT * FROM GRAPH_TABLE (lightrag_graph
"has_node": """SELECT * FROM GRAPH_TABLE (lightrag_graph
MATCH (a)
WHERE a.workspace='{workspace}' AND a.name='{node_id}'
COLUMNS (a.name))""",
"has_edge":
"""SELECT * FROM GRAPH_TABLE (lightrag_graph
"has_edge": """SELECT * FROM GRAPH_TABLE (lightrag_graph
MATCH (a) -[e]-> (b)
WHERE e.workspace='{workspace}' and a.workspace='{workspace}' and b.workspace='{workspace}'
AND a.name='{source_node_id}' AND b.name='{target_node_id}'
COLUMNS (e.source_name,e.target_name) )""",
"node_degree":
"""SELECT count(1) as degree FROM GRAPH_TABLE (lightrag_graph
"node_degree": """SELECT count(1) as degree FROM GRAPH_TABLE (lightrag_graph
MATCH (a)-[e]->(b)
WHERE a.workspace='{workspace}' and a.workspace='{workspace}' and b.workspace='{workspace}'
AND a.name='{node_id}' or b.name = '{node_id}'
COLUMNS (a.name))""",
"get_node":
"""SELECT t1.name,t2.entity_type,t2.source_chunk_id as source_id,NVL(t2.description,'') AS description
"get_node": """SELECT t1.name,t2.entity_type,t2.source_chunk_id as source_id,NVL(t2.description,'') AS description
FROM GRAPH_TABLE (lightrag_graph
MATCH (a)
WHERE a.workspace='{workspace}' AND a.name='{node_id}'
COLUMNS (a.name)
) t1 JOIN LIGHTRAG_GRAPH_NODES t2 on t1.name=t2.name
WHERE t2.workspace='{workspace}'""",
"get_edge":
"""SELECT t1.source_id,t2.weight,t2.source_chunk_id as source_id,t2.keywords,
"get_edge": """SELECT t1.source_id,t2.weight,t2.source_chunk_id as source_id,t2.keywords,
NVL(t2.description,'') AS description,NVL(t2.KEYWORDS,'') AS keywords
FROM GRAPH_TABLE (lightrag_graph
MATCH (a)-[e]->(b)
@@ -659,15 +679,12 @@ SQL_TEMPLATES = {
AND a.name='{source_node_id}' and b.name = '{target_node_id}'
COLUMNS (e.id,a.name as source_id)
) t1 JOIN LIGHTRAG_GRAPH_EDGES t2 on t1.id=t2.id""",
"get_node_edges":
"""SELECT source_name,target_name
"get_node_edges": """SELECT source_name,target_name
FROM GRAPH_TABLE (lightrag_graph
MATCH (a)-[e]->(b)
WHERE e.workspace='{workspace}' and a.workspace='{workspace}' and b.workspace='{workspace}'
AND a.name='{source_node_id}'
COLUMNS (a.name as source_name,b.name as target_name))""",
"merge_node": """MERGE INTO LIGHTRAG_GRAPH_NODES a
USING DUAL
ON (a.workspace = '{workspace}' and a.name='{name}' and a.source_chunk_id='{source_chunk_id}')
@@ -679,5 +696,5 @@ SQL_TEMPLATES = {
ON (a.workspace = '{workspace}' and a.source_name='{source_name}' and a.target_name='{target_name}' and a.source_chunk_id='{source_chunk_id}')
WHEN NOT MATCHED THEN
INSERT(workspace,source_name,target_name,weight,keywords,description,source_chunk_id,content,content_vector)
values (:1,:2,:3,:4,:5,:6,:7,:8,:9) """
values (:1,:2,:3,:4,:5,:6,:7,:8,:9) """,
}

38
lightrag/lightrag.py
View File

@@ -42,11 +42,7 @@ from .storage import (
from .kg.neo4j_impl import Neo4JStorage
from .kg.oracle_impl import (
OracleKVStorage,
OracleGraphStorage,
OracleVectorDBStorage
)
from .kg.oracle_impl import OracleKVStorage, OracleGraphStorage, OracleVectorDBStorage
# future KG integrations
@@ -54,6 +50,7 @@ from .kg.oracle_impl import (
# GraphStorage as ArangoDBStorage
# )
def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
try:
return asyncio.get_event_loop()
@@ -134,18 +131,25 @@ class LightRAG:
# @TODO: should move all storage setup here to leverage initial start params attached to self.
self.key_string_value_json_storage_cls: Type[BaseKVStorage] = self._get_storage_class()[self.kv_storage]
self.vector_db_storage_cls: Type[BaseVectorStorage] = self._get_storage_class()[self.vector_storage]
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[self.graph_storage]
self.key_string_value_json_storage_cls: Type[BaseKVStorage] = (
self._get_storage_class()[self.kv_storage]
)
self.vector_db_storage_cls: Type[BaseVectorStorage] = self._get_storage_class()[
self.vector_storage
]
self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
self.graph_storage
]
if not os.path.exists(self.working_dir):
logger.info(f"Creating working directory {self.working_dir}")
os.makedirs(self.working_dir)
self.llm_response_cache = (
self.key_string_value_json_storage_cls(
namespace="llm_response_cache", global_config=asdict(self),embedding_func=None
namespace="llm_response_cache",
global_config=asdict(self),
embedding_func=None,
)
if self.enable_llm_cache
else None
@@ -159,13 +163,19 @@ class LightRAG:
# add embedding func by walter
####
self.full_docs = self.key_string_value_json_storage_cls(
namespace="full_docs", global_config=asdict(self), embedding_func=self.embedding_func
namespace="full_docs",
global_config=asdict(self),
embedding_func=self.embedding_func,
)
self.text_chunks = self.key_string_value_json_storage_cls(
namespace="text_chunks", global_config=asdict(self), embedding_func=self.embedding_func
namespace="text_chunks",
global_config=asdict(self),
embedding_func=self.embedding_func,
)
self.chunk_entity_relation_graph = self.graph_storage_cls(
namespace="chunk_entity_relation", global_config=asdict(self), embedding_func=self.embedding_func
namespace="chunk_entity_relation",
global_config=asdict(self),
embedding_func=self.embedding_func,
)
####
# add embedding func by walter over
@@ -202,11 +212,9 @@ class LightRAG:
# kv storage
"JsonKVStorage": JsonKVStorage,
"OracleKVStorage": OracleKVStorage,
# vector storage
"NanoVectorDBStorage": NanoVectorDBStorage,
"OracleVectorDBStorage": OracleVectorDBStorage,
# graph storage
"NetworkXStorage": NetworkXStorage,
"Neo4JStorage": Neo4JStorage,

2
lightrag/operate.py
View File

@@ -16,7 +16,7 @@ from .utils import (
split_string_by_multi_markers,
truncate_list_by_token_size,
process_combine_contexts,
locate_json_string_body_from_string
locate_json_string_body_from_string,
)
from .base import (
BaseGraphStorage,

24
requirements.txt
View File

@@ -1,22 +1,22 @@
accelerate
aioboto3
aiohttp
# database packages
graspologic
hnswlib
nano-vectordb
neo4j
networkx
ollama
openai
oracledb
pyvis
tenacity
xxhash
# lmdeploy[all]
# LLM packages
tiktoken
torch
transformers
aioboto3
ollama
openai
# database packages
graspologic
hnswlib
networkx
oracledb
nano-vectordb
neo4j
xxhash