src/server.py

from src.agents.chat_agent import BaseChatAgent
from src.utils.utils import save_as_json

from flask import Flask, request, jsonify
from flask_cors import CORS
import os
import sys
import time
import traceback
import boto3
import argparse
import pytz
import json
from datetime import datetime

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

class Server:
    def __init__(self) -> None:
        self.patient_info = ""

        self.conversation = []

        self.patient_out = None
        self.doctor_out = None
        self.patient = None
        self.doctor = None

        self.conversation_round = 0
        self.interview_protocol_index = None

    def set_timestamp(self):
        self.timestamp = datetime.now(pytz.timezone('US/Eastern')).strftime("%m/%d/%Y-%H:%M:%S")

    def set_patient(self, patient):
        self.patient = patient
        self.patient_info = {
            "patient_model_config": patient.agent_config,
        }

    def set_doctor(self, doctor):
        self.doctor = doctor

    def set_interview_protocol_index(self, interview_protocol_index):
        self.interview_protocol_index = interview_protocol_index

    def generate_doctor_response(self):
        '''
        Must be called after setting the patient and doctor
        '''
        self.doctor_out = self.doctor.talk_to_user(
            self.patient_out, conversations=self.conversation)[0]
        return self.doctor_out

    def submit_doctor_response(self, response):
        self.conversation.append(("doctor", response))
        self.doctor.context.add_assistant_prompt(response)

    def submit_patient_response(self, response):
        self.conversation.append(("patient", response))
        self.patient.context.add_assistant_prompt(response)

    def get_response(self, patient_prompt):
        self.patient_out = patient_prompt
        self.submit_patient_response(patient_prompt)
        print(f'Round {self.conversation_round} Patient: {patient_prompt}')
    
        if patient_prompt is not None:
            self.conversation_round += 1

        doctor_out = self.generate_doctor_response()
        self.submit_doctor_response(doctor_out)
        print(f'Round {self.conversation_round} Doctor: {doctor_out}')

        return {"response": doctor_out}

    def to_dict(self):
        return {
            'time_stamp': self.timestamp,
            'patient': {
                'patient_user_id': self.patient.patient_id,
                'patient_info': self.patient_info,
                'patient_context': self.patient.context.msg
            },
            'doctor': {
                'doctor_model_config': self.doctor.agent_config,
                'doctor_context': self.doctor.context.msg
            },
            "conversation": self.conversation,
            "interview_protocol_index": self.interview_protocol_index
        }

    def __json__(self):
        return self.to_dict()

    def reset(self):
        self.conversation = []
        self.conversation_round = 0
        if hasattr(self.doctor, 'reset') and callable(getattr(self.doctor, 'reset')):
            self.doctor.reset()
        if hasattr(self.patient, 'reset') and callable(getattr(self.patient, 'reset')):
            self.patient.reset()


def create_app():
    app = Flask(__name__)
    CORS(app)
    app.user_servers = {}
    return app


def configure_routes(app, args):

    @app.route('/', methods=['GET'])
    def home():
        '''
        This api will return the default prompts used in the backend, including system prompt, autobiography generation prompt, therapy prompt, and conversation instruction prompt
        Return:
            {
            system_prompt: String,
            autobio_generation_prompt: String,
            therapy_prompt: String,
            conv_instruction_prompt: String
            }
        '''


        return jsonify({
        }), 200


    @app.route('/api/initialization', methods=['POST'])
    def initialization():
        '''
        This API processes user configurations to initialize conversation states. It specifically accepts the following parameters:
        api_key, username, chapter_name, topic_name, and prompts. The API will then:
        1. Initialize a Server() instance for managing conversations and sessions.
        2. Configure the user-defined prompts.
        3. Set up the chapter and topic for the conversation.
        4. Configure the save path for both local storage and Amazon S3.
        '''
        data = request.get_json()
        username = data.get('username')
        api_key = data.get('api_key')
        if api_key and isinstance(api_key, str) and api_key.strip():
            os.environ["OPENAI_API_KEY"] = api_key.strip()

        # initialize
        # server.patient.patient_id = username
        counselor = BaseChatAgent(config_path=args.counselor_config_path)
        print(counselor)
        server = Server()
        # server.set_doctor = counselor
        server.doctor = counselor
        app.user_servers[username] = server

        return jsonify({"message": "API key set successfully"}), 200

    @app.route('/save/download_conversations', methods=['POST'])
    def download_conversations():
        """
        This API retrieves the user's conversation history based on their username and returns the conversation data to the frontend.
        Return:
            conversations: List[String]
        """
        data = request.get_json()
        username = data.get('username')
        chatbot_type = data.get('chatbot_type')
        if not username:
            return jsonify({'error': 'Username not provided'}), 400
        if not chatbot_type:
            return jsonify({'error': 'Chatbot type not provided'}), 400

        conversation_dir = os.path.join('user_data', chatbot_type, username, 'conversation')
        if not os.path.exists(conversation_dir):
            return jsonify({'error': 'User not found or no conversations available'}), 404

        # Llist all files in the conversation directory
        files = os.listdir(conversation_dir)
        conversations = []

        # read each conversation file and append the conversation data to the list
        for file_name in files:
            file_path = os.path.join(conversation_dir, file_name)
            try:
                with open(file_path, 'r') as f:
                    conversation_data = json.load(f)
                    # extract the 'conversation' from the JSON
                    conversation_content = conversation_data.get('conversation', [])
                    conversations.append({
                        'file_name': file_name,
                        'conversation': conversation_content
                    })
            except Exception as e:
                print(f"Error reading {file_name}: {e}")
                continue

        return jsonify(conversations), 200

    @app.route('/save/end_and_save', methods=['POST'])
    def save_conversation_memory():
        """
        This API saves the current conversation history and memory events to the backend, then synchronizes the data with the Amazon S3 server.
        """
        data = request.get_json()
        username = data.get('username')
        chatbot_type = data.get('chatbot_type')

        if not username:
            return jsonify({"error": "Username not provided"}), 400
        if not chatbot_type:
            return jsonify({"error": "Chatbot type not provided"}), 400
        server = app.user_servers.get(username)
        if not server:
            return jsonify({"error": "User session not found"}), 400

        # save conversation history
        server.set_timestamp()
        save_name = f'{server.chapter_name}-{server.topic_name}-{server.timestamp}.json'
        save_name = save_name.replace(' ', '-').replace('/', '-')
        print(save_name)

        # save to local file
        local_conv_file_path = os.path.join(server.patient.conv_history_path, save_name)
        save_as_json(local_conv_file_path, server.to_dict())

        local_memory_graph_file = os.path.join(server.patient.memory_graph_path, save_name)
        # if the chatbot type is 'baseline', create a dummy memory graph file
        if chatbot_type == 'baseline':
            save_as_json(local_memory_graph_file, {'time_indexed_memory_chain': []})
        else:
            # save memory graph
            server.doctor.memory_graph.save(local_memory_graph_file)
        
        # Auto-upload to Google Drive if authenticated
        try:
            import sys
            sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
            from google_drive_sync import auto_upload_to_drive
            
            # Upload conversation file
            auto_upload_to_drive(local_conv_file_path, user_id=username, folder_name="Chatbot_Conversations")
            # Upload memory graph file
            auto_upload_to_drive(local_memory_graph_file, user_id=username, folder_name="Chatbot_Conversations")
        except Exception as e:
            # Fail silently if Google Drive upload fails
            print(f"Google Drive auto-upload failed (non-critical): {str(e)}")
        
        return jsonify({"message": "Current conversation and memory graph are saved!"}), 200

    @app.route('/responses/doctor', methods=['POST'])
    def get_response():
        """
        This API retrieves the chatbot's response and returns both the response and updated memory events to the frontend.
        Return:
            {
            doctor_response: String,
            memory_events: List[dict]
            }
        """
        data = request.get_json()
        username = data.get('username')

        # patient_prompt = data.get('patient_prompt')
        # chatbot_type = data.get('chatbot_type')
        # if not username or not patient_prompt:
        #     return jsonify({"error": "Username or patient prompt not provided"}), 400
        # if not chatbot_type:
        #     return jsonify({"error": "Chatbot type not provided"}), 400
        # if not
        # server = app.user_servers.get(username)
        # if not server:
        #     return jsonify({"error": "User session not found"}), 400

        # print(server.patient.patient_id, server.chapter_name, server.topic_name)
        # doctor_response = server.get_response(patient_prompt=patient_prompt)
        
        # if chatbot_type == 'baseline':
        #     memory_events = []
        # else:
        #     memory_events = server.doctor.memory_graph.to_list()
        print('username', username)
        server = app.user_servers.get(username)
        llm_chatbot = server.doctor
        response = llm_chatbot.talk_to_user(data)

        return jsonify({'doctor_response': response})


def main():
    parser = argparse.ArgumentParser()
    # parser.add_argument('--patient-config-path', type=str,
    #                     default='./src/configs/patient_config.yaml')
    parser.add_argument('--counselor-config-path', type=str,
                        default='./src/configs/counselor_config.yaml')
    # parser.add_argument('--retriever-config-path', type=str,
    #                     default='./src/configs/retrievers/faiss_retriever.yaml')
    parser.add_argument('--store-dir',
                        type=str, default='./user_data')
    # parser.add_argument('--memory-graph-config', default='./src/configs/memory_graph_config.yaml')
    # parser.add_argument('--num-conversation-round', type=int, default=30)
    args = parser.parse_args()

    app = create_app()
    configure_routes(app, args)

    port = int(os.environ.get('PORT', 8080))
    app.run(port=port, host='0.0.0.0', debug=False)


if __name__ == '__main__':
    main()