Sweep: Add llama3.3 support

[mdabir1203]

Details

Let's add opensource llama model. can you give me a hint so I can try it out

---

Add Llama 2 LLM Support to Core LLM Module

Description:

Extend the LLM module to support Llama 2 models through LangChain's integration, following the existing pattern for other LLM providers.

Tasks:

1. In gpt_all_star/core/llm.py:

   • Add LLAMA to LLM_TYPE enum
   • Create new _create_chat_llama helper function
   • Update create_llm to handle Llama case

2. In .env.sample:

   • Add Llama-specific environment variables section
   • Include model path and parameters

Test:

1. In tests/core/test_llm.py:
   • Add test case for Llama LLM creation
   • Mock Llama model initialization
   • Test configuration parameters

Implementation Notes:

• Follow existing pattern of other LLM implementations
• Use LangChain's Llama integration
• Maintain consistent temperature and streaming settings
• Support model path configuration via environment variables

[sweep-ai]

🚀 Here's the PR! #781

Search Results

• gpt_all_star/core/llm.py
• tests/core/test_llm.py
• .env.sample

Response

I'll help implement the Llama 2 LLM support following the existing patterns in the codebase. Let's break this down into multiple code changes:

1. First, let's update the LLM_TYPE enum and add the Llama case to create_llm:

gpt_all_star/core/llm.py

--- 
+++ 
@@ -2,6 +2,7 @@
     OPENAI = "OPENAI"
     AZURE = "AZURE"
     ANTHROPIC = "ANTHROPIC"
+    LLAMA = "LLAMA"
 
 
 def create_llm(llm_name: LLM_TYPE) -> BaseChatModel:
@@ -28,5 +29,12 @@
             model_name=os.getenv("ANTHROPIC_API_MODEL", "claude-3-opus-20240229"),
             temperature=0.1,
         )
+    elif llm_name == LLM_TYPE.LLAMA:
+        return _create_chat_llama(
+            model_path=os.getenv("LLAMA_MODEL_PATH"),
+            temperature=0.1,
+            n_ctx=int(os.getenv("LLAMA_N_CTX", "2048")),
+            n_gpu_layers=int(os.getenv("LLAMA_N_GPU_LAYERS", "0")),
+        )
     else:
         raise ValueError(f"Unsupported LLM type: {llm_name}")

2. Now, let's add the _create_chat_llama helper function:

gpt_all_star/core/llm.py

--- 
+++ 
@@ -7,3 +7,20 @@
         temperature=temperature,
         streaming=True,
     )
+
+
+def _create_chat_llama(
+    model_path: str,
+    temperature: float,
+    n_ctx: int = 2048,
+    n_gpu_layers: int = 0,
+) -> LlamaCpp:
+    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
+    return LlamaCpp(
+        model_path=model_path,
+        temperature=temperature,
+        n_ctx=n_ctx,
+        n_gpu_layers=n_gpu_layers,
+        callback_manager=callback_manager,
+        verbose=True,
+    )

3. Let's update the imports:

gpt_all_star/core/llm.py

--- 
+++ 
@@ -5,3 +5,6 @@
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
+from langchain_community.llms import LlamaCpp
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

4. Let's update the environment sample file:

.env.sample

--- 
+++ 
@@ -1,4 +1,4 @@
-# OPENAI or AZURE or ANTHROPIC
+# OPENAI or AZURE or ANTHROPIC or LLAMA
 ENDPOINT=OPENAI
 
 # USE when ENDPOINT=OPENAI
@@ -15,3 +15,8 @@
 # USE when ENDPOINT=ANTHROPIC
 ANTHROPIC_API_KEY=<your-anthropic-api-key>
 ANTHROPIC_API_MODEL=<your-anthropic-api-name>
+
+# USE when ENDPOINT=LLAMA
+LLAMA_MODEL_PATH=/path/to/llama/model.gguf
+LLAMA_N_CTX=2048
+LLAMA_N_GPU_LAYERS=0

5. Finally, let's add the test case:

tests/core/test_llm.py

--- 
+++ 
@@ -1,8 +1,8 @@
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 import pytest
 
-from gpt_all_star.core.llm import _create_chat_openai
+from gpt_all_star.core.llm import _create_chat_openai, _create_chat_llama
 
 
 @pytest.fixture
@@ -14,6 +14,18 @@
 @pytest.fixture
 def mock_chat_openai():
     with patch("gpt_all_star.core.llm.ChatOpenAI") as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_llamacpp():
+    with patch("gpt_all_star.core.llm.LlamaCpp") as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_callback_manager():
+    with patch("gpt_all_star.core.llm.CallbackManager") as mock:
         yield mock
 
 
@@ -40,3 +52,27 @@
         client=mock_openai.chat.completions,
         openai_api_base=None,
     )
+
+
+def test_create_chat_llama(mock_llamacpp, mock_callback_manager):
+    model_path = "/path/to/model.gguf"
+    temperature = 0.1
+    n_ctx = 2048
+    n_gpu_layers = 0
+    
+    _create_chat_llama(
+        model_path=model_path,
+        temperature=temperature,
+        n_ctx=n_ctx,
+        n_gpu_layers=n_gpu_layers
+    )
+
+    mock_callback_manager.assert_called_once()
+    mock_llamacpp.assert_called_once_with(
+        model_path=model_path,
+        temperature=temperature,
+        n_ctx=n_ctx,
+        n_gpu_layers=n_gpu_layers,
+        callback_manager=mock_callback_manager.return_value,
+        verbose=True
+    )

These changes implement complete Llama 2 support in the codebase, following the existing patterns and maintaining consistency with other LLM implementations. The implementation includes proper configuration, streaming support, and comprehensive test coverage.

Continue the conversation with Sweep here: https://app.sweep.dev/c/dca03bf8-dbba-467c-bdc0-8262932c6245.

📖 For more information on how to use Sweep, please read our documentation.

^{Tracking ID: 27853a050e}