Update processing.py

processing.py

@@ -6,8 +6,8 @@ from llm_loader import load_model, count_tokens
 from config import openai_api_key
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
 from langchain.schema.runnable import RunnablePassthrough
-from langchain_core.runnables import RunnableLambda
 from typing import List, Any, Optional
 from pydantic import Field
 from langchain_core.callbacks import CallbackManagerForRetrieverRun
@@ -35,7 +35,7 @@ for key, file_path in knowledge_files.items():
 text_faiss_index = FAISS.from_texts(documents, embedding_model)
 
 # Load pre-existing FAISS indexes
-attachments_faiss_index = FAISS.load_local("knowledge/faiss_index_Attachments_db", embedding_model, allow_dangerous_deserialization=True)
+attachments_faiss_index = FAISS.load_local("knowledge/faiss_index_Attachment_db", embedding_model, allow_dangerous_deserialization=True)
 personalities_faiss_index = FAISS.load_local("knowledge/faiss_index_Personalities_db", embedding_model, allow_dangerous_deserialization=True)
 
 # Initialize LLM
@@ -75,12 +75,15 @@ combined_retriever = CombinedRetriever(retrievers=[text_retriever, attachments_r
 
 # Create prompt template for query generation
 prompt_template = PromptTemplate(
-    input_variables=["question"],
-    template="Generate multiple search queries for the following question: {question}"
+    input_variables=["question", "context"],
+    template="Use the following context to answer the question: {context}\n\nQuestion: {question}\nAnswer:"
 )
 
 # Create query generation chain
-query_generation_chain = prompt_template | llm
+query_generation_chain = PromptTemplate(
+    input_variables=["question"],
+    template="Generate multiple search queries for the following question: {question}"
+) | llm
 
 # Create multi-query retrieval chain
 def generate_queries(input):
@@ -97,10 +100,10 @@ def multi_query_retrieve(input):
 
 multi_query_retriever = RunnableLambda(multi_query_retrieve)
 
-# Create QA chain with multi-query retriever
 def format_docs(docs):
     return "\n\n".join(doc.page_content for doc in docs["documents"])
 
+# Create QA chain with multi-query retriever
 qa_chain = (
     {
         "context": multi_query_retriever | format_docs,
@@ -129,25 +132,33 @@ def process_input(input_text: str, llm):
 
     truncated_input = truncate_text(input_text)
 
-    relevant_docs = qa_chain.invoke({"query": truncated_input})
+    # Get the retrieval results and LLM output
+    retrieval_result = multi_query_retrieve(truncated_input)
     
     # Print the generated queries
     print("Generated Queries:")
-    for query in relevant_docs["retrieval_results"]["queries"]:
+    for query in retrieval_result["queries"]:
         print(f"- {query}")
 
     # Print the retrieved documents
     print("\nRetrieved Documents:")
-    for i, doc in enumerate(relevant_docs["retrieval_results"]["documents"], 1):
+    for i, doc in enumerate(retrieval_result["documents"], 1):
         print(f"Document {i}:")
-        print(f"Content: {doc.page_content}...")  # Print first 200 characters
+        print(f"Content: {doc.page_content")
         print("-" * 50)
-        
-    retrieved_knowledge = str(relevant_docs["llm_output"])
+    
+    # Format the retrieved documents
+    formatted_docs = format_docs(retrieval_result)
+    
+    # Generate the LLM response
+    llm_input = prompt_template.format(question=truncated_input, context=formatted_docs)
+    llm_output = llm.invoke(llm_input)
+    
+    retrieved_knowledge = str(llm_output.content)
 
     prompt = f"""
 {general_task}
-Genral Impression Task:
+General Impression Task:
 {general_impression_task}
 Attachment Styles Task:
 {attachments_task}
@@ -165,11 +176,6 @@ 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 four aspects mentioned above, however, General impressions must not be in json or dict format.
 Analysis:"""
 
-    #truncated_input_tokents_count = count_tokens(truncated_input)
-    #print('truncated_input_tokents_count:', truncated_input_tokents_count)
-    #input_tokens_count = count_tokens(prompt)
-    #print('input_tokens_count', input_tokens_count)
-    
     response = llm.invoke(prompt)
     
     print("Raw LLM Model Output:")
@@ -190,7 +196,7 @@ Analysis:"""
             speaker_id = f"Speaker {i}"
             parsed_analysis = output_parser.parse_speaker_analysis(speaker_analysis)
 
-                # Convert general_impression to string if it's a dict or JSON object
+            # Convert general_impression to string if it's a dict or JSON object
             general_impression = parsed_analysis.general_impression
             if isinstance(general_impression, dict):
                 general_impression = json.dumps(general_impression)
@@ -229,4 +235,5 @@ Analysis:"""
             'attachments': empty_analysis.attachment_style,
             'bigfive': empty_analysis.big_five_traits,
             'personalities': empty_analysis.personality_disorder
-        }}
+        }}
+