kssrag 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

kssrag-0.2.0.dist-info/RECORD

@@ -0,0 +1,33 @@
+kssrag/__init__.py,sha256=N1XfR8IRKtEJAzcOVyHnKXtgx-ztlrSFtFwiVkGbAX8,2041
+kssrag/cli.py,sha256=9AbtUEV9X63bhRj4EU-eHhud8iPM7LJAGSbu_IPlMUE,9703
+kssrag/config.py,sha256=zd978GZQ66TlLZnk9yP7uvoXoWD89BS8VHi7w_yGXrM,6529
+kssrag/kssrag.py,sha256=vy3oCHeHFAp_dJW0JjLbTxeEwCcwtXuOL_Ejmv0qz8Y,5251
+kssrag/server.py,sha256=CbnC0WhIKIi6iJ3q448swEdLDcvmUf80lsdlSKp0GpM,5942
+kssrag/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kssrag/core/agents.py,sha256=5zRSudh_4tbp4lDfAhaczeOe-INCpgfm6OJEhE6Ut4I,5421
+kssrag/core/chunkers.py,sha256=HmWL3y2DhhobV5zIlIdZP2KK2N7TASqeirPqmc3_inI,7324
+kssrag/core/retrievers.py,sha256=1e9c7ukUD4pFSVasOMTXSKoz_rapXQTl-FrSHK6Osqg,3037
+kssrag/core/vectorstores.py,sha256=H8hTpjc6hAFMhqAO2Cjq-Jp6xrJhsJKiRN9qxb_-6XM,21003
+kssrag/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kssrag/models/local_llms.py,sha256=IsthEwiNG1QcvHrTpQWdd1kZuHa4-0bfGTxHe8F3i2M,1178
+kssrag/models/openrouter.py,sha256=tplACtQ5J9YTemk0616dhg6H81_eAdsfeLs3AEytKE0,6429
+kssrag/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kssrag/utils/document_loaders.py,sha256=grXtU4sy8j23zJwanadO8rFXlsRJ2e2JF6MvoczsmqQ,4192
+kssrag/utils/helpers.py,sha256=MoTZRgTTiHDnbELHLrDKOBoxxiwFyLKJXUnZeWOtHFg,3806
+kssrag/utils/ocr.py,sha256=T2yZM-Z8B-1Y7K7CKxM5BrvNFPB5Cx0vjlk-XZnc3p8,1425
+kssrag/utils/ocr_loader.py,sha256=0RvY56aSNulo4U1eHwSzOUBZUIzc1nBLt7395OYdkXM,6930
+kssrag/utils/preprocessors.py,sha256=_kbeZOWnbqbKKSBiyRP8QZAKx9uYMXgHfARcWBqC3JU,938
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/test_basic.py,sha256=JdBBRpP9wOo4BvvZTisidP40gGyK_azUoewJpoJaa5M,1275
+tests/test_bm25s.py,sha256=tfvhWGxippGmNiLujc2OiaFewvvkJoOwrGXzBGZQMtU,2749
+tests/test_config.py,sha256=zIawdV9xb-EuDl1BXKKOvgZY-uUc5Q0KeyJHBP85eIE,1398
+tests/test_image_chunker.py,sha256=7cY3HucIFdNzcOVI2WA0nY5QmGcsv5umfE4c_yNnLfw,741
+tests/test_integration.py,sha256=TY7MrTcAiu1KG4MlgIC7VVlzUTnOoqp9pieK8rhBNrg,1059
+tests/test_ocr.py,sha256=PoGKLNISpAwaoPvGuS7qiOf6dsVnsFRFtYkG1WFi6TU,6202
+tests/test_streaming.py,sha256=rMQ0w8_HQFFV0PbHDqQXRBqaNfbd3WqJVNT2hKVbsqw,1442
+tests/test_vectorstores.py,sha256=YOwI2bfqprzbq8ahIw4pbbbEOaKGcg-XPcLCO7WiLxE,1474
+kssrag-0.2.0.dist-info/METADATA,sha256=MK2r6XV1oT76WuQJ01vdAbJxER9ZkCDalAITwMo3tNg,24016
+kssrag-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+kssrag-0.2.0.dist-info/entry_points.txt,sha256=g4tQj5YUqPK3Osb9BI85tsErxleSBUENiqlnX0fWK5M,43
+kssrag-0.2.0.dist-info/top_level.txt,sha256=sO9LGINa0GEjLoHTtufpz01yM5SmeTw6M4zWHEF0R2s,13
+kssrag-0.2.0.dist-info/RECORD,,

tests/test_bm25s.py

@@ -0,0 +1,74 @@
+import pytest
+import numpy as np
+import tempfile
+import os
+from kssrag.core.vectorstores import BM25SVectorStore
+
+def test_bm25s_vector_store_basic():
+    """Test BM25S vector store basic functionality"""
+    documents = [
+        {"content": "This is a test document about Python programming.", "metadata": {"source": "test1"}},
+        {"content": "Another document about machine learning and AI.", "metadata": {"source": "test2"}},
+        {"content": "A third document on web development with JavaScript.", "metadata": {"source": "test3"}},
+    ]
+    
+    vector_store = BM25SVectorStore()
+    vector_store.add_documents(documents)
+    
+    results = vector_store.retrieve("Python programming", top_k=2)
+    
+    assert len(results) == 2
+    assert "Python" in results[0]["content"]
+    assert all("metadata" in result for result in results)
+
+def test_bm25s_persistence():
+    """Test BM25S vector store persistence"""
+    documents = [
+        {"content": "Test document for persistence.", "metadata": {"source": "test1"}},
+        {"content": "Another test document.", "metadata": {"source": "test2"}},
+    ]
+    
+    with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False) as f:
+        temp_file = f.name
+    
+    try:
+        # Create and persist
+        vector_store = BM25SVectorStore(persist_path=temp_file)
+        vector_store.add_documents(documents)
+        vector_store.persist()
+        
+        # Load and verify
+        new_vector_store = BM25SVectorStore(persist_path=temp_file)
+        new_vector_store.load()
+        
+        results = new_vector_store.retrieve("persistence", top_k=1)
+        assert len(results) == 1
+        assert "persistence" in results[0]["content"]
+        
+    finally:
+        if os.path.exists(temp_file):
+            os.unlink(temp_file)
+
+def test_bm25s_empty_query():
+    """Test BM25S with empty query"""
+    documents = [
+        {"content": "Test document.", "metadata": {"source": "test1"}},
+    ]
+    
+    vector_store = BM25SVectorStore()
+    vector_store.add_documents(documents)
+    
+    results = vector_store.retrieve("", top_k=1)
+    # BM25S may return documents even with empty query, but they should have low scores
+    # Let's check that the behavior is consistent
+    if len(results) > 0:
+        # If it returns results, they should be the documents we added
+        assert results[0]["content"] == "Test document."
+    # Either behavior is acceptable for this test
+
+def test_bm25s_no_documents():
+    """Test BM25S with no documents added"""
+    vector_store = BM25SVectorStore()
+    
+    with pytest.raises(ValueError, match="BM25S index not initialized"):
+        vector_store.retrieve("test query")

tests/test_config.py

@@ -0,0 +1,42 @@
+import pytest
+import os
+from kssrag.config import Config, VectorStoreType, ChunkerType
+
+def test_config_new_options():
+    """Test new configuration options"""
+    config = Config(
+        OCR_DEFAULT_MODE="handwritten",
+        ENABLE_STREAMING=True,
+        VECTOR_STORE_TYPE=VectorStoreType.BM25S
+    )
+    
+    assert config.OCR_DEFAULT_MODE == "handwritten"
+    assert config.ENABLE_STREAMING == True
+    assert config.VECTOR_STORE_TYPE == VectorStoreType.BM25S
+
+def test_config_vector_store_types():
+    """Test all vector store types including BM25S"""
+    config = Config(VECTOR_STORE_TYPE=VectorStoreType.BM25S)
+    assert config.VECTOR_STORE_TYPE == "bm25s"
+    
+    config = Config(VECTOR_STORE_TYPE=VectorStoreType.BM25)
+    assert config.VECTOR_STORE_TYPE == "bm25"
+
+def test_config_chunker_types():
+    """Test all chunker types including image"""
+    config = Config(CHUNKER_TYPE=ChunkerType.IMAGE)
+    assert config.CHUNKER_TYPE == "image"
+
+def test_config_environment_variables():
+    """Test new environment variables"""
+    os.environ["OCR_DEFAULT_MODE"] = "handwritten"
+    os.environ["ENABLE_STREAMING"] = "true"
+    
+    config = Config()
+    
+    assert config.OCR_DEFAULT_MODE == "handwritten"
+    assert config.ENABLE_STREAMING == True
+    
+    # Cleanup
+    del os.environ["OCR_DEFAULT_MODE"]
+    del os.environ["ENABLE_STREAMING"]

tests/test_image_chunker.py

@@ -0,0 +1,17 @@
+import pytest
+from kssrag.core.chunkers import ImageChunker, OCR_AVAILABLE
+
+@pytest.mark.skipif(not OCR_AVAILABLE, reason="OCR dependencies not available")
+def test_image_chunker_basic():
+    """Basic test for ImageChunker - just check it initializes"""
+    chunker = ImageChunker(ocr_mode="typed")
+    assert chunker.ocr_mode == "typed"
+
+@pytest.mark.skipif(not OCR_AVAILABLE, reason="OCR dependencies not available") 
+def test_image_chunker_modes():
+    """Test that ImageChunker accepts valid modes"""
+    chunker_typed = ImageChunker(ocr_mode="typed")
+    chunker_handwritten = ImageChunker(ocr_mode="handwritten")
+    
+    assert chunker_typed.ocr_mode == "typed"
+    assert chunker_handwritten.ocr_mode == "handwritten"

tests/test_integration.py

@@ -0,0 +1,35 @@
+import pytest
+import tempfile
+import os
+from kssrag import KSSRAG, Config
+
+def test_bm25s_integration():
+    """Test BM25S integration with KSSRAG"""
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
+        f.write("Test document about Python programming and machine learning.")
+        temp_file = f.name
+    
+    try:
+        config = Config(
+            VECTOR_STORE_TYPE="bm25s",
+            MAX_DOCS_FOR_TESTING=1
+        )
+        
+        rag = KSSRAG(config=config)
+        rag.load_document(temp_file, format="text")
+        
+        response = rag.query("Python programming")
+        
+        assert isinstance(response, str)
+        assert len(response) > 0
+        
+    finally:
+        os.unlink(temp_file)
+
+def test_streaming_integration():
+    """Test streaming integration (mock test)"""
+    config = Config(ENABLE_STREAMING=True)
+    
+    # This is a basic test that config is accepted
+    # Actual streaming would require API calls
+    assert config.ENABLE_STREAMING == True

tests/test_ocr.py

@@ -0,0 +1,142 @@
+import pytest
+import tempfile
+import os
+from unittest.mock import Mock, patch, MagicMock
+from kssrag.utils.ocr_loader import OCRLoader
+
+def test_ocr_loader_initialization():
+    """Test OCRLoader initialization with mocked PaddleOCR"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        mock_instance = Mock()
+        mock_paddle.return_value = mock_instance
+        
+        loader = OCRLoader()
+        assert loader.paddle_ocr == mock_instance
+
+def test_ocr_loader_invalid_mode():
+    """Test OCRLoader with invalid mode"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        mock_instance = Mock()
+        mock_paddle.return_value = mock_instance
+        
+        loader = OCRLoader()
+        
+        with pytest.raises(ValueError, match="Invalid OCR mode"):
+            loader.extract_text("test.jpg", "invalid_mode")
+
+def test_ocr_loader_file_not_found():
+    """Test OCRLoader with non-existent file"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        mock_instance = Mock()
+        mock_paddle.return_value = mock_instance
+        
+        loader = OCRLoader()
+        
+        with pytest.raises(FileNotFoundError):
+            loader.extract_text("nonexistent.jpg", "typed")
+
+def test_ocr_loader_typed_mode():
+    """Test OCRLoader typed mode"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        with patch('kssrag.utils.ocr_loader.pytesseract') as mock_tesseract:
+            with patch('kssrag.utils.ocr_loader.Image') as mock_image:
+                with patch('kssrag.utils.ocr_loader.os.path.exists') as mock_exists:
+                    # Mock file exists
+                    mock_exists.return_value = True
+                    
+                    # Mock image opening
+                    mock_img_instance = MagicMock()
+                    mock_image.open.return_value = mock_img_instance
+                    
+                    # Mock OCR result
+                    mock_tesseract.image_to_string.return_value = "Typed text content"
+                    
+                    mock_paddle_instance = Mock()
+                    mock_paddle.return_value = mock_paddle_instance
+                    
+                    loader = OCRLoader()
+                    
+                    result = loader.extract_text("test.jpg", "typed")
+                    assert result == "Typed text content"
+                    mock_tesseract.image_to_string.assert_called_once_with(mock_img_instance)
+
+def test_ocr_loader_handwritten_mode():
+    """Test OCRLoader handwritten mode"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        with patch('kssrag.utils.ocr_loader.cv2') as mock_cv2:
+            with patch('kssrag.utils.ocr_loader.os.path.exists') as mock_exists:
+                # Mock file exists
+                mock_exists.return_value = True
+                
+                # Mock image reading
+                mock_cv2.imread.return_value = "mock_image"
+                
+                # Mock OCR result
+                mock_paddle_instance = Mock()
+                mock_paddle_instance.ocr.return_value = [[[None, ["Handwritten text", 0.9]]]]
+                mock_paddle.return_value = mock_paddle_instance
+                
+                loader = OCRLoader()
+                
+                result = loader.extract_text("test.jpg", "handwritten")
+                assert result == "Handwritten text"
+                mock_paddle_instance.ocr.assert_called_once_with("mock_image", cls=True)
+
+def test_ocr_loader_paddle_not_initialized():
+    """Test OCRLoader when PaddleOCR is not initialized"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        mock_paddle.return_value = None  # Simulate initialization failure
+        
+        loader = OCRLoader()
+        loader.paddle_ocr = None  # Force the failure state
+        
+        with pytest.raises(RuntimeError, match="PaddleOCR not initialized"):
+            loader.extract_text("test.jpg", "handwritten")
+
+def test_ocr_loader_empty_text():
+    """Test OCRLoader when no text is extracted"""
+    with patch('kssrag.utils.ocr_loader.PaddleOCR') as mock_paddle:
+        with patch('kssrag.utils.ocr_loader.pytesseract') as mock_tesseract:
+            with patch('kssrag.utils.ocr_loader.Image') as mock_image:
+                with patch('kssrag.utils.ocr_loader.os.path.exists') as mock_exists:
+                    # Mock file exists
+                    mock_exists.return_value = True
+                    
+                    # Mock image opening
+                    mock_img_instance = MagicMock()
+                    mock_image.open.return_value = mock_img_instance
+                    
+                    # Mock empty OCR result
+                    mock_tesseract.image_to_string.return_value = "   "  # Only whitespace
+                    
+                    mock_paddle_instance = Mock()
+                    mock_paddle.return_value = mock_paddle_instance
+                    
+                    loader = OCRLoader()
+                    
+                    result = loader.extract_text("test.jpg", "typed")
+                    assert result == ""  # Should return empty string
+
+@pytest.mark.skipif(not os.getenv('TEST_OCR'), reason="OCR tests require actual OCR dependencies")
+def test_ocr_loader_integration():
+    """Integration test for OCRLoader with actual image"""
+    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
+        temp_file = f.name
+    
+    try:
+        # Create a simple test image with text
+        from PIL import Image, ImageDraw
+        img = Image.new('RGB', (400, 100), color='white')
+        draw = ImageDraw.Draw(img)
+        draw.text((50, 40), "OCR Test Text", fill='black')
+        img.save(temp_file)
+        
+        loader = OCRLoader()
+        text = loader.extract_text(temp_file, "typed")
+        
+        assert isinstance(text, str)
+        # Note: Actual OCR text recognition might vary
+        
+    finally:
+        if os.path.exists(temp_file):
+            os.unlink(temp_file)

tests/test_streaming.py

@@ -0,0 +1,41 @@
+import pytest
+import asyncio
+from kssrag.models.openrouter import OpenRouterLLM
+
+def test_openrouter_streaming_initialization():
+    """Test OpenRouterLLM streaming initialization"""
+    llm = OpenRouterLLM(stream=True)
+    assert llm.stream == True
+
+def test_openrouter_non_streaming_initialization():
+    """Test OpenRouterLLM non-streaming initialization"""
+    llm = OpenRouterLLM(stream=False)
+    assert llm.stream == False
+
+def test_streaming_generator():
+    """Test streaming generator interface"""
+    # Mock the predict_stream method for testing
+    class TestOpenRouterLLM(OpenRouterLLM):
+        def predict_stream(self, messages):
+            yield "Hello "
+            yield "World"
+            yield "!"
+    
+    llm = TestOpenRouterLLM(stream=True)
+    messages = [{"role": "user", "content": "test"}]
+    
+    chunks = list(llm.predict_stream(messages))
+    assert chunks == ["Hello ", "World", "!"]
+
+def test_streaming_fallback_to_non_streaming():
+    """Test that streaming falls back to non-streaming when no chunks"""
+    class TestOpenRouterLLM(OpenRouterLLM):
+        def predict_stream(self, messages):
+            # Simulate no chunks returned
+            if False:
+                yield "test"
+    
+    llm = TestOpenRouterLLM(stream=True)
+    # This should not raise an error
+    result = llm.predict([{"role": "user", "content": "test"}])
+    assert isinstance(result, str)