cai / main.py
tachibanaa710's picture
Update main.py
dd88fc6 verified
import asyncio
from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel
from typing import Optional
from PyCharacterAI import get_client
from PyCharacterAI.exceptions import SessionClosedError, RequestError
import uvicorn
import os
app = FastAPI()
DEFAULT_TOKEN = os.getenv("token")
DEFAULT_CHARACTER_ID = "smtV3Vyez6ODkwS8BErmBAdgGNj-1XWU73wIFVOY1hQ"
DEFAULT_VOICE_ID = "26edcb78-7e4d-4d81-970e-6fdaf81ed778" # Replace with the desired voice ID
# This will hold the client instance, which should be initialized once
client = None
@app.on_event("startup")
async def startup_event():
global client
# Initialize client during startup
client = await get_client(token=DEFAULT_TOKEN)
@app.on_event("shutdown")
async def shutdown_event():
# Close the client session during shutdown
await client.close_session()
async def initialize_client(token: str):
"""Helper function to reinitialize the client connection."""
global client
if client is not None:
await client.close_session() # Ensure any existing session is closed
client = await get_client(token=token)
@app.get("/send_message")
async def send_message(
message: str,
token: Optional[str] = Query(None, description="API token for authentication."),
character_id: str = Query(DEFAULT_CHARACTER_ID, description="Character ID for the chat session."),
chat_id: Optional[str] = Query(None, description="ID of the existing chat session, if available."),
voice_id: str = Query(DEFAULT_VOICE_ID, description="Voice ID for generating speech."),
voice: bool = Query(True, description="Set to true to generate voice, false to skip voice generation.")
):
"""
Send a message to the character. If no chat_id is provided, initialize a new chat session.
Optionally generate voice for the response.
"""
try:
token = token or DEFAULT_TOKEN
# Ensure client connection is active
if client is None:
await initialize_client(token)
# If chat_id is not provided, create a new chat session
if not chat_id:
chat, greeting_message = await client.chat.create_chat(character_id)
chat_id = chat.chat_id
greeting_text = {
"author": greeting_message.author_name,
"text": greeting_message.get_primary_candidate().text
}
else:
greeting_text = None # No greeting if we’re using an existing chat
# Send the message and get the response
answer = await client.chat.send_message(character_id, chat_id, message)
response_text = answer.get_primary_candidate().text
speech_url = None # Default to None
if voice: # Only generate voice if the voice parameter is true
try:
speech_url = await client.utils.generate_speech(
chat_id,
answer.turn_id,
answer.get_primary_candidate().candidate_id,
voice_id,
return_url=True
)
except Exception as e:
# Log the voice generation error and continue without the speech URL
print(f"Voice generation failed: {e}")
# Prepare the response data
response_data = {
"chat_id": chat_id,
"author": answer.author_name,
"response": response_text,
"voice_url": speech_url if voice else None
}
# Include greeting message only if a new chat was created
if greeting_text:
response_data["greeting_message"] = greeting_text
return response_data
except (SessionClosedError, RequestError):
# Attempt to reinitialize the client on connection closure
print("Connection closed. Reinitializing client.")
await initialize_client(token)
raise HTTPException(status_code=500, detail="Connection was closed. Please try again.")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# Root endpoint
@app.get("/")
def root():
return {"message": "Welcome to the Character AI API with FastAPI!"}
# Run with: uvicorn main:app --host 0.0.0.0 --port 8000
if __name__ == "__main__":
uvicorn.run("main:app", host="0.0.0.0", port=7860, workers=8, timeout_keep_alive=60000)