feat: Ollama support #148
feat: Ollama support #148
Conversation
AAClause
force-pushed
the
ollamaProvider
branch
from
2b5ee46 to
4fc7777
Compare
AAClause
force-pushed
the
ollamaProvider
branch
2 times, most recently
from
6be1b51 to
01b03c0
Compare
AAClause
force-pushed
the
ollamaProvider
branch
2 times, most recently
from
4926864 to
5943160
Compare
AAClause
force-pushed
the
ollamaProvider
branch
from
0e0e06a to
8c276dd
Compare
AAClause
force-pushed
the
ollamaProvider
branch
from
8c276dd to
676770f
Compare
AAClause
force-pushed
the
customBaseURL
branch
from
bc5ddf8 to
d669042
Compare
AAClause
force-pushed
the
ollamaProvider
branch
from
676770f to
e0f8d23
Compare
AAClause
force-pushed
the
ollamaProvider
branch
from
bf61c49 to
1dcfef0
Compare
clementb49
force-pushed
the
ollamaProvider
branch
from
1dcfef0 to
dfbae03
Compare
WalkthroughThis update adds support for integrating with the Ollama API. A new provider instance for "Ollama" is registered in the provider module, and an accompanying engine class ( Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant E as OllamaEngine
participant C as Ollama Client
participant A as Ollama API
U->>E: Request chat completion
E->>E: Prepare message request
E->>C: Call chat(parameters)
C->>A: Send chat request
A-->>C: Return response (streaming/non-streaming)
C-->>E: Deliver chat response
E-->>U: Provide chat completion result
Possibly related PRs
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (3)
basilisk/provider.py(1 hunks)basilisk/provider_engine/ollama_engine.py(1 hunks)pyproject.toml(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
basilisk/provider_engine/ollama_engine.py
1-1: No such file or directory (os error 2)
(E902)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: build_app / build_windows (x64)
🔇 Additional comments (8)
basilisk/provider_engine/ollama_engine.py (6)
1-2: Well-structured module docstring.
This docstring clearly states the purpose of the Ollama provider engine.🧰 Tools
🪛 Ruff (0.8.2)
1-1: No such file or directory (os error 2)
(E902)
23-25: Class definition appears clear and concise.
The OllamaEngine class is a straightforward extension of BaseEngine, with a docstring that provides a succinct overview.
26-30: Confirm IMAGE support alignment.
You've declaredProviderCapability.IMAGEhere. However, the engine only supports local image attachments in practice. If that's intentional, this is fine; otherwise, consider clarifying the partial image support in the class docstring.
64-75: Simple and effective property for client creation.
This property correctly handles the base URL override. No immediate concerns.
144-157: Straightforward streaming response handling.
This method efficiently iterates and yields content chunks. Implementation looks clean and correct for streaming.
158-176: Non-streaming response assignment is clear.
The method cleanly updates theMessageBlockwith the assistant's content. No issues spotted.basilisk/provider.py (1)
119-129: New provider registration looks good.
Ollama is set to run locally on the default port. If this is intentional, great. Otherwise, consider verifying that this is the desired default base URL or if HTTPS is possible.pyproject.toml (1)
54-54: New dependency version check.
Verify^0.4.7meets all required capabilities and security patches for your application.
| @cached_property | ||
| @measure_time | ||
| def models(self) -> list[ProviderAIModel]: | ||
| """Get Ollama models. | ||
|
|
||
| Returns: | ||
| A list of provider AI models. | ||
| """ | ||
| models = [] | ||
| models_list = self.client.list().models | ||
| for model in models_list: | ||
| info = self.client.show(model.model) | ||
| context_length = 0 | ||
| description = json.dumps(info.modelinfo, indent=2) | ||
| description += f"\n\n{info.license}" | ||
| for k, v in info.modelinfo.items(): | ||
| if k.endswith("context_length"): | ||
| context_length = v | ||
| models.append( | ||
| ProviderAIModel( | ||
| id=model.model, | ||
| name=model.model, | ||
| description=description, | ||
| context_window=context_length, | ||
| max_output_tokens=0, | ||
| max_temperature=2, | ||
| default_temperature=1, | ||
| vision=True, | ||
| ) | ||
| ) | ||
|
|
||
| return models | ||
|
|
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Validate handling of nullable or missing fields.
The logic for parsing info.license and info.modelinfo (especially context length) looks good. However, consider robust error handling or a helpful fallback if these fields are missing or unexpected in the Ollama response.
| def completion( | ||
| self, | ||
| new_block: MessageBlock, | ||
| conversation: Conversation, | ||
| system_message: Message | None, | ||
| **kwargs, | ||
| ) -> ChatResponse | Iterator[ChatResponse]: | ||
| """Get completion from Ollama. | ||
|
|
||
| Args: | ||
| new_block: The new message block. | ||
| conversation: The conversation instance. | ||
| system_message: The system message, if any. | ||
| **kwargs: Additional keyword arguments. | ||
|
|
||
| Returns: | ||
| The chat response or an iterator of chat responses. | ||
| """ | ||
| super().completion(new_block, conversation, system_message, **kwargs) | ||
| params = { | ||
| "model": new_block.model.model_id, | ||
| "messages": self.get_messages( | ||
| new_block, conversation, system_message | ||
| ), | ||
| # "temperature": new_block.temperature, | ||
| # "top_p": new_block.top_p, | ||
| "stream": new_block.stream, | ||
| } | ||
| params.update(kwargs) | ||
| return self.client.chat(**params) | ||
|
|
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Consider exposing temperature or top_p parameters.
You've commented out the temperature and top_p fields. If future usage requires them, ensure they are implemented consistently or removed if not needed.
| def prepare_message_request(self, message: Message): | ||
| """Prepare message request for Ollama. | ||
|
|
||
| Args: | ||
| message: The message to prepare. | ||
|
|
||
| Returns: | ||
| The prepared message request. | ||
| """ | ||
| super().prepare_message_request(message) | ||
| images = [] | ||
| if message.attachments: | ||
| for attachment in message.attachments: | ||
| location = str(attachment.location) | ||
| # TODO: fix `attachment.type` (currently returns IMAGE_UNKNOWN instead of IMAGE_LOCAL) | ||
| # best way: `if attachment.type != ImageFileTypes.IMAGE_LOCAL:`` | ||
| if ( | ||
| location.startswith("http") | ||
| or location.startswith("data:") | ||
| or location.startswith("memory://") | ||
| ): | ||
| log.warning( | ||
| f"Received unsupported image type: {attachment.type}, {attachment.location}" | ||
| ) | ||
| raise NotImplementedError( | ||
| "Only local images are supported for Ollama" | ||
| ) | ||
| images.append(attachment.location) | ||
| return { | ||
| "role": message.role.value, | ||
| "content": message.content, | ||
| "images": images, | ||
| } | ||
|
|
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Partial local-image-only support.
The direct raise of NotImplementedError for non-local images is a pragmatic approach for now. If you plan to expand to remote images, consider returning a structured error or partial fallback.
clementb49
force-pushed
the
ollamaProvider
branch
from
dfbae03 to
6e5da4b
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (3)
basilisk/provider.py(1 hunks)basilisk/provider_engine/ollama_engine.py(1 hunks)pyproject.toml(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
basilisk/provider_engine/ollama_engine.py
1-1: No such file or directory (os error 2)
(E902)
🔇 Additional comments (5)
basilisk/provider_engine/ollama_engine.py (4)
42-49: Validate handling of nullable or missing fields.The logic here doesn't account for the possibility that
info.licenseorinfo.modelinfomight be absent or malformed in certain Ollama responses. Consider robust error handling or a fallback if these fields are missing or unexpected.
101-102: Consider exposing temperature or top_p parameters.You’ve commented out the
temperatureandtop_pfields. If future usage requires them, ensure they’re implemented consistently or removed if not needed.
125-135: Partial local-image-only support.Raising
NotImplementedErrorfor non-local images is a pragmatic approach for now. If you plan to expand support for remote images, consider returning a structured error or fallback for a more graceful user experience.
158-175: Verify existence of response keys to avoid KeyError.The code directly accesses
response["message"]["content"]without fallback checks. If the response is missing the "message" or "content" keys, it will crash. Consider safe lookups or fallback logic to ensure reliability.pyproject.toml (1)
54-54:✅ Verification successful
Verify the added Ollama dependency.
Ensure the version
^0.4.7is valid, up to date, and free of vulnerabilities.Use this script to check for the latest version and any advisories:
🏁 Script executed:
#!/bin/bash # Check the latest version from PyPI: curl -s https://pypi.org/pypi/ollama/json | jq '.info.version' # Check for security advisories: gh api graphql -f query=' { securityVulnerabilities(first: 5, ecosystem: PIP, package: "ollama") { nodes { advisory { summary severity publishedAt } vulnerableVersionRange firstPatchedVersion { identifier } } } } 'Length of output: 411
Ollama dependency verified: using version ^0.4.7 is up to date and secure.
The dependency check shows that the installed version matches the latest release from PyPI (
0.4.7) and no security advisories were reported via the GitHub API.No changes are necessary.
| Provider( | ||
| id="ollama", | ||
| name="Ollama", | ||
| base_url="http://127.0.0.1:11434", | ||
| api_type=ProviderAPIType.OLLAMA, | ||
| organization_mode_available=False, | ||
| require_api_key=False, | ||
| env_var_name_api_key="OLLAMA_API_KEY", | ||
| engine_cls_path="basilisk.provider_engine.ollama_engine.OllamaEngine", | ||
| allow_custom_base_url=True, | ||
| ), |
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Consider robust handling of local vs. remote usage.
With require_api_key=False and a local base URL, the provider relies on an accessible local Ollama instance. If a remote configuration is desired in the future, consider providing fallback or guidance for scenarios where the local server is not running, improving user experience.
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (1)
basilisk/provider_engine/ollama_engine.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
basilisk/provider_engine/ollama_engine.py
1-1: No such file or directory (os error 2)
(E902)
🔇 Additional comments (3)
basilisk/provider_engine/ollama_engine.py (3)
32-64: Handle missinginfo.licenseorinfo.modelinfo.
If thelicenseormodelinfofields are unexpectedly missing, this code might fail. Consider providing a fallback or more robust error handling.
57-58: Consider making temperature parameters user-configurable.These default temperature values are defined but never exposed. If you plan to leverage them, consider allowing developers to override or adjust them via configuration or parameters. Otherwise, remove them to reduce clutter.
107-133: Support or gracefully handle remote images.Currently, remote images cause an immediate
NotImplementedError. If you need future support for remote images, consider implementing a structured fallback or partial handling rather than a hard failure.
| class OllamaEngine(BaseEngine): | ||
| """Engine implementation for Ollama API integration.""" | ||
|
|
||
| capabilities: set[ProviderCapability] = { | ||
| ProviderCapability.TEXT, | ||
| ProviderCapability.IMAGE, | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick (assertive)
Add unit tests for the new engine.
This engine class is critical for Ollama API integration yet it lacks direct test coverage. Please consider adding corresponding tests to maintain code reliability and facilitate future refactoring.
Would you like me to provide an initial test file scaffold for OllamaEngine?
| role=MessageRoleEnum.ASSISTANT, | ||
| content=response["message"]["content"], | ||
| ) | ||
| return new_block |
There was a problem hiding this comment.
Add fallback or error handling for missing content in responses.
This line can raise a KeyError if "message" or "content" are absent from the response. Consider adding a safe fallback or raising a more descriptive error for debugging.
Here’s a potential fix:
- content=response["message"]["content"],
+ content=response.get("message", {}).get("content") or ""📝 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.
| role=MessageRoleEnum.ASSISTANT, | |
| content=response["message"]["content"], | |
| ) | |
| return new_block | |
| role=MessageRoleEnum.ASSISTANT, | |
| - content=response["message"]["content"], | |
| + content=response.get("message", {}).get("content") or "", | |
| ) | |
| return new_block |
Summary by CodeRabbit
New Features
Chores
ollamapackage for the new provider integration.