Update app.py

app.py

@@ -7,6 +7,8 @@ import subprocess
 import gc
 import requests
 import time
+import random
+import re
 from googletrans import Translator
 import asyncio
 from flask import Flask, request, jsonify, send_from_directory
@@ -33,10 +35,62 @@ HEYGEN_API_KEY = "NGM2N2VjNmM4NWM0NGQxMjkyNWFiMjg4OTdlMTI2MDItMTcyNDQ5ODM1MA=="
 HEYGEN_GENERATE_URL = "https://api.heygen.com/v2/video/generate"
 HEYGEN_STATUS_URL = "https://api.heygen.com/v1/video_status.get"
 
+# Initialize OpenAI client
+client = OpenAI(api_key="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")
+
 def clear_cuda_memory():
     torch.cuda.empty_cache()
     gc.collect()
 
+def openai_chat_avatar(text_prompt):
+    """Summarize text using OpenAI GPT-4o-mini"""
+    response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "Summarize the following paragraph into a complete and accurate single sentence with no more than 30 words. The summary should capture the gist of the paragraph and make sense and remove the citation and document name from the end."},
+            {"role": "user", "content": f"Please summarize the following paragraph into one sentence with 30 words or fewer, ensuring it makes sense and captures the gist and remove the citation from the end: {text_prompt}"},
+        ],
+        max_tokens = len(text_prompt),
+    )
+    return response
+
+def ryzedb_chat_avatar(question, app_id):
+    """Query RyzeDB API for response"""
+    url = "https://inference.dev.ryzeai.ai/v2/chat/stream"
+    print("ryze db question", question)
+    
+    payload = {
+        "input": {
+            "app_id": app_id,
+            "query": question,
+            "chat_history": []
+        },
+        "config": {
+            "thread_id": "123456"
+        }
+    }
+    
+    headers = {
+        'Content-Type': 'application/json'
+    }
+    
+    response = requests.post(url, json=payload, headers=headers, stream=True)
+    
+    try:
+        raw_text = response.text.strip()
+    
+        if raw_text.startswith("data:"):
+            raw_text = raw_text[len("data:"):].strip()
+    
+        json_data = json.loads(raw_text)
+    
+        response_content = json_data.get("content", "")
+        return response_content
+
+    except Exception as e:
+        print("Error parsing response:", e)
+        return ""
+
 def run_inference(video_path, audio_path, video_out_path,
                   inference_ckpt_path, unet_config_path="configs/unet/second_stage.yaml",
                   inference_steps=20, guidance_scale=1.0, seed=1247):
@@ -391,6 +445,7 @@ def download_heygen_video(video_url):
 def generate_video():
     global TEMP_DIR
     TEMP_DIR = create_temp_dir()
+    start_time = time.time()
 
     # Get form parameters
     text_prompt = request.form.get('text_prompt', '').strip()
@@ -399,18 +454,53 @@ def generate_video():
     
     print('Input text prompt:', text_prompt)
     
-    # Check if user wants to use HeyGen or local AI avatar
+    # Get processing parameters
     use_heygen = request.form.get('use_heygen', 'no').lower() == 'yes'
     voice_cloning = request.form.get('voice_cloning', 'no')
     target_language = request.form.get('target_language', 'original_text')
+    chat_model_used = request.form.get('chat_model_used', 'ryzedb')
+    app_id = request.form.get('app_id', '')
     
-    # Translate text if needed
-    if target_language != 'original_text':
-        translated_text = translate_text(text_prompt, target_language)
-        text_prompt = translated_text.strip()
-        print('Translated input text prompt:', text_prompt)
+    # Validate app_id if using RyzeDB
+    if chat_model_used == 'ryzedb' and not app_id:
+        return jsonify({'error': 'App ID cannot be blank when using RyzeDB'}), 400
 
     try:
+        # Process text prompt based on chat model selection
+        if chat_model_used == 'ryzedb':
+            start_time_ryze = time.time()
+            print("Processing text with RyzeDB...")
+            
+            # Get response from RyzeDB
+            ryze_response = ryzedb_chat_avatar(text_prompt, app_id)
+            print("Response from RyzeDB inference:", ryze_response)
+            
+            # Clean up response if needed
+            if "No information available" in ryze_response:
+                ryze_response = re.sub(r'\\+', '', ryze_response)
+            
+            # Summarize with OpenAI
+            openai_response = openai_chat_avatar(ryze_response)
+            text_prompt = openai_response.choices[0].message.content.strip()
+            
+            end_time_ryze = time.time()
+            ryze_processing_time = end_time_ryze - start_time_ryze
+            print(f'Final processed text prompt using RyzeDB + OpenAI: {text_prompt}')
+            print(f'Time to process with RyzeDB + OpenAI: {ryze_processing_time:.2f} seconds')
+            
+        elif chat_model_used == 'self':
+            print("Using original text prompt without processing...")
+            text_prompt = text_prompt.strip()
+        else:
+            print("Unknown chat model specified, using original text...")
+            text_prompt = text_prompt.strip()
+    
+        # Translate text if needed
+        if target_language != 'original_text':
+            translated_text = translate_text(text_prompt, target_language)
+            text_prompt = translated_text.strip()
+            print('Translated input text prompt:', text_prompt)
+
         if use_heygen:
             print("Using HeyGen API for video generation...")
             
@@ -476,10 +566,17 @@ def generate_video():
             
             processing_method = "HeyGen API" if use_heygen else "Local AI Avatar"
             
+            # Calculate total processing time
+            end_time = time.time()
+            total_time = end_time - start_time
+            
             return jsonify({
                 "message": f"Video processed successfully using {processing_method}.",
                 "output_video": video_url,
                 "processing_method": processing_method,
+                "text_prompt": text_prompt,
+                "chat_model_used": chat_model_used,
+                "time_taken": round(total_time, 2),
                 "status": "success"
             }), 200
         else: