File size: 4,834 Bytes
e454383 | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | import os
from dotenv import load_dotenv
from typing import List, Optional, Dict
from pydantic import BaseModel
from fastapi import Body
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from utils.ai import AiService
from utils.vector_db import VectorDB
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4_1_mini_complete, openai_embed
from utils.time_funcs import extract_timestamps
from utils.document_management import load_document, transcript_chunker, merge_chunks
from prompts import name_generation_prompt
# Load environment variables
load_dotenv()
# Initialize FastAPI app
app = FastAPI(title="Transcription API")
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"], # Allows all methods
allow_headers=["*"], # Allows all headers
)
# Initialize services
working_dir = os.path.join(os.getcwd(), "_backend", "data","honduras_data")
if os.getenv("ENVIRONMENT") == "production":
working_dir = os.path.join(os.getcwd(), "data","honduras_data")
transcription_data_dir = os.path.join(working_dir)
vector_db_name = "transcription"
vector_db_embedding_length = 1536
vector_db = VectorDB(
vector_db_name,
vector_db_embedding_length,
os.getenv("OPENAI_API_KEY"),
working_dir
)
ai_service = AiService(model="gpt-4.1")
lightrag_dir = os.path.join(working_dir, "lightrag")
rag = LightRAG(
working_dir=lightrag_dir,
llm_model_func=gpt_4_1_mini_complete,
embedding_func=openai_embed
)
class ChatRequest(BaseModel):
message: str
chat_history: List[Dict[str, str]]
class ChatResponse(BaseModel):
response: str
timestamps: List[Dict[str, str | int]] # Updated type annotation
class NameRequest(BaseModel):
user_query: str
assistant_response: str
class SearchRequest(BaseModel):
query: str
filters: Optional[Dict] = {}
class SessionInfoResponse(BaseModel):
session_info: Dict
video_url: str
# API Endpoints
@app.post("/api/v1/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
try:
translated_query = ai_service.generate_single_response(
request.message,
"Translate the following text to Spanish",
json_output=False
)
# Search documents
search_docs = vector_db.search(translated_query, limit=100,
# filter=filters,
score_threshold=0
)
# Get GraphML response using asyncio.create_task to avoid event loop issues
import asyncio
loop = asyncio.get_event_loop()
graphML_resp = await loop.run_in_executor(
None,
lambda: rag.query(
request.message,
param=QueryParam(
mode="hybrid",
)
)
)
# Generate response
assistant_response = ai_service.generate_response(
request.chat_history,
docs=search_docs,
extra_context=graphML_resp,
)
# Extract timestamps
timestamps = extract_timestamps(assistant_response, search_docs)
return ChatResponse(
response=assistant_response,
timestamps=timestamps # Remove json.dumps(), return list directly
)
except Exception as e:
print(e)
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/v1/generate_name", response_model=str)
async def generate_name_endpoint(request: NameRequest):
try:
messages = ai_service.init_messages(name_generation_prompt)
ai_service.add_user_message(messages, request.user_query)
# Generate response
response = ai_service.generate_response(messages)
return response
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/v1/transcript", response_model=List[Dict])
async def get_transcript_file(body: Dict = Body(...)):
try:
file_name = body.get("file_name")
if not file_name:
raise HTTPException(status_code=400, detail="file_name is required")
file_path = os.path.join(transcription_data_dir, file_name)
doc = load_document(file_path)
transcript = transcript_chunker(doc, filename=file_name)
transcript = merge_chunks(transcript)
return transcript
except FileNotFoundError:
raise HTTPException(status_code=404, detail="File not found")
@app.get("/api/v1/health")
async def health_check():
return {"status": "healthy"}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|