kssrag 0.1.0__py3-none-any.whl

tests/test_basic.py

@@ -0,0 +1,43 @@
+import pytest
+import os
+import tempfile
+from kssrag import KSSRAG, Config
+
+def test_text_rag():
+    """Test basic text RAG functionality"""
+    # Create a temporary text file
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
+        f.write("This is a test document about artificial intelligence and machine learning.")
+        temp_file = f.name
+    
+    try:
+        # Initialize with test config
+        config = Config(
+            OPENROUTER_API_KEY=os.getenv("OPENROUTER_API_KEY", "test_key"),
+            MAX_DOCS_FOR_TESTING=1
+        )
+        
+        rag = KSSRAG(config=config)
+        rag.load_document(temp_file, format="text")
+        
+        # Test query
+        response = rag.query("What is this document about?")
+        
+        assert isinstance(response, str)
+        assert len(response) > 0
+        
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+def test_config_validation():
+    """Test configuration validation"""
+    config = Config(
+        OPENROUTER_API_KEY="test_key",
+        CHUNK_SIZE=500,
+        CHUNK_OVERLAP=50
+    )
+    
+    assert config.OPENROUTER_API_KEY == "test_key"
+    assert config.CHUNK_SIZE == 500
+    assert config.CHUNK_OVERLAP == 50

tests/test_vectorstores.py

@@ -0,0 +1,35 @@
+import pytest
+import numpy as np
+from kssrag.core.vectorstores import BM25VectorStore, TFIDFVectorStore
+
+def test_bm25_vector_store():
+    """Test BM25 vector store 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 = BM25VectorStore()
+    vector_store.add_documents(documents)
+    
+    results = vector_store.retrieve("Python programming", top_k=2)
+    
+    assert len(results) == 2
+    assert "Python" in results[0]["content"]
+
+def test_tfidf_vector_store():
+    """Test TFIDF vector store 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 = TFIDFVectorStore()
+    vector_store.add_documents(documents)
+    
+    results = vector_store.retrieve("machine learning", top_k=2)
+    
+    assert len(results) == 2
+    assert "machine learning" in results[0]["content"].lower()