content: import logging
  from collections.abc import Callable
  from copy import deepcopy
  from typing import Any, ClassVar, Self, cast
  from uuid import UUID
  
  from aviary.core import (
      Environment,
      Frame,
      Message,
      Messages,
      Tool,
      ToolRequestMessage,
      ToolResponseMessage,
  )
  from aviary.env import ENV_REGISTRY
  from aviary.utils import MultipleChoiceQuestion
  from lmi import EmbeddingModel, LiteLLMModel
  
  from paperqa.docs import Docs
  from paperqa.settings import Settings
  from paperqa.sources.clinical_trials import (
      CLINICAL_TRIALS_BASE,
      partition_clinical_trials_by_source,
  )
  from paperqa.types import PQASession
  from paperqa.utils import get_year
  
  from .tools import (
      AVAILABLE_TOOL_NAME_TO_CLASS,
      DEFAULT_TOOL_NAMES,
      ClinicalTrialsSearch,
      Complete,
      EnvironmentState,
      GatherEvidence,
      GenerateAnswer,
      NamedTool,
      PaperSearch,
      Reset,
  )
  
  logger = logging.getLogger(__name__)
  
  POPULATE_FROM_SETTINGS = None
  
  
  def settings_to_tools(  # noqa: PLR0912
      settings: Settings,
      llm_model: LiteLLMModel | None = POPULATE_FROM_SETTINGS,
      summary_llm_model: LiteLLMModel | None = POPULATE_FROM_SETTINGS,
      embedding_model: EmbeddingModel | None = POPULATE_FROM_SETTINGS,
  ) -> list[Tool]:
      """
      Convert a Settings into tools, confirming the complete tool is present.
  
      NOTE: the last element of the return will always be Complete.
      """
      llm_model = llm_model or settings.get_llm()
      summary_llm_model = summary_llm_model or settings.get_summary_llm()
      embedding_model = embedding_model or settings.get_embedding_model()
      tools: list[Tool] = []
      for tool_type in (
          [AVAILABLE_TOOL_NAME_TO_CLASS[name] for name in DEFAULT_TOOL_NAMES]
          if settings.agent.tool_names is None
          else [
              AVAILABLE_TOOL_NAME_TO_CLASS[name]
              for name in set(settings.agent.tool_names)
          ]
      ):
  
          def make_tool(fn: Callable, tool_type: type[NamedTool] = tool_type) -> Tool:
              return Tool.from_function(fn, concurrency_safe=tool_type.CONCURRENCY_SAFE)
  
          if issubclass(tool_type, PaperSearch):
              tool = make_tool(
                  PaperSearch(
                      settings=settings, embedding_model=embedding_model
                  ).paper_search
              )
              for pname in ("min_year", "max_year"):
                  tool.info.get_properties()[pname]["description"] = cast(
                      "str", tool.info.get_properties()[pname]["description"]
                  ).format(current_year=get_year())
          elif issubclass(tool_type, GatherEvidence):
              gather_evidence_tool = GatherEvidence(
                  settings=settings,
                  summary_llm_model=summary_llm_model,
                  embedding_model=embedding_model,
              )
  
              # if we're using the SearchClinicalTrialsTool,
              # we override this tool's docstring/prompt
              # because the default prompt is unaware of the clinical trials tool
  
              if ClinicalTrialsSearch.TOOL_FN_NAME in (
                  settings.agent.tool_names or DEFAULT_TOOL_NAMES
              ):
                  gather_evidence_tool.gather_evidence.__func__.__doc__ = (  # type: ignore[attr-defined]
                      ClinicalTrialsSearch.GATHER_EVIDENCE_TOOL_PROMPT_OVERRIDE
                  )
                  gather_evidence_tool.partitioning_fn = (
                      partition_clinical_trials_by_source
                  )
  
              tool = make_tool(gather_evidence_tool.gather_evidence)
  
          elif issubclass(tool_type, GenerateAnswer):
              generate_answer_tool = GenerateAnswer(
                  settings=settings,
                  llm_model=llm_model,
                  summary_llm_model=summary_llm_model,
                  embedding_model=embedding_model,
              )
  
              if ClinicalTrialsSearch.TOOL_FN_NAME in (
                  settings.agent.tool_names or DEFAULT_TOOL_NAMES
              ):
                  generate_answer_tool.partitioning_fn = (
                      partition_clinical_trials_by_source
                  )
  
              tool = make_tool(generate_answer_tool.gen_answer)
  
          elif issubclass(tool_type, Reset):
              tool = make_tool(Reset().reset)
          elif issubclass(tool_type, Complete):
              tool = make_tool(Complete().complete)
          elif issubclass(tool_type, ClinicalTrialsSearch):
              tool = make_tool(
                  ClinicalTrialsSearch(
                      search_count=settings.agent.search_count, settings=settings
                  ).clinical_trials_search
              )
          else:
              raise NotImplementedError(f"Didn't handle tool type {tool_type}.")
          if tool.info.name == Complete.complete.__name__:
              tools.append(tool)  # Place at the end
          else:
              tools.insert(0, tool)
      return tools
  
  
  def make_clinical_trial_status(
      total_paper_count: int,
      relevant_paper_count: int,
      total_clinical_trials: int,
      relevant_clinical_trials: int,
      evidence_count: int,
      cost: float,
  ) -> str:
      return (
          f"Status: Paper Count={total_paper_count}"
          f" | Relevant Papers={relevant_paper_count}"
          f" | Clinical Trial Count={total_clinical_trials}"
          f" | Relevant Clinical Trials={relevant_clinical_trials}"
          f" | Current Evidence={evidence_count}"
          f" | Current Cost=${cost:.4f}"
      )
  
  
  # SEE: https://regex101.com/r/L0L5MH/1
  CLINICAL_STATUS_SEARCH_REGEX_PATTERN: str = (
      r"Status: Paper Count=(\d+) \| Relevant Papers=(\d+)(?:\s\|\sClinical Trial"
      r" Count=(\d+)\s\|\sRelevant Clinical Trials=(\d+))?\s\|\sCurrent Evidence=(\d+)"
  )
  
  
  def clinical_trial_status(state: "EnvironmentState") -> str:
      relevant_contexts = state.get_relevant_contexts()
      return make_clinical_trial_status(
          total_paper_count=len(
              {
                  d.dockey
                  for d in state.docs.docs.values()
                  if CLINICAL_TRIALS_BASE
                  not in getattr(d, "other", {}).get("client_source", [])
              }
          ),
          relevant_paper_count=len(
              {
                  c.text.doc.dockey
                  for c in relevant_contexts
                  if CLINICAL_TRIALS_BASE
                  not in getattr(c.text.doc, "other", {}).get("client_source", [])
              }
          ),
          total_clinical_trials=len(
              {
                  d.dockey
                  for d in state.docs.docs.values()
                  if CLINICAL_TRIALS_BASE
                  in getattr(d, "other", {}).get("client_source", [])
              }
          ),
          relevant_clinical_trials=len(
              {
                  c.text.doc.dockey
                  for c in relevant_contexts
                  if CLINICAL_TRIALS_BASE
                  in getattr(c.text.doc, "other", {}).get("client_source", [])
              }
          ),
          evidence_count=len(relevant_contexts),
          cost=state.session.cost,
      )
  
  
  class PaperQAEnvironment(Environment[EnvironmentState]):
      """Environment connecting paper-qa's tools with state."""
  
      def __init__(
          self,
          query: str | MultipleChoiceQuestion,
          settings: Settings,
          docs: Docs,
          llm_model: LiteLLMModel | None = POPULATE_FROM_SETTINGS,
          summary_llm_model: LiteLLMModel | None = POPULATE_FROM_SETTINGS,
          embedding_model: EmbeddingModel | None = POPULATE_FROM_SETTINGS,
          session_id: UUID | None = None,
          **env_kwargs,
      ):
          super().__init__(**env_kwargs)
          self._query = query
          self._settings = settings
          self._docs = docs
          self._llm_model = llm_model
          self._summary_llm_model = summary_llm_model
          self._embedding_model = embedding_model
          self._session_id = session_id
  
      @classmethod
      def from_task(cls, task: str) -> Self:
          return cls(query=task, settings=Settings(), docs=Docs())
  
      def make_tools(self) -> list[Tool]:
          return settings_to_tools(
              settings=self._settings,
              llm_model=self._llm_model,
              summary_llm_model=self._summary_llm_model,
              embedding_model=self._embedding_model,
          )
  
      async def _reset_docs(self) -> None:
          """Hook to reset the docs when creating the initial state."""
          self._docs.clear_docs()
  
      async def make_initial_state(self) -> EnvironmentState:
          await self._reset_docs()
  
          status_fn = None
  
          if ClinicalTrialsSearch.TOOL_FN_NAME in (
              self._settings.agent.tool_names or DEFAULT_TOOL_NAMES
          ):
              status_fn = clinical_trial_status
  
          session_kwargs: dict[str, Any] = {}
          if self._session_id:
              session_kwargs["id"] = self._session_id
          return EnvironmentState(
              docs=self._docs,
              session=PQASession(
                  question=(
                      self._query
                      if isinstance(self._query, str)
                      else self._query.question_prompt
                  ),
                  config_md5=self._settings.md5,
                  **session_kwargs,
              ),
              status_fn=status_fn,
          )
  
      async def reset(self) -> tuple[list[Message], list[Tool]]:
          # NOTE: don't build the index here, as sometimes we asyncio.gather over this
          # method, and our current design (as of v5.0.10) could hit race conditions
          # because index building does not use file locks
          self.state, self.tools = (await self.make_initial_state()), self.make_tools()
          return (
              [
                  Message(
                      content=self._settings.agent.agent_prompt.format(
                          question=self.state.session.question,
                          status=self.state.status,
                          complete_tool_name=Complete.TOOL_FN_NAME,
                      ),
                  )
              ],
              self.tools,
          )
  
      def export_frame(self) -> Frame:
          return Frame(state=self.state, info={"query": self._query})
  
      def _has_excess_answer_failures(self) -> bool:
          if self._settings.answer.max_answer_attempts is None:
              return False
          return (
              sum(
                  tn == GenerateAnswer.gen_answer.__name__
                  for s in self.state.session.tool_history
                  for tn in s
              )
              > self._settings.answer.max_answer_attempts
          )
  
      USE_POST_PROCESSED_REWARD: ClassVar[float] = 0.0
  
      async def step(
          self, action: ToolRequestMessage
      ) -> tuple[Messages, float, bool, bool]:
          self.state.record_action(action)
  
          response_messages = cast(
              "list[Message]",
              await self.exec_tool_calls(
                  action,
                  concurrency=True,  # We allow tools to define their own concurrency
                  state=self.state,
                  handle_tool_exc=True,
              ),
          ) or [Message(content=f"No tool calls input in tool request {action}.")]
          done = any(
              isinstance(msg, ToolResponseMessage)
              and msg.name == Complete.complete.__name__
              for msg in response_messages
          )
          if not done and self._has_excess_answer_failures():
              # If the caller set max_answer_attempts, and the agent has tried to answer
              # too many times, we consider this done, but we cannot determine success
              # because we're not calling the complete tool
              self.state.session.has_successful_answer = None
              done = True
          return (
              response_messages,
              self.USE_POST_PROCESSED_REWARD,
              done,
              False,  # Let caller determine truncations
          )
  
      async def get_id(self) -> str:
          if (
              isinstance(self._query, str)
              or self._query.question_id
              == MultipleChoiceQuestion.model_fields["question_id"].default
          ):
              details = (
                  ", as just a question was configured"
                  if isinstance(self._query, str)
                  else ", as the default ID remains present"
              )
              raise ValueError(f"No question ID was configured{details}.")
          return str(self._query.question_id)
  
      def __deepcopy__(self, memo) -> Self:
          copy_state = deepcopy(self.state, memo)
          # We don't know the side effects of deep copying a litellm.Router,
          # so we force a shallow copy of these LiteLLMModels
          env_model_kwargs: dict[str, Any] = {
              name: model if model is None else type(model)(**model.model_dump())
              for name, model in (
                  ("llm_model", self._llm_model),
                  ("summary_llm_model", self._summary_llm_model),
                  ("embedding_model", self._embedding_model),
              )
          }
          copy_self = type(self)(
              query=self._query,  # No need to copy since we read only
              settings=deepcopy(self._settings, memo),  # Deepcopy just to be safe
              docs=copy_state.docs,
              **env_model_kwargs,
          )
          copy_self.state = copy_state
          # Because we shallow copied the LiteLLMModels, we need to re-make the
          # tool functions within the tools
          copy_self.tools = copy_self.make_tools()
          return copy_self
  
  
  ENV_REGISTRY["paperqa"] = "paperqa.agents.env", PaperQAEnvironment.__name__
  