Skip to content

Commit

Permalink
create QuestionType enum
Browse files Browse the repository at this point in the history
* make `type` in PartialQuestions more strict

* apply QuestionType to the the `type` field of Conversation
  • Loading branch information
kantord authored Jan 24, 2025
1 parent ee9a590 commit 6db0657
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions src/codegate/api/v1_models.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import datetime
from enum import Enum
from typing import Any, List, Optional, Union

import pydantic
Expand Down Expand Up @@ -87,6 +88,11 @@ class QuestionAnswer(pydantic.BaseModel):
answer: Optional[ChatMessage]


class QuestionType(str, Enum):
chat = "chat"
fim = "fim"


class PartialQuestions(pydantic.BaseModel):
"""
Represents all user messages obtained from a DB row.
Expand All @@ -96,7 +102,7 @@ class PartialQuestions(pydantic.BaseModel):
timestamp: datetime.datetime
message_id: str
provider: Optional[str]
type: str
type: QuestionType


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

question_answers: List[QuestionAnswer]
provider: Optional[str]
type: str
type: QuestionType
chat_id: str
conversation_timestamp: datetime.datetime

Expand Down

0 comments on commit 6db0657

Please sign in to comment.