From 65565406554219bba8c7ed07261125741149aaa8 Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Thu, 19 Dec 2024 16:21:24 -0800
Subject: [PATCH] User interface prototype

---
 api/types.go     | 21 +++++++++++++--------
 server/routes.go | 40 +++++++++++++++++++++++++++++++++-------
 2 files changed, 46 insertions(+), 15 deletions(-)

diff --git a/api/types.go b/api/types.go
index c9a9a6041..d53742ee5 100644
--- a/api/types.go
+++ b/api/types.go
@@ -103,14 +103,18 @@ type ChatRequest struct {
 	// Tools is an optional list of tools the model has access to.
 	Tools `json:"tools,omitempty"`
 
-	// DryRun when true will prepare and validate the request but stop before sending it to the model.
-	// This allows inspecting how the prompt would be constructed and what parameters would be used.
-	DryRun bool `json:"dry_run,omitempty"`
+	Debug *Debug `json:"debug,omitempty"`
+
+	Dry bool `json:"dry,omitempty"`
 
 	// Options lists model-specific options.
 	Options map[string]interface{} `json:"options"`
 }
 
+type Debug struct {
+	Include []string `json:"include,omitempty"`
+}
+
 type Tools []Tool
 
 func (t Tools) String() string {
@@ -189,11 +193,12 @@ func (t *ToolFunction) String() string {
 // ChatResponse is the response returned by [Client.Chat]. Its fields are
 // similar to [GenerateResponse].
 type ChatResponse struct {
-	Model        string    `json:"model"`
-	CreatedAt    time.Time `json:"created_at"`
-	Message      Message   `json:"message"`
-	DoneReason   string    `json:"done_reason,omitempty"`
-	DryRunOutput string    `json:"dry_run_output,omitempty"`
+	Model      string    `json:"model"`
+	CreatedAt  time.Time `json:"created_at"`
+	Message    Message   `json:"message"`
+	DoneReason string    `json:"done_reason,omitempty"`
+
+	Debug map[string]any `json:"debug,omitempty"`
 
 	Done bool `json:"done"`
 
diff --git a/server/routes.go b/server/routes.go
index 06db0bbed..98be80d73 100644
--- a/server/routes.go
+++ b/server/routes.go
@@ -1539,14 +1539,30 @@ func (s *Server) ChatHandler(c *gin.Context) {
 		return
 	}
 
-	if req.DryRun {
+	if req.Dry {
+		var debug map[string]any
+		if req.Debug != nil && req.Debug.Include != nil && slices.Contains(req.Debug.Include, "prompt") {
+			debug = map[string]any{"prompt": prompt}
+		}
+		tokens, err := r.Tokenize(c.Request.Context(), prompt)
+		if err != nil {
+			slog.Error("tokenize error", "error", err)
+			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+			return
+		}
 		c.JSON(http.StatusOK, api.ChatResponse{
-			Model:        req.Model,
-			CreatedAt:    time.Now().UTC(),
-			Message:      api.Message{Role: "assistant", Content: ""},
-			Done:         true,
-			DoneReason:   "dry_run",
-			DryRunOutput: prompt,
+			Model:      req.Model,
+			CreatedAt:  time.Now().UTC(),
+			Message:    api.Message{Role: "assistant", Content: ""},
+			Done:       true,
+			DoneReason: "dry_run",
+			Debug:      debug,
+			Metrics: api.Metrics{
+				PromptEvalCount:    len(tokens),
+				PromptEvalDuration: 0,
+				EvalCount:          0,
+				EvalDuration:       0,
+			},
 		})
 		return
 	}
@@ -1583,6 +1599,16 @@ func (s *Server) ChatHandler(c *gin.Context) {
 				res.LoadDuration = checkpointLoaded.Sub(checkpointStart)
 			}
 
+			if req.Debug != nil && req.Debug.Include != nil && slices.Contains(req.Debug.Include, "prompt") {
+				res.Debug = map[string]any{"prompt": prompt}
+				if req.Stream != nil && !*req.Stream {
+					tempMsg := res.Message
+					res.Message = api.Message{Role: "assistant", Content: ""}
+					ch <- res
+					res.Message = tempMsg
+				}
+			}
+
 			// TODO: tool call checking and filtering should be moved outside of this callback once streaming
 			// however this was a simple change for now without reworking streaming logic of this (and other)
 			// handlers