rag-python 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

rag_python/__init__.py

@@ -2,18 +2,29 @@
 
 Quick start::
 
-    from rag_python import RAG
+    pip install rag-python
+    export OPENAI_API_KEY=sk-...
+
+    # CLI
+    rag-python ingest ./docs --reindex
+    rag-python query "What is our leave policy?"
+    rag-python docs quickstart
 
-    rag = RAG(llm_model="gpt-4o-mini")
+    # Python
+    from rag_python import RAG
+    rag = RAG()
     rag.ingest(["./docs"], reindex=True)
     print(rag.query("What is our leave policy?").text)
+
+Documentation: https://github.com/RaghavOG/rag-python/tree/main/docs
 """
 
-__version__ = "0.2.0"
+__version__ = "0.3.1"
 
 from .client import RAG, RAGAnswer
-from .rag_pipeline import ingest, query, RAGResponse
+from .rag_pipeline import ingest, query, query_stream, RAGResponse, RAGStream
 from .providers import make_llm_provider, make_embedding_provider
+from .log import configure_logging, get_logger
 from .options import (
     ChunkingConfig,
     DocumentConfig,
@@ -33,7 +44,11 @@ __all__ = [
     "QueryConfig",
     "ingest",
     "query",
+    "query_stream",
+    "RAGStream",
     "RAGResponse",
+    "configure_logging",
+    "get_logger",
     "make_llm_provider",
     "make_embedding_provider",
 ]

rag_python/cli.py

@@ -1,13 +1,18 @@
 """rag-python command-line interface."""
+from __future__ import annotations
+
 import argparse
+import json
+import sys
 from dataclasses import replace
 
 from . import __version__
 from .client import RAG
+from .help_text import CLI_EPILOG, list_topics, print_topic, print_topic_list
 
 
 def _build_rag(args: argparse.Namespace) -> RAG:
-    return RAG(
+    kwargs: dict = dict(
         llm_provider=args.llm_provider,
         llm_model=args.llm_model,
         embedding_provider=args.embedding_provider,
@@ -20,6 +25,20 @@ def _build_rag(args: argparse.Namespace) -> RAG:
         gemini_api_key=args.gemini_api_key,
         ollama_base_url=args.ollama_base_url,
     )
+    if getattr(args, "retriever", None):
+        kwargs["retriever"] = args.retriever
+    if getattr(args, "metadata_filter", None):
+        kwargs["metadata_filter"] = args.metadata_filter
+    return RAG(**kwargs)
+
+
+def _parse_metadata_filter(raw: str | None) -> dict | None:
+    if not raw:
+        return None
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"Invalid JSON for metadata filter: {e}") from e
 
 
 def _add_provider_args(parser: argparse.ArgumentParser) -> None:
@@ -27,43 +46,180 @@ def _add_provider_args(parser: argparse.ArgumentParser) -> None:
         "--llm-provider",
         default="openai",
         choices=["openai", "azure_openai", "anthropic", "gemini", "ollama"],
+        metavar="PROVIDER",
+        help="LLM backend (default: openai). See: rag-python docs providers",
+    )
+    parser.add_argument(
+        "--llm-model",
+        default=None,
+        metavar="MODEL",
+        help="LLM model or Azure deployment name (default: from env LLM_MODEL)",
     )
-    parser.add_argument("--llm-model", default=None)
     parser.add_argument(
         "--embedding-provider",
         default="openai",
         choices=["openai", "azure_openai", "ollama", "local"],
+        metavar="PROVIDER",
+        help="Embedding backend (default: openai). Use local for offline embeddings",
+    )
+    parser.add_argument(
+        "--embedding-model",
+        default=None,
+        metavar="MODEL",
+        help="Embedding model name (default: from env EMBEDDING_MODEL)",
+    )
+    parser.add_argument(
+        "--ollama-base-url",
+        default=None,
+        metavar="URL",
+        help="Ollama server URL (default: http://localhost:11434 or OLLAMA_BASE_URL)",
+    )
+    parser.add_argument("--azure-endpoint", default=None, help="Azure OpenAI endpoint URL")
+    parser.add_argument("--azure-api-key", default=None, help="Azure OpenAI API key")
+    parser.add_argument(
+        "--azure-api-version",
+        default=None,
+        help="Azure API version (default: 2023-09-01-preview)",
+    )
+    parser.add_argument("--openai-api-key", default=None, help="OpenAI API key (overrides env)")
+    parser.add_argument("--anthropic-api-key", default=None, help="Anthropic API key")
+    parser.add_argument("--gemini-api-key", default=None, help="Gemini API key")
+
+
+def _add_search_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--retriever",
+        choices=["vector", "multi_query", "hybrid"],
+        default=None,
+        metavar="MODE",
+        help=(
+            "Retrieval mode: vector (single query), multi_query (default, with rewriting), "
+            "or hybrid (BM25+vector; requires pip install rag-python[hybrid])"
+        ),
+    )
+    parser.add_argument(
+        "--metadata-filter",
+        type=_parse_metadata_filter,
+        default=None,
+        metavar="JSON",
+        help='Filter chunks by metadata, e.g. \'{"filename": "policy.pdf"}\'',
     )
-    parser.add_argument("--embedding-model", default=None)
-    parser.add_argument("--ollama-base-url", default=None)
-    parser.add_argument("--azure-endpoint", default=None)
-    parser.add_argument("--azure-api-key", default=None)
-    parser.add_argument("--azure-api-version", default=None)
-    parser.add_argument("--openai-api-key", default=None)
-    parser.add_argument("--anthropic-api-key", default=None)
-    parser.add_argument("--gemini-api-key", default=None)
 
 
-def main() -> None:
+def _make_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         prog="rag-python",
-        description="rag-python — modular RAG with query rewriting, reranking, guardrails, and multi-LLM support.",
+        description=(
+            "Production-grade RAG for Python — ingest documents, ask questions, "
+            "get grounded answers with multi-LLM support."
+        ),
+        epilog=CLI_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    parser.add_argument("--version", action="version", version=f"rag-python {__version__}")
-    sub = parser.add_subparsers(dest="command", required=True)
+    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True, metavar="COMMAND")
 
-    ing = sub.add_parser("ingest", help="Ingest files/folders into the vector store")
-    ing.add_argument("paths", nargs="+", help="Files or folders to ingest")
-    ing.add_argument("--reindex", action="store_true", help="Clear vector store and re-ingest")
+    ing = sub.add_parser(
+        "ingest",
+        help="Load files into the vector store (chunk + embed)",
+        description=(
+            "Ingest one or more files or directories into the ChromaDB vector store.\n"
+            "Supported formats: .txt .md .pdf .docx .csv .json .html"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=(
+            "examples:\n"
+            "  rag-python ingest ./data --reindex\n"
+            "  rag-python ingest policy.pdf handbook/ --embedding-provider local"
+        ),
+    )
+    ing.add_argument(
+        "paths",
+        nargs="+",
+        metavar="PATH",
+        help="File or directory paths to ingest",
+    )
+    ing.add_argument(
+        "--reindex",
+        action="store_true",
+        help="Delete existing vectors before ingesting (fresh index)",
+    )
     _add_provider_args(ing)
 
-    q = sub.add_parser("query", help="Ask a question against ingested documents")
-    q.add_argument("question", nargs="+", help="Question text")
-    q.add_argument("--no-multi-query", action="store_true")
-    q.add_argument("-v", "--verbose", action="store_true")
+    q = sub.add_parser(
+        "query",
+        help="Ask a question against ingested documents",
+        description=(
+            "Run the full RAG pipeline: retrieve relevant chunks, generate an answer, "
+            "optionally stream tokens and show sources."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=(
+            "examples:\n"
+            '  rag-python query "How many days of annual leave?"\n'
+            "  rag-python query \"PTO policy\" --stream -v\n"
+            '  rag-python query "benefits" --retriever hybrid --metadata-filter \'{"filename": "hr.pdf"}\''
+        ),
+    )
+    q.add_argument(
+        "question",
+        nargs="+",
+        metavar="QUESTION",
+        help="Question text (multiple words are joined)",
+    )
+    q.add_argument(
+        "--no-multi-query",
+        action="store_true",
+        help="Use single-query vector retrieval (same as --retriever vector)",
+    )
+    q.add_argument(
+        "--stream",
+        action="store_true",
+        help="Stream answer tokens to stdout as they are generated",
+    )
+    q.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="After the answer, print evaluation scores and top source paths",
+    )
     _add_provider_args(q)
+    _add_search_args(q)
+
+    docs = sub.add_parser(
+        "docs",
+        help="Show user documentation in the terminal",
+        description="Print built-in help topics. Full docs: https://github.com/RaghavOG/rag-python/tree/main/docs",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="topics: " + ", ".join(list_topics()),
+    )
+    docs.add_argument(
+        "topic",
+        nargs="?",
+        default="quickstart",
+        choices=list_topics(),
+        metavar="TOPIC",
+        help="Documentation topic (default: quickstart)",
+    )
+    docs.add_argument(
+        "--list",
+        action="store_true",
+        help="List all available documentation topics",
+    )
+
+    return parser
 
-    args = parser.parse_args()
+
+def main(argv: list[str] | None = None) -> None:
+    parser = _make_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "docs":
+        if args.list:
+            print_topic_list()
+        else:
+            print_topic(args.topic)
+        return
 
     if args.command == "ingest":
         rag = _build_rag(args)
@@ -74,10 +230,28 @@ def main() -> None:
     if args.command == "query":
         rag = _build_rag(args)
         question = " ".join(args.question)
+        retriever = args.retriever
+        if retriever is None and args.no_multi_query:
+            retriever = "vector"
         search = replace(
             rag.config.search,
-            retriever="vector" if args.no_multi_query else "multi_query",
+            retriever=retriever or rag.config.search.retriever,
+            metadata_filter=args.metadata_filter or rag.config.search.metadata_filter,
         )
+        if args.stream:
+            stream = rag.query_stream(question, search=search)
+            for token in stream:
+                print(token, end="", flush=True)
+            print()
+            result = stream.result
+            if args.verbose:
+                print("\n--- evaluation ---")
+                print(result.evaluation)
+                print("\n--- sources ---")
+                for s in result.sources[:5]:
+                    print(s.get("metadata", {}).get("source", ""), "score:", s.get("score"))
+            return
+
         ans = rag.query(question, search=search)
         print(ans.text)
         if args.verbose:
@@ -89,4 +263,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main(sys.argv[1:])

rag_python/client.py

@@ -30,7 +30,7 @@ from .options import (
     SearchConfig,
 )
 from .providers import make_llm_provider, make_embedding_provider
-from .rag_pipeline import ingest as _ingest, query as _query, RAGResponse
+from .rag_pipeline import ingest as _ingest, query as _query, query_stream as _query_stream, RAGResponse, RAGStream
 from .vector_store import set_persist_dir
 
 
@@ -60,6 +60,7 @@ class RAG:
         chunk_size: int | None = None,
         chunk_overlap: int | None = None,
         retriever: str | None = None,
+        metadata_filter: dict | None = None,
         top_k_retrieve: int | None = None,
         top_k_rerank: int | None = None,
         multi_query_n: int | None = None,
@@ -104,6 +105,8 @@ class RAG:
             self.config.search = replace(self.config.search, rerank_enabled=rerank_enabled)
         if document_extensions is not None:
             self.config.documents = replace(self.config.documents, extensions=document_extensions)
+        if metadata_filter is not None:
+            self.config.search = replace(self.config.search, metadata_filter=metadata_filter)
 
         self.llm = make_llm_provider(
             llm_provider,  # type: ignore[arg-type]
@@ -188,3 +191,21 @@ class RAG:
             evaluation=resp.evaluation,
             retried=resp.retried,
         )
+
+    def query_stream(
+        self,
+        question: str,
+        *,
+        search: SearchConfig | None = None,
+        query_config: QueryConfig | None = None,
+    ) -> RAGStream:
+        """Stream answer tokens; call ``stream.result`` after iterating."""
+        return _query_stream(
+            question,
+            search=search or self.config.search,
+            query_config=query_config or self.config.query,
+            llm_model=self.llm_model,
+            embedding_model=self.embedding_model,
+            llm=self.llm,
+            embedder=self.embedder,
+        )

rag_python/document_loaders.py

@@ -1,4 +1,7 @@
 """Document loaders: raw data → structured text + metadata."""
+import csv
+import json
+from html.parser import HTMLParser
 from pathlib import Path
 from dataclasses import dataclass
 from typing import Iterator
@@ -22,18 +25,85 @@ class LoadedDocument:
     metadata: dict
 
 
+class _HTMLTextExtractor(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.parts: list[str] = []
+
+    def handle_data(self, data: str) -> None:
+        text = data.strip()
+        if text:
+            self.parts.append(text)
+
+
+def _html_to_text(html: str) -> str:
+    parser = _HTMLTextExtractor()
+    parser.feed(html)
+    return "\n".join(parser.parts)
+
+
+def _load_csv(path: Path, metadata: dict) -> LoadedDocument | None:
+    rows: list[str] = []
+    with path.open(encoding="utf-8", errors="replace", newline="") as f:
+        reader = csv.DictReader(f)
+        if reader.fieldnames:
+            for row in reader:
+                rows.append(", ".join(f"{k}: {v}" for k, v in row.items() if v))
+        else:
+            f.seek(0)
+            for row in csv.reader(f):
+                rows.append(", ".join(row))
+    content = "\n".join(rows)
+    metadata["rows"] = len(rows)
+    return LoadedDocument(content=content, source=str(path), metadata=metadata) if content.strip() else None
+
+
+def _load_json(path: Path, metadata: dict) -> LoadedDocument | None:
+    data = json.loads(path.read_text(encoding="utf-8", errors="replace"))
+    if isinstance(data, list):
+        parts = []
+        for item in data:
+            if isinstance(item, dict) and "text" in item:
+                parts.append(str(item["text"]))
+            else:
+                parts.append(json.dumps(item, ensure_ascii=False))
+        content = "\n\n".join(parts)
+    elif isinstance(data, dict):
+        if "text" in data:
+            content = str(data["text"])
+        else:
+            content = json.dumps(data, ensure_ascii=False, indent=2)
+    else:
+        content = str(data)
+    return LoadedDocument(content=content, source=str(path), metadata=metadata) if content.strip() else None
+
+
 def load_file(path: Path) -> LoadedDocument | None:
-    """Load a single file (PDF, TXT, DOCX, MD) into text + metadata."""
+    """Load a single file (PDF, TXT, DOCX, MD, CSV, JSON, HTML) into text + metadata."""
     path = Path(path)
     if not path.exists():
         return None
     suffix = path.suffix.lower()
     metadata = {"source": str(path), "filename": path.name}
 
-    if suffix == ".txt" or suffix == ".md":
+    if suffix in (".txt", ".md"):
         content = path.read_text(encoding="utf-8", errors="replace")
         return LoadedDocument(content=content, source=str(path), metadata=metadata)
 
+    if suffix == ".html":
+        html = path.read_text(encoding="utf-8", errors="replace")
+        content = _html_to_text(html)
+        return LoadedDocument(content=content, source=str(path), metadata=metadata) if content.strip() else None
+
+    if suffix == ".csv":
+        return _load_csv(path, metadata)
+
+    if suffix == ".json":
+        try:
+            return _load_json(path, metadata)
+        except json.JSONDecodeError:
+            return None
+
     if suffix == ".pdf" and PdfReader:
         try:
             reader = PdfReader(path)
@@ -61,7 +131,10 @@ def load_file(path: Path) -> LoadedDocument | None:
     return None
 
 
-def load_directory(dir_path: Path, extensions: tuple = (".txt", ".md", ".pdf", ".docx")) -> Iterator[LoadedDocument]:
+def load_directory(
+    dir_path: Path,
+    extensions: tuple = (".txt", ".md", ".pdf", ".docx", ".csv", ".json", ".html"),
+) -> Iterator[LoadedDocument]:
     """Yield LoadedDocument for each supported file under dir_path."""
     dir_path = Path(dir_path)
     if not dir_path.is_dir():
@@ -71,4 +144,3 @@ def load_directory(dir_path: Path, extensions: tuple = (".txt", ".md", ".pdf", "
             doc = load_file(f)
             if doc and doc.content.strip():
                 yield doc
-

rag_python/generation.py

@@ -1,6 +1,9 @@
 """LLM generation with context (RAG)."""
+from collections.abc import Iterator
+
 from .config import LLM_MODEL
 from .providers import LLMProvider, make_llm_provider
+from .providers.streaming import stream_generate
 
 
 RAG_SYSTEM = (
@@ -10,6 +13,11 @@ RAG_SYSTEM = (
 )
 
 
+def _build_user_prompt(query: str, context_chunks: list[str]) -> str:
+    context = "\n\n---\n\n".join(context_chunks)
+    return f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"
+
+
 def generate(
     query: str,
     context_chunks: list[str],
@@ -20,12 +28,11 @@ def generate(
 ) -> str:
     """Generate answer from query and retrieved context."""
     llm = llm or make_llm_provider("openai")
-    context = "\n\n---\n\n".join(context_chunks)
     sys = system_prompt or RAG_SYSTEM
     try:
         return llm.generate(
             system=sys,
-            user=f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:",
+            user=_build_user_prompt(query, context_chunks),
             model=model or LLM_MODEL,
             temperature=0.2,
             max_tokens=1024,
@@ -33,3 +40,26 @@ def generate(
     except Exception as e:
         return f"[Generation error: {e}]"
 
+
+def generate_stream(
+    query: str,
+    context_chunks: list[str],
+    *,
+    model: str | None = None,
+    system_prompt: str | None = None,
+    llm: LLMProvider | None = None,
+) -> Iterator[str]:
+    """Stream answer tokens from query and retrieved context."""
+    llm = llm or make_llm_provider("openai")
+    sys = system_prompt or RAG_SYSTEM
+    try:
+        yield from stream_generate(
+            llm,
+            system=sys,
+            user=_build_user_prompt(query, context_chunks),
+            model=model or LLM_MODEL,
+            temperature=0.2,
+            max_tokens=1024,
+        )
+    except Exception as e:
+        yield f"[Generation error: {e}]"