output_parser.py

from typing import List, Optional
from pydantic import BaseModel, Field
import json

class AttachmentStyle(BaseModel):
    Speaker: str = "Unknown Speaker"
    Secured: float = Field(0.0, alias="Secured")
    Anxious_Preoccupied: float = Field(0.0, alias="Anxious-Preoccupied")
    Dismissive_Avoidant: float = Field(0.0, alias="Dismissive-Avoidant")
    Fearful_Avoidant: float = Field(0.0, alias="Fearful-Avoidant")
    Self: int = Field(0, alias="Self")
    Others: int = Field(0, alias="Others")
    Anxiety: int = 0
    Avoidance: int = 0
    Explanation: str = "No explanation provided"

class BigFiveTraits(BaseModel):
    Speaker: str = "Unknown Speaker"
    Extraversion: int = 0
    Agreeableness: int = 0
    Conscientiousness: int = 0
    Neuroticism: int = 0
    Openness: int = 0
    Explanation: str = "No explanation provided"

class PersonalityDisorder(BaseModel):
    Speaker: str = "Unknown Speaker"
    Depressed: int = 0
    Paranoid: int = 0
    Schizoid_Schizotypal: int = Field(0, alias="Schizoid-Schizotypal")
    Antisocial_Psychopathic: int = Field(0, alias="Antisocial-Psychopathic")
    Borderline_Dysregulated: int = Field(0, alias="Borderline-Dysregulated")
    Narcissistic: int = 0
    Anxious_Avoidant: int = Field(0, alias="Anxious-Avoidant")
    Dependent_Victimized: int = Field(0, alias="Dependent-Victimized")
    Obsessional: int = 0
    Explanation: str = "No explanation provided"

class AttachmentStyleOutputParser:
    def parse(self, text: str) -> List[AttachmentStyle]:
        try:
            data = json.loads(text)
            if isinstance(data, list):
                return [AttachmentStyle(**item) for item in data]
            elif isinstance(data, dict):
                return [AttachmentStyle(**data)]
            else:
                raise ValueError("Invalid JSON structure")
        except json.JSONDecodeError:
            raise ValueError(f"Invalid JSON: {text}")
        except Exception as e:
            raise ValueError(f"Error parsing AttachmentStyle: {str(e)}")

    def parse_object(self, obj: dict) -> AttachmentStyle:
        return AttachmentStyle(**obj)

    def get_format_instructions(self) -> str:
        return """
        The output should be a JSON object with a "speaker_analyses" key containing an array of objects, where each object represents a speaker and has the following format:
        {
          "speaker_analyses": [
            {
              "Speaker": "1",
              "attachment_styles": {
                "Secured": 0.5,
                "Anxious-Preoccupied": 0.2,
                "Dismissive-Avoidant": 0.1,
                "Fearful-Avoidant": 0.2,
                "Self": 7,
                "Others": 6,
                "Anxiety": 3,
                "Avoidance": 4,
                "Explanation": "Explanation of the attachment style for Speaker 1"
              },
              "big_five_traits": {
                "Extraversion": 5,
                "Agreeableness": 6,
                "Conscientiousness": 7,
                "Neuroticism": 4,
                "Openness": 5,
                "Explanation": "Explanation of the Big Five traits for Speaker 1"
              },
              "personality_disorders": {
                "Depressed": 0,
                "Paranoid": 0,
                "Schizoid-Schizotypal": 0,
                "Antisocial-Psychopathic": 0,
                "Borderline-Dysregulated": 0,
                "Narcissistic": 1,
                "Anxious-Avoidant": 2,
                "Dependent-Victimized": 0,
                "Obsessional": 1,
                "Explanation": "Explanation of the personality disorders for Speaker 1"
              }
            },
            {
              "Speaker": "2",
              ...
            }
          ]
        }
        """

class BigFiveTraitsOutputParser:
    def parse(self, text: str) -> List[BigFiveTraits]:
        try:
            data = json.loads(text)
            if isinstance(data, list):
                return [BigFiveTraits(**item) for item in data]
            elif isinstance(data, dict):
                return [BigFiveTraits(**data)]
            else:
                raise ValueError("Invalid JSON structure")
        except json.JSONDecodeError:
            raise ValueError(f"Invalid JSON: {text}")
        except Exception as e:
            raise ValueError(f"Error parsing BigFiveTraits: {str(e)}")

    def parse_object(self, obj: dict) -> BigFiveTraits:
        return BigFiveTraits(**obj)

    def get_format_instructions(self) -> str:
        return """
        The output should be a JSON object with a "speaker_analyses" key containing an array of objects, where each object represents a speaker and includes a "big_five_traits" section.
        """

class PersonalityDisorderOutputParser:
    def parse(self, text: str) -> List[PersonalityDisorder]:
        try:
            data = json.loads(text)
            if isinstance(data, list):
                return [PersonalityDisorder(**item) for item in data]
            elif isinstance(data, dict):
                return [PersonalityDisorder(**data)]
            else:
                raise ValueError("Invalid JSON structure")
        except json.JSONDecodeError:
            raise ValueError(f"Invalid JSON: {text}")
        except Exception as e:
            raise ValueError(f"Error parsing PersonalityDisorder: {str(e)}")

    def parse_object(self, obj: dict) -> PersonalityDisorder:
        return PersonalityDisorder(**obj)

    def get_format_instructions(self) -> str:
        return """
        The output should be a JSON object with a "speaker_analyses" key containing an array of objects, where each object represents a speaker and includes a "personality_disorders" section.
        """

attachment_parser = AttachmentStyleOutputParser()
bigfive_parser = BigFiveTraitsOutputParser()
personality_parser = PersonalityDisorderOutputParser()