-
Notifications
You must be signed in to change notification settings - Fork 1
/
llm.go
49 lines (40 loc) · 981 Bytes
/
llm.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
package node
import (
"context"
"github.com/primalmotion/simplai/engine"
)
// LLM is a node responsible for running its
// input into an inference engine.
type LLM struct {
engine engine.LLM
*BaseNode
options []engine.LLMOption
}
// NewLLM returns a new LLM using the given engine.
func NewLLM(info Info, engine engine.LLM, options ...engine.LLMOption) *LLM {
return &LLM{
BaseNode: New(info),
engine: engine,
options: options,
}
}
// Execute implements the Node interface.
func (n *LLM) Execute(ctx context.Context, input Input) (string, error) {
opts := append(n.options, input.LLMOptions()...)
if input.Debug() {
opts = append(opts, engine.OptionDebug(true))
}
output, err := n.engine.Infer(ctx, input.Input(), opts...)
if err != nil {
return "", NewError(n, "unable to run llm inference: %w", err)
}
if input.Debug() {
LogNode(n, "4", output)
}
return n.BaseNode.Execute(
ctx,
input.
WithInput(output).
ResetLLMOptions(),
)
}