Skip to content

Commit c484473

Browse files
committed
add kimi to llmclient
1 parent 101a51f commit c484473

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

psyflow/LLM.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ class LLMClient:
3333
- ``gemini`` : Google GenAI SDK
3434
- ``openai`` : Official OpenAI SDK
3535
- ``deepseek``: OpenAI SDK with custom base_url
36+
- ``moonshot`` : OpenAI SDK with custom base_url
3637
3738
Attributes:
3839
provider: Lowercase provider name.
@@ -64,14 +65,20 @@ def __init__(self, provider: str, api_key: str, model: str):
6465
self.last_response_token_count: Optional[int] = None
6566

6667
if self.provider == "gemini":
67-
from google import genai # noqa: F401
68-
from google.genai.types import GenerateContentConfig # noqa: F401
68+
from google import genai
69+
from google.genai.types import GenerateContentConfig
6970
self._sdk_client = genai.Client(api_key=self.api_key)
7071
self._GenerateContentConfig = GenerateContentConfig
7172

72-
elif self.provider in ("openai", "deepseek"):
73-
from openai import OpenAI # noqa: F401
74-
base_url = "https://api.deepseek.com/v1" if self.provider == "deepseek" else None
73+
elif self.provider in ("openai", "deepseek", "moonshot"):
74+
from openai import OpenAI # OpenAI-compatible SDK
75+
76+
base_urls = {
77+
"openai": None, # Default for OpenAI
78+
"deepseek": "https://api.deepseek.com/v1",
79+
"moonshot": "https://api.moonshot.cn/v1",
80+
}
81+
base_url = base_urls.get(self.provider)
7582
self._sdk_client = OpenAI(api_key=self.api_key, base_url=base_url)
7683

7784
else:

0 commit comments

Comments
 (0)