Handle multiple response formats in OllamaEmbeddingEngine

[weikao]

The OllamaEmbeddingEngine is compatible with OpenAI

Description

Type of Change

• [*] Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation update
• Code refactoring
• Performance improvement
• Other (please specify):

Screenshots/Videos (if applicable)

Pre-submission Checklist

• [* ] I have tested my changes thoroughly before submitting this PR
• [*] This PR contains minimal changes necessary to address the issue/feature
• [*] My code follows the project's coding standards and style guidelines
• I have added tests that prove my fix is effective or that my feature works
• I have added necessary documentation (if applicable)
• All new and existing tests pass
• I have searched existing PRs to ensure this change hasn't been submitted already
• I have linked any relevant issues in the description
• My commits have clear and descriptive messages

DCO Affirmation

I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.

[pull-checklist]

Please make sure all the checkboxes are checked:

• I have tested these changes locally.
• I have reviewed the code changes.
• I have added end-to-end and unit tests (if applicable).
• I have updated the documentation and README.md file (if necessary).
• I have removed unnecessary code and debug statements.
• PR title is clear and follows the convention.
• I have tagged reviewers or team members for feedback.

[coderabbitai]

Walkthrough

This PR adjusts response handling in OllamaEmbeddingEngine's _get_embedding method to support alternative response shapes from the Ollama API. It conditionally extracts embeddings based on the presence of an "embeddings" key in the response JSON, falling back to an alternative nested structure if needed.

Changes

Cohort / File(s)	Summary
Ollama Embedding Response Parsing cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py	Adjusts _get_embedding to handle alternative Ollama response shapes by checking for "embeddings" key and extracting embedding from either the root level or nested data[0]["embedding"] path.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

• Single file modification with localized logic adjustment
• Defensive response parsing to handle API format variations
• No public API or structural changes

Possibly related PRs

• #555 — Directly modifies the same _get_embedding method to handle different Ollama response shapes.
• #1330 — Also modifies OllamaEmbeddingEngine._get_embedding, adjusting request payload construction alongside response parsing.

Suggested labels

run-checks

Suggested reviewers

• Vasilije1990
• borisarzentar

Poem

🐰 Multiple shapes dance in the API night,
Our embedding engine adapts with delight,
One path, then another—both shine equally bright,
Flexible parsing makes everything right! ✨

Pre-merge checks and finishing touches

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	❓ Inconclusive	The description includes a brief note about OpenAI compatibility and marks the change as a bug fix, but lacks detail about the specific problem being solved, the implementation approach, and reasoning for the change.	Expand the description with a clear explanation of why multiple response formats are needed, what the compatibility issue is, and how the fix addresses it.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: handling multiple response formats in OllamaEmbeddingEngine, which matches the core modification described in the raw summary.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 1e53138 and 72c20c2.

📒 Files selected for processing (1)

• cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (2)

{cognee,cognee-mcp,distributed,examples,alembic}/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

{cognee,cognee-mcp,distributed,examples,alembic}/**/*.py: Use 4-space indentation; name modules and functions in snake_case; name classes in PascalCase (Python)
Adhere to ruff rules, including import hygiene and configured line length (100)
Keep Python lines ≤ 100 characters

Files:

• cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

cognee/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

cognee/**/*.py: Public APIs in the core library should be type-annotated where practical
Prefer explicit, structured error handling and use shared logging utilities from cognee.shared.logging_utils

Files:

• cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

🪛 Pylint (4.0.2)

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

[refactor] 127-130: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it

(R1705)

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}