From 7cfea2a480cb9eb4e0f6383150fc4d1d3cb26640 Mon Sep 17 00:00:00 2001 From: Aarushi Date: Tue, 8 Oct 2024 12:36:45 +0100 Subject: [PATCH] wip --- .../backend/backend/blocks/llm.py | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/autogpt_platform/backend/backend/blocks/llm.py b/autogpt_platform/backend/backend/blocks/llm.py index 1403baaffb..f2f4f9557e 100644 --- a/autogpt_platform/backend/backend/blocks/llm.py +++ b/autogpt_platform/backend/backend/blocks/llm.py @@ -53,14 +53,31 @@ class LlmModel(str, Enum): LLAMA3_1_405B = "llama-3.1-405b-reasoning" LLAMA3_1_70B = "llama-3.1-70b-versatile" LLAMA3_1_8B = "llama-3.1-8b-instant" + + @classmethod + def _missing_(cls, value): + if settings.config.app_env == AppEnvironment.LOCAL: + if value == "llama3": + return cls.OLLAMA_LLAMA3_8B + elif value == "llama3.1:405b": + return cls.OLLAMA_LLAMA3_405B + return None + + @property + def metadata(self) -> "ModelMetadata": + return MODEL_METADATA[self] + + @classmethod + def all_models(cls) -> List["LlmModel"]: + models = list(cls) + if settings.config.app_env == AppEnvironment.LOCAL: + models.extend([cls.OLLAMA_LLAMA3_8B, cls.OLLAMA_LLAMA3_405B]) + return models + # Ollama models OLLAMA_LLAMA3_8B = "llama3" OLLAMA_LLAMA3_405B = "llama3.1:405b" - @property - def metadata(self) -> ModelMetadata: - return MODEL_METADATA[self] - MODEL_METADATA = { LlmModel.O1_PREVIEW: ModelMetadata("openai", 32000, cost_factor=60), @@ -90,7 +107,7 @@ if settings.config.app_env == AppEnvironment.LOCAL: } ) -for model in LlmModel: +for model in LlmModel.all_models(): if model not in MODEL_METADATA: raise ValueError(f"Missing MODEL_METADATA metadata for model: {model}")