From 55bf516f802a07efde8adddec5b926743da04d37 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Thu, 19 Dec 2024 14:27:55 -0800 Subject: [PATCH] Allow persona definition with `None` system prompt --- neon_data_models/models/api/llm.py | 5 ++--- tests/models/api/test_llm.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/neon_data_models/models/api/llm.py b/neon_data_models/models/api/llm.py index ada21ba..099b807 100644 --- a/neon_data_models/models/api/llm.py +++ b/neon_data_models/models/api/llm.py @@ -37,7 +37,7 @@ class LLMPersona(BaseModel): name: str = Field(description="Unique name for this persona") description: Optional[str] = Field( None, description="Human-readable description of this persona") - system_prompt: str = Field( + system_prompt: Optional[str] = Field( None, description="System prompt associated with this persona. " "If None, `description` will be used.") enabled: bool = Field( @@ -48,8 +48,7 @@ class LLMPersona(BaseModel): @model_validator(mode='after') def validate_request(self): - assert any(x is not None for x in (self.description, self.system_prompt)) - if self.system_prompt is None: + if self.system_prompt is None and self.description: self.system_prompt = self.description return self diff --git a/tests/models/api/test_llm.py b/tests/models/api/test_llm.py index acf5f2b..cc9fc96 100644 --- a/tests/models/api/test_llm.py +++ b/tests/models/api/test_llm.py @@ -51,7 +51,7 @@ def test_llm_persona(self): "A customized bot for something") with self.assertRaises(ValidationError): - LLMPersona(name="underdefined persona") + LLMPersona(description="underdefined persona") def test_llm_request(self): from neon_data_models.models.api.llm import LLMRequest, LLMPersona