diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 012ec4c6..32a8c84a 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -407,7 +407,7 @@ func (h *APIHandlers) handleNonStreamingResponse(c *gin.Context, rawJson []byte) cliClient.RequestMutex.Lock() } - log.Debugf("Request use account: %s, project id: %s", cliClient.Email, cliClient.ProjectID) + log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID()) jsonTemplate := `{"id":"","object":"chat.completion","created":123456,"model":"model","choices":[{"index":0,"message":{"role":"assistant","content":null,"reasoning_content":null,"tool_calls":null},"finish_reason":null,"native_finish_reason":null}]}` respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, contents, tools) for { @@ -501,7 +501,7 @@ func (h *APIHandlers) handleStreamingResponse(c *gin.Context, rawJson []byte) { cliClient.RequestMutex.Lock() } - log.Debugf("Request use account: %s, project id: %s", cliClient.Email, cliClient.ProjectID) + log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID()) respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, contents, tools) for { select { diff --git a/internal/auth/auth.go b/internal/auth/auth.go index 979b403f..f1be254f 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -7,6 +7,7 @@ import ( "fmt" "github.com/luispater/CLIProxyAPI/internal/config" log "github.com/sirupsen/logrus" + "github.com/skratchdot/open-golang/open" "github.com/tidwall/gjson" "golang.org/x/net/proxy" "io" @@ -15,7 +16,6 @@ import ( "net/url" "time" - "github.com/skratchdot/open-golang/open" "golang.org/x/oauth2" "golang.org/x/oauth2/google" ) @@ -33,14 +33,6 @@ var ( } ) -type TokenStorage struct { - Token any `json:"token"` - ProjectID string `json:"project_id"` - Email string `json:"email"` - Auto bool `json:"auto"` - Checked bool `json:"checked"` -} - // GetAuthenticatedClient configures and returns an HTTP client with OAuth2 tokens. // It handles the entire flow: loading, refreshing, and fetching new tokens. func GetAuthenticatedClient(ctx context.Context, ts *TokenStorage, cfg *config.Config) (*http.Client, error) { @@ -199,7 +191,8 @@ func getTokenFromWeb(ctx context.Context, config *oauth2.Config) (*oauth2.Token, authURL := config.AuthCodeURL("state-token", oauth2.AccessTypeOffline, oauth2.SetAuthURLParam("prompt", "consent")) log.Debugf("CLI login required.\nAttempting to open authentication page in your browser.\nIf it does not open, please navigate to this URL:\n\n%s\n", authURL) - err := open.Run(authURL) + var err error + err = open.Run(authURL) if err != nil { log.Errorf("Failed to open browser: %v. Please open the URL manually.", err) } diff --git a/internal/auth/models.go b/internal/auth/models.go new file mode 100644 index 00000000..c90d47c5 --- /dev/null +++ b/internal/auth/models.go @@ -0,0 +1,9 @@ +package auth + +type TokenStorage struct { + Token any `json:"token"` + ProjectID string `json:"project_id"` + Email string `json:"email"` + Auto bool `json:"auto"` + Checked bool `json:"checked"` +} diff --git a/internal/client/client.go b/internal/client/client.go index c42e787a..f721c3d7 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -22,96 +22,16 @@ import ( "time" ) -// --- Constants --- const ( codeAssistEndpoint = "https://cloudcode-pa.googleapis.com" apiVersion = "v1internal" - pluginVersion = "1.0.0" + pluginVersion = "0.1.9" ) -type ErrorMessage struct { - StatusCode int - Error error -} - -type GCPProject struct { - Projects []GCPProjectProjects `json:"projects"` -} -type GCPProjectLabels struct { - GenerativeLanguage string `json:"generative-language"` -} -type GCPProjectProjects struct { - ProjectNumber string `json:"projectNumber"` - ProjectID string `json:"projectId"` - LifecycleState string `json:"lifecycleState"` - Name string `json:"name"` - Labels GCPProjectLabels `json:"labels"` - CreateTime time.Time `json:"createTime"` -} - -type Content struct { - Role string `json:"role"` - Parts []Part `json:"parts"` -} - -// Part represents a single part of a message's content. -type Part struct { - Text string `json:"text,omitempty"` - InlineData *InlineData `json:"inlineData,omitempty"` - FunctionCall *FunctionCall `json:"functionCall,omitempty"` - FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"` -} - -type InlineData struct { - MimeType string `json:"mime_type,omitempty"` - Data string `json:"data,omitempty"` -} - -// FunctionCall represents a tool call requested by the model. -type FunctionCall struct { - Name string `json:"name"` - Args map[string]interface{} `json:"args"` -} - -// FunctionResponse represents the result of a tool execution. -type FunctionResponse struct { - Name string `json:"name"` - Response map[string]interface{} `json:"response"` -} - -// GenerateContentRequest is the request payload for the streamGenerateContent endpoint. -type GenerateContentRequest struct { - Contents []Content `json:"contents"` - Tools []ToolDeclaration `json:"tools,omitempty"` - GenerationConfig `json:"generationConfig"` -} - -// GenerationConfig defines model generation parameters. -type GenerationConfig struct { - ThinkingConfig GenerationConfigThinkingConfig `json:"thinkingConfig,omitempty"` - Temperature float64 `json:"temperature,omitempty"` - TopP float64 `json:"topP,omitempty"` - TopK float64 `json:"topK,omitempty"` - // Temperature, TopP, TopK, etc. can be added here. -} - -type GenerationConfigThinkingConfig struct { - IncludeThoughts bool `json:"include_thoughts,omitempty"` -} - -// ToolDeclaration is the structure for declaring tools to the API. -// For now, we'll assume a simple structure. A more complete implementation -// would mirror the OpenAPI schema definition. -type ToolDeclaration struct { - FunctionDeclarations []interface{} `json:"functionDeclarations"` -} - // Client is the main client for interacting with the CLI API. type Client struct { httpClient *http.Client - ProjectID string RequestMutex sync.Mutex - Email string tokenStorage *auth.TokenStorage cfg *config.Config } @@ -145,9 +65,17 @@ func (c *Client) IsAuto() bool { return c.tokenStorage.Auto } +func (c *Client) GetEmail() string { + return c.tokenStorage.Email +} + +func (c *Client) GetProjectID() string { + return c.tokenStorage.ProjectID +} + // SetupUser performs the initial user onboarding and setup. -func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string, error) { - c.Email = email +func (c *Client) SetupUser(ctx context.Context, email, projectID string) error { + c.tokenStorage.Email = email log.Info("Performing user onboarding...") // 1. LoadCodeAssist @@ -161,7 +89,7 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string var loadAssistResp map[string]interface{} err := c.makeAPIRequest(ctx, "loadCodeAssist", "POST", loadAssistReqBody, &loadAssistResp) if err != nil { - return projectID, fmt.Errorf("failed to load code assist: %w", err) + return fmt.Errorf("failed to load code assist: %w", err) } // a, _ := json.Marshal(&loadAssistResp) @@ -197,14 +125,14 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string if onboardProjectID != "" { onboardReqBody["cloudaicompanionProject"] = onboardProjectID } else { - return projectID, fmt.Errorf("failed to start user onboarding, need define a project id") + return fmt.Errorf("failed to start user onboarding, need define a project id") } for { var lroResp map[string]interface{} err = c.makeAPIRequest(ctx, "onboardUser", "POST", onboardReqBody, &lroResp) if err != nil { - return projectID, fmt.Errorf("failed to start user onboarding: %w", err) + return fmt.Errorf("failed to start user onboarding: %w", err) } // a, _ := json.Marshal(&lroResp) // log.Debug(string(a)) @@ -214,12 +142,12 @@ func (c *Client) SetupUser(ctx context.Context, email, projectID string) (string if doneOk && done { if project, projectOk := lroResp["response"].(map[string]interface{})["cloudaicompanionProject"].(map[string]interface{}); projectOk { if projectID != "" { - c.ProjectID = projectID + c.tokenStorage.ProjectID = projectID } else { - c.ProjectID = project["id"].(string) + c.tokenStorage.ProjectID = project["id"].(string) } - log.Infof("Onboarding complete. Using Project ID: %s", c.ProjectID) - return c.ProjectID, nil + log.Infof("Onboarding complete. Using Project ID: %s", c.tokenStorage.ProjectID) + return nil } } else { log.Println("Onboarding in progress, waiting 5 seconds...") @@ -356,7 +284,7 @@ func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model st request.Tools = tools requestBody := map[string]interface{}{ - "project": c.ProjectID, // Assuming ProjectID is available + "project": c.tokenStorage.ProjectID, // Assuming ProjectID is available "request": request, "model": model, } diff --git a/internal/client/models.go b/internal/client/models.go new file mode 100644 index 00000000..23515000 --- /dev/null +++ b/internal/client/models.go @@ -0,0 +1,80 @@ +package client + +import "time" + +type ErrorMessage struct { + StatusCode int + Error error +} + +type GCPProject struct { + Projects []GCPProjectProjects `json:"projects"` +} +type GCPProjectLabels struct { + GenerativeLanguage string `json:"generative-language"` +} +type GCPProjectProjects struct { + ProjectNumber string `json:"projectNumber"` + ProjectID string `json:"projectId"` + LifecycleState string `json:"lifecycleState"` + Name string `json:"name"` + Labels GCPProjectLabels `json:"labels"` + CreateTime time.Time `json:"createTime"` +} + +type Content struct { + Role string `json:"role"` + Parts []Part `json:"parts"` +} + +// Part represents a single part of a message's content. +type Part struct { + Text string `json:"text,omitempty"` + InlineData *InlineData `json:"inlineData,omitempty"` + FunctionCall *FunctionCall `json:"functionCall,omitempty"` + FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"` +} + +type InlineData struct { + MimeType string `json:"mime_type,omitempty"` + Data string `json:"data,omitempty"` +} + +// FunctionCall represents a tool call requested by the model. +type FunctionCall struct { + Name string `json:"name"` + Args map[string]interface{} `json:"args"` +} + +// FunctionResponse represents the result of a tool execution. +type FunctionResponse struct { + Name string `json:"name"` + Response map[string]interface{} `json:"response"` +} + +// GenerateContentRequest is the request payload for the streamGenerateContent endpoint. +type GenerateContentRequest struct { + Contents []Content `json:"contents"` + Tools []ToolDeclaration `json:"tools,omitempty"` + GenerationConfig `json:"generationConfig"` +} + +// GenerationConfig defines model generation parameters. +type GenerationConfig struct { + ThinkingConfig GenerationConfigThinkingConfig `json:"thinkingConfig,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + TopP float64 `json:"topP,omitempty"` + TopK float64 `json:"topK,omitempty"` + // Temperature, TopP, TopK, etc. can be added here. +} + +type GenerationConfigThinkingConfig struct { + IncludeThoughts bool `json:"include_thoughts,omitempty"` +} + +// ToolDeclaration is the structure for declaring tools to the API. +// For now, we'll assume a simple structure. A more complete implementation +// would mirror the OpenAPI schema definition. +type ToolDeclaration struct { + FunctionDeclarations []interface{} `json:"functionDeclarations"` +} diff --git a/internal/cmd/login.go b/internal/cmd/login.go index 49809a4b..d71b85a9 100644 --- a/internal/cmd/login.go +++ b/internal/cmd/login.go @@ -29,7 +29,7 @@ func DoLogin(cfg *config.Config, projectID string) { // 3. Initialize CLI Client cliClient := client.NewClient(httpClient, &ts, cfg) - projectID, err = cliClient.SetupUser(clientCtx, ts.Email, projectID) + err = cliClient.SetupUser(clientCtx, ts.Email, projectID) if err != nil { if err.Error() == "failed to start user onboarding, need define a project id" { log.Error("failed to start user onboarding") @@ -52,8 +52,7 @@ func DoLogin(cfg *config.Config, projectID string) { log.Fatalf("failed to complete user setup: %v", err) } } else { - auto := ts.ProjectID == "" - cliClient.SetProjectID(projectID) + auto := projectID == "" cliClient.SetIsAuto(auto) if !cliClient.IsChecked() && !cliClient.IsAuto() { diff --git a/internal/cmd/run.go b/internal/cmd/run.go index 133a47da..8bf291bf 100644 --- a/internal/cmd/run.go +++ b/internal/cmd/run.go @@ -57,30 +57,7 @@ func StartService(cfg *config.Config) { // 3. Initialize CLI Client cliClient := client.NewClient(httpClient, &ts, cfg) - if _, err = cliClient.SetupUser(clientCtx, ts.Email, ts.ProjectID); err != nil { - if err.Error() == "failed to start user onboarding, need define a project id" { - log.Error("failed to start user onboarding") - project, errGetProjectList := cliClient.GetProjectList(clientCtx) - if errGetProjectList != nil { - log.Fatalf("failed to complete user setup: %v", err) - } else { - log.Infof("Your account %s needs specify a project id.", ts.Email) - log.Info("========================================================================") - for i := 0; i < len(project.Projects); i++ { - log.Infof("Project ID: %s", project.Projects[i].ProjectID) - log.Infof("Project Name: %s", project.Projects[i].Name) - log.Info("========================================================================") - } - log.Infof("Please run this command to login again:\n\n%s --login --project_id \n", os.Args[0]) - } - } else { - // Log as a warning because in some cases, the CLI might still be usable - // or the user might want to retry setup later. - log.Fatalf("failed to complete user setup: %v", err) - } - } else { - cliClients = append(cliClients, cliClient) - } + cliClients = append(cliClients, cliClient) } } return nil