File size: 4,043 Bytes
0703e2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """Command line: ``fgao serve``, ``fgao ask``, ``fgao check``.
fgao serve --config fgao.yaml --host 0.0.0.0 --port 8000 [--workers 1]
fgao serve --factory myapp.rag:build # build() returns an FGAORAG
fgao ask --config fgao.yaml "What is ...?" # one query, JSON to stdout
fgao check --config fgao.yaml # build the pipeline and print its layout
"""
from __future__ import annotations
import argparse
import asyncio
import json
import logging
import os
import sys
from typing import Any, Dict, List, Optional, Tuple
def _load(config: Optional[str], factory: Optional[str]) -> Tuple[Any, Dict[str, Any]]:
from .config import import_object, load_config
if factory:
return import_object(factory)(), {}
if not config:
raise SystemExit("give --config FILE or --factory module:function")
return load_config(config)
def app_from_env():
"""uvicorn factory (used for --workers > 1): reads FGAO_CONFIG / FGAO_FACTORY."""
from .app import create_app
from .config import server_kwargs
rag, settings = _load(os.environ.get("FGAO_CONFIG"), os.environ.get("FGAO_FACTORY"))
return create_app(rag, **server_kwargs(settings))
def main(argv: Optional[List[str]] = None) -> int:
ap = argparse.ArgumentParser(prog="fgao", description="FGAO orchestration service")
sub = ap.add_subparsers(dest="cmd", required=True)
for name in ("serve", "ask", "check"):
p = sub.add_parser(name)
p.add_argument("--config", help="YAML or JSON configuration file")
p.add_argument("--factory", help="module:function returning an FGAORAG")
if name == "serve":
p.add_argument("--host", default="127.0.0.1")
p.add_argument("--port", type=int, default=8000)
p.add_argument("--workers", type=int, default=1)
p.add_argument("--log-level", default="info")
p.add_argument("--keep-alive", type=int, default=75,
help="idle keep-alive seconds; keep it above your load balancer's idle timeout")
if name == "ask":
p.add_argument("question")
p.add_argument("--budget", type=float)
p.add_argument("--explain", action="store_true", help="print the decision trace instead of JSON")
a = ap.parse_args(argv)
if a.cmd == "check":
rag, settings = _load(a.config, a.factory)
print(json.dumps({"sources": [{"name": s.name, "k": s.k, "cost": s.cost, "family": s.family,
"rewrite": s.rewrite, "retriever": type(s.retriever).__name__}
for s in rag.sources],
"llm": type(rag.llm).__name__ if rag.llm is not None else None,
"budget": rag.budget, "on_error": rag.on_error,
"operations": [str(x) for x in rag.app.actions], "server": settings}, indent=2))
return 0
if a.cmd == "ask":
rag, _ = _load(a.config, a.factory)
r = asyncio.run(rag.ainvoke(a.question, budget=a.budget))
print(r.explain() if a.explain else json.dumps(r.to_dict(), indent=2))
return 0
import uvicorn
logging.basicConfig(level=a.log_level.upper(), format="%(message)s")
if a.workers > 1:
if a.config:
os.environ["FGAO_CONFIG"] = os.path.abspath(a.config)
if a.factory:
os.environ["FGAO_FACTORY"] = a.factory
uvicorn.run("fgao.server.cli:app_from_env", factory=True, host=a.host, port=a.port, workers=a.workers,
log_level=a.log_level, timeout_keep_alive=a.keep_alive)
return 0
from .app import create_app
from .config import server_kwargs
rag, settings = _load(a.config, a.factory)
uvicorn.run(create_app(rag, **server_kwargs(settings)), host=a.host, port=a.port, log_level=a.log_level,
timeout_keep_alive=a.keep_alive)
return 0
if __name__ == "__main__": # pragma: no cover
sys.exit(main())
|