kssrag 0.1.0__py3-none-any.whl

kssrag/core/vectorstores.py

@@ -0,0 +1,397 @@
+import json
+import os
+import re
+import pickle
+import numpy as np
+import faiss
+from rank_bm25 import BM25Okapi
+from sentence_transformers import SentenceTransformer
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import scipy.sparse as sp
+from typing import List, Dict, Any, Optional
+from ..utils.helpers import logger
+from ..config import config
+
+class BaseVectorStore:
+    """Base class for vector stores"""
+    
+    def __init__(self, persist_path: Optional[str] = None):
+        self.persist_path = persist_path
+        self.documents: List[Dict[str, Any]] = []
+        self.doc_texts: List[str] = []
+    
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        """Add documents to the vector store"""
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        """Retrieve documents based on query"""
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def persist(self):
+        """Persist the vector store to disk"""
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def load(self):
+        """Load the vector store from disk"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+class BM25VectorStore(BaseVectorStore):
+    """BM25 vector store implementation"""
+    
+    def __init__(self, persist_path: Optional[str] = "bm25_index.pkl"):
+        super().__init__(persist_path)
+        self.bm25 = None
+    
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        self.documents = documents
+        self.doc_texts = [doc["content"] for doc in documents]
+        
+        # Tokenize corpus for BM25
+        tokenized_corpus = [self._tokenize(doc) for doc in self.doc_texts]
+        self.bm25 = BM25Okapi(tokenized_corpus)
+        
+        logger.info(f"BM25 index created with {len(self.documents)} documents")
+    
+    def _tokenize(self, text: str) -> List[str]:
+        """Tokenize text for BM25"""
+        return re.findall(r'\w+', text.lower())
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        if not self.bm25:
+            raise ValueError("BM25 index not initialized. Call add_documents first.")
+        
+        try:
+            tokenized_query = self._tokenize(query)
+            doc_scores = self.bm25.get_scores(tokenized_query)
+            
+            # Get top-k indices
+            top_indices = np.argsort(doc_scores)[::-1][:top_k]
+            
+            # Filter out invalid indices
+            valid_indices = [i for i in top_indices if i < len(self.documents)]
+            
+            if not valid_indices:
+                logger.warning(f"No valid indices found for query: {query}")
+                return []
+                
+            return [self.documents[i] for i in valid_indices]
+            
+        except Exception as e:
+            logger.error(f"Error in BM25 retrieval: {str(e)}")
+            return []
+    
+    def persist(self):
+        if self.persist_path:
+            with open(self.persist_path, 'wb') as f:
+                pickle.dump({
+                    'documents': self.documents,
+                    'doc_texts': self.doc_texts,
+                    'bm25': self.bm25
+                }, f)
+            logger.info(f"BM25 index persisted to {self.persist_path}")
+    
+    def load(self):
+        if self.persist_path and os.path.exists(self.persist_path):
+            with open(self.persist_path, 'rb') as f:
+                data = pickle.load(f)
+                self.documents = data['documents']
+                self.doc_texts = data['doc_texts']
+                self.bm25 = data['bm25']
+            logger.info(f"BM25 index loaded from {self.persist_path}")
+
+import tempfile
+class FAISSVectorStore(BaseVectorStore):
+    def __init__(self, persist_path: Optional[str] = None, model_name: Optional[str] = None):
+        super().__init__(persist_path)
+        self.model_name = model_name or config.FAISS_MODEL_NAME
+        
+        # Handle cache directory permissions
+        try:
+            cache_dir = config.CACHE_DIR
+            os.makedirs(cache_dir, exist_ok=True)
+            # Test if we can write to the cache directory
+            test_file = os.path.join(cache_dir, 'write_test.txt')
+            with open(test_file, 'w') as f:
+                f.write('test')
+            os.remove(test_file)
+        except PermissionError:
+            logger.warning(f"Could not write to cache directory {cache_dir}. Using temp directory.")
+            cache_dir = tempfile.gettempdir()
+        
+        self.model = SentenceTransformer(self.model_name, cache_folder=cache_dir)
+        self.dimension = self.model.get_sentence_embedding_dimension()
+        self.index = faiss.IndexFlatL2(self.dimension)
+        self.metadata_path = persist_path + ".meta" if persist_path else None
+    
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        self.documents = documents
+        self.doc_texts = [doc["content"] for doc in documents]
+        
+        # Generate embeddings in batches
+        embeddings = []
+        batch_size = config.BATCH_SIZE
+        
+        for i in range(0, len(self.doc_texts), batch_size):
+            batch_texts = self.doc_texts[i:i+batch_size]
+            batch_embeddings = self.model.encode(batch_texts, show_progress_bar=False)
+            embeddings.append(batch_embeddings)
+            logger.info(f"Processed batch {i//batch_size + 1}/{(len(self.doc_texts)-1)//batch_size + 1}")
+        
+        embeddings = np.vstack(embeddings).astype('float32')
+        self.index.add(embeddings)
+        
+        logger.info(f"FAISS index created with {len(self.documents)} documents")
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        if not self.index or self.index.ntotal == 0:
+            raise ValueError("FAISS index not initialized. Call add_documents first.")
+        
+        try:
+            query_embedding = self.model.encode([query])
+            distances, indices = self.index.search(query_embedding.astype('float32'), top_k)
+            
+            # Filter out invalid indices (FAISS might return -1 for no results)
+            valid_indices = [i for i in indices[0] if i >= 0 and i < len(self.documents)]
+            
+            if not valid_indices:
+                logger.warning(f"No valid indices found for query: {query}")
+                return []
+                
+            return [self.documents[i] for i in valid_indices]
+        except Exception as e:
+            logger.error(f"Error in FAISS retrieval: {str(e)}")
+            return []
+    
+    def persist(self):
+        if self.persist_path and self.metadata_path:
+            faiss.write_index(self.index, self.persist_path)
+            
+            # Save metadata
+            with open(self.metadata_path, 'wb') as f:
+                pickle.dump({
+                    'documents': self.documents,
+                    'doc_texts': self.doc_texts
+                }, f)
+            logger.info(f"FAISS index persisted to {self.persist_path}")
+    
+    def load(self):
+        if (self.persist_path and os.path.exists(self.persist_path) and 
+            self.metadata_path and os.path.exists(self.metadata_path)):
+            
+            self.index = faiss.read_index(self.persist_path)
+            
+            # Load metadata
+            with open(self.metadata_path, 'rb') as f:
+                data = pickle.load(f)
+                self.documents = data['documents']
+                self.doc_texts = data['doc_texts']
+            logger.info(f"FAISS index loaded from {self.persist_path}")
+
+class TFIDFVectorStore(BaseVectorStore):
+    """TFIDF vector store implementation"""
+    
+    def __init__(self, persist_path: Optional[str] = "tfidf_index.pkl", max_features: int = 10000):
+        super().__init__(persist_path)
+        self.vectorizer = TfidfVectorizer(max_features=max_features, stop_words='english')
+        self.tfidf_matrix = None
+    
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        self.documents = documents
+        self.doc_texts = [doc["content"] for doc in documents]
+        
+        # Fit and transform the documents
+        self.tfidf_matrix = self.vectorizer.fit_transform(self.doc_texts)
+        
+        logger.info(f"TFIDF index created with {len(self.documents)} documents")
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        if self.tfidf_matrix is None:
+            raise ValueError("TFIDF index not initialized. Call add_documents first.")
+        
+        try:
+            # Transform the query
+            query_vec = self.vectorizer.transform([query])
+            
+            # Calculate cosine similarities
+            similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten()
+            
+            # Get top-k indices
+            top_indices = np.argsort(similarities)[::-1][:top_k]
+            
+            # Filter out invalid indices
+            valid_indices = [i for i in top_indices if i < len(self.documents)]
+            
+            if not valid_indices:
+                logger.warning(f"No valid indices found for query: {query}")
+                return []
+                
+            return [self.documents[i] for i in valid_indices]
+            
+        except Exception as e:
+            logger.error(f"Error in TFIDF retrieval: {str(e)}")
+            return []
+    
+    def persist(self):
+        if self.persist_path:
+            with open(self.persist_path, 'wb') as f:
+                pickle.dump({
+                    'documents': self.documents,
+                    'doc_texts': self.doc_texts,
+                    'vectorizer': self.vectorizer,
+                    'tfidf_matrix': self.tfidf_matrix
+                }, f)
+            logger.info(f"TFIDF index persisted to {self.persist_path}")
+    
+    def load(self):
+        if self.persist_path and os.path.exists(self.persist_path):
+            with open(self.persist_path, 'rb') as f:
+                data = pickle.load(f)
+                self.documents = data['documents']
+                self.doc_texts = data['doc_texts']
+                self.vectorizer = data['vectorizer']
+                self.tfidf_matrix = data['tfidf_matrix']
+            logger.info(f"TFIDF index loaded from {self.persist_path}")
+
+class HybridVectorStore(BaseVectorStore):
+    """Hybrid vector store combining BM25 and FAISS"""
+    
+    def __init__(self, persist_path: Optional[str] = "hybrid_index"):
+        super().__init__(persist_path)
+        self.bm25_store = BM25VectorStore(persist_path + "_bm25")
+        self.faiss_store = FAISSVectorStore(persist_path + "_faiss")
+    
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        self.documents = documents
+        self.bm25_store.add_documents(documents)
+        self.faiss_store.add_documents(documents)
+        logger.info(f"Hybrid index created with {len(self.documents)} documents")
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        try:
+            # Get results from both methods
+            bm25_results = self.bm25_store.retrieve(query, top_k * 2)
+            faiss_results = self.faiss_store.retrieve(query, top_k * 2)
+            
+            # Combine and deduplicate by content
+            combined = {}
+            for doc in bm25_results + faiss_results:
+                # Use a combination of content and metadata for deduplication
+                key = hash(doc["content"] + str(doc["metadata"]))
+                if key not in combined:
+                    combined[key] = doc
+            
+            all_results = list(combined.values())
+            
+            # If no results after deduplication
+            if not all_results:
+                return []
+            
+            # Rerank by relevance to query using FAISS similarity
+            query_embedding = self.faiss_store.model.encode(query)
+            scored_results = []
+            
+            for doc in all_results:
+                doc_embedding = self.faiss_store.model.encode(doc["content"])
+                similarity = np.dot(query_embedding, doc_embedding) / (
+                    np.linalg.norm(query_embedding) * np.linalg.norm(doc_embedding))
+                
+                scored_results.append((doc, similarity))
+            
+            scored_results.sort(key=lambda x: x[1], reverse=True)
+            return [doc for doc, _ in scored_results[:top_k]]
+            
+        except Exception as e:
+            logger.error(f"Error in hybrid retrieval: {str(e)}")
+            return []
+    
+    def persist(self):
+        self.bm25_store.persist()
+        self.faiss_store.persist()
+        logger.info(f"Hybrid index persisted")
+    
+    def load(self):
+        self.bm25_store.load()
+        self.faiss_store.load()
+        self.documents = self.bm25_store.documents
+        logger.info(f"Hybrid index loaded")
+
+class HybridOfflineVectorStore(BaseVectorStore):
+    """Hybrid offline vector store combining BM25 and TFIDF"""
+    
+    def __init__(self, persist_path: Optional[str] = "hybrid_offline_index"):
+        super().__init__(persist_path)
+        self.bm25_store = BM25VectorStore(persist_path + "_bm25")
+        self.tfidf_store = TFIDFVectorStore(persist_path + "_tfidf")
+        self.alpha = 0.5  # Weight for BM25 vs TFIDF
+    
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        self.documents = documents
+        self.bm25_store.add_documents(documents)
+        self.tfidf_store.add_documents(documents)
+        logger.info(f"Hybrid offline index created with {len(self.documents)} documents")
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        try:
+            # Get results from both methods
+            bm25_results = self.bm25_store.retrieve(query, top_k * 2)
+            tfidf_results = self.tfidf_store.retrieve(query, top_k * 2)
+            
+            # Combine and deduplicate by content
+            combined = {}
+            for doc in bm25_results + tfidf_results:
+                # Use a combination of content and metadata for deduplication
+                key = hash(doc["content"] + str(doc["metadata"]))
+                if key not in combined:
+                    combined[key] = doc
+            
+            all_results = list(combined.values())
+            
+            # If no results after deduplication
+            if not all_results:
+                return []
+            
+            # Score results based on both methods
+            scored_results = []
+            
+            for doc in all_results:
+                # Get BM25 score
+                bm25_score = 0
+                for i, bm25_doc in enumerate(bm25_results):
+                    if (doc["content"] == bm25_doc["content"] and 
+                        doc["metadata"] == bm25_doc["metadata"]):
+                        # Normalize score based on position
+                        bm25_score = (len(bm25_results) - i) / len(bm25_results)
+                        break
+                
+                # Get TFIDF score
+                tfidf_score = 0
+                for i, tfidf_doc in enumerate(tfidf_results):
+                    if (doc["content"] == tfidf_doc["content"] and 
+                        doc["metadata"] == tfidf_doc["metadata"]):
+                        # Normalize score based on position
+                        tfidf_score = (len(tfidf_results) - i) / len(tfidf_results)
+                        break
+                
+                # Combine scores
+                combined_score = self.alpha * bm25_score + (1 - self.alpha) * tfidf_score
+                scored_results.append((doc, combined_score))
+            
+            scored_results.sort(key=lambda x: x[1], reverse=True)
+            return [doc for doc, _ in scored_results[:top_k]]
+            
+        except Exception as e:
+            logger.error(f"Error in hybrid offline retrieval: {str(e)}")
+            return []
+    
+    def persist(self):
+        self.bm25_store.persist()
+        self.tfidf_store.persist()
+        logger.info(f"Hybrid offline index persisted")
+    
+    def load(self):
+        self.bm25_store.load()
+        self.tfidf_store.load()
+        self.documents = self.bm25_store.documents
+        logger.info(f"Hybrid offline index loaded")

kssrag/kssrag.py

@@ -0,0 +1,116 @@
+"""
+Main KSSRAG class that ties everything together for easy usage.
+"""
+from typing import Optional, List, Dict, Any
+import os
+from .core.chunkers import TextChunker, JSONChunker, PDFChunker
+from .core.vectorstores import BM25VectorStore, FAISSVectorStore, TFIDFVectorStore, HybridVectorStore, HybridOfflineVectorStore
+from .core.retrievers import SimpleRetriever, HybridRetriever
+from .core.agents import RAGAgent
+from .models.openrouter import OpenRouterLLM
+from .utils.document_loaders import load_document, load_json_documents
+from .config import Config, VectorStoreType, ChunkerType, RetrieverType
+from .utils.helpers import logger, validate_config, import_custom_component
+
+class KSSRAG:
+    """Main class for KSS RAG functionality"""
+    
+    def __init__(self, config: Optional[Config] = None):
+        self.config = config or Config()
+        self.vector_store = None
+        self.retriever = None
+        self.agent = None
+        self.documents = []
+        
+        # Validate configuration
+        validate_config()
+    
+    def load_document(self, file_path: str, format: Optional[str] = None, 
+                     chunker: Optional[Any] = None, metadata: Optional[Dict[str, Any]] = None):
+        """Load and process a document"""
+        if format is None:
+            # Auto-detect format
+            if file_path.endswith('.txt'):
+                format = 'text'
+            elif file_path.endswith('.json'):
+                format = 'json'
+            elif file_path.endswith('.pdf'):
+                format = 'pdf'
+            else:
+                raise ValueError(f"Unsupported file format: {file_path}")
+        
+        # Use custom chunker if provided
+        if chunker is None:
+            if format == 'text':
+                chunker = TextChunker(chunk_size=self.config.CHUNK_SIZE, overlap=self.config.CHUNK_OVERLAP)
+            elif format == 'json':
+                chunker = JSONChunker()
+            elif format == 'pdf':
+                chunker = PDFChunker(chunk_size=self.config.CHUNK_SIZE, overlap=self.config.CHUNK_OVERLAP)
+        
+        if metadata is None:
+            metadata = {"source": file_path}
+        
+        # Load and chunk document
+        if format == 'text':
+            content = load_document(file_path)
+            self.documents = chunker.chunk(content, metadata)
+        elif format == 'json':
+            data = load_json_documents(file_path)
+            self.documents = chunker.chunk(data)
+        elif format == 'pdf':
+            self.documents = chunker.chunk_pdf(file_path, metadata)
+        
+        # Create vector store
+        if self.config.CUSTOM_VECTOR_STORE:
+            vector_store_class = import_custom_component(self.config.CUSTOM_VECTOR_STORE)
+            self.vector_store = vector_store_class()
+        else:
+            if self.config.VECTOR_STORE_TYPE == VectorStoreType.BM25:
+                self.vector_store = BM25VectorStore()
+            elif self.config.VECTOR_STORE_TYPE == VectorStoreType.FAISS:
+                self.vector_store = FAISSVectorStore()
+            elif self.config.VECTOR_STORE_TYPE == VectorStoreType.TFIDF:
+                self.vector_store = TFIDFVectorStore()
+            elif self.config.VECTOR_STORE_TYPE == VectorStoreType.HYBRID_ONLINE:
+                self.vector_store = HybridVectorStore()
+            elif self.config.VECTOR_STORE_TYPE == VectorStoreType.HYBRID_OFFLINE:
+                self.vector_store = HybridOfflineVectorStore()
+        
+        self.vector_store.add_documents(self.documents)
+        
+        # Create retriever
+        if self.config.CUSTOM_RETRIEVER:
+            retriever_class = import_custom_component(self.config.CUSTOM_RETRIEVER)
+            self.retriever = retriever_class(self.vector_store)
+        else:
+            if self.config.RETRIEVER_TYPE == RetrieverType.SIMPLE:
+                self.retriever = SimpleRetriever(self.vector_store)
+            elif self.config.RETRIEVER_TYPE == RetrieverType.HYBRID:
+                # For hybrid retriever, extract entity names from documents if available
+                entity_names = []
+                if format == 'json':
+                    entity_names = [doc['metadata'].get('name', '') for doc in self.documents if doc['metadata'].get('name')]
+                self.retriever = HybridRetriever(self.vector_store, entity_names)
+        
+        # Create LLM
+        if self.config.CUSTOM_LLM:
+            llm_class = import_custom_component(self.config.CUSTOM_LLM)
+            llm = llm_class()
+        else:
+            llm = OpenRouterLLM()
+        
+        # Create agent
+        self.agent = RAGAgent(self.retriever, llm)
+    
+    def query(self, question: str, top_k: Optional[int] = None) -> str:
+        """Query the RAG system"""
+        if not self.agent:
+            raise ValueError("No documents loaded. Call load_document first.")
+        
+        return self.agent.query(question, top_k=top_k or self.config.TOP_K)
+    
+    def create_server(self, server_config=None):
+        """Create a FastAPI server for the RAG system"""
+        from .server import create_app
+        return create_app(self.agent, server_config)

kssrag/models/local_llms.py

@@ -0,0 +1,30 @@
+"""
+Local LLM implementations for offline usage
+"""
+from typing import List, Dict, Any
+from ..utils.helpers import logger
+
+class LocalLLM:
+    """Base class for local LLM implementations"""
+    
+    def predict(self, messages: List[Dict[str, str]]) -> str:
+        """Generate a response using a local model"""
+        raise NotImplementedError("Subclasses must implement this method")
+
+class MockLLM(LocalLLM):
+    """Mock LLM for testing without API calls"""
+    
+    def predict(self, messages: List[Dict[str, str]]) -> str:
+        """Generate a mock response for testing"""
+        logger.info("Using mock LLM for response generation")
+        
+        # Extract the last user message
+        user_message = ""
+        for msg in reversed(messages):
+            if msg["role"] == "user":
+                user_message = msg["content"]
+                break
+        
+        return f"This is a mock response to: {user_message}"
+
+# TO-DO - Add more local LLM implementations as needed - I hope i do not forget, someone if you are seeing this, the date i added this was 8th of september 2025, how long has it been? DM me on WhatsApp +2349019549473.

kssrag/models/openrouter.py

@@ -0,0 +1,85 @@
+import requests
+import json
+from typing import List, Dict, Any, Optional
+from ..utils.helpers import logger
+from ..config import config
+
+class OpenRouterLLM:
+    """OpenRouter LLM interface with fallback models"""
+    
+    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None, 
+                 fallback_models: Optional[List[str]] = None):
+        self.api_key = api_key or config.OPENROUTER_API_KEY
+        self.model = model or config.DEFAULT_MODEL
+        self.fallback_models = fallback_models or config.FALLBACK_MODELS
+        self.base_url = "https://openrouter.ai/api/v1/chat/completions"
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://github.com/Ksschkw/kssrag",
+            "X-Title": "KSS RAG Agent"
+        }
+    
+    def predict(self, messages: List[Dict[str, str]]) -> str:
+        """Generate a response using OpenRouter's API with fallbacks"""
+        logger.info(f"Attempting to generate response with {len(messages)} messages")
+        
+        for model in [self.model] + self.fallback_models:
+            payload = {
+                "model": model,
+                "messages": messages,
+                "temperature": 0.7,
+                "max_tokens": 1024,
+                "top_p": 1,
+                "stop": None,
+                "stream": False
+            }
+            
+            try:
+                logger.info(f"Trying model: {model}")
+                response = requests.post(
+                    self.base_url,
+                    headers=self.headers,
+                    json=payload,
+                    timeout=15
+                )
+                
+                # Check for HTTP errors
+                response.raise_for_status()
+                
+                # Parse JSON response
+                response_data = response.json()
+                
+                # Validate response structure
+                if ("choices" not in response_data or 
+                    len(response_data["choices"]) == 0 or
+                    "message" not in response_data["choices"][0] or
+                    "content" not in response_data["choices"][0]["message"]):
+                    
+                    logger.warning(f"Invalid response format from {model}: {response_data}")
+                    continue
+                
+                content = response_data["choices"][0]["message"]["content"]
+                logger.info(f"Successfully used model: {model}")
+                return content
+                
+            except requests.exceptions.Timeout:
+                logger.warning(f"Model {model} timed out")
+                continue
+            except requests.exceptions.RequestException as e:
+                logger.warning(f"Request error with model {model}: {str(e)}")
+                if hasattr(e, 'response') and e.response is not None:
+                    try:
+                        error_data = e.response.json()
+                        logger.warning(f"Error response: {error_data}")
+                    except:
+                        logger.warning(f"Error response text: {e.response.text}")
+                continue
+            except Exception as e:
+                logger.warning(f"Unexpected error with model {model}: {str(e)}")
+                continue
+        
+        # If all models fail, return a friendly error message
+        error_msg = "I'm having trouble connecting to the knowledge service right now. Please try again in a moment."
+        logger.error("All model fallbacks failed to respond")
+        return error_msg