1 files changed, 14 insertions, 12 deletions

diff --git a/g4f/Provider/Llama2.py b/g4f/Provider/Llama2.py
index 17969621..d1f8e194 100644
--- a/g4f/Provider/Llama2.py
+++ b/g4f/Provider/Llama2.py
@@ -3,18 +3,24 @@ from __future__ import annotations
 from aiohttp import ClientSession
 
 from ..typing import AsyncResult, Messages
-from .base_provider import AsyncGeneratorProvider
+from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
 
-models = {
-    "meta-llama/Llama-2-7b-chat-hf": "meta/llama-2-7b-chat",
-    "meta-llama/Llama-2-13b-chat-hf": "meta/llama-2-13b-chat",
-    "meta-llama/Llama-2-70b-chat-hf": "meta/llama-2-70b-chat",
-}
 
-class Llama2(AsyncGeneratorProvider):
+class Llama2(AsyncGeneratorProvider, ProviderModelMixin):
     url = "https://www.llama2.ai"
     working = True
     supports_message_history = True
+    default_model = "meta/llama-2-70b-chat"
+    models = [
+        "meta/llama-2-7b-chat",
+        "meta/llama-2-13b-chat",
+        "meta/llama-2-70b-chat",
+    ]
+    model_aliases = {
+        "meta-llama/Llama-2-7b-chat-hf": "meta/llama-2-7b-chat",
+        "meta-llama/Llama-2-13b-chat-hf": "meta/llama-2-13b-chat",
+        "meta-llama/Llama-2-70b-chat-hf": "meta/llama-2-70b-chat",
+    }
 
     @classmethod
     async def create_async_generator(
@@ -24,10 +30,6 @@ class Llama2(AsyncGeneratorProvider):
         proxy: str = None,
         **kwargs
     ) -> AsyncResult:
-        if not model:
-            model = "meta/llama-2-70b-chat"
-        elif model in models:
-            model = models[model]
         headers = {
             "User-Agent": "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/118.0",
             "Accept": "*/*",
@@ -48,7 +50,7 @@ class Llama2(AsyncGeneratorProvider):
             prompt = format_prompt(messages)
             data = {
                 "prompt": prompt,
-                "model": model,
+                "model": cls.get_model(model),
                 "systemPrompt": kwargs.get("system_message", "You are a helpful assistant."),
                 "temperature": kwargs.get("temperature", 0.75),
                 "topP": kwargs.get("top_p", 0.9),