mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
import json
|
|
from typing import Optional
|
|
|
|
import litellm
|
|
from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig
|
|
from litellm.types.llms.databricks import GenericStreamingChunk
|
|
|
|
|
|
class CodestralTextCompletionConfig(OpenAITextCompletionConfig):
|
|
"""
|
|
Reference: https://docs.mistral.ai/api/#operation/createFIMCompletion
|
|
"""
|
|
|
|
suffix: Optional[str] = None
|
|
temperature: Optional[int] = None
|
|
max_tokens: Optional[int] = None
|
|
min_tokens: Optional[int] = None
|
|
stream: Optional[bool] = None
|
|
random_seed: Optional[int] = None
|
|
|
|
def __init__(
|
|
self,
|
|
suffix: Optional[str] = None,
|
|
temperature: Optional[int] = None,
|
|
top_p: Optional[float] = None,
|
|
max_tokens: Optional[int] = None,
|
|
min_tokens: Optional[int] = None,
|
|
stream: Optional[bool] = None,
|
|
random_seed: Optional[int] = None,
|
|
stop: Optional[str] = None,
|
|
) -> None:
|
|
locals_ = locals().copy()
|
|
for key, value in locals_.items():
|
|
if key != "self" and value is not None:
|
|
setattr(self.__class__, key, value)
|
|
|
|
@classmethod
|
|
def get_config(cls):
|
|
return super().get_config()
|
|
|
|
def get_supported_openai_params(self, model: str):
|
|
return [
|
|
"suffix",
|
|
"temperature",
|
|
"top_p",
|
|
"max_tokens",
|
|
"max_completion_tokens",
|
|
"stream",
|
|
"seed",
|
|
"stop",
|
|
]
|
|
|
|
def map_openai_params(
|
|
self,
|
|
non_default_params: dict,
|
|
optional_params: dict,
|
|
model: str,
|
|
drop_params: bool,
|
|
) -> dict:
|
|
for param, value in non_default_params.items():
|
|
if param == "suffix":
|
|
optional_params["suffix"] = value
|
|
if param == "temperature":
|
|
optional_params["temperature"] = value
|
|
if param == "top_p":
|
|
optional_params["top_p"] = value
|
|
if param == "max_tokens" or param == "max_completion_tokens":
|
|
optional_params["max_tokens"] = value
|
|
if param == "stream" and value is True:
|
|
optional_params["stream"] = value
|
|
if param == "stop":
|
|
optional_params["stop"] = value
|
|
if param == "seed":
|
|
optional_params["random_seed"] = value
|
|
if param == "min_tokens":
|
|
optional_params["min_tokens"] = value
|
|
|
|
return optional_params
|
|
|
|
def _chunk_parser(self, chunk_data: str) -> GenericStreamingChunk:
|
|
text = ""
|
|
is_finished = False
|
|
finish_reason = None
|
|
logprobs = None
|
|
|
|
chunk_data = (
|
|
litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk_data) or ""
|
|
)
|
|
chunk_data = chunk_data.strip()
|
|
if len(chunk_data) == 0 or chunk_data == "[DONE]":
|
|
return {
|
|
"text": "",
|
|
"is_finished": is_finished,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
try:
|
|
chunk_data_dict = json.loads(chunk_data)
|
|
except json.JSONDecodeError:
|
|
return {
|
|
"text": "",
|
|
"is_finished": is_finished,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
|
|
original_chunk = litellm.ModelResponse(**chunk_data_dict, stream=True)
|
|
_choices = chunk_data_dict.get("choices", []) or []
|
|
_choice = _choices[0]
|
|
text = _choice.get("delta", {}).get("content", "")
|
|
|
|
if _choice.get("finish_reason") is not None:
|
|
is_finished = True
|
|
finish_reason = _choice.get("finish_reason")
|
|
logprobs = _choice.get("logprobs")
|
|
|
|
return GenericStreamingChunk(
|
|
text=text,
|
|
original_chunk=original_chunk,
|
|
is_finished=is_finished,
|
|
finish_reason=finish_reason,
|
|
logprobs=logprobs,
|
|
)
|