From f16b3db70c075b2db13e1947652a6496237ff788 Mon Sep 17 00:00:00 2001
From: Roy Han <royjhan02@gmail.com>
Date: Tue, 30 Jul 2024 11:29:44 -0700
Subject: [PATCH] oai compat

---
 openai/openai.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/openai/openai.go b/openai/openai.go
index 5bd806604..65033574a 100644
--- a/openai/openai.go
+++ b/openai/openai.go
@@ -192,9 +192,9 @@ func toolCallId() string {
 	return "call_" + strings.ToLower(string(b))
 }
 
-func toChatCompletion(id string, r api.ChatResponse) ChatCompletion {
-	toolCalls := make([]ToolCall, len(r.Message.ToolCalls))
-	for i, tc := range r.Message.ToolCalls {
+func parseToolCalls(respToolCalls []api.ToolCall) []ToolCall {
+	toolCalls := make([]ToolCall, len(respToolCalls))
+	for i, tc := range respToolCalls {
 		toolCalls[i].ID = toolCallId()
 		toolCalls[i].Type = "function"
 		toolCalls[i].Function.Name = tc.Function.Name
@@ -207,6 +207,11 @@ func toChatCompletion(id string, r api.ChatResponse) ChatCompletion {
 
 		toolCalls[i].Function.Arguments = string(args)
 	}
+	return toolCalls
+}
+
+func toChatCompletion(id string, r api.ChatResponse) ChatCompletion {
+	toolCalls := parseToolCalls(r.Message.ToolCalls)
 
 	return ChatCompletion{
 		Id:                id,
@@ -218,9 +223,6 @@ func toChatCompletion(id string, r api.ChatResponse) ChatCompletion {
 			Index:   0,
 			Message: Message{Role: r.Message.Role, Content: r.Message.Content, ToolCalls: toolCalls},
 			FinishReason: func(reason string) *string {
-				if len(toolCalls) > 0 {
-					reason = "tool_calls"
-				}
 				if len(reason) > 0 {
 					return &reason
 				}
@@ -236,6 +238,8 @@ func toChatCompletion(id string, r api.ChatResponse) ChatCompletion {
 }
 
 func toChunk(id string, r api.ChatResponse) ChatCompletionChunk {
+	toolCalls := parseToolCalls(r.Message.ToolCalls)
+
 	return ChatCompletionChunk{
 		Id:                id,
 		Object:            "chat.completion.chunk",
@@ -244,7 +248,7 @@ func toChunk(id string, r api.ChatResponse) ChatCompletionChunk {
 		SystemFingerprint: "fp_ollama",
 		Choices: []ChunkChoice{{
 			Index: 0,
-			Delta: Message{Role: "assistant", Content: r.Message.Content},
+			Delta: Message{Role: "assistant", Content: r.Message.Content, ToolCalls: toolCalls},
 			FinishReason: func(reason string) *string {
 				if len(reason) > 0 {
 					return &reason