Added support for gemini

README.md

@@ -20,9 +20,14 @@ commands that:
 Requirements:
 
 * Go ≥ 1.24
-* A valid OpenAI API key (export `OPENAI_API_KEY`)
+* A valid API key for your chosen provider (`OPENAI_API_KEY` for OpenAI,
+  `GEMINI_API_KEY` for Gemini).
 
 ```bash
+# set your API key
+$ export OPENAI_API_KEY="sk-..." # or...
+$ export GEMINI_API_KEY="..."
+
 # install the latest directly from github
 $ go install github.com/vybdev/vyb
 
@@ -94,7 +99,7 @@ edit it manually.
 CLI knows which LLM backend to call:
 
 ```yaml
-provider: openai
+provider: openai # or "gemini"
 ```
 
 Only one key is defined for now but the document might grow in the future
@@ -109,8 +114,9 @@ we use a two-part specification:
 * **Family** – logical grouping (`gpt`, `reasoning`, …)
 * **Size**   – `large` or `small`
 
-The active provider maps the tuple to its concrete model name.  For example
-the OpenAI implementation currently resolves to:
+The active provider maps the tuple to its concrete model name.
+
+For example, the **OpenAI** implementation currently resolves to:
 
 | Family / Size | Resolved model |
 |---------------|----------------|
@@ -119,6 +125,13 @@ the OpenAI implementation currently resolves to:
 | reasoning / large | o3         |
 | reasoning / small | o4-mini    |
 
+The **Gemini** provider maps both families to the same models:
+
+| Family / Size | Resolved model                 |
+|---------------|--------------------------------|
+| *any* / large | gemini-2.5-pro-preview-06-05   |
+| *any* / small | gemini-2.5-flash-preview-05-20 |
+
 This indirection keeps templates provider-agnostic and allows you to switch
 backends without touching prompt definitions.
 
@@ -140,15 +153,15 @@ into prompts to reduce the number of files that need to be submitted in each req
 ```
 cmd/            entry-points and Cobra command wiring
   template/     YAML + Mustache definitions used by AI commands
-llm/            OpenAI API wrapper + strongly typed JSON payloads
+llm/            LLM provider wrappers + strongly typed JSON payloads
 workspace/      file selection, .gitignore handling, metadata evolution
 ```
 
 Flow of an AI command (`vyb code` for instance):
 
 1. "template" loads the prompt YAML, computes inclusion/exclusion sets.
 2. "selector" walks the workspace to gather the right files.
-3. The user & system messages are built, then sent to `llm/openai`.
+3. The user & system messages are built, then sent to `llm`.
 4. The JSON reply is validated and applied to the working tree.
 
 ---
@@ -169,4 +182,4 @@ See `cmd/template/embedded/code.vyb` for the field reference.
 * Unit tests: `go test ./...`
 * Lint / CI:   see `.github/workflows/go.yml`
 
-Feel free to open issues or PRs – all contributions are welcome!
+Feel free to open issues or PRs – all contributions are welcome!

llm/README.md

@@ -1,7 +1,10 @@
 # llm Package
 
-`llm` wraps all interaction with OpenAI and exposes strongly typed data
-structures so the rest of the codebase never has to deal with raw JSON.
+`llm` wraps all interaction with LLM providers (currently OpenAI and Gemini)
+and exposes strongly typed data structures so the rest of the codebase never
+has to deal with raw JSON.
+
+The active provider is selected based on `.vyb/config.yaml`.
 
 ## Model abstractions ⚙️
 
@@ -11,11 +14,11 @@ structures so the rest of the codebase never has to deal with raw JSON.
 | `ModelSize`    | `large`, `small`         | Coarse size tier inside a family     |
 
 The `(family, size)` tuple is later resolved by the active provider into a
-concrete model string (e.g. `GPT+Large → "GPT-4.1"`).
+concrete model string (e.g. `GPT+Large → "GPT-4.1"` for OpenAI).
 
 ## Sub-packages
 
-### `llm/openai`
+### `llm/internal/openai`
 
 * Builds requests (`model`, messages, `response_format`).
 * Retries on `rate_limit_exceeded`.
@@ -28,6 +31,13 @@ debugging.
     contexts.
   * `GetModuleExternalContexts` – produces *external* contexts in bulk.
 
+### `llm/internal/gemini`
+
+* Builds requests (`model`, messages, `generationConfig`).
+* Dumps every request/response pair to a temporary JSON file for easy
+debugging.
+* Public helpers are the same as the OpenAI provider.
+
 ### `llm/payload`
 
 Pure data & helper utilities:
@@ -41,6 +51,9 @@ Pure data & helper utilities:
 ## JSON Schema enforcement
 
 The JSON responses expected from the LLM are described under
-`llm/openai/internal/schema/schemas/*.json`.  Each request sets the
-`json_schema` field so GPT returns **validatable, deterministic** output
-that can be unmarshalled straight into Go types.
+`llm/internal/<provider>/internal/schema/schemas/*.json`. Both providers
+enforce structured JSON output to ensure responses can be unmarshalled
+straight into Go types.
+
+* **OpenAI** uses the `response_format` field with a `json_schema`.
+* **Gemini** uses the `generationConfig` field with a `responseSchema`.

llm/dispatcher.go

@@ -1,10 +1,13 @@
 package llm
 
 import (
-	"fmt"
-	"github.com/vybdev/vyb/config"
-	"github.com/vybdev/vyb/llm/internal/openai"
-	"github.com/vybdev/vyb/llm/payload"
+    "fmt"
+    "strings"
+
+    "github.com/vybdev/vyb/config"
+    "github.com/vybdev/vyb/llm/internal/gemini"
+    "github.com/vybdev/vyb/llm/internal/openai"
+    "github.com/vybdev/vyb/llm/payload"
 )
 
 // provider captures the common operations expected from any LLM backend.
@@ -15,53 +18,88 @@ import (
 // Additional methods should be appended here whenever new high-level
 // helpers are added to the llm façade.
 type provider interface {
-	GetWorkspaceChangeProposals(fam config.ModelFamily, sz config.ModelSize, systemMessage, userMessage string) (*payload.WorkspaceChangeProposal, error)
-	GetModuleContext(systemMessage, userMessage string) (*payload.ModuleSelfContainedContext, error)
-	GetModuleExternalContexts(systemMessage, userMessage string) (*payload.ModuleExternalContextResponse, error)
+    GetWorkspaceChangeProposals(fam config.ModelFamily, sz config.ModelSize, systemMessage, userMessage string) (*payload.WorkspaceChangeProposal, error)
+    GetModuleContext(systemMessage, userMessage string) (*payload.ModuleSelfContainedContext, error)
+    GetModuleExternalContexts(systemMessage, userMessage string) (*payload.ModuleExternalContextResponse, error)
 }
 
 type openAIProvider struct{}
 
+type geminiProvider struct{}
+
 func (*openAIProvider) GetWorkspaceChangeProposals(fam config.ModelFamily, sz config.ModelSize, sysMsg, userMsg string) (*payload.WorkspaceChangeProposal, error) {
-	return openai.GetWorkspaceChangeProposals(fam, sz, sysMsg, userMsg)
+    return openai.GetWorkspaceChangeProposals(fam, sz, sysMsg, userMsg)
 }
 
 func (*openAIProvider) GetModuleContext(sysMsg, userMsg string) (*payload.ModuleSelfContainedContext, error) {
-	return openai.GetModuleContext(sysMsg, userMsg)
+    return openai.GetModuleContext(sysMsg, userMsg)
 }
 
 func (*openAIProvider) GetModuleExternalContexts(sysMsg, userMsg string) (*payload.ModuleExternalContextResponse, error) {
-	return openai.GetModuleExternalContexts(sysMsg, userMsg)
+    return openai.GetModuleExternalContexts(sysMsg, userMsg)
+}
+
+// -----------------------------------------------------------------------------
+//  Gemini provider implementation – WorkspaceChangeProposals hooked up
+// -----------------------------------------------------------------------------
+
+func mapGeminiModel(fam config.ModelFamily, sz config.ModelSize) (string, error) {
+    switch sz {
+    case config.ModelSizeSmall:
+        return "gemini-2.5-flash-preview-05-20", nil
+    case config.ModelSizeLarge:
+        return "gemini-2.5-pro-preview-06-05", nil
+    default:
+        return "", fmt.Errorf("gemini: unsupported model size %s", sz)
+    }
+}
+
+func (*geminiProvider) GetWorkspaceChangeProposals(fam config.ModelFamily, sz config.ModelSize, sysMsg, userMsg string) (*payload.WorkspaceChangeProposal, error) {
+    return gemini.GetWorkspaceChangeProposals(fam, sz, sysMsg, userMsg)
+}
+
+func (*geminiProvider) GetModuleContext(sysMsg, userMsg string) (*payload.ModuleSelfContainedContext, error) {
+    return gemini.GetModuleContext(sysMsg, userMsg)
 }
 
+func (*geminiProvider) GetModuleExternalContexts(sysMsg, userMsg string) (*payload.ModuleExternalContextResponse, error) {
+    return gemini.GetModuleExternalContexts(sysMsg, userMsg)
+}
+
+// -----------------------------------------------------------------------------
+//  Public façade helpers remain unchanged (dispatcher section).
+// -----------------------------------------------------------------------------
+
 func GetModuleExternalContexts(cfg *config.Config, sysMsg, userMsg string) (*payload.ModuleExternalContextResponse, error) {
-	if provider, err := resolveProvider(cfg); err != nil {
-		return nil, err
-	} else {
-		return provider.GetModuleExternalContexts(sysMsg, userMsg)
-	}
+    if provider, err := resolveProvider(cfg); err != nil {
+        return nil, err
+    } else {
+        return provider.GetModuleExternalContexts(sysMsg, userMsg)
+    }
 }
 
 func GetModuleContext(cfg *config.Config, sysMsg, userMsg string) (*payload.ModuleSelfContainedContext, error) {
-	if provider, err := resolveProvider(cfg); err != nil {
-		return nil, err
-	} else {
-		return provider.GetModuleContext(sysMsg, userMsg)
-	}
+    if provider, err := resolveProvider(cfg); err != nil {
+        return nil, err
+    } else {
+        return provider.GetModuleContext(sysMsg, userMsg)
+    }
 }
 func GetWorkspaceChangeProposals(cfg *config.Config, fam config.ModelFamily, sz config.ModelSize, sysMsg, userMsg string) (*payload.WorkspaceChangeProposal, error) {
-	if provider, err := resolveProvider(cfg); err != nil {
-		return nil, err
-	} else {
-		return provider.GetWorkspaceChangeProposals(fam, sz, sysMsg, userMsg)
-	}
+    if provider, err := resolveProvider(cfg); err != nil {
+        return nil, err
+    } else {
+        return provider.GetWorkspaceChangeProposals(fam, sz, sysMsg, userMsg)
+    }
 }
 
 func resolveProvider(cfg *config.Config) (provider, error) {
-	switch cfg.Provider {
-	case "openai":
-		return &openAIProvider{}, nil
-	default:
-		return nil, fmt.Errorf("unknown provider: %s", cfg.Provider)
-	}
+    switch strings.ToLower(cfg.Provider) {
+    case "openai":
+        return &openAIProvider{}, nil
+    case "gemini":
+        return &geminiProvider{}, nil
+    default:
+        return nil, fmt.Errorf("unknown provider: %s", cfg.Provider)
+    }
 }

llm/dispatcher_test.go

@@ -6,23 +6,35 @@ import (
     "github.com/vybdev/vyb/config"
 )
 
-// TestResolveProvider verifies that the dispatcher returns the expected
-// concrete implementation for known providers and fails for unknown ones.
-func TestResolveProvider(t *testing.T) {
-    // 1. Happy-path – "openai" should map to *openAIProvider.
-    cfg := &config.Config{Provider: "openai"}
+// TestMapGeminiModel ensures that the (family,size) tuple is translated to
+// the correct concrete model identifier and that unsupported sizes are
+// properly rejected.
+func TestMapGeminiModel(t *testing.T) {
+    t.Parallel()
 
-    p, err := resolveProvider(cfg)
-    if err != nil {
-        t.Fatalf("unexpected error resolving provider: %v", err)
+    cases := []struct {
+        fam  config.ModelFamily
+        size config.ModelSize
+        want string
+    }{
+        {config.ModelFamilyGPT, config.ModelSizeSmall, "gemini-2.5-flash-preview-05-20"},
+        {config.ModelFamilyGPT, config.ModelSizeLarge, "gemini-2.5-pro-preview-06-05"},
+        {config.ModelFamilyReasoning, config.ModelSizeSmall, "gemini-2.5-flash-preview-05-20"},
+        {config.ModelFamilyReasoning, config.ModelSizeLarge, "gemini-2.5-pro-preview-06-05"},
     }
-    if _, ok := p.(*openAIProvider); !ok {
-        t.Fatalf("resolveProvider returned %T, want *openAIProvider", p)
+
+    for _, c := range cases {
+        got, err := mapGeminiModel(c.fam, c.size)
+        if err != nil {
+            t.Fatalf("mapGeminiModel(%s,%s) returned unexpected error: %v", c.fam, c.size, err)
+        }
+        if got != c.want {
+            t.Fatalf("mapGeminiModel(%s,%s) = %q, want %q", c.fam, c.size, got, c.want)
+        }
     }
 
-    // 2. Unknown provider should surface an error.
-    cfg.Provider = "doesnotexist"
-    if _, err := resolveProvider(cfg); err == nil {
-        t.Fatalf("expected error for unknown provider, got nil")
+    // Ensure an unsupported size triggers an error.
+    if _, err := mapGeminiModel(config.ModelFamilyGPT, config.ModelSize("medium")); err == nil {
+        t.Fatalf("expected error for unsupported model size, got nil")
     }
 }