Skip to content

Commit ba72cc2

Browse files
roh26itroh26it
committed
Merge branch 'main' into feat/models-endpoint
2 parents 625cfb1 + 598ab6f commit ba72cc2

File tree

10 files changed

+368
-81
lines changed

10 files changed

+368
-81
lines changed

‎plugins/types.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
export interface PluginContext {
22
[key: string]: any;
3-
requestType?: 'complete' | 'chatComplete' | 'embed';
3+
requestType?: 'complete' | 'chatComplete' | 'embed' | 'messages';
44
provider?: string;
55
metadata?: Record<string, any>;
66
}

‎plugins/utils.ts‎

Lines changed: 120 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -36,18 +36,19 @@ export class TimeoutError extends Error {
3636
}
3737
}
3838

39+
/**
40+
* Helper function to get the text from the current content part of a request/response context
41+
* @param context - The plugin context containing request/response data
42+
* @param eventType - The type of hook event (beforeRequestHook or afterRequestHook)
43+
* @returns The text from the current content part of the request/response context
44+
*/
3945
export const getText = (
4046
context: PluginContext,
4147
eventType: HookEventType
4248
): string => {
43-
switch (eventType) {
44-
case 'beforeRequestHook':
45-
return context.request?.text;
46-
case 'afterRequestHook':
47-
return context.response?.text;
48-
default:
49-
throw new Error('Invalid hook type');
50-
}
49+
return getCurrentContentPart(context, eventType)
50+
.textArray.filter((text) => text)
51+
.join('\n');
5152
};
5253

5354
/**
@@ -66,33 +67,52 @@ export const getCurrentContentPart = (
6667
// Determine if we're handling request or response data
6768
const target = eventType === 'beforeRequestHook' ? 'request' : 'response';
6869
const json = context[target].json;
70+
71+
if (target === 'request') {
72+
return getRequestContentPart(json, context.requestType!);
73+
} else {
74+
return getResponseContentPart(json, context.requestType || '');
75+
}
76+
};
77+
78+
const getRequestContentPart = (json: any, requestType: string) => {
79+
let content: Array<any> | string | Record<string, any> | null = null;
6980
let textArray: Array<string> = [];
81+
if (requestType === 'chatComplete' || requestType === 'messages') {
82+
content = json.messages[json.messages.length - 1].content;
83+
textArray = Array.isArray(content)
84+
? content.map((item: any) => item.text || '')
85+
: [content];
86+
} else if (requestType === 'complete') {
87+
content = json.prompt;
88+
textArray = Array.isArray(content)
89+
? content.map((item: any) => item)
90+
: [content];
91+
} else if (requestType === 'embed') {
92+
content = json.input;
93+
textArray = Array.isArray(content) ? content : [content];
94+
}
95+
return { content, textArray };
96+
};
97+
98+
const getResponseContentPart = (json: any, requestType: string) => {
7099
let content: Array<any> | string | Record<string, any> | null = null;
100+
let textArray: Array<string> = [];
71101

72-
// Handle chat completion request/response format
73-
if (context.requestType === 'chatComplete') {
74-
if (target === 'request') {
75-
// Get the last message's content from the chat history
76-
content = json.messages[json.messages.length - 1].content;
77-
textArray = Array.isArray(content)
78-
? content.map((item: any) => item.text || '')
79-
: [content];
80-
} else {
81-
// Get the content from the last choice in the response
82-
content = json.choices[json.choices.length - 1].message.content as string;
83-
textArray = [content];
84-
}
85-
} else if (context.requestType === 'complete') {
86-
if (target === 'request') {
87-
// Handle completions format
88-
content = json.prompt;
89-
textArray = Array.isArray(content)
90-
? content.map((item: any) => item)
91-
: [content];
92-
} else {
93-
content = json.choices[json.choices.length - 1].text as string;
94-
textArray = [content];
95-
}
102+
// This can happen for streaming mode.
103+
if (!json) {
104+
return { content: null, textArray: [] };
105+
}
106+
107+
if (requestType === 'chatComplete') {
108+
content = json.choices[0].message.content as string;
109+
textArray = [content];
110+
} else if (requestType === 'complete') {
111+
content = json.choices[0].text as string;
112+
textArray = [content];
113+
} else if (requestType === 'messages') {
114+
content = json.content;
115+
textArray = (content as Array<any>).map((item: any) => item.text || '');
96116
}
97117
return { content, textArray };
98118
};
@@ -114,58 +134,79 @@ export const setCurrentContentPart = (
114134
const target = eventType === 'beforeRequestHook' ? 'request' : 'response';
115135
const json = context[target].json;
116136

117-
// Create shallow copy of the json
137+
if (textArray?.length === 0 || !textArray) {
138+
return;
139+
}
140+
141+
if (target === 'request') {
142+
setRequestContentPart(json, requestType!, textArray, transformedData);
143+
} else {
144+
setResponseContentPart(json, requestType!, textArray, transformedData);
145+
}
146+
};
147+
148+
function setRequestContentPart(
149+
json: any,
150+
requestType: string,
151+
textArray: Array<string | null>,
152+
transformedData: Record<string, any>
153+
) {
154+
// Create a safe to use shallow copy of the json
118155
const updatedJson = { ...json };
119156

120-
// Handle updating text fields if provided
121-
if (textArray?.length) {
122-
if (requestType === 'chatComplete') {
123-
if (target === 'request') {
124-
const currentContent =
125-
updatedJson.messages[updatedJson.messages.length - 1].content;
126-
updatedJson.messages = [...json.messages];
127-
updatedJson.messages[updatedJson.messages.length - 1] = {
128-
...updatedJson.messages[updatedJson.messages.length - 1],
129-
};
130-
131-
if (Array.isArray(currentContent)) {
132-
updatedJson.messages[updatedJson.messages.length - 1].content =
133-
currentContent.map((item: any, index: number) => ({
134-
...item,
135-
text: textArray[index] || item.text,
136-
}));
137-
} else {
138-
updatedJson.messages[updatedJson.messages.length - 1].content =
139-
textArray[0] || currentContent;
140-
}
141-
transformedData.request.json = updatedJson;
142-
} else {
143-
updatedJson.choices = [...json.choices];
144-
const lastChoice = {
145-
...updatedJson.choices[updatedJson.choices.length - 1],
146-
};
147-
lastChoice.message = {
148-
...lastChoice.message,
149-
content: textArray[0] || lastChoice.message.content,
150-
};
151-
updatedJson.choices[updatedJson.choices.length - 1] = lastChoice;
152-
transformedData.response.json = updatedJson;
153-
}
157+
if (requestType === 'chatComplete' || requestType === 'messages') {
158+
updatedJson.messages = [...json.messages];
159+
const lastMessage = {
160+
...updatedJson.messages[updatedJson.messages.length - 1],
161+
};
162+
const originalContent = lastMessage.content;
163+
if (Array.isArray(originalContent)) {
164+
lastMessage.content = originalContent.map((item: any, index: number) => ({
165+
...item,
166+
text: textArray[index] || item.text,
167+
}));
154168
} else {
155-
if (target === 'request') {
156-
updatedJson.prompt = Array.isArray(updatedJson.prompt)
157-
? textArray.map((text, index) => text || updatedJson.prompt[index])
158-
: textArray[0];
159-
transformedData.request.json = updatedJson;
160-
} else {
161-
updatedJson.choices = [...json.choices];
162-
updatedJson.choices[json.choices.length - 1].text =
163-
textArray[0] || json.choices[json.choices.length - 1].text;
164-
transformedData.response.json = updatedJson;
165-
}
169+
lastMessage.content = textArray[0] || originalContent;
166170
}
171+
updatedJson.messages[updatedJson.messages.length - 1] = lastMessage;
172+
} else if (requestType === 'complete') {
173+
updatedJson.prompt = Array.isArray(updatedJson.prompt)
174+
? textArray.map((text, index) => text || updatedJson.prompt[index])
175+
: textArray[0];
167176
}
168-
};
177+
transformedData.request.json = updatedJson;
178+
}
179+
180+
function setResponseContentPart(
181+
json: any,
182+
requestType: string,
183+
textArray: Array<string | null>,
184+
transformedData: Record<string, any>
185+
) {
186+
// Create a safe to use shallow copy of the json
187+
const updatedJson = { ...json };
188+
189+
if (requestType === 'chatComplete') {
190+
updatedJson.choices = [...json.choices];
191+
const firstChoice = {
192+
...updatedJson.choices[0],
193+
};
194+
firstChoice.message = {
195+
...firstChoice.message,
196+
content: textArray[0] || firstChoice.message.content,
197+
};
198+
updatedJson.choices[0] = firstChoice;
199+
} else if (requestType === 'complete') {
200+
updatedJson.choices = [...json.choices];
201+
updatedJson.choices[json.choices.length - 1].text =
202+
textArray[0] || json.choices[json.choices.length - 1].text;
203+
} else if (requestType === 'messages') {
204+
updatedJson.content = textArray.map(
205+
(text, index) => text || updatedJson.content[index]
206+
);
207+
}
208+
transformedData.response.json = updatedJson;
209+
}
169210

170211
/**
171212
* Sends a POST request to the specified URL with the given data and timeout.

‎src/globals.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ export const BYTEZ: string = 'bytez';
101101
export const FEATHERLESS_AI: string = 'featherless-ai';
102102
export const KRUTRIM: string = 'krutrim';
103103
export const QDRANT: string = 'qdrant';
104+
export const THREE_ZERO_TWO_AI: string = '302ai';
104105

105106
export const VALID_PROVIDERS = [
106107
ANTHROPIC,
@@ -165,6 +166,7 @@ export const VALID_PROVIDERS = [
165166
FEATHERLESS_AI,
166167
KRUTRIM,
167168
QDRANT,
169+
THREE_ZERO_TWO_AI,
168170
];
169171

170172
export const CONTENT_TYPES = {

‎src/middlewares/hooks/index.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -423,7 +423,9 @@ export class HooksManager {
423423
private shouldSkipHook(span: HookSpan, hook: HookObject): boolean {
424424
const context = span.getContext();
425425
return (
426-
!['chatComplete', 'complete', 'embed'].includes(context.requestType) ||
426+
!['chatComplete', 'complete', 'embed', 'messages'].includes(
427+
context.requestType
428+
) ||
427429
(context.requestType === 'embed' &&
428430
hook.eventType !== 'beforeRequestHook') ||
429431
(context.requestType === 'embed' && hook.type === HookType.MUTATOR) ||

‎src/providers/302ai/api.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
import { ProviderAPIConfig } from '../types';
2+
3+
const AI302APIConfig: ProviderAPIConfig = {
4+
getBaseURL: () => 'https://api.302.ai',
5+
headers: ({ providerOptions }) => {
6+
return { Authorization: `Bearer ${providerOptions.apiKey}` };
7+
},
8+
getEndpoint: ({ fn }) => {
9+
switch (fn) {
10+
case 'chatComplete':
11+
return '/v1/chat/completions';
12+
default:
13+
return '';
14+
}
15+
},
16+
};
17+
18+
export default AI302APIConfig;

0 commit comments

Comments
 (0)