Focus 라우팅과 인용 정합성 — 도구가 문자열을 반환하는 순간 잃는 것

from typing import Annotated, Literal
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

class State(TypedDict):
    messages: Annotated[list, add_messages]
    focus: Literal["web", "academic", "video", "math"]

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.utilities import ArxivAPIWrapper
from langchain_community.utilities.wolfram_alpha import WolframAlphaAPIWrapper
from langchain_core.tools import tool

web_tool = TavilySearchResults(max_results=2)

@tool
def academic_tool(query: str):
    """academic paper search tool"""
    arxiv = ArxivAPIWrapper()
    return arxiv.run(query)

@tool
def math_tool(query: str):
    """math tool"""
    wolfram = WolframAlphaAPIWrapper()
    return wolfram.run(query)

youtube_search_tool = YouTubeSearchTool()

@tool
def video_tool(query: str) -> str:
    """Retriever tool for the transcript of a YouTube video."""
    urls = ast.literal_eval(youtube_search_tool.run(query))   # 1) 영상 URL 검색

    docs = []
    for url in urls:                                          # 2) 자막 로드
        loader = YoutubeLoader.from_youtube_url(
            url, add_video_info=True, language=["en", "ko"])
        scripts = loader.load()
        docs.append(Document(
            page_content=scripts[0].page_content,
            metadata={"source": url,
                      "title": scripts[0].metadata["title"],
                      "author": scripts[0].metadata["author"]}))

    text_splitter = RecursiveCharacterTextSplitter(           # 3) 청킹
        separators=["\n\n", "\n", ".", ",", " ", ""],
        chunk_size=1000, chunk_overlap=0)
    texts = text_splitter.split_documents(docs)

    db = Chroma.from_documents(texts, OpenAIEmbeddings())     # 4) 인메모리 인덱싱
    retrieved_docs = db.as_retriever().invoke(query)          # 5) 유사도 검색

    return "\n\n".join(                                       # 6) 문자열로 평탄화
        f"Title: {d.metadata.get('title')}\n"
        f"Author: {d.metadata.get('author')}\n"
        f"Transcript:\n{d.page_content}"
        for d in retrieved_docs)

tools = {
    "web": [web_tool],
    "academic": [academic_tool],
    "video": [video_tool],
    "math": [math_tool],
}
tool_nodes = {focus: ToolNode(tools[focus]) for focus in tools}

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

def chatbot(state: State):
    llm_with_tools = llm.bind_tools(tools[state["focus"]])   # Focus별 동적 바인딩
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot)
for focus, tool_node in tool_nodes.items():
    graph_builder.add_node(f"{focus}_tools", tool_node)

def focus_condition(state):
    if state["messages"][-1].tool_calls:      # 도구 호출이 있으면 Focus 노드로
        return f"{state['focus']}_tools"
    return END                                # 없으면 종료

graph_builder.add_conditional_edges(
    "chatbot", focus_condition,
    {"web_tools": "web_tools", "academic_tools": "academic_tools",
     "video_tools": "video_tools", "math_tools": "math_tools", END: END})

for focus in tools:
    graph_builder.add_edge(f"{focus}_tools", "chatbot")

graph_builder.set_entry_point("chatbot")
graph = graph_builder.compile()

for chunk in graph.stream(
        {"messages": [{"role": "user", "content": "AI 에이전트에는 어떤 프레임워크가 있어?"}],
         "focus": "video"},
        stream_mode="values"):
    print(chunk["messages"][-1].content)

class State(TypedDict):
    messages: Annotated[list, add_messages]
    focus: Literal["web", "academic", "video", "math"]
    sources: Annotated[list, operator.add]   # 검색 노드가 누적

if "messages" not in st.session_state:
    st.session_state.messages = []

focus_areas = {"Web Search": "web", "Academic Search": "academic",
               "Video Search": "video", "Math": "math"}
selected_focus = [key for area, key in focus_areas.items()
                  if st.checkbox(area, key=f"checkbox_{key}")]

if prompt := st.chat_input("Type your message here"):
    if not selected_focus:
        st.warning("Please select at least one focus area.")
    else:
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").markdown(prompt)

        with st.status("Processing your request...", expanded=True) as status:
            responses = []
            for focus in selected_focus:                 # Focus마다 순차 실행
                status.write(f"Searching {focus.capitalize()}...")
                result = graph.invoke({
                    "messages": st.session_state.messages,
                    "focus": focus})
                responses.append((focus, result["messages"][-1].content))
                status.write(f"Completed {focus.capitalize()} search.")
            status.update(label="Processing complete!", state="complete", expanded=False)

        for focus, response in responses:                 # 태그를 붙여 히스토리에 적재
            tagged = f"[{focus.upper()} SEARCH]\n\n{response}"
            st.session_state.messages.append({"role": "assistant", "content": tagged})
            st.chat_message("assistant").markdown(tagged)

@st.cache_resource
def get_graph():
    return build_graph()          # 도구·LLM·그래프 조립을 여기 안으로