...

lib/clients/openai/completions.v

@@ -2,63 +2,61 @@ module openai
 
 import json
 
-pub struct ChatCompletion {
-pub mut:
-	id      string
-	object  string
-	created u32
-	choices []Choice
-	usage   Usage
-}
-
-pub struct Choice {
-pub mut:
-	index         int
-	message       MessageRaw
-	finish_reason string
-}
-
-pub struct Message {
-pub mut:
-	role    RoleType
-	content string
-}
-
-pub struct Usage {
-pub mut:
-	prompt_tokens     int
-	completion_tokens int
-	total_tokens      int
-}
-
-pub struct Messages {
-pub mut:
-	messages []Message
-}
-
-pub struct MessageRaw {
-pub mut:
-	role    string
-	content string
-}
-
-struct ChatMessagesRaw {
-mut:
-	model                 string
-	messages              []MessageRaw
-	temperature           f64 = 0.5
-	max_completion_tokens int = 32000
-}
-
 @[params]
 pub struct CompletionArgs {
 pub mut:
 	model                 string
-	msgs                  Messages
-	temperature           f64 = 0.5
+	messages              []Message // optional because we can use message, which means we just pass a string
+	message               string
+	temperature           f64 = 0.2
 	max_completion_tokens int = 32000
 }
 
+struct Message {
+mut:
+	role    RoleType
+	content string
+}
+
+pub enum RoleType {
+	system
+	user
+	assistant
+	function
+}
+
+fn roletype_str(x RoleType) string {
+	return match x {
+		.system {
+			'system'
+		}
+		.user {
+			'user'
+		}
+		.assistant {
+			'assistant'
+		}
+		.function {
+			'function'
+		}
+	}
+}
+
+struct Usage {
+mut:
+	prompt_tokens     int
+	completion_tokens int
+	total_tokens      int
+}
+
+struct ChatCompletion {
+mut:
+	id      string
+	created u32
+	result  string
+	usage   Usage
+}
+
 // creates a new chat completion given a list of messages
 // each message consists of message content and the role of the author
 pub fn (mut f OpenAI) chat_completion(args_ CompletionArgs) !ChatCompletion {
@@ -71,19 +69,38 @@ pub fn (mut f OpenAI) chat_completion(args_ CompletionArgs) !ChatCompletion {
 		temperature:           args.temperature
 		max_completion_tokens: args.max_completion_tokens
 	}
-	for msg in args.msgs.messages {
+	for msg in args.messages {
 		mr := MessageRaw{
 			role:    roletype_str(msg.role)
 			content: msg.content
 		}
 		m.messages << mr
 	}
+	if args.message != '' {
+		mr := MessageRaw{
+			role:    'user'
+			content: args.message
+		}
+		m.messages << mr
+	}
 	data := json.encode(m)
 	// println('data: ${data}')
 	mut conn := f.connection()!
 	r := conn.post_json_str(prefix: 'chat/completions', data: data)!
 	// println('res: ${r}')
 
-	res := json.decode(ChatCompletion, r)!
-	return res
+	res := json.decode(ChatCompletionRaw, r)!
+
+	mut result := ''
+	for choice in res.choices {
+		result += choice.message.content
+	}
+
+	mut chat_completion_result := ChatCompletion{
+		id:      res.id
+		created: res.created
+		result:  result
+		usage:   res.usage
+	}
+	return chat_completion_result
 }