fix name typo, add unit test to ollama

run-llama · Feb 29, 2024 · e6437b2 · e6437b2
1 parent f508073
commit e6437b2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 7 deletions.
diff --git a/llama-index-integrations/llms/llama-index-llms-ollama/llama_index/llms/ollama/base.py b/llama-index-integrations/llms/llama-index-llms-ollama/llama_index/llms/ollama/base.py
@@ -20,7 +20,7 @@
 DEFAULT_REQUEST_TIMEOUT = 30.0
 
 
-def get_addtional_kwargs(
+def get_additional_kwargs(
     response: Dict[str, Any], exclude: Tuple[str, ...]
 ) -> Dict[str, Any]:
     return {k: v for k, v in response.items() if k not in exclude}
@@ -109,12 +109,12 @@ def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
                 message=ChatMessage(
                     content=message.get("content"),
                     role=MessageRole(message.get("role")),
-                    additional_kwargs=get_addtional_kwargs(
+                    additional_kwargs=get_additional_kwargs(
                         message, ("content", "role")
                     ),
                 ),
                 raw=raw,
-                additional_kwargs=get_addtional_kwargs(raw, ("message",)),
+                additional_kwargs=get_additional_kwargs(raw, ("message",)),
             )
 
     @llm_chat_callback()
@@ -156,13 +156,15 @@ def stream_chat(
                             message=ChatMessage(
                                 content=text,
                                 role=MessageRole(message.get("role")),
-                                additional_kwargs=get_addtional_kwargs(
+                                additional_kwargs=get_additional_kwargs(
                                     message, ("content", "role")
                                 ),
                             ),
                             delta=delta,
                             raw=chunk,
-                            additional_kwargs=get_addtional_kwargs(chunk, ("message",)),
+                            additional_kwargs=get_additional_kwargs(
+                                chunk, ("message",)
+                            ),
                         )
 
     @llm_completion_callback()
@@ -188,7 +190,7 @@ def complete(
             return CompletionResponse(
                 text=text,
                 raw=raw,
-                additional_kwargs=get_addtional_kwargs(raw, ("response",)),
+                additional_kwargs=get_additional_kwargs(raw, ("response",)),
             )
 
     @llm_completion_callback()
@@ -220,7 +222,7 @@ def stream_complete(
                             delta=delta,
                             text=text,
                             raw=chunk,
-                            additional_kwargs=get_addtional_kwargs(
+                            additional_kwargs=get_additional_kwargs(
                                 chunk, ("response",)
                             ),
                         )
diff --git a/llama-index-integrations/llms/llama-index-llms-ollama/tests/test_utils.py b/llama-index-integrations/llms/llama-index-llms-ollama/tests/test_utils.py
@@ -0,0 +1,12 @@
+from llama_index.llms.ollama.base import get_additional_kwargs
+
+
+def test_get_additional_kwargs():
+    response = {"key1": "value1", "key2": "value2", "exclude_me": "value3"}
+    exclude = ("exclude_me", "exclude_me_too")
+
+    expected = {"key1": "value1", "key2": "value2"}
+
+    actual = get_additional_kwargs(response, exclude)
+
+    assert actual == expected