48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
"""One small LLM request, without creating or changing case/workspace files."""
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
|
|
from .llm import GatewayUnavailable, LLMError, PromptBuilder
|
|
from .trial import preflight
|
|
from .telemetry.safety import safe_code
|
|
|
|
|
|
async def probe(gateway):
|
|
request = PromptBuilder().build(
|
|
task='Return exactly the JSON object {"answer": "ok"}.',
|
|
data={}, privacy="configured",
|
|
output_schema={"type": "object", "properties": {"answer": {"const": "ok"}},
|
|
"required": ["answer"], "additionalProperties": False},
|
|
)
|
|
try:
|
|
result = await gateway.complete_round(request)
|
|
return {"status": "SUCCESS", "answer": result["document"], "metadata": result["metadata"]}
|
|
except GatewayUnavailable as error:
|
|
return {"status": "UNAVAILABLE", "attempts": error.attempts}
|
|
except LLMError as error:
|
|
result = {"status": "ERROR", "code": safe_code(error.code), "http_status": error.http_status}
|
|
if error.metadata is not None:
|
|
result['metadata'] = error.metadata
|
|
if error.diagnostics:
|
|
result['diagnostics'] = error.diagnostics
|
|
return result
|
|
|
|
|
|
def main(argv=None):
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--llm-config", default="config/llm.local.toml")
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
result = asyncio.run(probe(preflight(args.llm_config)))
|
|
except KeyboardInterrupt:
|
|
result = {"status": "INTERRUPTED"}
|
|
except Exception as error:
|
|
result = {"status": "ERROR", "code": safe_code(getattr(error, "code", None))}
|
|
print(json.dumps(result))
|
|
return 0 if result["status"] == "SUCCESS" else 130 if result["status"] == "INTERRUPTED" else 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|