raglite-toolkit 1.0.0__py3-none-any.whl

raglite/__init__.py

@@ -0,0 +1,114 @@
+from .constants import PACKAGE_NAME, PACKAGE_VERSION
+
+VERSION = PACKAGE_VERSION
+
+from .errors import (
+    RagLiteError,
+    UnsupportedFileTypeError,
+    FileNotIndexedError,
+    LoaderError,
+    ChunkingError,
+    EmbeddingError,
+    VectorDBError,
+    LLMError,
+    ConfigError,
+)
+
+from .types import (
+    LLMProviderName,
+    EmbeddingProviderName,
+    LLMProviderConfig,
+    EmbeddingProviderConfig,
+    ChunkMetadata,
+    StoredChunk,
+    SearchResult,
+    AnswerResult,
+    IndexMetadata,
+)
+
+from .config import DocumentOptions, resolve_config
+
+from .utils.logger import create_logger
+
+from .loaders import (
+    BaseLoader,
+    TxtLoader,
+    MarkdownLoader,
+    JsonLoader,
+    PdfLoader,
+    DocxLoader,
+    get_loader,
+)
+
+from .chunking import BaseChunker, RecursiveChunker
+
+from .vectordb import MemoryVectorStore
+
+from .retrieval import Retriever
+
+from .embeddings import (
+    DEFAULT_EMBEDDING_MODELS,
+    LocalEmbedder,
+    RemoteEmbedder,
+    create_embedder,
+)
+
+from .llm import (
+    DEFAULT_LLM_MODELS,
+    build_system_prompt,
+    build_user_prompt,
+    create_llm,
+    generate_answer,
+    stream_answer,
+)
+
+from .core.document import Document
+
+__all__ = [
+    "PACKAGE_NAME",
+    "PACKAGE_VERSION",
+    "VERSION",
+    "RagLiteError",
+    "UnsupportedFileTypeError",
+    "FileNotIndexedError",
+    "LoaderError",
+    "ChunkingError",
+    "EmbeddingError",
+    "VectorDBError",
+    "LLMError",
+    "ConfigError",
+    "LLMProviderName",
+    "EmbeddingProviderName",
+    "LLMProviderConfig",
+    "EmbeddingProviderConfig",
+    "ChunkMetadata",
+    "StoredChunk",
+    "SearchResult",
+    "AnswerResult",
+    "IndexMetadata",
+    "DocumentOptions",
+    "resolve_config",
+    "create_logger",
+    "BaseLoader",
+    "TxtLoader",
+    "MarkdownLoader",
+    "JsonLoader",
+    "PdfLoader",
+    "DocxLoader",
+    "get_loader",
+    "BaseChunker",
+    "RecursiveChunker",
+    "MemoryVectorStore",
+    "Retriever",
+    "DEFAULT_EMBEDDING_MODELS",
+    "LocalEmbedder",
+    "RemoteEmbedder",
+    "create_embedder",
+    "DEFAULT_LLM_MODELS",
+    "build_system_prompt",
+    "build_user_prompt",
+    "create_llm",
+    "generate_answer",
+    "stream_answer",
+    "Document",
+]

raglite/api/__init__.py

@@ -0,0 +1,4 @@
+from .schemas import SearchRequest, AskRequest
+from .server import build_app, create_server, ServerHandle
+
+__all__ = ["SearchRequest", "AskRequest", "build_app", "create_server", "ServerHandle"]

raglite/api/schemas.py

@@ -0,0 +1,26 @@
+from typing import Optional
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class SearchRequest(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    query: str = Field(..., min_length=1)
+    topK: Optional[int] = Field(default=None, alias="topK", gt=0, le=50)
+    scoreThreshold: Optional[float] = Field(
+        default=None, alias="scoreThreshold", ge=-1.0, le=1.0
+    )
+
+
+class AskRequest(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    question: str = Field(..., min_length=1)
+    topK: Optional[int] = Field(default=None, alias="topK", gt=0, le=50)
+    scoreThreshold: Optional[float] = Field(
+        default=None, alias="scoreThreshold", ge=-1.0, le=1.0
+    )
+    includeCitations: Optional[bool] = Field(
+        default=None, alias="includeCitations"
+    )
+    stream: Optional[bool] = None

raglite/api/server.py

@@ -0,0 +1,198 @@
+import time
+import threading
+from typing import Optional, Dict, Any, Union
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request, Depends
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import ValidationError
+
+from ..constants import PACKAGE_VERSION, DEFAULT_HOST, DEFAULT_PORT
+from ..errors import RagLiteError
+from ..embeddings.models import DEFAULT_EMBEDDING_MODELS
+from .schemas import SearchRequest, AskRequest
+
+security = HTTPBearer(auto_error=False)
+
+
+class ServerHandle:
+    def __init__(
+        self,
+        server_thread: threading.Thread,
+        uvicorn_server: uvicorn.Server,
+        url: str,
+    ):
+        self.server_thread = server_thread
+        self.uvicorn_server = uvicorn_server
+        self.url = url
+
+    def close(self) -> None:
+        """Stop the background API server."""
+        self.uvicorn_server.should_exit = True
+        self.server_thread.join()
+
+
+def build_app(document: Any, options: Dict[str, Any]) -> FastAPI:
+    app = FastAPI(title="RAGLite REST API")
+
+    bearer_token = options.get("bearerToken") or options.get("bearer_token")
+
+    if options.get("requestLogging"):
+
+        @app.middleware("http")
+        async def log_requests(request: Request, call_next):
+            start = time.time()
+            response = await call_next(request)
+            duration = int((time.time() - start) * 1000)
+            print(
+                f"[raglite] {request.method} {request.url.path} -> {response.status_code} ({duration}ms)"
+            )
+            return response
+
+    def authorize(
+        credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
+    ):
+        if bearer_token:
+            if not credentials or credentials.credentials != bearer_token:
+                raise HTTPException(status_code=401, detail="Unauthorized")
+        return True
+
+    @app.exception_handler(ValidationError)
+    async def validation_exception_handler(request: Request, exc: ValidationError):
+        return JSONResponse(
+            status_code=400,
+            content={"error": "ValidationError", "details": exc.errors()},
+        )
+
+    @app.exception_handler(RagLiteError)
+    async def raglite_exception_handler(request: Request, exc: RagLiteError):
+        return JSONResponse(
+            status_code=400,
+            content={"error": exc.__class__.__name__, "message": str(exc)},
+        )
+
+    @app.exception_handler(Exception)
+    async def general_exception_handler(request: Request, exc: Exception):
+        print("[raglite] unhandled error", exc)
+        return JSONResponse(
+            status_code=500,
+            content={"error": "InternalServerError"},
+        )
+
+    @app.get("/health")
+    async def health():
+        return {
+            "status": "ok",
+            "version": PACKAGE_VERSION,
+            "chunks": document.chunk_count,
+            "namespace": document.store_namespace,
+        }
+
+    dependencies = [Depends(authorize)] if bearer_token else []
+
+    @app.get("/info", dependencies=dependencies)
+    async def info():
+        cfg = document.resolved_config
+        llm_config = options.get("llm") or cfg.llm
+        return {
+            "version": PACKAGE_VERSION,
+            "chunkSize": cfg.chunkSize,
+            "overlap": cfg.overlap,
+            "topK": cfg.topK,
+            "embeddings": {
+                "provider": cfg.embeddings.provider,
+                "model": cfg.embeddings.model
+                or DEFAULT_EMBEDDING_MODELS.get(cfg.embeddings.provider),
+            },
+            "llmProvider": llm_config.provider if llm_config else None,
+            "chunks": document.chunk_count,
+        }
+
+    @app.post("/search", dependencies=dependencies)
+    async def search(req: SearchRequest):
+        results = document.search(
+            req.query,
+            top_k=req.topK,
+            score_threshold=req.scoreThreshold,
+        )
+        return {
+            "results": [r.model_dump(by_alias=True) for r in results]
+        }
+
+    ask_provider = options.get("llm")
+
+    if ask_provider:
+
+        @app.post("/ask", dependencies=dependencies)
+        async def ask(req: AskRequest):
+            ask_options = {
+                "llm": ask_provider,
+                "topK": req.topK,
+                "scoreThreshold": req.scoreThreshold,
+                "includeCitations": req.includeCitations,
+            }
+
+            if req.stream:
+
+                def event_generator():
+                    for chunk in document.ask_stream(req.question, ask_options):
+                        yield chunk
+
+                return StreamingResponse(
+                    event_generator(), media_type="text/plain"
+                )
+
+            answer = document.ask(req.question, ask_options)
+            return answer.model_dump(by_alias=True)
+
+    else:
+
+        @app.post("/ask", dependencies=dependencies)
+        async def ask(req: AskRequest):
+            return JSONResponse(
+                status_code=503,
+                content={
+                    "error": "AskDisabled",
+                    "message": "No LLM provider configured on the server.",
+                },
+            )
+
+    return app
+
+
+def create_server(
+    document: Any,
+    options: Dict[str, Any],
+    *,
+    host: Optional[str] = None,
+    port: Optional[int] = None,
+    bearer_token: Optional[str] = None,
+) -> ServerHandle:
+    srv_host = host or options.get("host") or DEFAULT_HOST
+    srv_port = port or options.get("port") or DEFAULT_PORT
+
+    # Create server options clone
+    server_options = dict(options)
+    if bearer_token:
+        server_options["bearerToken"] = bearer_token
+
+    app = build_app(document, server_options)
+
+    config = uvicorn.Config(
+        app=app, host=srv_host, port=srv_port, log_level="warning"
+    )
+    server = uvicorn.Server(config)
+
+    thread = threading.Thread(target=server.run)
+    thread.daemon = True
+    thread.start()
+
+    url = f"http://{srv_host}:{srv_port}"
+
+    # Wait up to 2 seconds for server to boot
+    retries = 20
+    while retries > 0 and not server.started:
+        time.sleep(0.1)
+        retries -= 1
+
+    return ServerHandle(server_thread=thread, uvicorn_server=server, url=url)

raglite/chunking/__init__.py

@@ -0,0 +1,4 @@
+from .base import BaseChunker
+from .recursive import RecursiveChunker
+
+__all__ = ["BaseChunker", "RecursiveChunker"]

raglite/chunking/base.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+
+class BaseChunker(ABC):
+    def __init__(self, chunk_size: int, overlap: int):
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    @abstractmethod
+    def split(self, text: str) -> List[str]:
+        """Split text into chunks."""
+        pass

raglite/chunking/recursive.py

@@ -0,0 +1,34 @@
+import re
+from typing import List
+from .base import BaseChunker
+from ..errors import ChunkingError
+
+
+class RecursiveChunker(BaseChunker):
+    def split(self, text: str) -> List[str]:
+        if not text.strip():
+            return []
+        if self.overlap >= self.chunk_size:
+            raise ChunkingError(
+                f"overlap ({self.overlap}) must be smaller than chunkSize ({self.chunk_size})"
+            )
+
+        words = [w for w in re.split(r"\s+", text) if w]
+        if len(words) <= self.chunk_size:
+            return [" ".join(words)]
+
+        step = self.chunk_size - self.overlap
+        chunks: List[str] = []
+
+        start = 0
+        while start < len(words):
+            end = start + self.chunk_size
+            slice_words = words[start:end]
+            if not slice_words:
+                break
+            chunks.append(" ".join(slice_words))
+            if end >= len(words):
+                break
+            start += step
+
+        return chunks

raglite/cli.py

@@ -0,0 +1,212 @@
+import sys
+import json
+import argparse
+import time
+from typing import Optional
+
+from .constants import PACKAGE_VERSION
+from .core.document import Document
+
+HELP = f"""raglite v{PACKAGE_VERSION}
+
+Usage:
+  raglite index <file>   [--chunk-size N] [--overlap N] [--embed-provider P] [--embed-model M] [--embed-key K] [--rebuild]
+  raglite search <file> "query"   [--top-k N]
+  raglite ask <file> "question"   --llm-provider P [--llm-model M] [--llm-key K] [--stream]
+  raglite serve <file>            --llm-provider P [--llm-key K] [--host H] [--port N] [--token T]
+  raglite --help
+  raglite --version
+
+Providers:
+  LLM:        openai, anthropic, google, mistral, cohere, groq, xai, ollama
+  Embeddings: openai, google, mistral, cohere, voyage, ollama, local
+"""
+
+
+def parse_common_embedding(args_dict: dict) -> dict:
+    provider = args_dict.get("embed_provider") or "local"
+    config = {"provider": provider}
+    if args_dict.get("embed_model"):
+        config["model"] = args_dict["embed_model"]
+    if args_dict.get("embed_key"):
+        config["apiKey"] = args_dict["embed_key"]
+    return config
+
+
+def parse_llm(args_dict: dict) -> Optional[dict]:
+    provider = args_dict.get("llm_provider")
+    if not provider:
+        return None
+    config = {"provider": provider}
+    if args_dict.get("llm_model"):
+        config["model"] = args_dict["llm_model"]
+    if args_dict.get("llm_key"):
+        config["apiKey"] = args_dict["llm_key"]
+    return config
+
+
+def run_index(args):
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("file")
+    parser.add_argument("--chunk-size", type=int)
+    parser.add_argument("--overlap", type=int)
+    parser.add_argument("--embed-provider")
+    parser.add_argument("--embed-model")
+    parser.add_argument("--embed-key")
+    parser.add_argument("--rebuild", action="store_true")
+
+    parsed = parser.parse_args(args)
+
+    embeddings = parse_common_embedding(vars(parsed))
+    doc = Document(parsed.file, {"embeddings": embeddings})
+
+    build_opts = {}
+    if parsed.chunk_size is not None:
+        build_opts["chunkSize"] = parsed.chunk_size
+    if parsed.overlap is not None:
+        build_opts["overlap"] = parsed.overlap
+    if parsed.rebuild:
+        build_opts["rebuild"] = True
+
+    result = doc.build(build_opts)
+    sys.stdout.write(f"{json.dumps(result, indent=2)}\n")
+
+
+def run_search(args):
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("file")
+    parser.add_argument("query")
+    parser.add_argument("--top-k", type=int)
+    parser.add_argument("--embed-provider")
+    parser.add_argument("--embed-model")
+    parser.add_argument("--embed-key")
+
+    parsed = parser.parse_args(args)
+
+    embeddings = parse_common_embedding(vars(parsed))
+    doc = Document(parsed.file, {"embeddings": embeddings})
+
+    search_opts = {}
+    if parsed.top_k is not None:
+        search_opts["topK"] = parsed.top_k
+
+    results = doc.search(parsed.query, search_opts)
+    serialized = [r.model_dump(by_alias=True) for r in results]
+    sys.stdout.write(f"{json.dumps(serialized, indent=2)}\n")
+
+
+def run_ask(args):
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("file")
+    parser.add_argument("question")
+    parser.add_argument("--top-k", type=int)
+    parser.add_argument("--embed-provider")
+    parser.add_argument("--embed-model")
+    parser.add_argument("--embed-key")
+    parser.add_argument("--llm-provider", required=True)
+    parser.add_argument("--llm-model")
+    parser.add_argument("--llm-key")
+    parser.add_argument("--stream", action="store_true")
+
+    parsed = parser.parse_args(args)
+
+    embeddings = parse_common_embedding(vars(parsed))
+    llm = parse_llm(vars(parsed))
+
+    doc = Document(parsed.file, {"embeddings": embeddings, "llm": llm})
+
+    opts = {}
+    if parsed.top_k is not None:
+        opts["topK"] = parsed.top_k
+
+    if parsed.stream:
+        for chunk in doc.ask_stream(parsed.question, opts):
+            sys.stdout.write(chunk)
+            sys.stdout.flush()
+        sys.stdout.write("\n")
+    else:
+        answer = doc.ask(parsed.question, opts)
+        sys.stdout.write(f"{answer.text}\n")
+
+
+def run_serve(args):
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("file")
+    parser.add_argument("--embed-provider")
+    parser.add_argument("--embed-model")
+    parser.add_argument("--embed-key")
+    parser.add_argument("--llm-provider")
+    parser.add_argument("--llm-model")
+    parser.add_argument("--llm-key")
+    parser.add_argument("--host")
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--token")
+
+    parsed = parser.parse_args(args)
+
+    embeddings = parse_common_embedding(vars(parsed))
+    llm = parse_llm(vars(parsed))
+
+    doc = Document(
+        parsed.file,
+        {"embeddings": embeddings, **({"llm": llm} if llm else {})},
+    )
+    doc.build()
+
+    serve_opts = {}
+    if llm:
+        serve_opts["llm"] = llm
+    if parsed.host:
+        serve_opts["host"] = parsed.host
+    if parsed.port is not None:
+        serve_opts["port"] = parsed.port
+    if parsed.token:
+        serve_opts["bearerToken"] = parsed.token
+
+    handle = doc.serve(serve_opts)
+    sys.stdout.write(f"RagLite listening on {handle.url}\n")
+    sys.stdout.flush()
+
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        handle.close()
+        sys.exit(0)
+
+
+def main():
+    argv = sys.argv[1:]
+    if not argv or argv[0] in ("--help", "-h", "help"):
+        sys.stdout.write(HELP)
+        sys.stdout.flush()
+        sys.exit(0)
+    if argv[0] in ("--version", "-v", "version"):
+        sys.stdout.write(f"{PACKAGE_VERSION}\n")
+        sys.stdout.flush()
+        sys.exit(0)
+
+    command = argv[0]
+    args = argv[1:]
+
+    try:
+        if command == "index":
+            run_index(args)
+        elif command == "search":
+            run_search(args)
+        elif command == "ask":
+            run_ask(args)
+        elif command == "serve":
+            run_serve(args)
+        else:
+            sys.stderr.write(f"Unknown command: {command}\n\n{HELP}")
+            sys.stderr.flush()
+            sys.exit(2)
+    except Exception as e:
+        sys.stderr.write(f"Error: {e}\n")
+        sys.stderr.flush()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

raglite/config.py

@@ -0,0 +1,65 @@
+from typing import Optional, Union, Dict, Any
+from pydantic import BaseModel, Field, ConfigDict
+
+from .types import EmbeddingProviderConfig, LLMProviderConfig
+from .constants import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_CHUNK_OVERLAP,
+    DEFAULT_TOP_K,
+    DEFAULT_SCORE_THRESHOLD,
+    DEFAULT_STORE_DIRNAME,
+)
+
+
+class DocumentOptions(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    chunkSize: Optional[int] = Field(default=None, alias="chunkSize")
+    overlap: Optional[int] = None
+    topK: Optional[int] = Field(default=None, alias="topK")
+    scoreThreshold: Optional[float] = Field(default=None, alias="scoreThreshold")
+    storeDir: Optional[str] = Field(default=None, alias="storeDir")
+    embeddings: Optional[EmbeddingProviderConfig] = None
+    llm: Optional[LLMProviderConfig] = None
+    logLevel: Optional[str] = Field(default=None, alias="logLevel")
+
+
+class ResolvedConfig(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="allow")
+
+    chunkSize: int = Field(..., alias="chunkSize")
+    overlap: int
+    topK: int = Field(..., alias="topK")
+    scoreThreshold: float = Field(..., alias="scoreThreshold")
+    storeDir: str = Field(..., alias="storeDir")
+    embeddings: EmbeddingProviderConfig
+    llm: Optional[LLMProviderConfig] = None
+    logLevel: str = Field(..., alias="logLevel")
+
+
+def resolve_config(options: Optional[Union[DocumentOptions, Dict[str, Any]]] = None) -> ResolvedConfig:
+    if options is None:
+        opts = DocumentOptions()
+    elif isinstance(options, dict):
+        opts = DocumentOptions.model_validate(options)
+    else:
+        opts = options
+
+    chunk_size = opts.chunkSize if opts.chunkSize is not None else DEFAULT_CHUNK_SIZE
+    overlap = opts.overlap if opts.overlap is not None else DEFAULT_CHUNK_OVERLAP
+    top_k = opts.topK if opts.topK is not None else DEFAULT_TOP_K
+    score_threshold = opts.scoreThreshold if opts.scoreThreshold is not None else DEFAULT_SCORE_THRESHOLD
+    store_dir = opts.storeDir if opts.storeDir is not None else DEFAULT_STORE_DIRNAME
+    embeddings = opts.embeddings if opts.embeddings is not None else EmbeddingProviderConfig(provider="local")
+    log_level = opts.logLevel if opts.logLevel is not None else "info"
+
+    return ResolvedConfig(
+        chunkSize=chunk_size,
+        overlap=overlap,
+        topK=top_k,
+        scoreThreshold=score_threshold,
+        storeDir=store_dir,
+        embeddings=embeddings,
+        llm=opts.llm,
+        logLevel=log_level
+    )

raglite/constants.py

@@ -0,0 +1,18 @@
+PACKAGE_NAME = "raglite-toolkit"
+PACKAGE_VERSION = "0.1.0"
+
+SUPPORTED_EXTENSIONS = {".pdf", ".txt", ".json", ".md", ".markdown", ".docx"}
+
+DEFAULT_CHUNK_SIZE = 500
+DEFAULT_CHUNK_OVERLAP = 50
+
+DEFAULT_TOP_K = 5
+DEFAULT_SCORE_THRESHOLD = 0.0
+
+DEFAULT_TEMPERATURE = 0.0
+
+DEFAULT_STORE_DIRNAME = ".raglite"
+DEFAULT_COLLECTION_NAME = "default"
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8085

raglite/core/__init__.py

@@ -0,0 +1 @@
+# core module