Update processing.py

processing.py

@@ -1,32 +1,108 @@
-from langchain.schema import HumanMessage
+from langchain.schema import HumanMessage, BaseRetriever, Document
 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
+from langchain.prompts import PromptTemplate
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
+from typing import List, Any, Optional
+from pydantic import Field
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
 import os
 import json
 
+# Initialize embedding model
 embedding_model = OpenAIEmbeddings(openai_api_key=openai_api_key)
 
+# Define knowledge files
 knowledge_files = {
     "attachments": "knowledge/bartholomew_attachments_definitions.txt",
     "bigfive": "knowledge/bigfive_definitions.txt",
     "personalities": "knowledge/personalities_definitions.txt"
 }
 
+# Load text-based knowledge
 documents = []
 for key, file_path in knowledge_files.items():
     with open(file_path, 'r', encoding='utf-8') as file:
         content = file.read().strip()
         documents.append(content)
 
-faiss_index = FAISS.from_texts(documents, embedding_model)
+# Create FAISS index from text documents
+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)
+personalities_faiss_index = FAISS.load_local("knowledge/faiss_index_Personalities_db", embedding_model, allow_dangerous_deserialization=True)
+
+# Initialize LLM
 llm = load_model(openai_api_key)
 
-qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=faiss_index.as_retriever())
+# Create retrievers for each index
+text_retriever = text_faiss_index.as_retriever()
+attachments_retriever = attachments_faiss_index.as_retriever()
+personalities_retriever = personalities_faiss_index.as_retriever()
+
+class CombinedRetriever(BaseRetriever):
+    retrievers: List[BaseRetriever] = Field(default_factory=list)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None
+    ) -> List[Document]:
+        combined_docs = []
+        for retriever in self.retrievers:
+            docs = retriever.get_relevant_documents(query, run_manager=run_manager)
+            combined_docs.extend(docs)
+        return combined_docs
+
+    async def _aget_relevant_documents(
+        self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None
+    ) -> List[Document]:
+        combined_docs = []
+        for retriever in self.retrievers:
+            docs = await retriever.aget_relevant_documents(query, run_manager=run_manager)
+            combined_docs.extend(docs)
+        return combined_docs
+
+# Create an instance of the combined retriever
+combined_retriever = CombinedRetriever(retrievers=[text_retriever, attachments_retriever, personalities_retriever])
+
+# Create prompt template for query generation
+prompt_template = PromptTemplate(
+    input_variables=["question"],
+    template="Generate multiple search queries for the following question: {question}"
+)
+
+# Create query generation chain
+query_generation_chain = prompt_template | llm
+
+# 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()]
+
+def multi_query_retrieve(input):
+    queries = generate_queries(input)
+    all_docs = []
+    for query in queries:
+        docs = combined_retriever.get_relevant_documents(query)
+        all_docs.extend(docs)
+    return all_docs
+
+multi_query_retriever = RunnableLambda(multi_query_retrieve)
+
+# Create QA chain with 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:
@@ -48,10 +124,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:
@@ -70,8 +143,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)
@@ -116,4 +188,10 @@ Analysis:"""
             'attachments': empty_analysis.attachment_style,
             '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))