-
Notifications
You must be signed in to change notification settings - Fork 11
Expand file tree
/
Copy pathtest_agents.py
More file actions
136 lines (114 loc) · 4.22 KB
/
Copy pathtest_agents.py
File metadata and controls
136 lines (114 loc) · 4.22 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
from __future__ import annotations
import asyncio
import pytest
from contextrie import (
ComposerAgent,
DocumentSource,
IndexingAgent,
JudgeAgent,
JudgeTask,
ListSource,
Metadata,
)
class FakeModel:
def __init__(self) -> None:
self.prompts: list[tuple[str, type]] = []
async def generate_object(self, *, prompt: str, schema: type):
self.prompts.append((prompt, schema))
if schema.__name__ == "Metadata":
return {
"title": "Generated title",
"description": "Generated description.",
"keypoints": ["first", "second", "third"],
}
if schema.__name__ == "JudgeDecision":
return {"score": 0.75, "reason": "Matches the task."}
if schema.__name__ == "ComposerOutput":
return {
"source_ids": ["architecture", "missing"],
"context": " The architecture source explains the indexing flow. ",
}
raise AssertionError(f"Unexpected schema: {schema.__name__}")
def test_indexing_agent_assigns_metadata() -> None:
agent = IndexingAgent(FakeModel())
source = DocumentSource("architecture", None, "Indexing content")
indexed = asyncio.run(agent.add(source).run())
assert indexed == [source]
assert source.metadata is not None
assert source.metadata.title == "Generated title"
def test_judge_agent_handles_shallow_and_deep_sources() -> None:
model = FakeModel()
shallow = DocumentSource(
"architecture",
Metadata(
title="Architecture",
description="System design notes.",
keypoints=["indexing", "judgment", "composition"],
),
"Architecture content",
)
deep = ListSource("roadmap", None, ["alpha", "beta", "gamma"])
results = asyncio.run(
JudgeAgent(model)
.from_sources(shallow)
.from_deep(deep)
.run(JudgeTask(objective="response", input="What matters?"))
)
assert results["architecture"].score == pytest.approx(0.75)
assert results["roadmap"].reason == "Matches the task."
prompt_text = "\n".join(prompt for prompt, _ in model.prompts)
assert "Source title: Architecture" in prompt_text
assert "Source content:\nalpha\nbeta\ngamma" in prompt_text
def test_composer_agent_returns_trimmed_context() -> None:
model = FakeModel()
source = DocumentSource(
"architecture",
Metadata(
title="Architecture",
description="System design notes.",
keypoints=["indexing", "judgment", "composition"],
),
"Architecture content",
path="core/index.ts",
)
context = asyncio.run(
ComposerAgent(model)
.from_sources(
{
"architecture": {
"source": source,
"decision": {"score": 0.8, "reason": "Useful for the answer."},
}
}
)
.run(JudgeTask(objective="response", input="What matters?"))
)
assert context == "The architecture source explains the indexing flow."
def test_composer_agent_rejects_invalid_selected_source_ids() -> None:
class InvalidComposerModel(FakeModel):
async def generate_object(self, *, prompt: str, schema: type):
if schema.__name__ == "ComposerOutput":
return {"source_ids": ["missing"], "context": "Invalid"}
return await super().generate_object(prompt=prompt, schema=schema)
source = DocumentSource(
"architecture",
Metadata(
title="Architecture",
description="System design notes.",
keypoints=["indexing", "judgment", "composition"],
),
"Architecture content",
)
with pytest.raises(RuntimeError, match="did not select any valid source IDs"):
asyncio.run(
ComposerAgent(InvalidComposerModel())
.from_sources(
{
"architecture": {
"source": source,
"decision": {"score": 0.8, "reason": "Useful for the answer."},
}
}
)
.run(JudgeTask(objective="response", input="What matters?"))
)