Update processing.py

processing.py

@@ -8,10 +8,8 @@ from langchain.chains import RetrievalQA
 import os
 import json
 
-# Initialize embeddings and FAISS index
 embedding_model = OpenAIEmbeddings(openai_api_key=openai_api_key)
 
-# Load the content of knowledge files
 knowledge_files = {
     "attachments": "knowledge/bartholomew_attachments_definitions.txt",
     "bigfive": "knowledge/bigfive_definitions.txt",
@@ -24,13 +22,10 @@ for key, file_path in knowledge_files.items():
         content = file.read().strip()
         documents.append(content)
 
-# Create a FAISS index from the knowledge documents
 faiss_index = FAISS.from_texts(documents, embedding_model)
 
-# Load the LLM
 llm = load_model(openai_api_key)
 
-# Initialize the retrieval chain
 qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=faiss_index.as_retriever())
 
 def load_text(file_path: str) -> str:
@@ -40,10 +35,7 @@ def load_text(file_path: str) -> str:
 def truncate_text(text: str, max_tokens: int = 10000) -> str:
     words = text.split()
     if len(words) > max_tokens:
-        truncated_text = ' '.join(words[:max_tokens])
-        print(f"Text truncated from {len(words)} to {max_tokens} words")
-        return truncated_text
-    print(f"Text not truncated, contains {len(words)} words")
+        return ' '.join(words[:max_tokens])
     return text
 
 def process_input(input_text: str, llm):
@@ -54,16 +46,13 @@ def process_input(input_text: str, llm):
 
     truncated_input = truncate_text(input_text)
 
-    # Perform retrieval to get the most relevant context
     relevant_docs = qa_chain.invoke({"query": truncated_input})
     
-    # Extract the retrieved knowledge
     if isinstance(relevant_docs, dict) and 'result' in relevant_docs:
         retrieved_knowledge = relevant_docs['result']
     else:
         retrieved_knowledge = str(relevant_docs)
 
-    # Combine all tasks and knowledge into a single prompt
     prompt = f"""{general_task}
 
 Attachment Styles Task:
@@ -86,45 +75,28 @@ Please provide a comprehensive analysis for each speaker, including:
 
 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.
 
-The speaker format should be "Speaker [number]", for example:
-Speaker [1]
-Secured: [probability]
-Anxious-Preoccupied: [probability]
-Dismissive-Avoidant: [probability]
-Fearful-Avoidant: [probability]
-Self: [rating]
-Others: [rating]
-Anxiety: [rating]
-Avoidance: [rating]
-Explanation: [very brief explanation]
-
 Analysis:"""
 
     messages = [HumanMessage(content=prompt)]
     response = llm.invoke(messages)
     
+    print("Raw LLM Model Output:")
+    print(response.content)
+
     try:
-        # Remove code block markers if present
         content = response.content
         if content.startswith("```json"):
             content = content.split("```json", 1)[1]
         if content.endswith("```"):
             content = content.rsplit("```", 1)[0]
         
-        # Parse the JSON
         parsed_json = json.loads(content.strip())
         
-        # Process the parsed JSON
         results = {}
         speaker_analyses = parsed_json.get('speaker_analyses', [])
         for speaker_analysis in speaker_analyses:
             speaker_id = speaker_analysis.get('speaker', 'Unknown Speaker')
-            
-            # Extract speaker number
-            speaker_number = speaker_id.split('[')[-1].split(']')[0]
-            speaker_key = f"Speaker {speaker_number}"
-            
-            results[speaker_key] = {
+            results[speaker_id] = {
                 'attachments': attachment_parser.parse_object(speaker_analysis.get('attachment_styles', {})),
                 'bigfive': bigfive_parser.parse_object(speaker_analysis.get('big_five_traits', {})),
                 'personalities': personality_parser.parse_object(speaker_analysis.get('personality_disorders', {}))
@@ -139,18 +111,10 @@ Analysis:"""
             }}
         
         return results
-    except json.JSONDecodeError as e:
-        print(f"Error parsing JSON: {e}")
-        print("Raw content causing the error:")
-        print(response.content)
     except Exception as e:
-        print(f"Unexpected error: {e}")
-        print("Raw content:")
-        print(response.content)
-    
-    # If any error occurs, return a default result
-    return {"Unknown Speaker": {
-        'attachments': attachment_parser.parse_object({}),
-        'bigfive': bigfive_parser.parse_object({}),
-        'personalities': personality_parser.parse_object({})
-    }}
+        print(f"Error processing input: {e}")
+        return {"Unknown Speaker": {
+            'attachments': attachment_parser.parse_object({}),
+            'bigfive': bigfive_parser.parse_object({}),
+            'personalities': personality_parser.parse_object({})
+        }}