Repo iniciado

shell/llm/openai.go

@@ -0,0 +1,76 @@
+package llm
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	openai "github.com/sashabaranov/go-openai"
+
+	coretypes "github.com/enmanuel/agents/pkg/llm"
+)
+
+// NewOpenAIComplete returns a CompleteFunc backed by the OpenAI-compatible API.
+// Works with OpenAI, Ollama, vLLM, LMStudio — just change baseURL.
+func NewOpenAIComplete(apiKeyEnv, baseURL string) coretypes.CompleteFunc {
+	return func(ctx context.Context, req coretypes.CompletionRequest) (coretypes.CompletionResponse, error) {
+		apiKey := os.Getenv(apiKeyEnv)
+		if apiKey == "" {
+			apiKey = "ollama" // Ollama doesn't require a real key
+		}
+
+		cfg := openai.DefaultConfig(apiKey)
+		if baseURL != "" {
+			cfg.BaseURL = baseURL
+		}
+		client := openai.NewClientWithConfig(cfg)
+
+		msgs := make([]openai.ChatCompletionMessage, 0, len(req.Messages)+1)
+		if req.SystemPrompt != "" {
+			msgs = append(msgs, openai.ChatCompletionMessage{
+				Role:    openai.ChatMessageRoleSystem,
+				Content: req.SystemPrompt,
+			})
+		}
+		for _, m := range req.Messages {
+			role := openai.ChatMessageRoleUser
+			switch m.Role {
+			case coretypes.RoleAssistant:
+				role = openai.ChatMessageRoleAssistant
+			case coretypes.RoleSystem:
+				role = openai.ChatMessageRoleSystem
+			case coretypes.RoleTool:
+				role = openai.ChatMessageRoleTool
+			}
+			msgs = append(msgs, openai.ChatCompletionMessage{
+				Role:    role,
+				Content: m.Content,
+			})
+		}
+
+		openReq := openai.ChatCompletionRequest{
+			Model:       req.Model,
+			Messages:    msgs,
+			MaxTokens:   req.MaxTokens,
+			Temperature: float32(req.Temperature),
+		}
+
+		resp, err := client.CreateChatCompletion(ctx, openReq)
+		if err != nil {
+			return coretypes.CompletionResponse{}, fmt.Errorf("openai completion: %w", err)
+		}
+		if len(resp.Choices) == 0 {
+			return coretypes.CompletionResponse{}, fmt.Errorf("openai: empty choices")
+		}
+
+		return coretypes.CompletionResponse{
+			Content:      resp.Choices[0].Message.Content,
+			FinishReason: string(resp.Choices[0].FinishReason),
+			Usage: coretypes.TokenUsage{
+				InputTokens:  resp.Usage.PromptTokens,
+				OutputTokens: resp.Usage.CompletionTokens,
+				TotalTokens:  resp.Usage.TotalTokens,
+			},
+		}, nil
+	}
+}