lint

src/canopy/chat_engine/chat_engine.py

@@ -154,7 +154,8 @@ def __init__(self,
         if max_context_tokens + self._system_prompt_tokens > max_prompt_tokens:
             raise ValueError(
                 f"Not enough token budget for knowledge base context. The system prompt"
-                f" is taking {self._system_prompt_tokens } tokens, and together with the "
+                f" is taking {self._system_prompt_tokens } tokens,"
+                f" and together with the "
                 f"configured max context tokens {max_context_tokens} it exceeds "
                 f"max_prompt_tokens of {self.max_prompt_tokens}"
             )

src/canopy/chat_engine/history_pruner/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Tuple, Optional
+from typing import Optional
 
 from canopy.tokenizer import Tokenizer
 from canopy.models.data_models import Messages, SystemMessage
@@ -30,7 +30,9 @@ def _max_tokens_history(self,
                             system_prompt: Optional[str] = None,
                             context: Optional[str] = None, ) -> int:
         if system_prompt is not None:
-            max_tokens -= self._tokenizer.messages_token_count([SystemMessage(content=system_prompt)])
+            max_tokens -= self._tokenizer.messages_token_count(
+                [SystemMessage(content=system_prompt)]
+            )
 
         if context is not None:
             max_tokens -= self._tokenizer.token_count(context)

src/canopy/chat_engine/history_pruner/raising.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Optional
+from typing import Optional
 
 from canopy.chat_engine.history_pruner.base import HistoryPruner
 from canopy.models.data_models import Messages
@@ -11,12 +11,16 @@ def build(self,
               max_tokens: int,
               system_prompt: Optional[str] = None,
               context: Optional[str] = None, ) -> Messages:
-        max_tokens_history = self._max_tokens_history(max_tokens, system_prompt, context)
+        max_tokens_history = self._max_tokens_history(max_tokens,
+                                                      system_prompt,
+                                                      context)
         token_count = self._tokenizer.messages_token_count(chat_history)
         if token_count > max_tokens:
             raise ValueError(f"The history require {token_count} tokens, "
                              f"which exceeds the calculated limit for history "
-                             f"of {max_tokens_history} tokens left for history out of {max_tokens} tokens allowed in context window.")
+                             f"of {max_tokens_history} tokens left for"
+                             f" history out of {max_tokens} tokens"
+                             f" allowed in context window.")
         return chat_history
 
     async def abuild(self,

src/canopy/chat_engine/history_pruner/recent.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Optional
+from typing import Optional
 
 from canopy.chat_engine.history_pruner.base import HistoryPruner
 from canopy.models.data_models import Messages
@@ -17,7 +17,9 @@ def build(self,
               system_prompt: Optional[str] = None,
               context: Optional[str] = None,
               ) -> Messages:
-        max_tokens_history = self._max_tokens_history(max_tokens, system_prompt, context)
+        max_tokens_history = self._max_tokens_history(max_tokens,
+                                                      system_prompt,
+                                                      context)
         token_count = self._tokenizer.messages_token_count(chat_history)
         if token_count < max_tokens:
             return chat_history
@@ -27,7 +29,9 @@ def build(self,
         if token_count > max_tokens:
             raise ValueError(f"The {self._min_history_messages} most recent messages in"
                              f" history require {token_count} tokens, which exceeds the"
-                             f" calculated limit for history of {max_tokens_history} tokens out of total {max_tokens} allowed in context window.")
+                             f" calculated limit for history of {max_tokens_history}"
+                             f" tokens out of total {max_tokens} allowed"
+                             f" in context window.")
 
         for message in reversed(chat_history[:-self._min_history_messages]):
             token_count = self._tokenizer.messages_token_count(

src/canopy/llm/openai.py

@@ -14,7 +14,7 @@
 from canopy.llm import BaseLLM
 from canopy.llm.models import Function
 from canopy.models.api_models import ChatResponse, StreamingChatChunk
-from canopy.models.data_models import Messages, Context, SystemMessage, StringContextContent
+from canopy.models.data_models import Messages, Context, SystemMessage
 
 
 def _format_openai_error(e):
@@ -121,7 +121,8 @@ def chat_completion(self,
             system_message = system_prompt
         else:
             system_message = system_prompt + f"\nContext: {context.to_text()}"
-        messages = [SystemMessage(content=system_message).dict()] + [m.dict() for m in chat_history]
+        messages = [SystemMessage(content=system_message).dict()
+                    ] + [m.dict() for m in chat_history]
         try:
             response = self._client.chat.completions.create(model=self.model_name,
                                                             messages=messages,
@@ -211,7 +212,8 @@ def enforced_function_call(self,
         function_dict = cast(ChatCompletionToolParam,
                              {"type": "function", "function": function.dict()})
 
-        messages = [SystemMessage(content=system_prompt).dict()] + [m.dict() for m in chat_history]
+        messages = [SystemMessage(content=system_prompt).dict()
+                    ] + [m.dict() for m in chat_history]
         try:
             chat_completion = self._client.chat.completions.create(
                 messages=messages,
@@ -248,7 +250,6 @@ async def achat_completion(self,
                                           Iterable[StreamingChatChunk]]:
         raise NotImplementedError()
 
-
     async def aenforced_function_call(self,
                                       system_prompt: str,
                                       chat_history: Messages,

src/canopy_server/app.py

@@ -238,7 +238,10 @@ async def health_check() -> HealthStatus:
 
     try:
         msg = UserMessage(content="This is a health check. Are you alive? Be concise")
-        await run_in_threadpool(llm.chat_completion, system_prompt="hi", chat_history=[msg], max_tokens=5)
+        await run_in_threadpool(llm.chat_completion,
+                                system_prompt="hi",
+                                chat_history=[msg],
+                                max_tokens=5)
     except Exception as e:
         err_msg = f"Failed to communicate with {llm.__class__.__name__}"
         logger.exception(err_msg)

tests/system/llm/test_anyscale.py

@@ -88,9 +88,11 @@ def test_chat_completion_with_context(anyscale_llm, messages):
     response = anyscale_llm.chat_completion(system_prompt=SYSTEM_PROMPT,
                                             chat_history=messages,
                                             context=Context(
-                                                content=StringContextContent(__root__="context from kb"),
-                                                num_tokens=5
-                                          ))
+                                                content=StringContextContent(
+                                                    __root__="context from kb"
+                                                ),
+                                                num_tokens=5)
+                                            )
     assert_chat_completion(response)
 
 

tests/system/llm/test_openai.py

@@ -37,6 +37,7 @@ def assert_function_call_format(result):
 def model_name():
     return "gpt-3.5-turbo-0613"
 
+
 @pytest.fixture
 def messages():
     # Create a list of MessageBase objects
@@ -46,6 +47,7 @@ def messages():
                     content="Hello, user. How can I assist you?")
     ]
 
+
 @pytest.fixture
 def function_query_knowledgebase():
     return Function(
@@ -104,7 +106,9 @@ def test_chat_completion_with_context(openai_llm, messages):
     response = openai_llm.chat_completion(system_prompt=SYSTEM_PROMPT,
                                           chat_history=messages,
                                           context=Context(
-                                              content=StringContextContent(__root__="context from kb"),
+                                              content=StringContextContent(
+                                                  __root__="context from kb"
+                                              ),
                                               num_tokens=5
                                           ))
     assert_chat_completion(response)

tests/system/query_generator/test_query_generator_integration.py

@@ -1,5 +1,3 @@
-from unittest.mock import create_autospec
-
 import pytest
 
 from canopy.tokenizer.tokenizer import Tokenizer # noqa

tests/unit/chat_engine/test_chat_engine.py

@@ -10,7 +10,6 @@
                                                             ContextQueryResult,
                                                             StuffingContextContent, )
 from canopy.llm import BaseLLM
-from canopy.models.data_models import SystemMessage
 from canopy.models.api_models import ChatResponse, _Choice, TokenCounts
 from canopy.models.data_models import MessageBase, Query, Context, Role
 from .. import random_words
@@ -73,7 +72,6 @@ def _get_inputs_and_expected(self,
             num_tokens=1  # TODO: This is a dummy value. Need to improve.
         )
 
-
         mock_chat_response = ChatResponse(
             id='chatcmpl-7xuuGZzniUGiqxDSTJnqwb0l1xtfp',
             object='chat.completion',