code: """Base agent type."""
  
  import uuid
  from abc import abstractmethod
  from typing import Any, Dict, List, Optional
  
  from llama_index.legacy.bridge.pydantic import BaseModel, Field
  from llama_index.legacy.callbacks import CallbackManager, trace_method
  from llama_index.legacy.chat_engine.types import (
      BaseChatEngine,
      StreamingAgentChatResponse,
  )
  from llama_index.legacy.core.base_query_engine import BaseQueryEngine
  from llama_index.legacy.core.llms.types import ChatMessage
  from llama_index.legacy.core.response.schema import RESPONSE_TYPE, Response
  from llama_index.legacy.memory.types import BaseMemory
  from llama_index.legacy.prompts.mixin import (
      PromptDictType,
      PromptMixin,
      PromptMixinType,
  )
  from llama_index.legacy.schema import QueryBundle
  
  
  class BaseAgent(BaseChatEngine, BaseQueryEngine):
      """Base Agent."""
  
      def _get_prompts(self) -> PromptDictType:
          """Get prompts."""
          # TODO: the ReAct agent does not explicitly specify prompts, would need a
          # refactor to expose those prompts
          return {}
  
      def _get_prompt_modules(self) -> PromptMixinType:
          """Get prompt modules."""
          return {}
  
      def _update_prompts(self, prompts: PromptDictType) -> None:
          """Update prompts."""
  
      # ===== Query Engine Interface =====
      @trace_method("query")
      def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
          agent_response = self.chat(
              query_bundle.query_str,
              chat_history=[],
          )
          return Response(
              response=str(agent_response), source_nodes=agent_response.source_nodes
          )
  
      @trace_method("query")
      async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
          agent_response = await self.achat(
              query_bundle.query_str,
              chat_history=[],
          )
          return Response(
              response=str(agent_response), source_nodes=agent_response.source_nodes
          )
  
      def stream_chat(
          self, message: str, chat_history: Optional[List[ChatMessage]] = None
      ) -> StreamingAgentChatResponse:
          raise NotImplementedError("stream_chat not implemented")
  
      async def astream_chat(
          self, message: str, chat_history: Optional[List[ChatMessage]] = None
      ) -> StreamingAgentChatResponse:
          raise NotImplementedError("astream_chat not implemented")
  
  
  class TaskStep(BaseModel):
      """Agent task step.
  
      Represents a single input step within the execution run ("Task") of an agent
      given a user input.
  
      The output is returned as a `TaskStepOutput`.
  
      """
  
      task_id: str = Field(..., diescription="Task ID")
      step_id: str = Field(..., description="Step ID")
      input: Optional[str] = Field(default=None, description="User input")
      # memory: BaseMemory = Field(
      #     ..., type=BaseMemory, description="Conversational Memory"
      # )
      step_state: Dict[str, Any] = Field(
          default_factory=dict, description="Additional state for a given step."
      )
  
      # NOTE: the state below may change throughout the course of execution
      # this tracks the relationships to other steps
      next_steps: Dict[str, "TaskStep"] = Field(
          default_factory=dict, description="Next steps to be executed."
      )
      prev_steps: Dict[str, "TaskStep"] = Field(
          default_factory=dict,
          description="Previous steps that were dependencies for this step.",
      )
      is_ready: bool = Field(
          default=True, description="Is this step ready to be executed?"
      )
  
      def get_next_step(
          self,
          step_id: str,
          input: Optional[str] = None,
          step_state: Optional[Dict[str, Any]] = None,
      ) -> "TaskStep":
          """Convenience function to get next step.
  
          Preserve task_id, memory, step_state.
  
          """
          return TaskStep(
              task_id=self.task_id,
              step_id=step_id,
              input=input,
              # memory=self.memory,
              step_state=step_state or self.step_state,
          )
  
      def link_step(
          self,
          next_step: "TaskStep",
      ) -> None:
          """Link to next step.
  
          Add link from this step to next, and from next step to current.
  
          """
          self.next_steps[next_step.step_id] = next_step
          next_step.prev_steps[self.step_id] = self
  
  
  class TaskStepOutput(BaseModel):
      """Agent task step output."""
  
      output: Any = Field(..., description="Task step output")
      task_step: TaskStep = Field(..., description="Task step input")
      next_steps: List[TaskStep] = Field(..., description="Next steps to be executed.")
      is_last: bool = Field(default=False, description="Is this the last step?")
  
      def __str__(self) -> str:
          """String representation."""
          return str(self.output)
  
  
  class Task(BaseModel):
      """Agent Task.
  
      Represents a "run" of an agent given a user input.
  
      """
  
      class Config:
          arbitrary_types_allowed = True
  
      task_id: str = Field(
          default_factory=lambda: str(uuid.uuid4()), type=str, description="Task ID"
      )
      input: str = Field(..., type=str, description="User input")
  
      # NOTE: this is state that may be modified throughout the course of execution of the task
      memory: BaseMemory = Field(
          ...,
          type=BaseMemory,
          description=(
              "Conversational Memory. Maintains state before execution of this task."
          ),
      )
  
      callback_manager: CallbackManager = Field(
          default_factory=CallbackManager,
          exclude=True,
          description="Callback manager for the task.",
      )
  
      extra_state: Dict[str, Any] = Field(
          default_factory=dict,
          description=(
              "Additional user-specified state for a given task. "
              "Can be modified throughout the execution of a task."
          ),
      )
  
  
  class BaseAgentWorker(PromptMixin):
      """Base agent worker."""
  
      class Config:
          arbitrary_types_allowed = True
  
      def _get_prompts(self) -> PromptDictType:
          """Get prompts."""
          # TODO: the ReAct agent does not explicitly specify prompts, would need a
          # refactor to expose those prompts
          return {}
  
      def _get_prompt_modules(self) -> PromptMixinType:
          """Get prompt modules."""
          return {}
  
      def _update_prompts(self, prompts: PromptDictType) -> None:
          """Update prompts."""
  
      @abstractmethod
      def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep:
          """Initialize step from task."""
  
      @abstractmethod
      def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
          """Run step."""
  
      @abstractmethod
      async def arun_step(
          self, step: TaskStep, task: Task, **kwargs: Any
      ) -> TaskStepOutput:
          """Run step (async)."""
          raise NotImplementedError
  
      @abstractmethod
      def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
          """Run step (stream)."""
          # TODO: figure out if we need a different type for TaskStepOutput
          raise NotImplementedError
  
      @abstractmethod
      async def astream_step(
          self, step: TaskStep, task: Task, **kwargs: Any
      ) -> TaskStepOutput:
          """Run step (async stream)."""
          raise NotImplementedError
  
      @abstractmethod
      def finalize_task(self, task: Task, **kwargs: Any) -> None:
          """Finalize task, after all the steps are completed."""
  
      def set_callback_manager(self, callback_manager: CallbackManager) -> None:
          """Set callback manager."""
          # TODO: make this abstractmethod (right now will break some agent impls)
  
apis: llama_index.legacy.callbacks.trace_method llama_index.legacy.bridge.pydantic.Field
extract_api: [((1310, 1331), 'llama_index.legacy.callbacks.trace_method', 'trace_method', (['"""query"""'], {}), "('query')\n", (1322, 1331), False, 'from llama_index.legacy.callbacks import CallbackManager, trace_method\n'), ((1633, 1654), 'llama_index.legacy.callbacks.trace_method', 'trace_method', (['"""query"""'], {}), "('query')\n", (1645, 1654), False, 'from llama_index.legacy.callbacks import CallbackManager, trace_method\n'), ((2613, 2647), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'diescription': '"""Task ID"""'}), "(..., diescription='Task ID')\n", (2618, 2647), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((2667, 2700), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""Step ID"""'}), "(..., description='Step ID')\n", (2672, 2700), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((2728, 2773), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'None', 'description': '"""User input"""'}), "(default=None, description='User input')\n", (2733, 2773), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((2917, 2994), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Additional state for a given step."""'}), "(default_factory=dict, description='Additional state for a given step.')\n", (2922, 2994), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((3175, 3244), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Next steps to be executed."""'}), "(default_factory=dict, description='Next steps to be executed.')\n", (3180, 3244), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((3299, 3399), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Previous steps that were dependencies for this step."""'}), "(default_factory=dict, description=\n    'Previous steps that were dependencies for this step.')\n", (3304, 3399), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((3439, 3508), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(True)', 'description': '"""Is this step ready to be executed?"""'}), "(default=True, description='Is this step ready to be executed?')\n", (3444, 3508), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((4404, 4446), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""Task step output"""'}), "(..., description='Task step output')\n", (4409, 4446), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((4473, 4514), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""Task step input"""'}), "(..., description='Task step input')\n", (4478, 4514), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((4548, 4600), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""Next steps to be executed."""'}), "(..., description='Next steps to be executed.')\n", (4553, 4600), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((4621, 4679), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(False)', 'description': '"""Is this the last step?"""'}), "(default=False, description='Is this the last step?')\n", (4626, 4679), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((5080, 5126), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'type': 'str', 'description': '"""User input"""'}), "(..., type=str, description='User input')\n", (5085, 5126), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((5247, 5364), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'type': 'BaseMemory', 'description': '"""Conversational Memory. Maintains state before execution of this task."""'}), "(..., type=BaseMemory, description=\n    'Conversational Memory. Maintains state before execution of this task.')\n", (5252, 5364), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((5456, 5559), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'CallbackManager', 'exclude': '(True)', 'description': '"""Callback manager for the task."""'}), "(default_factory=CallbackManager, exclude=True, description=\n    'Callback manager for the task.')\n", (5461, 5559), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((5621, 5775), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Additional user-specified state for a given task. Can be modified throughout the execution of a task."""'}), "(default_factory=dict, description=\n    'Additional user-specified state for a given task. Can be modified throughout the execution of a task.'\n    )\n", (5626, 5775), False, 'from llama_index.legacy.bridge.pydantic import BaseModel, Field\n'), ((5010, 5022), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (5020, 5022), False, 'import uuid\n')]