output_parser.py

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

class AttachmentStyle(BaseModel):
    Secured: float = 0.0
    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 = 0
    Others: int = 0
    Anxiety: int = 0
    Avoidance: int = 0
    Explanation: str = "No explanation provided"

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

class PersonalityDisorder(BaseModel):
    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 SpeakerAnalysis(BaseModel):
    Speaker: str = "Unknown Speaker"
    GeneralImpression: str = "No general impression provided"
    AttachmentStyle: AttachmentStyle = AttachmentStyle()
    BigFiveTraits: BigFiveTraits = BigFiveTraits()
    PersonalityDisorder: PersonalityDisorder = PersonalityDisorder()

class OutputParser:
    def parse(self, text: str) -> List[SpeakerAnalysis]:
        try:
            data = json.loads(text)
            if "speaker_analyses" in data and isinstance(data["speaker_analyses"], list):
                return [self.parse_speaker_analysis(item) for item in data["speaker_analyses"]]
            else:
                raise ValueError("Invalid JSON structure: missing or invalid 'speaker_analyses' key")
        except json.JSONDecodeError:
            raise ValueError(f"Invalid JSON: {text}")
        except Exception as e:
            raise ValueError(f"Error parsing output: {str(e)}")

    def parse_speaker_analysis(self, obj: dict) -> SpeakerAnalysis:
        return SpeakerAnalysis(
            Speaker=obj.get("Speaker", "Unknown Speaker"),
            GeneralImpression=obj.get("General Impression", "No general impression provided"),
            AttachmentStyle=AttachmentStyle(**(obj.get("Attachment Styles", {}))),
            BigFiveTraits=BigFiveTraits(**(obj.get("Big Five Traits", {}))),
            PersonalityDisorder=PersonalityDisorder(**(obj.get("Personality Disorders", {})))
        )

    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",
              "GeneralImpression": "General impression of 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",
              ...
            }
          ]
        }
        """

output_parser = OutputParser()