File size: 3,359 Bytes
cde68dd ed30b5b cde68dd ed30b5b cde68dd |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import os
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Dict
import google.generativeai as genai
from google.generativeai.types import HarmBlockThreshold, HarmCategory
from dotenv import load_dotenv
load_dotenv()
app = FastAPI()
class ChatMessage(BaseModel):
role: str
content: str
class ChatRequest(BaseModel):
messages: List[ChatMessage]
api_key: str = None
chat_history = []
async def stream_response(response):
for chunk in response:
yield chunk.text
@app.post("/api/think")
async def think(request: ChatRequest):
api_key = request.api_key or os.getenv("GOOGLE_API_KEY")
if not api_key:
raise HTTPException(status_code=400, detail="API key is required")
genai.configure(api_key=api_key)
model = genai.GenerativeModel(
model_name="gemini-1.5-flash",
generation_config={
"temperature": 1,
"max_output_tokens": 8192,
"response_mime_type": "application/json",
},
safety_settings={
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE
},
system_instruction="You are a helpful AI assistant. Analyze the user's request and provide your thoughts in XML tags. Use <thinking> for overall analysis, <step> for individual steps, and <count> for remaining steps."
)
messages = [{"role": m.role, "parts": [m.content]} for m in request.messages]
response = await model.generate_content_async(messages, stream=True)
return StreamingResponse(stream_response(response), media_type="text/plain")
@app.post("/api/chat")
async def chat(request: ChatRequest):
api_key = request.api_key or os.getenv("GOOGLE_API_KEY")
if not api_key:
raise HTTPException(status_code=400, detail="API key is required")
genai.configure(api_key=api_key)
model = genai.GenerativeModel(
model_name="gemini-1.5-flash",
generation_config={
"temperature": 1,
"max_output_tokens": 8192,
"response_mime_type": "application/json",
},
safety_settings={
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE
},
system_instruction="You are a helpful AI assistant. Provide a concise summary of the thinking process and give a final answer."
)
messages = [{"role": m.role, "parts": [m.content]} for m in request.messages]
response = await model.generate_content_async(messages, stream=True)
return StreamingResponse(stream_response(response), media_type="text/plain")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000) |