151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
import pytest
|
|
|
|
from src.notebookllama.models import (
|
|
Notebook,
|
|
)
|
|
from src.notebookllama.utils import MindMap, Node, Edge
|
|
from src.notebookllama.audio import MultiTurnConversation, ConversationTurn
|
|
from pydantic import ValidationError
|
|
|
|
|
|
def test_notebook() -> None:
|
|
n1 = Notebook(
|
|
summary="This is a summary",
|
|
questions=[
|
|
"What is the capital of Spain?",
|
|
"What is the capital of France?",
|
|
"What is the capital of Italy?",
|
|
"What is the capital of Portugal?",
|
|
"What is the capital of Germany?",
|
|
],
|
|
answers=[
|
|
"Madrid",
|
|
"Paris",
|
|
"Rome",
|
|
"Lisbon",
|
|
"Berlin",
|
|
],
|
|
highlights=["This", "is", "a", "summary"],
|
|
)
|
|
assert n1.summary == "This is a summary"
|
|
assert n1.questions[0] == "What is the capital of Spain?"
|
|
assert n1.answers[0] == "Madrid"
|
|
assert n1.highlights[0] == "This"
|
|
# Fewer answers than questions
|
|
with pytest.raises(ValidationError):
|
|
Notebook(
|
|
summary="This is a summary",
|
|
questions=[
|
|
"What is the capital of France?",
|
|
"What is the capital of Italy?",
|
|
"What is the capital of Portugal?",
|
|
"What is the capital of Germany?",
|
|
],
|
|
answers=[
|
|
"Paris",
|
|
"Rome",
|
|
"Lisbon",
|
|
],
|
|
highlights=["This", "is", "a", "summary"],
|
|
)
|
|
# Fewer highlights than required
|
|
with pytest.raises(ValidationError):
|
|
Notebook(
|
|
summary="This is a summary",
|
|
questions=[
|
|
"What is the capital of Spain?",
|
|
"What is the capital of France?",
|
|
"What is the capital of Italy?",
|
|
"What is the capital of Portugal?",
|
|
"What is the capital of Germany?",
|
|
],
|
|
answers=[
|
|
"Madrid",
|
|
"Paris",
|
|
"Rome",
|
|
"Lisbon",
|
|
"Berlin",
|
|
],
|
|
highlights=["This", "is"],
|
|
)
|
|
|
|
|
|
def test_mind_map() -> None:
|
|
m1 = MindMap(
|
|
nodes=[
|
|
Node(id="A", content="Auxin is released"),
|
|
Node(id="B", content="Travels to the roots"),
|
|
Node(id="C", content="Root cells grow"),
|
|
],
|
|
edges=[
|
|
Edge(from_id="A", to_id="B"),
|
|
Edge(from_id="A", to_id="C"),
|
|
Edge(from_id="B", to_id="C"),
|
|
],
|
|
)
|
|
assert m1.nodes[0].id == "A"
|
|
assert m1.nodes[0].content == "Auxin is released"
|
|
assert m1.edges[0].from_id == "A"
|
|
assert m1.edges[0].to_id == "B"
|
|
|
|
with pytest.raises(ValidationError):
|
|
MindMap(
|
|
nodes=[
|
|
Node(id="A", content="Auxin is released"),
|
|
Node(id="B", content="Travels to the roots"),
|
|
Node(id="C", content="Root cells grow"),
|
|
],
|
|
edges=[
|
|
Edge(from_id="A", to_id="B"),
|
|
Edge(from_id="A", to_id="D"), # "D" does not exist
|
|
Edge(from_id="B", to_id="C"),
|
|
],
|
|
)
|
|
|
|
|
|
def test_multi_turn_conversation() -> None:
|
|
turns = [
|
|
ConversationTurn(speaker="speaker1", content="Hello, who are you?"),
|
|
ConversationTurn(speaker="speaker2", content="I am very well, how about you?"),
|
|
ConversationTurn(speaker="speaker1", content="I am well too, thanks!"),
|
|
]
|
|
assert turns[0].speaker == "speaker1"
|
|
assert turns[0].content == "Hello, who are you?"
|
|
conversation = MultiTurnConversation(
|
|
conversation=turns,
|
|
)
|
|
assert isinstance(conversation.conversation, list)
|
|
assert isinstance(conversation.conversation[0], ConversationTurn)
|
|
wrong_turns = [
|
|
ConversationTurn(speaker="speaker1", content="Hello, who are you?"),
|
|
ConversationTurn(speaker="speaker2", content="I am very well, how about you?"),
|
|
]
|
|
wrong_turns1 = [
|
|
ConversationTurn(speaker="speaker2", content="Hello, who are you?"),
|
|
ConversationTurn(speaker="speaker1", content="I am very well, how about you?"),
|
|
ConversationTurn(speaker="speaker2", content="I am well too!"),
|
|
]
|
|
wrong_turns2 = [
|
|
ConversationTurn(speaker="speaker1", content="Hello, who are you?"),
|
|
ConversationTurn(speaker="speaker1", content="How is your life going?"),
|
|
ConversationTurn(
|
|
speaker="speaker2",
|
|
content="What is all this interest in me all of a sudden?!",
|
|
),
|
|
]
|
|
wrong_turns3 = [
|
|
ConversationTurn(speaker="speaker1", content="Hello, who are you?"),
|
|
ConversationTurn(speaker="speaker2", content="I'm well! But..."),
|
|
ConversationTurn(
|
|
speaker="speaker2",
|
|
content="...What is all this interest in me all of a sudden?!",
|
|
),
|
|
]
|
|
with pytest.raises(ValidationError):
|
|
MultiTurnConversation(conversation=wrong_turns)
|
|
with pytest.raises(ValidationError):
|
|
MultiTurnConversation(conversation=wrong_turns1)
|
|
with pytest.raises(ValidationError):
|
|
MultiTurnConversation(conversation=wrong_turns2)
|
|
with pytest.raises(ValidationError):
|
|
MultiTurnConversation(conversation=wrong_turns3)
|