Kang Suhyun commited on
Commit
3230abf
1 Parent(s): d962d72

[#90] Separate history into different tables based on category (#91)

Browse files

This change separates the history into different tables based on their category (summarize or translation).

Files changed (1) hide show
  1. response.py +12 -4
response.py CHANGED
@@ -15,8 +15,16 @@ from model import Model
15
  from model import supported_models
16
 
17
 
18
- def create_history(model_name: str, instruction: str, prompt: str,
19
- response: str):
 
 
 
 
 
 
 
 
20
  doc_id = uuid4().hex
21
 
22
  doc = {
@@ -28,7 +36,7 @@ def create_history(model_name: str, instruction: str, prompt: str,
28
  "timestamp": firestore.SERVER_TIMESTAMP
29
  }
30
 
31
- doc_ref = db.collection("arena-history").document(doc_id)
32
  doc_ref.set(doc)
33
 
34
 
@@ -70,7 +78,7 @@ def get_responses(prompt: str, category: str, source_lang: str,
70
  "role": "user",
71
  "content": prompt
72
  }])
73
- create_history(model.name, instruction, prompt, response)
74
  responses.append(response)
75
 
76
  # TODO(#1): Narrow down the exception type.
 
15
  from model import supported_models
16
 
17
 
18
+ def get_history_collection(category: str):
19
+ if category == Category.SUMMARIZE.value:
20
+ return db.collection("arena-summarization-history")
21
+
22
+ if category == Category.TRANSLATE.value:
23
+ return db.collection("arena-translation-history")
24
+
25
+
26
+ def create_history(category: str, model_name: str, instruction: str,
27
+ prompt: str, response: str):
28
  doc_id = uuid4().hex
29
 
30
  doc = {
 
36
  "timestamp": firestore.SERVER_TIMESTAMP
37
  }
38
 
39
+ doc_ref = get_history_collection(category).document(doc_id)
40
  doc_ref.set(doc)
41
 
42
 
 
78
  "role": "user",
79
  "content": prompt
80
  }])
81
+ create_history(category, model.name, instruction, prompt, response)
82
  responses.append(response)
83
 
84
  # TODO(#1): Narrow down the exception type.