From fd35c3a95191f4a80d8ea8f24ef721b33ee1a047 Mon Sep 17 00:00:00 2001
From: Darren Shepherd <darren@acorn.io>
Date: Wed, 9 Oct 2024 21:46:50 -0700
Subject: [PATCH] chore: return toolMapping in chat progress message

---
 pkg/openai/client.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/pkg/openai/client.go b/pkg/openai/client.go
index 61a7ec77..af518f93 100644
--- a/pkg/openai/client.go
+++ b/pkg/openai/client.go
@@ -356,6 +356,7 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
 		}
 	}
 
+	toolMapping := map[string]string{}
 	for _, tool := range messageRequest.Tools {
 		var params any = tool.Function.Parameters
 		if tool.Function.Parameters == nil || len(tool.Function.Parameters.Properties) == 0 {
@@ -365,6 +366,10 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
 			}
 		}
 
+		if tool.Function.ToolID != "" {
+			toolMapping[tool.Function.Name] = tool.Function.ToolID
+		}
+
 		request.Tools = append(request.Tools, openai.Tool{
 			Type: openai.ToolTypeFunction,
 			Function: &openai.FunctionDefinition{
@@ -378,7 +383,10 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
 	id := counter.Next()
 	status <- types.CompletionStatus{
 		CompletionID: id,
-		Request:      request,
+		Request: map[string]any{
+			"chatCompletion": request,
+			"toolMapping":    toolMapping,
+		},
 	}
 
 	var cacheResponse bool