From a49e83179117ec451f473be55207056f3e42a3c6 Mon Sep 17 00:00:00 2001
From: ivynya <ivy@ivy.direct>
Date: Sat, 30 Dec 2023 19:08:34 -0800
Subject: [PATCH] fix: summarize return info

---
 client/summarize.go | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/client/summarize.go b/client/summarize.go
index 635e02a..1cbf40e 100644
--- a/client/summarize.go
+++ b/client/summarize.go
@@ -1,10 +1,12 @@
 package main
 
 import (
+	"encoding/json"
 	"strconv"
 
 	"github.com/gorilla/websocket"
 	"github.com/ivynya/illm/internal"
+	"github.com/ivynya/illm/ollama"
 	"github.com/kkdai/youtube/v2"
 )
 
@@ -22,13 +24,22 @@ func summarize(c *websocket.Conn, req *internal.Request) (bool, error) {
 		return false, err
 	}
 
-	resp, err := encodeRequest(req.Tag, "response", "Video: "+video.Title+"\nTranscript length: "+strconv.Itoa(len(transcript.String()))+"\n\n")
+	info := &ollama.GenerateResponse{
+		Model:    req.Generate.Model,
+		Response: "Video: `" + video.Title + "`\nTranscript length: `" + strconv.Itoa(len(transcript.String())) + "`\n\n",
+		Done:     false,
+	}
+	infoJson, err := json.Marshal(info)
+	if err != nil {
+		return false, err
+	}
+	infoResp, err := encodeRequest(req.Tag, "response", string(infoJson))
 	if err != nil {
 		return false, err
 	}
-	c.WriteMessage(websocket.TextMessage, resp)
+	c.WriteMessage(websocket.TextMessage, infoResp)
 
-	req.Generate.Prompt = "Summarize the following video: " + video.Title + "\n\n" + transcript.String() + "\n\nSummary:"
+	req.Generate.Prompt = "Summarize the following video. Only include information from the video in your response. Video: " + video.Title + "\n\n" + transcript.String() + "\n\nSummary:"
 	req.Generate.Context = []int{}
 
 	complete, err := generate(c, req)