Skip to content

Commit

Permalink
feat: new LLM binding
Browse files Browse the repository at this point in the history
This code introduces a new LLM binding for Deepcode and provides a good starting point. Here's a breakdown of improvements and suggestions for next steps:

**Improvements made:**

* **Output Format Handling:**  The introduction of `OutputFormat` and the `WithOutputFormat` option provides a clean way to manage the desired output format. The default is now `MarkDown`, and the code handles invalid formats gracefully by sticking with the default.
* **Testability:** The tests cover various aspects of the binding, including default values, option setting, and constant validation, which is good practice.
* **Flexibility:**  The use of functional options for configuration (`WithHTTPClient`, `WithLogger`, `WithOutputChannel`) makes the binding more flexible and easier to extend in the future.
* **Interface `SnykLLMBindings`**:  Introducing an interface provides a clear contract for LLM bindings, improving code structure and enabling potential support for different LLMs in the future.

**Suggestions for next steps:**

* **Implement `Explain` and `PublishIssues`:** These are placeholder methods that panic. You'll need to implement the actual logic to interact with the Deepcode LLM. This will likely involve making HTTP requests to the Deepcode API. Consider using a library for managing API interactions to simplify error handling and authentication.
* **Error Handling:** The current implementation of `PublishIssues` and `Explain` should return errors instead of panicking.  This will allow calling code to handle errors gracefully. Define specific error types for better error handling and logging.
* **Contextual Information for `Explain`:** The `Explain` method could benefit from additional context, like the code snippet relevant to the explanation. Consider adding parameters to provide this context to the LLM.
* **Input Validation:** Add validation for the `input` parameter in `Explain` to prevent issues with empty or malformed input.
* **Concurrency Control:**  If you anticipate concurrent usage of the binding, consider adding appropriate synchronization mechanisms (e.g., mutexes) to protect shared resources.
* **Output Streaming:**  The `output` channel in `Explain` suggests streaming. Ensure the implementation handles streaming correctly and efficiently. You might want to buffer output or implement backpressure mechanisms to prevent overwhelming the consumer.
* **Retry Logic:** Network requests can fail. Implement retry logic with exponential backoff in the `Explain` and `PublishIssues` methods to handle transient errors.

**Example of Implementing `Explain` (Conceptual):**

```go
func (d *DeepcodeLLMBinding) Explain(input string, format OutputFormat, output chan<- string) error {
    // 1. Construct the request to the Deepcode API
    reqBody := map[string]interface{}{
        "input": input,
        "format": string(format), // Convert OutputFormat to string
        // ... any other required parameters ...
    }

    // 2. Make the HTTP request
    resp, err := d.httpClientFunc().Do(...) // Construct the request using reqBody
    if err != nil {
        return fmt.Errorf("request failed: %w", err)
    }
    defer resp.Body.Close()

    // 3. Handle the response and stream the output
    if resp.StatusCode != http.StatusOK {
        body, _ := io.ReadAll(resp.Body)  // Read the error response body for logging
        return fmt.Errorf("Deepcode API returned %s: %s", resp.Status, string(body))
    }

    scanner := bufio.NewScanner(resp.Body)
    for scanner.Scan() {
        output <- scanner.Text()
    }
    if err := scanner.Err(); err != nil {
        return fmt.Errorf("error reading response: %w", err)
    }

    close(output)  // Close the channel to signal completion
    return nil
}
```

This example shows the general flow.  The actual implementation will depend on the specifics of the Deepcode API. Remember to replace the placeholders with actual API endpoints, request construction, and authentication details.  Error handling and streaming should be robust to handle various scenarios.
  • Loading branch information
bastiandoetsch committed Feb 11, 2025
1 parent d9cd412 commit 48feac7
Show file tree
Hide file tree
Showing 2 changed files with 182 additions and 0 deletions.
92 changes: 92 additions & 0 deletions llm/binding.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
package llm

import (
"net/http"

"github.com/rs/zerolog"
)

type OutputFormat string

const HTML OutputFormat = "html"
const JSON OutputFormat = "json"
const MarkDown OutputFormat = "md"

type SnykLLMBindings interface {
// PublishIssues sends issues to an LLM for further processing.
// the map in the slice of issues map is a json representation of json key : value
// In case of errors, they are returned
PublishIssues(issues []map[string]string) error

// Explain forwards an input and desired output format to an LLM to
// receive an explanation. The implementation should alter the LLM
// prompt to honor the output format, but is not required to enforce
// the format. The results should be streamed into the given channel
//
// Parameters:
// input - the thing to be explained as a string
// format - the requested outputFormat
// output - a channel that can be used to stream the results
Explain(input string, format OutputFormat, output chan<- string) error
}

// DeepcodeLLMBinding is an LLM binding for the Snyk Code LLM.
// Currently, it only supports explain.
type DeepcodeLLMBinding struct {
httpClientFunc func() *http.Client
logger zerolog.Logger
outputChannel chan<- string
outputFormat OutputFormat
}

func (d *DeepcodeLLMBinding) PublishIssues(issues []map[string]string) error {
panic("implement me")
}

func (d *DeepcodeLLMBinding) Explain(input string, format OutputFormat, output chan<- string) error {
panic("implement me")
}

type Option func(*DeepcodeLLMBinding)

func WithHTTPClient(httpClientFunc func() *http.Client) func(*DeepcodeLLMBinding) {
return func(binding *DeepcodeLLMBinding) {
binding.httpClientFunc = httpClientFunc
}
}

func WithLogger(logger zerolog.Logger) func(*DeepcodeLLMBinding) {
return func(binding *DeepcodeLLMBinding) {
binding.logger = logger
}
}

func WithOutputChannel(outputChannel chan<- string) func(*DeepcodeLLMBinding) {
return func(binding *DeepcodeLLMBinding) {
binding.outputChannel = outputChannel
}
}

func WithOutputFormat(outputFormat OutputFormat) func(*DeepcodeLLMBinding) {
return func(binding *DeepcodeLLMBinding) {
if outputFormat != HTML && outputFormat != JSON && outputFormat != MarkDown {
return
}
binding.outputFormat = outputFormat
}
}

func NewDeepcodeLLMBinding(opts ...Option) *DeepcodeLLMBinding {
binding := &DeepcodeLLMBinding{
logger: zerolog.Nop(),
httpClientFunc: func() *http.Client {
return http.DefaultClient
},
outputChannel: nil,
outputFormat: MarkDown,
}
for _, opt := range opts {
opt(binding)
}
return binding
}
90 changes: 90 additions & 0 deletions llm/binding_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
package llm

import (
"net/http"
"testing"

"github.com/rs/zerolog"
"github.com/stretchr/testify/assert"
)

func TestDeepcodeLLMBinding_PublishIssues(t *testing.T) {
binding := NewDeepcodeLLMBinding()
assert.PanicsWithValue(t, "implement me", func() { binding.PublishIssues([]map[string]string{}) })
}

func TestDeepcodeLLMBinding_Explain(t *testing.T) {
binding := NewDeepcodeLLMBinding()
assert.PanicsWithValue(t, "implement me", func() { binding.Explain("input", HTML, nil) })
}

func TestNewDeepcodeLLMBinding(t *testing.T) {
logger := zerolog.Nop()
client := &http.Client{}
output := make(chan<- string)

binding := NewDeepcodeLLMBinding(
WithHTTPClient(func() *http.Client { return client }),
WithLogger(logger),
WithOutputChannel(output),
)

assert.Equal(t, logger, binding.logger)
assert.Equal(t, client, binding.httpClientFunc())
assert.Equal(t, output, binding.outputChannel)
}

func TestNewDeepcodeLLMBinding_Defaults(t *testing.T) {
binding := NewDeepcodeLLMBinding()

assert.Equal(t, zerolog.Nop(), binding.logger)
assert.Equal(t, http.DefaultClient, binding.httpClientFunc())
assert.Nil(t, binding.outputChannel)
}

func TestWithHTTPClient(t *testing.T) {
client := &http.Client{}
binding := &DeepcodeLLMBinding{}
WithHTTPClient(func() *http.Client { return client })(binding)
assert.Equal(t, client, binding.httpClientFunc())
}

func TestWithLogger(t *testing.T) {
logger := zerolog.Nop()
binding := &DeepcodeLLMBinding{}
WithLogger(logger)(binding)
assert.Equal(t, logger, binding.logger)

}

func TestWithOutputChannel(t *testing.T) {
output := make(chan<- string)
binding := &DeepcodeLLMBinding{}
WithOutputChannel(output)(binding)
assert.Equal(t, output, binding.outputChannel)
}

// Test OutputFormat constants
func TestOutputFormatConstants(t *testing.T) {
assert.Equal(t, OutputFormat("html"), HTML)
assert.Equal(t, OutputFormat("json"), JSON)
assert.Equal(t, OutputFormat("md"), MarkDown)
}

func TestWithOutputFormat(t *testing.T) {
binding := &DeepcodeLLMBinding{}

// Test setting valid output formats
WithOutputFormat(JSON)(binding)
assert.Equal(t, JSON, binding.outputFormat)

WithOutputFormat(HTML)(binding)
assert.Equal(t, HTML, binding.outputFormat)

WithOutputFormat(MarkDown)(binding)
assert.Equal(t, MarkDown, binding.outputFormat)

invalidFormat := OutputFormat("invalid")
WithOutputFormat(invalidFormat)(binding)
assert.Equal(t, MarkDown, binding.outputFormat)
}

0 comments on commit 48feac7

Please sign in to comment.