151 lines
4.3 KiB
Python
151 lines
4.3 KiB
Python
"""Scaleway API client with async support and model selection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import httpx
|
|
|
|
from migration_pipeline.config import (
|
|
API_KEY,
|
|
API_TIMEOUT,
|
|
BASE_URL,
|
|
DEFAULT_MODEL_FAMILY,
|
|
MODEL_GLM,
|
|
MODEL_QWEN,
|
|
)
|
|
|
|
MODEL_FAMILY_GLM = "glm"
|
|
MODEL_FAMILY_QWEN = "qwen"
|
|
_SELECTED_MODEL_FAMILY = DEFAULT_MODEL_FAMILY
|
|
|
|
_MODEL_PAYLOADS: dict[str, dict] = {
|
|
MODEL_FAMILY_GLM: {
|
|
"model": MODEL_GLM,
|
|
"max_tokens": 16384,
|
|
"temperature": 1,
|
|
"top_p": 0.95,
|
|
"presence_penalty": 0,
|
|
"stream": True,
|
|
"response_format": {"type": "text"},
|
|
},
|
|
MODEL_FAMILY_QWEN: {
|
|
"model": MODEL_QWEN,
|
|
"max_tokens": 16384,
|
|
"temperature": 0.6,
|
|
"top_p": 0.95,
|
|
"presence_penalty": 0,
|
|
"stream": True,
|
|
"response_format": {"type": "text"},
|
|
},
|
|
}
|
|
|
|
_MODEL_TO_FAMILY = {
|
|
profile["model"]: family for family, profile in _MODEL_PAYLOADS.items()
|
|
}
|
|
|
|
|
|
def set_model_family(family: str) -> None:
|
|
"""Set active model family for the current pipeline run."""
|
|
normalized = family.strip().lower()
|
|
if normalized not in _MODEL_PAYLOADS:
|
|
raise ValueError(f"Unsupported model family: {family}")
|
|
global _SELECTED_MODEL_FAMILY
|
|
_SELECTED_MODEL_FAMILY = normalized
|
|
|
|
|
|
def get_selected_model_family() -> str:
|
|
return _SELECTED_MODEL_FAMILY
|
|
|
|
|
|
def select_model(line_count: int) -> str:
|
|
"""Return model configured for the selected family (line_count kept for compatibility)."""
|
|
_ = line_count
|
|
return _MODEL_PAYLOADS[_SELECTED_MODEL_FAMILY]["model"]
|
|
|
|
|
|
async def _read_streaming_content(response: httpx.Response) -> str:
|
|
content_parts: list[str] = []
|
|
async for raw_line in response.aiter_lines():
|
|
line = raw_line.strip()
|
|
if not line or line == "data: [DONE]" or not line.startswith("data: "):
|
|
continue
|
|
payload = line[len("data: "):]
|
|
try:
|
|
data = json.loads(payload)
|
|
except ValueError:
|
|
continue
|
|
choices = data.get("choices", [])
|
|
if not choices:
|
|
continue
|
|
delta = choices[0].get("delta", {})
|
|
token = delta.get("content")
|
|
if token:
|
|
content_parts.append(token)
|
|
return "".join(content_parts)
|
|
|
|
|
|
async def _send_messages(
|
|
client: httpx.AsyncClient,
|
|
model: str,
|
|
messages: list[dict],
|
|
) -> str:
|
|
"""Core HTTP helper: POST *messages* to the chat-completions endpoint.
|
|
|
|
Raises ``httpx.HTTPStatusError`` on non-2xx responses so the caller can
|
|
handle retries / quarantine.
|
|
"""
|
|
family = _MODEL_TO_FAMILY.get(model)
|
|
if not family:
|
|
raise ValueError(f"Unsupported model id: {model}")
|
|
model_payload = _MODEL_PAYLOADS[family]
|
|
payload = {**model_payload, "messages": messages, "model": model}
|
|
headers = {
|
|
"Authorization": "Bearer " + API_KEY,
|
|
"Content-Type": "application/json",
|
|
}
|
|
if payload.get("stream"):
|
|
async with client.stream(
|
|
"POST",
|
|
BASE_URL,
|
|
json=payload,
|
|
headers=headers,
|
|
timeout=API_TIMEOUT,
|
|
) as response:
|
|
response.raise_for_status()
|
|
return await _read_streaming_content(response)
|
|
response = await client.post(BASE_URL, json=payload, headers=headers, timeout=API_TIMEOUT)
|
|
response.raise_for_status()
|
|
return response.json()["choices"][0]["message"]["content"]
|
|
|
|
|
|
async def call_api(
|
|
client: httpx.AsyncClient,
|
|
model: str,
|
|
system_prompt: str,
|
|
user_content: str,
|
|
) -> str:
|
|
"""Stateless single-turn call (backward-compatible).
|
|
|
|
Builds a minimal ``[system, user]`` message list and delegates to
|
|
:func:`_send_messages`.
|
|
"""
|
|
messages: list[dict] = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_content},
|
|
]
|
|
return await _send_messages(client, model, messages)
|
|
|
|
|
|
async def call_api_multiturn(
|
|
client: httpx.AsyncClient,
|
|
model: str,
|
|
messages: list[dict],
|
|
) -> str:
|
|
"""Multi-turn call: pass the full conversation history and return the next reply.
|
|
|
|
The caller is responsible for maintaining and extending *messages*
|
|
(appending ``{"role": "assistant", "content": reply}`` after each call).
|
|
"""
|
|
return await _send_messages(client, model, messages)
|