74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
|
|
|
|
|
|
class TestCodeExecution:
|
|
def test_create_successful_execution(self):
|
|
execution = CodeExecution(
|
|
code="print('hello')",
|
|
stdout="hello\n",
|
|
stderr="",
|
|
success=True,
|
|
)
|
|
assert execution.code == "print('hello')"
|
|
assert execution.stdout == "hello\n"
|
|
assert execution.stderr == ""
|
|
assert execution.success is True
|
|
|
|
def test_create_failed_execution(self):
|
|
execution = CodeExecution(
|
|
code="1/0",
|
|
stdout="",
|
|
stderr="ZeroDivisionError: division by zero",
|
|
success=False,
|
|
)
|
|
assert execution.success is False
|
|
assert "ZeroDivisionError" in execution.stderr
|
|
|
|
|
|
class TestRLMResult:
|
|
def test_create_result_with_answer_only(self):
|
|
result = RLMResult(answer="The answer is 42")
|
|
assert result.answer == "The answer is 42"
|
|
assert result.citations == []
|
|
assert result.code_executions == []
|
|
|
|
def test_create_result_with_code_executions(self):
|
|
executions = [
|
|
CodeExecution(
|
|
code="x = 1 + 1",
|
|
stdout="",
|
|
stderr="",
|
|
success=True,
|
|
),
|
|
CodeExecution(
|
|
code="print(x)",
|
|
stdout="2\n",
|
|
stderr="",
|
|
success=True,
|
|
),
|
|
]
|
|
result = RLMResult(
|
|
answer="x equals 2",
|
|
code_executions=executions,
|
|
)
|
|
assert len(result.code_executions) == 2
|
|
assert result.code_executions[1].stdout == "2\n"
|
|
|
|
def test_create_result_with_citations(self):
|
|
from haiku.rag.agents.research.models import Citation
|
|
|
|
citations = [
|
|
Citation(
|
|
document_id="doc1",
|
|
chunk_id="chunk1",
|
|
document_uri="file://test.pdf",
|
|
document_title="Test Doc",
|
|
content="Some content",
|
|
)
|
|
]
|
|
result = RLMResult(
|
|
answer="Found in Test Doc",
|
|
citations=citations,
|
|
)
|
|
assert len(result.citations) == 1
|
|
assert result.citations[0].document_title == "Test Doc"
|