app.py

import os
import sys
import json
import base64
import wave
import signal
import argparse
import threading
import logging
import time
from queue import Queue, Empty
from datetime import datetime
import httpx
from fastapi import FastAPI, WebSocket, Request
from fastapi.responses import JSONResponse
import requests

from deepgram import (
    DeepgramClient,
    DeepgramClientOptions,
    PrerecordedOptions,
    FileSource,
)
from deepgram.utils import verboselogs

from fastapi import FastAPI, WebSocket
import uvicorn

app = FastAPI()

def transcribe_audio_file(audio_file_path: str):
    lang_file_updated_to_hindi = False
    try:
        if not os.path.exists(audio_file_path):
            print(f"Error: File '{audio_file_path}' does not exist for transcription.")
            return

        # Create Deepgram client with the API key
        config = DeepgramClientOptions(
            verbose=verboselogs.SPAM,  # Maximum logging for debugging
        )
        deepgram = DeepgramClient("30880849b38adbbf45da95e3c7cb0f18ca89e6ea", config)

        # Read the audio file
        print(f"Reading audio file for transcription: {audio_file_path}...")
        with open(audio_file_path, "rb") as file:
            buffer_data = file.read()

        # Prepare the payload
        payload: FileSource = {
            "buffer": buffer_data,
        }

        # Configure transcription options
        options = PrerecordedOptions(
            model="general",  # Using general model for Hindi
            smart_format=True,  # Enable smart formatting
            utterances=True,  # Split audio into utterances
            punctuate=True,  # Add punctuation
            diarize=True,  # Speaker diarization
            language="hi",  # Set language to Hindi
        )

        # Transcribe the file
        print("Transcribing...")
        before = datetime.now()
        response = deepgram.listen.rest.v("1").transcribe_file(
            payload, options, timeout=httpx.Timeout(300.0, connect=10.0)
        )
        after = datetime.now()

        # Check for Hindi characters in the entire response
        response_json_str = response.to_json()
        has_hindi_ha = "\\u0939" in response_json_str  # Check for ह
        has_hindi_na = "\\u0928" in response_json_str  # Check for न
        
        # Only update to Hindi if ह is present but न is not present
        if has_hindi_ha and not has_hindi_na:
            print("[*] Detected Hindi character '\u0939' without '\u0928' in transcription response.")
            with open("LANG.txt", "w") as f:
                f.write("HINDI")
            lang_file_updated_to_hindi = True
            print("[*] Updated LANG.txt with HINDI.")
        else:
            if has_hindi_na:
                print("[*] Detected Hindi character '\u0928' - not updating to Hindi.")
            else:
                print("[*] Did not detect required Hindi characters in transcription response.")

        # Print results
        print("\nTranscription Results:")
        print("=" * 50)
        print(response.to_json(indent=4))
        print("\nTime taken:", (after - before).seconds, "seconds")

    except Exception as e:
        print(f"Error during transcription: {e}")
    finally:
        # This block ensures LANG.txt is cleared if HINDI was not successfully written
        if not lang_file_updated_to_hindi:
            try:
                with open("LANG.txt", "w") as f:
                    f.write("")  # Clear the file
                print("[*] LANG.txt cleared (either not Hindi or error occurred).")
            except IOError:
                print("[*] Could not clear LANG.txt in finally block.")

class Stream:
    def __init__(self, rate, channels, sample_width):
        self.rate = rate
        self.channels = channels
        self.sample_width = sample_width
        self.buff = Queue()
        self.closed = False
        self.output_file = "output.wav"
        self.frames = []
        self.start_time = None
        self.recording_duration = 5  # Record for 5 seconds

    def fill_buffer(self, chunk):
        if self.start_time is None:
            self.start_time = time.time()
        
        # Only add chunks if we're within the recording duration
        if time.time() - self.start_time <= self.recording_duration:
            self.buff.put(chunk)
        else:
            if not self.closed:
                self.closed = True
                print(f"[*] Recording stopped after {self.recording_duration} seconds")

    def save_audio(self):
        print("[*] Saving audio to", self.output_file)
        with wave.open(self.output_file, 'wb') as wf:
            wf.setnchannels(self.channels)
            wf.setsampwidth(self.sample_width)
            wf.setframerate(self.rate)
            wf.writeframes(b''.join(self.frames))
        print("[*] Audio saved successfully.")
        
        # Add transcription after saving
        if os.path.exists(self.output_file):
            print(f"[*] Initiating transcription for {self.output_file}")
            transcribe_audio_file(self.output_file)
        else:
            print(f"[*] Transcription skipped: {self.output_file} not found after saving.")

    def consume_audio(self):
        while not self.closed:
            try:
                chunk = self.buff.get(timeout=1)
                if chunk is None:
                    break
                self.frames.append(chunk)
            except Empty:
                continue

@app.websocket("/media")
async def media_endpoint(websocket: WebSocket):
    await websocket.accept()
    print("[*] WebSocket connection accepted")

    has_seen_media = False
    message_count = 0

    # Setup stream object
    stream = Stream(rate=8000, channels=1, sample_width=2)
    consumer_thread = threading.Thread(target=stream.consume_audio)
    consumer_thread.start()

    while True:
        try:
            message = await websocket.receive_text()
        except:
            break

        if message is None:
            print("[*] No message received")
            continue

        data = json.loads(message)

        if data["event"] == "connected":
            print("[*] Connected Message:", message)
        elif data["event"] == "start":
            print("[*] Start Message:", message)
        elif data["event"] == "media":
            payload = data["media"]["payload"]
            chunk = base64.b64decode(payload)
            stream.fill_buffer(chunk)

            if not has_seen_media:
                print("[*] First media message received. Suppressing further logs.")
                has_seen_media = True

        elif data["event"] == "stop":
            print("[*] Stop Message:", message)
            break

        message_count += 1

        # Check if we've reached the 5-second limit
        if stream.closed:
            break

    print(f"[*] WebSocket connection closed. Total messages: {message_count}")
    stream.closed = True
    consumer_thread.join()
    stream.save_audio()

def signal_handler(sig, frame):
    print("Exiting gracefully...")
    sys.exit(0)

# --- Health check endpoints ---
@app.get("/health")
async def health_check_get():
    try:
        with open("LANG.txt", "r") as f:
            content = f.read().strip()
        if content == "HINDI":
            return JSONResponse(content={"message": "GET OK - Hindi language detected"}, status_code=200)
    except FileNotFoundError:
        pass
    except IOError as e:
        print(f"Error reading LANG.txt for GET health check: {e}")
    return JSONResponse(content={"message": "GET OK - Hindi language not detected"}, status_code=404)

@app.post("/health")
async def health_check_post(request: Request):
    try:
        with open("LANG.txt", "r") as f:
            content = f.read().strip()
        if content == "HINDI":
            return JSONResponse(content={"message": "POST OK - Hindi language detected"}, status_code=200)
    except FileNotFoundError:
        pass
    except IOError as e:
        print(f"Error reading LANG.txt for POST health check: {e}")
    return JSONResponse(content={"message": "POST OK - Hindi language not detected"}, status_code=404)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='FastAPI WebSocket Audio Saver')
    parser.add_argument('--port', type=int, default=8000, help='Port to run the server on')
    args = parser.parse_args()

    signal.signal(signal.SIGINT, signal_handler)

    print(f"[*] Server running at ws://localhost:{args.port}/media")
    uvicorn.run("app:app", host="0.0.0.0", port=args.port, log_level="info")