Files
VianoBugAnalysis/src/viano_bug_analysis/connection_test.py
T

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())