forked from railwayapp-templates/fastapi
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
53 lines (46 loc) · 1.35 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import os
import requests
import json
from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
from typing import Dict
app = FastAPI()
class PromptRequest(BaseModel):
token: str
text: str
@app.post("/generate-prompt")
async def generate_prompt(request: PromptRequest):
# 验证token
validation_token = os.getenv('VALIDATION_TOKEN')
if request.token != validation_token:
raise HTTPException(status_code=401, detail="Invalid token")
# 从环境变量中获取API密钥
api_key = os.getenv('API_KEY')
if not api_key:
raise HTTPException(status_code=500, detail="API key not found")
# 设置请求头和数据
headers = {
'Content-Type': 'application/json'
}
payload = {
'contents': [
{
'parts': [
{
'text': request.text
}
]
}
]
}
# 发送POST请求
response = requests.post(
f'https://generativelanguage.googleapis.com/v1/models/gemini-1.5-flash:generateContent?key={api_key}',
headers=headers,
data=json.dumps(payload)
)
# 返回响应内容
if response.status_code == 200:
return response.json()
else:
raise HTTPException(status_code=response.status_code, detail=response.text)