feat(langchain): Add Google Gemini , Vertex AI and Ollama model support#30
feat(langchain): Add Google Gemini , Vertex AI and Ollama model support#30Thareesha98 wants to merge 7 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis pull request extends AFM's model provider support by adding Gemini and Ollama integrations. The Model schema is extended with optional Sequence DiagramsequenceDiagram
participant Client
participant create_model_provider
participant _create_gemini_model
participant _create_ollama_model
participant ChatGoogleGenerativeAI
participant ChatOllama
Client->>create_model_provider: Model with provider="gemini"
create_model_provider->>_create_gemini_model: afm_model
_create_gemini_model->>ChatGoogleGenerativeAI: model, base_url, project, location, vertexai, api_key
ChatGoogleGenerativeAI-->>Client: ChatGoogleGenerativeAI instance
Client->>create_model_provider: Model with provider="ollama"
create_model_provider->>_create_ollama_model: afm_model
_create_ollama_model->>ChatOllama: model, base_url
ChatOllama-->>Client: ChatOllama instance
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
|
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python-interpreter/packages/afm-langchain/src/afm_langchain/providers.py`:
- Around line 130-143: In _create_gemini_model, avoid unconditionally calling
_get_api_key and always injecting "api_key" into kwargs: only call
_get_api_key(api sources) when afm_model.project is falsy (non-Vertex path) so
Vertex (afm_model.project → vertexai=True) can rely on ADC; update kwargs
construction to include "api_key" only if a value was returned; alternatively
check GEMINI_API_KEY as a fallback env var before calling _get_api_key or simply
omit "api_key" when using Vertex so ChatGoogleGenerativeAI initialization can
handle its own credential fallback. Ensure changes reference
_create_gemini_model, afm_model.project, api_key, kwargs, _get_api_key, and
ChatGoogleGenerativeAI.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 4faf53c3-0919-477b-9c2c-39873b5686f3
⛔ Files ignored due to path filters (1)
python-interpreter/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (3)
python-interpreter/packages/afm-core/src/afm/models.pypython-interpreter/packages/afm-langchain/pyproject.tomlpython-interpreter/packages/afm-langchain/src/afm_langchain/providers.py
Thareesha98
force-pushed
the
feat/add-model-support
branch
from
f486733 to
2294d7b
Compare
Thareesha98
changed the title
MaryamZi
left a comment
MaryamZi
left a comment
There was a problem hiding this comment.
Shall we stick to imperative mood for commit messages?
| import ballerinax/ai.anthropic; | ||
| import ballerinax/ai.openai; | ||
|
|
||
| import ballerinax/ai.ollama; |
There was a problem hiding this comment.
| import ballerinax/ai.ollama; | |
| import ballerinax/ai.ollama; | |
| import ballerina/io; | ||
| import ballerina/lang.runtime; | ||
| import ballerina/test; | ||
| import ballerina/ai; |
| // ============================================ | ||
| // Array Schema End-to-End Tests | ||
| // ============================================ | ||
|
|
There was a problem hiding this comment.
WOuldn't change unrelated code.
| }; | ||
| var result = getModel(model); | ||
| test:assertTrue(result is ai:ModelProvider); | ||
| } No newline at end of file |
There was a problem hiding this comment.
Let's add an error test for Gemini also.
There was a problem hiding this comment.
Error test for gemini is already added before. (332 - 341 Lines)
| return new ollama:ModelProvider( | ||
| check name.ensureType(), | ||
| model.url ?: "http://localhost:11434" | ||
| ); |
There was a problem hiding this comment.
| return new ollama:ModelProvider( | |
| check name.ensureType(), | |
| model.url ?: "http://localhost:11434" | |
| ); | |
| string? url = model.url; | |
| return url is string ? | |
| new ollama:ModelProvider(check name.ensureType(), url) : | |
| new ollama:ModelProvider(check name.ensureType()); |
to rely on the default from the model provider itself.
| [ballerina] | ||
| dependencies-toml-version = "2" | ||
| distribution-version = "2201.12.10" | ||
| distribution-version = "2201.13.4" |
There was a problem hiding this comment.
Let's not bump the Ballerina version for now.
| project: str | None = None | ||
| location: str | None = None |
There was a problem hiding this comment.
Since these are relevant only for Gemini, adding them in the base Model is not quite correct. Let's consider an alternative to this.
There was a problem hiding this comment.
I think we can set extra="allow" in model_config = ConfigDict(extra="allow") and access those extra fields by extra_attrs = afm_model.model_extra
There was a problem hiding this comment.
🧹 Nitpick comments (2)
python-interpreter/packages/afm-langchain/tests/test_providers.py (2)
177-195: ⚡ Quick winAssert
locationin the explicit Vertex test.This test pins
project,vertexai, andapi_key, but it does not verify thatlocationis still forwarded on the explicit-credentials Vertex path. Adding that assertion would fully cover the Vertex contract exercised here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/tests/test_providers.py` around lines 177 - 195, Update the test_create_gemini_vertex_with_explicit_credentials test to also assert that the forwarded location is correct: after calling create_model_provider and grabbing kwargs = mock_chat_gemini.call_args[1], add an assertion that kwargs["location"] == "us-central1" so the explicit-credentials Vertex path verifies the location is passed through; the change is in the test function that references Model, ClientAuthentication, create_model_provider and mock_chat_gemini.
82-124: ⚡ Quick winAdd no-credential failure coverage for Anthropic and Gemini Studio.
Only the OpenAI path currently verifies the expected
ProviderErrorwhen neither AFM metadata nor environment variables provide credentials. Adding the same negative-path test for Anthropic and Gemini Studio would catch provider-specific regressions such as an incorrect env var lookup or an unintended silent fallback.Also applies to: 126-157
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python-interpreter/packages/afm-langchain/tests/test_providers.py` around lines 82 - 124, Add negative-path unit tests for Anthropic and Gemini Studio mirroring the OpenAI no-credential test: create Model(provider="anthropic", name="claude-sonnet-4-5") and Model(provider="gemini-studio", name="...") with no authentication, use patch.dict(os.environ, {}, clear=True) to ensure no env creds, then assert that create_model_provider(...) raises ProviderError; name the tests test_anthropic_no_credentials_failure and test_gemini_studio_no_credentials_failure and reference create_model_provider and ProviderError so CI will catch missing env-var lookups or silent fallbacks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@python-interpreter/packages/afm-langchain/tests/test_providers.py`:
- Around line 177-195: Update the
test_create_gemini_vertex_with_explicit_credentials test to also assert that the
forwarded location is correct: after calling create_model_provider and grabbing
kwargs = mock_chat_gemini.call_args[1], add an assertion that kwargs["location"]
== "us-central1" so the explicit-credentials Vertex path verifies the location
is passed through; the change is in the test function that references Model,
ClientAuthentication, create_model_provider and mock_chat_gemini.
- Around line 82-124: Add negative-path unit tests for Anthropic and Gemini
Studio mirroring the OpenAI no-credential test: create
Model(provider="anthropic", name="claude-sonnet-4-5") and
Model(provider="gemini-studio", name="...") with no authentication, use
patch.dict(os.environ, {}, clear=True) to ensure no env creds, then assert that
create_model_provider(...) raises ProviderError; name the tests
test_anthropic_no_credentials_failure and
test_gemini_studio_no_credentials_failure and reference create_model_provider
and ProviderError so CI will catch missing env-var lookups or silent fallbacks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 5f4066d4-20b5-493c-833c-3a73196991d2
📒 Files selected for processing (3)
ballerina-interpreter/Dependencies.tomlballerina-interpreter/tests/agent_test.balpython-interpreter/packages/afm-langchain/tests/test_providers.py
✅ Files skipped from review due to trivial changes (1)
- ballerina-interpreter/Dependencies.toml
Thareesha98
force-pushed
the
feat/add-model-support
branch
2 times, most recently
from
bd1f791 to
bfb9410
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python-interpreter/packages/afm-core/src/afm/models.py`:
- Line 60: The Model class has been relaxed by setting model_config =
ConfigDict(extra="allow"), which permits unknown keys and hides frontmatter
typos; revert this to keep Model strict by removing or changing the model_config
to the default (no extra="allow") so only the explicit Vertex fields (e.g.,
project, location) are accepted; update the Model class definition (inspect the
model_config variable and ConfigDict usage) to restore strict validation and
rely on the explicit fields consumed downstream (see providers that read project
and location).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 884aaa3e-079c-4b9f-a5df-d2536cde24c3
📒 Files selected for processing (3)
ballerina-interpreter/agent.balballerina-interpreter/tests/agent_test.balpython-interpreter/packages/afm-core/src/afm/models.py
🚧 Files skipped from review as they are similar to previous changes (1)
- ballerina-interpreter/agent.bal
|
|
||
| class Model(BaseModel): | ||
| model_config = ConfigDict(extra="forbid") | ||
| model_config = ConfigDict(extra="allow") |
There was a problem hiding this comment.
Keep Model strict and rely on the explicit Vertex fields.
Changing Model to extra="allow" broadens the AFM schema from rejecting unknown keys to silently accepting them. In this PR, the downstream Gemini path only reads the explicit project and location fields, so relaxing validation is not needed for the new provider support and makes frontmatter typos harder to catch. Cross-file evidence: python-interpreter/packages/afm-langchain/src/afm_langchain/providers.py:125-160 only consumes those explicit fields.
Suggested change
class Model(BaseModel):
- model_config = ConfigDict(extra="allow")
+ model_config = ConfigDict(extra="forbid")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| model_config = ConfigDict(extra="allow") | |
| model_config = ConfigDict(extra="forbid") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python-interpreter/packages/afm-core/src/afm/models.py` at line 60, The Model
class has been relaxed by setting model_config = ConfigDict(extra="allow"),
which permits unknown keys and hides frontmatter typos; revert this to keep
Model strict by removing or changing the model_config to the default (no
extra="allow") so only the explicit Vertex fields (e.g., project, location) are
accepted; update the Model class definition (inspect the model_config variable
and ConfigDict usage) to restore strict validation and rely on the explicit
fields consumed downstream (see providers that read project and location).
bfb9410 to
85a487c
Compare
Thareesha98
force-pushed
the
feat/add-model-support
branch
from
4909811 to
24659e9
Compare
3 participants
Purpose
The Python execution backend (
afm-langchain) currently only supports OpenAI and Anthropic models out of the box. With the recent rapid adoption of Google Gemini models, enterprise-grade architecture migrations across Google Cloud Platform (GCP), and the growing demand for completely local, privacy-focused execution using open-source weights, developers using Agent-Flavored Markdown (AFM) lack a native mechanism to route their agent definitions to either Google's GenAI ecosystem or local endpoints like Ollama.Furthermore, enterprise configurations often require provisioning models within strict cloud compliance boundaries (e.g., Google Vertex AI) which demand declarative parameters such as specific project identifiers and localized cloud regions straight from the source file metadata rather than scattered across global system environment settings.
Resolves #6
Goals
langchain-google-genai.langchain-ollama.ballerinax/ai.ollamaconnector, mirroring the Python provider structure.Approach
packages/afm-core/src/afm/models.py):Modified the strict
ModelPydantic class to include optionalproject: str | Noneandlocation: str | Noneproperties, allowing these keys to bypass theextra="forbid"structural layout validation block cleanly.packages/afm-langchain/pyproject.toml):Added production runtime tracking for both
langchain-google-genaiandlangchain-ollamainside package dependencies to safely tie model availability into the workspace toolchain layer.packages/afm-langchain/src/afm_langchain/providers.py):_create_gemini_model(afm_model: Model)and_create_ollama_model(afm_model: Model)utilizing safe lazy-import parsing to shield the runner from boot overhead unless explicitly called.kwargs["vertexai"] = True) triggered instantly upon discovering a declaredprojectmetadata string, ensuring data traffic smoothly hits GCP compliance zones.api_keyextraction ifprojectis defined, natively enabling Google Cloud Application Default Credentials (ADC) fallbacks for enterprise deployments.ballerina-interpreter/agent_creator.bal):Wired the
ollamamatch block to initializeollama:ModelProviderwith custom URL fallbacks, and explicitly mappedgeminito raise a descriptive "not yet supported" error.User stories
geminiwith model configurations likegemini-2.5-flashso that my agent uses Google's latest reasoning engines seamlessly..mdfile frontmatter so that execution logic runs deterministically in a secured cloud space using API keys or Application Default Credentials (ADC) without modifying system-level variables.ollamaand run models likellama3completely locally and offline without requiring API keys or incurring external transactional costs.Release note
Introduced native Google Gemini and local Ollama model support within the Python LangChain execution backend (
afm-langchain). This includes complete support for the Gemini Developer API (Google AI Studio), declarative metadata-driven onboarding for the enterprise Google Cloud Vertex AI Platform (supporting both explicit API keys and native Application Default Credentials), and zero-dependency local execution for open-source weights using Ollama.Unit tests
Python Interpreter (
packages/afm-langchain/tests/test_providers.py):Implemented mock-isolated pytest suites utilizing
unittest.mock.patchto verify configuration mapping logic (OpenAI, Anthropic, Gemini AI Studio, Vertex ADC, and Ollama) completely offline without invoking actual runtime sockets.Run via:
Security checks
Sample 1 : Enterprise Cloud Infrastructure Access (Google Cloud Vertex AI Platform)
Sample 2 : Standard Developer Access (Google AI Studio / Gemini Developer API)
Sample 3 : Ollama Model Access