gsas-query 0.2.0__py3-none-any.whl

gsas_query/__init__.py

@@ -0,0 +1,20 @@
+"""
+GSAS-II Documentation Assistant
+
+Semantic search + AI answers over GSAS-II tutorials, help pages, and PDFs.
+All embedding and retrieval runs locally; the LLM backend is configurable.
+
+Quick start:
+    from gsas_query.rag import answer_question
+    result = answer_question("How do I set up a sequential refinement?", [])
+
+wxPython GUI (for GSAS-II Help menu integration):
+    from gsas_query.gui import show_assistant
+    show_assistant(parent_wx_window)
+"""
+
+from gsas_query._paths import get_chroma_path, get_data_dir, get_static_dir
+from gsas_query.rag import answer_question
+
+__all__ = ["answer_question", "get_chroma_path", "get_data_dir", "get_static_dir"]
+__version__ = "0.2.0"

gsas_query/_paths.py

@@ -0,0 +1,24 @@
+"""Central path management for user data (chroma index) and package assets."""
+
+import os
+from pathlib import Path
+
+# Suppress HuggingFace tokenizer fork warning and OpenMP threading —
+# both must be set before sentence-transformers / loky are imported.
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+
+
+def get_data_dir() -> Path:
+    """Return (and create) the user data directory for gsas_query."""
+    d = Path(os.environ.get("GSAS_QUERY_DATA_DIR", Path.home() / ".gsas_query"))
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def get_chroma_path() -> str:
+    return str(get_data_dir() / "chroma_db")
+
+
+def get_static_dir() -> Path:
+    return Path(__file__).parent / "static"

gsas_query/_web.py

@@ -0,0 +1,16 @@
+"""Entry point for `gsas-query-web` console script."""
+
+import os
+
+
+def main():
+    import uvicorn
+    from .app import app
+
+    host = os.environ.get("HOST", "0.0.0.0")
+    port = int(os.environ.get("PORT", "8000"))
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()

gsas_query/app.py

@@ -0,0 +1,154 @@
+"""
+FastAPI web server for the GSAS-II documentation chatbot.
+
+  GET  /         -> chat UI (static/index.html)
+  GET  /health   -> {"status": "ok"}
+  GET  /stats    -> {"chunks_indexed": N, "llm_backend": "..."}
+  POST /chat     -> RAG query, returns {answer, sources, citations}
+  POST /ingest   -> trigger re-indexing (requires X-Admin-Key header)
+"""
+
+import os
+import subprocess
+import sys
+import time
+from typing import Optional
+
+from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, JSONResponse
+from pydantic import BaseModel, Field
+
+from ._paths import get_static_dir
+
+STATIC_DIR = get_static_dir()
+
+app = FastAPI(
+    title="GSAS-II Documentation Assistant",
+    description="RAG chatbot over GSAS-II tutorials and help documentation",
+    version="0.2.0",
+)
+
+origins = os.environ.get("ALLOWED_ORIGINS", "*").split(",")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_methods=["GET", "POST"],
+    allow_headers=["*"],
+)
+
+
+@app.exception_handler(Exception)
+async def generic_exception_handler(request: Request, exc: Exception):
+    return JSONResponse(status_code=500, content={"detail": str(exc)})
+
+
+# --------------------------------------------------------------------------- #
+# Request / response models                                                    #
+# --------------------------------------------------------------------------- #
+
+class HistoryTurn(BaseModel):
+    role: str
+    content: str
+
+
+class ChatRequest(BaseModel):
+    message: str = Field(..., max_length=2000)
+    history: list[HistoryTurn] = Field(default_factory=list, max_length=20)
+
+
+class Source(BaseModel):
+    title: str
+    section: str
+    url: str
+    category: str
+    relevance: float
+
+
+class Citation(BaseModel):
+    title: str
+    section: str
+    url: str
+
+
+class ChatResponse(BaseModel):
+    answer: str
+    sources: list[Source]
+    citations: dict[str, Citation]
+    elapsed_ms: int
+
+
+# --------------------------------------------------------------------------- #
+# Simple rate limiting (in-memory, per IP)                                     #
+# --------------------------------------------------------------------------- #
+
+_rate_store: dict[str, list[float]] = {}
+RATE_LIMIT = int(os.environ.get("RATE_LIMIT_RPM", "30"))
+
+
+def check_rate_limit(request: Request):
+    if RATE_LIMIT <= 0:
+        return
+    ip = request.client.host if request.client else "unknown"
+    now = time.time()
+    timestamps = [t for t in _rate_store.get(ip, []) if now - t < 60]
+    if len(timestamps) >= RATE_LIMIT:
+        raise HTTPException(status_code=429, detail="Rate limit exceeded. Please slow down.")
+    timestamps.append(now)
+    _rate_store[ip] = timestamps
+
+
+# --------------------------------------------------------------------------- #
+# Routes                                                                       #
+# --------------------------------------------------------------------------- #
+
+@app.get("/", include_in_schema=False)
+async def serve_ui():
+    return FileResponse(str(STATIC_DIR / "index.html"))
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@app.get("/stats")
+async def stats():
+    from .rag import _get_collection
+    try:
+        col = _get_collection()
+        count = col.count()
+    except Exception:
+        count = -1
+    backend = os.environ.get("LLM_BACKEND", "ollama")
+    return {"chunks_indexed": count, "llm_backend": backend}
+
+
+@app.post("/chat", response_model=ChatResponse)
+async def chat(
+    request: ChatRequest,
+    http_request: Request,
+    _: None = Depends(check_rate_limit),
+):
+    from .rag import answer_question
+
+    t0 = time.monotonic()
+    history = [h.model_dump() for h in request.history]
+    result = answer_question(request.message, history)
+    elapsed = int((time.monotonic() - t0) * 1000)
+
+    return ChatResponse(
+        answer=result["answer"],
+        sources=result["sources"],
+        citations=result.get("citations", {}),
+        elapsed_ms=elapsed,
+    )
+
+
+@app.post("/ingest")
+async def trigger_ingest(http_request: Request):
+    admin_key = os.environ.get("ADMIN_KEY", "")
+    if admin_key and http_request.headers.get("X-Admin-Key") != admin_key:
+        raise HTTPException(status_code=403, detail="Invalid admin key.")
+    subprocess.Popen([sys.executable, "-m", "gsas_query.ingest"])
+    return {"status": "ingestion started in background"}

gsas_query/cli.py

@@ -0,0 +1,208 @@
+"""
+GSAS-II Documentation Assistant — command-line interface.
+
+First-time setup:
+    gsas-query --setup
+    gsas-query --setup --reset
+    gsas-query --setup --html-only
+
+Ask a single question:
+    gsas-query "How do I set up a sequential refinement?"
+
+Interactive REPL:
+    gsas-query
+"""
+
+import os
+import sys
+import textwrap
+
+from ._paths import get_chroma_path
+
+WIDTH = 80
+
+
+def _hr(char="─", width=WIDTH):
+    print(char * width)
+
+
+def _wrap(text: str, indent: int = 0) -> str:
+    prefix = " " * indent
+    return textwrap.fill(text, width=WIDTH, initial_indent=prefix, subsequent_indent=prefix)
+
+
+def _print_answer(result: dict):
+    answer = result.get("answer", "")
+    sources = result.get("sources", [])
+
+    print()
+    for para in answer.split("\n"):
+        if para.strip():
+            print(_wrap(para))
+        else:
+            print()
+
+    if sources:
+        print()
+        _hr("·")
+        print("Sources:")
+        seen = set()
+        for s in sources:
+            key = s["url"]
+            if key in seen:
+                continue
+            seen.add(key)
+            rel = int(s.get("relevance", 0) * 100)
+            label = f"  [{rel}%] {s['title']}"
+            if s.get("section") and s["section"] != s["title"]:
+                label += f"  ›  {s['section']}"
+            print(label)
+            print(f"         {s['url']}")
+    print()
+
+
+def _collection_count() -> int:
+    try:
+        import chromadb
+        client = chromadb.PersistentClient(path=get_chroma_path())
+        return client.get_or_create_collection("gsasii_docs").count()
+    except Exception:
+        return 0
+
+
+def run_setup(reset: bool = False, html_only: bool = False):
+    from . import ingest
+    import argparse
+
+    print("Starting ingestion. This may take 10–20 minutes.")
+    print(f"Index will be stored at: {get_chroma_path()}\n")
+
+    argv_backup = sys.argv
+    sys.argv = ["gsas-query"]
+    if reset:
+        sys.argv.append("--reset")
+    if html_only:
+        sys.argv.append("--html-only")
+    try:
+        ingest.main()
+    finally:
+        sys.argv = argv_backup
+
+
+def ask(question: str, history: list | None = None) -> dict:
+    from .rag import answer_question
+    return answer_question(question, history or [])
+
+
+def interactive():
+    count = _collection_count()
+    if count == 0:
+        print("\nWarning: the knowledge base is empty.")
+        print("Run:  gsas-query --setup\n")
+    else:
+        print(f"\nKnowledge base: {count:,} indexed chunks.")
+
+    backend = os.environ.get("LLM_BACKEND", "ollama")
+    print(f"LLM backend: {backend}")
+    print("Type your question and press Enter. 'clear' resets history, 'quit' exits.\n")
+    _hr()
+
+    history: list[dict] = []
+
+    while True:
+        try:
+            question = input("You: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye.")
+            break
+
+        if not question:
+            continue
+        if question.lower() in {"quit", "exit", "q"}:
+            print("Goodbye.")
+            break
+        if question.lower() in {"clear", "reset"}:
+            history.clear()
+            print("(History cleared)\n")
+            continue
+
+        print("Thinking…", end="", flush=True)
+        try:
+            result = ask(question, history)
+        except Exception as e:
+            print(f"\rError: {e}\n")
+            continue
+
+        print("\r" + " " * 12 + "\r", end="")
+        print("Assistant:", end="")
+        _print_answer(result)
+
+        history.append({"role": "user", "content": question})
+        history.append({"role": "assistant", "content": result.get("answer", "")})
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="GSAS-II Documentation Assistant",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument("question", nargs="?", help="Question to ask (omit for interactive mode)")
+    parser.add_argument("--setup", action="store_true", help="Index documentation (first-time setup)")
+    parser.add_argument("--reset", action="store_true", help="Drop and rebuild the index")
+    parser.add_argument("--html-only", action="store_true", help="Skip PDFs during setup")
+    parser.add_argument("--backend", choices=["ollama", "anthropic", "retrieval"],
+                        help="Override LLM_BACKEND env var")
+    parser.add_argument("--model", help="Override OLLAMA_MODEL or ANTHROPIC_MODEL")
+    parser.add_argument("--stats", action="store_true", help="Show index statistics and exit")
+    parser.add_argument("--gui", action="store_true", help="Open the wxPython desktop assistant")
+    args = parser.parse_args()
+
+    if args.backend:
+        os.environ["LLM_BACKEND"] = args.backend
+    if args.model:
+        key = "OLLAMA_MODEL" if os.environ.get("LLM_BACKEND", "ollama") == "ollama" else "ANTHROPIC_MODEL"
+        os.environ[key] = args.model
+
+    print("GSAS-II Documentation Assistant")
+    _hr("═")
+
+    if args.stats:
+        count = _collection_count()
+        print(f"Indexed chunks : {count:,}")
+        print(f"LLM backend    : {os.environ.get('LLM_BACKEND', 'ollama')}")
+        print(f"Chroma DB      : {get_chroma_path()}")
+        return
+
+    if args.setup:
+        run_setup(reset=args.reset, html_only=args.html_only)
+        return
+
+    if args.gui:
+        from .gui import show_assistant
+        try:
+            import wx
+        except ImportError:
+            print("wxPython is required for the GUI. Install it or use GSAS-II's Python.")
+            sys.exit(1)
+        app = wx.App(False)
+        show_assistant()
+        app.MainLoop()
+        return
+
+    if args.question:
+        count = _collection_count()
+        if count == 0:
+            print("Knowledge base is empty. Run:  gsas-query --setup")
+            sys.exit(1)
+        print(f"({count:,} chunks indexed)\n")
+        result = ask(args.question)
+        _print_answer(result)
+    else:
+        interactive()
+
+
+if __name__ == "__main__":
+    main()