111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
"""
|
|
Debug script to test LLM function configuration in LightRAG
|
|
"""
|
|
|
|
import asyncio
|
|
import sys
|
|
import os
|
|
import logging
|
|
|
|
# Add the current directory to the path so we can import LightRAG
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def create_mock_llm_and_embedding():
|
|
"""Create mock LLM and embedding functions for testing"""
|
|
|
|
# Set up mock LLM function
|
|
async def mock_llm_func(prompt, system_prompt=None, **kwargs):
|
|
logger.info(f"Mock LLM called with prompt: {prompt[:100]}...")
|
|
return "Mock LLM response for testing"
|
|
|
|
# Set up mock embedding function wrapped with EmbeddingFunc
|
|
async def mock_embedding_func(texts):
|
|
logger.info(f"Mock embedding called with {len(texts)} texts")
|
|
import numpy as np
|
|
# Return random embeddings of dimension 384
|
|
return [np.random.rand(384).astype(np.float32) for _ in texts]
|
|
|
|
llm_func = mock_llm_func
|
|
|
|
# Wrap the embedding function with EmbeddingFunc class
|
|
from lightrag.utils import EmbeddingFunc
|
|
embedding_func = EmbeddingFunc(
|
|
embedding_dim=384,
|
|
func=mock_embedding_func
|
|
)
|
|
|
|
return llm_func, embedding_func
|
|
|
|
async def test_llm_function_setup():
|
|
"""Test LLM function setup and LightRAG initialization"""
|
|
|
|
llm_func, embedding_func = create_mock_llm_and_embedding()
|
|
|
|
logger.info(f"LLM function type: {type(llm_func)}")
|
|
logger.info(f"LLM function callable: {callable(llm_func)}")
|
|
logger.info(f"Embedding function type: {type(embedding_func)}")
|
|
logger.info(f"Embedding function callable: {callable(embedding_func)}")
|
|
logger.info(f"Embedding function embedding_dim: {embedding_func.embedding_dim}")
|
|
|
|
try:
|
|
from lightrag import LightRAG
|
|
|
|
# Test creating LightRAG instance
|
|
logger.info("Creating LightRAG instance...")
|
|
rag = LightRAG(
|
|
working_dir="./debug_test_storage",
|
|
kv_storage="JsonKVStorage",
|
|
vector_storage="NanoVectorDBStorage",
|
|
graph_storage="NetworkXStorage",
|
|
doc_status_storage="JsonDocStatusStorage",
|
|
llm_model_func=llm_func,
|
|
embedding_func=embedding_func
|
|
)
|
|
|
|
logger.info("LightRAG instance created successfully!")
|
|
|
|
# Test storage initialization
|
|
logger.info("Initializing storages...")
|
|
await rag.initialize_storages()
|
|
logger.info("Storages initialized successfully!")
|
|
|
|
# Initialize pipeline status
|
|
logger.info("Initializing pipeline status...")
|
|
from lightrag.kg.shared_storage import initialize_pipeline_status
|
|
await initialize_pipeline_status()
|
|
logger.info("Pipeline status initialized successfully!")
|
|
|
|
# Test basic insert
|
|
logger.info("Testing document insertion...")
|
|
test_document = "This is a test document about artificial intelligence and machine learning."
|
|
track_id = await rag.ainsert([test_document])
|
|
logger.info(f"Document inserted with track_id: {track_id}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Test failed: {e}")
|
|
import traceback
|
|
logger.error(traceback.format_exc())
|
|
return False
|
|
|
|
async def main():
|
|
"""Main test execution"""
|
|
logger.info("Starting LLM function debug test...")
|
|
|
|
success = await test_llm_function_setup()
|
|
|
|
if success:
|
|
logger.info("LLM function test PASSED!")
|
|
return 0
|
|
else:
|
|
logger.error("LLM function test FAILED!")
|
|
return 1
|
|
|
|
if __name__ == "__main__":
|
|
exit_code = asyncio.run(main())
|
|
sys.exit(exit_code) |