From 1ee5e39a5ef2f3fb2585f0863c3c2d30b569c4ce Mon Sep 17 00:00:00 2001 From: Chris Steege Date: Tue, 20 Jun 2023 18:26:47 -0500 Subject: [PATCH 1/5] Improve model organization --- mindflow/core/types/definitions/object.py | 37 --------- mindflow/core/types/mindflow_model.py | 25 +++++-- mindflow/core/types/model.py | 91 +++++++++++++++-------- 3 files changed, 81 insertions(+), 72 deletions(-) delete mode 100644 mindflow/core/types/definitions/object.py diff --git a/mindflow/core/types/definitions/object.py b/mindflow/core/types/definitions/object.py deleted file mode 100644 index 14f526d..0000000 --- a/mindflow/core/types/definitions/object.py +++ /dev/null @@ -1,37 +0,0 @@ -from enum import Enum - -from mindflow.core.types.document import Document -from mindflow.core.types.mindflow_model import MindFlowModel, MindFlowModelConfig -from mindflow.core.types.model import Model, ModelConfig -from mindflow.core.types.service import Service, ServiceConfig - - -class ObjectID(Enum): - SERVICE = "service" - SERVICE_CONFIG = "service_config" - MODEL = "model" - MODEL_CONFIG = "model_config" - MINDFLOW_MODEL_CONFIG = "mindflow_model_config" - DOCUMENT = "document" - - -class ObjectName(Enum): - SERVICE = "Service" - SERVICE_CONFIG = "Service Configurations" - MODEL = "Model" - MODEL_CONFIG = "Model Configurations" - MINDFLOW_MODEL_CONFIG = "Mindflow Model Configurations" - DOCUMENT = "Document" - - -class ObjectType(Enum): - SERVICE = Service - SERVICE_CONFIG = ServiceConfig - MODEL = Model - MODEL_CONFIG = ModelConfig - MINDFLOW_MODEL = MindFlowModel - MINDFLOW_MODEL_CONFIG = MindFlowModelConfig - DOCUMENT = Document - - -# OBJECT_UNION = Union[Service, ServiceConfig, Model, ModelConfig, MindFlowModelConfig, Document] diff --git a/mindflow/core/types/mindflow_model.py b/mindflow/core/types/mindflow_model.py index 3820cd1..2140e6a 100644 --- a/mindflow/core/types/mindflow_model.py +++ b/mindflow/core/types/mindflow_model.py @@ -1,9 +1,15 @@ import sys -from typing import Dict +from typing import Dict, Generic, TypeVar, cast +from mindflow.core.types.definitions.model import ModelID from mindflow.core.types.store_traits.static import StaticStore from mindflow.core.types.store_traits.json import JsonStore -from mindflow.core.types.model import ConfiguredModel +from mindflow.core.types.model import ( + ConfiguredModel, + ConfiguredOpenAIChatCompletionModel, + ConfiguredAnthropicChatCompletionModel, + ConfiguredOpenAITextEmbeddingModel, +) from mindflow.core.types.service import ConfiguredServices from mindflow.core.types.definitions.mind_flow_model import MindFlowModelID from mindflow.core.types.definitions.service import ( @@ -24,11 +30,13 @@ class MindFlowModelConfig(JsonStore): model: str -class ConfiguredMindFlowModel: +T = TypeVar("T", bound="ConfiguredModel") + +class ConfiguredMindFlowModel(Generic[T]): id: str # index, query, embedding name: str defaults: Dict[str, str] - model: ConfiguredModel + model: T def __init__(self, mindflow_model_id: str, configured_services: ConfiguredServices): self.id = mindflow_model_id @@ -44,7 +52,14 @@ def __init__(self, mindflow_model_id: str, configured_services: ConfiguredServic ) is None: model_id = self.get_default_model_id(mindflow_model_id, configured_services) - self.model = ConfiguredModel(model_id) + if model_id in [ModelID.GPT_3_5_TURBO.value, ModelID.GPT_4.value]: + self.model = cast(T, ConfiguredOpenAIChatCompletionModel(model_id)) + elif model_id in [ModelID.CLAUDE_INSTANT_V1.value, ModelID.CLAUDE_V1.value]: + self.model = cast(T, ConfiguredAnthropicChatCompletionModel(model_id)) + elif model_id == ModelID.TEXT_EMBEDDING_ADA_002.value: + self.model = cast(T, ConfiguredOpenAITextEmbeddingModel(model_id)) + else: + raise Exception("Unsupported model: " + model_id) def get_default_model_id( self, mindflow_model_id: str, configured_services: ConfiguredServices diff --git a/mindflow/core/types/model.py b/mindflow/core/types/model.py index 40582a9..f16ec0d 100644 --- a/mindflow/core/types/model.py +++ b/mindflow/core/types/model.py @@ -1,14 +1,12 @@ +from abc import ABC, abstractmethod import time -from typing import Optional, Union +from typing import Optional, TypeVar, Union import openai import anthropic import numpy as np from traitlets import Callable - -from mindflow.core.types.definitions.model_type import ModelType - import tiktoken from mindflow.core.types.store_traits.json import JsonStore @@ -39,8 +37,7 @@ class ModelConfig(JsonStore): id: str soft_token_limit: int - -class ConfiguredModel(Callable): +class ConfiguredModel(ABC, Callable): id: str name: str service: str @@ -80,7 +77,28 @@ def __init__(self, model_id: str): except NameError: pass - def openai_chat_completion( + @abstractmethod + def __call__(self, *args, **kwargs): + pass + + +class ConfiguredOpenAIChatCompletionModel(ConfiguredModel): + id: str + name: str + service: str + model_type: str + + tokenizer: tiktoken.Encoding + + hard_token_limit: int + token_cost: int + token_cost_unit: str + + # Config + soft_token_limit: int + api_key: str + + def __call__( self, messages: list, temperature: float = 0.0, @@ -106,7 +124,24 @@ def openai_chat_completion( return ModelError(error_message) - def anthropic_chat_completion( + +class ConfiguredAnthropicChatCompletionModel(ConfiguredModel): + id: str + name: str + service: str + model_type: str + + tokenizer: tiktoken.Encoding + + hard_token_limit: int + token_cost: int + token_cost_unit: str + + # Config + soft_token_limit: int + api_key: str + + def __call__( self, prompt: str, temperature: float = 0.0, @@ -131,7 +166,24 @@ def anthropic_chat_completion( return ModelError(error_message) - def openai_embedding(self, text: str) -> Union[np.ndarray, ModelError]: + +class ConfiguredOpenAITextEmbeddingModel(ConfiguredModel): + id: str + name: str + service: str + model_type: str + + tokenizer: tiktoken.Encoding + + hard_token_limit: int + token_cost: int + token_cost_unit: str + + # Config + soft_token_limit: int + api_key: str + + def __call__(self, text: str) -> Union[np.ndarray, ModelError]: try_count = 0 error_message = "" while try_count < 5: @@ -146,24 +198,3 @@ def openai_embedding(self, text: str) -> Union[np.ndarray, ModelError]: time.sleep(5) return ModelError(error_message) - - def __call__(self, prompt, *args, **kwargs): - service_model_mapping = { - ( - ServiceID.OPENAI.value, - ModelType.TEXT_COMPLETION.value, - ): self.openai_chat_completion, - ( - ServiceID.OPENAI.value, - ModelType.TEXT_EMBEDDING.value, - ): self.openai_embedding, - ( - ServiceID.ANTHROPIC.value, - ModelType.TEXT_COMPLETION.value, - ): self.anthropic_chat_completion, - } - if ( - func := service_model_mapping.get((self.service, self.model_type)) - ) is not None: - return func(prompt, *args, **kwargs) - raise NotImplementedError(f"Service {self.service} not implemented.") From 2ecf739f60b7f468898aaccb0372db117af2d2e6 Mon Sep 17 00:00:00 2001 From: Chris Steege Date: Tue, 20 Jun 2023 18:28:14 -0500 Subject: [PATCH 2/5] Lint --- mindflow/core/types/mindflow_model.py | 1 + mindflow/core/types/model.py | 1 + 2 files changed, 2 insertions(+) diff --git a/mindflow/core/types/mindflow_model.py b/mindflow/core/types/mindflow_model.py index 2140e6a..d55e4fc 100644 --- a/mindflow/core/types/mindflow_model.py +++ b/mindflow/core/types/mindflow_model.py @@ -32,6 +32,7 @@ class MindFlowModelConfig(JsonStore): T = TypeVar("T", bound="ConfiguredModel") + class ConfiguredMindFlowModel(Generic[T]): id: str # index, query, embedding name: str diff --git a/mindflow/core/types/model.py b/mindflow/core/types/model.py index f16ec0d..07cc4aa 100644 --- a/mindflow/core/types/model.py +++ b/mindflow/core/types/model.py @@ -37,6 +37,7 @@ class ModelConfig(JsonStore): id: str soft_token_limit: int + class ConfiguredModel(ABC, Callable): id: str name: str From 27fcc6bb5999e00b993741128f6718b373304cdb Mon Sep 17 00:00:00 2001 From: Chris Steege Date: Tue, 20 Jun 2023 18:30:55 -0500 Subject: [PATCH 3/5] Remove some unnecessary types --- mindflow/core/types/definitions/conversation.py | 8 -------- mindflow/core/types/definitions/document.py | 4 ---- .../core/types/definitions/mind_flow_model.py | 14 -------------- mindflow/core/types/definitions/model.py | 13 ------------- mindflow/core/types/definitions/service.py | 16 ---------------- 5 files changed, 55 deletions(-) diff --git a/mindflow/core/types/definitions/conversation.py b/mindflow/core/types/definitions/conversation.py index feb41f4..885c248 100644 --- a/mindflow/core/types/definitions/conversation.py +++ b/mindflow/core/types/definitions/conversation.py @@ -2,19 +2,11 @@ class ConversationParameterKey(Enum): - """ - Document argument enum - """ - ID: str = "id" MESSAGES: str = "messages" TOTAL_TOKENS: str = "total_tokens" class ConversationID(Enum): - """ - Conversation ID enum - """ - CHAT_0: str = "chat_0" CODE_GEN_0: str = "code_gen_0" diff --git a/mindflow/core/types/definitions/document.py b/mindflow/core/types/definitions/document.py index cee392a..23356ee 100644 --- a/mindflow/core/types/definitions/document.py +++ b/mindflow/core/types/definitions/document.py @@ -2,10 +2,6 @@ class DocumentType(Enum): - """ - Document type enum - """ - FILE: str = "file" SLACK: str = "slack" GITHUB: str = "github" diff --git a/mindflow/core/types/definitions/mind_flow_model.py b/mindflow/core/types/definitions/mind_flow_model.py index b4524b7..0a13635 100644 --- a/mindflow/core/types/definitions/mind_flow_model.py +++ b/mindflow/core/types/definitions/mind_flow_model.py @@ -67,17 +67,3 @@ class MindFlowModelDescription(Enum): MindFlowModelParameterKey.DESCRIPTION.value: MindFlowModelDescription.EMBEDDING.value, }, } - -# MindFlowModelUnion = Union[ -# MindFlowModelID, -# MindFlowModelDefaults, -# MindFlowModelName, -# MindFlowModelType, -# MindFlowModelDescription, -# ] - - -# def get_mind_flow_model_static( -# static: Type[MindFlowModelUnion], key: MindFlowModelUnion -# ) -> MindFlowModelUnion: -# return static.__members__[key.name] diff --git a/mindflow/core/types/definitions/model.py b/mindflow/core/types/definitions/model.py index dc7120f..cf8ca72 100644 --- a/mindflow/core/types/definitions/model.py +++ b/mindflow/core/types/definitions/model.py @@ -282,16 +282,3 @@ class ModelAnthropic(Enum): ModelParameterKey.CONFIG_DESCRIPTION.value: ModelConfigDescription.TEXT_EMBEDDING_ADA_002.value, }, } - - -# ModelUnion = Union[ -# ModelID, -# ModelParameterKey, -# ModelName, -# ModelHardTokenLimit, -# ModelDescription, -# ] - - -# def get_model_static(static: Type[ModelUnion], key: ModelUnion) -> ModelUnion: -# return static.__members__[key.name] diff --git a/mindflow/core/types/definitions/service.py b/mindflow/core/types/definitions/service.py index 7973059..f67fb49 100644 --- a/mindflow/core/types/definitions/service.py +++ b/mindflow/core/types/definitions/service.py @@ -118,19 +118,3 @@ class ServiceModel(Enum): ServiceParameterKey.API_URL.value: ServiceURL.PINECONE.value, }, } - - -# ServiceUnion = Union[ -# ServiceID, -# ServiceParameterKey, -# ServiceConfigParameterKey, -# ServiceName, -# ServiceURL, -# ServiceModel, -# ServiceModelTypeTextEmbedding, -# ServiceModelTypeTextCompletion, -# ] - - -# def get_service_static(static: Type[ServiceUnion], key: ServiceUnion) -> ServiceUnion: -# return static.__members__[key.name] From 3a64120d8696855e34969647da9bd2c473aed5dc Mon Sep 17 00:00:00 2001 From: Chris Steege Date: Tue, 20 Jun 2023 18:31:33 -0500 Subject: [PATCH 4/5] Update version --- mindflow/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindflow/__init__.py b/mindflow/__init__.py index dd9b22c..7225152 100644 --- a/mindflow/__init__.py +++ b/mindflow/__init__.py @@ -1 +1 @@ -__version__ = "0.5.1" +__version__ = "0.5.2" From ff9741d708aeb85c069a6fc6b44e3e721619b591 Mon Sep 17 00:00:00 2001 From: Chris Steege Date: Tue, 20 Jun 2023 18:34:57 -0500 Subject: [PATCH 5/5] Remove unused TypeVar import from model.py: Co-authored-by: MindFlow --- mindflow/core/types/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindflow/core/types/model.py b/mindflow/core/types/model.py index 07cc4aa..fcafc39 100644 --- a/mindflow/core/types/model.py +++ b/mindflow/core/types/model.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod import time -from typing import Optional, TypeVar, Union +from typing import Optional, Union import openai import anthropic