Skip to content

Commit

Permalink
refactor the creation of messages for openAI to reduce repetition and…
Browse files Browse the repository at this point in the history
… increase readability
  • Loading branch information
InfraK committed May 5, 2024
1 parent 9232142 commit b1ac3fe
Showing 1 changed file with 28 additions and 30 deletions.
58 changes: 28 additions & 30 deletions src/KurtOpenAI.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,46 +38,27 @@ export class KurtOpenAI implements Kurt {
generateNaturalLanguage(
options: KurtGenerateNaturalLanguageOptions
): KurtResult {
const systemPrompt = options.systemPrompt ?? this.options.systemPrompt
const prompt = options.prompt
const extraMessages = options.extraMessages ?? []

return this.handleStream(
undefined,
this.options.openAI.chat.completions.create({
stream: true,
model: this.options.model,
messages: [
...(systemPrompt
? [{ role: "system", content: systemPrompt } as const]
: []),
{ role: "user", content: prompt },
...toOpenAIMessages(extraMessages),
],
messages: this.toOpenAIMessages(options),
})
)
}

generateStructuredData<T extends KurtSchemaInner>(
options: KurtGenerateStructuredDataOptions<T>
): KurtResult<T> {
const systemPrompt = options.systemPrompt ?? this.options.systemPrompt
const prompt = options.prompt
const schema = options.schema
const extraMessages = options.extraMessages ?? []

return this.handleStream(
schema as KurtSchemaMaybe<T>,
this.options.openAI.chat.completions.create({
stream: true,
model: this.options.model,
messages: [
...(systemPrompt
? [{ role: "system", content: systemPrompt } as const]
: []),
{ role: "user", content: prompt },
...toOpenAIMessages(extraMessages),
],
messages: this.toOpenAIMessages(options),
tool_choice: {
type: "function",
function: { name: "structured_data" },
Expand Down Expand Up @@ -144,14 +125,31 @@ export class KurtOpenAI implements Kurt {

return new KurtResult<T>(generator())
}
}

function toOpenAIMessages(messages: KurtMessage[]): OpenAIMessage[] {
return messages.map((message) => {
const { role, text } = message
return {
role: role === "model" ? "assistant" : role,
content: text,
}
})
private toOpenAIMessages = ({
prompt,
systemPrompt = this.options.systemPrompt,
extraMessages = [],
}: KurtGenerateNaturalLanguageOptions): OpenAIMessage[] => {
const systemMessage: OpenAIMessage[] = systemPrompt
? [toOpenAIMessage({ role: "system", text: systemPrompt })]
: []

const userMessage = toOpenAIMessage({ role: "user", text: prompt })

const extras = extraMessages.map(toOpenAIMessage)

return systemMessage.concat(userMessage, extras)
}
}

const toOpenAIMessage = ({ role, text }: KurtMessage): OpenAIMessage => ({
role: openAIRoleMapping[role],
content: text,
})

const openAIRoleMapping = {
model: "assistant",
system: "system",
user: "user",
} as const satisfies Record<KurtMessage["role"], OpenAIMessage["role"]>

0 comments on commit b1ac3fe

Please sign in to comment.