Update processing.py

processing.py

@@ -7,7 +7,6 @@ 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 typing import List, Any, Optional
 from pydantic import Field
 from langchain_core.callbacks import CallbackManagerForRetrieverRun
@@ -75,15 +74,12 @@ combined_retriever = CombinedRetriever(retrievers=[text_retriever, attachments_r
 
 # Create prompt template for query generation
 prompt_template = PromptTemplate(
-    input_variables=["question", "context"],
-    template="Use the following context to answer the question: {context}\n\nQuestion: {question}\nAnswer:"
+    input_variables=["question"],
+    template="Generate multiple search queries for the following question: {question}"
 )
 
 # Create query generation chain
-query_generation_chain = PromptTemplate(
-    input_variables=["question"],
-    template="Generate multiple search queries for the following question: {question}"
-) | llm
+query_generation_chain = prompt_template | llm
 
 # Create multi-query retrieval chain
 def generate_queries(input):
@@ -96,19 +92,13 @@ def multi_query_retrieve(input):
     for query in queries:
         docs = combined_retriever.get_relevant_documents(query)
         all_docs.extend(docs)
-    return {"queries": queries, "documents": all_docs}
+    return all_docs
 
 multi_query_retriever = RunnableLambda(multi_query_retrieve)
 
-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,
-        "question": RunnablePassthrough()
-    }
+    {"context": multi_query_retriever, "question": RunnablePassthrough()}
     | prompt_template
     | llm
 )
@@ -132,33 +122,13 @@ def process_input(input_text: str, llm):
 
     truncated_input = truncate_text(input_text)
 
-    # Get the retrieval results and LLM output
-    retrieval_result = multi_query_retrieve(truncated_input)
+    relevant_docs = qa_chain.invoke({"query": truncated_input})
     
-    # Print the generated queries
-    print("Generated Queries:")
-    for query in retrieval_result["queries"]:
-        print(f"- {query}")
-
-    # Print the retrieved documents
-    print("\nRetrieved Documents:")
-    for i, doc in enumerate(retrieval_result["documents"], 1):
-        print(f"Document {i}:")
-        print(f"Content: {doc.page_content}")
-        print("-" * 50)
-    
-    # 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)
+    retrieved_knowledge = str(relevant_docs)
 
     prompt = f"""
 {general_task}
-General Impression Task:
+Genral Impression Task:
 {general_impression_task}
 Attachment Styles Task:
 {attachments_task}
@@ -176,6 +146,11 @@ 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:")
@@ -196,7 +171,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)
@@ -235,5 +210,4 @@ Analysis:"""
             'attachments': empty_analysis.attachment_style,
             'bigfive': empty_analysis.big_five_traits,
             'personalities': empty_analysis.personality_disorder
-        }}
-
+        }}