코딩 에이전트의 동적 RAG 라우팅 — 검색할 질문과 그냥 답할 질문을 가른다

class AgentState(TypedDict):
    """The state of our agent."""
    question: str          # 원 질문 — 모든 노드가 읽는 유일한 입력
    certainty_score: int   # LLM의 자기 확신도 1~100
    search_results: list   # Tavily 웹 검색 결과 원본
    web_score: str         # 'yes' | 'no' — 웹 결과로 해결 가능한가
    repo_name: str         # 선택된 깃헙 레포 'owner/repo'
    generation: str        # 최종 답변 (세 경로가 모두 여기에 씀)
    github_chunks: list    # 검색된 코드 청크

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def check_certainty(state: AgentState) -> AgentState:
    """Evaluate certainty score for the query."""
    question = state["question"]

    class CertaintyScoreResponse(BaseModel):
        score: int = Field(description="Certainty score from 1 to 100. Higher is better.")

    certainty_scorer = llm.with_structured_output(CertaintyScoreResponse)
    score_response = certainty_scorer.invoke(question)

    return {"certainty_score": score_response.score}


def route_based_on_certainty(state: AgentState) -> Literal["web_search", "direct_response"]:
    """Route to appropriate node based on certainty score."""
    score = state["certainty_score"]
    if score != 100:
        return "web_search"
    else:
        return "direct_response"

def web_search(state: AgentState) -> AgentState:
    """Perform web search and evaluate results"""
    question = state["question"]

    search_tool = TavilySearchResults(max_results=3)
    search_results = search_tool.invoke(question)

    class answer_availability(BaseModel):
        """Binary score for answer availability."""
        binary_score: str = Field(description="""
            If web search result can solve the user's ask, answer 'yes'.
            If user's ask is related with github or search_results are insufficient, answer 'no'""")

    evaluator = llm.with_structured_output(answer_availability)
    eval_prompt = ChatPromptTemplate.from_messages([
        ("system", """Evaluate if these search results can answer the user's question
            with a simple yes/no. If user ask github related info, then it is not
            sufficient with web search so you should answer with no."""),
        ("user", "Question: {question}\nSearch Results: {results}\nCan these results answer adequately?")
    ])
    evaluation = evaluator.invoke(
        eval_prompt.format(
            question=question,
            results="\n".join(f"- {r['content']}" for r in search_results)
        )
    )
    return {"search_results": search_results, "web_score": evaluation.binary_score}

def web_generate(state: AgentState):
    question, web_results = state["question"], state["search_results"]

    def format_web_results(results):   # 출처 URL을 본문과 함께 주입
        return "\n".join(
            f"Source {i}:\nURL: {r['url']}\nContent: {r['content']}\n"
            for i, r in enumerate(results, 1)
        )

    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a helpful assistant that generates comprehensive answers
        based on web search results. Make sure to synthesize information from multiple
        sources when possible. If the search results don't contain enough information
        to fully answer the question, acknowledge this limitation."""),
        ("user", "Question: {question}\n\nSearch Results:\n{web_results}\n\nAnswer in Korean")
    ])
    chain = ({"question": lambda x: x["question"],
              "web_results": lambda x: format_web_results(x["web_results"])}
             | prompt | llm | StrOutputParser())
    return {"generation": chain.invoke({"question": question, "web_results": web_results})}

def github_generate(state: AgentState) -> AgentState:
    class GitHubRepo(BaseModel):
        repo_name: str = Field(description="Full repository name in format 'owner/repo'")

    question = state["question"]
    # 1. 깃헙으로 한정한 재검색  2. 후보 중 단 하나의 레포를 구조화 출력으로 확정
    search_results = TavilySearchResults(max_results=5).invoke(
        f"github repository {question} site:github.com")
    repo_extractor = llm.with_structured_output(GitHubRepo)
    repo_name = repo_extractor.invoke(
        eval_prompt.format(question=question, results=...)).repo_name

    # 3. 레포를 벡터화하고 RAG
    vectorstore = git_vector_embedding(repo_name)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
    result = rag_chain.invoke(question)
    return {"repo_name": repo_name, "generation": result, "github_chunks": retrieved_chunks}

def git_vector_embedding(repo_name):
    client = chromadb.Client(Settings(is_persistent=True, persist_directory="./chroma_db"))
    collection_name = repo_name.split("/")[1]

    if collection_name in [c.name for c in client.list_collections()]:
        vectorstore = Chroma(client=client, collection_name=collection_name,
                             embedding_function=OpenAIEmbeddings())   # 캐시 히트
    else:
        try:
            git_docs = git_loader(repo_name, "master")   # 기본 브랜치 이름 폴백
        except:
            git_docs = git_loader(repo_name, "main")
        doc_splits = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
            chunk_size=500, chunk_overlap=50
        ).split_documents(git_docs)
        vectorstore = Chroma.from_documents(documents=doc_splits,
                                            collection_name=collection_name,
                                            embedding=OpenAIEmbeddings(), client=client)
    return vectorstore

workflow = StateGraph(AgentState)

workflow.add_node("check_certainty", check_certainty)
workflow.add_node("direct_response", direct_response)
workflow.add_node("web_search", web_search)
workflow.add_node("web_generate", web_generate)
workflow.add_node("github_generate", github_generate)

workflow.add_edge(START, "check_certainty")

workflow.add_conditional_edges(
    "check_certainty", route_based_on_certainty,
    {"web_search": "web_search", "direct_response": "direct_response"},
)
workflow.add_conditional_edges(
    "web_search", route_after_search,
    {"web_generate": "web_generate", "github_generate": "github_generate"},
)

workflow.add_edge("direct_response", END)
workflow.add_edge("web_generate", END)
workflow.add_edge("github_generate", END)

graph = workflow.compile()

for step in graph.stream(initial_state, config={"recursion_limit": 100}):
    for node_name, state in step.items():
        if 'certainty_score' in state:   # "제가 스스로 답할 수 있는지 고민중이에요..."
            st.write(f"🤔 LLM의 확신 정도: {state['certainty_score']}/100")
        if 'web_score' in state:         # "웹 검색만으로는 어려워요. 깃헙을 찾아볼게요!"
            ...
        if 'repo_name' in state:
            st.write(f"📚 참고한 GitHub 저장소: {state['repo_name']}")
        if 'generation' in state:
            st.markdown(state['generation'])