|
|
|
@ -155,13 +155,38 @@ def create_embedding_func():
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def rerank_func(query: str, documents: list, top_k: int = None, **kwargs):
|
|
|
|
|
return await jina_rerank(
|
|
|
|
|
query=query,
|
|
|
|
|
documents=documents,
|
|
|
|
|
model=Config.Config.RERANK_MODEL,
|
|
|
|
|
base_url=Config.Config.RERANK_BASE_URL,
|
|
|
|
|
api_key=Config.Config.RERANK_BINDING_API_KEY,
|
|
|
|
|
top_k=top_k or 10,
|
|
|
|
|
**kwargs
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# AGE
|
|
|
|
|
os.environ["POSTGRES_HOST"] = POSTGRES_HOST
|
|
|
|
|
os.environ["POSTGRES_PORT"] = str(POSTGRES_PORT)
|
|
|
|
|
os.environ["POSTGRES_USER"] = POSTGRES_USER
|
|
|
|
|
os.environ["POSTGRES_PASSWORD"] = POSTGRES_PASSWORD
|
|
|
|
|
os.environ["POSTGRES_DATABASE"] = POSTGRES_DATABASE
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def initialize_pg_rag(WORKING_DIR, workspace='default'):
|
|
|
|
|
rag = LightRAG(
|
|
|
|
|
working_dir=WORKING_DIR,
|
|
|
|
|
llm_model_func=llm_model_func,
|
|
|
|
|
llm_model_name=LLM_MODEL_NAME,
|
|
|
|
|
llm_model_max_async=4,
|
|
|
|
|
llm_model_max_token_size=32768,
|
|
|
|
|
enable_llm_cache_for_entity_extract=True,
|
|
|
|
|
embedding_func=EmbeddingFunc(
|
|
|
|
|
embedding_dim=EMBED_DIM,
|
|
|
|
|
max_token_size=EMBED_MAX_TOKEN_SIZE,
|
|
|
|
|
func=embedding_func
|
|
|
|
|
),
|
|
|
|
|
rerank_model_func=rerank_func,
|
|
|
|
|
kv_storage="PGKVStorage",
|
|
|
|
|
doc_status_storage="PGDocStatusStorage",
|
|
|
|
|
graph_storage="PGGraphStorage",
|
|
|
|
|
vector_storage="PGVectorStorage",
|
|
|
|
|
auto_manage_storages_states=False,
|
|
|
|
|
vector_db_storage_cls_kwargs={"workspace": workspace}
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
await rag.initialize_storages()
|
|
|
|
|
await initialize_pipeline_status()
|
|
|
|
|
|
|
|
|
|
return rag
|