Add llmgw package

pkg/agent/agent.go

@@ -0,0 +1,54 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package agent
+
+type (
+	Agent struct {
+		Name         string
+		Instructions string
+		Model        string
+		OutputType   string
+	}
+
+	AgentOption func(*Agent)
+)
+
+func WithInstructions(instructions string) AgentOption {
+	return func(a *Agent) {
+		a.Instructions = instructions
+	}
+}
+
+func WithModel(model string) AgentOption {
+	return func(a *Agent) {
+		a.Model = model
+	}
+}
+
+func WithOutputType(outputType string) AgentOption {
+	return func(a *Agent) {
+		a.OutputType = outputType
+	}
+}
+
+func NewAgent(name string, opts ...AgentOption) *Agent {
+	agent := &Agent{Name: name}
+
+	for _, opt := range opts {
+		opt(agent)
+	}
+
+	return agent
+}

pkg/agent/tool.go

@@ -0,0 +1,26 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package agent
+
+import "context"
+
+type (
+	Tool[I any, O any] interface {
+		Name() string
+		Description() string
+		Parameters() I
+		Execute(ctx context.Context, input I) (O, error)
+	}
+)

pkg/agents/vetra/ventra.go

@@ -0,0 +1,15 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package vetra

pkg/llmgw/gateway.go

@@ -0,0 +1,54 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package llmgw
+
+import (
+	"context"
+	"fmt"
+)
+
+type (
+	Gateway struct {
+		providers map[string]Provider
+	}
+
+	ErrProviderNotFound struct {
+		ProviderName string
+	}
+)
+
+func (e *ErrProviderNotFound) Error() string {
+	return fmt.Sprintf("provider %s not found", e.ProviderName)
+}
+
+func NewGateway(providers map[string]Provider) *Gateway {
+	return &Gateway{providers: providers}
+}
+
+func (g *Gateway) Generate(ctx context.Context, providerName string, req GenerateRequest) (*GenerateResponse, error) {
+	provider, ok := g.providers[providerName]
+	if !ok {
+		return nil, &ErrProviderNotFound{ProviderName: providerName}
+	}
+	return provider.Generate(ctx, req)
+}
+
+func (g *Gateway) Chat(ctx context.Context, providerName string, req ChatRequest) (*ChatResponse, error) {
+	provider, ok := g.providers[providerName]
+	if !ok {
+		return nil, &ErrProviderNotFound{ProviderName: providerName}
+	}
+	return provider.Chat(ctx, req)
+}

pkg/llmgw/message.go

@@ -0,0 +1,30 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package llmgw
+
+type (
+	Role string
+
+	Message struct {
+		Role    Role   `json:"role"`
+		Content string `json:"content"`
+	}
+)
+
+const (
+	RoleUser      Role = "user"
+	RoleSystem    Role = "system"
+	RoleAssistant Role = "assistant"
+)

pkg/llmgw/models.go

@@ -0,0 +1,23 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package llmgw
+
+type (
+	Model string
+)
+
+const (
+	ModelGPT4o Model = "gpt-4o"
+)

pkg/llmgw/provider.go

@@ -0,0 +1,46 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package llmgw
+
+import "context"
+
+type (
+	ChatRequest struct {
+		Model       Model     `json:"model"`
+		Messages    []Message `json:"messages"`
+		MaxTokens   int       `json:"max_tokens"`
+		Temperature float64   `json:"temperature"`
+	}
+
+	GenerateRequest struct {
+		Model       Model   `json:"model"`
+		Prompt      string  `json:"prompt"`
+		MaxTokens   int     `json:"max_tokens"`
+		Temperature float64 `json:"temperature"`
+	}
+
+	ChatResponse struct {
+		Text string `json:"text"`
+	}
+
+	GenerateResponse struct {
+		Text string `json:"text"`
+	}
+
+	Provider interface {
+		Generate(ctx context.Context, req GenerateRequest) (*GenerateResponse, error)
+		Chat(ctx context.Context, req ChatRequest) (*ChatResponse, error)
+	}
+)

pkg/llmgw/providers/openai/openai.go

@@ -0,0 +1,90 @@
+// Copyright (c) 2025 Probo Inc <[email protected]>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package openai
+
+import (
+	"context"
+
+	"github.com/getprobo/probo/pkg/llmgw"
+	"github.com/openai/openai-go"
+)
+
+type (
+	Provider struct {
+		client *openai.Client
+	}
+)
+
+func New(client *openai.Client) *Provider {
+	return &Provider{client: client}
+}
+
+func (g *Provider) Generate(ctx context.Context, req llmgw.GenerateRequest) (*llmgw.GenerateResponse, error) {
+	model := openai.ChatModel(req.Model)
+
+	params := openai.ChatCompletionNewParams{
+		Model:       model,
+		MaxTokens:   openai.Int(int64(req.MaxTokens)),
+		Temperature: openai.Float(req.Temperature),
+		Messages: []openai.ChatCompletionMessageParamUnion{
+			openai.UserMessage(req.Prompt),
+		},
+	}
+
+	response, err := g.client.Chat.Completions.New(ctx, params)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(response.Choices) == 0 {
+		return &llmgw.GenerateResponse{}, nil
+	}
+
+	return &llmgw.GenerateResponse{Text: response.Choices[0].Message.Content}, nil
+}
+
+func (g *Provider) Chat(ctx context.Context, req llmgw.ChatRequest) (*llmgw.ChatResponse, error) {
+	model := openai.ChatModel(req.Model)
+
+	var messages []openai.ChatCompletionMessageParamUnion
+	for _, msg := range req.Messages {
+		switch msg.Role {
+		case llmgw.RoleUser:
+			messages = append(messages, openai.UserMessage(msg.Content))
+		case llmgw.RoleSystem:
+			messages = append(messages, openai.SystemMessage(msg.Content))
+		case llmgw.RoleAssistant:
+			messages = append(messages, openai.AssistantMessage(msg.Content))
+		}
+	}
+
+	params := openai.ChatCompletionNewParams{
+		Model:       model,
+		MaxTokens:   openai.Int(int64(req.MaxTokens)),
+		Temperature: openai.Float(req.Temperature),
+		Messages:    messages,
+	}
+
+	response, err := g.client.Chat.Completions.New(ctx, params)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(response.Choices) == 0 {
+		return &llmgw.ChatResponse{}, nil
+	}
+
+	return &llmgw.ChatResponse{Text: response.Choices[0].Message.Content}, nil
+}