Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions config/config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ preferences:

models:
- name: gpt-4.1
provider: openai
endpoint: https://api.openai.com/v1/chat/completions
auth_env_var: OPENAI_API_KEY
org_env_var: OPENAI_ORG_ID
Expand All @@ -15,6 +16,7 @@ models:
content: "```bash\necho \"hi\"\n```"

- name: gpt-4.1-mini
provider: openai
endpoint: https://api.openai.com/v1/chat/completions
auth_env_var: OPENAI_API_KEY
org_env_var: OPENAI_ORG_ID
Expand Down
38 changes: 36 additions & 2 deletions llm/llm.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,24 +31,43 @@ func NewLLMClient(config ModelConfig) *LLMClient {
}
}

func (c *LLMClient) provider() string {
if c.config.Provider == "" {
return "openai"
}
return strings.ToLower(c.config.Provider)
}

func (c *LLMClient) isAnthropic() bool {
return c.provider() == "anthropic"
}

func (c *LLMClient) createRequest(payload Payload) (*http.Request, error) {
payloadBytes, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("failed to marshal payload: %w", err)
}

req, err := http.NewRequest("POST", c.config.Endpoint, bytes.NewBuffer(payloadBytes))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
if strings.Contains(c.config.Endpoint, "openai.azure.com") {

if c.isAnthropic() {
req.Header.Set("x-api-key", c.config.Auth)
req.Header.Set("anthropic-version", "2023-06-01")
} else if strings.Contains(c.config.Endpoint, "openai.azure.com") {
req.Header.Set("Api-Key", c.config.Auth)
} else {
req.Header.Set("Authorization", "Bearer "+c.config.Auth)
}
if c.config.OrgID != "" {

if c.provider() == "openai" && c.config.OrgID != "" {
req.Header.Set("OpenAI-Organization", c.config.OrgID)
}

req.Header.Set("Content-Type", "application/json")

return req, nil
}

Expand All @@ -67,23 +86,29 @@ func (c *LLMClient) Query(query string) (string, error) {
if err != nil {
return "", err
}

c.messages = append(c.messages, message)

return message.Content, nil
}

func (c *LLMClient) processStream(resp *http.Response) (string, error) {
counter := 0
streamReader := bufio.NewReader(resp.Body)
totalData := ""

for {
line, err := streamReader.ReadString('\n')
if err != nil {
break
}

line = strings.TrimSpace(line)

if line == "data: [DONE]" {
break
}

if strings.HasPrefix(line, "data:") {
payload := strings.TrimPrefix(line, "data:")

Expand All @@ -93,18 +118,23 @@ func (c *LLMClient) processStream(resp *http.Response) (string, error) {
fmt.Println("Error parsing data:", err)
continue
}

if len(responseData.Choices) == 0 {
continue
}

content := responseData.Choices[0].Delta.Content

if counter < 2 && strings.Count(content, "\n") > 0 {
continue
}

totalData += content
c.StreamCallback(totalData, nil)
counter++
}
}

return totalData, nil
}

Expand All @@ -113,15 +143,19 @@ func (c *LLMClient) callStream(payload Payload) (Message, error) {
if err != nil {
return Message{}, fmt.Errorf("failed to create the request: %w", err)
}

resp, err := c.httpClient.Do(req)
if err != nil {
return Message{}, fmt.Errorf("failed to make the API request: %w", err)
}

defer resp.Body.Close()

if resp.StatusCode != 200 {
return Message{}, fmt.Errorf("API request failed: %s", resp.Status)
}

content, err := c.processStream(resp)

return Message{Role: "assistant", Content: content}, err
}
1 change: 1 addition & 0 deletions types/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package types

type ModelConfig struct {
ModelName string `yaml:"name"`
Provider string `yaml:"provider,omitempty"`
Endpoint string `yaml:"endpoint"`
Auth string `yaml:"auth_env_var"`
OrgID string `yaml:"org_env_var,omitempty"`
Expand Down