app.py

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 pathlib import Path
from openai import OpenAI
import soundfile as sf
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 datetime import datetime

os.environ["OPENAI_API_KEY"] = "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"
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")

class UserIntent(BaseModel):
 intent: str = Field(description= "Intent of the user")

question_parser = JsonOutputParser(pydantic_object=GenerateQuestion)
starting_question_parser = JsonOutputParser(pydantic_object=StartingQuestion)
intent_parser = JsonOutputParser(pydantic_object=UserIntent)

@chain
def input_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()},
             ])]
             )
 return msg.content

@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.
"""

user_intent_prompt = """
  The system and user are engaged in a conversation about a photo uploaded by the user.
  The system asks questions related to the photograph, and the user responds.
  Your task is to analyze the user's input to accurately determine their intent.

  The possible intents are:
  1. "change photo" - The user explicitly or implicitly indicates they want to move on to the next photo, do not wish to discuss the current photo, or directly state a desire to stop talking about the current photograph.
  2. "change topic" - The user expresses a desire to talk about something else within the context of the current photograph or shows disinterest in the current line of questioning but doesn't want to change the photograph itself.
  3. "continue" - The user is comfortable with the current conversation and wants to continue discussing the current photo.

  Here is the user's input:
  {input}

  Provide:
  1. The intent of the user.
"""


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 | parser
   return vision_chain.invoke({'image_path': f'{image_path}', 'prompt': prompt, 'parser':parser})



def get_intent(user_input) -> dict:
   parser = intent_parser
   prompt = user_intent_prompt.format(input=user_input)
   intent_chain = input_model | parser
   return intent_chain.invoke({'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 image_handler(user_name, image_path):

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

    if not image_path:
      message = "Great! Please enter your name and upload a photo to share your story."
      return   None, message, transform_text_to_speech(message, user_name)


    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):
      shutil.copy(image_path, new_image_path)
      iter = increment_counts(count_file_path)
      handle_duration(duration_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 res, transform_text_to_speech(res, user_name)

    else:
      res = retrieve_memory(input_filepath, user_name)
      if res["speaker"] == AI_CHARACTER:
        qa = "Continuing from where we left off— "+ res["reply"]
        return  qa, transform_text_to_speech(qa , 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 res, transform_text_to_speech(res, user_name)




def audio_handler(user_name, image_path, audio):

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

    if not image_path:
      message = "Great! Please enter your name and upload a photo to share your story."
      return   None, message, transform_text_to_speech(message, user_name)


    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 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)

      output =  get_intent(user_input)

      if(output["intent"].strip() == "change photo"):
        res = "Press the 'Tell a New Story' button to move to a new photograph."
      else:
        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  None, None, None



    # 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:
          return  None, res["reply"], transform_text_to_speech(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 None, res, transform_text_to_speech(res, user_name)




# Backend function to clear inputs
def clear_inputs(user_name, image_path):
  message = "Great! Please enter your name and upload a photo to share your story."
  return None, None, message, 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", value=None)


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

    image_input.change(image_handler, inputs=[username, image_input], outputs=[stud_output, audio_output])
    audio_input.change(audio_handler, inputs=[username, image_input, audio_input], outputs=[audio_input, stud_output, audio_output])



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

# Launch the interface
demo.launch(share=True, debug=True)