diff --git a/neon_data_models/models/api/mq/brainforge.py b/neon_data_models/models/api/mq/brainforge.py index 5f484a7..437cbad 100644 --- a/neon_data_models/models/api/mq/brainforge.py +++ b/neon_data_models/models/api/mq/brainforge.py @@ -28,6 +28,7 @@ from pydantic import Field from neon_data_models.models.base.contexts import MQContext +from neon_data_models.models.base.mq import MQResponse from neon_data_models.models.api.llm import BrainForgeLLM, LLMRequest, LLMResponse, LLMPersona from neon_data_models.models.api.http.brainforge import LLMGetModelsHttpResponse, LLMGetPersonasHttpRequest @@ -37,7 +38,7 @@ class LLMGetModels(MQContext): description="ID of user to get models for") -class LLMGetModelsResponse(MQContext, LLMGetModelsHttpResponse): +class LLMGetModelsResponse(MQResponse, LLMGetModelsHttpResponse): pass @@ -51,7 +52,7 @@ def model_version(self): return self.model_id.split("@", 1)[1] -class LLMGetPersonasResponse(MQContext): +class LLMGetPersonasResponse(MQResponse): model: Optional[BrainForgeLLM] = Field( "Full configuration of requested model if model is loaded and access " "is allowed, else None.") @@ -64,7 +65,7 @@ def personas(self) -> List[LLMPersona]: return self.model.personas if self.model else [] -class LLMGetInference(LLMRequest, MQContext): +class LLMGetInference(LLMRequest, MQResponse): user_id: str = Field("ID of user making the request") def as_llm_request(self): @@ -80,7 +81,7 @@ class LLMGetCompletion(LLMGetModels): description="Dictionary of kwargs to pass to OpenAI completion request") -class LLMGetCompletionResponse(MQContext): +class LLMGetCompletionResponse(MQResponse): openai_response: dict = Field( description="OpenAI ChatCompletion model") @@ -98,12 +99,12 @@ class LLMGetTokenizerChatTemplatedString(LLMGetModels): "else a single string") -class LLMGetTokenizerChatTemplatedStringResponse(MQContext): +class LLMGetTokenizerChatTemplatedStringResponse(MQResponse): prompt: Union[List[str], str] = Field( description="Prompt generated by the tokenizer") -class LLMGetInferenceResponse(LLMResponse, MQContext): +class LLMGetInferenceResponse(LLMResponse, MQResponse): pass