Skip to content

Repository files navigation

KSS RAG

Python Version License Version Framework

Stand up a streaming RAG API from a single document with one command.

Quick StartWhy KSS RAGCLIPython APIArchitectureConfiguration

Overview

KSS RAG is a Retrieval-Augmented Generation framework built around one idea: getting from "I have a document" to "I have a live, streaming Q&A API over it" should take a single command — no glue code, no notebook, no orchestration boilerplate.

Point the CLI at a source-of-truth document, hand it a system prompt, and it loads, chunks, indexes, and serves a FastAPI endpoint with Server-Sent Events streaming and per-session conversation memory. The same pipeline is available as a Python API and a one-shot CLI query when you don't need a server.

It's provider-flexible by design: point it at any LLM provider — hosted (OpenRouter, OpenAI, Groq, Together, DeepSeek, Anthropic, and more) or local (Ollama, LM Studio, vLLM) — by setting one env var or one CLI flag, with automatic fallback to backup models when one is unavailable. See LLM Providers.

Why KSS RAG

Most RAG frameworks are libraries — powerful, but they hand you primitives and expect you to assemble the server, the streaming, and the session handling yourself. KSS RAG ships that assembly as a first-class feature:

  • One command to a running API. kssrag server --file docs.pdf --system-prompt prompt.txt gives you /query, /stream (SSE), /health, and session management out of the box.
  • Bring your own prompt and source of truth. The system prompt and the document are inputs, not code changes.
  • Pluggable everything. Six vector stores, two retrievers, multiple chunkers — selected by config, or replaced entirely with your own classes via an import path (no forking required).
  • Streaming that doesn't leak internals. Token-by-token SSE with marker-aware buffering (see Conversation memory).
  • Rolling conversation memory. The agent compresses history into running summaries to keep context bounded across long conversations.

If you want a RAG service rather than a RAG toolkit, that's the niche this fills.

Quick Start

Install

pip install kssrag

# Optional extras
pip install kssrag[ocr]      # PaddleOCR (handwritten) + Tesseract (typed)
pip install kssrag[office]   # DOCX / Excel / PowerPoint loaders
pip install kssrag[faiss]    # FAISS + sentence-transformers (semantic stores)
pip install kssrag[gpu]      # GPU FAISS
pip install kssrag[all]      # everything

Set your key (see .env.example for all options):

echo "OPENROUTER_API_KEY=your_key_here" > .env

Serve a RAG API in one command

python -m kssrag.cli server \
    --file knowledge_base.txt \
    --system-prompt "You are a support assistant. Answer only from the provided context." \
    --vector-store hybrid_offline \
    --host 0.0.0.0 --port 8000

Then query it:

curl -X POST http://localhost:8000/stream \
    -H "Content-Type: application/json" \
    -d '{"query": "How do I reset my password?", "session_id": "user-123"}'

--system-prompt accepts either an inline string or a path to a prompt file.

CLI

Two subcommands: query (one-shot) and server (persistent API).

# One-shot query with streaming output
python -m kssrag.cli query \
    --file report.pdf \
    --format pdf \
    --query "Summarize the key risks." \
    --vector-store hybrid_online \
    --top-k 8 \
    --stream

# OCR an image, then query it
python -m kssrag.cli query \
    --file scanned_notes.png \
    --format image \
    --ocr-mode handwritten \
    --query "What are the action items?"
Flag Applies to Description
--file both Path to the source document (required)
--query query The question to ask (required)
--format both text, json, pdf, image, docx, excel, pptx
--vector-store both bm25, bm25s, faiss, tfidf, hybrid_online, hybrid_offline
--system-prompt both Inline prompt text, or a path to a prompt file
--stream query Stream the response token-by-token
--top-k query Number of chunks to retrieve
--ocr-mode query typed (Tesseract) or handwritten (PaddleOCR)
--host / --port server Server bind address

Note: the server subcommand currently loads text, json, and pdf formats. Image and Office formats are supported by the query subcommand.

LLM Providers

KSS RAG talks to any LLM provider through a single --provider flag (or the PROVIDER env var). OpenRouter is the default.

# Groq (hosted, OpenAI-compatible)
kssrag query --file docs.txt --query "..." --provider groq --model llama-3.3-70b-versatile

# OpenAI
kssrag query --file docs.txt --query "..." --provider openai --model gpt-4o

# Anthropic (native Messages API)
kssrag query --file docs.txt --query "..." --provider anthropic --model claude-sonnet-4-6

# Local Ollama — no API key needed
kssrag query --file docs.txt --query "..." --provider ollama --model llama3

# Any custom OpenAI-compatible endpoint
kssrag query --file docs.txt --query "..." --provider custom \
    --base-url http://my-host:8000/v1/chat/completions --model my-model

Supported providers

Kind Providers
Hosted (OpenAI-compatible) openrouter, openai, groq, together, deepseek, fireworks, mistral, perplexity, xai, deepinfra, anyscale
Native protocol anthropic (Claude Messages API), ollama (/api/chat)
Local (OpenAI-compatible, no key) ollama-openai, lmstudio, vllm, llamacpp
Anything else custom (supply --base-url)

API keys

Set the key via LLM_API_KEY, or the provider's own env var (GROQ_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, ...), or --api-key. Local providers need no key. Selection precedence: --api-key > LLM_API_KEY > provider env var > OPENROUTER_API_KEY.

From Python

from kssrag import create_llm, RAGAgent

llm = create_llm(provider="groq", model="llama-3.3-70b-versatile")
# or: create_llm(provider="ollama", model="llama3")
# or: create_llm(provider="custom", base_url="http://localhost:8000/v1/chat/completions", model="m")

All providers share one interface (predict / predict_stream), so streaming, fallback models, and conversation memory work identically regardless of provider.

Cross-provider fallback

FALLBACK_MODELS only covers other models on the same provider. For real resilience across providers, set FALLBACK_PROVIDERS (comma-separated preset names) to try other providers in order when the primary fails — each fallback uses its own API key:

PROVIDER=groq
GROQ_API_KEY=your_groq_key
FALLBACK_PROVIDERS=openrouter

Env-driven fallbacks use DEFAULT_MODEL for every provider, so only list providers where that model id is valid (OpenRouter-style ids like deepseek/deepseek-chat work on OpenRouter, but not on Anthropic). Use the programmatic form to give each fallback its own model:

or programmatically:

llm = create_llm(
    provider="groq", model="openai/gpt-oss-20b",
    fallback_providers=[
        ("openrouter", {"model": "deepseek/deepseek-chat"}),
        {"provider": "anthropic", "model": "claude-sonnet-4-6"},
    ],
)

Each entry is a provider name, a (name, kwargs) tuple, or a dict that includes provider. A fallback provider that isn't configured (e.g. no key) is skipped with a warning instead of breaking the chain. Rate-limit (HTTP 429) responses are retried on the same model before moving on, honoring the server's Retry-After / rate-limit-reset headers (LLM_MAX_RETRIES, LLM_RETRY_BACKOFF).

Python API

from kssrag import KSSRAG, Config, VectorStoreType

config = Config(
    OPENROUTER_API_KEY="your-key",
    VECTOR_STORE_TYPE=VectorStoreType.HYBRID_OFFLINE,
    CHUNK_SIZE=800,
    TOP_K=8,
)

rag = KSSRAG(config=config)
rag.load_document("technical_docs.pdf")

# Blocking query
print(rag.query("What are the technical specifications?"))

# Filter retrieved context by document metadata (AND semantics)
print(rag.query("Specifications", metadata_filter={"source": "docs"}))

# Streaming query
for chunk in rag.agent.query_stream("Walk me through the architecture.", top_k=8):
    print(chunk, end="", flush=True)

Note: the KSSRAG class auto-detects .txt, .json, and .pdf. For images and Office documents, use the CLI or construct the matching chunker (ImageChunker, OfficeChunker) directly.

Custom components

Any pipeline stage can be replaced with your own class — no fork needed. Point config at an import path:

config = Config(
    CUSTOM_VECTOR_STORE="my_module.MyVectorStore",
    CUSTOM_RETRIEVER="my_module.MyRetriever",
    CUSTOM_LLM="my_module.MyLLM",
)

Custom vector stores subclass BaseVectorStore (add_documents / retrieve / persist / load); retrievers subclass BaseRetriever (retrieve).

Serve as an embedded API

from kssrag import KSSRAG
import uvicorn

rag = KSSRAG()
rag.load_document("knowledge.txt")
app, server_config = rag.create_server()
uvicorn.run(app, host="0.0.0.0", port=8000)

Endpoints

Endpoint Method Description
/query POST Query the RAG system (query, session_id, optional metadata_filter)
/stream POST Streaming query via Server-Sent Events
/v1/chat/completions POST OpenAI-compatible chat completions (plain + SSE streaming)
/health GET Health check
/config GET Active server configuration
/sessions/{id}/clear GET Clear a session's conversation history

Each session_id gets its own conversation state (held in memory for the life of the server process). CORS is configurable via environment variables.

Architecture

The pipeline: load → chunk → vector store → retriever → agent → LLM. Each stage is swappable by config or replaceable with a custom class.

Document ──> Chunker ──> Vector Store ──> Retriever ──┐
                                                       ├──> RAG Agent ──> LLM (any provider) ──> Response (stream / blocking)
                              Query ───────────────────┘

Vector stores

Store Method Needs model download? Best for
bm25 Keyword (BM25Okapi) No Fast keyword search
bm25s Stemmed BM25 (bm25s lib) No Faster BM25 with stemming
tfidf TF-IDF + cosine No Statistical relevance
faiss Dense embeddings (SentenceTransformers) Yes Semantic search
hybrid_online BM25 + FAISS, embedding-reranked Yes Best semantic quality
hybrid_offline BM25 + TF-IDF, score-fused No Semantic-ish quality, zero downloads, air-gapped

hybrid_offline is the default — it needs no network access or model download, which makes it a solid choice for restricted environments.

Chunkers

TextChunker (character windows with overlap) is the base. SentenceChunker packs whole sentences into chunks so fixed windows never cut mid-thought — great for persona/memory documents (select via CHUNKER_TYPE=sentence or --chunker sentence). PDFChunker, ImageChunker (OCR), and OfficeChunker extract text and delegate to a text chunker; JSONChunker flattens records keyed on a name field. Every chunk is a {"content", "metadata"} dict carried through the whole pipeline.

Conversation memory

To keep long conversations from blowing up context windows, the agent maintains rolling summaries: after a couple of exchanges it asks the model to append a compact [SUMMARY_START]...[SUMMARY_END] block to each response, extracts and stores it, and strips it before the user ever sees it. Streaming is marker-aware — it buffers around partial markers at chunk boundaries so a summary can never leak mid-stream. Older raw turns are trimmed while their summaries are retained, so the agent "remembers" the gist without paying for the full transcript.

FAISS loading

FAISS is imported lazily and only when a FAISS-backed store is actually used. It probes AVX512 → AVX2 → standard builds in order, so it runs on machines without AVX2 (including many Windows setups) instead of hard-failing at import.

Configuration

Everything is configurable through environment variables (.env) or the Config object. Highlights:

OPENROUTER_API_KEY=your_key
DEFAULT_MODEL=deepseek/deepseek-chat-v3.1:free
FALLBACK_MODELS=deepseek/deepseek-r1:free,deepseek/deepseek-chat
FALLBACK_PROVIDERS=openrouter,anthropic   # optional cross-provider chain
LLM_MAX_RETRIES=3                          # 429 rate-limit retries
LLM_RETRY_BACKOFF=2.0                      # base backoff seconds (doubles per attempt)

CHUNK_SIZE=500
CHUNK_OVERLAP=50
VECTOR_STORE_TYPE=hybrid_offline
RETRIEVER_TYPE=simple
TOP_K=5

SERVER_HOST=localhost
SERVER_PORT=8000
CORS_ORIGINS=*

See .env.example for the complete list, including OCR mode, batch size, fuzzy-match threshold, CORS details, and custom-component import paths.

Development

git clone https://github.com/Ksschkw/kssrag
cd kssrag
pip install -e .[dev,ocr,all]

python -m pytest tests/ -v        # run tests
python -m pytest tests/test_basic.py::test_text_rag -v   # single test

black kssrag/ tests/             # format
flake8 kssrag/                   # lint
mypy kssrag/                     # type-check

Acknowledgments

Built on FAISS, PaddleOCR, SentenceTransformers, bm25s, and a range of LLM providers via OpenRouter and OpenAI-compatible / native APIs.

Links

License

MIT — see LICENSE.

About

An AI RAG agent package.

Topics

Resources

Code of conduct

Contributing

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages