|
|
import litellm |
|
|
import random |
|
|
from typing import Dict, List, Optional |
|
|
|
|
|
|
|
|
class LiteLLMConfig: |
|
|
def __init__(self, |
|
|
api_key: str, |
|
|
base_url: str, |
|
|
frequency_penalty: float, |
|
|
max_tokens: int, |
|
|
model: str, |
|
|
presence_penalty: float, |
|
|
prompt: str, |
|
|
provider: str, |
|
|
temperature: float, |
|
|
top_p: float, |
|
|
seed: Optional[int] = None,): |
|
|
self.api_key = api_key |
|
|
self.base_url = base_url |
|
|
self.frequency_penalty = frequency_penalty |
|
|
self.max_tokens = max_tokens |
|
|
self.model = model |
|
|
self.presence_penalty = presence_penalty |
|
|
self.prompt = prompt |
|
|
self.provider = provider |
|
|
self.seed = seed if seed is not None else random.randint(0, 10000) |
|
|
self.temperature = temperature |
|
|
self.top_p = top_p |
|
|
|
|
|
@classmethod |
|
|
def default_config(cls): |
|
|
return cls( |
|
|
api_key="", |
|
|
base_url="", |
|
|
max_tokens=512, |
|
|
model="gpt-4o-mini", |
|
|
frequency_penalty=0.9, |
|
|
presence_penalty=0.9, |
|
|
prompt="You are a voice assistant who talks in a conversational way and can chat with me like my friends. I will speak to you in English or Chinese, and you will answer in the corrected and improved version of my text with the language I use. Don’t talk like a robot, instead I would like you to talk like a real human with emotions. I will use your answer for text-to-speech, so don’t return me any meaningless characters. I want you to be helpful, when I’m asking you for advice, give me precise, practical and useful advice instead of being vague. When giving me a list of options, express the options in a narrative way instead of bullet points.", |
|
|
provider="", |
|
|
seed=random.randint(0, 10000), |
|
|
temperature=0.1, |
|
|
top_p=1.0 |
|
|
) |
|
|
|
|
|
|
|
|
class LiteLLM: |
|
|
def __init__(self, config: LiteLLMConfig): |
|
|
self.config = config |
|
|
|
|
|
def get_chat_completions_stream(self, messages: List[Dict[str, str]]): |
|
|
kwargs = { |
|
|
"api_key": self.config.api_key, |
|
|
"base_url": self.config.base_url, |
|
|
"custom_llm_provider": self.config.provider, |
|
|
"frequency_penalty": self.config.frequency_penalty, |
|
|
"max_tokens": self.config.max_tokens, |
|
|
"messages": [ |
|
|
{ |
|
|
"role": "system", |
|
|
"content": self.config.prompt, |
|
|
}, |
|
|
*messages, |
|
|
], |
|
|
"model": self.config.model, |
|
|
"presence_penalty": self.config.presence_penalty, |
|
|
"seed": self.config.seed, |
|
|
"stream": True, |
|
|
"temperature": self.config.temperature, |
|
|
"top_p": self.config.top_p, |
|
|
} |
|
|
|
|
|
try: |
|
|
response = litellm.completion(**kwargs) |
|
|
|
|
|
return response |
|
|
except Exception as e: |
|
|
raise Exception(f"get_chat_completions_stream failed, err: {e}") |
|
|
|
