Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 6db0657

Browse files
authored
create QuestionType enum
* make `type` in PartialQuestions more strict * apply QuestionType to the the `type` field of Conversation
1 parent ee9a590 commit 6db0657

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

src/codegate/api/v1_models.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import datetime
2+
from enum import Enum
23
from typing import Any, List, Optional, Union
34

45
import pydantic
@@ -87,6 +88,11 @@ class QuestionAnswer(pydantic.BaseModel):
8788
answer: Optional[ChatMessage]
8889

8990

91+
class QuestionType(str, Enum):
92+
chat = "chat"
93+
fim = "fim"
94+
95+
9096
class PartialQuestions(pydantic.BaseModel):
9197
"""
9298
Represents all user messages obtained from a DB row.
@@ -96,7 +102,7 @@ class PartialQuestions(pydantic.BaseModel):
96102
timestamp: datetime.datetime
97103
message_id: str
98104
provider: Optional[str]
99-
type: str
105+
type: QuestionType
100106

101107

102108
class PartialQuestionAnswer(pydantic.BaseModel):
@@ -115,7 +121,7 @@ class Conversation(pydantic.BaseModel):
115121

116122
question_answers: List[QuestionAnswer]
117123
provider: Optional[str]
118-
type: str
124+
type: QuestionType
119125
chat_id: str
120126
conversation_timestamp: datetime.datetime
121127

0 commit comments

Comments
 (0)