# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import Any
from typing import Optional

from google.genai import types as genai_types
from pydantic import alias_generators
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field


class EvalBaseModel(BaseModel):
  model_config = ConfigDict(
      alias_generator=alias_generators.to_camel,
      populate_by_name=True,
  )


class IntermediateData(EvalBaseModel):
  """Container for intermediate data that an agent would generate as it responds with a final answer."""

  tool_uses: list[genai_types.FunctionCall] = []
  """Tool use trajectory in chronological order."""

  tool_responses: list[genai_types.FunctionResponse] = []
  """Tool response trajectory in chronological order."""

  intermediate_responses: list[tuple[str, list[genai_types.Part]]] = []
  """Intermediate responses generated by sub-agents to convey progress or status
  in a multi-agent system, distinct from the final response.

  This is expressed as a tuple of:
    - Author: Usually the sub-agent name that generated the intermediate
      response.

    - A list of Parts that comprise of the response.
  """


class Invocation(EvalBaseModel):
  """Represents a single invocation."""

  invocation_id: str = ''
  """Unique identifier for the invocation."""

  user_content: genai_types.Content
  """Content provided by the user in this invocation."""

  final_response: Optional[genai_types.Content] = None
  """Final response from the agent."""

  intermediate_data: Optional[IntermediateData] = None
  """Intermediate steps generated as a part of Agent execution.

  For a multi-agent system, it is also helpful to inspect the route that
  the agent took to generate final response.
  """

  creation_timestamp: float = 0.0
  """Timestamp for the current invocation, primarily intended for debugging purposes."""


class SessionInput(EvalBaseModel):
  """Values that help initialize a Session."""

  app_name: str
  """The name of the app."""

  user_id: str
  """The user id."""

  state: dict[str, Any] = Field(default_factory=dict)
  """The state of the session."""


class EvalCase(EvalBaseModel):
  """An eval case."""

  eval_id: str
  """Unique identifier for the evaluation case."""

  conversation: list[Invocation]
  """A conversation between the user and the Agent. The conversation can have any number of invocations."""

  session_input: Optional[SessionInput] = None
  """Session input that will be passed on to the Agent during eval.
     It is common for Agents state to be initialized to some initial/default value,
     for example, your agent may need to know today's date.
  """

  creation_timestamp: float = 0.0
  """The time at which this eval case was created."""
