Skip to content

Commit 95c8479

Browse files
authored
feat(llm): Add LiteLLM (Proxy) Support (#630)
* Add LiteLLM * Use existing impl * Use OpenAI Config * Add Docs for LiteLLM
1 parent 316cdf3 commit 95c8479

File tree

5 files changed

+91
-0
lines changed

5 files changed

+91
-0
lines changed

docs/docs/ai/llm.mdx

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,3 +121,68 @@ cocoindex.LlmSpec(
121121

122122
You can find the full list of models supported by Anthropic [here](https://docs.anthropic.com/en/docs/about-claude/models/all-models).
123123

124+
### LiteLLM
125+
126+
To use the LiteLLM API, you need to set the environment variable `LITELLM_API_KEY`.
127+
128+
#### 1. Install LiteLLM Proxy
129+
130+
```bash
131+
pip install 'litellm[proxy]'
132+
```
133+
134+
#### 2. Create a `config.yml` for LiteLLM
135+
136+
**Example for OpenAI:**
137+
```yaml
138+
model_list:
139+
- model_name: "*"
140+
litellm_params:
141+
model: openai/*
142+
api_key: os.environ/LITELLM_API_KEY
143+
```
144+
145+
**Example for DeepSeek:**
146+
147+
First, pull the DeepSeek model with Ollama:
148+
```bash
149+
ollama pull deepseek-r1
150+
```
151+
Then run it if it's not running:
152+
```bash
153+
ollama run deepseek-r1
154+
```
155+
156+
Then, use this in your `config.yml`:
157+
```yaml
158+
model_list:
159+
- model_name: "deepseek-r1"
160+
litellm_params:
161+
model: "ollama_chat/deepseek-r1"
162+
api_base: "http://localhost:11434"
163+
```
164+
165+
#### 3. Run LiteLLM Proxy
166+
167+
```bash
168+
litellm --config config.yml
169+
```
170+
171+
#### 4. A Spec for LiteLLM will look like this:
172+
173+
<Tabs>
174+
<TabItem value="python" label="Python" default>
175+
176+
```python
177+
cocoindex.LlmSpec(
178+
api_type=cocoindex.LlmApiType.LITELLM,
179+
model="deepseek-r1",
180+
address="http://127.0.0.1:4000", # default url of LiteLLM
181+
)
182+
```
183+
184+
</TabItem>
185+
</Tabs>
186+
187+
You can find the full list of models supported by LiteLLM [here](https://docs.litellm.ai/docs/providers).
188+

python/cocoindex/llm.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ class LlmApiType(Enum):
99
OLLAMA = "Ollama"
1010
GEMINI = "Gemini"
1111
ANTHROPIC = "Anthropic"
12+
LITELLM = "LiteLlm"
1213

1314

1415
@dataclass

src/llm/litellm.rs

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
use async_openai::config::OpenAIConfig;
2+
use async_openai::Client as OpenAIClient;
3+
4+
pub use super::openai::Client;
5+
6+
impl Client {
7+
pub async fn new_litellm(spec: super::LlmSpec) -> anyhow::Result<Self> {
8+
let address = spec.address.clone().unwrap_or_else(|| "http://127.0.0.1:4000".to_string());
9+
let api_key = std::env::var("LITELLM_API_KEY").ok();
10+
let mut config = OpenAIConfig::new().with_api_base(address);
11+
if let Some(api_key) = api_key {
12+
config = config.with_api_key(api_key);
13+
}
14+
Ok(Client::from_parts(OpenAIClient::with_config(config), spec.model))
15+
}
16+
}

src/llm/mod.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ pub enum LlmApiType {
1313
OpenAi,
1414
Gemini,
1515
Anthropic,
16+
LiteLlm,
1617
}
1718

1819
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -56,6 +57,7 @@ mod anthropic;
5657
mod gemini;
5758
mod ollama;
5859
mod openai;
60+
mod litellm;
5961

6062
pub async fn new_llm_generation_client(spec: LlmSpec) -> Result<Box<dyn LlmGenerationClient>> {
6163
let client = match spec.api_type {
@@ -71,6 +73,9 @@ pub async fn new_llm_generation_client(spec: LlmSpec) -> Result<Box<dyn LlmGener
7173
LlmApiType::Anthropic => {
7274
Box::new(anthropic::Client::new(spec).await?) as Box<dyn LlmGenerationClient>
7375
}
76+
LlmApiType::LiteLlm => {
77+
Box::new(litellm::Client::new_litellm(spec).await?) as Box<dyn LlmGenerationClient>
78+
}
7479
};
7580
Ok(client)
7681
}

src/llm/openai.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ pub struct Client {
2020
}
2121

2222
impl Client {
23+
pub(crate) fn from_parts(client: async_openai::Client<OpenAIConfig>, model: String) -> Self {
24+
Self { client, model }
25+
}
26+
2327
pub async fn new(spec: super::LlmSpec) -> Result<Self> {
2428
if let Some(address) = spec.address {
2529
api_bail!("OpenAI doesn't support custom API address: {address}");

0 commit comments

Comments
 (0)