Update output_parser.py

output_parser.py

@@ -1,40 +1,164 @@
-from pydantic import BaseModel
-from langchain.output_parsers import PydanticOutputParser
+# output_parser.py
+
+from typing import List
+from pydantic import BaseModel, Field
+import json
 
 class AttachmentStyle(BaseModel):
     speaker: str
-    secured: float
-    anxious_preoccupied: float
-    dismissive_avoidant: float
-    fearful_avoidant: float
-    self_rating: int
-    others_rating: int
-    anxiety: int
-    avoidance: int
+    secured: float = Field(..., ge=0, le=1)
+    anxious_preoccupied: float = Field(..., ge=0, le=1)
+    dismissive_avoidant: float = Field(..., ge=0, le=1)
+    fearful_avoidant: float = Field(..., ge=0, le=1)
+    self_rating: int = Field(..., ge=0, le=10)
+    others_rating: int = Field(..., ge=0, le=10)
+    anxiety: int = Field(..., ge=0, le=10)
+    avoidance: int = Field(..., ge=0, le=10)
     explanation: str
 
 class BigFiveTraits(BaseModel):
     speaker: str
-    extraversion: int
-    agreeableness: int
-    conscientiousness: int
-    neuroticism: int
-    openness: int
+    extraversion: int = Field(..., ge=-10, le=10)
+    agreeableness: int = Field(..., ge=-10, le=10)
+    conscientiousness: int = Field(..., ge=-10, le=10)
+    neuroticism: int = Field(..., ge=-10, le=10)
+    openness: int = Field(..., ge=-10, le=10)
     explanation: str
 
 class PersonalityDisorder(BaseModel):
     speaker: str
-    depressed: int
-    paranoid: int
-    schizoid_schizotypal: int
-    antisocial_psychopathic: int
-    borderline_dysregulated: int
-    narcissistic: int
-    anxious_avoidant: int
-    dependent_victimized: int
-    obsessional: int
+    depressed: int = Field(..., ge=0, le=5)
+    paranoid: int = Field(..., ge=0, le=5)
+    schizoid_schizotypal: int = Field(..., ge=0, le=5)
+    antisocial_psychopathic: int = Field(..., ge=0, le=5)
+    borderline_dysregulated: int = Field(..., ge=0, le=5)
+    narcissistic: int = Field(..., ge=0, le=5)
+    anxious_avoidant: int = Field(..., ge=0, le=5)
+    dependent_victimized: int = Field(..., ge=0, le=5)
+    obsessional: int = Field(..., ge=0, le=5)
     explanation: str
 
-attachment_parser = PydanticOutputParser(pydantic_object=AttachmentStyle)
-bigfive_parser = PydanticOutputParser(pydantic_object=BigFiveTraits)
-personality_parser = PydanticOutputParser(pydantic_object=PersonalityDisorder)
+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 list of objects, where each object represents a speaker and has the following format:
+        [
+          {
+            "speaker": "Speaker 1",
+            "secured": 0.5,
+            "anxious_preoccupied": 0.2,
+            "dismissive_avoidant": 0.1,
+            "fearful_avoidant": 0.2,
+            "self_rating": 7,
+            "others_rating": 6,
+            "anxiety": 3,
+            "avoidance": 4,
+            "explanation": "Explanation of the attachment style for Speaker 1"
+          },
+          {
+            "speaker": "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 list of objects, where each object represents a speaker and has the following format:
+        [
+          {
+            "speaker": "Speaker 1",
+            "extraversion": 5,
+            "agreeableness": 4,
+            "conscientiousness": 7,
+            "neuroticism": 3,
+            "openness": 6,
+            "explanation": "Explanation of the Big Five traits for Speaker 1"
+          },
+          {
+            "speaker": "Speaker 2",
+            ...
+          }
+        ]
+        """
+
+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 list of objects, where each object represents a speaker and has the following format:
+        [
+          {
+            "speaker": "Speaker 1",
+            "depressed": 1,
+            "paranoid": 0,
+            "schizoid_schizotypal": 0,
+            "antisocial_psychopathic": 0,
+            "borderline_dysregulated": 0,
+            "narcissistic": 2,
+            "anxious_avoidant": 3,
+            "dependent_victimized": 1,
+            "obsessional": 2,
+            "explanation": "Explanation of the personality disorders for Speaker 1"
+          },
+          {
+            "speaker": "Speaker 2",
+            ...
+          }
+        ]
+        """
+
+attachment_parser = AttachmentStyleOutputParser()
+bigfive_parser = BigFiveTraitsOutputParser()
+personality_parser = PersonalityDisorderOutputParser()