让 OpenAI 更 Open,在 ChatGPT 里自由接入数据源( 三 )


下面我们聊聊,如何封装这样一个简单的数据源,让能够输出一些不一样的东西 。
封装自定义数据源:Flag
Flag的数据源封装实现,存放在后端项目的 /flag- 里,关键实现代码行数不到 200 行 。
参考官方文档,一个完整的 Flag图片生成流程中,需要根据我们申请的 API Key 去换服务调用所需要的 Token,最后携带 Token 去调用图片生成接口即可 。
我们先来实现根据 API Key 换 Token 的逻辑:
package FlagStudioimport ("encoding/json""fmt""io""net/http")const API_GET_TOKEN = "https://flagopen.baai.ac.cn/flagStudio/auth/getToken"type ResponseToken struct {Code int `json:"code"`Data struct {Token string `json:"token"`} `json:"data"`}// parseToken parses the token from the response bodyfunc parseToken(buf []byte) (string, error) {var data ResponseTokenerr := json.Unmarshal(buf, &data)if err != nil {return "", err}if data.Code != 200 || data.Data.Token == "" {return "", fmt.Errorf("FlagStudio API, Get Token error, Code %d\n, Token: %s", data.Code, data.Data.Token)}return data.Data.Token, nil}// get token from the APIfunc GetToken(apikey string) (string, error) {req, err := http.NewRequest("GET", API_GET_TOKEN, nil)if err != nil {return "", fmt.Errorf("FlagStudio API, Error initializing network components, err: %v", err)}req.Header.Set("Accept", "application/json")req.Header.Set("Content-Type", "application/json")q := req.URL.Query()q.Add("apikey", apikey)req.URL.RawQuery = q.Encode()client := &http.Client{}resp, err := client.Do(req)if err != nil {fmt.Println(err)return "", fmt.Errorf("FlagStudio API, Error sending request, err: %v", err)}defer resp.Body.Close()body, err := io.ReadAll(resp.Body)if err != nil {fmt.Println(err)return "", fmt.Errorf("FlagStudio API, Error reading response, err: %v", err)}token, err := parseToken(body)if err != nil {fmt.Println(err)return "", fmt.Errorf("FlagStudio API, Error parsing response, err: %v", err)}return token, nil}
上面的代码中,我们实现了一个非常基础的 HTTP 调用,以及对服务端返回的 JSON 内容的解析,如果 API Key 正确、网络没有异常的情况下,函数运行结束,我们将得到生成图片所需要的 Token 。
接下来,我们来实现主要逻辑,图片生成接口调用:
package FlagStudioimport ("encoding/json""fmt""io""net/http""strings""github.com/soulteary/sparrow/internal/define")type TextToImage struct {Promptstring`json:"prompt"`GuidanceScalefloat64 `json:"guidance_scale"`Heightint`json:"height"`NegativePrompts string`json:"negative_prompts"`Samplerstring`json:"sampler"`Seedint`json:"seed"`Stepsint`json:"steps"`Stylestring`json:"style"`Upsampleint`json:"upsample"`Widthint`json:"width"`}const API_TEXT_TO_IMAGE = "https://flagopen.baai.ac.cn/flagStudio/v1/text2img"var FS_STYLES = []string{"国画", "写实主义", "虚幻引擎", "黑白插画", "版绘", "低聚", "工业霓虹", "电影艺术", "史诗大片", "暗黑", "涂鸦", "漫画场景", "特写", "儿童画", "油画", "水彩画", "素描", "卡通画", "浮世绘", "赛博朋克", "吉卜力", "哑光", "现代中式", "相机", "CG渲染", "动漫", "霓虹游戏", "蒸汽波", "宝可梦", "火影忍者", "圣诞老人", "个人特效", "通用漫画", "Momoko", "MJ风格", "剪纸", "齐白石", "张大千", "丰子恺", "毕加索", "梵高", "塞尚", "莫奈", "马克·夏加尔", "丢勒", "米开朗基罗", "高更", "爱德华·蒙克", "托马斯·科尔", "安迪·霍尔", "新海诚", "倪传婧", "村上隆", "黄光剑", "吴冠中", "林风眠", "木内达朗", "萨雷尔", "杜拉克", "比利宾", "布拉德利", "普罗旺森", "莫比乌斯", "格里斯利", "比普", "卡尔·西松", "玛丽·布莱尔", "埃里克·卡尔", "扎哈·哈迪德", "包豪斯", "英格尔斯", "RHADS", "阿泰·盖兰", "俊西", "坎皮恩", "德尚鲍尔", "库沙特", "雷诺阿"}func GetRandomStyle() string {return FS_STYLES[define.GetRandomNumber(0, len(FS_STYLES)-1)]}func GenerateImageByText(s string) string {data := TextToImage{Prompt:s,GuidanceScale:7.5,Width:512,Height:512,NegativePrompts: "",Sampler:"ddim",Seed:1024,Steps:50,Style:GetRandomStyle(),Upsample:1,}payload, err := define.MakeJSON(data)if err != nil {return fmt.Sprintf("FlagStudio API, An error occurred while preparing to enter data: %v", err)}token, err := GetToken(define.FLAGSTUDIO_API_KEY)if err != nil {return fmt.Sprintf("FlagStudio API, An error occurred while getting the token: %v", err)}req, err := http.NewRequest("POST", API_TEXT_TO_IMAGE, strings.NewReader(payload))if err != nil {return fmt.Sprintf("FlagStudio API, An error occurred while initializing network components: %v", err)}req.Header.Set("Accept", "application/json")req.Header.Set("Content-Type", "application/json")req.Header.Add("token", token)client := &http.Client{}resp, err := client.Do(req)if err != nil {return fmt.Sprintf("FlagStudio API, An error occurred while sending request: %v", err)}defer resp.Body.Close()body, err := io.ReadAll(resp.Body)if err != nil {return fmt.Sprintf("FlagStudio API, An error occurred while reading response: %v", err)}base64Image, err := parseTextToImage(body)if err != nil {return fmt.Sprintf("FlagStudio API, An error occurred while parsing response: %v", err)}return `![](data:image/png;base64,` + base64Image + `)`}type ResponseTextToImage struct {Code int`json:"code"`Data string `json:"data"`Nsfw int`json:"nsfw"`}// parseToken parses the token from the response bodyfunc parseTextToImage(buf []byte) (string, error) {var data ResponseTextToImageerr := json.Unmarshal(buf, &data)if err != nil {return "", err}if data.Code != 200 || data.Data =http://www.kingceram.com/post/="" {return "", fmt.Errorf("FlagStudio API, Get Result error, Code %d", data.Code)}if data.Nsfw != 0 {return "", fmt.Errorf("FlagStudio API, Get Token error, Code %d\n, NSFW: %d", data.Code, data.Nsfw)}return data.Data, nil}