From e3fffe325b65a9fd86611cf518cb7d016c449640 Mon Sep 17 00:00:00 2001
From: devjpt23 <devpatel232408@gmail.com>
Date: Sat, 22 Feb 2025 08:47:53 -0600
Subject: [PATCH] Added Anthropic and Groq support

Signed-off-by: devjpt23 <devpatel232408@gmail.com>
---
 .trunk/configs/custom-words.txt       |  3 ++
 kai/kai_config.py                     |  2 +
 kai/llm_interfacing/model_provider.py | 59 +++++++++++++++++++++++++++
 kai/rpc_server/server.py              |  2 +-
 pyproject.toml                        |  2 +
 5 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/.trunk/configs/custom-words.txt b/.trunk/configs/custom-words.txt
index b3d6e646c..ebb9b77e7 100644
--- a/.trunk/configs/custom-words.txt
+++ b/.trunk/configs/custom-words.txt
@@ -42,6 +42,9 @@ getpid
 gomod
 gopls
 gpgsign
+groq
+Groq
+GROQ
 htmlcov
 httpserver
 httpsnoop
diff --git a/kai/kai_config.py b/kai/kai_config.py
index 04019392f..ca725d9ea 100644
--- a/kai/kai_config.py
+++ b/kai/kai_config.py
@@ -150,11 +150,13 @@ class KaiConfigIncidentStore(BaseModel):
 
 
 class SupportedModelProviders(StrEnum):
+    CHAT_ANTHROPIC = "ChatAnthropic"
     CHAT_OLLAMA = "ChatOllama"
     CHAT_OPENAI = "ChatOpenAI"
     CHAT_BEDROCK = "ChatBedrock"
     FAKE_LIST_CHAT_MODEL = "FakeListChatModel"
     CHAT_GOOGLE_GENERATIVE_AI = "ChatGoogleGenerativeAI"
+    CHAT_GROQ = "ChatGroq"
     AZURE_CHAT_OPENAI = "AzureChatOpenAI"
     CHAT_DEEP_SEEK = "ChatDeepSeek"
 
diff --git a/kai/llm_interfacing/model_provider.py b/kai/llm_interfacing/model_provider.py
index 6df814b42..1798f08f7 100644
--- a/kai/llm_interfacing/model_provider.py
+++ b/kai/llm_interfacing/model_provider.py
@@ -4,6 +4,7 @@
 from abc import abstractmethod
 from typing import Any, Iterator, Optional, Sequence, assert_never, cast, override
 
+from langchain_anthropic import ChatAnthropic
 from langchain_aws import ChatBedrock
 from langchain_community.chat_models.fake import FakeListChatModel
 from langchain_core.language_models.base import LanguageModelInput
@@ -13,6 +14,7 @@
 from langchain_core.runnables import ConfigurableField, Runnable, RunnableConfig
 from langchain_deepseek import ChatDeepSeek
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
 from langchain_ollama import ChatOllama
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 from opentelemetry import trace
@@ -49,6 +51,10 @@ def from_config(
                 return ModelProviderAzureChatOpenAI(config, demo_mode, cache)
             case SupportedModelProviders.CHAT_DEEP_SEEK:
                 return ModelProviderChatDeepSeek(config, demo_mode, cache)
+            case SupportedModelProviders.CHAT_ANTHROPIC:
+                return ModelProviderChatAnthropic(config, demo_mode, cache)
+            case SupportedModelProviders.CHAT_GROQ:
+                return ModelProviderChatGroq(config, demo_mode, cache)
             case _:
                 assert_never(config.provider)
 
@@ -473,3 +479,56 @@ def prepare_model_args(
         config_args: dict[str, Any],
     ) -> tuple[dict[str, Any], str]:
         return deep_update(defaults, config_args), config_args["model"]
+
+
+class ModelProviderChatAnthropic(ModelProvider):
+    def __init__(self, config: KaiConfigModels, demo_mode: bool, cache: Cache | None):
+        super().__init__(
+            config=config,
+            demo_mode=demo_mode,
+            cache=cache,
+            model_class=ChatAnthropic,
+            defaults={
+                "model": "claude-3-7-sonnet-20250219",
+                "temperature": 0,
+                "timeout": None,
+                "max_retries": 2,
+            },
+        )
+
+    def validate_environment(self) -> None:
+        self.default_challenge("max_tokens")
+
+    def prepare_model_args(
+        self,
+        defaults: dict[str, Any],
+        config_args: dict[str, Any],
+    ) -> tuple[dict[str, Any], str]:
+        return deep_update(defaults, config_args), config_args["model"]
+
+
+class ModelProviderChatGroq(ModelProvider):
+    def __init__(self, config: KaiConfigModels, demo_mode: bool, cache: Cache | None):
+        super().__init__(
+            config=config,
+            demo_mode=demo_mode,
+            cache=cache,
+            model_class=ChatGroq,
+            defaults={
+                "model": "llama-3.3-70b-versatile",
+                "temperature": 0,
+                "max_tokens": None,
+                "timeout": None,
+                "max_retries": 2,
+            },
+        )
+
+    def validate_environment(self) -> None:
+        self.default_challenge("max_tokens")
+
+    def prepare_model_args(
+        self,
+        defaults: dict[str, Any],
+        config_args: dict[str, Any],
+    ) -> tuple[dict[str, Any], str]:
+        return deep_update(defaults, config_args), config_args["model"]
diff --git a/kai/rpc_server/server.py b/kai/rpc_server/server.py
index 6209f35d5..30df39c83 100644
--- a/kai/rpc_server/server.py
+++ b/kai/rpc_server/server.py
@@ -464,7 +464,7 @@ class GetCodeplanAgentSolutionParams(BaseModel):
     max_depth: Optional[int] = None
     max_priority: Optional[int] = None
 
-    chat_token: str
+    chat_token: Optional[str] = None
 
 
 class GetCodeplanAgentSolutionResult(BaseModel):
diff --git a/pyproject.toml b/pyproject.toml
index e1c418364..497af08d5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,10 +27,12 @@ dependencies = [
   "python-dateutil==2.8.2",
   "Jinja2==3.1.4",
   "langchain==0.3.19",
+  "langchain-anthropic==0.3.7",
   "langchain-community==0.3.1",
   "langchain-openai==0.3.3",
   "langchain-ollama==0.2.3",
   "langchain-google-genai==2.0.9",
+  "langchain-groq==0.2.4",
   "langchain-aws==0.2.11",
   "langchain-experimental==0.3.2",
   "langchain-deepseek-official==0.1.0",