Backend: mock unit tests (#908)

* chore(backend): add type hints to test files

* feat(backend): improved typehints and decorators for deployments

* feat(backend): removed dependency on cohere api key for testing

Author: ezawadski

.github/workflows/backend_integration_tests.yml

@@ -13,40 +13,47 @@ on:
 jobs:
   pytest:
     permissions: write-all
-    environment: development
+    # environment: development
     runs-on: ubuntu-latest
 
     steps:
       - name: Checkout repo
         uses: actions/checkout@v3
+
       - uses: actions/setup-python@v5
         with:
           python-version: '3.11'
           cache: 'pip'
+
       - name: Install poetry
         uses: snok/install-poetry@v1
         with:
           virtualenvs-create: true
           virtualenvs-in-project: true
           virtualenvs-path: .venv
           installer-parallel: true
+
       - name: Load cached venv
         id: cached-poetry-dependencies
         uses: actions/cache@v4
         with:
           path: .venv
           key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
+
       - name: Install dependencies
         if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
         run: poetry install --with dev --no-interaction --no-root
+
       - name: Setup test DB container
         run: make test-db
+
       - name: Test with pytest
         if: github.actor != 'dependabot[bot]'
         run: |
           make run-integration-tests
         env:
           PYTHONPATH: src
+
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v4.0.1
         with:

.github/workflows/backend_unit_tests.yml

@@ -19,34 +19,41 @@ jobs:
     steps:
       - name: Checkout repo
         uses: actions/checkout@v3
+
       - uses: actions/setup-python@v5
         with:
           python-version: '3.11'
           cache: 'pip'
+
       - name: Install poetry
         uses: snok/install-poetry@v1
         with:
           virtualenvs-create: true
           virtualenvs-in-project: true
           virtualenvs-path: .venv
           installer-parallel: true
+
       - name: Load cached venv
         id: cached-poetry-dependencies
         uses: actions/cache@v4
         with:
           path: .venv
           key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
+
       - name: Install dependencies
         if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
         run: poetry install --with dev --no-interaction --no-root
+
       - name: Setup test DB container
         run: make test-db
+
       - name: Test with pytest
         if: github.actor != 'dependabot[bot]'
         run: |
           make run-unit-tests
         env:
           PYTHONPATH: src
+
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v4.0.1
         with:

src/backend/chat/collate.py

@@ -60,7 +60,7 @@ async def rerank_and_chunk(
             reranked_results[tool_call_hashable] = tool_result
             continue
 
-        chunked_outputs = []
+        chunked_outputs: list[dict[str, Any]] = []
         for output in tool_result["outputs"]:
             text = output.get("text")
 
@@ -78,7 +78,7 @@ async def rerank_and_chunk(
 
         res = await model.invoke_rerank(
             query=query,
-            documents=chunked_outputs,
+            documents=[output["text"] for output in chunked_outputs],
             ctx=ctx,
         )
 
@@ -102,7 +102,12 @@ async def rerank_and_chunk(
     return list(reranked_results.values())
 
 
-def chunk(content, compact_mode=False, soft_word_cut_off=100, hard_word_cut_off=300):
+def chunk(
+    content: str,
+    compact_mode: bool = False,
+    soft_word_cut_off: int = 100,
+    hard_word_cut_off: int = 300,
+) -> list[str]:
     if compact_mode:
         content = content.replace("\n", " ")
 
@@ -139,7 +144,7 @@ def chunk(content, compact_mode=False, soft_word_cut_off=100, hard_word_cut_off=
     return chunks
 
 
-def to_dict(obj):
+def to_dict(obj) -> dict:
     return json.loads(
         json.dumps(
             obj, default=lambda o: o.__dict__ if hasattr(o, "__dict__") else str(o)

src/backend/chat/custom/utils.py

@@ -24,7 +24,7 @@ def get_deployment(name: str, session: Session, ctx: Context, **kwargs: Any) ->
     kwargs["ctx"] = ctx
     try:
         deployment = deployment_service.get_deployment_instance_by_name(session, name, **kwargs)
-    except DeploymentNotFoundError:
+    except (DeploymentNotFoundError, Exception):
         deployment = deployment_service.get_default_deployment_instance(**kwargs)
 
     return deployment

src/backend/model_deployments/azure.py

@@ -1,4 +1,4 @@
-from typing import Any, AsyncGenerator, Dict, List
+from typing import Any, AsyncGenerator
 
 import cohere
 
@@ -43,33 +43,33 @@ def __init__(self, **kwargs: Any):
             base_url=self.chat_endpoint_url, api_key=self.api_key
         )
 
-    @classmethod
-    def name(cls) -> str:
+    @staticmethod
+    def name() -> str:
         return "Azure"
 
-    @classmethod
-    def env_vars(cls) -> List[str]:
+    @staticmethod
+    def env_vars() -> list[str]:
         return [AZURE_API_KEY_ENV_VAR, AZURE_CHAT_URL_ENV_VAR]
 
-    @classmethod
-    def rerank_enabled(cls) -> bool:
+    @staticmethod
+    def rerank_enabled() -> bool:
         return False
 
     @classmethod
-    def list_models(cls) -> List[str]:
+    def list_models(cls) -> list[str]:
         if not cls.is_available():
             return []
 
         return cls.DEFAULT_MODELS
 
-    @classmethod
-    def is_available(cls) -> bool:
+    @staticmethod
+    def is_available() -> bool:
         return (
             AzureDeployment.default_api_key is not None
             and AzureDeployment.default_chat_endpoint_url is not None
         )
 
-    async def invoke_chat(self, chat_request: CohereChatRequest) -> Any:
+    async def invoke_chat(self, chat_request: CohereChatRequest, **kwargs) -> Any:
         response = self.client.chat(
             **chat_request.model_dump(exclude={"stream", "file_ids", "agent_id"}),
         )
@@ -86,6 +86,6 @@ async def invoke_chat_stream(
             yield to_dict(event)
 
     async def invoke_rerank(
-        self, query: str, documents: List[Dict[str, Any]], ctx: Context
+        self, query: str, documents: list[str], ctx: Context, **kwargs
     ) -> Any:
         return None

src/backend/model_deployments/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, AsyncGenerator, Dict, List
+from typing import Any
 
 from backend.config.settings import Settings
 from backend.schemas.cohere_chat import CohereChatRequest
@@ -25,32 +25,32 @@ def __init__(self, db_id=None, **kwargs: Any):
     def id(cls) -> str:
         return cls.db_id if cls.db_id else cls.name().replace(" ", "_").lower()
 
-    @classmethod
+    @staticmethod
     @abstractmethod
-    def name(cls) -> str: ...
+    def name() -> str: ...
 
-    @classmethod
+    @staticmethod
     @abstractmethod
-    def env_vars(cls) -> List[str]: ...
+    def env_vars() -> list[str]: ...
 
-    @classmethod
+    @staticmethod
     @abstractmethod
-    def rerank_enabled(cls) -> bool: ...
+    def rerank_enabled() -> bool: ...
 
     @classmethod
     @abstractmethod
-    def list_models(cls) -> List[str]: ...
+    def list_models(cls) -> list[str]: ...
 
-    @classmethod
+    @staticmethod
     @abstractmethod
-    def is_available(cls) -> bool: ...
+    def is_available() -> bool: ...
 
     @classmethod
     def is_community(cls) -> bool:
         return False
 
     @classmethod
-    def config(cls) -> Dict[str, Any]:
+    def config(cls) -> dict[str, Any]:
         config = Settings().get(f"deployments.{cls.id()}")
 
         if not config:
@@ -81,9 +81,9 @@ async def invoke_chat(
     @abstractmethod
     async def invoke_chat_stream(
         self, chat_request: CohereChatRequest, ctx: Context, **kwargs: Any
-    ) -> AsyncGenerator[Any, Any]: ...
+    ) -> Any: ...
 
     @abstractmethod
     async def invoke_rerank(
-        self, query: str, documents: List[Dict[str, Any]], ctx: Context, **kwargs: Any
+        self, query: str, documents: list[str], ctx: Context, **kwargs: Any
     ) -> Any: ...

src/backend/model_deployments/bedrock.py

@@ -1,4 +1,4 @@
-from typing import Any, AsyncGenerator, Dict, List
+from typing import Any, AsyncGenerator
 
 import cohere
 
@@ -42,40 +42,40 @@ def __init__(self, **kwargs: Any):
             ),
         )
 
-    @classmethod
-    def name(cls) -> str:
+    @staticmethod
+    def name() -> str:
         return "Bedrock"
 
-    @classmethod
-    def env_vars(cls) -> List[str]:
+    @staticmethod
+    def env_vars() -> list[str]:
         return [
             BEDROCK_ACCESS_KEY_ENV_VAR,
             BEDROCK_SECRET_KEY_ENV_VAR,
             BEDROCK_SESSION_TOKEN_ENV_VAR,
             BEDROCK_REGION_NAME_ENV_VAR,
         ]
 
-    @classmethod
-    def rerank_enabled(cls) -> bool:
+    @staticmethod
+    def rerank_enabled() -> bool:
         return False
 
     @classmethod
-    def list_models(cls) -> List[str]:
+    def list_models(cls) -> list[str]:
         if not cls.is_available():
             return []
 
         return cls.DEFAULT_MODELS
 
-    @classmethod
-    def is_available(cls) -> bool:
+    @staticmethod
+    def is_available() -> bool:
         return (
             BedrockDeployment.access_key is not None
             and BedrockDeployment.secret_access_key is not None
             and BedrockDeployment.session_token is not None
             and BedrockDeployment.region_name is not None
         )
 
-    async def invoke_chat(self, chat_request: CohereChatRequest) -> Any:
+    async def invoke_chat(self, chat_request: CohereChatRequest, **kwargs: Any) -> Any:
         # bedrock accepts a subset of the chat request fields
         bedrock_chat_req = chat_request.model_dump(
             exclude={"tools", "conversation_id", "model", "stream"}, exclude_none=True
@@ -101,6 +101,6 @@ async def invoke_chat_stream(
             yield to_dict(event)
 
     async def invoke_rerank(
-        self, query: str, documents: List[Dict[str, Any]], ctx: Context
+        self, query: str, documents: list[str], ctx: Context, **kwargs: Any
     ) -> Any:
         return None

src/backend/model_deployments/cohere_platform.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any
 
 import cohere
 import requests
@@ -29,20 +29,20 @@ def __init__(self, **kwargs: Any):
         )
         self.client = cohere.Client(api_key, client_name=self.client_name)
 
-    @classmethod
-    def name(cls) -> str:
+    @staticmethod
+    def name() -> str:
         return "Cohere Platform"
 
-    @classmethod
-    def env_vars(cls) -> List[str]:
+    @staticmethod
+    def env_vars() -> list[str]:
         return [COHERE_API_KEY_ENV_VAR]
 
-    @classmethod
-    def rerank_enabled(cls) -> bool:
+    @staticmethod
+    def rerank_enabled() -> bool:
         return True
 
     @classmethod
-    def list_models(cls) -> List[str]:
+    def list_models(cls) -> list[str]:
         logger = LoggerFactory().get_logger()
         if not CohereDeployment.is_available():
             return []
@@ -64,12 +64,12 @@ def list_models(cls) -> List[str]:
         models = response.json()["models"]
         return [model["name"] for model in models if model.get("endpoints") and "chat" in model["endpoints"]]
 
-    @classmethod
-    def is_available(cls) -> bool:
+    @staticmethod
+    def is_available() -> bool:
         return CohereDeployment.api_key is not None
 
     async def invoke_chat(
-        self, chat_request: CohereChatRequest, ctx: Context, **kwargs: Any
+        self, chat_request: CohereChatRequest, **kwargs: Any
     ) -> Any:
         response = self.client.chat(
             **chat_request.model_dump(exclude={"stream", "file_ids", "agent_id"}),
@@ -99,7 +99,7 @@ async def invoke_chat_stream(
             yield event_dict
 
     async def invoke_rerank(
-        self, query: str, documents: List[Dict[str, Any]], ctx: Context, **kwargs: Any
+        self, query: str, documents: list[str], ctx: Context, **kwargs: Any
     ) -> Any:
         response = self.client.rerank(
             query=query, documents=documents, model=DEFAULT_RERANK_MODEL