Update processing.py

processing.py

@@ -2,9 +2,10 @@ from langchain.schema import HumanMessage
 from output_parser import attachment_parser, bigfive_parser, personality_parser
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import FAISS
-from llm_loader import load_model  # Import the function to load the model
-from config import openai_api_key  # Import the API key from config.py
+from llm_loader import load_model
+from config import openai_api_key
 from langchain.chains import RetrievalQA
+import json
 import os
 
 # Initialize embeddings and FAISS index
@@ -57,8 +58,8 @@ def process_task(llm, input_text: str, general_task: str, specific_task: str, ou
     truncated_input = truncate_text(input_text)
 
     # Perform retrieval to get the most relevant context
-    query = f"{general_task}\n\n{specific_task}\n\n{truncated_input}"
-    retrieved_knowledge = qa_chain.run(query)
+    relevant_docs = qa_chain({"query": truncated_input})
+    retrieved_knowledge = "\n".join([doc.page_content for doc in relevant_docs['documents']])
 
     # Combine the retrieved knowledge with the original prompt
     prompt = f"""{general_task}
@@ -78,7 +79,11 @@ Analysis:"""
     print(response)
 
     try:
-        parsed_output = output_parser.parse(response.content)
+        # Parse the response as JSON
+        parsed_json = json.loads(response.content)
+        
+        # Validate and convert each item in the list
+        parsed_output = [output_parser.parse_object(item) for item in parsed_json]
         return parsed_output
     except Exception as e:
         print(f"Error parsing output: {e}")
@@ -99,17 +104,11 @@ def process_input(input_text: str, llm):
         specific_task = load_text(task_file)
         task_result = process_task(llm, input_text, general_task, specific_task, parser, qa_chain)
         
-        # If the result is a list, assume it's for multiple speakers
-        if isinstance(task_result, list):
+        if task_result:
             for i, speaker_result in enumerate(task_result):
                 speaker_id = f"Speaker {i+1}"
                 if speaker_id not in results:
                     results[speaker_id] = {}
                 results[speaker_id][task_name] = speaker_result
-        else:
-            # If it's not a list, assume it's for a single speaker
-            if "Speaker 1" not in results:
-                results["Speaker 1"] = {}
-            results["Speaker 1"][task_name] = task_result
 
     return results