Update rag_system.py

rag_system.py

@@ -8,11 +8,6 @@ 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
-from langchain.prompts import PromptTemplate
 
 # Load environment variables
 load_dotenv()
@@ -33,40 +28,11 @@ def load_retrieval_qa_chain():
     # Initialize ChatOpenAI model
     llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)  # "gpt-4o-mini
 
-    # 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()
-    )
-
-    # Define your instruction/prompt
-    instruction = """당신은 RAG(Retrieval-Augmented Generation) 기반 AI 어시스턴트입니다. 다음 지침을 따라 사용자 질문에 답하세요:
-
-1. ��색 결과 활용: 제공된 검색 결과를 분석하고 관련 정보를 사용해 답변하세요.
-2. 정확성 유지: 정보의 정확성을 확인하고, 불확실한 경우 이를 명시하세요.
-3. 간결한 응답: 질문에 직접 답하고 핵심 내용에 집중하세요.
-4. 추가 정보 제안: 관련된 추가 정보가 있다면 언급하세요.
-5. 윤리성 고려: 객관적이고 중립적인 태도를 유지하세요.
-6. 한계 인정: 답변할 수 없는 경우 솔직히 인정하세요.
-7. 대화 유지: 자연스럽게 대화를 이어가고, 필요시 후속 질문을 제안하세요.
-
-항상 정확하고 유용한 정보를 제공하는 것을 목표로 하세요."""
-
-    # Create a prompt template
-    prompt_template = PromptTemplate(
-        input_variables=["context", "question"],
-        template=instruction + "\n\nContext: {context}\n\nQuestion: {question}\n\nAnswer:"
-    )
-
-    # Create ConversationalRetrievalChain with the new retriever and prompt
+    # Create ConversationalRetrievalChain
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
-        retriever=compression_retriever,
-        return_source_documents=True,
-        combine_docs_chain_kwargs={"prompt": prompt_template}
+        vectorstore.as_retriever(),
+        return_source_documents=True
     )
 
     return qa_chain
@@ -116,69 +82,41 @@ def update_embeddings():
             documents.extend(result)
     vectorstore.add_documents(documents)
 
-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):
-    try:
-        response = rag_chain({"question": query, "chat_history": chat_history})
-        
-        if not response or "answer" not in response:
-            return {
-                "answer": "죄송합니다. 답변을 생성할 수 없었습니다. 질문을 다시 표현해 주시겠습니까?",
-                "sources": []
-            }
-        
-        sources = response.get("sources", [])
-        
-        return {
-            "answer": response["answer"],
-            "sources": sources
-        }
-    except Exception as e:
-        print(f"Error in get_answer: {str(e)}")
-        return {
-            "answer": "답변 생성 중 오류가 발생했습니다. 다시 시도해 주세요.",
-            "sources": []
-        }
+# Generate answer for a query
+def get_answer(qa_chain, 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"]
+    
+    source_docs = response.get("source_documents", [])
+    source_texts = [f"{os.path.basename(doc.metadata['source'])} (Page {doc.metadata['page']})" for doc in source_docs]
+    
+    return {"answer": answer, "sources": source_texts}
 
 # Example usage
 if __name__ == "__main__":
     update_embeddings()  # Update embeddings with new documents
-    question = "RAG 시스템에 대해 설명해주세요."
-    response = get_answer(question, [])
+    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, [])
     print(f"Question: {question}")
     print(f"Answer: {response['answer']}")
-    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}")
+    print(f"Sources: {response['sources']}")