Files

188 lines
6.3 KiB
Python
Raw Permalink Normal View History

import os
from typing import Optional
from fastapi import FastAPI, HTTPException, Security, Depends
from fastapi.security import APIKeyHeader
from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv
import json
from pydantic import BaseModel
from fastapi import HTTPException
import os
from scripts.transcriber import transcribe_media, group_words_into_sentences
from scripts.generate_summary import general_summary, custom_summary
from src.models import PlanTier, PlanLimits
# Load environment variables
load_dotenv()
API_KEY = os.getenv("API_KEY_ACCESS")
# Initialize FastAPI app
app = FastAPI(
title="Microdot AI API",
description="API For fire fighter",
version="1.0.0"
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Setup API key authentication
api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
async def get_api_key(api_key_header: str = Security(api_key_header)) -> str:
"""Validate API key from header"""
if not api_key_header or not api_key_header.startswith('Bearer '):
raise HTTPException(
status_code=401,
detail={"error": "Unauthorized", "message": "API key is missing or invalid."}
)
token = api_key_header.split(' ')[1]
if token != API_KEY:
raise HTTPException(
status_code=401,
detail={"error": "Unauthorized", "message": "API key does not match."}
)
return token
class TranscribeRequest(BaseModel):
media_url: Optional[str] = None
media_type: Optional[str] # Corrected type hint for media_type
plan_tier: Optional[str] = "freemium" # Default to freemium plan if not specified
class ChatResp(BaseModel): # Added BaseModel inheritance
error: Optional[str] = None
class TranscriptResponse(BaseModel):
transcript: dict # Changed type hint for transcript to a dictionary
class GeneralSummaryRequest(BaseModel):
transcript: Optional[str] = None
plan_tier: Optional[str] = "pro" # Default to pro plan if not specified
class TemplateSummaryRequest(BaseModel):
transcript: Optional[str] = None
template: Optional[str] = None
@app.post("/microdot-ai/transcribe")
async def chat_endpoint(
request: TranscribeRequest,
api_key: str = Depends(get_api_key)
):
try:
# Get the plan tier from the request or default to freemium
plan_tier = request.plan_tier.lower() if request.plan_tier else "freemium"
# Validate plan tier using our PlanTier enum
valid_tiers = [t.value for t in PlanTier]
if plan_tier not in valid_tiers:
plan_tier = PlanTier.FREEMIUM.value # Default to freemium if invalid tier
# Check if the plan includes speaker identification
include_speakers = PlanLimits.get_limit(plan_tier, "speaker_identification")
# Use the transcribe_media function to transcribe the media
if request.media_url:
transcription_response = transcribe_media(request.media_url, media_type=request.media_type)
if transcription_response is None:
raise HTTPException(status_code=500, detail="Transcription failed.")
print(f"Transcription response: {transcription_response}") # Debugging print
# Parse response
words = transcription_response["results"]["channels"][0]["alternatives"][0]["words"]
transcript = group_words_into_sentences(words=words, include_speakers=include_speakers)
return TranscriptResponse(
transcript=transcript, # Corrected to return the transcript
error=None
)
except Exception as e:
print(f"Error processing chat request: {str(e)}") # Print statement added
raise HTTPException(
status_code=500,
detail=f"Error processing chat request: {str(e)}"
)
@app.post("/microdot-ai/general-summary")
async def general_summary_endpoint(
request: GeneralSummaryRequest,
api_key: str = Depends(get_api_key)
):
try:
if not request.transcript:
raise HTTPException(status_code=400, detail="Transcript is required.")
# Get the plan tier from the request or default to pro
plan_tier = request.plan_tier.lower() if request.plan_tier else "pro"
# Validate plan tier using our PlanTier enum
valid_tiers = [t.value for t in PlanTier]
if plan_tier not in valid_tiers:
plan_tier = PlanTier.PRO.value # Default to pro if invalid tier
# Get the appropriate summary type for this plan tier
summary_type = PlanLimits.get_limit(plan_tier, "summary_type")
# Generate the summary based on the plan tier
response = general_summary(json.loads(request.transcript), plan_tier=plan_tier)
return TranscriptResponse(
transcript=response
)
except Exception as e:
print(f"Error processing general summary request: {str(e)}")
raise HTTPException(
status_code=500,
detail=f"Error processing general summary request: {str(e)}"
)
@app.post("/microdot-ai/template-summary")
async def template_summary_endpoint( # Corrected function name to avoid conflict
request: TemplateSummaryRequest,
api_key: str = Depends(get_api_key)
):
try:
if not request.transcript:
raise HTTPException(status_code=400, detail="Transcript is required.")
if not request.template:
raise HTTPException(status_code=400, detail="Template is required.")
transcript = json.loads(request.transcript)
template = json.loads(request.template) # Removed the check for missing template as it's now required
response = custom_summary(template, transcript)
return TranscriptResponse(
transcript=response
)
except Exception as e:
print(f"Error processing template summary request: {str(e)}") # Updated print statement for clarity
raise HTTPException(
status_code=500,
detail=f"Error processing template summary request: {str(e)}"
)
@app.on_event("startup")
async def startup_event():
"""Initialize required components on startup"""
pass
if __name__ == "__main__":
import uvicorn
uvicorn.run("app:app", host="0.0.0.0", port=5056, reload=True)