import json
from openai import OpenAI
class ModelRouter:
    def __init__(self, base_url, api_key):
        self.client = OpenAI(
            base_url=base_url,
            api_key=api_key,  # ModelScope Token
            timeout=18000,
        )

    def select_model(self, question, available_models):
        # Prepare the prompt with dynamic question and models
        prompt = f"""
You are a model router designed to select the most suitable AI model from a given list of models ({json.dumps(available_models)}) based on the input question ({json.dumps(question)}). The selection should be made by evaluating the model's strengths across six dimensions: coding, knowledge, math, reasoning, roleplay, and writing. Use the following performance insights to guide your decision:

- gpt4o: Strong in writing, with moderate general knowledge.
- Qwen2.5-7B-Instruct: Excels in reasoning and roleplay, with solid writing skills.
- gemma-3-4b-it: Good at roleplay and writing, with decent reasoning capabilities.
- Kimi-K2-Instruct: Stands out in coding and knowledge, with strong roleplay performance.
- Llama-3.1-8B-Instruct: Effective in writing, with reasonable knowledge and reasoning.
- GPT5: Highly capable in reasoning, knowledge, and math, with good overall balance.
- gpt-3.5-turbo-1106: Strong in writing and math, with moderate coding skills.
- claude-3.7-sonnet-thinking: Excels in knowledge, reasoning, and coding.
- o3mini: Outstanding in reasoning and coding, with strong writing skills.
- deepseek-r1: Strong in coding and writing, with good reasoning and roleplay.
- qwen2.5-72b-instruct: Balanced performance, particularly strong in math and writing.
- gemini-2.5-pro: Exceptional in math, with solid coding and reasoning.
- gemma-3-27b-it: Excels in coding and reasoning, with good writing skills.
- Qwen3-32B: Strong in math and reasoning, with balanced overall performance.
- claude-3.5-sonnet-20241022: Outstanding in roleplay and knowledge, with strong reasoning.
- QwQ-32B: Good in writing and math, with solid reasoning skills.
- llama3.3-70B-instruct: Excels in writing, with decent roleplay and reasoning.
- doubao-1-5-thinking-pro-250415: Exceptional in writing, with good roleplay skills.
- gemini-2.5-flash: Top performer in math, writing, and reasoning, with strong overall ability.

### Routing Logic:
1. Analyze the {question} to identify the primary focus (e.g., coding, knowledge, math, reasoning, roleplay, or writing). If the question involves multiple aspects, prioritize based on the most dominant theme.
2. Select the model with the strongest performance in the identified focus area. If multiple models are strong, consider their overall balance across dimensions.
3. If no clear focus is identifiable, choose the model with the broadest and strongest overall capabilities.
4. Return the selected model's name as the output.

### Output Format:
- Return only the model name (e.g., "gemini-2.5-flash") with no additional text or explanation.
"""

        # Prepare messages for the API call
        messages = [{"role": "system", "content": prompt},
                   {"role": "user", "content": f"Question: {question}, Available models: {json.dumps(available_models)}"}]

        # Call the API to get the model recommendation
        response = self.client.chat.completions.create(
            model="deepseek-r1",  # Replace with appropriate ModelScope Model-Id
            messages=messages,
            temperature=0.7,
            stream=False,
        )

        # Extract and return the recommended model name
        return response.choices[0].message.content.strip()

# Example usage
if __name__ == "__main__":
    router = ModelRouter(base_url="", api_key="adjlhadjkl")
    question = ""
    models = ["gpt4o", "Qwen2.5-7B-Instruct", "gemma-3-4b-it", "Kimi-K2-Instruct", "Llama-3.1-8B-Instruct", "GPT5", "gpt-3.5-turbo-1106", "claude-3.7-sonnet-thinking", "o3mini", "deepseek-r1", "qwen2.5-72b-instruct", "gemini-2.5-pro", "gemma-3-27b-it", "Qwen3-32B", "claude-3.5-sonnet-20241022", "QwQ-32B", "llama3.3-70B-instruct", "doubao-1-5-thinking-pro-250415", "gemini-2.5-flash"]
    selected_model = router.select_model(question, models)
    print(f"Selected model: {selected_model}")