Add Support for Local LLama Models via LlamaCpp Integration

.env.sample

@@ -1,4 +1,4 @@
-# OPENAI or AZURE or ANTHROPIC
+# OPENAI or AZURE or ANTHROPIC or LLAMA
 ENDPOINT=OPENAI
 
 # USE when ENDPOINT=OPENAI
@@ -16,6 +16,11 @@ AZURE_OPENAI_ENDPOINT=https://<your-azure-openai-endpoint>.openai.azure.com/
 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
+
 # LangSmith
 LANGCHAIN_TRACING_V2=false
 LANGCHAIN_ENDPOINT=https://api.smith.langchain.com

gpt_all_star/core/llm.py

@@ -5,12 +5,16 @@
 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
 
 
 class LLM_TYPE(str, Enum):
     OPENAI = "OPENAI"
     AZURE = "AZURE"
     ANTHROPIC = "ANTHROPIC"
+    LLAMA = "LLAMA"
 
 
 def create_llm(llm_name: LLM_TYPE) -> BaseChatModel:
@@ -37,6 +41,13 @@ def create_llm(llm_name: LLM_TYPE) -> BaseChatModel:
             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}")
 
@@ -83,3 +94,20 @@ def _create_chat_anthropic(
         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,
+    )

tests/core/test_llm.py

@@ -2,7 +2,7 @@
 
 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
@@ -17,6 +17,18 @@ def mock_chat_openai():
         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
+
+
 def test_create_chat_openai_with_base_url(mock_openai, mock_chat_openai):
     base_url = "https://custom-openai-api.com/v1"
     _create_chat_openai(model_name="gpt-4", temperature=0.1, base_url=base_url)
@@ -40,3 +52,27 @@ def test_create_chat_openai_without_base_url(mock_openai, mock_chat_openai):
         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,
+    )