literal_thread_manager.py

from literalai import LiteralClient
from dotenv import load_dotenv
import os
from typing import List, Dict, Any
load_dotenv()


class LiteralThreadManager:
    """
        The LiteralThreadManager is responsible for managing and extracting conversation threads from LiteralAI db.
        It includes methods for retrieving threads, messages, processing chat history, and ensuring efficient and clear
        extraction of conversation details. This class is designed to support functions that handle user and
        assistant interactions within HarmonyAI.
    """

    def __init__(self, api_key: str):
        """
        Initialize the LiteralThreadManager with a LiteralClient.

        Args:
            api_key (str): The API key for the LiteralClient.
        """
        self.literal_client = LiteralClient(api_key=api_key)

    @staticmethod
    def threads_to_dict(threads_input) -> List[Dict[str, Any]]:
        """
        Convert a list of threads to a list of dictionaries.

        Args:
            threads_input: The input threads to be converted.

        Returns:
            List[Dict[str, Any]]: A list of dictionaries representing the threads.
        """
        return [a_thread.to_dict() for a_thread in threads_input.data]

    def filter_threads_by_participant(self, participant_name: str) -> List[Dict[str, Any]]:
        """
        Filter threads by participant name.

        Args:
            participant_name (str): The name of the participant.

        Returns:
            List[Dict[str, Any]]: A list of dictionaries representing the filtered threads.
        """
        self._all_threads = self.literal_client.api.get_threads()
        self.thread_dict_list = self.threads_to_dict(self._all_threads)
        return [thread for thread in self.thread_dict_list if
                thread['participant']['identifier'].lower() == participant_name.lower()]

    def filter_thread_by_id(self, thread_id: str) -> List[Dict[str, Any]]:
        """
        Filter a thread by its ID.

        Args:
            thread_id (str): The ID of the thread.

        Returns:
            List[Dict[str, Any]]: A list containing the dictionary representation of the thread.
        """
        return [self.literal_client.api.get_thread(thread_id).to_dict()]

    def get_other_partner_thread_id(self, thread_id: str) -> str:
        """
        Get the partner thread ID of a given thread.

        Args:
            thread_id (str): The ID of the thread.

        Returns:
            str: The partner thread ID.
        """
        current_thread = self.filter_thread_by_id(thread_id)
        partner_thread_id = current_thread[0].get("metadata")["partner_thread_id"]
        return partner_thread_id

    def get_messages_from_thread(self, input_thread) -> List[str]:
        """
        Get all messages from a thread.

        Args:
            input_thread: The input thread to retrieve messages from.

        Returns:
            List[str]: A list of messages from the thread.
        """
        steps_in_thread = input_thread[0]['steps']
        return steps_in_thread

    def get_user_name_from_thread(self, thread) -> str:
        """
        Get the participant name from a thread.

        Args:
            thread: The input thread to retrieve the participant name from.

        Returns:
            str: The participant's name.
        """
        return thread[0]['participant']['identifier']

    def is_conflict_resolved(self, thread_id: str) -> bool:
        """
        Check if the conflict in a thread is resolved.

        Args:
            thread_id (str): The ID of the thread.

        Returns:
            bool: True if the conflict is resolved, False otherwise.
        """
        if self.literal_client.api.get_thread(thread_id).metadata.get('isResolved') is None:
            self.set_conflict_resolved(thread_id, False)
        return self.literal_client.api.get_thread(thread_id).metadata['isResolved']

    def set_conflict_resolved(self, thread_id: str, is_resolved: bool):
        """
        Set the conflict resolution status of a thread.

        Args:
            thread_id (str): The ID of the thread.
            is_resolved (bool): The resolution status to be set.
        """
        thread = self.literal_client.api.get_thread(thread_id)
        thread.metadata['isResolved'] = is_resolved
        self.literal_client.api.update_thread(id=thread.id, metadata=thread.metadata)

    def extract_chat_history_from_thread(self, input_thread) -> List[Dict[str, str]]:
        """
        Extract chat history from a thread.

        Args:
            input_thread: The input thread to extract chat history from.

        Returns:
            List[Dict[str, str]]: A list of dictionaries representing the chat history.
        """
        conversation = self.get_messages_from_thread(input_thread)
        name = self.get_user_name_from_thread(input_thread)
        chat_history = []

        for convo in conversation:
            if convo['type'] in ['user_message', 'assistant_message']:
                role = 'user' if convo['type'] == 'user_message' else 'assistant'
                content = convo['output']['content']
                chat_history.append({
                    'role': role,
                    'name': name if role == 'user' else 'HarmonyAI',
                    'content': content
                })

            if 'generation' in convo and convo['generation'] is not None:
                if 'messages' in convo['generation']:
                    for message in convo['generation']['messages']:
                        if message['role'] in ['user', 'assistant']:
                            chat_history.append({
                                'role': message['role'],
                                'name': name if message['role'] == 'user' else 'HarmonyAI',
                                'content': message['content']
                            })
                if 'messageCompletion' in convo['generation'] and convo['generation']['messageCompletion'] is not None:
                    message_completion = convo['generation']['messageCompletion']
                    chat_history.append({
                        'role': message_completion['role'],
                        'name': 'HarmonyAI',
                        'content': message_completion['content']
                    })

        filtered_chat_history = []
        for message in chat_history:
            if not any(
                    m['content'] == message['content'] and m['role'] == message['role'] for m in filtered_chat_history):
                filtered_chat_history.append(message)
        # print(conversation)

        for message in conversation:
            if message.get('type') == 'user_message':
                last_user_message = message['output']['content']
        filtered_chat_history.append({
            'role': 'user',
            'name': name,
            'content': last_user_message
        })

        return filtered_chat_history

    def count_llm_messages(self, thread_id: str) -> int:
        """
        Count the number of AI-generated messages in a thread.

        Args:
            thread_id (str): The ID of the thread.

        Returns:
            int: The number of AI-generated messages in the thread.
        """
        thread = self.literal_client.api.get_thread(thread_id)
        return sum(step.type == 'llm' for step in thread.steps)

    def send_message(self, thread_id: str, message: str):
        """
        Send a message to a thread.

        Args:
            thread_id (str): The ID of the thread.
            message (str): The message content to be sent.
        """
        self.literal_client.api.create_step(thread_id=thread_id, type='assistant_message', name='HarmonyAI',
                                            output={'content': message})


# Example usage:
if __name__ == "__main__":
    manager = LiteralThreadManager(api_key=os.getenv("LITERAL_API_KEY"))

    # Filter threads for "tom"
    # all_threads = manager.literal_client.api.get_threads()
    # print(all_threads)
    # exit()

    # tom_threads = manager.filter_threads_by_participant("tom")
    #print("Threads involving Tom:")
    #print(tom_threads)

    # Filter thread by ID
    thread_id = '1bb44dc5-0b81-42e9-84a2-cd34b6fe8480'
    linda_thrad_id = '76d03507-4084-46ba-ba7e-734be1f58304'
    thread_content = manager.filter_thread_by_id(linda_thrad_id)
    #print(thread_content)
    #exit()

    chat = manager.extract_chat_history_from_thread(thread_content)
    print(f"\nChat history for thread ID {thread_id}:")
    for message in chat:
        print(f"{message['role']} - {message['name']}: {message['content']}")