kssrag 0.1.0__py3-none-any.whl

kssrag/__init__.py

@@ -0,0 +1,66 @@
+"""
+KSS RAG - A flexible Retrieval-Augmented Generation framework by Ksschkw
+"""
+from .kssrag import KSSRAG
+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 .utils.helpers import logger, validate_config
+from .config import Config, VectorStoreType, ChunkerType, RetrieverType
+from .server import create_app, ServerConfig
+from .cli import main
+
+__version__ = "0.1.0"
+__author__ = "Ksschkw"
+__license__ = "MIT"
+
+# Your footprint - include your GitHub username and a signature
+__signature__ = "Built with HATE by Ksschkw (github.com/Ksschkw)"
+
+# Export the main classes for easy access
+__all__ = [
+    'KSSRAG',
+    'TextChunker',
+    'JSONChunker',
+    'PDFChunker',
+    'BM25VectorStore',
+    'FAISSVectorStore',
+    'TFIDFVectorStore',
+    'HybridVectorStore',
+    'HybridOfflineVectorStore',
+    'SimpleRetriever',
+    'HybridRetriever',
+    'RAGAgent',
+    'OpenRouterLLM',
+    'load_document',
+    'load_json_documents',
+    'Config',
+    'VectorStoreType',
+    'ChunkerType',
+    'RetrieverType',
+    'ServerConfig',
+    'create_app',
+    'main',
+    'logger',
+    'validate_config'
+]
+
+# Initialize configuration validation on import
+validate_config()
+
+import platform
+from pathlib import Path
+import os
+
+# Windows-specific cache directory handling
+if platform.system() == "Windows":
+    # Use local app data directory instead of home directory
+    cache_base = os.getenv('LOCALAPPDATA', os.path.expanduser('~'))
+    config.CACHE_DIR = os.path.join(cache_base, '.kssrag', 'cache')
+else:
+    config.CACHE_DIR = os.path.join(os.path.expanduser('~'), '.cache', 'kssrag')
+
+os.makedirs(config.CACHE_DIR, exist_ok=True)

kssrag/cli.py

@@ -0,0 +1,142 @@
+import argparse
+import sys
+from .utils.document_loaders import load_document, load_json_documents
+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 .config import config
+from .utils.helpers import logger, validate_config
+
+def main():
+    """Command-line interface for KSS RAG"""
+    parser = argparse.ArgumentParser(description="KSS RAG - Retrieval-Augmented Generation Framework")
+    subparsers = parser.add_subparsers(dest="command", help="Command to execute")
+    
+    # Query command
+    query_parser = subparsers.add_parser("query", help="Query the RAG system")
+    query_parser.add_argument("--file", type=str, required=True, help="Path to document file")
+    query_parser.add_argument("--query", type=str, required=True, help="Query to ask")
+    query_parser.add_argument("--format", type=str, default="text", choices=["text", "json", "pdf"], 
+                             help="Document format")
+    query_parser.add_argument("--vector-store", type=str, default=config.VECTOR_STORE_TYPE,
+                             choices=["bm25", "faiss", "tfidf", "hybrid_online", "hybrid_offline"], 
+                             help="Vector store type")
+    query_parser.add_argument("--top-k", type=int, default=config.TOP_K, help="Number of results to retrieve")
+    
+    # Server command
+    server_parser = subparsers.add_parser("server", help="Start the RAG API server")
+    server_parser.add_argument("--file", type=str, required=True, help="Path to document file")
+    server_parser.add_argument("--format", type=str, default="text", choices=["text", "json", "pdf"], 
+                              help="Document format")
+    server_parser.add_argument("--vector-store", type=str, default=config.VECTOR_STORE_TYPE,
+                              choices=["bm25", "faiss", "tfidf", "hybrid_online", "hybrid_offline"], 
+                              help="Vector store type")
+    server_parser.add_argument("--port", type=int, default=config.SERVER_PORT, help="Port to run server on")
+    server_parser.add_argument("--host", type=str, default=config.SERVER_HOST, help="Host to run server on")
+    
+    args = parser.parse_args()
+    
+    # Validate config
+    validate_config()
+    
+    if args.command == "query":
+        # Load and process document
+        if args.format == "text":
+            content = load_document(args.file)
+            chunker = TextChunker(chunk_size=config.CHUNK_SIZE, overlap=config.CHUNK_OVERLAP)
+            documents = chunker.chunk(content, {"source": args.file})
+        elif args.format == "json":
+            data = load_json_documents(args.file)
+            chunker = JSONChunker()
+            documents = chunker.chunk(data)
+        elif args.format == "pdf":
+            chunker = PDFChunker(chunk_size=config.CHUNK_SIZE, overlap=config.CHUNK_OVERLAP)
+            documents = chunker.chunk_pdf(args.file, {"source": args.file})
+        else:
+            logger.error(f"Unsupported format: {args.format}")
+            return 1
+        
+        # Create vector store
+        if args.vector_store == "bm25":
+            vector_store = BM25VectorStore()
+        elif args.vector_store == "faiss":
+            vector_store = FAISSVectorStore()
+        elif args.vector_store == "tfidf":
+            vector_store = TFIDFVectorStore()
+        elif args.vector_store == "hybrid_online":
+            vector_store = HybridVectorStore()
+        elif args.vector_store == "hybrid_offline":
+            vector_store = HybridOfflineVectorStore()
+        else:
+            logger.error(f"Unsupported vector store: {args.vector_store}")
+            return 1
+        
+        vector_store.add_documents(documents)
+        
+        # Create retriever and agent
+        retriever = SimpleRetriever(vector_store)
+        llm = OpenRouterLLM()
+        agent = RAGAgent(retriever, llm)
+        
+        # Query and print response
+        response = agent.query(args.query, top_k=args.top_k)
+        print(f"Query: {args.query}")
+        print(f"Response: {response}")
+        
+    elif args.command == "server":
+        # Load and process document
+        if args.format == "text":
+            content = load_document(args.file)
+            chunker = TextChunker(chunk_size=config.CHUNK_SIZE, overlap=config.CHUNK_OVERLAP)
+            documents = chunker.chunk(content, {"source": args.file})
+        elif args.format == "json":
+            data = load_json_documents(args.file)
+            chunker = JSONChunker()
+            documents = chunker.chunk(data)
+        elif args.format == "pdf":
+            chunker = PDFChunker(chunk_size=config.CHUNK_SIZE, overlap=config.CHUNK_OVERLAP)
+            documents = chunker.chunk_pdf(args.file, {"source": args.file})
+        else:
+            logger.error(f"Unsupported format: {args.format}")
+            return 1
+        
+        # Create vector store
+        if args.vector_store == "bm25":
+            vector_store = BM25VectorStore()
+        elif args.vector_store == "faiss":
+            vector_store = FAISSVectorStore()
+        elif args.vector_store == "tfidf":
+            vector_store = TFIDFVectorStore()
+        elif args.vector_store == "hybrid_online":
+            vector_store = HybridVectorStore()
+        elif args.vector_store == "hybrid_offline":
+            vector_store = HybridOfflineVectorStore()
+        else:
+            logger.error(f"Unsupported vector store: {args.vector_store}")
+            return 1
+        
+        vector_store.add_documents(documents)
+        
+        # Create retriever and agent
+        retriever = SimpleRetriever(vector_store)
+        llm = OpenRouterLLM()
+        agent = RAGAgent(retriever, llm)
+        
+        # Create and run server
+        from .server import create_app
+        import uvicorn
+        
+        app, server_config = create_app(agent)
+        logger.info(f"Starting server on {args.host}:{args.port}")
+        uvicorn.run(app, host=args.host, port=args.port)
+        
+    else:
+        parser.print_help()
+        return 1
+    
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

kssrag/config.py

@@ -0,0 +1,193 @@
+import os
+from dotenv import load_dotenv
+from typing import List, Optional, Dict, Any
+from pydantic import Field, validator
+from pydantic_settings import BaseSettings
+from enum import Enum
+
+load_dotenv()
+
+class VectorStoreType(str, Enum):
+    BM25 = "bm25"
+    FAISS = "faiss"
+    TFIDF = "tfidf"
+    HYBRID_ONLINE = "hybrid_online"
+    HYBRID_OFFLINE = "hybrid_offline"
+    CUSTOM = "custom"
+
+class ChunkerType(str, Enum):
+    TEXT = "text"
+    JSON = "json"
+    PDF = "pdf"
+    CUSTOM = "custom"
+
+class RetrieverType(str, Enum):
+    SIMPLE = "simple"
+    HYBRID = "hybrid"
+    CUSTOM = "custom"
+
+class Config(BaseSettings):
+    """Configuration settings for KSS RAG with extensive customization options"""
+    
+    # OpenRouter settings
+    OPENROUTER_API_KEY: str = Field(
+        default=os.getenv("OPENROUTER_API_KEY", ""),
+        description="Your OpenRouter API key for accessing LLMs"
+    )
+    
+    DEFAULT_MODEL: str = Field(
+        default=os.getenv("DEFAULT_MODEL", "deepseek/deepseek-chat-v3.1:free"),
+        description="Default model to use for LLM responses"
+    )
+    
+    FALLBACK_MODELS: List[str] = Field(
+        default=os.getenv("FALLBACK_MODELS", "deepseek/deepseek-r1-0528:free,deepseek/deepseek-chat,deepseek/deepseek-r1:free").split(","),
+        description="List of fallback models to try if the default model fails"
+    )
+    
+    # Chunking settings
+    CHUNK_SIZE: int = Field(
+        default=int(os.getenv("CHUNK_SIZE", 500)),
+        ge=100,
+        le=2000,
+        description="Size of text chunks in characters"
+    )
+    
+    CHUNK_OVERLAP: int = Field(
+        default=int(os.getenv("CHUNK_OVERLAP", 50)),
+        ge=0,
+        le=500,
+        description="Overlap between chunks in characters"
+    )
+    
+    CHUNKER_TYPE: ChunkerType = Field(
+        default=os.getenv("CHUNKER_TYPE", ChunkerType.TEXT),
+        description="Type of chunker to use"
+    )
+    
+    # Vector store settings
+    VECTOR_STORE_TYPE: VectorStoreType = Field(
+        default=os.getenv("VECTOR_STORE_TYPE", VectorStoreType.HYBRID_OFFLINE),
+        description="Type of vector store to use"
+    )
+    
+    FAISS_MODEL_NAME: str = Field(
+        default=os.getenv("FAISS_MODEL_NAME", "sentence-transformers/all-MiniLM-L6-v2"),
+        description="SentenceTransformer model name for FAISS embeddings"
+    )
+    
+    # Retrieval settings
+    RETRIEVER_TYPE: RetrieverType = Field(
+        default=os.getenv("RETRIEVER_TYPE", RetrieverType.SIMPLE),
+        description="Type of retriever to use"
+    )
+    
+    TOP_K: int = Field(
+        default=int(os.getenv("TOP_K", 5)),
+        ge=1,
+        le=20,
+        description="Number of results to retrieve"
+    )
+    
+    FUZZY_MATCH_THRESHOLD: int = Field(
+        default=int(os.getenv("FUZZY_MATCH_THRESHOLD", 80)),
+        ge=0,
+        le=100,
+        description="Threshold for fuzzy matching (0-100)"
+    )
+    
+    # Performance settings
+    BATCH_SIZE: int = Field(
+        default=int(os.getenv("BATCH_SIZE", 64)),
+        ge=1,
+        le=256,
+        description="Batch size for processing documents"
+    )
+    
+    MAX_DOCS_FOR_TESTING: Optional[int] = Field(
+        default=os.getenv("MAX_DOCS_FOR_TESTING"),
+        description="Limit documents for testing (None for all)"
+    )
+    
+    # Server settings
+    SERVER_HOST: str = Field(
+        default=os.getenv("SERVER_HOST", "localhost"),
+        description="Host to run the server on"
+    )
+    
+    SERVER_PORT: int = Field(
+        default=int(os.getenv("SERVER_PORT", 8000)),
+        ge=1024,
+        le=65535,
+        description="Port to run the server on"
+    )
+    
+    CORS_ORIGINS: List[str] = Field(
+        default=os.getenv("CORS_ORIGINS", "*").split(","),
+        description="List of CORS origins"
+    )
+    
+    CORS_ALLOW_CREDENTIALS: bool = Field(
+        default=os.getenv("CORS_ALLOW_CREDENTIALS", "True").lower() == "true",
+        description="Whether to allow CORS credentials"
+    )
+    
+    CORS_ALLOW_METHODS: List[str] = Field(
+        default=os.getenv("CORS_ALLOW_METHODS", "GET,POST,PUT,DELETE,OPTIONS").split(","),
+        description="List of allowed CORS methods"
+    )
+    
+    CORS_ALLOW_HEADERS: List[str] = Field(
+        default=os.getenv("CORS_ALLOW_HEADERS", "Content-Type,Authorization").split(","),
+        description="List of allowed CORS headers"
+    )
+    
+    # Advanced settings
+    ENABLE_CACHE: bool = Field(
+        default=os.getenv("ENABLE_CACHE", "True").lower() == "true",
+        description="Whether to enable caching for vector stores"
+    )
+    
+    CACHE_DIR: str = Field(
+        default=os.getenv("CACHE_DIR", ".cache"),
+        description="Directory to store cache files"
+    )
+    
+    LOG_LEVEL: str = Field(
+        default=os.getenv("LOG_LEVEL", "INFO"),
+        description="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)"
+    )
+    
+    # Custom components (for advanced users)
+    CUSTOM_CHUNKER: Optional[str] = Field(
+        default=os.getenv("CUSTOM_CHUNKER"),
+        description="Import path to a custom chunker class"
+    )
+    
+    CUSTOM_VECTOR_STORE: Optional[str] = Field(
+        default=os.getenv("CUSTOM_VECTOR_STORE"),
+        description="Import path to a custom vector store class"
+    )
+    
+    CUSTOM_RETRIEVER: Optional[str] = Field(
+        default=os.getenv("CUSTOM_RETRIEVER"),
+        description="Import path to a custom retriever class"
+    )
+    
+    CUSTOM_LLM: Optional[str] = Field(
+        default=os.getenv("CUSTOM_LLM"),
+        description="Import path to a custom LLM class"
+    )
+    
+    class Config:
+        env_file = ".env"
+        case_sensitive = False
+        use_enum_values = True
+        
+    @validator('FALLBACK_MODELS', 'CORS_ORIGINS', 'CORS_ALLOW_METHODS', 'CORS_ALLOW_HEADERS', pre=True)
+    def split_string(cls, v):
+        if isinstance(v, str):
+            return [item.strip() for item in v.split(',')]
+        return v
+
+config = Config()

kssrag/core/agents.py

@@ -0,0 +1,68 @@
+from typing import List, Dict, Any, Optional
+from ..utils.helpers import logger
+
+class RAGAgent:
+    """RAG agent implementation"""
+    
+    def __init__(self, retriever, llm, system_prompt: Optional[str] = None, 
+                 conversation_history: Optional[List[Dict[str, str]]] = None):
+        self.retriever = retriever
+        self.llm = llm
+        self.conversation = conversation_history or []
+        self.system_prompt = system_prompt or """You are a helpful AI assistant. Use the following context to answer the user's question. 
+        If you don't know the answer based on the context, say so."""
+        
+        # Initialize with system message if not already present
+        if not any(msg.get("role") == "system" for msg in self.conversation):
+            self.add_message("system", self.system_prompt)
+    
+    def add_message(self, role: str, content: str):
+        """Add a message to the conversation history"""
+        self.conversation.append({"role": role, "content": content})
+        
+        # Keep conversation manageable (last 10 messages)
+        if len(self.conversation) > 10:
+            # Always keep the system message
+            system_msg = next((msg for msg in self.conversation if msg["role"] == "system"), None)
+            other_msgs = [msg for msg in self.conversation if msg["role"] != "system"]
+            
+            # Keep the most recent messages
+            self.conversation = [system_msg] + other_msgs[-9:] if system_msg else other_msgs[-10:]
+    
+    def query(self, question: str, top_k: int = 5, include_context: bool = True) -> str:
+        """Process a query and return a response"""
+        try:
+            # Retrieve relevant context
+            context_docs = self.retriever.retrieve(question, top_k)
+            
+            if not context_docs and include_context:
+                logger.warning(f"No context found for query: {question}")
+                return "I couldn't find relevant information to answer your question."
+            
+            # Format context
+            context = ""
+            if include_context and context_docs:
+                context = "Relevant information:\n"
+                for i, doc in enumerate(context_docs, 1):
+                    context += f"\n--- Document {i} ---\n{doc['content']}\n"
+            
+            # Add user query with context
+            user_message = f"{context}\n\nQuestion: {question}" if context else question
+            self.add_message("user", user_message)
+            
+            # Generate response
+            response = self.llm.predict(self.conversation)
+            
+            # Add assistant response to conversation
+            self.add_message("assistant", response)
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"Error processing query: {str(e)}")
+            return "I encountered an issue processing your query. Please try again."
+    
+    def clear_conversation(self):
+        """Clear conversation history except system message"""
+        system_msg = next((msg for msg in self.conversation if msg["role"] == "system"), None)
+        self.conversation = [system_msg] if system_msg else []

kssrag/core/chunkers.py

@@ -0,0 +1,100 @@
+import json
+import re
+from typing import List, Dict, Any, Optional
+import pypdf
+from ..utils.helpers import logger
+
+class BaseChunker:
+    """Base class for document chunkers"""
+    
+    def __init__(self, chunk_size: int = 500, overlap: int = 50):
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+    
+    def chunk(self, content: Any, metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+class TextChunker(BaseChunker):
+    """Chunker for plain text documents"""
+    
+    def chunk(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        """Split text into chunks with overlap"""
+        if metadata is None:
+            metadata = {}
+        
+        chunks = []
+        start = 0
+        content_length = len(content)
+        
+        while start < content_length:
+            end = start + self.chunk_size
+            if end > content_length:
+                end = content_length
+            
+            chunk = content[start:end]
+            chunk_metadata = metadata.copy()
+            chunk_metadata["chunk_id"] = len(chunks)
+            
+            chunks.append({
+                "content": chunk,
+                "metadata": chunk_metadata
+            })
+            
+            start += self.chunk_size - self.overlap
+        
+        logger.info(f"Created {len(chunks)} chunks from text")
+        return chunks
+
+class JSONChunker(BaseChunker):
+    """Chunker for JSON documents (like your drug data)"""
+    
+    def chunk(self, data: List[Dict[str, Any]], metadata_field: str = "name") -> List[Dict[str, Any]]:
+        """Create chunks from JSON data"""
+        chunks = []
+        
+        for item in data:
+            if metadata_field not in item:
+                continue
+                
+            # Create a comprehensive text representation
+            item_text = f"Item Name: {item.get(metadata_field, 'N/A')}\n"
+            
+            for key, value in item.items():
+                if key != metadata_field and value:
+                    if isinstance(value, str):
+                        item_text += f"{key.replace('_', ' ').title()}: {value}\n"
+                    elif isinstance(value, list):
+                        item_text += f"{key.replace('_', ' ').title()}: {', '.join(value)}\n"
+            
+            chunks.append({
+                "content": item_text,
+                "metadata": {
+                    "name": item[metadata_field],
+                    "source": item.get('url', 'N/A')
+                }
+            })
+        
+        logger.info(f"Created {len(chunks)} chunks from JSON data")
+        return chunks
+
+class PDFChunker(TextChunker):
+    """Chunker for PDF documents"""
+    
+    def extract_text(self, pdf_path: str) -> str:
+        """Extract text from PDF file"""
+        text = ""
+        try:
+            with open(pdf_path, 'rb') as f:
+                reader = pypdf.PdfReader(f)
+                for page in reader.pages:
+                    text += page.extract_text() + "\n"
+        except Exception as e:
+            logger.error(f"Failed to extract text from PDF: {str(e)}")
+            raise
+        
+        return text
+    
+    def chunk_pdf(self, pdf_path: str, metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        """Extract text from PDF and chunk it"""
+        text = self.extract_text(pdf_path)
+        return self.chunk(text, metadata)

kssrag/core/retrievers.py

@@ -0,0 +1,74 @@
+from typing import List, Dict, Any, Optional
+from ..utils.helpers import logger
+
+class BaseRetriever:
+    """Base class for retrievers"""
+    
+    def __init__(self, vector_store):
+        self.vector_store = vector_store
+    
+    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")
+
+class SimpleRetriever(BaseRetriever):
+    """Simple retriever using only vector store"""
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        return self.vector_store.retrieve(query, top_k)
+
+class HybridRetriever(BaseRetriever):
+    """Hybrid retriever with fuzzy matching for specific entities"""
+    
+    def __init__(self, vector_store, entity_names: Optional[List[str]] = None):
+        super().__init__(vector_store)
+        self.entity_names = entity_names or []
+    
+    def extract_entities(self, query: str) -> List[str]:
+        """Extract entities from query using fuzzy matching"""
+        from rapidfuzz import process, fuzz
+        
+        extracted_entities = []
+        query_lower = query.lower()
+        
+        # Check for exact matches first
+        for entity in self.entity_names:
+            if entity in query_lower:
+                extracted_entities.append(entity)
+        
+        # Use fuzzy matching for partial matches
+        if not extracted_entities and self.entity_names:
+            matches = process.extract(query, self.entity_names, scorer=fuzz.partial_ratio, limit=5)
+            extracted_entities = [match[0] for match in matches if match[1] > 80]
+        
+        return extracted_entities
+    
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        # First get results from vector store
+        results = self.vector_store.retrieve(query, top_k * 2)
+        
+        # If we have entity names, boost documents that mention extracted entities
+        if self.entity_names:
+            extracted_entities = self.extract_entities(query)
+            
+            if extracted_entities:
+                # Boost scores for documents containing extracted entities
+                scored_results = []
+                
+                for doc in results:
+                    score = 1.0
+                    content_lower = doc["content"].lower()
+                    
+                    # Boost score if any extracted entity is mentioned
+                    for entity in extracted_entities:
+                        if entity in content_lower:
+                            score += 0.5
+                            break
+                    
+                    scored_results.append((doc, score))
+                
+                # Sort by score and return top_k
+                scored_results.sort(key=lambda x: x[1], reverse=True)
+                return [doc for doc, _ in scored_results[:top_k]]
+        
+        return results[:top_k]