add llm adapter interface and implementations
This commit is contained in:
leonardotrapani
@@ -0,0 +1,73 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
// GroqAdapter implements Adapter using Groq's OpenAI-compatible API
|
||||
type GroqAdapter struct {
|
||||
client *openai.Client
|
||||
config Config
|
||||
}
|
||||
|
||||
// NewGroqAdapter creates a new Groq LLM adapter
|
||||
func NewGroqAdapter(cfg Config) *GroqAdapter {
|
||||
clientConfig := openai.DefaultConfig(cfg.APIKey)
|
||||
clientConfig.BaseURL = "https://api.groq.com/openai/v1"
|
||||
return &GroqAdapter{
|
||||
client: openai.NewClientWithConfig(clientConfig),
|
||||
config: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *GroqAdapter) Process(ctx context.Context, text string) (string, error) {
|
||||
if text == "" {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
opts := PostProcessingOptions{
|
||||
RemoveStutters: a.config.RemoveStutters,
|
||||
AddPunctuation: a.config.AddPunctuation,
|
||||
FixGrammar: a.config.FixGrammar,
|
||||
RemoveFillerWords: a.config.RemoveFillerWords,
|
||||
}
|
||||
|
||||
systemPrompt := BuildSystemPrompt(opts, a.config.Keywords)
|
||||
userPrompt := BuildUserPrompt(text, a.config.CustomPrompt)
|
||||
|
||||
model := a.config.Model
|
||||
if model == "" {
|
||||
model = "llama-3.3-70b-versatile"
|
||||
}
|
||||
|
||||
req := openai.ChatCompletionRequest{
|
||||
Model: model,
|
||||
Messages: []openai.ChatCompletionMessage{
|
||||
{Role: openai.ChatMessageRoleSystem, Content: systemPrompt},
|
||||
{Role: openai.ChatMessageRoleUser, Content: userPrompt},
|
||||
},
|
||||
Temperature: 0.3, // Low temperature for consistent cleanup
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := a.client.CreateChatCompletion(ctx, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("groq-llm-adapter: API call failed after %v: %v", duration, err)
|
||||
return "", fmt.Errorf("groq chat completion: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Choices) == 0 {
|
||||
return "", fmt.Errorf("groq chat completion: no response choices")
|
||||
}
|
||||
|
||||
result := resp.Choices[0].Message.Content
|
||||
log.Printf("groq-llm-adapter: processed in %v: %q -> %q", duration, text, result)
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
// OpenAIAdapter implements Adapter using OpenAI's chat completions API
|
||||
type OpenAIAdapter struct {
|
||||
client *openai.Client
|
||||
config Config
|
||||
}
|
||||
|
||||
// NewOpenAIAdapter creates a new OpenAI LLM adapter
|
||||
func NewOpenAIAdapter(cfg Config) *OpenAIAdapter {
|
||||
return &OpenAIAdapter{
|
||||
client: openai.NewClient(cfg.APIKey),
|
||||
config: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *OpenAIAdapter) Process(ctx context.Context, text string) (string, error) {
|
||||
if text == "" {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
opts := PostProcessingOptions{
|
||||
RemoveStutters: a.config.RemoveStutters,
|
||||
AddPunctuation: a.config.AddPunctuation,
|
||||
FixGrammar: a.config.FixGrammar,
|
||||
RemoveFillerWords: a.config.RemoveFillerWords,
|
||||
}
|
||||
|
||||
systemPrompt := BuildSystemPrompt(opts, a.config.Keywords)
|
||||
userPrompt := BuildUserPrompt(text, a.config.CustomPrompt)
|
||||
|
||||
model := a.config.Model
|
||||
if model == "" {
|
||||
model = "gpt-4o-mini"
|
||||
}
|
||||
|
||||
req := openai.ChatCompletionRequest{
|
||||
Model: model,
|
||||
Messages: []openai.ChatCompletionMessage{
|
||||
{Role: openai.ChatMessageRoleSystem, Content: systemPrompt},
|
||||
{Role: openai.ChatMessageRoleUser, Content: userPrompt},
|
||||
},
|
||||
Temperature: 0.3, // Low temperature for consistent cleanup
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := a.client.CreateChatCompletion(ctx, req)
|
||||
duration := time.Since(start)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("openai-llm-adapter: API call failed after %v: %v", duration, err)
|
||||
return "", fmt.Errorf("openai chat completion: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Choices) == 0 {
|
||||
return "", fmt.Errorf("openai chat completion: no response choices")
|
||||
}
|
||||
|
||||
result := resp.Choices[0].Message.Content
|
||||
log.Printf("openai-llm-adapter: processed in %v: %q -> %q", duration, text, result)
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Adapter interface for LLM text processing
|
||||
type Adapter interface {
|
||||
Process(ctx context.Context, text string) (string, error)
|
||||
}
|
||||
|
||||
// Config holds LLM adapter configuration
|
||||
type Config struct {
|
||||
Provider string
|
||||
APIKey string
|
||||
Model string
|
||||
RemoveStutters bool
|
||||
AddPunctuation bool
|
||||
FixGrammar bool
|
||||
RemoveFillerWords bool
|
||||
CustomPrompt string
|
||||
Keywords []string
|
||||
}
|
||||
|
||||
// NewAdapter creates an LLM adapter based on the provider
|
||||
func NewAdapter(cfg Config) (Adapter, error) {
|
||||
switch cfg.Provider {
|
||||
case "openai":
|
||||
if cfg.APIKey == "" {
|
||||
return nil, fmt.Errorf("OpenAI API key required")
|
||||
}
|
||||
return NewOpenAIAdapter(cfg), nil
|
||||
case "groq":
|
||||
if cfg.APIKey == "" {
|
||||
return nil, fmt.Errorf("Groq API key required")
|
||||
}
|
||||
return NewGroqAdapter(cfg), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported LLM provider: %s", cfg.Provider)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildSystemPrompt(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts PostProcessingOptions
|
||||
keywords []string
|
||||
contains []string
|
||||
}{
|
||||
{
|
||||
name: "all options enabled",
|
||||
opts: PostProcessingOptions{
|
||||
RemoveStutters: true,
|
||||
AddPunctuation: true,
|
||||
FixGrammar: true,
|
||||
RemoveFillerWords: true,
|
||||
},
|
||||
keywords: nil,
|
||||
contains: []string{
|
||||
"Remove stutters",
|
||||
"Add proper punctuation",
|
||||
"Fix grammar",
|
||||
"Remove filler words",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "only grammar",
|
||||
opts: PostProcessingOptions{
|
||||
FixGrammar: true,
|
||||
},
|
||||
keywords: nil,
|
||||
contains: []string{
|
||||
"Fix grammar",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "with keywords",
|
||||
opts: PostProcessingOptions{
|
||||
RemoveStutters: true,
|
||||
},
|
||||
keywords: []string{"Kubernetes", "TypeScript", "hyprvoice"},
|
||||
contains: []string{
|
||||
"Kubernetes",
|
||||
"TypeScript",
|
||||
"hyprvoice",
|
||||
"Context keywords",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "no options - should have default",
|
||||
opts: PostProcessingOptions{},
|
||||
keywords: nil,
|
||||
contains: []string{
|
||||
"Clean up the text",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
result := BuildSystemPrompt(tc.opts, tc.keywords)
|
||||
for _, expected := range tc.contains {
|
||||
if !strings.Contains(result, expected) {
|
||||
t.Errorf("expected prompt to contain %q, got: %s", expected, result)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildUserPrompt(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
customPrompt string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "no custom prompt",
|
||||
text: "hello world",
|
||||
customPrompt: "",
|
||||
expected: "hello world",
|
||||
},
|
||||
{
|
||||
name: "with custom prompt",
|
||||
text: "hello world",
|
||||
customPrompt: "Format as a haiku",
|
||||
expected: "Format as a haiku\n\nText to process:\nhello world",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
result := BuildUserPrompt(tc.text, tc.customPrompt)
|
||||
if result != tc.expected {
|
||||
t.Errorf("expected %q, got %q", tc.expected, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAdapter(t *testing.T) {
|
||||
// Test OpenAI adapter creation
|
||||
openaiCfg := Config{
|
||||
Provider: "openai",
|
||||
APIKey: "sk-test-key",
|
||||
Model: "gpt-4o-mini",
|
||||
}
|
||||
adapter, err := NewAdapter(openaiCfg)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create openai adapter: %v", err)
|
||||
}
|
||||
if _, ok := adapter.(*OpenAIAdapter); !ok {
|
||||
t.Error("expected OpenAIAdapter type")
|
||||
}
|
||||
|
||||
// Test Groq adapter creation
|
||||
groqCfg := Config{
|
||||
Provider: "groq",
|
||||
APIKey: "gsk_test-key",
|
||||
Model: "llama-3.3-70b-versatile",
|
||||
}
|
||||
adapter, err = NewAdapter(groqCfg)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create groq adapter: %v", err)
|
||||
}
|
||||
if _, ok := adapter.(*GroqAdapter); !ok {
|
||||
t.Error("expected GroqAdapter type")
|
||||
}
|
||||
|
||||
// Test missing API key
|
||||
noKeyCfg := Config{
|
||||
Provider: "openai",
|
||||
APIKey: "",
|
||||
}
|
||||
_, err = NewAdapter(noKeyCfg)
|
||||
if err == nil {
|
||||
t.Error("expected error for missing API key")
|
||||
}
|
||||
|
||||
// Test unsupported provider
|
||||
badCfg := Config{
|
||||
Provider: "unsupported",
|
||||
APIKey: "key",
|
||||
}
|
||||
_, err = NewAdapter(badCfg)
|
||||
if err == nil {
|
||||
t.Error("expected error for unsupported provider")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// PostProcessingOptions controls which cleanup operations to request
|
||||
type PostProcessingOptions struct {
|
||||
RemoveStutters bool
|
||||
AddPunctuation bool
|
||||
FixGrammar bool
|
||||
RemoveFillerWords bool
|
||||
}
|
||||
|
||||
// BuildSystemPrompt generates the system prompt for text cleanup
|
||||
func BuildSystemPrompt(opts PostProcessingOptions, keywords []string) string {
|
||||
var tasks []string
|
||||
|
||||
if opts.RemoveStutters {
|
||||
tasks = append(tasks, "Remove stutters and repeated words/phrases")
|
||||
}
|
||||
if opts.AddPunctuation {
|
||||
tasks = append(tasks, "Add proper punctuation")
|
||||
}
|
||||
if opts.FixGrammar {
|
||||
tasks = append(tasks, "Fix grammar errors")
|
||||
}
|
||||
if opts.RemoveFillerWords {
|
||||
tasks = append(tasks, "Remove filler words (um, uh, like, you know, etc.)")
|
||||
}
|
||||
|
||||
// If no tasks, just clean up generally
|
||||
if len(tasks) == 0 {
|
||||
tasks = append(tasks, "Clean up the text while preserving meaning")
|
||||
}
|
||||
|
||||
prompt := "You are a text cleanup assistant. Your job is to clean up speech-to-text transcriptions.\n\n"
|
||||
prompt += "Tasks:\n"
|
||||
for _, task := range tasks {
|
||||
prompt += fmt.Sprintf("- %s\n", task)
|
||||
}
|
||||
|
||||
prompt += "\nRules:\n"
|
||||
prompt += "- Preserve the original meaning and intent\n"
|
||||
prompt += "- Keep the same language as the input\n"
|
||||
prompt += "- Do not add any new information\n"
|
||||
prompt += "- Do not remove meaningful content\n"
|
||||
prompt += "- Output ONLY the cleaned text, nothing else\n"
|
||||
prompt += "- If the input is empty or nonsensical, return it as-is\n"
|
||||
|
||||
if len(keywords) > 0 {
|
||||
prompt += fmt.Sprintf("\nContext keywords (use correct spelling for these terms): %s\n", strings.Join(keywords, ", "))
|
||||
}
|
||||
|
||||
return prompt
|
||||
}
|
||||
|
||||
// BuildUserPrompt generates the user prompt with the text to process
|
||||
func BuildUserPrompt(text string, customPrompt string) string {
|
||||
if customPrompt != "" {
|
||||
return fmt.Sprintf("%s\n\nText to process:\n%s", customPrompt, text)
|
||||
}
|
||||
return text
|
||||
}
|
||||
Reference in New Issue
Block a user