probo/pkg/llm/llm.go

// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// 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
}