143 lines
3.6 KiB
Go
143 lines
3.6 KiB
Go
// 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
|
|
}
|