app.py

# App.py File.
import base64
import logging
import numpy as np
import os
import langchain
import base64
import gradio as gr
import shutil
import json
import re
from datetime import datetime
from pathlib import Path
from openai import OpenAI
import soundfile as sf
from datetime import datetime
from pydub import AudioSegment
# from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.chains import TransformChain
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langchain import globals
from langchain_core.runnables import chain
from langchain_core.output_parsers import JsonOutputParser
from langchain.memory import ConversationSummaryBufferMemory, ConversationBufferMemory
from pydantic import BaseModel, Field


os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
client = OpenAI()

def transform_text_to_speech(text: str, user):
  # Generate speech from transcription
  speech_file_path_mp3 = Path.cwd() / f"{user}-speech.mp3"
  speech_file_path_wav = Path.cwd() / f"{user}-speech.wav"
  response = client.audio.speech.create(
                model="tts-1",
                voice="onyx",
                input=text
            )

  with open(speech_file_path_mp3, "wb") as f:
      f.write(response.content)

  # Convert mp3 to wav
  audio = AudioSegment.from_mp3(speech_file_path_mp3)
  audio.export(speech_file_path_wav, format="wav")

  # Read the audio file and encode it to base64
  with open(speech_file_path_wav, "rb") as audio_file:
      audio_data = audio_file.read()
      audio_base64 = base64.b64encode(audio_data).decode('utf-8')

  # Create an HTML audio player with autoplay
  audio_html = f"""
  <audio controls autoplay>
      <source src="data:audio/wav;base64,{audio_base64}" type="audio/wav">
      Your browser does not support the audio element.
  </audio>
  """
  return audio_html



def transform_speech_to_text(audio, user):
  file_path = f"{user}-saved_audio.wav"
  sample_rate, audio_data = audio
  sf.write(file_path, audio_data, sample_rate)
  # Transcribe audio
  with open(file_path, "rb") as audio_file:
      transcription = client.audio.transcriptions.create(
          model="whisper-1",
          file=audio_file
      )
  return transcription.text



def load_image(inputs: dict) -> dict:
    """Load image from file and encode it as base64."""
    image_path = inputs["image_path"]

    def encode_image(image_path):
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    image_base64 = encode_image(image_path)
    return {"image": image_base64}

from langchain.chains import TransformChain
load_image_chain = TransformChain(
    input_variables=["image_path"],
    output_variables=["image"],
    transform=load_image
)

class GenerateQuestion(BaseModel):
 """Information about an image."""
 question: str = Field(description= "Respond to the user input and ask a follow back question, using conversation and photo provded as a guide.")

class StartingQuestion(BaseModel):
 """Information about an image."""
 question: str = Field(description= "A question to start converstion around the photograph")


# Pass the schemas to JsonOutputParser
question_parser = JsonOutputParser(pydantic_object=GenerateQuestion)
starting_question_parser = JsonOutputParser(pydantic_object=StartingQuestion)


@chain
def image_model(inputs: dict) -> str | list[str] | dict:
 """Invoke model with image and prompt."""
 model = ChatOpenAI(temperature=0.5, model="gpt-4o", max_tokens=1024)
 msg = model.invoke(
             [HumanMessage(
             content=[
             {"type": "text", "text": inputs["prompt"]},
             {"type": "text", "text": inputs["parser"].get_format_instructions()},
             {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{inputs['image']}"}},
             ])]
             )
 return msg.content


AI_CHARACTER = "Good Friend"

CONVERSATION_STARTER_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  Provide:
  - A question to start the conversation around the photograph.
  Note:
  1. You first want to know about the photo. Focus on the people in the photo and the relationships.  For example, ask who is in it, where it was taken, or if it was a special occasion.
  2. Avoid questions about emotions or feelings. Keep questions simple and open ended.
  3. Don't ask about things in the photo.
  4. Ask them if there are topics they would like to talk about.

"""

CONVERSATION_STARTER2_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  Here is the conversation history about the image between the user and you ({character}):
  {history}
  Provide:
  - A question about the contents of the photograph.

  Note:
  1. You first want to know about the contents of the photo. For example, ask who is in it, where it was taken, or if it was a special occasion.
  2. Do not repeat questions.  Don't ask a very personal question.
  3.  Focus on people and relationships. Keep questions simple and open ended.
  4. Use conversation history.
  5. Ask about people mentioned in user's responses not just in photo.
  6. Don't ask about things in the photo unless the person brings them up
  7. Avoid platitudes

"""


CONVERSATION_EXPANDING_PROMPT = """
  You are playing the role of a good friend who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  You are currently in the middle of a conversation with the user.
  Here is the conversation history about the image between the user and you (the good friend), reflecting the ongoing dialogue:
  {history}
  Provide:
  - A reply to the user's most recent input and a follow-up question that encourages them to expand on their answer about the people in the photograph or mentioned in their response
  Notes:
  1- Ask them about any stories they are reminded of and please ask only one question.
  2- Do not repeat questions or ask about information that has already been covered.
  3- Encourage full responses by asking open-ended questions that invite further elaboration.
  5- Use the conversation history to inform your question, while maintaining the flow of the ongoing conversation.
  6- Don't ask about things in the photo unless the person brings them up
  7- Avoid platitudes.

"""


CONVERSATION_ENDING_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  Here is the conversation history about the image between the user and you ({character}): reflecting the ongoing dialogue:
  {history}
  Provide:
  - A reply to the user's most recent input and a follow-up question that encourages them to share more about the story depicted in the photograph,
    discuss anything that the photograph reminds them of, or move on to another photograph or stop reminiscing.
  Notes:
  1- Ask them if they want to keep talking about the photo or move onto another photo.
  2- please ask only one question.
  3- Ask them to summarize how they feel about the photo
  4- Do not repeat questions or ask about information already covered in the conversation.
  5- Encourage full responses by asking open-ended questions that invite further elaboration.
"""

def get_prompt(image_path: str, iter: int, memory: str, firstname: str) -> dict:

   if iter == 1:
    parser = starting_question_parser
    prompt = CONVERSATION_STARTER_PROMPT.format(character=AI_CHARACTER)
   elif  iter >= 2 and iter <= 3:
    parser = starting_question_parser
    prompt =   CONVERSATION_STARTER2_PROMPT.format(history=memory, character=AI_CHARACTER)
   elif  iter > 3 and iter <= 9:
    parser = question_parser
    prompt=   CONVERSATION_EXPANDING_PROMPT.format(history=memory, character=AI_CHARACTER)
   else:
    parser = question_parser
    prompt=  CONVERSATION_ENDING_PROMPT.format(history=memory, character=AI_CHARACTER)

   vision_chain = load_image_chain | image_model | question_parser
   return vision_chain.invoke({'image_path': f'{image_path}', 'prompt': prompt, 'parser':parser})



def retrieve_memory(input_filepath, name):
  with open(input_filepath, 'r') as f:
    conversation = f.read()
  lines = conversation.strip().split('\n')
  last_reply = None

  # Loop through the lines from the end
  for line in reversed(lines):
      if re.match(r'(' + re.escape(AI_CHARACTER) + '|' + re.escape(name) + '):', line):
        last_reply = line
        break

  # Determine who made the last reply, split it based on the colon, and return JSON
  if last_reply:
      speaker, message = last_reply.split(":", 1)
      result = {
          "speaker": speaker.strip(),
          "reply": message.strip()
      }
      return result
  else:
      result = {
          "speaker": "",
          "reply": ""
      }
      return result

def load_counts(count_file_path):
    if os.path.exists(count_file_path):
        with open(count_file_path, 'r') as f:
            return json.load(f)
    return {"count": 0}

def save_counts(count_file_path, counts):
    with open(count_file_path, 'w') as f:
        json.dump(counts, f)

def increment_counts(count_file_path):
    counts = load_counts(count_file_path)
    counts["count"] += 1
    save_counts(count_file_path, counts)
    return counts["count"]

def handle_duration(duration_file_path: str):
  data = {
    "starttime": datetime.now().isoformat(),
    "endtime": None,  # Assuming endtime is None initially
    "duration": 0
  }
  with open(duration_file_path, 'w') as json_file:
    json.dump(data, json_file, indent=4)  

def calculate_duration(duration_file_path: str):
  with open(duration_file_path, 'r') as json_file:
    data = json.load(json_file)
  data['endtime'] = datetime.now().isoformat()   
  starttime = datetime.fromisoformat(data['starttime'])
  endtime = datetime.fromisoformat(data['endtime'])
  data['duration'] = (endtime - starttime).total_seconds() / 60  # Convert seconds to minutes
  with open(duration_file_path, 'w') as json_file:
    json.dump(data, json_file, indent=4)    



def pred(user_name, image_path, audio):

    if user_name.strip() == "":
      user_name = "Unknown"

    if image_path:
        current_date = datetime.now().strftime("%Y-%m-%d")
        base_dir = f"/data/{current_date}"
        os.makedirs(base_dir, exist_ok=True)
        user_name = user_name.strip()
        image_name = image_path.split("/")[-1]
        new_image_name = f"{user_name}-{image_name}"
        new_image_path = os.path.join(base_dir, new_image_name)
        input_filename = f"{new_image_name}-conversation-memory.txt"
        input_filepath = os.path.join(base_dir, input_filename)
        count_file_path = os.path.join(base_dir, f"{new_image_name}-tracking.json")
        duration_file_path = os.path.join(base_dir, f"{new_image_name}-duration.json")

        if not os.path.exists(new_image_path):
          handle_duration(duration_file_path)  
          shutil.copy(image_path, new_image_path)
          iter = increment_counts(count_file_path)
          output = get_prompt(new_image_path, iter, None, user_name)
          res = output["question"]
          with open(input_filepath, 'w') as f:
            f.write(AI_CHARACTER + ": " + res)
          return "New Photo Uploaded" , res, transform_text_to_speech(res, user_name)

        else:

          if audio is not None:
            user_input = transform_speech_to_text(audio, user_name)

            iter = increment_counts(count_file_path)
            with open(input_filepath, 'a') as f:
                f.write("\n" + user_name + ": " + user_input)
            with open(input_filepath, 'r') as f:
                content = f.read()
            output = get_prompt(new_image_path, iter, content, user_name)
            res = output["question"]
            with open(input_filepath, 'a') as f:
                f.write("\n" + AI_CHARACTER + ": "+ res)
            calculate_duration(duration_file_path)    
            return  user_input, res, transform_text_to_speech(res, user_name)



          # decide the path from the contents of the conversation memory.
          if  os.path.exists(input_filepath):
              res = retrieve_memory(input_filepath, user_name)
              if res["speaker"] == AI_CHARACTER:
                  if audio is None:
                    message = ""
                    return  "Loading response, please wait...", "Loading response, please wait...", None
                  prefix = "Continuing from where we left off: "
                  return   "" ,  prefix+ res["reply"], transform_text_to_speech(prefix+res["reply"], user_name)
              else:
                  with open(input_filepath, 'a') as f:
                      f.write("\n" + user_name + ": " + "I'd like to continue our conversation about this photograph.")
                  with open(input_filepath, 'r') as f:
                      content = f.read()
                  iter = increment_counts(count_file_path)
                  output = get_prompt(new_image_path, iter, content, user_name)
                  res = output["question"]
                  with open(input_filepath, 'a') as f:
                      f.write("\n" + AI_CHARACTER + ": " + res)
                  return  "I'd like to continue our conversation about this photograph.", res, transform_text_to_speech(res, user_name)

    message = "Great! Please upload a photo to tell your story."
    return   "", message, None

# Backend function to clear inputs
def clear_inputs(user_name, image_path):
  message = "Great! Please upload a photo to tell your story."
  if user_name.strip() == "" or image_path == None:
      return None, None, "", message, transform_text_to_speech(message, user_name)

  # image_name = image_path.split("/")[-1]
  # input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
  # input_filepath = f"/data/{input_filename}"
  # if  os.path.exists(input_filepath):
  #   with open(input_filepath, 'a') as f:
  #     f.write("\n" + f"{user_name}: " + "new photo uploaded")

  message = "Great!"
  return None, None, "", "",  None

# Gradio Interface
with gr.Blocks(title = "KitchenTable.AI") as demo:
    with gr.Row():
        with gr.Column():
            clear_button = gr.Button("Tell a new Story", elem_id="clear-button")
            username = gr.Textbox(label="Enter your first name")
            image_input = gr.Image(type="filepath", label="Upload an Image")
            audio_input = gr.Audio(sources="microphone", type="numpy", label="")
            

        with gr.Column():
            # Output fields
            user_input_output = gr.Textbox(label="User Input")
            stud_output = gr.Textbox(label="Good Friend")
            audio_output = gr.HTML(label="Audio Player")

    audio_input.change(pred, inputs=[username, image_input, audio_input], outputs=[ user_input_output, stud_output, audio_output])
    image_input.change(pred, inputs=[username, image_input, audio_input], outputs=[ user_input_output, stud_output, audio_output])

   
    clear_button.click(fn=clear_inputs, inputs=[username, image_input], outputs=[image_input, audio_input, user_input_output, stud_output, audio_output])

# Launch the interface
demo.launch()