Feat: Extend available LLMs

.env.example

@@ -16,6 +16,12 @@ OPENAI_API_KEY= # OpenAI API key
 OPENAI_ORGANIZATION= # OpenAI organization ID
 ENABLE_LLM_CACHE=false # Set to true to enable LLM cache in redis
 
+#############################################
+# Other LLM API keys
+#############################################
+# ANTHROPIC_API_KEY = "<dummy>" # must be set to real key if using an Anthropic LLM
+# GOOGLE_API_KEY = "<dummy>" # must be set to real key if using a Google LLM
+
 #############################################
 # Authentication variables
 #############################################

backend/app/app/core/config.py

@@ -19,6 +19,8 @@ class Settings(BaseSettings):
  OPENAI_API_KEY: str
  OPENAI_ORGANIZATION: Optional[str] = None
  OPENAI_API_BASE: Optional[str] = None
+ ANTHROPIC_API_KEY: Optional[str] = None
+ GOOGLE_API_KEY: Optional[str] = None
  DATABASE_USER: str
  DATABASE_PASSWORD: str
  DATABASE_HOST: str

backend/app/app/schemas/tool_schema.py

@@ -10,6 +10,9 @@
  "gpt-3.5-turbo",
  "azure-4-32k",
  "azure-3.5",
+ "claude-3-opus",
+ "claude-3.5-sonnet",
+ "claude-3-sonnet",
 ]
 
 

backend/app/app/services/chat_agent/helpers/llm.py

@@ -7,6 +7,7 @@
 
 import tiktoken
 from langchain.base_language import BaseLanguageModel
+from langchain_anthropic import ChatAnthropic
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 
 from app.core.config import settings
@@ -58,6 +59,27 @@ def get_llm(
  openai_api_key=api_key if api_key is not None else settings.OPENAI_API_KEY,
  streaming=True,
  )
+ case "claude-3-opus":
+ return ChatAnthropic(
+ temperature=0,
+ model_name="claude-3-opus-20240229",
+ anthropic_api_key=settings.ANTHROPIC_API_KEY,
+ streaming=True,
+ )
+ case "claude-3.5-sonnet":
+ return ChatAnthropic(
+ temperature=0,
+ model_name="claude-3-5-sonnet-20240620",
+ anthropic_api_key=settings.ANTHROPIC_API_KEY,
+ streaming=True,
+ )
+ case "claude-3-sonnet":
+ return ChatAnthropic(
+ temperature=0,
+ model_name="claude-3-sonnet-20240229",
+ anthropic_api_key=settings.ANTHROPIC_API_KEY,
+ streaming=True,
+ )
  # If an exact match is not confirmed, this last case will be used if provided
  case _:
  logger.warning(f"LLM {llm} not found, using default LLM")

backend/app/app/services/chat_agent/meta_agent.py

@@ -34,8 +34,10 @@ def get_conv_token_buffer_memory(
  ConversationTokenBufferMemory: The ConversationTokenBufferMemory object.
  """
  agent_config = get_agent_config()
+
+ # We use gpt-4 in ConversationTokenBufferMemory to standardize tokenization
  llm = get_llm(
- agent_config.common.llm,
+ "gpt-4",
  api_key=api_key,
  )
  chat_history = ChatMessageHistory()

backend/app/pyproject.toml

@@ -71,6 +71,7 @@ tiktoken = "^0.5.2"
 watchfiles = "^0.18.1"
 wheel = "^0.40.0"
 isort = "^5.13.2"
+langchain-anthropic = "0.1.4"
 
 [tool.poetry.group.dev]
 optional = true

docs/docusaurus/docs/advanced/llms.md

@@ -0,0 +1,26 @@
+## Configuring Different LLMs with AgentKit
+
+While the default implementation of AgentKit uses GPT-4, it is fairly straightforward to configure a different LLM provider, e.g. Anthropic. 
+
+Follow these steps to change the LLM: 
+
+* Add a Langchain LLM to `llm.py`: You can add any LLM here as long as it's supported by a Langchain chat model. For example: 
+```python
+case "claude-3-opus":
+ return ChatAnthropic(
+ temperature=0,
+ model_name="claude-3-opus-20240229",
+ anthropic_api_key=settings.ANTHROPIC_API_KEY,
+ streaming=True,
+ )
+```
+You will need to set `streaming=True` and ensure that the given chat model supports asynchronous streaming. 
+
+* Add an API key (or other environment variables needed to make an LLM call) to your `.env`. You should read through the Langchain chat model documentation to determine what other inputs you might need. 
+
+* Go to `schemas/tool_schema.py` and update the list of acceptable LLM inputs by changing the `LLMTypeObject`
+
+* Finally, update the LLM being used in `config/agent.yml`, setting it to match the name which you set in `llm.py`. 
+
+This should be all, but you should note that not all Langchain versions will support asynchronous streaming for all LLMs, even if they offer a chat model wrapper. For example, in the Langchain 0.1.x series, Google's gemini models don't support asynchronous streaming and hence do not integrate well with AgentKit's off-the-shelf frontend. Getting this to work will require a bugfix on Langchain's part in a future version (which has been implemented in Lanchain 0.2+), or a lower-level change to the AgentKit code. 
+