From d4f0a5355ebdb20f54ed59fa3d927adb4f65439a Mon Sep 17 00:00:00 2001
From: Yang Luo <hsluoyz@qq.com>
Date: Sat, 30 Sep 2023 17:29:31 +0800
Subject: [PATCH] Fix bug in getNearestVectors()

---
 model/openai.go               |  8 ++++++--
 model/openrouter.go           |  8 ++++++--
 object/search_default_util.go |  3 ++-
 web/src/MessageListPage.js    | 20 +++++++++++++++++++-
 4 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/model/openai.go b/model/openai.go
index 71454db..a0241e7 100644
--- a/model/openai.go
+++ b/model/openai.go
@@ -104,12 +104,16 @@ func (p *OpenAiModelProvider) QueryText(question string, writer io.Writer, build
 
 	// https://platform.openai.com/tokenizer
 	// https://github.com/pkoukk/tiktoken-go#available-encodings
-	promptTokens, err := GetTokenSize(model, question)
+	tokenCount, err := GetTokenSize(model, question)
 	if err != nil {
 		return err
 	}
 
-	maxTokens := p.GetMaxTokens() - promptTokens
+	maxTokens := p.GetMaxTokens() - tokenCount
+	if maxTokens < 0 {
+		return fmt.Errorf("The token count: [%d] exceeds the model: [%s]'s maximum token count: [%d]", tokenCount, model, p.GetMaxTokens())
+	}
+
 	temperature := p.temperature
 	topP := p.topP
 	frequencyPenalty := p.frequencyPenalty
diff --git a/model/openrouter.go b/model/openrouter.go
index 888e3f5..80d276d 100644
--- a/model/openrouter.go
+++ b/model/openrouter.go
@@ -72,12 +72,16 @@ func (p *OpenRouterModelProvider) QueryText(question string, writer io.Writer, b
 		model = openrouter.Gpt35Turbo
 	}
 
-	promptTokens, err := GetTokenSize(model, question)
+	tokenCount, err := GetTokenSize(model, question)
 	if err != nil {
 		return err
 	}
 
-	maxTokens := 4097 - promptTokens
+	maxTokens := 4097 - tokenCount
+	if maxTokens < 0 {
+		return fmt.Errorf("The token count: [%d] exceeds the model: [%s]'s maximum token count: [%d]", tokenCount, model, 4097)
+	}
+
 	temperature := p.temperature
 	topP := p.topP
 
diff --git a/object/search_default_util.go b/object/search_default_util.go
index 8f4d656..59a61f2 100644
--- a/object/search_default_util.go
+++ b/object/search_default_util.go
@@ -70,5 +70,6 @@ func getNearestVectors(target []float32, vectors [][]float32, n int) []Similarit
 		n = len(vectors)
 	}
 
-	return similarities
+	res := similarities[:n]
+	return res
 }
diff --git a/web/src/MessageListPage.js b/web/src/MessageListPage.js
index bfc1479..e3a356c 100644
--- a/web/src/MessageListPage.js
+++ b/web/src/MessageListPage.js
@@ -14,7 +14,7 @@
 
 import React from "react";
 import {Link} from "react-router-dom";
-import {Button, Popconfirm, Table} from "antd";
+import {Button, Popconfirm, Table, Tag} from "antd";
 import * as Setting from "./Setting";
 import * as MessageBackend from "./backend/MessageBackend";
 import moment from "moment";
@@ -175,6 +175,24 @@ class MessageListPage extends React.Component {
         width: "400px",
         sorter: (a, b) => a.text.localeCompare(b.text),
       },
+      {
+        title: i18next.t("message:Knowledge"),
+        dataIndex: "knowledge",
+        key: "knowledge",
+        width: "200px",
+        sorter: (a, b) => a.knowledge.localeCompare(b.knowledge),
+        render: (text, record, index) => {
+          return record.vectorScores?.map(vectorScore => {
+            return (
+              <Link key={vectorScore.vector} to={`/vectors/${vectorScore.vector}`}>
+                <Tag color={"processing"}>
+                  {vectorScore.score}
+                </Tag>
+              </Link>
+            );
+          });
+        },
+      },
       {
         title: i18next.t("general:Action"),
         dataIndex: "action",