diff --git a/internal/llama/client.go b/internal/llama/client.go new file mode 100644 index 0000000..b79292e --- /dev/null +++ b/internal/llama/client.go @@ -0,0 +1,103 @@ +package llama + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "time" +) + +type Client struct { + baseURL string + httpClient *http.Client +} + +type CompletionRequest struct { + Prompt string `json:"prompt"` + Temperature float64 `json:"temperature"` + TopP float64 `json:"top_p"` + MaxTokens int `json:"n_predict"` + Stream bool `json:"stream"` + CachePrompt bool `json:"cache_prompt"` +} + +type CompletionResponse struct { + Content string `json:"content"` + Tokens int `json:"tokens_predicted"` + Timings struct { + PredictedPerSecond float64 `json:"predicted_per_second"` + } `json:"timings"` +} + +type EmbeddingRequest struct { + Content string `json:"content"` +} + +type EmbeddingResponse struct { + Embedding []float64 `json:"embedding"` +} + +func NewClient(baseURL string) *Client { + return &Client{ + baseURL: baseURL, + httpClient: &http.Client{ + Timeout: 120 * time.Second, + }, + } +} + +func (c *Client) Health() error { + resp, err := c.httpClient.Get(c.baseURL + "/health") + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("llama.cpp health check failed: %s", resp.Status) + } + return nil +} + +func (c *Client) Complete(req *CompletionRequest) (*CompletionResponse, error) { + body, err := json.Marshal(req) + if err != nil { + return nil, err + } + resp, err := c.httpClient.Post(c.baseURL+"/completion", "application/json", bytes.NewReader(body)) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + data, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("llama.cpp completion failed (%d): %s", resp.StatusCode, string(data)) + } + var res CompletionResponse + if err := json.NewDecoder(resp.Body).Decode(&res); err != nil { + return nil, err + } + return &res, nil +} + +func (c *Client) Embed(content string) ([]float64, error) { + body, err := json.Marshal(EmbeddingRequest{Content: content}) + if err != nil { + return nil, err + } + resp, err := c.httpClient.Post(c.baseURL+"/embedding", "application/json", bytes.NewReader(body)) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + data, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("llama.cpp embedding failed (%d): %s", resp.StatusCode, string(data)) + } + var res EmbeddingResponse + if err := json.NewDecoder(resp.Body).Decode(&res); err != nil { + return nil, err + } + return res.Embedding, nil +}