|
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) |
