Update processing.py

processing.py

@@ -1,12 +1,12 @@
 from langchain.schema import HumanMessage, BaseRetriever
 from output_parser import output_parser
-from langchain_openai import OpenAIEmbeddings
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_community.vectorstores import FAISS
 from llm_loader import load_model
 from config import openai_api_key
-from langchain.chains import RetrievalQA, LLMChain
-from langchain.retrievers import MultiQueryRetriever
+from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
 from typing import List
 from pydantic import Field
 import os
@@ -51,17 +51,10 @@ class CombinedRetriever(BaseRetriever):
     class Config:
         arbitrary_types_allowed = True
 
-    def get_relevant_documents(self, query: str):
+    def invoke(self, input):
         combined_docs = []
         for retriever in self.retrievers:
-            docs = retriever.get_relevant_documents(query)
-            combined_docs.extend(docs)
-        return combined_docs
-
-    async def aget_relevant_documents(self, query: str):
-        combined_docs = []
-        for retriever in self.retrievers:
-            docs = await retriever.aget_relevant_documents(query)
+            docs = retriever.invoke(input)
             combined_docs.extend(docs)
         return combined_docs
 
@@ -74,18 +67,22 @@ prompt_template = PromptTemplate(
     template="Generate multiple search queries for the following question: {question}"
 )
 
-# Create LLM chain for query generation
-llm_chain = LLMChain(llm=llm, prompt=prompt_template)
+# Create query generation chain
+query_generation_chain = prompt_template | llm
 
-# Initialize MultiQueryRetriever
-multi_query_retriever = MultiQueryRetriever(
-    retriever=combined_retriever,
-    llm_chain=llm_chain,
-    parser_key="lines"  # Assuming the LLM outputs queries line by line
-)
+# Create multi-query retrieval chain
+def generate_queries(input):
+    queries = query_generation_chain.invoke({"question": input}).content.split('\n')
+    return [query.strip() for query in queries if query.strip()]
+
+multi_query_retriever = RunnableLambda(generate_queries) | combined_retriever
 
 # Create QA chain with multi-query retriever
-qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=multi_query_retriever)
+qa_chain = (
+    {"context": multi_query_retriever, "question": RunnablePassthrough()}
+    | prompt_template
+    | llm
+)
 
 def load_text(file_path: str) -> str:
     with open(file_path, 'r', encoding='utf-8') as file:
@@ -107,10 +104,7 @@ def process_input(input_text: str, llm):
 
     relevant_docs = qa_chain.invoke({"query": truncated_input})
     
-    if isinstance(relevant_docs, dict) and 'result' in relevant_docs:
-        retrieved_knowledge = relevant_docs['result']
-    else:
-        retrieved_knowledge = str(relevant_docs)
+    retrieved_knowledge = str(relevant_docs)
 
     prompt = f"""{general_task}
 Attachment Styles Task:
@@ -129,8 +123,7 @@ 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.
 Analysis:"""
 
-    messages = [HumanMessage(content=prompt)]
-    response = llm.invoke(messages)
+    response = llm.invoke(prompt)
     
     print("Raw LLM Model Output:")
     print(response.content)
@@ -176,3 +169,9 @@ Analysis:"""
             'bigfive': empty_analysis.big_five_traits,
             'personalities': empty_analysis.personality_disorder
         }}
+
+# Example usage
+if __name__ == "__main__":
+    input_text = "Your input text here"
+    result = process_input(input_text, llm)
+    print(json.dumps(result, indent=2))