Update rag_system.py

rag_system.py

@@ -8,6 +8,10 @@ from langchain.docstore.document import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 import pdfplumber
 from concurrent.futures import ThreadPoolExecutor
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import LLMChainExtractor
+from langgraph.graph import Graph
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
 
 # Load environment variables
 load_dotenv()
@@ -28,10 +32,19 @@ def load_retrieval_qa_chain():
     # Initialize ChatOpenAI model
     llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)  # "gpt-4o-mini
 
-    # Create ConversationalRetrievalChain
+    # Create a compressor for re-ranking
+    compressor = LLMChainExtractor.from_llm(llm)
+    
+    # Create a ContextualCompressionRetriever
+    compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor,
+        base_retriever=vectorstore.as_retriever()
+    )
+
+    # Create ConversationalRetrievalChain with the new retriever
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
-        vectorstore.as_retriever(),
+        retriever=compression_retriever,
         return_source_documents=True
     )
 
@@ -82,41 +95,58 @@ def update_embeddings():
             documents.extend(result)
     vectorstore.add_documents(documents)
 
-# Generate answer for a query
-def get_answer(qa_chain, query, chat_history):
+def create_rag_graph():
+    qa_chain = load_retrieval_qa_chain()
+
+    def retrieve_and_generate(inputs):
+        question = inputs["question"]
+        chat_history = inputs["chat_history"]
+        result = qa_chain({"question": question, "chat_history": chat_history})
+        
+        # Ensure source documents have the correct metadata
+        sources = []
+        for doc in result.get("source_documents", []):
+            if "source" in doc.metadata and "page" in doc.metadata:
+                sources.append(f"{os.path.basename(doc.metadata['source'])} (Page {doc.metadata['page']})")
+            else:
+                print(f"Warning: Document missing metadata: {doc.metadata}")
+        
+        return {
+            "answer": result["answer"],
+            "sources": sources
+        }
+
+    workflow = Graph()
+    workflow.add_node("retrieve_and_generate", retrieve_and_generate)
+    workflow.set_entry_point("retrieve_and_generate")
+
+    chain = workflow.compile()
+    return chain
+
+rag_chain = create_rag_graph()
+
+def get_answer(query, chat_history):
     formatted_history = [(q, a) for q, a in zip(chat_history[::2], chat_history[1::2])]
     
-    response = qa_chain.invoke({"question": query, "chat_history": formatted_history})
-    
-    answer = response["answer"]
+    response = rag_chain.invoke({"question": query, "chat_history": formatted_history})
     
-    source_docs = response.get("source_documents", [])
-    source_texts = [f"{os.path.basename(doc.metadata['source'])} (Page {doc.metadata['page']})" for doc in source_docs]
+    # Validate response format
+    if "answer" not in response or "sources" not in response:
+        print("Warning: Unexpected response format")
+        return {"answer": "Error in processing", "sources": []}
     
-    return {"answer": answer, "sources": source_texts}
+    return {"answer": response["answer"], "sources": response["sources"]}
 
 # Example usage
 if __name__ == "__main__":
     update_embeddings()  # Update embeddings with new documents
-    qa_chain = load_retrieval_qa_chain()
-    question = """당신은 RAG(Retrieval-Augmented Generation) 기반 AI 어시스턴트입니다. 다음 지침을 따라 사용자 질문에 답하세요:
-
-1. 검색 결과 활용: 제공된 검색 결과를 분석하고 관련 정보를 사용해 답변하세요.
-
-2. 정확성 유지: 정보의 정확성을 확인하고, 불확실한 경우 이를 명시하세요.
-
-3. 간결한 응답: 질문에 직접 답하고 핵심 내용에 집중하세요.
-
-4. 추가 정보 제안: 관련된 추가 정보가 있다면 언급하세요.
-
-5. 윤리성 고려: 객관적이고 중립적인 태도를 유지하세요.
-
-6. 한계 인정: 답변할 수 없는 경우 솔직히 인정하세요.
-
-7. 대화 유지: 자연스럽게 대화를 이어가고, 필요시 후속 질문을 제안하세요.
-항상 정확하고 유용한 정보를 제공하는 것을 목표로 하세요."""
-
-    response = get_answer(qa_chain, question, [])
+    question = "RAG 시스템에 대해 설명해주세요."
+    response = get_answer(question, [])
     print(f"Question: {question}")
     print(f"Answer: {response['answer']}")
-    print(f"Sources: {response['sources']}")
+    print(f"Sources: {response['sources']}")
+    
+    # Validate source format
+    for source in response['sources']:
+        if not (source.endswith(')') and ' (Page ' in source):
+            print(f"Warning: Unexpected source format: {source}")