Adapt to Gemini

.gitignore

@@ -1,3 +1,4 @@
 api.json
 api.yaml
-.env
+.env
+__pycache__

json_str/gemini/request.json

@@ -0,0 +1,52 @@
+{
+    "contents": [
+        {
+            "role": "user",
+            "parts": [
+                {
+                    "text": "hi"
+                }
+            ]
+        },
+        {
+            "role": "model",
+            "parts": [
+                {
+                    "text": "Hi! \n\nHow are you today? What can I do for you? \n"
+                }
+            ]
+        },
+        {
+            "role": "user",
+            "parts": [
+                {
+                    "text": "怎么解决"
+                },
+                {
+                    "inlineData": {
+                        "mimeType": "image/jpeg",
+                        "data": "/9j/***"
+                    }
+                }
+            ]
+        }
+    ],
+    "safetySettings": [
+        {
+            "category": "HARM_CATEGORY_HARASSMENT",
+            "threshold": "BLOCK_NONE"
+        },
+        {
+            "category": "HARM_CATEGORY_HATE_SPEECH",
+            "threshold": "BLOCK_NONE"
+        },
+        {
+            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+            "threshold": "BLOCK_NONE"
+        },
+        {
+            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+            "threshold": "BLOCK_NONE"
+        }
+    ]
+}

main.py

@@ -2,29 +2,19 @@ import os
 import json
 import httpx
 import yaml
+import traceback
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, HTTPException, Depends
 from fastapi.responses import StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
-from pydantic import BaseModel
-from typing import List, Dict, Any, Optional, Union
+from models import RequestModel
+from request import get_payload
+from response import fetch_response, fetch_response_stream
 
-# 模拟存储API Key的数据库
-api_keys_db = {
-    "sk-KjjI60Yf0JFcsvgRmXqFwgGmWUd9GZnmi3KlvowmRWpWpQRo": "user1",
-    # 可以添加更多的API Key
-}
-
-# 安全性依赖
-security = HTTPBearer()
-
-def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    token = credentials.credentials
-    if token not in api_keys_db:
-        raise HTTPException(status_code=403, detail="Invalid or missing API Key")
-    return token
+from typing import List, Dict
+from urllib.parse import urlparse
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -36,6 +26,15 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
+# 模拟存储API Key的数据库
+api_keys_db = {
+    "sk-KjjI60Yf0JFcsvgRmXqFwgGmWUd9GZnmi3KlvowmRWpWpQRo": "user1",
+    # 可以添加更多的API Key
+}
+
+# 安全性依赖
+security = HTTPBearer()
+
 # 读取YAML配置文件
 def load_config():
     try:
@@ -51,107 +50,29 @@ def load_config():
 config = load_config()
 # print(config)
 
-# 定义 Function 参数模型
-class FunctionParameter(BaseModel):
-    type: str
-    properties: Dict[str, Dict[str, str]]
-    required: List[str]
-
-# 定义 Function 模型
-class Function(BaseModel):
-    name: str
-    description: str
-    parameters: FunctionParameter
-
-# 定义 Tool 模型
-class Tool(BaseModel):
-    type: str
-    function: Function
-
-class ImageUrl(BaseModel):
-    url: str
-
-class ContentItem(BaseModel):
-    type: str
-    text: Optional[str] = None
-    image_url: Optional[ImageUrl] = None
-
-class Message(BaseModel):
-    role: str
-    name: Optional[str] = None
-    content: Union[str, List[ContentItem]]
-
-class RequestModel(BaseModel):
-    model: str
-    messages: List[Message]
-    logprobs: Optional[bool] = None
-    top_logprobs: Optional[int] = None
-    stream: Optional[bool] = None
-    include_usage: Optional[bool] = None
-    temperature: Optional[float] = 0.5
-    top_p: Optional[float] = 1.0
-    max_tokens: Optional[int] = None
-    presence_penalty: Optional[float] = 0.0
-    frequency_penalty: Optional[float] = 0.0
-    n: Optional[int] = 1
-    user: Optional[str] = None
-    tool_choice: Optional[str] = None
-    tools: Optional[List[Tool]] = None
-
-async def fetch_response_stream(client, url, headers, payload):
-    async with client.stream('POST', url, headers=headers, json=payload) as response:
-        async for chunk in response.aiter_bytes():
-            print(chunk.decode('utf-8'))
-            yield chunk
-
-async def fetch_response(client, url, headers, payload):
-    response = await client.post(url, headers=headers, json=payload)
-    return response.json()
-
 async def process_request(request: RequestModel, provider: Dict):
     print("provider: ", provider['provider'])
     url = provider['base_url']
-    headers = {
-        'Authorization': f"Bearer {provider['api']}",
-        'Content-Type': 'application/json'
-    }
-
-    # 转换消息格式
-    messages = []
-    for msg in request.messages:
-        if isinstance(msg.content, list):
-            content = []
-            for item in msg.content:
-                if item.type == "text":
-                    content.append({"type": "text", "text": item.text})
-                elif item.type == "image_url":
-                    content.append({"type": "image_url", "image_url": item.image_url.dict()})
-        else:
-            content = msg.content
-            name = msg.name
-        if name:
-            messages.append({"role": msg.role, "name": name, "content": content})
-        else:
-            messages.append({"role": msg.role, "content": content})
-
-
-    payload = {
-        "model": request.model,
-        "messages": messages
-    }
-
-    for field, value in request.dict(exclude_unset=True).items():
-        if field not in ['model', 'messages'] and value is not None:
-            payload[field] = value
-
-    request_info = {
-        "url": url,
-        "headers": headers,
-        "payload": payload
-    }
-    print(f"Request details: {json.dumps(request_info, indent=2, ensure_ascii=False)}")
+    parsed_url = urlparse(url)
+    engine = None
+    if parsed_url.netloc == 'generativelanguage.googleapis.com':
+        engine = "gemini"
+    elif parsed_url.netloc == 'api.anthropic.com':
+        engine = "claude"
+    else:
+        engine = "gpt"
+
+    url, headers, payload = await get_payload(request, engine, provider)
+
+    # request_info = {
+    #     "url": url,
+    #     "headers": headers,
+    #     "payload": payload
+    # }
+    # print(f"Request details: {json.dumps(request_info, indent=2, ensure_ascii=False)}")
+
     if request.stream:
-        return StreamingResponse(fetch_response_stream(app.state.client, url, headers, payload), media_type="text/event-stream")
+        return StreamingResponse(fetch_response_stream(app.state.client, url, headers, payload, engine, request.model), media_type="text/event-stream")
     else:
         return await fetch_response(app.state.client, url, headers, payload)
 
@@ -191,13 +112,22 @@ class ModelRequestHandler:
                 response = await process_request(request, provider)
                 return response
             except Exception as e:
+                print('\033[31m')
                 print(f"Error with provider {provider['provider']}: {str(e)}")
+                traceback.print_exc()
+                print('\033[0m')
                 continue
 
         raise HTTPException(status_code=500, detail="All providers failed")
 
 model_handler = ModelRequestHandler()
 
+def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    token = credentials.credentials
+    if token not in api_keys_db:
+        raise HTTPException(status_code=403, detail="Invalid or missing API Key")
+    return token
+
 @app.post("/v1/chat/completions")
 async def request_model(request: RequestModel, token: str = Depends(verify_api_key)):
     return await model_handler.request_model(request, token)

models.py

@@ -0,0 +1,48 @@
+from pydantic import BaseModel
+from typing import List, Dict, Optional, Union
+
+class FunctionParameter(BaseModel):
+    type: str
+    properties: Dict[str, Dict[str, str]]
+    required: List[str]
+
+# 定义 Function 模型
+class Function(BaseModel):
+    name: str
+    description: str
+    parameters: FunctionParameter
+
+# 定义 Tool 模型
+class Tool(BaseModel):
+    type: str
+    function: Function
+
+class ImageUrl(BaseModel):
+    url: str
+
+class ContentItem(BaseModel):
+    type: str
+    text: Optional[str] = None
+    image_url: Optional[ImageUrl] = None
+
+class Message(BaseModel):
+    role: str
+    name: Optional[str] = None
+    content: Union[str, List[ContentItem]]
+
+class RequestModel(BaseModel):
+    model: str
+    messages: List[Message]
+    logprobs: Optional[bool] = None
+    top_logprobs: Optional[int] = None
+    stream: Optional[bool] = None
+    include_usage: Optional[bool] = None
+    temperature: Optional[float] = 0.5
+    top_p: Optional[float] = 1.0
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    n: Optional[int] = 1
+    user: Optional[str] = None
+    tool_choice: Optional[str] = None
+    tools: Optional[List[Tool]] = None

request.py

@@ -0,0 +1,163 @@
+from models import RequestModel
+
+async def get_image_message(base64_image, engine = None):
+    if "gpt" == engine:
+        return {
+            "type": "image_url",
+            "image_url": {
+                "url": base64_image,
+            }
+        }
+    if "claude" == engine:
+        return {
+            "type": "image",
+            "source": {
+                "type": "base64",
+                "media_type": "image/jpeg",
+                "data": base64_image.split(",")[1],
+            }
+        }
+    if "gemini" == engine:
+        return {
+            "inlineData": {
+                "mimeType": "image/jpeg",
+                "data": base64_image.split(",")[1],
+            }
+        }
+    raise ValueError("Unknown engine")
+
+async def get_text_message(role, message, engine = None):
+    if "gpt" == engine or "claude" == engine:
+        return {"type": "text", "text": message}
+    if "gemini" == engine:
+        return {"text": message}
+    raise ValueError("Unknown engine")
+
+async def get_gemini_payload(request, engine, provider):
+    headers = {
+        'Content-Type': 'application/json'
+    }
+    url = provider['base_url']
+    if request.stream:
+        gemini_stream = "streamGenerateContent"
+    url = url.format(model=request.model, stream=gemini_stream, api_key=provider['api'])
+
+    messages = []
+    for msg in request.messages:
+        if isinstance(msg.content, list):
+            content = []
+            for item in msg.content:
+                if item.type == "text":
+                    text_message = await get_text_message(msg.role, item.text, engine)
+                    # print("text_message", text_message)
+                    content.append(text_message)
+                elif item.type == "image_url":
+                    image_message = await get_image_message(item.image_url.url, engine)
+                    content.append(image_message)
+        else:
+            content = msg.content
+        if msg.role != "system":
+            messages.append({"role": msg.role, "parts": content})
+
+
+    payload = {
+        "contents": messages,
+        "safetySettings": [
+            {
+                "category": "HARM_CATEGORY_HARASSMENT",
+                "threshold": "BLOCK_NONE"
+            },
+            {
+                "category": "HARM_CATEGORY_HATE_SPEECH",
+                "threshold": "BLOCK_NONE"
+            },
+            {
+                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                "threshold": "BLOCK_NONE"
+            },
+            {
+                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                "threshold": "BLOCK_NONE"
+            }
+        ]
+    }
+
+    miss_fields = [
+        'model',
+        'messages',
+        'stream',
+        'tools',
+        'tool_choice',
+        'temperature',
+        'top_p',
+        'max_tokens',
+        'presence_penalty',
+        'frequency_penalty',
+        'n',
+        'user',
+        'include_usage',
+        'logprobs',
+        'top_logprobs'
+    ]
+
+    for field, value in request.model_dump(exclude_unset=True).items():
+        if field not in miss_fields and value is not None:
+            payload[field] = value
+
+    return url, headers, payload
+
+async def get_gpt_payload(request, engine, provider):
+    headers = {
+        'Authorization': f"Bearer {provider['api']}",
+        'Content-Type': 'application/json'
+    }
+    url = provider['base_url']
+    url = url.format(model=request.model, stream=request.stream, api_key=provider['api'])
+
+    messages = []
+    for msg in request.messages:
+        if isinstance(msg.content, list):
+            content = []
+            for item in msg.content:
+                if item.type == "text":
+                    text_message = await get_text_message(msg.role, item.text, engine)
+                    content.append(text_message)
+                elif item.type == "image_url":
+                    image_message = await get_image_message(item.image_url.url, engine)
+                    content.append(image_message)
+        else:
+            content = msg.content
+            name = msg.name
+        if name:
+            messages.append({"role": msg.role, "name": name, "content": content})
+        else:
+            messages.append({"role": msg.role, "content": content})
+
+    payload = {
+        "model": request.model,
+        "messages": messages,
+    }
+
+    miss_fields = [
+        'model',
+        'messages'
+    ]
+
+    for field, value in request.model_dump(exclude_unset=True).items():
+        if field not in miss_fields and value is not None:
+            payload[field] = value
+
+    return url, headers, payload
+
+async def get_claude_payload(request, engine, provider):
+    pass
+
+async def get_payload(request: RequestModel, engine, provider):
+    if engine == "gemini":
+        return await get_gemini_payload(request, engine, provider)
+    elif engine == "claude":
+        return await get_claude_payload(request, engine, provider)
+    elif engine == "gpt":
+        return await get_gpt_payload(request, engine, provider)
+    else:
+        raise ValueError("Unknown payload")

response.py

@@ -0,0 +1,135 @@
+from datetime import datetime
+import json
+import httpx
+
+async def generate_sse_response(timestamp, model, content):
+    sample_data = {
+        "id": "chatcmpl-9ijPeRHa0wtyA2G8wq5z8FC3wGMzc",
+        "object": "chat.completion.chunk",
+        "created": timestamp,
+        "model": model,
+        "system_fingerprint": "fp_d576307f90",
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"content": content},
+                "logprobs": None,
+                "finish_reason": None
+            }
+        ],
+        "usage": None
+    }
+    json_data = json.dumps(sample_data, ensure_ascii=False)
+
+    # 构建SSE响应
+    sse_response = f"data: {json_data}\n\n"
+
+    return sse_response
+
+async def fetch_gemini_response_stream(client, url, headers, payload, model):
+    try:
+        timestamp = datetime.timestamp(datetime.now())
+        async with client.stream('POST', url, headers=headers, json=payload) as response:
+            buffer = ""
+            async for chunk in response.aiter_text():
+                buffer += chunk
+                while "\n" in buffer:
+                    line, buffer = buffer.split("\n", 1)
+                    print(line)
+                    if line and '\"text\": \"' in line:
+                        try:
+                            json_data = json.loads( "{" + line + "}")
+                            content = json_data.get('text', '')
+                            content = "\n".join(content.split("\\n"))
+                            sse_string = await generate_sse_response(timestamp, model, content)
+                            yield sse_string
+                        except json.JSONDecodeError:
+                            print(f"无法解析JSON: {line}")
+
+            # 处理缓冲区中剩余的内容
+            if buffer:
+                # print(buffer)
+                if '\"text\": \"' in buffer:
+                    try:
+                        json_data = json.loads(buffer)
+                        content = json_data.get('text', '')
+                        content = "\n".join(content.split("\\n"))
+                        sse_string = await generate_sse_response(timestamp, model, content)
+                        yield sse_string
+                    except json.JSONDecodeError:
+                        print(f"无法解析JSON: {buffer}")
+
+            yield "data: [DONE]\n\n"
+    except httpx.ConnectError as e:
+        print(f"连接错误： {e}")
+
+async def fetch_gpt_response_stream(client, url, headers, payload):
+    try:
+        async with client.stream('POST', url, headers=headers, json=payload) as response:
+            async for chunk in response.aiter_bytes():
+                print(chunk.decode('utf-8'))
+                yield chunk
+    except httpx.ConnectError as e:
+        print(f"连接错误： {e}")
+
+async def fetch_claude_response_stream(client, url, headers, payload, engine, model):
+    try:
+        timestamp = datetime.timestamp(datetime.now())
+        async with client.stream('POST', url, headers=headers, json=payload) as response:
+            buffer = ""
+            async for chunk in response.aiter_text():
+                buffer += chunk
+                while "\n" in buffer:
+                    line, buffer = buffer.split("\n", 1)
+                    # print(line)
+                    if engine == "gemini":
+                        if line and '\"text\": \"' in line:
+                            try:
+                                json_data = json.loads( "{" + line + "}")
+                                content = json_data.get('text', '')
+                                content = "\n".join(content.split("\\n"))
+                                sse_string = await generate_sse_response(timestamp, model, content)
+                                yield sse_string
+                            except json.JSONDecodeError:
+                                print(f"无法解析JSON: {line}")
+                    else:
+                        yield line + "\n"
+
+            # 处理缓冲区中剩余的内容
+            if buffer:
+                # print(buffer)
+                if engine == "gemini":
+                    if '\"text\": \"' in buffer:
+                        try:
+                            json_data = json.loads(buffer)
+                            content = json_data.get('text', '')
+                            content = "\n".join(content.split("\\n"))
+                            sse_string = await generate_sse_response(timestamp, model, content)
+                            yield sse_string
+                        except json.JSONDecodeError:
+                            print(f"无法解析JSON: {buffer}")
+                else:
+                    yield buffer
+
+            if engine == "gemini":
+                yield "data: [DONE]\n\n"
+    except httpx.ConnectError as e:
+        print(f"连接错误： {e}")
+
+async def fetch_response(client, url, headers, payload):
+    response = await client.post(url, headers=headers, json=payload)
+    return response.json()
+
+async def fetch_response_stream(client, url, headers, payload, engine, model):
+    print(f"Engine: {engine}")
+    if engine == "gemini":
+        async for chunk in fetch_gemini_response_stream(client, url, headers, payload, model):
+            yield chunk
+    elif engine == "claude":
+        async for chunk in fetch_claude_response_stream(client, url, headers, payload, engine, model):
+            yield chunk
+    elif engine == "gpt":
+        async for chunk in fetch_gpt_response_stream(client, url, headers, payload):
+            yield chunk
+    else:
+        raise ValueError("Unknown response")