content: from __future__ import annotations
  
  import importlib
  import itertools
  import json
  import logging
  import re
  import shutil
  import tempfile
  import time
  import zlib
  from functools import wraps
  from pathlib import Path
  from typing import cast
  from unittest.mock import AsyncMock, patch
  from uuid import uuid4
  
  import ldp.agent
  import pytest
  from aviary.core import (
      Environment,
      Tool,
      ToolRequestMessage,
      ToolResponseMessage,
      ToolsAdapter,
      ToolSelector,
  )
  from ldp.agent import MemoryAgent, SimpleAgent
  from ldp.graph.memory import Memory, UIndexMemoryModel
  from ldp.graph.ops import OpResult
  from lmi import CommonLLMNames, EmbeddingModel, LiteLLMModel
  from pytest_subtests import SubTests
  from tantivy import Index
  from tenacity import Retrying, retry_if_exception_type, stop_after_attempt
  
  from paperqa.agents import SearchIndex, agent_query
  from paperqa.agents.env import (
      CLINICAL_STATUS_SEARCH_REGEX_PATTERN,
      PaperQAEnvironment,
      clinical_trial_status,
      settings_to_tools,
  )
  from paperqa.agents.main import FAKE_AGENT_TYPE, run_agent
  from paperqa.agents.models import AgentStatus, AnswerResponse
  from paperqa.agents.search import (
      FAILED_DOCUMENT_ADD_ID,
      get_directory_index,
      maybe_get_manifest,
  )
  from paperqa.agents.tools import (
      ClinicalTrialsSearch,
      Complete,
      EnvironmentState,
      GatherEvidence,
      GenerateAnswer,
      PaperSearch,
      Reset,
      make_status,
  )
  from paperqa.docs import Docs
  from paperqa.prompts import CANNOT_ANSWER_PHRASE, CONTEXT_INNER_PROMPT_NOT_DETAILED
  from paperqa.settings import AgentSettings, IndexSettings, Settings
  from paperqa.types import Context, Doc, DocDetails, PQASession, Text
  from paperqa.utils import compute_unique_doc_id, extract_thought, get_year, md5sum
  
  
  @pytest.mark.asyncio
  async def test_get_directory_index(
      subtests: SubTests, agent_test_settings: Settings
  ) -> None:
      # Since agent_test_settings is used by other tests, we use a tempdir so we
      # can delete files without affecting concurrent tests
      with tempfile.TemporaryDirectory() as tempdir:
          shutil.copytree(
              agent_test_settings.agent.index.paper_directory, tempdir, dirs_exist_ok=True
          )
          paper_dir = agent_test_settings.agent.index.paper_directory = Path(tempdir)
  
          index_name = agent_test_settings.agent.index.name = (
              f"stub{uuid4()}"  # Unique across test invocations
          )
          index = await get_directory_index(settings=agent_test_settings)
          assert (
              index.index_name == index_name
          ), "Index name should match its specification"
          assert index.fields == [
              "file_location",
              "body",
              "title",
              "year",
          ], "Incorrect fields in index"
          assert not index.changed, "Expected index to not have changes at this point"
          # bates.txt + empty.txt + flag_day.html + gravity_hill.md + influence.pdf
          # + obama.txt + paper.pdf + pasa.pdf + duplicate_media.pdf
          # + dummy.docx + dummy.pptx + dummy.xlsx,
          # but empty.txt fails to be added
          path_to_id = await index.index_files
          assert (
              sum(id_ != FAILED_DOCUMENT_ADD_ID for id_ in path_to_id.values()) == 12
          ), "Incorrect number of parsed index files"
  
          with subtests.test(msg="check-txt-query"):
              results = await index.query(query="who is Frederick Bates?", min_score=5)
              assert results
              target_doc_path = (paper_dir / "bates.txt").absolute()
              assert results[0].docs.keys() == {md5sum(target_doc_path)}, (
                  f"Expected to find {target_doc_path.name!r}, got citations"
                  f" {[d.formatted_citation for d in results[0].docs.values()]}."
              )
  
              # Check single quoted text in the query doesn't crash us
              results = await index.query(query="Who is 'Bates'")
              assert results
  
              # Check possessive in the query doesn't crash us
              results = await index.query(query="What is Bates' first name")
              assert results
  
          with subtests.test(msg="check-md-query"):
              results = await index.query(query="what is a gravity hill?", min_score=5)
              assert results
              first_result = results[0]
              assert len(first_result.docs) == 1, "Expected one result (gravity_hill.md)"
              target_doc_path = (paper_dir / "gravity_hill.md").absolute()
              expected_ids = {
                  compute_unique_doc_id(
                      None,
                      md5sum(target_doc_path),  # What we actually expect
                  ),
                  compute_unique_doc_id(
                      "10.2307/j.ctt5vkfh7.11",  # Crossref may match this Gravity Hill poem, lol
                      next(iter(first_result.docs.values())).content_hash,
                  ),
              }
              for expected_id in expected_ids:
                  if expected_id in set(first_result.docs.keys()):
                      break
              else:
                  raise AssertionError(
                      f"Failed to match an ID in {expected_ids}, got citations"
                      f" {[d.formatted_citation for d in first_result.docs.values()]}."
                  )
              assert (
                  "gravity hill"
                  in first_result.docs[expected_id].formatted_citation.lower()
              )
  
          # Check getting the same index name will not reprocess files
          with patch.object(Docs, "aadd") as mock_aadd:
              index = await get_directory_index(settings=agent_test_settings)
          assert len(await index.index_files) == len(path_to_id)
          mock_aadd.assert_not_awaited(), "Expected we didn't re-add files"
  
          # Now we actually remove (but not add!) a file from the paper directory,
          # and we still don't reprocess files
          (paper_dir / "obama.txt").unlink()
          with (
              patch.object(
                  Docs, "aadd", autospec=True, side_effect=Docs.aadd
              ) as mock_aadd,
              patch.object(
                  agent_test_settings.agent.index,
                  "files_filter",
                  lambda f: f.suffix in {".txt", ".pdf", ".md"},  # Also, exclude HTML
              ),
          ):
              index = await get_directory_index(settings=agent_test_settings)
          # Subtract 5 for the removed obama.txt,
          # dummy.docx, dummy_jap.docx, dummy.pptx, and dummy.xlsx files,
          # and another 1 for the filtered out flag_day.html
          assert len(await index.index_files) == len(path_to_id) - 5 - 1
          mock_aadd.assert_not_awaited(), "Expected we didn't re-add files"
  
          # Note let's delete files.zip, and confirm we can't load the index
          await (await index.file_index_filename).unlink()
          with pytest.raises(RuntimeError, match="please rebuild"):
              await get_directory_index(settings=agent_test_settings, build=False)
  
  
  @pytest.mark.asyncio
  async def test_resuming_crashed_index_build(agent_test_settings: Settings) -> None:
      index_settings = agent_test_settings.agent.index
      crash_threshold, index_settings.concurrency = 3, 2
      num_source_files = len(
          [
              x
              for x in cast("Path", index_settings.paper_directory).iterdir()
              # Filter out .csv and .DS_Store files
              if x.suffix != ".csv" and agent_test_settings.agent.index.files_filter(x)
          ]
      )
      assert (
          num_source_files >= 5
      ), "Less source files than this test was designed to work with"
      call_count = 0
      original_docs_aadd = Docs.aadd
  
      async def crashing_aadd(*args, **kwargs) -> str | None:
          nonlocal call_count
          if call_count == crash_threshold:
              raise RuntimeError("Unexpected crash.")
          call_count += 1
          return await original_docs_aadd(*args, **kwargs)
  
      # 1. Try to build an index, and crash halfway through
      with (
          pytest.raises(ExceptionGroup, match="unhandled"),
          patch.object(
              Docs, "aadd", side_effect=crashing_aadd, autospec=True
          ) as mock_aadd,
      ):
          await get_directory_index(settings=agent_test_settings)
      mock_aadd.assert_awaited()
  
      # 2. Resume and complete building the index
      for attempt in Retrying(
          stop=stop_after_attempt(3),
          # zlib.error: Error -5 while decompressing data: incomplete or truncated stream
          retry=retry_if_exception_type(zlib.error),
      ):
          with (
              attempt,
              patch.object(
                  Docs, "aadd", autospec=True, side_effect=Docs.aadd
              ) as mock_aadd,
          ):
              index = await get_directory_index(settings=agent_test_settings)
  
      assert len(await index.index_files) == num_source_files
      assert (
          mock_aadd.await_count < num_source_files
      ), "Should not rebuild the whole index"
  
  
  @pytest.mark.asyncio
  async def test_getting_manifest(
      agent_test_settings: Settings, stub_data_dir: Path, caplog
  ) -> None:
      agent_test_settings.agent.index.manifest_file = "stub_manifest.csv"
  
      # Since stub_manifest.csv is used by other tests, we use a tempdir so we
      # can modify it without affecting concurrent tests
      with tempfile.TemporaryDirectory() as tempdir, caplog.at_level(logging.WARNING):
          shutil.copytree(stub_data_dir, tempdir, dirs_exist_ok=True)
          agent_test_settings.agent.index.paper_directory = tempdir
          manifest_filepath = (
              await agent_test_settings.agent.index.finalize_manifest_file()
          )
          assert manifest_filepath
          assert await maybe_get_manifest(manifest_filepath)
          assert not caplog.records
  
          # If a header line isn't present, our manifest extraction should fail
          original_manifest_lines = (await manifest_filepath.read_text()).splitlines()
          await manifest_filepath.write_text(data="\n".join(original_manifest_lines[1:]))
          await maybe_get_manifest(manifest_filepath)
          assert len(caplog.records) == 1
          assert caplog.records[0].levelno == logging.ERROR
  
  
  EXPECTED_STUB_DATA_FILES = {
      "bates.txt",
      "duplicate_media.pdf",
      "empty.txt",
      "flag_day.html",
      "gravity_hill.md",
      "influence.pdf",
      "obama.txt",
      "paper.pdf",
      "pasa.pdf",
      "dummy.docx",
      "dummy_jap.docx",
      "dummy.pptx",
      "dummy.xlsx",
  }
  
  
  @pytest.mark.asyncio
  async def test_get_directory_index_w_manifest(agent_test_settings: Settings) -> None:
      # Set the paper_directory to be a relative path as starting point to confirm this
      # won't trip us up, and set the manifest file too
      abs_paper_dir = cast("Path", agent_test_settings.agent.index.paper_directory)
      agent_test_settings.agent.index.paper_directory = abs_paper_dir.relative_to(
          Path.cwd()
      )
      agent_test_settings.agent.index.manifest_file = "stub_manifest.csv"
  
      # Now set up both relative and absolute test settings
      relative_test_settings = agent_test_settings.model_copy(deep=True)
      absolute_test_settings = agent_test_settings.model_copy(deep=True)
      absolute_test_settings.agent.index.use_absolute_paper_directory = True
      assert (
          relative_test_settings != absolute_test_settings
      ), "We need to be able to differentiate between relative and absolute settings"
      del agent_test_settings
  
      relative_index = await get_directory_index(settings=relative_test_settings)
      assert (
          set((await relative_index.index_files).keys()) == EXPECTED_STUB_DATA_FILES
      ), "Incorrect index files, should be relative to share indexes across machines"
      absolute_index = await get_directory_index(settings=absolute_test_settings)
      assert set((await absolute_index.index_files).keys()) == {
          str(abs_paper_dir / f) for f in EXPECTED_STUB_DATA_FILES
      }, (
          "Incorrect index files, should be absolute to deny sharing indexes across"
          " machines"
      )
      for index in (relative_index, absolute_index):
          assert index.fields == [
              "file_location",
              "body",
              "title",
              "year",
          ], "Incorrect fields in index"
  
          results = await index.query(query="who is Frederick Bates?")
          top_result = next(iter(results[0].docs.values()))
  
          # note: we get every possible field from the manifest constructed in maybe_get_manifest,
          # and then DocDetails construction sets the dockey to the doc_id.
          assert top_result.dockey == top_result.doc_id
          # note: this title comes from the manifest, so we know it worked
          assert top_result.title == "Frederick Bates (Wikipedia article)"
  
          assert "wikipedia article" in top_result.citation.lower(), (
              "Other tests check we can override citation,"
              " so here we check here it's actually populated"
          )
  
  
  @pytest.mark.asyncio
  async def test_get_directory_index_w_no_citations(
      agent_test_settings: Settings,
  ) -> None:
      agent_test_settings.agent.index.manifest_file = "stub_manifest_nocitation.csv"
      index = await get_directory_index(settings=agent_test_settings)
  
      results = await index.query(query="who is Frederick Bates?")
      top_result = next(iter(results[0].docs.values()))
  
      assert not top_result.citation
  
  
  @pytest.mark.flaky(reruns=2, only_rerun=["AssertionError", "httpx.RemoteProtocolError"])
  @pytest.mark.parametrize("agent_type", [FAKE_AGENT_TYPE, ToolSelector, SimpleAgent])
  @pytest.mark.parametrize("llm_name", ["gpt-4o", "gemini/gemini-2.0-flash-lite"])
  @pytest.mark.asyncio
  async def test_agent_types(
      agent_test_settings: Settings,
      agent_type: str | type,
      llm_name: str,
      subtests: SubTests,
  ) -> None:
      question = "How can you use XAI for chemical property prediction?"
      # make sure agent_llm is different from default, so we can correctly track tokens
      # for agent
      agent_test_settings.agent.agent_llm = llm_name
      agent_test_settings.llm = "gpt-4o-mini"
      agent_test_settings.summary_llm = "gpt-4o-mini"
      agent_test_settings.agent.agent_prompt += (
          "\n\nCall each tool once in appropriate order and"
          " accept the answer for now, as we're in debug mode."
      )
      with patch.object(
          Index, "open", side_effect=Index.open, autospec=True
      ) as mock_open:
          response = await agent_query(
              question, agent_test_settings, agent_type=agent_type
          )
      assert (
          mock_open.call_count <= 1
      ), "Expected one Index.open call, or possibly zero if multiprocessing tests"
      assert response.session.answer, "Answer not generated"
      assert response.session.answer != CANNOT_ANSWER_PHRASE, "Answer not generated"
      assert response.session.context, "No contexts were found"
      assert response.session.question == question
      agent_llm = agent_test_settings.agent.agent_llm
      # TODO: once LDP can track tokens, we can remove this check
      if agent_type not in {FAKE_AGENT_TYPE, SimpleAgent}:
          assert (
              response.session.token_counts[agent_llm][0] > 500
          ), "Expected many prompt tokens"
          assert (
              response.session.token_counts[agent_llm][1] > 30
          ), "Expected many completion tokens"
          assert response.session.cost > 0, "Expected nonzero cost"
  
      with subtests.test("Test citation formatting"):
          citation_w_et_al = r"\b[\w\-]+\set\sal\.\s\([0-9]{4}\)"
          assert not re.search(
              citation_w_et_al, response.session.answer
          ), "Answer contains citation with et al. instead of citation key"
  
          missing_pages_regex = r"\b([a-zA-Z]+\d{4}[a-zA-Z]*\s+\d+-\d+)\b"
          assert not re.search(
              missing_pages_regex, response.session.answer
          ), "Answer contains citation with missing 'pages' keyword"
  
  
  @pytest.mark.asyncio
  async def test_successful_memory_agent(agent_test_settings: Settings) -> None:
      tic = time.perf_counter()
      memory_id = "call_Wtmv95JbNcQ2nRQCZBoOfcJy"  # Stub value
      memory = Memory(
          query=(
              "Use the tools to answer the question: How can you use XAI for chemical"
              " property prediction?\n\nWhen the answer looks sufficient,"
              " you can terminate by calling the {complete_tool_name} tool."
              " If the answer does not look sufficient,"
              " and you have already tried to answer several times,"
              " you can terminate by calling the {complete_tool_name} tool."
              " The current status of evidence/papers/cost is "
              f"{make_status(total_paper_count=0, relevant_paper_count=0, evidence_count=0, cost=0.0)}"  # Started 0  # noqa: E501
              "\n\nTool request message '' for tool calls: paper_search(query='XAI for"
              " chemical property prediction', min_year='2018', max_year='2024')"
              f" [id={memory_id}]\n\nTool response message '"
              f"{make_status(total_paper_count=2, relevant_paper_count=0, evidence_count=0, cost=0.0)}"  # Found 2  # noqa: E501
              f"' for tool call ID {memory_id} of tool 'paper_search'"
          ),
          input=(
              "Use the tools to answer the question: How can you use XAI for chemical"
              " property prediction?\n\nWhen the answer looks sufficient,"
              " you can terminate by calling the {complete_tool_name} tool."
              " If the answer does not look sufficient,"
              " and you have already tried to answer several times,"
              " you can terminate by calling the {complete_tool_name} tool."
              " The current status of evidence/papers/cost is "
              f"{make_status(total_paper_count=0, relevant_paper_count=0, evidence_count=0, cost=0.0)}"
          ),
          output=(
              "Tool request message '' for tool calls: paper_search(query='XAI for"
              " chemical property prediction', min_year='2018', max_year='2024')"
              f" [id={memory_id}]"
          ),
          value=5.0,  # Stub value
          template="Input: {input}\n\nOutput: {output}\n\nDiscounted Reward: {value}",
      )
      memory_model = UIndexMemoryModel(
          embedding_model=EmbeddingModel.from_name("text-embedding-3-small")
      )
      await memory_model.add_memory(memory)
      serialized_memory_model = memory_model.model_dump(exclude_none=True)
      query = "How can you use XAI for chemical property prediction?"
      # NOTE: use Claude 3 for its <thinking> feature, testing regex replacement of it
      agent_test_settings.agent.agent_llm = CommonLLMNames.CLAUDE_37_SONNET.value
      agent_test_settings.agent.agent_config = {
          "memories": serialized_memory_model.pop("memories"),
          "memory_model": serialized_memory_model,
      }
  
      thoughts: list[str] = []
      orig_llm_model_call = LiteLLMModel.call
  
      async def on_agent_action(  # noqa: RUF029
          action: OpResult[ToolRequestMessage], *_
      ) -> None:
          thoughts.append(extract_thought(content=action.value.content))
  
      async def llm_model_call(*args, **kwargs):
          # NOTE: "required" will not lead to thoughts being emitted, it has to be "auto"
          # https://docs.anthropic.com/en/docs/build-with-claude/tool-use#chain-of-thought
          args = args[:-1]  # removing last element (tool_choice) from args
          return await orig_llm_model_call(*args, tool_choice="auto", **kwargs)  # type: ignore[misc]
  
      with patch.object(LiteLLMModel, "call", side_effect=llm_model_call, autospec=True):
          response = await agent_query(
              query,
              agent_test_settings,
              Docs(),
              agent_type=f"{ldp.agent.__name__}.{MemoryAgent.__name__}",
              on_agent_action_callback=on_agent_action,
          )
      assert response.status == AgentStatus.SUCCESS, "Agent did not succeed"
      assert (
          time.perf_counter() - tic <= agent_test_settings.agent.timeout
      ), "Agent should not have timed out"
      assert all(thought and "<thinking>" not in thought for thought in thoughts)
  
  
  @pytest.mark.parametrize("agent_type", [ToolSelector, SimpleAgent])
  @pytest.mark.asyncio
  async def test_timeout(agent_test_settings: Settings, agent_type: str | type) -> None:
      agent_test_settings.prompts.pre = None
      agent_test_settings.agent.timeout = 0.05  # Give time for Environment.reset()
      agent_test_settings.llm = "gpt-4o-mini"
      agent_test_settings.agent.tool_names = {"gen_answer", "complete"}
      orig_exec_tool_calls = PaperQAEnvironment.exec_tool_calls
      tool_responses: list[list[ToolResponseMessage]] = []
  
      async def spy_exec_tool_calls(*args, **kwargs) -> list[ToolResponseMessage]:
          responses = await orig_exec_tool_calls(*args, **kwargs)
          tool_responses.append(responses)
          return responses
  
      with patch.object(PaperQAEnvironment, "exec_tool_calls", spy_exec_tool_calls):
          response = await agent_query(
              query="Are COVID-19 vaccines effective?",
              settings=agent_test_settings,
              agent_type=agent_type,
          )
      # Ensure that GenerateAnswerTool was called in truncation's failover
      assert response.status == AgentStatus.TRUNCATED, "Agent did not timeout"
      assert CANNOT_ANSWER_PHRASE in response.session.answer
      (last_response,) = tool_responses[-1]
      assert (
          "no papers" in last_response.content
      ), "Expecting agent to been shown specifics on the failure"
  
  
  @pytest.mark.flaky(reruns=5, only_rerun=["AssertionError"])
  @pytest.mark.asyncio
  async def test_propagate_options(agent_test_settings: Settings) -> None:
      llm_name = "gpt-4o-mini"
      default_llm_names = {
          cls.model_fields[name].default
          for name, cls in itertools.product(("llm", "summary_llm"), (Settings,))
      }
      assert (
          llm_name not in default_llm_names
      ), f"Assertions require not matching a default LLM name in {default_llm_names}."
  
      agent_test_settings.llm = llm_name
      agent_test_settings.answer.answer_max_sources = 5
      agent_test_settings.answer.evidence_k = 6
      agent_test_settings.answer.answer_length = "400 words"
      agent_test_settings.prompts.pre = None
      agent_test_settings.prompts.system = "End all responses with ###"
      agent_test_settings.prompts.context_inner = CONTEXT_INNER_PROMPT_NOT_DETAILED
      agent_test_settings.answer.evidence_skip_summary = True
  
      docs = Docs()
      response = await agent_query(
          query="What is a self-explanatory model?",
          settings=agent_test_settings,
          docs=docs,
          agent_type=FAKE_AGENT_TYPE,
      )
      assert response.status == AgentStatus.SUCCESS, "Agent did not succeed"
      result = response.session
      assert len(result.answer) > 200, "Answer did not return any results"
      assert "###" in result.answer, "Answer did not propagate system prompt"
      assert docs.docs, "Expected docs to have been added"
      assert all(isinstance(d, DocDetails) for d in docs.docs.values())
      assert all(
          d.file_location for d in docs.docs.values()  # type: ignore[union-attr]
      ), "Expected file location to be populated"
      assert len(result.contexts) >= 2, "Test expects a few contexts"
      # Subtract 2 to allow tolerance for chunks with leading/trailing whitespace
      num_contexts_sufficient_length = sum(
          len(c.context) >= agent_test_settings.parsing.reader_config["chunk_chars"] - 2
          for c in result.contexts
      )
      # Check most contexts have the expected length
      assert (
          num_contexts_sufficient_length >= len(result.contexts) - 1
      ), "Summary was not skipped"
  
  
  @pytest.mark.asyncio
  async def test_gather_evidence_rejects_empty_docs(
      agent_test_settings: Settings,
  ) -> None:
  
      @wraps(GenerateAnswer.gen_answer)
      async def gen_answer(self, state) -> str:  # noqa: ARG001, RUF029
          return f"{CANNOT_ANSWER_PHRASE}."
  
      # Patch GenerateAnswerTool.gen_answer so that if this tool is chosen first,
      # we keep running until we get truncated
      with (
          patch(
              "paperqa.agents.env.settings_to_tools",
              side_effect=[
                  [
                      Tool.from_function(
                          GatherEvidence(
                              settings=agent_test_settings,
                              summary_llm_model=agent_test_settings.get_summary_llm(),
                              embedding_model=agent_test_settings.get_embedding_model(),
                          ).gather_evidence,
                          concurrency_safe=GatherEvidence.CONCURRENCY_SAFE,
                      ),
                      Tool.from_function(
                          GenerateAnswer(
                              settings=agent_test_settings,
                              llm_model=agent_test_settings.get_llm(),
                              summary_llm_model=agent_test_settings.get_summary_llm(),
                              embedding_model=agent_test_settings.get_embedding_model(),
                          ).gen_answer,
                          concurrency_safe=GenerateAnswer.CONCURRENCY_SAFE,
                      ),
                  ]
              ],
          ),
          patch.object(GenerateAnswer, "gen_answer", gen_answer),
      ):
          agent_test_settings.agent = AgentSettings(
              max_timesteps=3,
              search_count=agent_test_settings.agent.search_count,
              index=IndexSettings(
                  paper_directory=agent_test_settings.agent.index.paper_directory,
                  index_directory=agent_test_settings.agent.index.index_directory,
              ),
          )
          response = await agent_query(
              query="Are COVID-19 vaccines effective?",
              settings=agent_test_settings,
              docs=Docs(),
          )
      assert (
          response.status == AgentStatus.TRUNCATED
      ), "Agent should have hit its max timesteps"
  
  
  @pytest.mark.parametrize("callback_type", [None, "async"])
  @pytest.mark.flaky(reruns=3, only_rerun=["AssertionError", "EmptyDocsError"])
  @pytest.mark.asyncio
  async def test_agent_sharing_state(
      agent_test_settings: Settings, subtests: SubTests, callback_type: str | None
  ) -> None:
      SAVE_API_COSTS_FILES_TO_EXCLUDE = {
          "pasa.pdf",
          *(f"dummy{x}" for x in (".docx", "_jap.docx", ".pptx", ".xlsx")),
      }
  
      def files_filter(f) -> bool:
          return (
              f.name not in SAVE_API_COSTS_FILES_TO_EXCLUDE
              and IndexSettings.model_fields["files_filter"].default(f)
          )
  
      agent_test_settings.agent.index.files_filter = files_filter
      agent_test_settings.agent.search_count = 3  # Keep low for speed
      agent_test_settings.answer.evidence_k = 2
      agent_test_settings.answer.answer_max_sources = 1
      llm_model = agent_test_settings.get_llm()
      summary_llm_model = agent_test_settings.get_summary_llm()
      embedding_model = agent_test_settings.get_embedding_model()
  
      callbacks = {}
      if callback_type == "async":
          gen_answer_initialized_callback = AsyncMock()
          gen_answer_completed_callback = AsyncMock()
          gather_evidence_initialized_callback = AsyncMock()
          gather_evidence_completed_callback = AsyncMock()
  
          callbacks = {
              "gen_answer_initialized": [gen_answer_initialized_callback],
              "gen_answer_completed": [gen_answer_completed_callback],
              "gather_evidence_initialized": [gather_evidence_initialized_callback],
              "gather_evidence_completed": [gather_evidence_completed_callback],
          }
  
      agent_test_settings.agent.callbacks = callbacks
  
      session = PQASession(question="What is a self-explanatory model?")
      env_state = EnvironmentState(docs=Docs(), session=session)
      built_index = await get_directory_index(settings=agent_test_settings)
      assert await built_index.count, "Index build did not work"
  
      with subtests.test(msg="Custom and default environment status"):
          assert re.search(
              pattern=EnvironmentState.STATUS_SEARCH_REGEX_PATTERN,
              string=env_state.status,
          ), "Default status not formatted correctly"
          # override the status function with a new one
  
          def new_status(state: EnvironmentState) -> str:
              return f"Custom status: paper count = {len(state.docs.docs)}"
  
          env_state.status_fn = new_status
          assert env_state.status == new_status(
              env_state
          ), "Custom status not set correctly."
          env_state.status_fn = None
  
      # run an initial complete tool to see that the answer object is populated by it
      # this simulates if no gen_answer tool was called
      with subtests.test(msg=Complete.__name__):
          complete_tool = Complete()
          await complete_tool.complete(state=env_state, has_successful_answer=False)
          assert (
              env_state.session.answer == Complete.NO_ANSWER_PHRASE
          ), "Complete did not succeed"
          # now we wipe the answer for further tests
          env_state.session.answer = ""
  
      with subtests.test(msg=PaperSearch.__name__):
          search_tool = PaperSearch(
              settings=agent_test_settings, embedding_model=embedding_model
          )
          with (
              patch.object(
                  SearchIndex, "save_index", wraps=SearchIndex.save_index, autospec=True
              ) as mock_save_index,
              patch.object(
                  Index, "open", side_effect=Index.open, autospec=True
              ) as mock_open,
          ):
              await search_tool.paper_search(
                  "XAI self explanatory model",
                  min_year=None,
                  max_year=None,
                  state=env_state,
              )
              assert env_state.docs.docs, "Search did not add any papers"
              assert (
                  mock_open.call_count <= 1
              ), "Expected one Index.open call, or possibly zero if multiprocessing tests"
              assert all(
                  isinstance(d, Doc) for d in env_state.docs.docs.values()
              ), "Document type or DOI propagation failure"
  
              await search_tool.paper_search(
                  "XAI for chemical property prediction",
                  min_year=2018,
                  max_year=2024,
                  state=env_state,
              )
              assert (
                  mock_open.call_count <= 1
              ), "Expected one Index.open call, or possibly zero if multiprocessing tests"
              mock_save_index.assert_not_awaited()
  
      with subtests.test(msg=GatherEvidence.__name__):
          assert not session.contexts, "No contexts is required for a later assertion"
  
          gather_evidence_tool = GatherEvidence(
              settings=agent_test_settings,
              summary_llm_model=summary_llm_model,
              embedding_model=embedding_model,
          )
  
          response = await gather_evidence_tool.gather_evidence(
              session.question, state=env_state
          )
  
          if callback_type == "async":
              gather_evidence_initialized_callback.assert_awaited_once_with(env_state)
              gather_evidence_completed_callback.assert_awaited_once_with(env_state)
  
          split = re.split(r"(\d+) pieces of evidence", response, maxsplit=1)
          assert len(split) == 3, "Unexpected response shape"
          total_added_1 = int(split[1])
          assert total_added_1 > 0, "Expected non-negative added evidence count"
          assert len(env_state.get_relevant_contexts()) == total_added_1
          assert (
              response.count("\n- ") == 1
          ), "Expected exactly one best evidence to be shown"
  
          # now adjust to give the agent 2x pieces of evidence
          gather_evidence_tool.settings.agent.agent_evidence_n = 2
          # also reset the question to ensure that contexts are
          # only returned to the agent for the new question
          new_question = "How does XAI relate to a self-explanatory model?"
          response = await gather_evidence_tool.gather_evidence(
              new_question, state=env_state
          )
          assert len({c.question for c in session.contexts}) == 2, "Expected 2 questions"
          # now we make sure this is only for the old question
          for context in session.contexts:
              if context.question != new_question:
                  assert (
                      context.context[:50] not in response
                  ), "gather_evidence should not return any contexts for the old question"
          assert (
              sum(
                  (1 if (context.context[:30] in response) else 0)
                  for context in session.contexts
                  if context.question == new_question
              )
              == 2
          ), "gather_evidence should only return 2 contexts for the new question"
          split = re.split(r"(\d+) pieces of evidence", response, maxsplit=1)
          assert len(split) == 3, "Unexpected response shape"
          total_added_2 = int(split[1])
          assert total_added_2 > 0, "Expected non-negative added evidence count"
  