From a1e9e7dc73f14fb71ae783fe7abfe8d8711dc623 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=ADsa=20Moura?= Date: Thu, 11 Jul 2024 19:52:53 +0100 Subject: [PATCH] Metrics: test deployment without metrics (#400) --- src/backend/main.py | 2 +- src/backend/model_deployments/cohere_platform.py | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/backend/main.py b/src/backend/main.py index b969626cbc..b190d1cc47 100644 --- a/src/backend/main.py +++ b/src/backend/main.py @@ -74,7 +74,7 @@ def create_app(): allow_headers=["*"], ) app.add_middleware(LoggingMiddleware) - app.add_middleware(MetricsMiddleware) + # app.add_middleware(MetricsMiddleware) return app diff --git a/src/backend/model_deployments/cohere_platform.py b/src/backend/model_deployments/cohere_platform.py index 15a0c5e2f9..0195959e10 100644 --- a/src/backend/model_deployments/cohere_platform.py +++ b/src/backend/model_deployments/cohere_platform.py @@ -69,14 +69,14 @@ def list_models(cls) -> List[str]: def is_available(cls) -> bool: return all([os.environ.get(var) is not None for var in COHERE_ENV_VARS]) - @collect_metrics_chat + # @collect_metrics_chat async def invoke_chat(self, chat_request: CohereChatRequest, **kwargs: Any) -> Any: response = self.client.chat( **chat_request.model_dump(exclude={"stream", "file_ids", "agent_id"}), ) yield to_dict(response) - @collect_metrics_chat_stream + # @collect_metrics_chat_stream async def invoke_chat_stream( self, chat_request: CohereChatRequest, **kwargs: Any ) -> AsyncGenerator[Any, Any]: @@ -87,10 +87,12 @@ async def invoke_chat_stream( for event in stream: yield to_dict(event) - @collect_metrics_rerank + # @collect_metrics_rerank async def invoke_rerank( self, query: str, documents: List[Dict[str, Any]], **kwargs: Any ) -> Any: - return self.client.rerank( + response = self.client.rerank( query=query, documents=documents, model=DEFAULT_RERANK_MODEL ) + + return to_dict(response)