AI GDE의 시선으로 바라본 Google I/O Connect China 2025 후기

import os
# 이 가이드는 JAX 백엔드에서만 실행할 수 있습니다.
os.environ["KERAS_BACKEND"] = "jax"
import jax
import keras
import numpy as np

class CustomModel(keras.Model):
    def compute_loss_and_updates(
        self,
        trainable_variables,
        non_trainable_variables,
        x,
        y,
        training=False,
    ):
        y_pred, non_trainable_variables = self.stateless_call(
            trainable_variables,
            non_trainable_variables,
            x,
            training=training,
        )
        loss = self.compute_loss(x, y, y_pred)
        return loss, (y_pred, non_trainable_variables)

    def train_step(self, state, data):
        (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        ) = state
        x, y = data

        # 그라디언트 함수를 가져옵니다.
        grad_fn = jax.value_and_grad(self.compute_loss_and_updates, has_aux=True)

        # 그라디언트를 계산합니다.
        (loss, (y_pred, non_trainable_variables)), grads = grad_fn(
            trainable_variables,
            non_trainable_variables,
            x,
            y,
            training=True,
        )

        # 훈련 가능한 변수와 옵티마이저 변수를 업데이트합니다.
        (
            trainable_variables,
            optimizer_variables,
        ) = self.optimizer.stateless_apply(
            optimizer_variables, grads, trainable_variables
        )

        # 메트릭을 업데이트합니다.
        new_metrics_vars = []
        logs = {}
        for metric in self.metrics:
            this_metric_vars = metrics_variables[
                len(new_metrics_vars) : len(new_metrics_vars) + len(metric.variables)
            ]
            if metric.name == "loss":
                this_metric_vars = metric.stateless_update_state(this_metric_vars, loss)
            else:
                this_metric_vars = metric.stateless_update_state(
                    this_metric_vars, y, y_pred
                )
            logs[metric.name] = metric.stateless_result(this_metric_vars)
            new_metrics_vars += this_metric_vars

        # 메트릭 로그와 업데이트된 상태 변수를 반환합니다.
        state = (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            new_metrics_vars,
        )
        return logs, state

# CustomModel의 인스턴스를 생성하고 컴파일합니다.
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])

# 평소처럼 `fit`을 사용합니다.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=3)

# 행사 스케줄 정보를 담은 가상 데이터베이스
IO_CONNECT_SCHEDULE = {
    "JAX 실제 활용 사례": {"time": "14:25", "track": "AI", "desc": "Gemini와 같은 모델을 구축하는 데 사용된 JAX와 Flax 라이브러리 소개"},
    "Keras 3에서 model.fit(...)을 사용하여 JAX 모델 훈련하기": {"time": "15:20", "track": "AI", "desc": "친숙한 Keras API를 사용하여 JAX로 모델을 훈련하는 방법"},
    "필요한 것은 에이전트뿐": {"time": "10:30", "track": "Cloud", "desc": "ADK 및 Agent Engine으로 다중 에이전트 시스템을 구축하는 방법"}
}

# 도구 함수 정의: 세션 제목으로 정보를 찾는 함수
def find_session_details(session_title: str) -> str:
  """Google I/O Connect China 2025 세션 제목으로 시간, 트랙, 설명을 검색합니다."""
  if session_title in IO_CONNECT_SCHEDULE:
    return str(IO_CONNECT_SCHEDULE[session_title])
  return f"'{session_title}' 세션을 찾을 수 없습니다. 제목을 다시 확인해주세요."

# '행사 안내 에이전트' 정의
io_schedule_agent = LlmAgent(
    model="gemini-2.0-flash",
    name="io_schedule_agent",
    description="Google I/O Connect China 2025의 세션 정보를 안내합니다.",
    instruction="""당신은 I/O Connect 행사 안내원입니다.
사용자가 세션 정보를 물어보면, 세션의 전체 제목을 파악하여 `find_session_details` 도구를 사용하세요.
그리고 찾은 정보를 바탕으로 사용자에게 친절하게 답변해주세요.
""",
    tools=[find_session_details] # 함수를 직접 도구로 제공
)

# 1. 각 작업을 전담할 하위 에이전트(전문가 팀) 정의

# '세션 검색 담당' 에이전트 (병렬 처리)
ai_track_search_agent = LlmAgent(
    name="AITrackSearcher",
    model="gemini-2.0-flash",
    instruction="""AI 트랙 세션 목록에서 'JAX' 또는 'TPU' 키워드가 포함된 세션을 모두 찾아주세요.
    검색 대상: {session_list}""",
    output_key="ai_session_results"
)

cloud_track_search_agent = LlmAgent(
    name="CloudTrackSearcher",
    model="gemini-2.0-flash",
    instruction="""Cloud 트랙 세션 목록에서 'ADK' 또는 'Agent' 키워드가 포함된 세션을 모두 찾아주세요.
    검색 대상: {session_list}""",
    output_key="cloud_session_results"
)

# 검색 작업을 병렬로 실행할 '검색팀' 에이전트
parallel_search_agent = ParallelAgent(
    name="ParallelSearchTeam",
    sub_agents=[ai_track_search_agent, cloud_track_search_agent],
)


# '추천 목록 생성 담당' 에이전트
recommendation_agent = LlmAgent(
    name="RecommendationGenerator",
    model="gemini-2.5-pro-preview-03-25", # 추천 요약은 더 강력한 모델 사용
    instruction="""당신은 I/O Connect의 맞춤형 가이드입니다.
    아래 두 검색 결과를 바탕으로, 사용자에게 추천할 세션 목록을 하나의 깔끔한 목록으로 종합해주세요.
    각 세션의 시간과 핵심 내용을 포함하여 흥미롭게 요약해야 합니다.

    AI 트랙 검색 결과:
    {ai_session_results}

    Cloud 트랙 검색 결과:
    {cloud_session_results}
    """,
    output_key="recommended_list",
)

# 2. SequentialAgent를 사용하여 '세션 추천 파이프라인' 생성
session_recommendation_pipeline = SequentialAgent(
    name="SessionRecommendationPipeline",
    sub_agents=[parallel_search_agent, recommendation_agent],
    description="관련 세션을 병렬로 검색한 후, 그 결과를 종합하여 최종 추천 목록을 생성합니다.",
)