-
Notifications
You must be signed in to change notification settings - Fork 335
/
answer.py
65 lines (55 loc) · 2 KB
/
answer.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
import json
from typing import List, Optional
from .source_document import SourceDocument
class Answer:
def __init__(
self,
question: str,
answer: str,
source_documents: List[SourceDocument] = [],
prompt_tokens: Optional[int] = 0,
completion_tokens: Optional[int] = 0,
):
self.question = question
self.answer = answer
self.source_documents = source_documents
self.prompt_tokens = prompt_tokens
self.completion_tokens = completion_tokens
def __eq__(self, value: object) -> bool:
if not isinstance(value, Answer):
return False
return (
self.question == value.question
and self.answer == value.answer
and self.source_documents == value.source_documents
and self.prompt_tokens == value.prompt_tokens
and self.completion_tokens == value.completion_tokens
)
def to_json(self):
return json.dumps(self, cls=AnswerEncoder)
@classmethod
def from_json(cls, json_string):
return json.loads(json_string, cls=AnswerDecoder)
class AnswerEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, Answer):
return {
"question": obj.question,
"answer": obj.answer,
"source_documents": [doc.to_json() for doc in obj.source_documents],
"prompt_tokens": obj.prompt_tokens,
"completion_tokens": obj.completion_tokens,
}
return super().default(obj)
class AnswerDecoder(json.JSONDecoder):
def decode(self, s, **kwargs):
obj = super().decode(s, **kwargs)
return Answer(
question=obj["question"],
answer=obj["answer"],
source_documents=[
SourceDocument.from_json(doc) for doc in obj["source_documents"]
],
prompt_tokens=obj["prompt_tokens"],
completion_tokens=obj["completion_tokens"],
)