Update processing.py

processing.py

@@ -7,6 +7,7 @@ from config import openai_api_key
 from langchain.chains import RetrievalQA
 import os
 import json
+import re
 
 # Initialize embeddings and FAISS index
 embedding_model = OpenAIEmbeddings(openai_api_key=openai_api_key)
@@ -46,6 +47,13 @@ def truncate_text(text: str, max_tokens: int = 10000) -> str:
     print(f"Text not truncated, contains {len(words)} words")
     return text
 
+def extract_json_from_text(text):
+    # Find the first occurrence of a JSON-like structure
+    match = re.search(r'\{(?:[^{}]|(?R))*\}', text, re.DOTALL)
+    if match:
+        return match.group(0)
+    return None
+
 def process_input(input_text: str, llm):
     general_task = load_text("tasks/general_task.txt")
     attachments_task = load_text("tasks/Attachments_task.txt")
@@ -84,30 +92,34 @@ Please provide a comprehensive analysis for each speaker, including:
 2. Big Five traits (use the format from the Big Five Traits Task)
 3. Personality disorders (use the format from the Personality Disorders Task)
 
-Respond with a JSON object containing an array of speaker analyses. Each speaker analysis should include all three aspects mentioned above.
+Respond with a JSON object containing an array of speaker analyses under the key 'speaker_analyses'. Each speaker analysis should include all three aspects mentioned above.
 
 Analysis:"""
 
     messages = [HumanMessage(content=prompt)]
     response = llm.invoke(messages)
     
-    # Print the raw LLM model output
     print("Raw LLM Model Output:")
     print(response.content)
     print("\n" + "-"*50 + "\n")  # Separator for readability
 
     try:
-        # Parse the response as JSON
-        parsed_json = json.loads(response.content)
+        # Extract JSON from the response
+        json_str = extract_json_from_text(response.content)
+        if not json_str:
+            raise ValueError("No valid JSON structure found in the response")
+        
+        # Parse the JSON
+        parsed_json = json.loads(json_str)
         
-        # Print the parsed JSON before further processing
         print("Parsed JSON Output:")
         print(json.dumps(parsed_json, indent=2))
         print("\n" + "-"*50 + "\n")  # Separator for readability
         
         # Process the parsed JSON
         results = {}
-        for speaker_analysis in parsed_json:
+        speaker_analyses = parsed_json.get('speaker_analyses', [])
+        for speaker_analysis in speaker_analyses:
             speaker_id = speaker_analysis.get('speaker', 'Unknown Speaker')
             results[speaker_id] = {
                 'attachments': attachment_parser.parse_object(speaker_analysis.get('attachment_styles', {})),
@@ -115,7 +127,36 @@ Analysis:"""
                 'personalities': personality_parser.parse_object(speaker_analysis.get('personality_disorders', {}))
             }
         
+        if not results:
+            print("Warning: No speaker analyses found in the parsed JSON.")
+            return {"Unknown Speaker": {
+                'attachments': attachment_parser.parse_object({}),
+                'bigfive': bigfive_parser.parse_object({}),
+                'personalities': personality_parser.parse_object({})
+            }}
+        
         return results
+    except json.JSONDecodeError as e:
+        print(f"Error parsing JSON: {e}")
+        print("Raw content causing the error:")
+        print(response.content)
+        return {"Unknown Speaker": {
+            'attachments': attachment_parser.parse_object({}),
+            'bigfive': bigfive_parser.parse_object({}),
+            'personalities': personality_parser.parse_object({})
+        }}
     except Exception as e:
-        print(f"Error parsing output: {e}")
-        return None
+        print(f"Unexpected error: {e}")
+        return {"Unknown Speaker": {
+            'attachments': attachment_parser.parse_object({}),
+            'bigfive': bigfive_parser.parse_object({}),
+            'personalities': personality_parser.parse_object({})
+        }}
+
+# For testing purposes
+if __name__ == "__main__":
+    test_input = "This is a test input for processing."
+    result = process_input(test_input, llm)
+    print("\nProcessing completed.")
+    print("Final results:")
+    print(json.dumps(result, indent=2))