imbajin commented on code in PR #335:
URL: https://github.com/apache/hugegraph-ai/pull/335#discussion_r3271797167
##########
hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py:
##########
@@ -79,16 +79,37 @@ def get_texts_embeddings(self, texts: List[str],
batch_size: int = 32) -> List[L
all_embeddings.extend([list(inner_sequence) for inner_sequence in
response])
return all_embeddings
+ def _get_embeddings_from_response(self, response) -> List[List[float]]:
Review Comment:
⚠️ 这里新增了 `_get_embeddings_from_response()` 来统一校验 `embeddings`
缺失/为空，但同步路径仍然直接索引 `self.client.embed(...)["embeddings"]`。这样 sync/async
两条路径的错误契约会继续分叉，后续如果 Ollama 响应形状变化，容易只修到 async 一侧。建议同步路径也复用这个 helper，并补一条 sync
路径缺 key / 空列表的回归测试。
##########
hugegraph-llm/src/tests/models/llms/test_litellm_client.py:
##########
@@ -0,0 +1,90 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+import unittest
+from unittest.mock import AsyncMock, patch
+
+from litellm.exceptions import APIError, BudgetExceededError
+from tenacity.retry import retry_if_exception_type
+
+from hugegraph_llm.models.llms.litellm import LiteLLMClient
+
+
+class TestLiteLLMClient(unittest.TestCase):
+ def test_budget_exceeded_error_is_not_retried(self):
+ client = LiteLLMClient(model_name="openai/gpt-4.1-mini")
+ error = BudgetExceededError(current_cost=2.0, max_budget=1.0)
+
+ with patch("hugegraph_llm.models.llms.litellm.completion",
side_effect=error) as mock_completion:
+ with self.assertRaises(BudgetExceededError):
+ client.generate(prompt="hello")
+
+ mock_completion.assert_called_once()
+
+ def test_generate_retries_api_error_and_reraises_original_exception(self):
+ client = LiteLLMClient(model_name="openai/gpt-4.1-mini")
+ error = APIError(status_code=500, message="upstream failed",
llm_provider="openai", model="gpt-4.1-mini")
+
+ with patch("hugegraph_llm.models.llms.litellm.completion",
side_effect=error) as mock_completion:
+ with self.assertRaises(APIError):
+ client.generate(prompt="hello")
+
+ self.assertEqual(mock_completion.call_count, 2)
+
+ def test_generate_retry_policy_excludes_budget_exceeded_error(self):
+ retry_policy = LiteLLMClient().generate.retry.retry
Review Comment:
⚠️ 这个测试直接断言 tenacity 的内部结构
`generate.retry.retry.exception_types`，和具体实现绑定得比较紧；一旦 tenacity
内部表示调整，即使运行时行为没变也可能失败。前面的 `mock_completion.assert_called_once()` 已经从行为上证明
`BudgetExceededError` 不会重试，建议删除这个实现细节断言，或改成只通过调用次数 / 抛出原异常来验证行为。
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
