feat: initial scaffold with context adapters and litellm pkg
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Mathias Bergqvist
250
pkg/litellm/model.go
Normal file
250
pkg/litellm/model.go
Normal file
@@ -0,0 +1,250 @@
|
||||
package litellm
|
||||
|
||||
// Model implements google.golang.org/adk/model.LLM against any
|
||||
// OpenAI-compatible endpoint (LiteLLM, Ollama, vLLM, etc.).
|
||||
//
|
||||
// The official Go ADK v1.x ships only Gemini adapters. This adapter
|
||||
// implements the official interface directly via net/http — no extra deps.
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"iter"
|
||||
"net/http"
|
||||
|
||||
adkmodel "google.golang.org/adk/model"
|
||||
"google.golang.org/genai"
|
||||
)
|
||||
|
||||
// Model is an ADK-compatible LLM backed by an OpenAI-compatible endpoint.
|
||||
type Model struct {
|
||||
name string
|
||||
baseURL string
|
||||
apiKey string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// New creates an OpenAI-compatible ADK model.
|
||||
// name is the model identifier sent in requests (e.g. "berget/llama-3.3-70b").
|
||||
// baseURL is the API base without path (e.g. "https://llm-api.d-ma.be/v1").
|
||||
func New(name, baseURL, apiKey string) *Model {
|
||||
return &Model{name: name, baseURL: baseURL, apiKey: apiKey, client: &http.Client{}}
|
||||
}
|
||||
|
||||
func (m *Model) Name() string { return m.name }
|
||||
|
||||
// --- OpenAI wire types (minimal subset ADK uses) ---
|
||||
|
||||
type oaiMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
ToolCalls []oaiToolCall `json:"tool_calls,omitempty"`
|
||||
}
|
||||
|
||||
type oaiToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function oaiFunctionCall `json:"function"`
|
||||
}
|
||||
|
||||
type oaiFunctionCall struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}
|
||||
|
||||
type oaiTool struct {
|
||||
Type string `json:"type"`
|
||||
Function oaiFunctionDef `json:"function"`
|
||||
}
|
||||
|
||||
type oaiFunctionDef struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters json.RawMessage `json:"parameters,omitempty"`
|
||||
}
|
||||
|
||||
type oaiRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []oaiMessage `json:"messages"`
|
||||
Tools []oaiTool `json:"tools,omitempty"`
|
||||
}
|
||||
|
||||
type oaiChoice struct {
|
||||
Message oaiMessage `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type oaiResponse struct {
|
||||
Choices []oaiChoice `json:"choices"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func (m *Model) GenerateContent(ctx context.Context, req *adkmodel.LLMRequest, _ bool) iter.Seq2[*adkmodel.LLMResponse, error] {
|
||||
return func(yield func(*adkmodel.LLMResponse, error) bool) {
|
||||
msgs := contentsToMessages(req.Contents)
|
||||
tools := adk2oaiTools(req)
|
||||
|
||||
oaiReq := oaiRequest{Model: m.name, Messages: msgs, Tools: tools}
|
||||
|
||||
body, err := json.Marshal(oaiReq)
|
||||
if err != nil {
|
||||
yield(nil, fmt.Errorf("marshal: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
|
||||
m.baseURL+"/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
yield(nil, fmt.Errorf("new request: %w", err))
|
||||
return
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Authorization", "Bearer "+m.apiKey)
|
||||
|
||||
resp, err := m.client.Do(httpReq)
|
||||
if err != nil {
|
||||
yield(nil, fmt.Errorf("http: %w", err))
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
yield(nil, fmt.Errorf("read body: %w", err))
|
||||
return
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
yield(nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(raw)))
|
||||
return
|
||||
}
|
||||
|
||||
var oaiResp oaiResponse
|
||||
if err := json.Unmarshal(raw, &oaiResp); err != nil {
|
||||
yield(nil, fmt.Errorf("unmarshal: %w", err))
|
||||
return
|
||||
}
|
||||
if oaiResp.Error != nil {
|
||||
yield(nil, fmt.Errorf("api error: %s", oaiResp.Error.Message))
|
||||
return
|
||||
}
|
||||
if len(oaiResp.Choices) == 0 {
|
||||
yield(nil, fmt.Errorf("no choices in response"))
|
||||
return
|
||||
}
|
||||
|
||||
content := oaiChoiceToContent(oaiResp.Choices[0])
|
||||
yield(&adkmodel.LLMResponse{Content: content, TurnComplete: true}, nil)
|
||||
}
|
||||
}
|
||||
|
||||
func contentsToMessages(contents []*genai.Content) []oaiMessage {
|
||||
var msgs []oaiMessage
|
||||
for _, c := range contents {
|
||||
if c == nil {
|
||||
continue
|
||||
}
|
||||
var textBuf bytes.Buffer
|
||||
var toolCalls []oaiToolCall
|
||||
|
||||
for _, p := range c.Parts {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
if p.Text != "" {
|
||||
textBuf.WriteString(p.Text)
|
||||
}
|
||||
if p.FunctionCall != nil {
|
||||
argBytes, _ := json.Marshal(p.FunctionCall.Args)
|
||||
toolCalls = append(toolCalls, oaiToolCall{
|
||||
ID: p.FunctionCall.ID,
|
||||
Type: "function",
|
||||
Function: oaiFunctionCall{
|
||||
Name: p.FunctionCall.Name,
|
||||
Arguments: string(argBytes),
|
||||
},
|
||||
})
|
||||
}
|
||||
if p.FunctionResponse != nil {
|
||||
respBytes, _ := json.Marshal(p.FunctionResponse.Response)
|
||||
msgs = append(msgs, oaiMessage{
|
||||
Role: "tool",
|
||||
Content: string(respBytes),
|
||||
ToolCallID: p.FunctionResponse.ID,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if len(toolCalls) > 0 || textBuf.Len() > 0 {
|
||||
msg := oaiMessage{Role: c.Role}
|
||||
if c.Role == "model" {
|
||||
msg.Role = "assistant"
|
||||
}
|
||||
msg.Content = textBuf.String()
|
||||
msg.ToolCalls = toolCalls
|
||||
msgs = append(msgs, msg)
|
||||
}
|
||||
}
|
||||
return msgs
|
||||
}
|
||||
|
||||
func adk2oaiTools(req *adkmodel.LLMRequest) []oaiTool {
|
||||
if len(req.Tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
var tools []oaiTool
|
||||
for name, def := range req.Tools {
|
||||
raw, _ := json.Marshal(def)
|
||||
var m map[string]json.RawMessage
|
||||
_ = json.Unmarshal(raw, &m)
|
||||
var desc string
|
||||
if d, ok := m["description"]; ok {
|
||||
_ = json.Unmarshal(d, &desc)
|
||||
}
|
||||
params := m["parameters"]
|
||||
// Some endpoints (e.g. Berget) reject null parameters for zero-arg tools.
|
||||
if len(params) == 0 || string(params) == "null" {
|
||||
params = json.RawMessage(`{"type":"object","properties":{}}`)
|
||||
}
|
||||
tools = append(tools, oaiTool{
|
||||
Type: "function",
|
||||
Function: oaiFunctionDef{
|
||||
Name: name,
|
||||
Description: desc,
|
||||
Parameters: params,
|
||||
},
|
||||
})
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
func oaiChoiceToContent(choice oaiChoice) *genai.Content {
|
||||
msg := choice.Message
|
||||
var parts []*genai.Part
|
||||
|
||||
if msg.Content != "" {
|
||||
parts = append(parts, &genai.Part{Text: msg.Content})
|
||||
}
|
||||
for _, tc := range msg.ToolCalls {
|
||||
var args map[string]any
|
||||
_ = json.Unmarshal([]byte(tc.Function.Arguments), &args)
|
||||
parts = append(parts, &genai.Part{
|
||||
FunctionCall: &genai.FunctionCall{
|
||||
ID: tc.ID,
|
||||
Name: tc.Function.Name,
|
||||
Args: args,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
role := msg.Role
|
||||
if role == "assistant" {
|
||||
role = "model"
|
||||
}
|
||||
return &genai.Content{Role: role, Parts: parts}
|
||||
}
|
||||
71
pkg/litellm/telemetry.go
Normal file
71
pkg/litellm/telemetry.go
Normal file
@@ -0,0 +1,71 @@
|
||||
package litellm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
|
||||
"go.opentelemetry.io/otel/sdk/resource"
|
||||
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
||||
semconv "go.opentelemetry.io/otel/semconv/v1.36.0"
|
||||
"google.golang.org/adk/telemetry"
|
||||
)
|
||||
|
||||
// SetupTelemetry wires ADK OTLP tracing from environment variables.
|
||||
//
|
||||
// Reads:
|
||||
// - OTLP_ENDPOINT full URL base, e.g. http://jaeger.d-ma.be:4318 (skip if empty)
|
||||
// - ADK_SERVICE_NAME service name in Jaeger (default: "agent")
|
||||
// - ADK_SERVICE_VERSION semver label (default: "0.1.0")
|
||||
//
|
||||
// Returns a shutdown func to call on exit with a short-timeout context.
|
||||
// No-op (nil error, noop shutdown) when OTLP_ENDPOINT is unset.
|
||||
func SetupTelemetry(ctx context.Context) (shutdown func(context.Context) error, err error) {
|
||||
endpoint := os.Getenv("OTLP_ENDPOINT")
|
||||
if endpoint == "" {
|
||||
return func(context.Context) error { return nil }, nil
|
||||
}
|
||||
|
||||
svcName := os.Getenv("ADK_SERVICE_NAME")
|
||||
if svcName == "" {
|
||||
svcName = "agent"
|
||||
}
|
||||
svcVersion := os.Getenv("ADK_SERVICE_VERSION")
|
||||
if svcVersion == "" {
|
||||
svcVersion = "0.1.0"
|
||||
}
|
||||
|
||||
exporter, err := otlptracehttp.New(ctx,
|
||||
otlptracehttp.WithEndpointURL(endpoint+"/v1/traces"),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("otlp exporter: %w", err)
|
||||
}
|
||||
|
||||
res, err := resource.New(ctx,
|
||||
resource.WithAttributes(
|
||||
semconv.ServiceName(svcName),
|
||||
semconv.ServiceVersion(svcVersion),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("resource: %w", err)
|
||||
}
|
||||
|
||||
providers, err := telemetry.New(ctx,
|
||||
telemetry.WithSpanProcessors(sdktrace.NewBatchSpanProcessor(exporter)),
|
||||
telemetry.WithResource(res),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("telemetry.New: %w", err)
|
||||
}
|
||||
providers.SetGlobalOtelProviders()
|
||||
|
||||
return func(ctx context.Context) error {
|
||||
shutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
return providers.Shutdown(shutCtx)
|
||||
}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user