108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
|
"""LLM service for AI-powered analysis using OpenAI-compatible APIs (StepFun, ZhipuAI)."""
|
||
|
|
import logging
|
||
|
|
from openai import AsyncOpenAI
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
async def get_llm_client(settings: dict, use_fallback: bool = False) -> tuple[AsyncOpenAI, str]:
|
||
|
|
"""Create an OpenAI-compatible client based on settings.
|
||
|
|
Returns (client, model_name) tuple.
|
||
|
|
"""
|
||
|
|
if use_fallback:
|
||
|
|
base_url = settings.get("ai_fallback_api_base_url", "")
|
||
|
|
api_key = settings.get("ai_fallback_api_key", "")
|
||
|
|
model = settings.get("ai_fallback_model_name", "codegeex-4")
|
||
|
|
else:
|
||
|
|
base_url = settings.get("ai_api_base_url", "")
|
||
|
|
api_key = settings.get("ai_api_key", "")
|
||
|
|
model = settings.get("ai_model_name", "step-2-16k")
|
||
|
|
|
||
|
|
client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=30.0)
|
||
|
|
return client, model
|
||
|
|
|
||
|
|
|
||
|
|
async def chat_completion(
|
||
|
|
messages: list[dict],
|
||
|
|
settings: dict,
|
||
|
|
temperature: float | None = None,
|
||
|
|
max_tokens: int | None = None,
|
||
|
|
) -> str:
|
||
|
|
"""Send a chat completion request with automatic fallback."""
|
||
|
|
temp = temperature or float(settings.get("ai_temperature", "0.7"))
|
||
|
|
tokens = max_tokens or int(settings.get("ai_max_tokens", "2000"))
|
||
|
|
|
||
|
|
# Try primary
|
||
|
|
try:
|
||
|
|
client, model = await get_llm_client(settings, use_fallback=False)
|
||
|
|
response = await client.chat.completions.create(
|
||
|
|
model=model,
|
||
|
|
messages=messages,
|
||
|
|
temperature=temp,
|
||
|
|
max_tokens=tokens,
|
||
|
|
)
|
||
|
|
return response.choices[0].message.content or ""
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning(f"Primary LLM failed: {e}")
|
||
|
|
|
||
|
|
# Try fallback if enabled
|
||
|
|
if settings.get("ai_fallback_enabled") == "true":
|
||
|
|
try:
|
||
|
|
client, model = await get_llm_client(settings, use_fallback=True)
|
||
|
|
response = await client.chat.completions.create(
|
||
|
|
model=model,
|
||
|
|
messages=messages,
|
||
|
|
temperature=temp,
|
||
|
|
max_tokens=tokens,
|
||
|
|
)
|
||
|
|
return response.choices[0].message.content or ""
|
||
|
|
except Exception as e2:
|
||
|
|
logger.error(f"Fallback LLM also failed: {e2}")
|
||
|
|
raise Exception(f"All LLM providers failed. Primary: {e}, Fallback: {e2}")
|
||
|
|
else:
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
async def test_connection(settings: dict) -> dict:
|
||
|
|
"""Test connection to the configured LLM provider."""
|
||
|
|
results = {"primary": {"status": "unknown"}, "fallback": {"status": "unknown"}}
|
||
|
|
|
||
|
|
# Test primary
|
||
|
|
try:
|
||
|
|
client, model = await get_llm_client(settings, use_fallback=False)
|
||
|
|
response = await client.chat.completions.create(
|
||
|
|
model=model,
|
||
|
|
messages=[{"role": "user", "content": "你好，请回复'连接成功'"}],
|
||
|
|
max_tokens=20,
|
||
|
|
temperature=0,
|
||
|
|
)
|
||
|
|
reply = response.choices[0].message.content or ""
|
||
|
|
results["primary"] = {
|
||
|
|
"status": "success",
|
||
|
|
"model": model,
|
||
|
|
"reply": reply[:100],
|
||
|
|
}
|
||
|
|
except Exception as e:
|
||
|
|
results["primary"] = {"status": "error", "error": str(e)[:200]}
|
||
|
|
|
||
|
|
# Test fallback
|
||
|
|
if settings.get("ai_fallback_enabled") == "true":
|
||
|
|
try:
|
||
|
|
client, model = await get_llm_client(settings, use_fallback=True)
|
||
|
|
response = await client.chat.completions.create(
|
||
|
|
model=model,
|
||
|
|
messages=[{"role": "user", "content": "你好，请回复'连接成功'"}],
|
||
|
|
max_tokens=20,
|
||
|
|
temperature=0,
|
||
|
|
)
|
||
|
|
reply = response.choices[0].message.content or ""
|
||
|
|
results["fallback"] = {
|
||
|
|
"status": "success",
|
||
|
|
"model": model,
|
||
|
|
"reply": reply[:100],
|
||
|
|
}
|
||
|
|
except Exception as e:
|
||
|
|
results["fallback"] = {"status": "error", "error": str(e)[:200]}
|
||
|
|
|
||
|
|
return results
|