Add LLM interface

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2026-03-11 22:33:12 +01:00
parent 307ca48394
commit 158c36d9ab
22 changed files with 2822 additions and 91 deletions
--- a/pkg/llm/llm.go
+++ b/pkg/llm/llm.go
@@ -0,0 +1,137 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// 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 llm
+
+import (
+	"context"
+	"io"
+	"time"
+
+	"go.gearno.de/kit/log"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
+)
+
+var tracerName = "go.probo.inc/probo/pkg/llm"
+
+type (
+	Option func(*Client)
+
+	Client struct {
+		provider       Provider
+		system         string
+		logger         *log.Logger
+		tracerProvider trace.TracerProvider
+		tracer         trace.Tracer
+	}
+)
+
+func WithLogger(l *log.Logger) Option {
+	return func(c *Client) {
+		c.logger = l
+	}
+}
+
+func WithTracerProvider(tp trace.TracerProvider) Option {
+	return func(c *Client) {
+		c.tracerProvider = tp
+	}
+}
+
+// NewClient creates a new instrumented LLM client.
+// The system parameter identifies the provider for the OTel gen_ai.provider.name
+// attribute (e.g., "openai", "anthropic", "aws.bedrock").
+func NewClient(provider Provider, system string, opts ...Option) *Client {
+	c := &Client{
+		provider:       provider,
+		system:         system,
+		logger:         log.NewLogger(log.WithOutput(io.Discard)),
+		tracerProvider: otel.GetTracerProvider(),
+	}
+
+	for _, opt := range opts {
+		opt(c)
+	}
+
+	c.logger = c.logger.Named("llm").With(log.String("system", system))
+	c.tracer = c.tracerProvider.Tracer(tracerName)
+
+	return c
+}
+
+func (c *Client) ChatCompletion(ctx context.Context, req *ChatCompletionRequest) (*ChatCompletionResponse, error) {
+	ctx, span := startChatSpan(ctx, c.tracer, c.system, req)
+
+	c.logger.InfoCtx(
+		ctx,
+		"chat completion request",
+		log.String("model", req.Model),
+		log.Int("message_count", len(req.Messages)),
+		log.Int("tool_count", len(req.Tools)),
+	)
+
+	start := time.Now()
+	resp, err := c.provider.ChatCompletion(ctx, req)
+	duration := time.Since(start)
+
+	if err != nil {
+		c.logger.ErrorCtx(
+			ctx,
+			"chat completion failed",
+			log.String("model", req.Model),
+			log.Duration("duration", duration),
+			log.Error(err),
+		)
+		endChatSpan(span, nil, err)
+		return nil, err
+	}
+
+	c.logger.InfoCtx(
+		ctx,
+		"chat completion response",
+		log.String("model", resp.Model),
+		log.Int("input_tokens", resp.Usage.InputTokens),
+		log.Int("output_tokens", resp.Usage.OutputTokens),
+		log.String("finish_reason", string(resp.FinishReason)),
+		log.Duration("duration", duration),
+	)
+
+	endChatSpan(span, resp, nil)
+	return resp, nil
+}
+
+func (c *Client) ChatCompletionStream(ctx context.Context, req *ChatCompletionRequest) (ChatCompletionStream, error) {
+	ctx, span := startChatSpan(ctx, c.tracer, c.system, req)
+
+	c.logger.InfoCtx(ctx, "chat completion stream request",
+		log.String("model", req.Model),
+		log.Int("message_count", len(req.Messages)),
+		log.Int("tool_count", len(req.Tools)),
+	)
+
+	stream, err := c.provider.ChatCompletionStream(ctx, req)
+	if err != nil {
+		c.logger.ErrorCtx(
+			ctx,
+			"chat completion stream failed",
+			log.String("model", req.Model),
+			log.Error(err),
+		)
+		endChatSpan(span, nil, err)
+		return nil, err
+	}
+
+	return newTracedStream(stream, span), nil
+}