From e5b72c1280b63266bb1b40d810faecf85bee1e27 Mon Sep 17 00:00:00 2001
From: Gaurav Kumbhat <kumbhat.gaurav@gmail.com>
Date: Wed, 19 Jun 2024 12:00:18 -0500
Subject: [PATCH] :bug: Fix model id configuration bug for nlp client and
 rename tgis-router to common-router (#89)

* :bug: Fix model id configuration bug for nlp client and rename tgis-router to common-router

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

* :art: Fix formatting

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

* :art::truck: Move common router key as const and re-use that

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

---------

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>
---
 src/clients.rs      |  1 +
 src/clients/nlp.rs  |  5 ++++-
 src/clients/tgis.rs |  5 +++--
 src/orchestrator.rs | 12 +++++++++---
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/clients.rs b/src/clients.rs
index d687ba46..74482f50 100644
--- a/src/clients.rs
+++ b/src/clients.rs
@@ -28,6 +28,7 @@ pub const DEFAULT_TGIS_PORT: u16 = 8033;
 pub const DEFAULT_CAIKIT_NLP_PORT: u16 = 8085;
 pub const DEFAULT_CHUNKER_PORT: u16 = 8085;
 pub const DEFAULT_DETECTOR_PORT: u16 = 8080;
+pub const COMMON_ROUTER_KEY: &str = "common-router";
 const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
 
diff --git a/src/clients/nlp.rs b/src/clients/nlp.rs
index bfc3499f..67f219ad 100644
--- a/src/clients/nlp.rs
+++ b/src/clients/nlp.rs
@@ -8,6 +8,7 @@ use tonic::Request;
 
 use super::{create_grpc_clients, Error};
 use crate::{
+    clients::COMMON_ROUTER_KEY,
     config::ServiceConfig,
     pb::{
         caikit::runtime::nlp::{
@@ -34,7 +35,9 @@ impl NlpClient {
         Self { clients }
     }
 
-    fn client(&self, model_id: &str) -> Result<NlpServiceClient<LoadBalancedChannel>, Error> {
+    fn client(&self, _model_id: &str) -> Result<NlpServiceClient<LoadBalancedChannel>, Error> {
+        // NOTE: We currently forward requests to common router, so we use a single client.
+        let model_id = COMMON_ROUTER_KEY;
         Ok(self
             .clients
             .get(model_id)
diff --git a/src/clients/tgis.rs b/src/clients/tgis.rs
index 33e6a67d..816337f7 100644
--- a/src/clients/tgis.rs
+++ b/src/clients/tgis.rs
@@ -7,6 +7,7 @@ use tokio_stream::wrappers::ReceiverStream;
 
 use super::{create_grpc_clients, Error};
 use crate::{
+    clients::COMMON_ROUTER_KEY,
     config::ServiceConfig,
     pb::fmaas::{
         generation_service_client::GenerationServiceClient, BatchedGenerationRequest,
@@ -32,8 +33,8 @@ impl TgisClient {
         &self,
         _model_id: &str,
     ) -> Result<GenerationServiceClient<LoadBalancedChannel>, Error> {
-        // NOTE: We currently forward requests to the tgis-router, so we use a single client.
-        let model_id = "tgis-router";
+        // NOTE: We currently forward requests to the common-router, so we use a single client.
+        let model_id = COMMON_ROUTER_KEY;
         Ok(self
             .clients
             .get(model_id)
diff --git a/src/orchestrator.rs b/src/orchestrator.rs
index 8f3b543c..131d1e0c 100644
--- a/src/orchestrator.rs
+++ b/src/orchestrator.rs
@@ -13,7 +13,7 @@ use uuid::Uuid;
 use crate::{
     clients::{
         self, detector::ContentAnalysisRequest, ChunkerClient, DetectorClient, GenerationClient,
-        NlpClient, TgisClient,
+        NlpClient, TgisClient, COMMON_ROUTER_KEY,
     },
     config::{GenerationProvider, OrchestratorConfig},
     models::{
@@ -601,7 +601,10 @@ async fn create_clients(
         GenerationProvider::Tgis => {
             let client = TgisClient::new(
                 clients::DEFAULT_TGIS_PORT,
-                &[("tgis-router".to_string(), config.generation.service.clone())],
+                &[(
+                    COMMON_ROUTER_KEY.to_string(),
+                    config.generation.service.clone(),
+                )],
             )
             .await;
             GenerationClient::Tgis(client)
@@ -609,7 +612,10 @@ async fn create_clients(
         GenerationProvider::Nlp => {
             let client = NlpClient::new(
                 clients::DEFAULT_CAIKIT_NLP_PORT,
-                &[("tgis-router".to_string(), config.generation.service.clone())],
+                &[(
+                    COMMON_ROUTER_KEY.to_string(),
+                    config.generation.service.clone(),
+                )],
             )
             .await;
             GenerationClient::Nlp(client)