Spaces:
Sleeping
Sleeping
| from litellm import completion, _turn_on_debug | |
| from dotenv import load_dotenv | |
| import random | |
| import logging | |
| load_dotenv() | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| models_low=["gemini/gemini-2.0-flash","openai/gpt-4.1-nano"] | |
| models_mid=["gemini/gemini-2.5-flash-preview-05-20","openai/gpt-4.1-mini"] | |
| models_high=["gemini/gemini-2.5-pro-preview-05-06","openai/gpt-4.1"] | |
| model_low="openai/gpt-4.1-nano" | |
| model_mid="openai/gpt-4.1-mini" | |
| model_high="openai/gpt-4.1" | |
| def call_llm(prompt: str, temperature: float, model_type: str, response_format=None, tools=None, shuffle=False, return_tokens=False, images=None) -> str: | |
| if shuffle: | |
| if model_type=="low": | |
| model = random.choice(models_low) | |
| elif model_type=="mid": | |
| model = random.choice(models_mid) | |
| elif model_type=="high": | |
| model = random.choice(models_high) | |
| logger.info(f"SHUFFLE. Using model: {model}") | |
| else: | |
| if model_type=="low": | |
| model = model_low | |
| elif model_type=="mid": | |
| model = model_mid | |
| elif model_type=="high": | |
| model = model_high | |
| # Create message content - support both text-only and multimodal | |
| if images: | |
| # Multimodal message with images | |
| content = [{"type": "text", "text": prompt}] | |
| for image_url in images: | |
| content.append({ | |
| "type": "image_url", | |
| "image_url": {"url": image_url} | |
| }) | |
| messages = [{"role": "user", "content": content}] | |
| else: | |
| # Text-only message | |
| messages = [{"role": "user", "content": prompt}] | |
| completion_args = { | |
| "model": model, | |
| "messages": messages, | |
| "temperature": temperature | |
| } | |
| if response_format: | |
| completion_args["response_format"] = response_format | |
| if tools: | |
| completion_args["tools"] = tools | |
| response = completion(**completion_args) | |
| response_str = response.choices[0].message.content | |
| if return_tokens: | |
| output_tokens = response.usage.completion_tokens | |
| input_tokens = response.usage.prompt_tokens | |
| return response_str,input_tokens,output_tokens | |
| else: | |
| return response_str |