From 507e4071277e9cd82671e9dc6e641878ee8a2619 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Fri, 14 Feb 2025 15:34:05 +0900
Subject: [PATCH] init: add codes

---
 README.md         |  13 ++++++
 api.go            |  37 +++++++++++++++++
 go.mod            |   3 ++
 model.go          |  12 ++++++
 model/deepseek.go | 100 ++++++++++++++++++++++++++++++++++++++++++++++
 model/utils.go    |  19 +++++++++
 6 files changed, 184 insertions(+)
 create mode 100644 api.go
 create mode 100644 go.mod
 create mode 100644 model.go
 create mode 100644 model/deepseek.go
 create mode 100644 model/utils.go

diff --git a/README.md b/README.md
index 06de139..3e60126 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,15 @@
 # deepinfra
 Call OpenAI compatible APIs, originally designed for DeepInfra.
+
+## Quick Start
+```go
+api := NewAPI(APIDeepInfra, "PUT YOUR API KEY HERE")
+txt, err := api.Request(model.NewDeepSeek(0.7, 0.9, 1024).
+    System("Be a good assistant.").User("Hello"),
+)
+if err != nil {
+    panic(err)
+}
+fmt.Println(txt)
+// Hello! How can I assist you today?
+```
diff --git a/api.go b/api.go
new file mode 100644
index 0000000..2646be6
--- /dev/null
+++ b/api.go
@@ -0,0 +1,37 @@
+package deepinfra
+
+import (
+	"net/http"
+)
+
+const (
+	APIDeepInfra = "https://api.deepinfra.com/v1/openai/chat/completions"
+)
+
+type API struct {
+	api string // api to call
+	key string // key in Authorization: Bearer
+}
+
+func NewAPI(api, key string) API {
+	return API{api: api, key: key}
+}
+
+func (api *API) Request(model Model) (string, error) {
+	req, err := http.NewRequest("POST", api.api, model.Body())
+	if err != nil {
+		return "", err
+	}
+	req.Header.Add("Content-Type", "application/json")
+	req.Header.Add("Authorization", "Bearer "+api.key)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+	err = model.Parse(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	return model.Output(), nil
+}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..570f13f
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,3 @@
+module github.com/fumiama/deepinfra
+
+go 1.23.2
diff --git a/model.go b/model.go
new file mode 100644
index 0000000..77e3af9
--- /dev/null
+++ b/model.go
@@ -0,0 +1,12 @@
+package deepinfra
+
+import (
+	"bytes"
+	"io"
+)
+
+type Model interface {
+	Body() *bytes.Buffer
+	Parse(io.Reader) error
+	Output() string
+}
diff --git a/model/deepseek.go b/model/deepseek.go
new file mode 100644
index 0000000..c584767
--- /dev/null
+++ b/model/deepseek.go
@@ -0,0 +1,100 @@
+package model
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+)
+
+const (
+	modelDeepDeek = "deepseek-ai/DeepSeek-R1"
+)
+
+// DeepSeek as an example.
+type DeepSeek struct {
+	// callback only
+	ID      string   `json:"id,omitempty"`
+	Object  string   `json:"object,omitempty"`
+	Created int      `json:"created,omitempty"`
+	Choices []Choice `json:"choices,omitempty"`
+	// callback/request
+	Model       string    `json:"model"`
+	Messages    []Message `json:"messages"`
+	Temperature float32   `json:"temperature"` // Temperature 0.7
+	TopP        float32   `json:"top_p"`       // TopP 0.9
+	MaxTokens   int       `json:"max_tokens"`  // MaxTokens 16384
+
+}
+
+type Message struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type Choice struct {
+	Index        int     `json:"index"`
+	Message      Message `json:"message"`
+	FinishReason string  `json:"finish_reason"`
+}
+
+// NewDeepSeek 0.7, 0.9
+func NewDeepSeek(temp, topp float32, maxn uint) *DeepSeek {
+	ds := new(DeepSeek)
+	ds.Model = modelDeepDeek
+	ds.Temperature = temp
+	ds.TopP = topp
+	ds.MaxTokens = int(maxn)
+	return ds
+}
+
+func (ds *DeepSeek) Parse(body io.Reader) error {
+	return json.NewDecoder(body).Decode(&ds)
+}
+
+func (ds *DeepSeek) Output() string {
+	if len(ds.Choices) == 0 {
+		return ""
+	}
+	return CutLast(ds.Choices[len(ds.Choices)-1].Message.Content, SeparatorThink)
+}
+
+func (ds *DeepSeek) OutputRaw() string {
+	if len(ds.Choices) == 0 {
+		return ""
+	}
+	return ds.Choices[len(ds.Choices)-1].Message.Content
+}
+
+func (ds *DeepSeek) System(prompt string) *DeepSeek {
+	ds.Messages = make([]Message, 1, 8)
+	ds.Messages[0] = Message{
+		Role:    "system",
+		Content: prompt,
+	}
+	return ds
+}
+
+func (ds *DeepSeek) User(prompt string) *DeepSeek {
+	ds.Messages = append(ds.Messages, Message{
+		Role:    "user",
+		Content: prompt,
+	})
+	return ds
+}
+
+func (ds *DeepSeek) Assistant(prompt string) *DeepSeek {
+	ds.Messages = append(ds.Messages, Message{
+		Role:    "assistant",
+		Content: prompt,
+	})
+	return ds
+}
+
+func (ds *DeepSeek) Body() *bytes.Buffer {
+	w := bytes.NewBuffer(make([]byte, 0, 16384))
+	err := json.NewEncoder(w).Encode(ds)
+	if err != nil {
+		panic(err)
+	}
+	return w
+}
diff --git a/model/utils.go b/model/utils.go
new file mode 100644
index 0000000..3e8d5e7
--- /dev/null
+++ b/model/utils.go
@@ -0,0 +1,19 @@
+package model
+
+import "strings"
+
+const (
+	SeparatorThink = "</think>"
+)
+
+func CutLast(txt, sep string) string {
+	a := strings.LastIndex(txt, sep)
+	if a < 0 {
+		return ""
+	}
+	a += len(sep)
+	if a >= len(txt) {
+		return ""
+	}
+	return strings.TrimSpace(txt[a:])
+}