fix: add unmarshall func to reponse

Signed-off-by: bjwswang <[email protected]>
kubeagi · Sep 4, 2023 · abcbdc4 · abcbdc4
1 parent aa617ad
commit abcbdc4
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/pkg/llms/llms.go b/pkg/llms/llms.go
@@ -41,6 +41,7 @@ type Response interface {
 	Type() LLMType
 	String() string
 	Bytes() []byte
+	Unmarshall([]byte) error
 }
 
 type UnknowLLM struct{}

diff --git a/pkg/llms/openai/response.go b/pkg/llms/openai/response.go
@@ -19,8 +19,9 @@ package openai
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/kubeagi/arcadia/pkg/llms"
 	"net/http"
+
+	"github.com/kubeagi/arcadia/pkg/llms"
 )
 
 type Response struct {
@@ -46,6 +47,10 @@ func (response *Response) String() string {
 	return string(response.Bytes())
 }
 
+func (response *Response) Unmarshall(bytes []byte) error {
+	return json.Unmarshal(bytes, response)
+}
+
 func parseHTTPResponse(resp *http.Response) (*Response, error) {
 	if resp.StatusCode != http.StatusOK {
 		return nil, fmt.Errorf("exception: %s", resp.Status)

diff --git a/pkg/llms/zhipuai/response.go b/pkg/llms/zhipuai/response.go
@@ -36,6 +36,10 @@ type EmbeddingResponse struct {
 	Success bool           `json:"success"`
 }
 
+func (response *Response) Unmarshall(bytes []byte) error {
+	return json.Unmarshal(response.Bytes(), response)
+}
+
 func (response *Response) Type() llms.LLMType {
 	return llms.ZhiPuAI
 }