Nik Afiq
ad50d641bd
- Updated LLMClient interface to support model-specific generation and model listing. - Integrated model store and validator into the command application for managing AI models. - Implemented commands for setting, getting, and listing active AI models in Discord. - Enhanced AI query handling to utilize the selected model and return model information in responses. - Added caching mechanism for model validation to improve performance. - Introduced gRPC methods for listing available AI models in the ai-gateway. - Updated protobuf definitions to include model-related fields and messages. - Added tests for model store and validator functionalities.
122 lines
3.0 KiB
Go
122 lines
3.0 KiB
Go
package ollama
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
|
|
)
|
|
|
|
type generateRequest struct {
|
|
Model string `json:"model"`
|
|
Prompt string `json:"prompt"`
|
|
Stream bool `json:"stream"`
|
|
Think *bool `json:"think,omitempty"`
|
|
}
|
|
|
|
type generateResponse struct {
|
|
Response string `json:"response"`
|
|
}
|
|
|
|
type listModelsResponse struct {
|
|
Models []struct {
|
|
Name string `json:"name"`
|
|
} `json:"models"`
|
|
}
|
|
|
|
// Client implements the LLM driven port with the Ollama generate API.
|
|
type Client struct {
|
|
baseURL string
|
|
http *http.Client
|
|
}
|
|
|
|
// New constructs an Ollama client with OTel-instrumented transport.
|
|
func New(baseURL string, httpClient *http.Client) *Client {
|
|
if httpClient == nil {
|
|
httpClient = &http.Client{Transport: otelhttp.NewTransport(http.DefaultTransport)}
|
|
}
|
|
if httpClient.Transport == nil {
|
|
httpClient.Transport = otelhttp.NewTransport(http.DefaultTransport)
|
|
}
|
|
return &Client{baseURL: baseURL, http: httpClient}
|
|
}
|
|
|
|
// Generate sends one non-streaming prompt to Ollama.
|
|
func (c *Client) Generate(ctx context.Context, model, prompt string) (string, error) {
|
|
reqBody := generateRequest{
|
|
Model: model,
|
|
Prompt: prompt,
|
|
Stream: false,
|
|
}
|
|
if isThinkingModel(model) {
|
|
disabled := false
|
|
reqBody.Think = &disabled
|
|
}
|
|
|
|
body, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return "", fmt.Errorf("marshal ollama request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/generate", bytes.NewReader(body))
|
|
if err != nil {
|
|
return "", fmt.Errorf("build ollama request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
resp, err := c.http.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("call ollama: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return "", fmt.Errorf("ollama returned status %s", resp.Status)
|
|
}
|
|
|
|
var out generateResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
|
return "", fmt.Errorf("decode ollama response: %w", err)
|
|
}
|
|
return out.Response, nil
|
|
}
|
|
|
|
// ListModels returns the installed model names from Ollama.
|
|
func (c *Client) ListModels(ctx context.Context) ([]string, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/api/tags", nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("build ollama list models request: %w", err)
|
|
}
|
|
|
|
resp, err := c.http.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list ollama models: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return nil, fmt.Errorf("ollama returned status %s", resp.Status)
|
|
}
|
|
|
|
var out listModelsResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
|
return nil, fmt.Errorf("decode ollama list models response: %w", err)
|
|
}
|
|
|
|
names := make([]string, 0, len(out.Models))
|
|
for _, model := range out.Models {
|
|
if model.Name != "" {
|
|
names = append(names, model.Name)
|
|
}
|
|
}
|
|
return names, nil
|
|
}
|
|
|
|
func isThinkingModel(name string) bool {
|
|
return strings.HasPrefix(name, "qwen3")
|
|
}
|