add --org flag to run and eval

Author: garman

cmd/eval/eval.go

@@ -48,6 +48,10 @@ type EvaluationResult struct {
 	Details       string  `json:"details,omitempty"`
 }
 
+type Organization struct {
+	Name string `json:"name"`
+}
+
 var FailedTests = errors.New("❌ Some tests failed.")
 
 // NewEvalCommand returns a new command to evaluate prompts against models
@@ -66,7 +70,7 @@ func NewEvalCommand(cfg *command.Config) *cobra.Command {
 
 			Example prompt.yml structure:
 			  name: My Evaluation
-			  model: gpt-4o
+			  model: openai/gpt-4o
 			  testData:
 			    - input: "Hello world"
 			      expected: "Hello there"
@@ -94,6 +98,9 @@ func NewEvalCommand(cfg *command.Config) *cobra.Command {
 				return err
 			}
 
+			// Get the org flag
+			org, _ := cmd.Flags().GetString("org")
+
 			// Load the evaluation prompt file
 			evalFile, err := loadEvaluationPromptFile(promptFilePath)
 			if err != nil {
@@ -106,6 +113,7 @@ func NewEvalCommand(cfg *command.Config) *cobra.Command {
 				client:     cfg.Client,
 				evalFile:   evalFile,
 				jsonOutput: jsonOutput,
+				org:        org,
 			}
 
 			err = handler.runEvaluation(cmd.Context())
@@ -120,6 +128,7 @@ func NewEvalCommand(cfg *command.Config) *cobra.Command {
 	}
 
 	cmd.Flags().Bool("json", false, "Output results in JSON format")
+	cmd.Flags().String("org", "", "Organization to attribute usage to (omitting will attribute usage to the current actor")
 	return cmd
 }
 
@@ -128,6 +137,7 @@ type evalCommandHandler struct {
 	client     azuremodels.Client
 	evalFile   *prompt.File
 	jsonOutput bool
+	org        string
 }
 
 func loadEvaluationPromptFile(filePath string) (*prompt.File, error) {
@@ -321,7 +331,7 @@ func (h *evalCommandHandler) templateString(templateStr string, data map[string]
 func (h *evalCommandHandler) callModel(ctx context.Context, messages []azuremodels.ChatMessage) (string, error) {
 	req := h.evalFile.BuildChatCompletionOptions(messages)
 
-	resp, err := h.client.GetChatCompletionStream(ctx, req)
+	resp, err := h.client.GetChatCompletionStream(ctx, req, h.org)
 	if err != nil {
 		return "", err
 	}
@@ -460,7 +470,7 @@ func (h *evalCommandHandler) runLLMEvaluator(ctx context.Context, name string, e
 		Stream:   false,
 	}
 
-	resp, err := h.client.GetChatCompletionStream(ctx, req)
+	resp, err := h.client.GetChatCompletionStream(ctx, req, h.org)
 	if err != nil {
 		return EvaluationResult{}, fmt.Errorf("failed to call evaluation model: %w", err)
 	}

cmd/run/run.go

@@ -216,6 +216,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 		Args: cobra.ArbitraryArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			filePath, _ := cmd.Flags().GetString("file")
+			org, _ := cmd.Flags().GetString("org")
 			var pf *prompt.File
 			if filePath != "" {
 				var err error
@@ -357,7 +358,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 				//nolint:gocritic,revive // TODO
 				defer sp.Stop()
 
-				reader, err := cmdHandler.getChatCompletionStreamReader(req)
+				reader, err := cmdHandler.getChatCompletionStreamReader(req, org)
 				if err != nil {
 					return err
 				}
@@ -408,6 +409,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 	cmd.Flags().String("temperature", "", "Controls randomness in the response, use lower to be more deterministic.")
 	cmd.Flags().String("top-p", "", "Controls text diversity by selecting the most probable words until a set probability is reached.")
 	cmd.Flags().String("system-prompt", "", "Prompt the system.")
+	cmd.Flags().String("org", "", "Organization to attribute usage to (omitting will attribute usage to the current actor")
 
 	return cmd
 }
@@ -522,8 +524,8 @@ func validateModelName(modelName string, models []*azuremodels.ModelSummary) (st
 	return modelName, nil
 }
 
-func (h *runCommandHandler) getChatCompletionStreamReader(req azuremodels.ChatCompletionOptions) (sse.Reader[azuremodels.ChatCompletion], error) {
-	resp, err := h.client.GetChatCompletionStream(h.ctx, req)
+func (h *runCommandHandler) getChatCompletionStreamReader(req azuremodels.ChatCompletionOptions, org string) (sse.Reader[azuremodels.ChatCompletion], error) {
+	resp, err := h.client.GetChatCompletionStream(h.ctx, req, org)
 	if err != nil {
 		return nil, err
 	}

internal/azuremodels/azure_client.go

@@ -40,7 +40,7 @@ func NewAzureClient(httpClient *http.Client, authToken string, cfg *AzureClientC
 }
 
 // GetChatCompletionStream returns a stream of chat completions using the given options.
-func (c *AzureClient) GetChatCompletionStream(ctx context.Context, req ChatCompletionOptions) (*ChatCompletionResponse, error) {
+func (c *AzureClient) GetChatCompletionStream(ctx context.Context, req ChatCompletionOptions, org string) (*ChatCompletionResponse, error) {
 	// Check for o1 models, which don't support streaming
 	if req.Model == "o1-mini" || req.Model == "o1-preview" || req.Model == "o1" {
 		req.Stream = false
@@ -55,7 +55,14 @@ func (c *AzureClient) GetChatCompletionStream(ctx context.Context, req ChatCompl
 
 	body := bytes.NewReader(bodyBytes)
 
-	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.InferenceURL, body)
+	var inferenceURL string
+	if org != "" {
+		inferenceURL = fmt.Sprintf("%s/orgs/%s/%s", c.cfg.InferenceRoot, org, c.cfg.InferencePath)
+	} else {
+		inferenceURL = c.cfg.InferenceRoot + "/" + c.cfg.InferencePath
+	}
+
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, inferenceURL, body)
 	if err != nil {
 		return nil, err
 	}

internal/azuremodels/azure_client_config.go

@@ -1,22 +1,25 @@
 package azuremodels
 
 const (
-	defaultInferenceURL     = "https://models.github.ai/inference/chat/completions"
+	defaultInferenceRoot    = "https://models.github.ai"
+	defaultInferencePath    = "inference/chat/completions"
 	defaultAzureAiStudioURL = "https://api.catalog.azureml.ms"
 	defaultModelsURL        = defaultAzureAiStudioURL + "/asset-gallery/v1.0/models"
 )
 
 // AzureClientConfig represents configurable settings for the Azure client.
 type AzureClientConfig struct {
-	InferenceURL     string
+	InferenceRoot    string
+	InferencePath    string
 	AzureAiStudioURL string
 	ModelsURL        string
 }
 
 // NewDefaultAzureClientConfig returns a new AzureClientConfig with default values for API URLs.
 func NewDefaultAzureClientConfig() *AzureClientConfig {
 	return &AzureClientConfig{
-		InferenceURL:     defaultInferenceURL,
+		InferenceRoot:    defaultInferenceRoot,
+		InferencePath:    defaultInferencePath,
 		AzureAiStudioURL: defaultAzureAiStudioURL,
 		ModelsURL:        defaultModelsURL,
 	}

internal/azuremodels/client.go

@@ -5,7 +5,7 @@ import "context"
 // Client represents a client for interacting with an API about models.
 type Client interface {
 	// GetChatCompletionStream returns a stream of chat completions using the given options.
-	GetChatCompletionStream(context.Context, ChatCompletionOptions) (*ChatCompletionResponse, error)
+	GetChatCompletionStream(context.Context, ChatCompletionOptions, string) (*ChatCompletionResponse, error)
 	// GetModelDetails returns the details of the specified model in a particular registry.
 	GetModelDetails(ctx context.Context, registry, modelName, version string) (*ModelDetails, error)
 	// ListModels returns a list of available models.

internal/azuremodels/mock_client.go

@@ -7,15 +7,15 @@ import (
 
 // MockClient provides a client for interacting with the Azure models API in tests.
 type MockClient struct {
-	MockGetChatCompletionStream func(context.Context, ChatCompletionOptions) (*ChatCompletionResponse, error)
+	MockGetChatCompletionStream func(context.Context, ChatCompletionOptions, string) (*ChatCompletionResponse, error)
 	MockGetModelDetails         func(context.Context, string, string, string) (*ModelDetails, error)
 	MockListModels              func(context.Context) ([]*ModelSummary, error)
 }
 
 // NewMockClient returns a new mock client for stubbing out interactions with the models API.
 func NewMockClient() *MockClient {
 	return &MockClient{
-		MockGetChatCompletionStream: func(context.Context, ChatCompletionOptions) (*ChatCompletionResponse, error) {
+		MockGetChatCompletionStream: func(context.Context, ChatCompletionOptions, string) (*ChatCompletionResponse, error) {
 			return nil, errors.New("GetChatCompletionStream not implemented")
 		},
 		MockGetModelDetails: func(context.Context, string, string, string) (*ModelDetails, error) {
@@ -28,8 +28,8 @@ func NewMockClient() *MockClient {
 }
 
 // GetChatCompletionStream calls the mocked function for getting a stream of chat completions for the given request.
-func (c *MockClient) GetChatCompletionStream(ctx context.Context, opt ChatCompletionOptions) (*ChatCompletionResponse, error) {
-	return c.MockGetChatCompletionStream(ctx, opt)
+func (c *MockClient) GetChatCompletionStream(ctx context.Context, opt ChatCompletionOptions, org string) (*ChatCompletionResponse, error) {
+	return c.MockGetChatCompletionStream(ctx, opt, org)
 }
 
 // GetModelDetails calls the mocked function for getting the details of the specified model in a particular registry.

internal/azuremodels/types.go

@@ -26,12 +26,13 @@ type ChatMessage struct {
 
 // ChatCompletionOptions represents available options for a chat completion request.
 type ChatCompletionOptions struct {
-	MaxTokens   *int          `json:"max_tokens,omitempty"`
-	Messages    []ChatMessage `json:"messages"`
-	Model       string        `json:"model"`
-	Stream      bool          `json:"stream,omitempty"`
-	Temperature *float64      `json:"temperature,omitempty"`
-	TopP        *float64      `json:"top_p,omitempty"`
+	MaxTokens    *int          `json:"max_tokens,omitempty"`
+	Messages     []ChatMessage `json:"messages"`
+	Model        string        `json:"model"`
+	Stream       bool          `json:"stream,omitempty"`
+	Temperature  *float64      `json:"temperature,omitempty"`
+	TopP         *float64      `json:"top_p,omitempty"`
+	Organization *string       `json:"organization,omitempty"`
 }
 
 // ChatChoiceMessage is a message from a choice in a chat conversation.

internal/azuremodels/unauthenticated_client.go

@@ -15,7 +15,7 @@ func NewUnauthenticatedClient() *UnauthenticatedClient {
 }
 
 // GetChatCompletionStream returns an error because this functionality requires authentication.
-func (c *UnauthenticatedClient) GetChatCompletionStream(ctx context.Context, opt ChatCompletionOptions) (*ChatCompletionResponse, error) {
+func (c *UnauthenticatedClient) GetChatCompletionStream(ctx context.Context, opt ChatCompletionOptions, org string) (*ChatCompletionResponse, error) {
 	return nil, errors.New("not authenticated")
 }