map.go

Overview

This file implements the parentmap package that provides a specialized mapping structure to represent and navigate hierarchical relationships among agents. Specifically, it constructs and manages a "parent map" for an agent tree, where each agent (except the root) has a single parent. This enables efficient retrieval of an agent's parent and traversal to the root agent. The package also supports storing and retrieving the parent map within a context.Context for convenient propagation across API boundaries.

The core functionality centers on verifying tree integrity (unique agent names, single-parent constraints) and facilitating upward navigation in the agent hierarchy. The file interacts with the agent.Agent interface from the google.golang.org/adk/agent package, which defines the agent abstraction including agent names and sub-agent relationships.

Types and Functions

Type: Map

type Map map[string]agent.Agent

• Description:
  Map is a custom type alias for a map from agent names (string) to their respective parent agent.Agent instances. It encapsulates the parent-child relationships within an agent tree, allowing quick lookup of a parent agent by its child's name.

• Usage:
  Use Map to query parent agents and to find the root agent of any subtree.

Function: New

func New(root agent.Agent) (Map, error)

• Purpose:
  Constructs a Map representing the parent relationships of an entire agent tree rooted at root.

• Parameters:

  • root: The root agent.Agent of the tree. This agent is considered the top-most node and has no parent.

• Returns:

  • Map: The populated parent map where keys are agent names and values are their parent agents.

  • error: Returns an error if the tree violates constraints:

    • An agent has more than one parent.

    • Duplicate agent names exist in the tree.

    • The root agent's name appears elsewhere in the tree.

• Behavior and Implementation Details:

  • Initializes an empty map res for agent name to parent mappings.

  • Uses an auxiliary map pointerMap to track which parent each agent has, ensuring no agent has multiple parents.

  • Recursively traverses the agent tree using an inner function f:

    • For each sub-agent of the current agent:

      • Validates uniqueness of the sub-agent's name.

      • Validates that the sub-agent has not been assigned a different parent before.

      • Records the current agent as the sub-agent's parent in res.

      • Recurses into the sub-agent.

  • Returns the constructed parent map and any error detected during traversal.

• Example:

  root := someAgent
  parentMap, err := New(root)
  if err != nil {
      // handle error: tree is invalid (duplicate names or multiple parents)
  }
  

Method: RootAgent

func (m Map) RootAgent(cur agent.Agent) agent.Agent

• Purpose:
  Returns the root agent of the tree to which the given agent cur belongs by following parent links upward.

• Parameters:

  • cur: The starting agent.Agent to find the root for.

• Returns:

  • The root agent.Agent if cur is non-nil.

  • nil if the input cur is nil.

• Behavior:

  • Iteratively follows the parent links in the Map until a node with no parent is found.

  • This node is the root agent.

• Example:

  root := parentMap.RootAgent(someAgent)
  

Function: ToContext

func ToContext(ctx context.Context, parents Map) context.Context

• Purpose:
  Embeds the parentmap.Map into a context.Context for propagation.

• Parameters:

  • ctx: The original context.

  • parents: The parent map to store in the context.

• Returns:

  • A new context.Context containing the parent map under an unexported key.

• Usage:
  Used to pass the parent map through call chains without explicit parameters.

Function: FromContext

func FromContext(ctx context.Context) Map

• Purpose:
  Extracts a parentmap.Map from a given context.Context.

• Parameters:

  • ctx: The context from which to retrieve the parent map.

• Returns:

  • The Map if present in the context.

  • nil if no parent map is found.

• Usage:
  Allows retrieval of the agent parent map previously stored via ToContext.

Internal Type: ctxKey and Constant mapCtxKey

type ctxKey int

const mapCtxKey ctxKey = 0

• Used as a private key type to avoid collisions when storing values in context.Context.

Interactions with Other Components

• agent.Agent Interface:
  The package depends on the agent.Agent interface, which exposes methods like Name() and SubAgents(). This interface defines the hierarchical structure of agents that the parent map represents.

• Context Propagation:
  By embedding the parent map in context.Context, this file enables other parts of the system that handle agents and their lifecycles (such as those covered under Agent Invocation Context) to access parent relationships seamlessly.

• Agent Lifecycle and Tree Integrity:
  The constraints enforced in New ensure that the agent tree conforms to a strict hierarchy, supporting reliable traversal and manipulation in workflows as managed by other components like Agent Workflow Management.

Implementation Details and Algorithms

• Tree Traversal and Validation:
  The New function performs a recursive depth-first traversal starting from the root agent. It validates uniqueness of agent names and prevents agents from having multiple parents by maintaining a map of visited agents and their parents.

• Parent Lookup Optimization:
  The parent map (Map) indexes parent agents by their child agent names, enabling O(1) parent retrieval.

• Root Agent Retrieval:
  The RootAgent method uses iterative traversal upward via the map until an agent without a parent is found, avoiding recursion and ensuring efficient root discovery.

• Context Key Safety:
  Defines a private type ctxKey and unexported constant key to protect against key collisions in context.Context usage.

Usage Examples

// Constructing a parent map for an agent tree.
rootAgent := getRootAgent()
parentMap, err := parentmap.New(rootAgent)
if err != nil {
    // handle invalid tree: duplicate names or multiple parents
}

// Find the root agent of any agent.
someAgent := getSomeAgent()
root := parentMap.RootAgent(someAgent)

// Passing parent map through context.
ctx := context.Background()
ctxWithMap := parentmap.ToContext(ctx, parentMap)

// Retrieving parent map from context.
retrievedMap := parentmap.FromContext(ctxWithMap)

Diagram: Structure and Main Function Relationships

classDiagram
class Map {
+RootAgent()
}
class New {
<<function>>
}
class ToContext {
<<function>>
}
class FromContext {
<<function>>
}
Map --> New : "created by"
ToContext --> Map : "stores in context"
FromContext --> Map : "retrieves from context"

This diagram shows the main entities and their relations: New creates a Map, which provides the RootAgent method, and ToContext/FromContext allow embedding and extracting the map from a context.Context.