Skip to content

Commit

Permalink
Exposed 'make_status' helper function, and integrated into tests
Browse files Browse the repository at this point in the history
  • Loading branch information
jamesbraza committed Sep 12, 2024
1 parent 9a2761e commit 17c06a4
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 13 deletions.
32 changes: 27 additions & 5 deletions paperqa/agents/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,16 @@
logger = logging.getLogger(__name__)


def make_status(
total_paper_count: int, relevant_paper_count: int, evidence_count: int, cost: float
) -> str:
return (
f"Status: Paper Count={total_paper_count}"
f" | Relevant Papers={relevant_paper_count} | Current Evidence={evidence_count}"
f" | Current Cost=${cost:.4f}"
)


class EnvironmentState(BaseModel):
"""State here contains documents and answer being populated."""

Expand All @@ -35,11 +45,23 @@ class EnvironmentState(BaseModel):
@computed_field # type: ignore[prop-decorator]
@property
def status(self) -> str:
return (
f"Status: Paper Count={len(self.docs.docs)} | Relevant Papers="
f"{len({c.text.doc.dockey for c in self.answer.contexts if c.score > self.RELEVANT_SCORE_CUTOFF})}"
f" | Current Evidence={len([c for c in self.answer.contexts if c.score > self.RELEVANT_SCORE_CUTOFF])}"
f" | Current Cost=${self.answer.cost:.4f}"
return make_status(
total_paper_count=len(self.docs.docs),
relevant_paper_count=len(
{
c.text.doc.dockey
for c in self.answer.contexts
if c.score > self.RELEVANT_SCORE_CUTOFF
}
),
evidence_count=len(
[
c
for c in self.answer.contexts
if c.score > self.RELEVANT_SCORE_CUTOFF
]
),
cost=self.answer.cost,
)


Expand Down
17 changes: 9 additions & 8 deletions tests/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
GatherEvidence,
GenerateAnswer,
PaperSearch,
make_status,
)
from paperqa.docs import Docs
from paperqa.settings import AgentSettings, Settings
Expand Down Expand Up @@ -121,14 +122,14 @@ async def test_successful_memory_agent() -> None:
" ampicillin\n\nThe gen_answer tool output is visible to the user, so"
" you do not need to restate the answer and can simply terminate if"
" the answer looks sufficient. The current status of"
" evidence/papers/cost is Status: Paper Count=0 | Relevant Papers=0 |"
" Current Evidence=0 | Current Cost=$0.00\n\nTool request message ''"
" evidence/papers/cost is Status: "
f"{make_status(total_paper_count=0, relevant_paper_count=0, evidence_count=0, cost=0.0)}" # Started 0
"\n\nTool request message ''"
f" for tool call: {memory_id} with content"
" 'paper_search(query='Acinetobacter lwoffii antibiotic resistance"
" evolution', min_year='None', max_year='None')'\n\nTool response"
" message 'Status: Paper Count=11 | Relevant Papers=0 | Current"
f" Evidence=0 | Current Cost=$0.0000' for tool call ID {memory_id} of"
" tool 'paper_search'"
" evolution', min_year='None', max_year='None')'\n\nTool response message '"
f"{make_status(total_paper_count=11, relevant_paper_count=0, evidence_count=0, cost=0.0)}" # Found 11
f"' for tool call ID {memory_id} of tool 'paper_search'"
),
input=(
"Use the tools to answer the question: Q: Acinetobacter lwoffii has"
Expand All @@ -138,8 +139,8 @@ async def test_successful_memory_agent() -> None:
" ampicillin\n\nThe gen_answer tool output is visible to the user, so"
" you do not need to restate the answer and can simply terminate if the"
" answer looks sufficient. The current status of evidence/papers/cost"
" is Status: Paper Count=0 | Relevant Papers=0 | Current Evidence=0 |"
" Current Cost=$0.0000"
" is Status: "
f"{make_status(total_paper_count=0, relevant_paper_count=0, evidence_count=0, cost=0.0)}"
),
output=(
f"Tool request message '' for tool call: {memory_id} with content"
Expand Down

0 comments on commit 17c06a4

Please sign in to comment.