(info: AiModelCreator, trans: Trans)
| 24 | @router.post("/status", include_in_schema=False) |
| 25 | @require_permissions(permission=SqlbotPermission(role=['admin'])) |
| 26 | async def check_llm(info: AiModelCreator, trans: Trans): |
| 27 | async def generate(): |
| 28 | try: |
| 29 | additional_params = {item.key: prepare_model_arg(item.val) for item in info.config_list if |
| 30 | item.key and item.val} |
| 31 | config = LLMConfig( |
| 32 | model_type="openai" if info.protocol == 1 else "vllm", |
| 33 | model_name=info.base_model, |
| 34 | api_key=info.api_key, |
| 35 | api_base_url=info.api_domain, |
| 36 | additional_params=additional_params, |
| 37 | ) |
| 38 | llm_instance = LLMFactory.create_llm(config) |
| 39 | async for chunk in llm_instance.llm.astream("1+1=?"): |
| 40 | SQLBotLogUtil.info(chunk) |
| 41 | if chunk and isinstance(chunk, str): |
| 42 | yield json.dumps({"content": chunk}) + "\n" |
| 43 | if chunk and isinstance(chunk, dict) and chunk.content: |
| 44 | yield json.dumps({"content": chunk.content}) + "\n" |
| 45 | |
| 46 | except Exception as e: |
| 47 | SQLBotLogUtil.error(f"Error checking LLM: {e}") |
| 48 | error_msg = trans('i18n_llm.validate_error', msg=str(e)) |
| 49 | yield json.dumps({"error": error_msg}) + "\n" |
| 50 | |
| 51 | return StreamingResponse(generate(), media_type="application/x-ndjson") |
| 52 | |
| 53 | |
| 54 | @router.get("/default", include_in_schema=False) |
nothing calls this directly
no test coverage detected