content: import asyncio
  import contextlib
  import csv
  import io
  import json
  import os
  import pathlib
  import pickle
  import random
  import re
  import string
  import sys
  from collections.abc import AsyncIterable, Sequence
  from copy import deepcopy
  from datetime import datetime, timedelta
  from io import BytesIO
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, cast
  from unittest.mock import MagicMock, call, patch
  from uuid import UUID, uuid4
  
  import anyio
  import httpx
  import litellm
  import litellm.llms.anthropic.common_utils
  import numpy as np
  import pytest
  import pytest_asyncio
  from aviary.core import Message
  from lmi import (
      CommonLLMNames,
      Embeddable,
      EmbeddingModel,
      HybridEmbeddingModel,
      LiteLLMEmbeddingModel,
      LiteLLMModel,
      LLMModel,
      LLMResult,
      SparseEmbeddingModel,
  )
  from lmi.llms import rate_limited
  from lmi.utils import VCR_DEFAULT_MATCH_ON, validate_image
  from paperqa_docling import parse_pdf_to_pages as docling_parse_pdf_to_pages
  from paperqa_nemotron import parse_pdf_to_pages as nemotron_parse_pdf_to_pages
  from paperqa_pymupdf import parse_pdf_to_pages as pymupdf_parse_pdf_to_pages
  from paperqa_pypdf import parse_pdf_to_pages as pypdf_parse_pdf_to_pages
  from pydantic import ValidationError
  from pytest_subtests import SubTests
  
  from paperqa import (
      Doc,
      DocDetails,
      Docs,
      NumpyVectorStore,
      PQASession,
      QdrantVectorStore,
      Settings,
      Text,
      VectorStore,
  )
  from paperqa.clients import CrossrefProvider
  from paperqa.clients.journal_quality import JournalQualityPostProcessor
  from paperqa.core import (
      LLMContextTimeoutError,
      _map_fxn_summary,
      llm_parse_json,
      map_fxn_summary,
  )
  from paperqa.prompts import CANNOT_ANSWER_PHRASE, summary_json_multimodal_system_prompt
  from paperqa.prompts import qa_prompt as default_qa_prompt
  from paperqa.readers import PDFParserFn, chunk_pdf, parse_image, read_doc
  from paperqa.settings import (
      AnswerSettings,
      AsyncContextSerializer,
      MultimodalOptions,
      ParsingSettings,
      PromptSettings,
  )
  from paperqa.types import (
      ChunkMetadata,
      Context,
      ParsedMedia,
      ParsedMetadata,
      ParsedText,
  )
  from paperqa.utils import (
      clean_possessives,
      encode_id,
      extract_score,
      maybe_get_date,
      maybe_is_html,
      maybe_is_text,
      name_in_text,
      strings_similarity,
      strip_citations,
  )
  
  if TYPE_CHECKING:
      import vcr.request
  
  THIS_MODULE = pathlib.Path(__file__)
  
  
  @pytest_asyncio.fixture(name="docs_fixture")
  async def fixture_docs_fixture(stub_data_dir: Path) -> Docs:
      docs = Docs()
      with (stub_data_dir / "paper.pdf").open("rb") as f:
          await docs.aadd_file(
              f,
              citation="Wellawatte et al, XAI Review, 2023",  # Skip citation inference
              doi="10.1021/acs.jctc.2c01235",  # Skip DOI inference
              title="A Perspective on Explanations of Molecular Prediction Models",  # Skip title inference
          )
      return docs
  
  
  def test_encode_id() -> None:
      assert (
          encode_id("10.1056/nejmoa2119451")
          == encode_id("10.1056/NEJMOA2119451")
          == "945f1f30b11bcae6"
      )
  
  
  def test_single_author() -> None:
      text = "This was first proposed by (Smith 1999)."
      assert strip_citations(text) == "This was first proposed by ."
  
  
  def test_multiple_authors() -> None:
      text = "Recent studies (Smith et al. 1999) show that this is true."
      assert strip_citations(text) == "Recent studies  show that this is true."
  
  
  def test_multiple_citations() -> None:
      text = "As discussed by several authors (Smith et al. 1999; Johnson 2001; Lee et al. 2003)."
      assert strip_citations(text) == "As discussed by several authors ."
  
  
  def test_citations_with_pages() -> None:
      text = "This is shown in (Smith et al. 1999, p. 150)."
      assert strip_citations(text) == "This is shown in ."
  
  
  def test_citations_without_space() -> None:
      text = "Findings by(Smith et al. 1999)were significant."
      assert strip_citations(text) == "Findings bywere significant."
  
  
  def test_citations_with_commas() -> None:
      text = "The method was adopted by (Smith, 1999, 2001; Johnson, 2002)."
      assert strip_citations(text) == "The method was adopted by ."
  
  
  def test_citations_with_text() -> None:
      text = "This was noted (see Smith, 1999, for a review)."
      assert strip_citations(text) == "This was noted ."
  
  
  def test_no_citations() -> None:
      text = "There are no references in this text."
      assert strip_citations(text) == "There are no references in this text."
  
  
  def test_malformed_citations() -> None:
      text = "This is a malformed citation (Smith 199)."
      assert strip_citations(text) == "This is a malformed citation (Smith 199)."
  
  
  def test_edge_case_citations() -> None:
      text = "Edge cases like (Smith et al.1999) should be handled."
      assert strip_citations(text) == "Edge cases like  should be handled."
  
  
  def test_citations_with_special_characters() -> None:
      text = "Some names have dashes (O'Neil et al. 2000; Smith-Jones 1998)."
      assert strip_citations(text) == "Some names have dashes ."
  
  
  def test_citations_with_nonstandard_chars() -> None:
      text = (
          "In non-English languages, citations might look different (Müller et al. 1999)."
      )
      assert (
          strip_citations(text)
          == "In non-English languages, citations might look different ."
      )
  
  
  @pytest.mark.vcr
  def test_maybe_is_text() -> None:
      assert maybe_is_text("This is a test. The sample conc. was 1.0 mM (at 245 ^F)")
      assert not maybe_is_text("\\C0\\C0\\B1\x00")
      # get front page of wikipedia
      r = httpx.get(
          "https://en.wikipedia.org/wiki/National_Flag_of_Canada_Day",
          headers={
              "User-Agent": "PaperQA testing (https://github.com/Future-House/paper-qa)"
          },
      )
      assert maybe_is_text(r.text)
  
      r_ja = httpx.get(
          "https://ja.wikipedia.org/wiki/%E6%97%A5%E6%9C%AC",
          headers={
              "User-Agent": "PaperQA testing (https://github.com/Future-House/paper-qa)"
          },
      )
      assert maybe_is_text(r_ja.text)
  
      assert maybe_is_html(BytesIO(r.text.encode()))
  
      # now force it to contain lots of weird encoding
      bad_text = r.text.encode("latin1", "ignore").decode("utf-16", "ignore")
      assert not maybe_is_text(bad_text)
  
      # account for possible spaces in the text due to tables or title pages
      assert maybe_is_text("entry1                    entry2                    entry3")
  
      # Test high entropy cases
      # English text with entropy > 6
      good_text_high_entropy = string.printable * 10
      assert maybe_is_text(good_text_high_entropy, thresh=6)
  
      # English text with entropy > 8
      bad_text_high_entropy = "".join([chr(i) for i in range(10000, 10300)])
      assert not maybe_is_text(bad_text_high_entropy)
  
      # Japanese text with entropy > 8
      jp_char_ranges = [
          (0x3040, 0x309F),  # Hiragana
          (0x30A0, 0x30FF),  # Katakana
          (0x4E00, 0x9FAF),  # CJK Unified Ideographs
      ]
      jp_chars = []
      for start, end in jp_char_ranges:
          jp_chars.extend([chr(i) for i in range(start, end + 1)])
  
      random.shuffle(jp_chars)
      bad_jp_text_high_entropy = "".join(jp_chars[:400])
      assert not maybe_is_text(bad_jp_text_high_entropy)
  
  
  def test_name_in_text() -> None:
      name1 = "FooBar2022"
      name2 = "FooBar2022a"
      name3 = "FooBar20"
  
      text1 = "As mentioned by FooBar2022, this is a great paper"
      assert name_in_text(name1, text1)
      assert not name_in_text(name2, text1)
      assert not name_in_text(name3, text1)
  
      text2 = "This is great, as found by FooBar20"
      assert name_in_text(name3, text2)
      assert not name_in_text(name1, text2)
      assert not name_in_text(name2, text2)
  
      text3 = "Per previous work (FooBar2022, FooBar2022a), this is great"
      assert name_in_text(name1, text3)
      assert name_in_text(name2, text3)
      assert not name_in_text(name3, text3)
  
      text4 = "Per previous work (Foo2022, Bar2023), this is great"
      assert not name_in_text(name1, text4)
      assert not name_in_text(name2, text4)
      assert not name_in_text(name3, text4)
  
      text5 = "Per previous work (FooBar2022; FooBar2022a), this is great"
      assert name_in_text(name1, text5)
      assert name_in_text(name2, text5)
      assert not name_in_text(name3, text5)
  
      text6 = "According to FooBar2022 and Foobars, this is great"
      assert name_in_text(name1, text6)
      assert not name_in_text(name2, text6)
      assert not name_in_text(name3, text6)
  
      text7 = "As stated by FooBar2022.\n\nThis is great"
      assert name_in_text(name1, text7)
      assert not name_in_text(name2, text7)
      assert not name_in_text(name3, text7)
  
  
  def test_extract_score() -> None:
      sample = """
      The text describes an experiment where different cell subtypes,
      including colorectal cancer-associated fibroblasts, were treated with
      oxaliplatin for 12 days. The concentration of oxaliplatin used was the
      EC50 for each cell subtype, which was determined individually.
      The media were changed every 3 days to avoid complete cell death.
      The text does not provide information about the percentage of colorectal
      cancer-associated fibroblasts that typically survive at 2 weeks when cultured
      with oxaliplatin. (0/10)
      """
      assert extract_score(sample) == 0
  
      sample = """
      COVID-19 vaccinations have been shown to be effective against hospitalization
      from the Omicron and Delta variants, though effectiveness may decrease over
      time. A study found that vaccine effectiveness against hospitalization peaked
      around 82-92% after a third dose but declined to 53-77% 15+ weeks after the third
      dose, depending on age group and hospitalization definition. Stricter
      definitions of hospitalization, like requiring oxygen use or ICU admission,
      showed higher and more sustained vaccine effectiveness. 8
      """
  
      assert extract_score(sample) == 8
  
      sample = """
      Here is a 100-word summary of the text:
      The text discusses a phase 3 trial of a combined
      vector vaccine based on rAd26 and rAd5 vectors carrying the
      SARS-CoV-2 spike protein gene. The trial aimed to assess the efficacy,
      immunogenicity and safety of the vaccine against COVID-19 in adults.
      The study design was a randomized, double-blind, placebo-controlled trial
      done at 25 hospitals in Moscow, Russia. Eligible participants were 18 years
      or older with no history of COVID-19. The exclusion criteria ensured
      participants were healthy and had no contraindications for vaccination.
      The trial aimed to determine if the vaccine could safely and effectively
      provide protection against COVID-19. Relevance score: 8
      """
  
      assert extract_score(sample) == 8
  
      sample = """
      Here is a 100-word summary of the provided text: The text details
      trial procedures for a COVID-19 vaccine, including screening
      visits, observation visits to assess vital signs, PCR testing, and
      telemedicine consultations. Participants who tested positive for
      COVID-19 during screening were excluded from the trial. During the trial
      , additional PCR tests were only done when COVID-19 symptoms were reported
      . An electronic health record platform was in place to record data from
      telemedicine consultations. The text details the screening and trial
      procedures but does not provide direct evidence regarding the
      effectiveness of COVID-19 vaccinations. Score: 3/10
      """
  
      assert extract_score(sample) == 3
  
      sample = """
      Here is a 100-word summary of the text: The text discusses a
      phase 3 trial of a COVID-19 vaccine in Russia. The vaccine
      uses a heterologous prime-boost regimen, providing robust
      immune responses. The vaccine can be stored at -18°C and
      2-8°C. The study reports 91.6% efficacy against COVID-19 based on
      interim analysis of over 21,000 participants. The authors
      compare their results to other published COVID-19 vaccine
      efficacy data. They previously published safety and immunogenicity
      results from phase 1/2 trials of the same vaccine. Relevance score:
      8/10. The text provides details on the efficacy and immune response
      generated by one COVID-19 vaccine in a large phase 3 trial, which is
      relevant evidence to help answer the question regarding effectiveness
      of COVID-19 vaccinations.
      """
  
      assert extract_score(sample) == 8
  
      sample = """
      Here is a 100-word summary of the text: The text discusses the safety and
      efficacy of the BNT162b2 mRNA Covid-19 vaccine. The study found that
      the vaccine was well tolerated with mostly mild to moderate side
      effects. The vaccine was found to be highly effective against Covid-19,
      with an observed vaccine efficacy of 90.5% after the second dose.
      Severe Covid-19 cases were also reduced among vaccine recipients.
      The vaccine showed an early protective effect after the first dose
      and reached full efficacy 7 days after the second dose. The favorable
      safety and efficacy results provide evidence that the BNT162b2 vaccine
      is effective against Covid-19. The text provides data on the efficacy
      and safety results from a clinical trial of the BNT162b2 Covid-19 vaccine,
      which is highly relevant to answering the question about the effectiveness
      of Covid-19 vaccinations.
      """
  
      with pytest.raises(ValueError, match="Failed to extract score"):
          extract_score(sample)
  
      sample = """
      Introduce dynamic elements such as moving nodes or edges to create a sense of activity within
      the network. 2. Add more nodes and connections to make the network
      appear more complex and interconnected. 3. Incorporate both red and
      green colors into the network, as the current screenshot only shows
      green lines. 4. Vary the thickness of the lines to add depth and
      visual interest. 5. Implement different shades of red and green to
      create a gradient effect for a more visually appealing experience.
      6. Consider adding a background color or pattern to enhance the
      contrast and make the network stand out. 7. Introduce interactive
      elements that allow users to manipulate the network, such as
      dragging nodes or zooming in/out. 8. Use animation effects like
      pulsing or changing colors to highlight certain parts of the network
      or to show activity. 9. Add labels or markers to provide information
        about the nodes or connections, if relevant to the purpose of the
          network visualization. 10. Consider the use of algorithms that
          organize the network in a visually appealing manner, such as
          force-directed placement or hierarchical layouts. 3/10 """
  
      assert extract_score(sample) == 3
  
      sample = (
          "The text mentions a work by Shozo Yokoyama titled "
          '"Evolution of Dim-Light and Color Vision Pigments". '
          "This work, published in the Annual Review of Genomics and "
          "Human Genetics, discusses the evolution of human color vision. "
          "However, the text does not provide specific details or findings "
          "from Yokoyama's work. \n"
          "Relevance Score: 7"
      )
  
      assert extract_score(sample) == 7
  
      sample = (
          "The evolution of human color vision is "
          "closely tied to theories about the nature "
          "of light, dating back to the 17th to 19th "
          "centuries. Initially, there was no clear distinction "
          "between the properties of light, the eye and retina, "
          "and color percepts. Major figures in science attempted "
          "to resolve these issues, with physicists leading most "
          "advances in color science into the 20th century. Prior "
          "to Newton, colors were viewed as stages between black "
          "and white. Newton was the first to describe colors in "
          "a modern sense, using prisms to disperse light into "
          "a spectrum of colors. He demonstrated that each color "
          "band could not be further divided and that different "
          "colors had different refrangibility. \n"
          "Relevance Score: 9.5"
      )
  
      assert extract_score(sample) == 9
  
  
  @pytest.mark.asyncio
  async def test_chain_completion(caplog) -> None:
      caplog.set_level(level="WARNING", logger="lmi.types")
  
      s = Settings(llm="babbage-002", temperature=0.2)
      outputs = []
  
      def accum(x) -> None:
          outputs.append(x)
  
      llm = s.get_llm()
      messages = [Message(content="The duck says")]
  
      # With callbacks, we use streaming
      completion = await llm.call_single(messages=messages, callbacks=[accum])
      first_id = completion.id
      assert isinstance(first_id, UUID)
      assert completion.text
      assert completion.seconds_to_first_token > 0
      assert completion.prompt_count is not None
      assert completion.prompt_count > 0
      assert completion.completion_count is not None
      assert completion.completion_count > 0
      assert completion.model == "babbage-002"
      assert str(completion) == "".join(outputs)
      assert completion.cost > 0
      assert not caplog.records
  
      # Without callbacks, we don't use streaming
      completion = await llm.call_single(messages=messages)
      assert isinstance(completion.id, UUID)
      assert completion.id != first_id, "Expected different response ID"
      assert completion.text
      assert completion.seconds_to_first_token == 0
      assert completion.seconds_to_last_token > 0
      assert completion.prompt_count is not None
      assert completion.prompt_count > 0
      assert completion.completion_count is not None
      assert completion.completion_count > 0
      try:
          assert completion.model == "babbage-002"
          assert completion.cost > 0
      except AssertionError:
          # Account for https://github.com/BerriAI/litellm/issues/10572
          assert any(
              "Failed to calculate cost".lower() in r.message.lower()
              for r in caplog.records
          )
  
  
  @pytest.mark.skipif(os.environ.get("ANTHROPIC_API_KEY") is None, reason="No API key")
  @pytest.mark.asyncio
  async def test_anthropic_chain(stub_data_dir: Path) -> None:
      anthropic_settings = Settings(llm=CommonLLMNames.ANTHROPIC_TEST.value)
      outputs: list[str] = []
  
      def accum(x) -> None:
          outputs.append(x)
  
      llm = anthropic_settings.get_llm()
      messages = [
          Message(content="The duck says"),
      ]
      completion = await llm.call_single(
          messages=messages,
          callbacks=[accum],
      )
      assert completion.seconds_to_first_token > 0
      assert completion.prompt_count is not None
      assert completion.prompt_count > 0
      assert completion.completion_count is not None
      assert completion.completion_count > 0
      assert str(completion) == "".join(outputs)
      assert isinstance(completion.text, str)
      assert completion.cost > 0
  
      completion = await llm.call_single(
          messages=messages,
      )
      assert completion.seconds_to_first_token == 0
      assert completion.seconds_to_last_token > 0
      assert isinstance(completion.text, str)
      assert completion.cost > 0
  
      docs = Docs()
      await docs.aadd(
          stub_data_dir / "flag_day.html",
          "National Flag of Canada Day",
          settings=anthropic_settings,
      )
      result = await docs.aget_evidence(
          "What is the national flag of Canada?", settings=anthropic_settings
      )
      assert result.cost > 0
  
  
  @pytest.mark.vcr(
      match_on=[*VCR_DEFAULT_MATCH_ON, "body"]  # body is needed for /embeddings
  )
  @pytest.mark.asyncio
  async def test_docs_lifecycle(subtests: SubTests, stub_data_dir: Path) -> None:
      docs = Docs()
      assert await docs.aadd(
          stub_data_dir / "flag_day.html",
          citation='"National Flag of Canada Day." WikiMedia Foundation, 2023, Accessed now',  # Skip citation inference
          title="National Flag of Canada Day",  # Skip title inference
          dockey="test",
      )
      grav_hill_docname = await docs.aadd(
          stub_data_dir / "gravity_hill.md",
          citation='"Gravity Hill." WikiMedia Foundation, 2023, Accessed now',  # Skip citation inference
          title="Gravity hill",  # Skip title inference
      )
      assert grav_hill_docname, "Test expects successful add"
  
      with subtests.test(msg="citation-creation"):
          assert docs.docs["test"].docname == "National2023"
  
      with subtests.test(msg="text-contains"):
          await docs.aget_evidence("What is the national flag of Canada?")
          assert docs.texts_index.texts_hashes
          assert docs.texts
          assert all(t in docs.texts_index for t in docs.texts)
  
      with subtests.test(msg="delete"):
          (grav_hill_details,) = (
              d for d in docs.docs.values() if d.docname == grav_hill_docname
          )
          prior_texts_index_size = len(docs.texts_index)
          docs.delete(docname=grav_hill_docname)
          assert grav_hill_details.dockey not in docs.docs, "Details should be gone"
          assert not [
              t for t in docs.texts if t.doc == grav_hill_details
          ], "Texts should be gone"
          if len(docs.texts_index) == prior_texts_index_size:
              pytest.xfail(
                  "Per https://github.com/Future-House/paper-qa/issues/1140"
                  " this can be improved"
              )
  
      with subtests.test(msg="cleanup"):
          docs.texts_index.clear()
          assert docs.texts
          assert all(t not in docs.texts_index for t in docs.texts)
  
  
  @pytest.mark.asyncio
  async def test_evidence(stub_data_dir: Path) -> None:
      debug_settings = Settings.from_name("debug")
      debug_settings.parsing.multimodal = False
  
      docs = Docs()
      assert await docs.aadd(
          stub_data_dir / "paper.pdf",
          citation="Wellawatte et al, XAI Review, 2023",  # Skip citation inference
          doi="10.1021/acs.jctc.2c01235",  # Skip DOI inference
          title="A Perspective on Explanations of Molecular Prediction Models",  # Skip title inference
          settings=debug_settings,
      )
      assert docs.texts, "Test expects texts to be added"
      assert all(
          not t.media for t in docs.texts
      ), "Expected no media to be parsed with multimodal disabled"
  
      evidence = (
          await docs.aget_evidence(
              PQASession(question="What does XAI stand for?"),
              settings=debug_settings,
          )
      ).contexts
      assert len(evidence) >= debug_settings.answer.evidence_k
      assert len({e.context for e in evidence}) == len(
          evidence
      ), "Expected unique contexts"
      texts = {c.text for c in evidence}
      assert texts, "Below assertions require at least one text to be used"
      orig_acompletion = litellm.acompletion
      has_made_scoreless_context = False
      no_score_context_body = "MAKEUNIQUE Explainable Artificial Intelligence (XAI)"
  
      async def acompletion_that_breaks_first_context(*args, **kwargs):
          completion = await orig_acompletion(*args, **kwargs)
          nonlocal has_made_scoreless_context
          if not has_made_scoreless_context:
              assert len(completion.choices) == 1, "Test expects one choice"
              completion.choices[0].message.content = no_score_context_body
              has_made_scoreless_context = True
          return completion
  
      # Let's check we are resilient to bad context creation
      with patch.object(litellm, "acompletion", acompletion_that_breaks_first_context):
          # Let's also check we can get other evidence using the same underlying sources
          other_evidence = (
              await docs.aget_evidence(
                  PQASession(question="What is an acronym for explainable AI?"),
                  settings=debug_settings,
              )
          ).contexts
      assert all(
          c.context != no_score_context_body for c in other_evidence
      ), "Expected context without score to be replaced via retrying"
      assert texts.intersection(
          {c.text for c in other_evidence}
      ), "We should be able to reuse sources across evidence calls"
  
  
  @pytest.mark.vcr
  @pytest.mark.asyncio
  async def test_nonduplicate_contexts() -> None:
      doc1 = Doc(docname="stub1", dockey="stub1", citation="Stub 1")
      text1 = Text(name="stub1", text="I like turtles", doc=doc1)
      text2 = Text(**text1.model_dump())
      question = "What do you like?"
  
      # Prior session with a context we want to dedupe against
      session = PQASession(
          question=question,
          contexts=[
              Context(
                  question=question,
                  context=(
                      "The excerpt states 'I like turtles,'"
                      " indicating a preference for turtles."
                  ),
                  text=text1,
                  score=10,
              )
          ],
      )
  
      # This pattern of pre-populating Docs, whereas it's not the
      # intended flow, it's technically possible
      docs = Docs(texts=[text2])
      assert await docs.aadd_texts(texts=[text1], doc=doc1)
      session = await docs.aget_evidence(session)
      assert len(session.contexts) == 1, "Expected just one context"
  
  
  @pytest.mark.asyncio
  async def test_json_evidence(docs_fixture: Docs) -> None:
      settings = Settings.from_name("fast")
      settings.prompts.use_json = True
      settings.prompts.summary_json_system = (
          "Provide a summary of the relevant information"
          " that could help answer the question based on the excerpt."
          " Your summary, combined with many others,"
          " will be given to the model to generate an answer."
          " Respond with the following JSON format:"
          '\n\n{{\n  "summary": "...",\n  "author_name": "...",\n  "relevance_score": 0-10,\n}}'
          "\n\nwhere `summary` is relevant information from the text - about 100 words."
          " `author_name` specifies the author."
          " `relevance_score` is an integer 0-10 for the relevance of `summary` to the question."
          "\n\nThe excerpt may or may not contain relevant information."
          " If not, leave `summary` empty, and make `relevance_score` be 0."
      )
      orig_acompletion = litellm.acompletion
      has_made_bad_json_context = False
      bad_json_context = (  # Broken summary and relevance_score
          '{\n  "summary": "Complete of th'
          '\n  "author_name": "Sentinel value.",'
          '\n  "relevance_score": "A"\n}'
      )
  
      async def acompletion_that_breaks_first_context(*args, **kwargs):
          completion = await orig_acompletion(*args, **kwargs)
          nonlocal has_made_bad_json_context
          if not has_made_bad_json_context:
              assert len(completion.choices) == 1, "Test expects one choice"
              completion.choices[0].message.content = bad_json_context
              has_made_bad_json_context = True
          return completion
  
      # Let's check we are resilient to bad context creation
      with patch.object(litellm, "acompletion", acompletion_that_breaks_first_context):
          evidence = (
              await docs_fixture.aget_evidence(
                  PQASession(question="Who wrote this article?"),
                  settings=settings,
              )
          ).contexts
      evidence_with_authors = [
          c for c in evidence if hasattr(c, "author_name") and c.author_name
      ]
      assert evidence_with_authors
      assert all(
          "sentinel" not in c.author_name.lower() for c in evidence_with_authors
      ), "Expected broken context retrying to work"
  
  
  @pytest.mark.asyncio
  async def test_ablations(docs_fixture: Docs) -> None:
      settings = Settings()
      settings.answer.evidence_skip_summary = True
      settings.answer.evidence_retrieval = False
      contexts = (
          await docs_fixture.aget_evidence(
              "Which page is the statement 'Deep learning (DL) is advancing the boundaries of"
              " computational chemistry because it can accurately model non-linear"
              " structure-function relationships.' on?",
              settings=settings,
          )
      ).contexts
      assert (
          contexts[0].text.text.strip() == contexts[0].context
      ), "summarization not ablated"
  
      assert len(contexts) == len(docs_fixture.texts), "evidence retrieval not ablated"
  
  
  @pytest.mark.asyncio
  async def test_location_awareness(stub_data_dir: Path) -> None:
      settings = Settings(
          answer=AnswerSettings(evidence_k=3),
          prompts=PromptSettings(
              use_json=False,
              system=(
                  "Answer either N/A, a page number, or a page range."
                  " For example N/A, Page 10, or Pages 10-12."
                  " Bibliography text is always N/A."
                  " If there are titles like 1Introduction in the paper excerpt,"
                  " there's likely a PDF page concatenation here such that"
                  " Introduction is actually on page 2, not page 1."
                  " For this reason, prefer pulling page or page ranges"
                  " from the citation over paper excerpt."
                  " The citation usually starts with name2023title pages X-Y,"
                  " so respond with Pages X-Y."
              ),
              summary=(
                  "## Paper Citation\n\n{citation}\n\n## Paper Excerpt\n\n{text}"
                  "\n\n## Question\n\n{question}"
              ),
          ),
          parsing=ParsingSettings(
              # Only read in first eight pages to save CI costs/runtime
              reader_config={"chunk_chars": 5000, "overlap": 250, "page_range": (1, 8)},
          ),
      )
  
      docs = Docs()
      assert await docs.aadd(
          stub_data_dir / "paper.pdf",
          citation="Wellawatte et al, XAI Review, 2023",  # Skip citation inference
          doi="10.1021/acs.jctc.2c01235",  # Skip DOI inference
          title="A Perspective on Explanations of Molecular Prediction Models",  # Skip title inference
          settings=settings,
      )
  
      session = await docs.aget_evidence(
          "Which page or page range has the full statement (insensitive to newlines)"
          " 'Deep learning (DL) is advancing the boundaries of computational chemistry"
          " because it can accurately model non-linear structure-function relationships."
          " Applications of DL can be found in a broad spectrum spanning"
          " from quantum computing to drug discovery to materials design.' on?"
          " If this statement is not present, just answer N/A.",
          settings=settings,
      )
  
      def to_pages(value: Context) -> str:
          cxt_val = value.context.lower().split("\n")[0]
          page_range = cxt_val.split("page's")[-1].split("pages")[-1].split("page")[-1]
          return page_range.strip().removesuffix(".")  # noqa: FURB184
  
      # NOTE: scores are useless here because we didn't describe them in the prompt
      locations = [to_pages(c) for c in session.contexts]
      try:
          # 2-3 is not strictly correct, but it's feasible enough that we allow it here
          assert any(
              x in locations for x in ("2", "1-3", "1 - 3", "2-3", "2 - 3")
          ), f"correct location not found in parsed evidence {locations}"
      except AssertionError:
          if "1" not in locations:
              # Fall 2025 LLMs are not smart enough yet :/ so just allow saying page 1
              raise
  
  
  @pytest.mark.asyncio
  async def test_query(docs_fixture) -> None:
      settings = Settings(prompts={"answer_iteration_prompt": None})
      await docs_fixture.aquery("Is XAI usable in chemistry?", settings=settings)
  
  
  @pytest.mark.asyncio
  async def test_custom_context_str_fn(docs_fixture) -> None:
      async def custom_context_str_fn(  # noqa: RUF029
          settings: Settings,  # noqa: ARG001
          contexts: list[Context],  # noqa: ARG001
          question: str,  # noqa: ARG001
          pre_str: str | None = None,  # noqa: ARG001
      ) -> str:
          return "TEST OVERRIDE"
  
      assert isinstance(custom_context_str_fn, AsyncContextSerializer)
  
      settings = Settings(
          custom_context_serializer=custom_context_str_fn,
          prompts={"answer_iteration_prompt": None},
      )
  
      session = await docs_fixture.aquery(
          "Is XAI usable in chemistry?", settings=settings
      )
      assert (
          session.context == "TEST OVERRIDE"
      ), "Expected custom context string to be returned."
  
  
  @pytest.mark.asyncio
  async def test_aquery_groups_contexts_by_question(docs_fixture) -> None:
      session = PQASession(question="What is the relationship between chemistry and AI?")
  
      doc = Doc(docname="test_doc", citation="Test Doc, 2025", dockey="key1")
      text1 = Text(text="XAI is useful for molecules.", name="t1", doc=doc)
      text2 = Text(text="Drug discovery uses AI.", name="t2", doc=doc)
      text3 = Text(text="Organic chemistry is a field.", name="t3", doc=doc)
  
      session.contexts = [
          Context(
              text=text1,
              context="Explanation about XAI and molecules (Smith 1999).",
              score=6,
              question="Is XAI usable in chemistry?",
          ),
          Context(
              text=text2,
  