Chat with a model using LiteLLM. Args: model: Model identifier (e.g., "gpt-3.5-turbo", "claude-3-sonnet-20240229") messages: List of message dictionaries with 'role' and 'content' keys temperature: Sampling temperature (0.0 to 2.0) ma
(
self,
model: str,
messages: List[Dict[str, str]],
temperature: Optional[float] = None,
max_tokens: Optional[int] = None,
top_p: Optional[float] = None,
stream: bool = False,
**kwargs
)
| 21 | self.config = kwargs |
| 22 | |
| 23 | def chat( |
| 24 | self, |
| 25 | model: str, |
| 26 | messages: List[Dict[str, str]], |
| 27 | temperature: Optional[float] = None, |
| 28 | max_tokens: Optional[int] = None, |
| 29 | top_p: Optional[float] = None, |
| 30 | stream: bool = False, |
| 31 | **kwargs |
| 32 | ) -> Any: |
| 33 | """ |
| 34 | Chat with a model using LiteLLM. |
| 35 | |
| 36 | Args: |
| 37 | model: Model identifier (e.g., "gpt-3.5-turbo", "claude-3-sonnet-20240229") |
| 38 | messages: List of message dictionaries with 'role' and 'content' keys |
| 39 | temperature: Sampling temperature (0.0 to 2.0) |
| 40 | max_tokens: Maximum tokens to generate |
| 41 | top_p: Nucleus sampling parameter |
| 42 | stream: Whether to stream the response |
| 43 | **kwargs: Additional provider-specific parameters |
| 44 | |
| 45 | Returns: |
| 46 | Response object from LiteLLM completion |
| 47 | |
| 48 | Example: |
| 49 | >>> llm = LiteLLMLLM(api_key="your-key") |
| 50 | >>> response = llm.chat( |
| 51 | ... model="gpt-3.5-turbo", |
| 52 | ... messages=[{"role": "user", "content": "Hello!"}] |
| 53 | ... ) |
| 54 | """ |
| 55 | # Set API key if provided |
| 56 | if self.api_key: |
| 57 | # LiteLLM uses environment variables for authentication |
| 58 | # The specific env var depends on the provider |
| 59 | # For now, we'll set OPENAI_API_KEY as a common case |
| 60 | # Users can override by setting the appropriate env var themselves |
| 61 | if "OPENAI_API_KEY" not in os.environ: |
| 62 | os.environ["OPENAI_API_KEY"] = self.api_key |
| 63 | |
| 64 | # Build completion parameters |
| 65 | completion_params = { |
| 66 | "model": model, |
| 67 | "messages": messages, |
| 68 | } |
| 69 | |
| 70 | # Add optional parameters if provided |
| 71 | if temperature is not None: |
| 72 | completion_params["temperature"] = temperature |
| 73 | if max_tokens is not None: |
| 74 | completion_params["max_tokens"] = max_tokens |
| 75 | if top_p is not None: |
| 76 | completion_params["top_p"] = top_p |
| 77 | if stream: |
| 78 | completion_params["stream"] = stream |
| 79 | |
| 80 | # Merge with any additional kwargs |