Skip to content

Commit b82c0b5

Browse files
committed
feat(wasip2/llm): add support for llm
Signed-off-by: Rajat Jindal <[email protected]>
1 parent 9f90e21 commit b82c0b5

File tree

6 files changed

+154
-0
lines changed

6 files changed

+154
-0
lines changed

v2/examples/llm/.gitignore

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
main.wasm
2+
.spin/

v2/examples/llm/go.mod

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
module github.com/spinframework/spin-go-sdk/v2/examples/llm
2+
3+
go 1.24.1
4+
5+
require github.com/spinframework/spin-go-sdk/v2 v2.0.0
6+
7+
require (
8+
github.com/julienschmidt/httprouter v1.3.0 // indirect
9+
go.bytecodealliance.org/cm v0.2.2 // indirect
10+
)
11+
12+
replace github.com/spinframework/spin-go-sdk/v2 => ../../

v2/examples/llm/go.sum

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
github.com/julienschmidt/httprouter v1.3.0 h1:U0609e9tgbseu3rBINet9P48AI/D3oJs4dN7jwJOQ1U=
2+
github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8IZAc4RVcycCCAKdM=
3+
go.bytecodealliance.org/cm v0.2.2 h1:M9iHS6qs884mbQbIjtLX1OifgyPG9DuMs2iwz8G4WQA=
4+
go.bytecodealliance.org/cm v0.2.2/go.mod h1:JD5vtVNZv7sBoQQkvBvAAVKJPhR/bqBH7yYXTItMfZI=

v2/examples/llm/main.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
7+
spinhttp "github.com/spinframework/spin-go-sdk/v2/http"
8+
"github.com/spinframework/spin-go-sdk/v2/llm"
9+
)
10+
11+
func init() {
12+
spinhttp.Handle(func(w http.ResponseWriter, r *http.Request) {
13+
result, err := llm.Infer("llama2-chat", "Tell me a joke", nil)
14+
if err != nil {
15+
http.Error(w, err.Error(), http.StatusInternalServerError)
16+
return
17+
}
18+
fmt.Printf("Prompt tokens: %d\n", result.Usage.PromptTokenCount)
19+
fmt.Printf("Generated tokens: %d\n", result.Usage.GeneratedTokenCount)
20+
fmt.Fprint(w, result.Text)
21+
fmt.Fprintf(w, "\n\n")
22+
23+
embeddings, err := llm.GenerateEmbeddings("all-minilm-l6-v2", []string{"Hello world"})
24+
if err != nil {
25+
http.Error(w, err.Error(), http.StatusInternalServerError)
26+
return
27+
}
28+
fmt.Printf("%d\n", len(embeddings.Embeddings[0]))
29+
fmt.Printf("Prompt Tokens: %d\n", embeddings.Usage.PromptTokenCount)
30+
31+
})
32+
}
33+
34+
func main() {}

v2/examples/llm/spin.toml

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
spin_manifest_version = 2
2+
3+
[application]
4+
authors = ["Fermyon Engineering <[email protected]>"]
5+
description = "Simple example using the llm sdk."
6+
name = "llm-example"
7+
version = "0.1.0"
8+
9+
[[trigger.http]]
10+
route = "/..."
11+
component = "llm"
12+
13+
[component.llm]
14+
source = "main.wasm"
15+
allowed_outbound_hosts = []
16+
ai_models = ["llama2-chat", "all-minilm-l6-v2"]
17+
[component.llm.build]
18+
command = "tinygo build -target=wasip2 --wit-package $(go list -mod=readonly -m -f '{{.Dir}}' github.com/spinframework/spin-go-sdk/v2)/wit --wit-world http-trigger -gc=leaking -no-debug -o main.wasm main.go"
19+
watch = ["**/*.go", "go.mod"]

v2/llm/llm.go

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package llm
2+
3+
import (
4+
"fmt"
5+
6+
"github.com/spinframework/spin-go-sdk/v2/internal/fermyon/spin/v2.0.0/llm"
7+
"go.bytecodealliance.org/cm"
8+
)
9+
10+
// The model use for inferencing
11+
const (
12+
Llama2Chat InferencingModel = "llama2-chat"
13+
CodellamaInstruct InferencingModel = "codellama-instruct"
14+
)
15+
16+
type InferencingParams llm.InferencingParams
17+
type InferencingResult llm.InferencingResult
18+
type InferencingModel llm.InferencingModel
19+
20+
type EmbeddingsResult struct {
21+
// Embeddings are the embeddings generated by the request.
22+
Embeddings [][]float32
23+
// Usage is usage related to an embeddings generation request.
24+
Usage *EmbeddingsUsage
25+
}
26+
27+
type EmbeddingsUsage struct {
28+
// PromptTokenCount is number of tokens in the prompt.
29+
PromptTokenCount int
30+
}
31+
32+
// Infer performs inferencing using the provided model and prompt with the
33+
// given optional parameters.
34+
func Infer(model string, prompt string, params *InferencingParams) (InferencingResult, error) {
35+
var iparams = cm.None[llm.InferencingParams]()
36+
if params != nil {
37+
iparams = cm.Some(llm.InferencingParams(*params))
38+
}
39+
40+
result := llm.Infer(llm.InferencingModel(model), prompt, iparams)
41+
if result.IsErr() {
42+
return InferencingResult{}, errorVariantToError(*result.Err())
43+
}
44+
45+
return InferencingResult(*result.OK()), nil
46+
}
47+
48+
// GenerateEmbeddings generates the embeddings for the supplied list of text.
49+
func GenerateEmbeddings(model InferencingModel, text []string) (*EmbeddingsResult, error) {
50+
result := llm.GenerateEmbeddings(llm.EmbeddingModel(model), cm.ToList(text))
51+
if result.IsErr() {
52+
return &EmbeddingsResult{}, errorVariantToError(*result.Err())
53+
}
54+
55+
llmEmbeddingResult := llm.EmbeddingsResult(*result.OK())
56+
57+
list := llmEmbeddingResult.Embeddings
58+
list2 := list.Slice()
59+
list3 := [][]float32{}
60+
for _, l := range list2 {
61+
list3 = append(list3, l.Slice())
62+
}
63+
64+
return &EmbeddingsResult{
65+
Embeddings: list3,
66+
Usage: &EmbeddingsUsage{
67+
PromptTokenCount: int(llmEmbeddingResult.Usage.PromptTokenCount),
68+
},
69+
}, nil
70+
}
71+
72+
func errorVariantToError(err llm.Error) error {
73+
switch {
74+
case llm.ErrorModelNotSupported() == err:
75+
return fmt.Errorf("model not supported")
76+
case err.RuntimeError() != nil:
77+
return fmt.Errorf("runtime error %s", *err.RuntimeError())
78+
case err.InvalidInput() != nil:
79+
return fmt.Errorf("invalid input %s", *err.InvalidInput())
80+
default:
81+
return fmt.Errorf("no error provided by host implementation")
82+
}
83+
}

0 commit comments

Comments
 (0)