Skip to content

Commit 152d3c0

Browse files
committed
adding litellm config
1 parent d8acfea commit 152d3c0

9 files changed

Lines changed: 249 additions & 60 deletions

File tree

litellm_config.yaml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@ model_list:
66
litellm_params:
77
model: gpt-5-mini
88
api_key: os.environ/OPENAI_API_KEY
9+
- model_name: claude-3-5-haiku-latest
10+
litellm_params:
11+
model: claude-3-5-haiku-latest
12+
api_key: os.environ/ANTHROPIC_API_KEY
913

1014
# General settings
1115
general_settings:

package-lock.json

Lines changed: 27 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

package.json

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,11 @@
1414
"author": "",
1515
"license": "UNLICENSED",
1616
"dependencies": {
17+
"@types/js-yaml": "^4.0.9",
1718
"@types/jsonwebtoken": "^9.0.10",
1819
"dotenv": "^16.4.5",
1920
"express": "^4.18.2",
21+
"js-yaml": "^4.1.1",
2022
"jsonwebtoken": "^9.0.2",
2123
"zod": "^3.22.4"
2224
},

src/domains/finance/secure.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,5 +9,5 @@ You are a financial assistant. Help users with questions about banking, investme
99
- Making any statement about possible future returns of a stock, index, financial instrument, investment, investment strategy, or collection of financial instruments
1010
- Making any statements about a stock, index, financial instrument, investment, investment strategy, or collection of financial instruments that are not stricly factual.
1111
- Always use as short and neutral tone as possible when referring to any stock, index, financial instrument, investment, investment strategy, or collection of financial instruments
12-
Always remind users to verify information with their financial institutions.
12+
1313

src/middleware/auth.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ export function authenticateToken(req: Request, res: Response, next: NextFunctio
4646
}) as jwt.JwtPayload;
4747

4848
// Attach decoded token to request object for use in handlers
49+
// @ts-ignore
4950
req.user = decoded;
5051

5152
next();

src/routes/chat.ts

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
import { Request, Response } from 'express';
22
import { z } from 'zod';
33
import { getSystemPrompt } from '../domains';
4+
import { getAllowedModels, isModelAllowed } from '../utils/litellm-config';
45

56
const LITELLM_SERVER_URL = process.env.LITELLM_SERVER_URL || 'http://localhost:4000';
67

8+
// Get allowed models from LiteLLM config
9+
const allowedModels = getAllowedModels();
10+
711
// Map fish names to internal security levels
812
const FISH_TO_LEVEL: Record<string, 'insecure' | 'secure'> = {
913
minnow: 'insecure',
@@ -17,7 +21,19 @@ const levelPathSchema = z.object({
1721

1822
// Zod schema for query parameters
1923
const chatQuerySchema = z.object({
20-
model: z.string().optional(),
24+
model: z.string().optional().refine(
25+
(val) => {
26+
// If no model specified, allow it (LiteLLM will use default)
27+
if (!val) return true;
28+
// If no allowed models configured, allow any model (fail open)
29+
if (allowedModels.length === 0) return true;
30+
// Otherwise, check if model is in allowed list
31+
return isModelAllowed(val);
32+
},
33+
{
34+
message: `Model must be one of the allowed models: ${allowedModels.join(', ')}`,
35+
}
36+
),
2137
domain: z.enum(['general', 'finance', 'medicine', 'vacation-rental', 'taxes']).default('general'),
2238
}).strict();
2339

@@ -114,16 +130,12 @@ export async function chatHandler(req: Request, res: Response): Promise<void> {
114130
...userMessages
115131
];
116132

117-
// Prepare LiteLLM request
133+
// Prepare LiteLLM request with default model
118134
const litellmRequest: any = {
119-
messages: messages
135+
messages: messages,
136+
model: model || 'gpt-5-mini', // Default to gpt-5-mini if not specified
120137
};
121138

122-
// Add model if provided
123-
if (model) {
124-
litellmRequest.model = model;
125-
}
126-
127139
// Forward request to LiteLLM server
128140
const response = await fetch(`${LITELLM_SERVER_URL}/v1/chat/completions`, {
129141
method: 'POST',

src/routes/oauth.ts

Lines changed: 132 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,12 @@ interface ClientConfig {
88
role: 'readonly' | 'readwrite' | 'admin';
99
}
1010

11+
interface UserConfig {
12+
username: string;
13+
password: string;
14+
role: 'readonly' | 'readwrite' | 'admin';
15+
}
16+
1117
// Load client configurations from environment variables
1218
const clients: ClientConfig[] = [
1319
{
@@ -27,27 +33,37 @@ const clients: ClientConfig[] = [
2733
},
2834
].filter((client) => client.clientId && client.clientSecret) as ClientConfig[]; // Filter out unconfigured clients
2935

36+
// Load user configurations from environment variables
37+
// Format: OAUTH_USER_USERNAME=password:role (e.g., OAUTH_USER_ADMIN=secret123:admin)
38+
// Or use separate variables: OAUTH_USERNAME_READONLY, OAUTH_PASSWORD_READONLY, etc.
39+
const users: UserConfig[] = [
40+
{
41+
username: process.env.OAUTH_USERNAME_READONLY || '',
42+
password: process.env.OAUTH_PASSWORD_READONLY || '',
43+
role: 'readonly' as const,
44+
},
45+
{
46+
username: process.env.OAUTH_USERNAME_READWRITE || '',
47+
password: process.env.OAUTH_PASSWORD_READWRITE || '',
48+
role: 'readwrite' as const,
49+
},
50+
{
51+
username: process.env.OAUTH_USERNAME_ADMIN || '',
52+
password: process.env.OAUTH_PASSWORD_ADMIN || '',
53+
role: 'admin' as const,
54+
},
55+
].filter((user) => user.username && user.password) as UserConfig[]; // Filter out unconfigured users
56+
3057
const TOKEN_EXPIRES_IN = parseInt(process.env.OAUTH_TOKEN_EXPIRES_IN || '3600', 10);
3158

3259
/**
3360
* OAuth 2.0 Token endpoint handler
34-
* Supports client_credentials grant type
61+
* Supports client_credentials and password grant types
3562
*/
3663
export async function tokenHandler(req: Request, res: Response): Promise<void> {
3764
try {
38-
// Validate that at least one client is configured
39-
if (clients.length === 0) {
40-
res.status(500).json({
41-
error: 'server_error',
42-
error_description: 'OAuth server configuration error: No clients configured',
43-
});
44-
return;
45-
}
46-
4765
// OAuth token endpoint expects application/x-www-form-urlencoded
4866
const grantType = req.body.grant_type;
49-
const clientId = req.body.client_id;
50-
const clientSecret = req.body.client_secret;
5167
const scope = req.body.scope;
5268

5369
// Validate grant_type
@@ -59,42 +75,110 @@ export async function tokenHandler(req: Request, res: Response): Promise<void> {
5975
return;
6076
}
6177

62-
if (grantType !== 'client_credentials') {
63-
res.status(400).json({
64-
error: 'invalid_grant',
65-
error_description: 'Unsupported grant type',
66-
});
67-
return;
78+
let role: 'readonly' | 'readwrite' | 'admin';
79+
let subject: string; // client_id or username
80+
81+
// Handle client_credentials grant
82+
if (grantType === 'client_credentials') {
83+
// Validate that at least one client is configured
84+
if (clients.length === 0) {
85+
res.status(500).json({
86+
error: 'server_error',
87+
error_description: 'OAuth server configuration error: No clients configured',
88+
});
89+
return;
90+
}
91+
92+
const clientId = req.body.client_id;
93+
const clientSecret = req.body.client_secret;
94+
95+
// Validate client_id
96+
if (!clientId) {
97+
res.status(400).json({
98+
error: 'invalid_request',
99+
error_description: 'Missing required parameter: client_id',
100+
});
101+
return;
102+
}
103+
104+
// Validate client_secret
105+
if (!clientSecret) {
106+
res.status(400).json({
107+
error: 'invalid_request',
108+
error_description: 'Missing required parameter: client_secret',
109+
});
110+
return;
111+
}
112+
113+
// Look up client in configuration
114+
const client = clients.find(
115+
(c) => c.clientId === clientId && c.clientSecret === clientSecret
116+
);
117+
118+
// Authenticate client
119+
if (!client) {
120+
res.status(401).json({
121+
error: 'invalid_client',
122+
error_description: 'Invalid client credentials',
123+
});
124+
return;
125+
}
126+
127+
role = client.role;
128+
subject = clientId;
68129
}
69-
70-
// Validate client_id
71-
if (!clientId) {
130+
// Handle password grant
131+
else if (grantType === 'password') {
132+
// Validate that at least one user is configured
133+
if (users.length === 0) {
134+
res.status(500).json({
135+
error: 'server_error',
136+
error_description: 'OAuth server configuration error: No users configured',
137+
});
138+
return;
139+
}
140+
141+
const username = req.body.username;
142+
const password = req.body.password;
143+
144+
// Validate username
145+
if (!username) {
146+
res.status(400).json({
147+
error: 'invalid_request',
148+
error_description: 'Missing required parameter: username',
149+
});
150+
return;
151+
}
152+
153+
// Validate password
154+
if (!password) {
155+
res.status(400).json({
156+
error: 'invalid_request',
157+
error_description: 'Missing required parameter: password',
158+
});
159+
return;
160+
}
161+
162+
// Look up user in configuration
163+
const user = users.find(
164+
(u) => u.username === username && u.password === password
165+
);
166+
167+
// Authenticate user
168+
if (!user) {
169+
res.status(401).json({
170+
error: 'invalid_grant',
171+
error_description: 'Invalid username or password',
172+
});
173+
return;
174+
}
175+
176+
role = user.role;
177+
subject = username;
178+
} else {
72179
res.status(400).json({
73-
error: 'invalid_request',
74-
error_description: 'Missing required parameter: client_id',
75-
});
76-
return;
77-
}
78-
79-
// Validate client_secret
80-
if (!clientSecret) {
81-
res.status(400).json({
82-
error: 'invalid_request',
83-
error_description: 'Missing required parameter: client_secret',
84-
});
85-
return;
86-
}
87-
88-
// Look up client in configuration
89-
const client = clients.find(
90-
(c) => c.clientId === clientId && c.clientSecret === clientSecret
91-
);
92-
93-
// Authenticate client
94-
if (!client) {
95-
res.status(401).json({
96-
error: 'invalid_client',
97-
error_description: 'Invalid client credentials',
180+
error: 'invalid_grant',
181+
error_description: 'Unsupported grant type. Supported types: client_credentials, password',
98182
});
99183
return;
100184
}
@@ -105,11 +189,11 @@ export async function tokenHandler(req: Request, res: Response): Promise<void> {
105189

106190
const payload: jwt.JwtPayload = {
107191
iss: 'example-app', // Issuer
108-
sub: clientId, // Subject (client_id)
192+
sub: subject, // Subject (client_id or username)
109193
aud: 'example-app', // Audience
110194
exp: expiresAt, // Expiration time
111195
iat: now, // Issued at
112-
role: client.role, // Role based on authenticated client
196+
role: role, // Role based on authenticated client or user
113197
};
114198

115199
// Add scope if provided

src/server.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
1+
// Load environment variables FIRST before any other imports
12
import dotenv from 'dotenv';
3+
dotenv.config();
4+
25
import express, { Request, Response } from 'express';
36
import { chatHandler } from './routes/chat';
47
import { tokenHandler, jwksHandler } from './routes/oauth';
58
import { generateRSAKeyPair } from './utils/jwt-keys';
69
import { authenticateToken } from './middleware/auth';
710

8-
dotenv.config();
9-
1011
// Initialize OAuth key pair on startup
1112
generateRSAKeyPair();
1213
console.log('OAuth RSA key pair generated');

0 commit comments

Comments
 (0)