111 lines
2.5 KiB
Go
111 lines
2.5 KiB
Go
|
package routing
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"os"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
// llmBaseURL is the base URL for the LLM API
|
||
|
|
llmBaseURL string
|
||
|
|
)
|
||
|
|
|
||
|
|
func init() {
|
||
|
|
llmBaseURL = os.Getenv("LOCAL_LLM_BASE_URL")
|
||
|
|
if llmBaseURL == "" {
|
||
|
|
llmBaseURL = "https://api.riotpiao.com"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// LLMClient is a simple LLM client for routing
|
||
|
|
type LLMClient struct {
|
||
|
|
baseURL string
|
||
|
|
httpClient *http.Client
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewLLMClient creates a new LLM client
|
||
|
|
func NewLLMClient() *LLMClient {
|
||
|
|
return &LLMClient{
|
||
|
|
baseURL: llmBaseURL,
|
||
|
|
httpClient: &http.Client{},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// llmRequest is the request body for the OpenAI-compatible API
|
||
|
|
type llmRequest struct {
|
||
|
|
Model string `json:"model"`
|
||
|
|
Messages []llmMessage `json:"messages"`
|
||
|
|
Stream bool `json:"stream"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type llmMessage struct {
|
||
|
|
Role string `json:"role"`
|
||
|
|
Content string `json:"content"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// llmResponse is the response from the OpenAI-compatible API
|
||
|
|
type llmResponse struct {
|
||
|
|
Choices []struct {
|
||
|
|
Message struct {
|
||
|
|
Content string `json:"content"`
|
||
|
|
} `json:"message"`
|
||
|
|
} `json:"choices"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// Chat sends a chat completion request
|
||
|
|
func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string) (string, error) {
|
||
|
|
req := llmRequest{
|
||
|
|
Model: "reasoning",
|
||
|
|
Messages: []llmMessage{
|
||
|
|
{Role: "system", Content: systemPrompt},
|
||
|
|
{Role: "user", Content: userMessage},
|
||
|
|
},
|
||
|
|
Stream: false,
|
||
|
|
}
|
||
|
|
|
||
|
|
reqBody, err := json.Marshal(req)
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("failed to marshal request: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
httpReq, err := http.NewRequestWithContext(ctx, "POST",
|
||
|
|
fmt.Sprintf("%s/v1/chat/completions", c.baseURL),
|
||
|
|
bytes.NewReader(reqBody))
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("failed to create HTTP request: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
||
|
|
|
||
|
|
resp, err := c.httpClient.Do(httpReq)
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("failed to connect to LLM API at %s: %w", c.baseURL, err)
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
respBody, err := io.ReadAll(resp.Body)
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("failed to read response body: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusOK {
|
||
|
|
return "", fmt.Errorf("LLM API returned status %d: %s", resp.StatusCode, string(respBody))
|
||
|
|
}
|
||
|
|
|
||
|
|
var respObj llmResponse
|
||
|
|
if err := json.Unmarshal(respBody, &respObj); err != nil {
|
||
|
|
return "", fmt.Errorf("failed to unmarshal response: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
if len(respObj.Choices) == 0 {
|
||
|
|
return "", fmt.Errorf("no choices in response from LLM API")
|
||
|
|
}
|
||
|
|
|
||
|
|
return respObj.Choices[0].Message.Content, nil
|
||
|
|
}
|