'commit'

dsRagAnything/Util/LightRagUtil.py

@@ -0,0 +1,185 @@
+import logging
+import logging.config
+import os
+
+import numpy as np
+from lightrag import LightRAG
+from lightrag.kg.shared_storage import initialize_pipeline_status
+from lightrag.llm.openai import openai_complete_if_cache, openai_embed
+from lightrag.rerank import custom_rerank
+from lightrag.utils import EmbeddingFunc, logger, set_verbose_debug
+
+import Config.Config
+from Config.Config import *
+
+
+async def print_stream(stream):
+    async for chunk in stream:
+        if chunk:
+            print(chunk, end="", flush=True)
+
+
+def configure_logging():
+    for logger_name in ["uvicorn", "uvicorn.access", "uvicorn.error", "lightrag"]:
+        logger_instance = logging.getLogger(logger_name)
+        logger_instance.handlers = []
+        logger_instance.filters = []
+
+    log_dir = os.getenv("LOG_DIR", os.getcwd())
+    log_file_path = os.path.abspath(
+        os.path.join(log_dir, "./Logs/lightrag.log")
+    )
+
+    print(f"\nLightRAG log file: {log_file_path}\n")
+    os.makedirs(os.path.dirname(log_dir), exist_ok=True)
+
+    log_max_bytes = int(os.getenv("LOG_MAX_BYTES", 10485760))
+    log_backup_count = int(os.getenv("LOG_BACKUP_COUNT", 5))
+
+    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,
+                },
+            },
+        }
+    )
+    logger.setLevel(logging.INFO)
+    set_verbose_debug(os.getenv("VERBOSE_DEBUG", "false").lower() == "true")
+
+
+async def llm_model_func(
+        prompt, system_prompt=None, history_messages=None, **kwargs
+) -> str:
+    return await openai_complete_if_cache(
+        os.getenv("LLM_MODEL", Config.Config.ALY_LLM_MODEL_NAME),
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        api_key=Config.Config.ALY_LLM_API_KEY,
+        base_url=Config.Config.ALY_LLM_BASE_URL,
+        **kwargs,
+    )
+
+
+async def embedding_func(texts: list[str]) -> np.ndarray:
+    return await openai_embed(
+        texts,
+        model=EMBED_MODEL_NAME,
+        api_key=EMBED_API_KEY,
+        base_url=EMBED_BASE_URL
+    )
+
+async def rerank_func(query: str, documents: list, top_k: int = None, **kwargs):
+    return await custom_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,
+    )
+
+async def initialize_rag(working_dir, graph_storage=None):
+    if graph_storage is None:
+        graph_storage = 'NetworkXStorage'
+    rag = LightRAG(
+        working_dir=working_dir,
+        llm_model_func=llm_model_func,
+        graph_storage=graph_storage,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=EMBED_DIM,
+            max_token_size=EMBED_MAX_TOKEN_SIZE,
+            func=embedding_func
+        ),
+        rerank_model_func=rerank_func,
+    )
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+    return rag
+
+
+def create_llm_model_func():
+    def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
+        return openai_complete_if_cache(
+            Config.Config.ALY_LLM_MODEL_NAME,
+            prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            api_key=Config.Config.ALY_LLM_API_KEY,
+            base_url=Config.Config.ALY_LLM_BASE_URL,
+            **kwargs,
+        )
+
+    return llm_model_func
+
+
+def create_embedding_func():
+    return EmbeddingFunc(
+        embedding_dim=Config.Config.EMBED_DIM,
+        max_token_size=Config.Config.EMBED_MAX_TOKEN_SIZE,
+        func=lambda texts: openai_embed(
+            texts,
+            model=Config.Config.EMBED_MODEL_NAME,
+            api_key=Config.Config.EMBED_API_KEY,
+            base_url=Config.Config.EMBED_BASE_URL,
+        ),
+    )
+
+
+
+async def initialize_pg_rag(WORKING_DIR, workspace='default'):
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=llm_model_func,
+        llm_model_name=Config.Config.ALY_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=Config.Config.EMBED_DIM,
+            max_token_size=Config.Config.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