pip install langgraph langchain langchain-openai langchain-community chromadb openai python-dotenv

[
  {
    "id": "company_overview",
    "topic": "회사 소개",
    "content": "테크노바는 2020년에 설립된 AI 기반 솔루션을 제공하는 스타트업입니다. 서울 강남구에 본사를 두고 있으며, 현재 직원 수는 약 50명입니다."
  },
  {
    "id": "main_product",
    "topic": "주력 제품",
    "content": "우리의 주력 제품은 'SmartAssist'라는 AI 챗봇 플랫폼입니다. 고객 서비스 자동화와 업무 효율성 향상을 위한 솔루션을 제공합니다."
  },
  {
    "id": "ceo_info",
    "topic": "대표이사",
    "content": "김철수 대표는 카이스트에서 인공지능을 전공했으며, 이전에 네이버에서 AI 연구원으로 근무했습니다. 10년 이상의 AI 개발 경험을 보유하고 있습니다."
  },
  {
    "id": "office_location",
    "topic": "사무실 위치",
    "content": "본사는 서울특별시 강남구 테헤란로 123에 위치해 있습니다. 지하철 2호선 강남역에서 도보 5분 거리이며, 15층에 사무실이 있습니다."
  },
  {
    "id": "work_hours",
    "topic": "근무 시간",
    "content": "근무 시간은 오전 9시부터 오후 6시까지이며, 코어타임은 오전 10시부터 오후 4시까지입니다. 주 5일 근무제를 시행하고 있습니다."
  },
  {
    "id": "benefits",
    "topic": "복지 제도",
    "content": "직원들에게 최신형 맥북 프로와 듀얼 모니터를 제공하며, 점심 식대 지원, 간식 무제한 제공, 연 1회 워크샵, 도서 구입비 지원 등의 복지를 제공합니다."
  },
  {
    "id": "tech_stack",
    "topic": "기술 스택",
    "content": "백엔드는 Python FastAPI와 Django를 사용하고, 프론트엔드는 React와 TypeScript를 사용합니다. AI 모델 학습에는 PyTorch를 주로 사용합니다."
  },
  {
    "id": "team_structure",
    "topic": "팀 구조",
    "content": "회사는 크게 개발팀, 디자인팀, 영업팀, 마케팅팀으로 구성되어 있습니다. 개발팀이 가장 크며 약 30명의 엔지니어가 근무하고 있습니다."
  },
  {
    "id": "funding",
    "topic": "투자 유치",
    "content": "2023년에 시리즈 A 라운드에서 50억원의 투자를 유치했습니다. 주요 투자사는 카카오벤처스와 네이버 D2SF입니다."
  },
  {
    "id": "clients",
    "topic": "주요 고객사",
    "content": "현재 삼성전자, LG전자, 현대자동차 등 대기업을 포함해 총 100여개의 고객사를 보유하고 있습니다."
  }
]

import os
import json
from dotenv import load_dotenv
from typing import TypedDict, Literal

# LangChain 관련
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document

# LangGraph 관련
from langgraph.graph import StateGraph, END

# 환경 변수 로드
load_dotenv()

print("✅ 환경 설정 완료!")

질문 → 검색 → 답변 생성 → 끝

질문 → 질문 평가 → ❌ → 질문 재작성 → 다시 평가
         ↓ ✅
       검색 → 검색 평가 → ❌ → 쿼리 재작성 → 다시 검색
         ↓ ✅
     답변 생성 → 답변 평가 → ❌ → 문제 진단 → 해당 단계 복귀
         ↓ ✅
       최종 답변

# JSON 파일에서 데이터 로드
with open('company_info.json', 'r', encoding='utf-8') as f:
    company_data = json.load(f)

# Document 객체로 변환
documents = []
for item in company_data:
    doc = Document(
        page_content=f"주제: {item['topic']}\n내용: {item['content']}",
        metadata={"id": item["id"], "topic": item["topic"]}
    )
    documents.append(doc)

# 임베딩 및 벡터 스토어
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings,
    collection_name="company_info"
)

retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})

# LLM 초기화
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)

print(f"✅ 준비 완료! 문서 수: {vectorstore._collection.count()}")

class SelfCorrectingRAGState(TypedDict):
    """Self-Correcting RAG 워크플로우 상태"""
    # 입력
    original_question: str  # 원본 질문
    
    # 질문 처리
    current_question: str  # 현재 처리 중인 질문 (재작성될 수 있음)
    question_quality: str  # "good" or "bad"
    question_rewrite_count: int  # 질문 재작성 횟수
    
    # 검색 처리
    search_query: str  # 검색 쿼리
    retrieved_docs: list  # 검색된 문서들
    retrieval_quality: str  # "relevant" or "irrelevant"
    search_retry_count: int  # 검색 재시도 횟수
    
    # 답변 처리
    answer: str  # 생성된 답변
    answer_quality: str  # "good" or "bad"
    problem_diagnosis: str  # 문제 진단 결과
    
    # 제어
    max_retries: int  # 최대 재시도 횟수

def evaluate_question(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """질문이 명확하고 구체적인지 평가합니다."""
    question = state["current_question"]
    
    evaluation_prompt = f"""
다음 질문이 명확하고 구체적인지 평가해주세요.

질문: "{question}"

평가 기준:
1. 질문이 무엇을 묻는지 명확한가?
2. 문맥 없이도 이해 가능한가?
3. 회사 정보에 대한 질문인가?

"good" 또는 "bad" 중 하나로만 답변하세요.
""".strip()
    
    response = llm.invoke(evaluation_prompt)
    quality = response.content.strip().lower()
    
    if "good" in quality:
        quality = "good"
    else:
        quality = "bad"
    
    print(f"🔍 [질문 평가] {question}")
    print(f"   결과: {quality}")
    
    return {
        **state,
        "question_quality": quality
    }

def rewrite_question(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """질문을 더 명확하고 구체적으로 재작성합니다."""
    question = state["current_question"]
    
    rewrite_prompt = f"""다음 질문을 더 명확하고 구체적으로 재작성해주세요.

원본 질문: "{question}"

재작성 시 고려사항:
1. 회사 정보를 묻는 것이 명확하게 드러나도록
2. 애매한 표현 제거
3. 구체적인 정보 요청으로 변환

재작성된 질문만 답변하세요.
"""
    
    response = llm.invoke(rewrite_prompt)
    rewritten = response.content.strip()
    
    print(f"✏️ [질문 재작성]")
    print(f"   원본: {question}")
    print(f"   재작성: {rewritten}")
    
    return {
        **state,
        "current_question": rewritten,
        "question_rewrite_count": state.get("question_rewrite_count", 0) + 1
    }

def retrieve_documents(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """벡터 DB에서 관련 문서를 검색합니다."""
    query = state.get("search_query") or state["current_question"]
    
    docs = retriever.invoke(query)
    
    print(f"📚 [문서 검색] 쿼리: {query}")
    print(f"   검색된 문서 수: {len(docs)}")
    for i, doc in enumerate(docs, 1):
        print(f"   {i}. {doc.metadata.get('topic', 'N/A')}")
    
    return {
        **state,
        "retrieved_docs": docs,
        "search_query": query
    }

def evaluate_retrieval(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """검색된 문서가 질문과 관련있는지 평가합니다."""
    question = state["current_question"]
    docs = state["retrieved_docs"]
    
    docs_content = "\n\n".join([f"문서 {i+1}: {doc.page_content[:100]}..." 
                                 for i, doc in enumerate(docs)])
    
    evaluation_prompt = f"""다음 질문에 대해 검색된 문서들이 관련성이 있는지 평가해주세요.

질문: "{question}"

검색된 문서들:
{docs_content}

평가 기준:
1. 문서에 질문에 답할 수 있는 정보가 포함되어 있는가?
2. 문서와 질문의 주제가 일치하는가?

"relevant" 또는 "irrelevant" 중 하나로만 답변하세요.
"""
    
    response = llm.invoke(evaluation_prompt)
    quality = response.content.strip().lower()
    
    if "relevant" in quality:
        quality = "relevant"
    else:
        quality = "irrelevant"
    
    print(f"✅ [검색 평가] 결과: {quality}")
    
    return {
        **state,
        "retrieval_quality": quality
    }

def rewrite_search_query(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """더 나은 검색 결과를 위해 쿼리를 재작성합니다."""
    question = state["current_question"]
    previous_query = state.get("search_query", question)
    
    rewrite_prompt = f"""검색 결과가 부적절했습니다. 더 나은 검색 결과를 위해 쿼리를 재작성해주세요.

원본 질문: "{question}"
이전 검색 쿼리: "{previous_query}"

재작성 시 고려사항:
1. 핵심 키워드 강조
2. 동의어나 관련 용어 포함
3. 더 구체적인 표현 사용

재작성된 검색 쿼리만 답변하세요.
"""
    
    response = llm.invoke(rewrite_prompt)
    new_query = response.content.strip()
    
    print(f"🔄 [쿼리 재작성]")
    print(f"   이전: {previous_query}")
    print(f"   재작성: {new_query}")
    
    return {
        **state,
        "search_query": new_query,
        "search_retry_count": state.get("search_retry_count", 0) + 1
    }

def generate_answer(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """검색된 문서를 바탕으로 답변을 생성합니다."""
    question = state["current_question"]
    docs = state["retrieved_docs"]
    
    context = "\n\n".join([doc.page_content for doc in docs])
    
    answer_prompt = f"""당신은 회사 정보를 안내하는 전문 AI 어시스턴트입니다.
아래 컨텍스트를 바탕으로 질문에 정확하고 친절하게 답변해주세요.

컨텍스트:
{context}

질문: {question}

답변:"""
    
    response = llm.invoke(answer_prompt)
    answer = response.content
    
    print(f"💬 [답변 생성] 완료")
    
    return {
        **state,
        "answer": answer
    }

def evaluate_answer(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """생성된 답변의 품질을 평가합니다."""
    question = state["current_question"]
    answer = state["answer"]
    
    evaluation_prompt = f"""다음 답변의 품질을 평가해주세요.

질문: "{question}"
답변: "{answer}"

평가 기준:
1. 질문에 직접적으로 답변하는가?
2. 컨텍스트의 정보를 정확히 사용했는가?
3. 답변이 구체적이고 유용한가?
4. 환각(hallucination)이 없는가?

"good" 또는 "bad" 중 하나로만 답변하세요.
"""
    
    response = llm.invoke(evaluation_prompt)
    quality = response.content.strip().lower()
    
    if "good" in quality:
        quality = "good"
    else:
        quality = "bad"
    
    print(f"⭐ [답변 평가] 결과: {quality}")
    
    return {
        **state,
        "answer_quality": quality
    }

def diagnose_problem(state: SelfCorrectingRAGState) -> SelfCorrectingRAGState:
    """답변에 문제가 있을 때, 어디에서 문제가 발생했는지 진단합니다."""
    question = state["current_question"]
    docs = state["retrieved_docs"]
    answer = state["answer"]
    
    docs_summary = "\n".join([f"- {doc.metadata.get('topic', 'N/A')}" for doc in docs])
    
    diagnosis_prompt = f"""답변에 문제가 있습니다. 어디에서 문제가 발생했는지 진단해주세요.

질문: "{question}"
검색된 문서들:
{docs_summary}
생성된 답변: "{answer}"

다음 중 하나를 선택하세요:
1. "question_issue" - 질문 자체에 문제가 있음 (애매하거나 불명확)
2. "retrieval_issue" - 검색된 문서가 부적절함
3. "generation_issue" - 답변 생성 과정에서 문제 발생

진단 결과만 답변하세요 (question_issue, retrieval_issue, generation_issue 중 하나).
"""
    
    response = llm.invoke(diagnosis_prompt)
    diagnosis = response.content.strip().lower()
    
    if "question" in diagnosis:
        diagnosis = "question_issue"
    elif "retrieval" in diagnosis:
        diagnosis = "retrieval_issue"
    else:
        diagnosis = "generation_issue"
    
    print(f"🔧 [문제 진단] {diagnosis}")
    
    return {
        **state,
        "problem_diagnosis": diagnosis
    }

def route_after_question_eval(state: SelfCorrectingRAGState) -> str:
    """질문 평가 후 라우팅"""
    if state["question_quality"] == "good":
        return "retrieve"
    elif state.get("question_rewrite_count", 0) >= state.get("max_retries", 2):
        return "retrieve"  # 최대 재시도 도달, 그냥 진행
    else:
        return "rewrite_question"


def route_after_retrieval_eval(state: SelfCorrectingRAGState) -> str:
    """검색 평가 후 라우팅"""
    if state["retrieval_quality"] == "relevant":
        return "generate"
    elif state.get("search_retry_count", 0) >= state.get("max_retries", 2):
        return "generate"  # 최대 재시도 도달, 그냥 진행
    else:
        return "rewrite_query"


def route_after_answer_eval(state: SelfCorrectingRAGState) -> str:
    """답변 평가 후 라우팅"""
    if state["answer_quality"] == "good":
        return "end"
    else:
        return "diagnose"


def route_after_diagnosis(state: SelfCorrectingRAGState) -> str:
    """문제 진단 후 라우팅"""
    diagnosis = state["problem_diagnosis"]
    
    if diagnosis == "question_issue":
        return "rewrite_question"
    elif diagnosis == "retrieval_issue":
        return "rewrite_query"
    else:  # generation_issue
        return "generate"  # 답변만 다시 생성

from IPython.display import Image, display

# 그래프 생성
workflow = StateGraph(SelfCorrectingRAGState)

# 노드 추가
workflow.add_node("evaluate_question", evaluate_question)
workflow.add_node("rewrite_question", rewrite_question)
workflow.add_node("retrieve", retrieve_documents)
workflow.add_node("evaluate_retrieval", evaluate_retrieval)
workflow.add_node("rewrite_query", rewrite_search_query)
workflow.add_node("generate", generate_answer)
workflow.add_node("evaluate_answer", evaluate_answer)
workflow.add_node("diagnose", diagnose_problem)

# 시작점 설정
workflow.set_entry_point("evaluate_question")

# 엣지 연결
# 1. 질문 평가 → 조건부 분기
workflow.add_conditional_edges(
    "evaluate_question",
    route_after_question_eval,
    {
        "retrieve": "retrieve",
        "rewrite_question": "rewrite_question"
    }
)

# 2. 질문 재작성 → 다시 평가
workflow.add_edge("rewrite_question", "evaluate_question")

# 3. 검색 → 검색 평가
workflow.add_edge("retrieve", "evaluate_retrieval")

# 4. 검색 평가 → 조건부 분기
workflow.add_conditional_edges(
    "evaluate_retrieval",
    route_after_retrieval_eval,
    {
        "generate": "generate",
        "rewrite_query": "rewrite_query"
    }
)

# 5. 쿼리 재작성 → 다시 검색
workflow.add_edge("rewrite_query", "retrieve")

# 6. 답변 생성 → 답변 평가
workflow.add_edge("generate", "evaluate_answer")

# 7. 답변 평가 → 조건부 분기
workflow.add_conditional_edges(
    "evaluate_answer",
    route_after_answer_eval,
    {
        "end": END,
        "diagnose": "diagnose"
    }
)

# 8. 문제 진단 → 조건부 분기 (해당 단계로 복귀)
workflow.add_conditional_edges(
    "diagnose",
    route_after_diagnosis,
    {
        "rewrite_question": "rewrite_question",
        "rewrite_query": "rewrite_query",
        "generate": "generate"
    }
)

# 컴파일
app = workflow.compile()

print("✅ Self-Correcting RAG 그래프 구성 완료!")

display(Image(app.get_graph().draw_mermaid_png()))

[시작]
  ↓
[질문 평가] → bad? → [질문 재작성] → [질문 평가]
  ↓ good
[문서 검색]
  ↓
[검색 평가] → irrelevant? → [쿼리 재작성] → [문서 검색]
  ↓ relevant
[답변 생성]
  ↓
[답변 평가] → bad? → [문제 진단] → question? → [질문 재작성]
  ↓ good              ↓              retrieval? → [쿼리 재작성]
[종료]                                generation? → [답변 생성]

def ask_with_self_correction(question: str, max_retries: int = 2):
    """Self-Correcting RAG로 질문에 답변"""
    print("\n" + "="*80)
    print(f"💭 원본 질문: {question}")
    print("="*80 + "\n")
    
    initial_state = {
        "original_question": question,
        "current_question": question,
        "question_quality": "",
        "question_rewrite_count": 0,
        "search_query": "",
        "retrieved_docs": [],
        "retrieval_quality": "",
        "search_retry_count": 0,
        "answer": "",
        "answer_quality": "",
        "problem_diagnosis": "",
        "max_retries": max_retries
    }
    
    # 실행
    result = app.invoke(initial_state)
    
    print("\n" + "="*80)
    print(f"✨ 최종 답변:")
    print(f"{result['answer']}")
    print("="*80)
    print(f"\n📊 통계:")
    print(f"  - 질문 재작성 횟수: {result['question_rewrite_count']}")
    print(f"  - 검색 재시도 횟수: {result['search_retry_count']}")
    print(f"  - 최종 질문: {result['current_question']}")
    
    return result

ask_with_self_correction("테크노바의 근무시간은 어떻게 돼?")