-
Notifications
You must be signed in to change notification settings - Fork 139
/
test_chatbot.py
92 lines (76 loc) · 2.78 KB
/
test_chatbot.py
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
from deepeval.tracing import trace, TraceType
from openai import OpenAI
client = OpenAI()
class Chatbot:
def __init__(self):
pass
@trace(type=TraceType.LLM, name="OpenAI", model="gpt-4")
def llm(self, input):
response = client.chat.completions.create(
model="gpt-4",
messages=[
{
"role": "system",
"content": "You are a helpful assistant.",
},
{"role": "user", "content": input},
],
)
return response.choices[0].message.content
@trace(
type=TraceType.EMBEDDING,
name="Embedding",
model="text-embedding-ada-002",
)
def get_embedding(self, input):
response = (
client.embeddings.create(
input=input, model="text-embedding-ada-002"
)
.data[0]
.embedding
)
return response
@trace(type=TraceType.RETRIEVER, name="Retriever")
def retriever(self, input=input):
embedding = self.get_embedding(input)
# Replace this with an actual vector search that uses embedding
list_of_retrieved_nodes = ["Retrieval Node 1", "Retrieval Node 2"]
return list_of_retrieved_nodes
@trace(type=TraceType.TOOL, name="Search")
def search(self, input):
# Replace this with an actual function that searches the web
title_of_the_top_search_results = "Search Result: " + input
return title_of_the_top_search_results
@trace(type=TraceType.TOOL, name="Format")
def format(self, retrieval_nodes, input):
prompt = "You are a helpful assistant, based on the following information: \n"
for node in retrieval_nodes:
prompt += node + "\n"
prompt += "Generate an unbiased response for " + input + "."
return prompt
@trace(type=TraceType.AGENT, name="Chatbot")
def query(self, user_input=input):
top_result_title = self.search(user_input)
retrieval_results = self.retriever(top_result_title)
prompt = self.format(retrieval_results, top_result_title)
return self.llm(prompt)
import pytest
from deepeval import assert_test
from deepeval.test_case import LLMTestCase
from deepeval.metrics import HallucinationMetric
chatbot = Chatbot()
def test_hallucination():
context = [
"Be a natural-born citizen of the United States.",
"Be at least 35 years old.",
"Have been a resident of the United States for 14 years.",
]
input = "What are the requimrents to be president?"
metric = HallucinationMetric(threshold=0.8)
test_case = LLMTestCase(
input=input,
actual_output=chatbot.query(user_input=input),
context=context,
)
assert_test(test_case, [metric])