parrot Public

Watch 0 Fork 0 Star 0
Go · 3715 bytes Raw Blame History
1 package llm
2
3 import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "fmt"
8 "net/http"
9 "time"
10 )
11
12 type APIClient struct {
13 Endpoint string
14 APIKey string
15 Model string
16 client *http.Client
17 }
18
19 type ChatMessage struct {
20 Role string `json:"role"`
21 Content string `json:"content"`
22 }
23
24 type ChatRequest struct {
25 Model string `json:"model"`
26 Messages []ChatMessage `json:"messages"`
27 MaxTokens int `json:"max_tokens,omitempty"`
28 Temperature float64 `json:"temperature,omitempty"`
29 }
30
31 type ChatChoice struct {
32 Message ChatMessage `json:"message"`
33 FinishReason string `json:"finish_reason"`
34 }
35
36 type ChatResponse struct {
37 Choices []ChatChoice `json:"choices"`
38 Error *APIError `json:"error,omitempty"`
39 }
40
41 type APIError struct {
42 Message string `json:"message"`
43 Type string `json:"type"`
44 Code string `json:"code"`
45 }
46
47 func NewAPIClient(endpoint, apiKey, model string, timeout int) *APIClient {
48 return &APIClient{
49 Endpoint: endpoint,
50 APIKey: apiKey,
51 Model: model,
52 client: &http.Client{
53 // Use a generous timeout; actual timeout is controlled by context
54 Timeout: 60 * time.Second,
55 },
56 }
57 }
58
59 func (c *APIClient) Generate(ctx context.Context, prompt string) (string, error) {
60 if c.APIKey == "" {
61 return "", fmt.Errorf("API key not configured")
62 }
63
64 // Build chat request
65 req := ChatRequest{
66 Model: c.Model,
67 Messages: []ChatMessage{
68 {
69 Role: "user",
70 Content: prompt,
71 },
72 },
73 MaxTokens: 150, // Keep responses concise
74 Temperature: 0.8, // Creative but focused
75 }
76
77 reqBody, err := json.Marshal(req)
78 if err != nil {
79 return "", fmt.Errorf("failed to marshal request: %w", err)
80 }
81
82 // Create HTTP request
83 endpoint := c.Endpoint + "/chat/completions"
84 httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(reqBody))
85 if err != nil {
86 return "", fmt.Errorf("failed to create request: %w", err)
87 }
88
89 // Set headers
90 httpReq.Header.Set("Content-Type", "application/json")
91 httpReq.Header.Set("Authorization", "Bearer "+c.APIKey)
92
93 // Send request
94 resp, err := c.client.Do(httpReq)
95 if err != nil {
96 return "", fmt.Errorf("failed to send request: %w", err)
97 }
98 defer resp.Body.Close()
99
100 // Parse response
101 var chatResp ChatResponse
102 if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
103 return "", fmt.Errorf("failed to decode response: %w", err)
104 }
105
106 // Check for API errors
107 if chatResp.Error != nil {
108 return "", fmt.Errorf("API error: %s", chatResp.Error.Message)
109 }
110
111 // Extract response
112 if len(chatResp.Choices) == 0 {
113 return "", fmt.Errorf("no response choices returned")
114 }
115
116 response := chatResp.Choices[0].Message.Content
117 if response == "" {
118 return "", fmt.Errorf("empty response from API")
119 }
120
121 return response, nil
122 }
123
124 func (c *APIClient) IsAvailable() bool {
125 if c.APIKey == "" {
126 return false
127 }
128
129 // Simple check - try to create a request
130 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
131 defer cancel()
132
133 // Create a minimal test request
134 req := ChatRequest{
135 Model: c.Model,
136 Messages: []ChatMessage{
137 {Role: "user", Content: "test"},
138 },
139 MaxTokens: 1,
140 }
141
142 reqBody, err := json.Marshal(req)
143 if err != nil {
144 return false
145 }
146
147 endpoint := c.Endpoint + "/chat/completions"
148 httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(reqBody))
149 if err != nil {
150 return false
151 }
152
153 httpReq.Header.Set("Content-Type", "application/json")
154 httpReq.Header.Set("Authorization", "Bearer "+c.APIKey)
155
156 resp, err := c.client.Do(httpReq)
157 if err != nil {
158 return false
159 }
160 defer resp.Body.Close()
161
162 // Consider 2xx status codes as available
163 return resp.StatusCode >= 200 && resp.StatusCode < 300
164 }