From 8329b1893e584681e049c1e760b38e8415feb0af Mon Sep 17 00:00:00 2001 From: Daniel Kantor Date: Fri, 24 Jan 2025 11:09:44 +0100 Subject: [PATCH 1/2] make `type` in PartialQuestions more strict --- src/codegate/api/v1_models.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index 9b23f74e..ea52fe34 100644 --- a/src/codegate/api/v1_models.py +++ b/src/codegate/api/v1_models.py @@ -1,4 +1,5 @@ import datetime +from enum import Enum from typing import Any, List, Optional, Union import pydantic @@ -87,6 +88,11 @@ class QuestionAnswer(pydantic.BaseModel): answer: Optional[ChatMessage] +class PartialQuestionsType(str, Enum): + chat = "chat" + fim = "fim" + + class PartialQuestions(pydantic.BaseModel): """ Represents all user messages obtained from a DB row. @@ -96,7 +102,7 @@ class PartialQuestions(pydantic.BaseModel): timestamp: datetime.datetime message_id: str provider: Optional[str] - type: str + type: PartialQuestionsType class PartialQuestionAnswer(pydantic.BaseModel): From 146fc14d7fe07752baa53c3eb64c7ae5d7d5c006 Mon Sep 17 00:00:00 2001 From: Daniel Kantor Date: Fri, 24 Jan 2025 11:19:17 +0100 Subject: [PATCH 2/2] apply QuestionType to the the `type` field of Conversation --- src/codegate/api/v1_models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index ea52fe34..2eba1d87 100644 --- a/src/codegate/api/v1_models.py +++ b/src/codegate/api/v1_models.py @@ -88,7 +88,7 @@ class QuestionAnswer(pydantic.BaseModel): answer: Optional[ChatMessage] -class PartialQuestionsType(str, Enum): +class QuestionType(str, Enum): chat = "chat" fim = "fim" @@ -102,7 +102,7 @@ class PartialQuestions(pydantic.BaseModel): timestamp: datetime.datetime message_id: str provider: Optional[str] - type: PartialQuestionsType + type: QuestionType class PartialQuestionAnswer(pydantic.BaseModel): @@ -121,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