from pydantic import BaseModel, Field
from typing import Optional, Union, Generator, Iterator

import os
import requests


class Pipe:
    class Valves(BaseModel):
        NAME_PREFIX: str = Field(
            default="",
            description="Prefix to be added before model names.",
        )
        OPENAI_API_BASE_URL: str = Field(
            default="",
            description="Base URL for accessing OpenAI API endpoints.",
        )
        OPENAI_API_KEY: str = Field(
            default="",
            description="",
        )
        pass

    def __init__(self):
        self.type = "manifold"
        self.valves = self.Valves()
        pass

    def pipes(self):
        return [
            {
                "id": "qwen-long",
                "name": "qwen-long",
            },
            {
                "id": "qwen-turbo",
                "name": "qwen-turbo（便宜）",
            },
            {
                "id": "qwen-plus",
                "name": "qwen-plus",
            },
            {
                "id": "qwen-max",
                "name": "qwen-max（最贵）",
            },
            {
                "id": "qwen-vl-max",
                "name": "qwen-vl-max（视觉模型）（最贵）",
            },
            {
                "id": "qwen-vl-max",
                "name": "qwen-vl-plus（视觉模型）",
            },
            {"id": "llama3.1-405b-instruct", "name": "llama3.1-405b-instruct(free)"},
            {
                "id": "qwen2.5-coder-1.5b-instruct",
                "name": "qwen2.5-coder-1.5b-instruct(free)",
            },
            {"id": "llama3-70b-instruct", "name": "llama3-70b-instruct(free)"},
            {"id": "qwen-coder-turbo", "name": "qwen-coder-turbo（代码）"},
            {"id": "qwen-math-plus", "name": "qwen-math-plus（数学）"},
        ]

    def pipe(self, body: dict, __user__: dict) -> Union[str, Generator, Iterator]:
        # This is where you can add your custom pipelines like RAG.
        print(f"pipe:{__name__}")
        headers = {}
        headers["Authorization"] = f"Bearer {self.valves.OPENAI_API_KEY}"
        headers["Content-Type"] = "application/json"

        model_id = body["model"][body["model"].find(".") + 1 :]
        payload = {**body, "model": model_id}
        print(payload)

        try:
            r = requests.post(
                url=f"{self.valves.OPENAI_API_BASE_URL}/chat/completions",
                json=payload,
                headers=headers,
                stream=True,
            )

            r.raise_for_status()

            if body["stream"]:
                return r.iter_lines()
            else:
                return r.json()
        except Exception as e:
            return f"Error: {e}"