feat: add support for Moonshot AI's Kimi K2 Thinking model

- Introduced MOONSHOT_API_KEY in settings.py for API integration.
- Updated providers.py to include KIMI_K2_THINKING model under Moonshot AI.
- Enhanced client.py to handle requests for the Moonshot AI provider.
- Added comprehensive tests for Kimi K2 Thinking model, covering streaming, response models, and tool calling functionalities.

Author: dragos199993

notdiamond/llms/client.py

@@ -1653,6 +1653,19 @@ def _llm_from_config(
                     replicate_api_key=provider.api_key,
                     **passed_kwargs,
                 )
+            if provider.provider == "moonshotai":
+                ChatOpenAI = _module_check(
+                    "langchain_openai.chat_models",
+                    "ChatOpenAI",
+                    provider.provider,
+                )
+                return ChatOpenAI(
+                    openai_api_key=provider.api_key,
+                    model_name=provider.model,
+                    openai_api_base="https://api.moonshot.cn/v1",
+                    callbacks=callbacks,
+                    **passed_kwargs,
+                )
             raise ValueError(f"Unsupported provider: {provider.provider}")
 
         def verify_against_response_model(self) -> bool:

notdiamond/llms/providers.py

@@ -90,6 +90,8 @@ class NDLLMProviders(Enum):
             model served via Replicate
 
         SONAR (NDLLMProvider): refers to "sonar" model by Perplexity
+
+        KIMI_K2_THINKING (NDLLMProvider): refers to "kimi-k2-thinking" model by Moonshot AI
     """
 
     GPT_3_5_TURBO = ("openai", "gpt-3.5-turbo")
@@ -207,6 +209,8 @@ class NDLLMProviders(Enum):
         "meta-llama-3.1-405b-instruct",
     )
 
+    KIMI_K2_THINKING = ("moonshotai", "kimi-k2-thinking")
+
     def __new__(cls, provider, model):
         return LLMConfig(provider=provider, model=model)
 

notdiamond/settings.py

@@ -16,6 +16,7 @@
 TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY", default="")
 PPLX_API_KEY = os.getenv("PPLX_API_KEY", default="")
 REPLICATE_API_KEY = os.getenv("REPLICATE_API_KEY", default="")
+MOONSHOT_API_KEY = os.getenv("MOONSHOT_API_KEY", default="")
 
 
 NOTDIAMOND_API_URL = os.getenv(
@@ -434,6 +435,24 @@
             "meta-llama-3.1-405b-instruct": {"input": 9.5, "output": 9.5},
         },
     },
+    "moonshotai": {
+        "models": [
+            "kimi-k2-thinking",
+        ],
+        "api_key": MOONSHOT_API_KEY,
+        "support_tools": [
+            "kimi-k2-thinking",
+        ],
+        "support_response_model": [
+            "kimi-k2-thinking",
+        ],
+        "openrouter_identifier": {
+            "kimi-k2-thinking": "moonshotai/kimi-k2-thinking",
+        },
+        "price": {
+            "kimi-k2-thinking": {"input": 0.6, "output": 2.5},
+        },
+    },
 }
 
 

tests/test_llm_calls/test_moonshotai.py

@@ -0,0 +1,71 @@
+import pytest
+from helpers import astream_chunks, stream_chunks
+
+from notdiamond.llms.client import NotDiamond
+from notdiamond.llms.providers import NDLLMProviders
+
+
+@pytest.mark.longrun
+@pytest.mark.vcr
+class Test_Moonshotai_LLMs:
+    def test_kimi_k2_thinking_with_streaming(self, prompt):
+        provider = NDLLMProviders.KIMI_K2_THINKING
+        nd_llm = NotDiamond(
+            llm_configs=[provider], latency_tracking=False, hash_content=True
+        )
+        stream_chunks(nd_llm.stream(prompt))
+
+    @pytest.mark.asyncio
+    async def test_kimi_k2_thinking_with_async_streaming(self, prompt):
+        provider = NDLLMProviders.KIMI_K2_THINKING
+        nd_llm = NotDiamond(
+            llm_configs=[provider], latency_tracking=False, hash_content=True
+        )
+
+        await astream_chunks(nd_llm.astream(prompt))
+
+    def test_kimi_k2_thinking_response_model(self, response_model):
+        provider = NDLLMProviders.KIMI_K2_THINKING
+        provider.kwargs = {"max_tokens": 200}
+        nd_llm = NotDiamond(
+            llm_configs=[provider], latency_tracking=False, hash_content=True
+        )
+        result, _, _ = nd_llm.invoke(
+            [{"role": "user", "content": "Tell me a joke"}],
+            response_model=response_model,
+        )
+
+        assert isinstance(result, response_model)
+        assert result.setup
+        assert result.punchline
+
+    def test_kimi_k2_thinking_with_tool_calling(self, tools_fixture):
+        provider = NDLLMProviders.KIMI_K2_THINKING
+        provider.kwargs = {"max_tokens": 200}
+        nd_llm = NotDiamond(
+            llm_configs=[provider], latency_tracking=False, hash_content=True
+        )
+        nd_llm = nd_llm.bind_tools(tools_fixture)
+        result, session_id, _ = nd_llm.invoke(
+            [{"role": "user", "content": "How much is 3 + 5?"}]
+        )
+
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0]["name"] == "add_fct"
+
+    def test_kimi_k2_thinking_with_openai_tool_calling(
+        self, openai_tools_fixture
+    ):
+        provider = NDLLMProviders.KIMI_K2_THINKING
+        provider.kwargs = {"max_tokens": 200}
+        nd_llm = NotDiamond(
+            llm_configs=[provider], latency_tracking=False, hash_content=True
+        )
+        nd_llm = nd_llm.bind_tools(openai_tools_fixture)
+        result, session_id, _ = nd_llm.invoke(
+            [{"role": "user", "content": "How much is 3 + 5?"}]
+        )
+
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0]["name"] == "add_fct"
+