app.ai_providers.kimi_provider

  1"""Moonshot Kimi API provider implementation.
  2
  3Uses the ``openai`` Python SDK pointed at the Moonshot Kimi API base URL.
  4Supports synchronous and streaming generation, CSV file attachments via
  5the Responses API, and automatic model-alias mapping for deprecated Kimi
  6model identifiers.
  7
  8Attributes:
  9    logger: Module-level logger for Kimi provider operations.
 10"""
 11
 12from __future__ import annotations
 13
 14import logging
 15from typing import Any, Callable, Iterator, Mapping
 16
 17from .base import (
 18    AIProvider,
 19    AIProviderError,
 20    DEFAULT_CLOUD_REQUEST_TIMEOUT_SECONDS,
 21    DEFAULT_KIMI_BASE_URL,
 22    DEFAULT_KIMI_FILE_UPLOAD_PURPOSE,
 23    DEFAULT_KIMI_MODEL,
 24    DEFAULT_MAX_TOKENS,
 25    _is_attachment_unsupported_error,
 26    _is_context_length_error,
 27    _is_kimi_model_not_available_error,
 28    _normalize_api_key_value,
 29    _normalize_kimi_model_name,
 30    _normalize_openai_compatible_base_url,
 31    _resolve_timeout_seconds,
 32    _run_with_rate_limit_retries,
 33    _T,
 34)
 35from .utils import (
 36    _extract_openai_delta_text,
 37    _extract_openai_text,
 38    _inline_attachment_data_into_prompt,
 39    upload_and_request_via_responses_api,
 40)
 41
 42logger = logging.getLogger(__name__)
 43
 44
 45class KimiProvider(AIProvider):
 46    """Moonshot Kimi API provider implementation.
 47
 48    Attributes:
 49        api_key (str): The Moonshot/Kimi API key.
 50        model (str): The Kimi model identifier.
 51        base_url (str): The normalized Kimi API base URL.
 52        attach_csv_as_file (bool): Whether to upload CSV artifacts as
 53            file attachments.
 54        request_timeout_seconds (float): HTTP timeout in seconds.
 55        client: The ``openai.OpenAI`` SDK client instance configured for Kimi.
 56    """
 57
 58    def __init__(
 59        self,
 60        api_key: str,
 61        model: str = DEFAULT_KIMI_MODEL,
 62        base_url: str = DEFAULT_KIMI_BASE_URL,
 63        attach_csv_as_file: bool = True,
 64        request_timeout_seconds: float = DEFAULT_CLOUD_REQUEST_TIMEOUT_SECONDS,
 65    ) -> None:
 66        """Initialize the Kimi provider.
 67
 68        Args:
 69            api_key: Moonshot/Kimi API key. Must be non-empty.
 70            model: Kimi model identifier. Deprecated aliases are mapped.
 71            base_url: Kimi API base URL.
 72            attach_csv_as_file: If ``True``, attempt file uploads.
 73            request_timeout_seconds: HTTP timeout in seconds.
 74
 75        Raises:
 76            AIProviderError: If the ``openai`` SDK is not installed or
 77                the API key is empty.
 78        """
 79        try:
 80            import openai
 81        except ImportError as error:
 82            raise AIProviderError(
 83                "openai SDK is not installed. Install it with `pip install openai`."
 84            ) from error
 85
 86        normalized_api_key = _normalize_api_key_value(api_key)
 87        if not normalized_api_key:
 88            raise AIProviderError(
 89                "Kimi API key is not configured. "
 90                "Set `ai.kimi.api_key` in config.yaml or the MOONSHOT_API_KEY environment variable."
 91            )
 92
 93        self._openai = openai
 94        self.api_key = normalized_api_key
 95        self.model = _normalize_kimi_model_name(model)
 96        self.base_url = _normalize_openai_compatible_base_url(
 97            base_url=base_url,
 98            default_base_url=DEFAULT_KIMI_BASE_URL,
 99        )
100        self.attach_csv_as_file = bool(attach_csv_as_file)
101        self._csv_attachment_supported: bool | None = None
102        self.request_timeout_seconds = _resolve_timeout_seconds(
103            request_timeout_seconds,
104            DEFAULT_CLOUD_REQUEST_TIMEOUT_SECONDS,
105        )
106        self.client = openai.OpenAI(
107            api_key=normalized_api_key,
108            base_url=self.base_url,
109            timeout=self.request_timeout_seconds,
110        )
111        logger.info("Initialized Kimi provider at %s with model %s (timeout %.1fs)", self.base_url, self.model, self.request_timeout_seconds)
112
113    def analyze(
114        self,
115        system_prompt: str,
116        user_prompt: str,
117        max_tokens: int = DEFAULT_MAX_TOKENS,
118    ) -> str:
119        """Send a prompt to Kimi and return the generated text.
120
121        Args:
122            system_prompt: The system-level instruction text.
123            user_prompt: The user-facing prompt with investigation context.
124            max_tokens: Maximum completion tokens.
125
126        Returns:
127            The generated analysis text.
128
129        Raises:
130            AIProviderError: On any API or network failure.
131        """
132        return self.analyze_with_attachments(
133            system_prompt=system_prompt,
134            user_prompt=user_prompt,
135            attachments=None,
136            max_tokens=max_tokens,
137        )
138
139    def analyze_stream(
140        self,
141        system_prompt: str,
142        user_prompt: str,
143        max_tokens: int = DEFAULT_MAX_TOKENS,
144    ) -> Iterator[str]:
145        """Stream generated text chunks from Kimi.
146
147        Args:
148            system_prompt: The system-level instruction text.
149            user_prompt: The user-facing prompt with investigation context.
150            max_tokens: Maximum completion tokens.
151
152        Yields:
153            Text chunk strings as they are generated.
154
155        Raises:
156            AIProviderError: On empty response or API failure.
157        """
158        def _stream() -> Iterator[str]:
159            stream = self._run_kimi_request(
160                lambda: self.client.chat.completions.create(
161                    model=self.model,
162                    max_tokens=max_tokens,
163                    messages=[
164                        {"role": "system", "content": system_prompt},
165                        {"role": "user", "content": user_prompt},
166                    ],
167                    stream=True,
168                )
169            )
170            emitted = False
171            try:
172                for chunk in stream:
173                    choices = getattr(chunk, "choices", None)
174                    if not choices:
175                        continue
176                    choice = choices[0]
177                    delta = getattr(choice, "delta", None)
178                    if delta is None and isinstance(choice, dict):
179                        delta = choice.get("delta")
180                    chunk_text = _extract_openai_delta_text(
181                        delta,
182                        ("content", "reasoning_content", "reasoning", "refusal"),
183                    )
184                    if not chunk_text:
185                        continue
186                    emitted = True
187                    yield chunk_text
188            except AIProviderError:
189                raise
190            except self._openai.APIConnectionError as error:
191                raise AIProviderError(
192                    "Unable to connect to Kimi API. Check `ai.kimi.base_url` and network access."
193                ) from error
194            except self._openai.AuthenticationError as error:
195                raise AIProviderError(
196                    "Kimi authentication failed. Check `ai.kimi.api_key`, MOONSHOT_API_KEY, or KIMI_API_KEY."
197                ) from error
198            except self._openai.BadRequestError as error:
199                if _is_context_length_error(error):
200                    raise AIProviderError(
201                        "Kimi request exceeded the model context length. Reduce prompt size and retry."
202                    ) from error
203                raise AIProviderError(f"Kimi request was rejected: {error}") from error
204            except self._openai.APIError as error:
205                if _is_kimi_model_not_available_error(error):
206                    raise AIProviderError(
207                        "Kimi rejected the configured model. "
208                        f"Current model: `{self.model}`. "
209                        "Set `ai.kimi.model` to a model enabled for your Moonshot account "
210                        "(for example `kimi-k2-turbo-preview`) and retry."
211                    ) from error
212                raise AIProviderError(f"Kimi API error: {error}") from error
213            except Exception as error:
214                raise AIProviderError(f"Unexpected Kimi provider error: {error}") from error
215
216            if not emitted:
217                raise AIProviderError("Kimi returned an empty response.")
218
219        return _stream()
220
221    def analyze_with_attachments(
222        self,
223        system_prompt: str,
224        user_prompt: str,
225        attachments: list[Mapping[str, str]] | None,
226        max_tokens: int = DEFAULT_MAX_TOKENS,
227    ) -> str:
228        """Analyze with optional CSV file attachments via the Kimi Responses API.
229
230        Args:
231            system_prompt: The system-level instruction text.
232            user_prompt: The user-facing prompt with investigation context.
233            attachments: Optional list of attachment descriptors.
234            max_tokens: Maximum completion tokens.
235
236        Returns:
237            The generated analysis text.
238
239        Raises:
240            AIProviderError: On any API or network failure.
241        """
242        def _request() -> str:
243            return self._request_non_stream(
244                system_prompt=system_prompt,
245                user_prompt=user_prompt,
246                max_tokens=max_tokens,
247                attachments=attachments,
248            )
249
250        return self._run_kimi_request(_request)
251
252    def _run_kimi_request(self, request_fn: Callable[[], _T]) -> _T:
253        """Execute a Kimi request with rate-limit retries and error mapping.
254
255        Args:
256            request_fn: A zero-argument callable that performs the request.
257
258        Returns:
259            The return value of ``request_fn`` on success.
260
261        Raises:
262            AIProviderError: On any OpenAI SDK error (with Kimi messages).
263        """
264        try:
265            return _run_with_rate_limit_retries(
266                request_fn=request_fn,
267                rate_limit_error_type=self._openai.RateLimitError,
268                provider_name="Kimi",
269            )
270        except AIProviderError:
271            raise
272        except self._openai.APIConnectionError as error:
273            raise AIProviderError(
274                "Unable to connect to Kimi API. Check `ai.kimi.base_url` and network access."
275            ) from error
276        except self._openai.AuthenticationError as error:
277            raise AIProviderError(
278                "Kimi authentication failed. Check `ai.kimi.api_key`, MOONSHOT_API_KEY, or KIMI_API_KEY."
279            ) from error
280        except self._openai.BadRequestError as error:
281            if _is_context_length_error(error):
282                raise AIProviderError(
283                    "Kimi request exceeded the model context length. Reduce prompt size and retry."
284                ) from error
285            raise AIProviderError(f"Kimi request was rejected: {error}") from error
286        except self._openai.APIError as error:
287            if _is_kimi_model_not_available_error(error):
288                raise AIProviderError(
289                    "Kimi rejected the configured model. "
290                    f"Current model: `{self.model}`. "
291                    "Set `ai.kimi.model` to a model enabled for your Moonshot account "
292                    "(for example `kimi-k2-turbo-preview`) and retry."
293                ) from error
294            raise AIProviderError(f"Kimi API error: {error}") from error
295        except Exception as error:
296            raise AIProviderError(f"Unexpected Kimi provider error: {error}") from error
297
298    def _request_non_stream(
299        self,
300        system_prompt: str,
301        user_prompt: str,
302        max_tokens: int,
303        attachments: list[Mapping[str, str]] | None = None,
304    ) -> str:
305        """Perform a non-streaming Kimi request with attachment handling.
306
307        Args:
308            system_prompt: The system-level instruction text.
309            user_prompt: The user-facing prompt text.
310            max_tokens: Maximum completion tokens.
311            attachments: Optional list of attachment descriptors.
312
313        Returns:
314            The generated analysis text.
315
316        Raises:
317            AIProviderError: If the response is empty.
318        """
319        attachment_response = self._request_with_csv_attachments(
320            system_prompt=system_prompt,
321            user_prompt=user_prompt,
322            max_tokens=max_tokens,
323            attachments=attachments,
324        )
325        if attachment_response:
326            return attachment_response
327
328        prompt_for_completion = user_prompt
329        if attachments:
330            prompt_for_completion, inlined = _inline_attachment_data_into_prompt(
331                user_prompt=user_prompt,
332                attachments=attachments,
333            )
334            if inlined:
335                logger.info("Kimi attachment fallback inlined attachment data into prompt.")
336
337        response = self.client.chat.completions.create(
338            model=self.model,
339            max_tokens=max_tokens,
340            messages=[
341                {"role": "system", "content": system_prompt},
342                {"role": "user", "content": prompt_for_completion},
343            ],
344        )
345        text = _extract_openai_text(response)
346        if not text:
347            raise AIProviderError("Kimi returned an empty response.")
348        return text
349
350    def _request_with_csv_attachments(
351        self,
352        system_prompt: str,
353        user_prompt: str,
354        max_tokens: int,
355        attachments: list[Mapping[str, str]] | None,
356    ) -> str | None:
357        """Attempt to send a request with CSV files via the Kimi Responses API.
358
359        Args:
360            system_prompt: The system-level instruction text.
361            user_prompt: The user-facing prompt text.
362            max_tokens: Maximum completion tokens.
363            attachments: Optional list of attachment descriptors.
364
365        Returns:
366            The generated text if succeeded, or ``None`` if skipped.
367        """
368        normalized_attachments = self._prepare_csv_attachments(
369            attachments,
370            supports_file_attachments=hasattr(self.client, "files") and hasattr(self.client, "responses"),
371        )
372        if not normalized_attachments:
373            return None
374
375        try:
376            text = upload_and_request_via_responses_api(
377                client=self.client,
378                openai_module=self._openai,
379                model=self.model,
380                normalized_attachments=normalized_attachments,
381                system_prompt=system_prompt,
382                user_prompt=user_prompt,
383                max_tokens=max_tokens,
384                provider_name="Kimi",
385                upload_purpose=DEFAULT_KIMI_FILE_UPLOAD_PURPOSE,
386                convert_csv_to_txt=False,
387            )
388            self._csv_attachment_supported = True
389            return text
390        except Exception as error:
391            if _is_attachment_unsupported_error(error):
392                self._csv_attachment_supported = False
393                logger.info(
394                    "Kimi endpoint does not support CSV attachments via /files + /responses; "
395                    "falling back to chat.completions text mode."
396                )
397                return None
398            raise
399
400    def get_model_info(self) -> dict[str, str]:
401        """Return Kimi provider and model metadata.
402
403        Returns:
404            A dictionary with ``"provider"`` and ``"model"`` keys.
405        """
406        return {"provider": "kimi", "model": self.model}