From 179b5da06b39a14863ca1e499f79e91e82cd25a5 Mon Sep 17 00:00:00 2001
From: Liang Lv <liang1.lv@intel.com>
Date: Thu, 9 Jan 2025 11:26:03 +0800
Subject: [PATCH] Refactor prompt registry microservice  (#1124)

Signed-off-by: lvliang-intel <liang1.lv@intel.com>
---
 .../docker/compose/prompt_registry-compose.yaml      |  6 +++---
 comps/prompt_registry/README.md                      |  2 +-
 .../compose_prompt_registry_mongo.yaml}              |  2 +-
 .../prompt_registry/deployment/kubernetes/README.md  |  0
 comps/prompt_registry/{mongo => src}/Dockerfile      |  6 +++---
 comps/prompt_registry/{mongo => src}/README.md       |  4 ++--
 comps/prompt_registry/src/__init__.py                |  2 ++
 comps/prompt_registry/src/integrations/__init__.py   |  2 ++
 .../src/integrations/mongo/__init__.py               |  2 ++
 .../{ => src/integrations}/mongo/config.py           |  0
 .../{ => src/integrations}/mongo/mongo_conn.py       |  3 ++-
 .../prompt.py => src/opea_prompt_microservice.py}    | 12 ++++++------
 .../{mongo/mongo_store.py => src/prompt_store.py}    |  9 +++++----
 .../prompt_registry/{mongo => src}/requirements.txt  |  0
 tests/prompt_registry/test_prompt_registry_mongo.sh  | 10 +++++-----
 15 files changed, 34 insertions(+), 26 deletions(-)
 rename comps/prompt_registry/{mongo/docker-compose-prompt-registry-mongo.yaml => deployment/docker_compose/compose_prompt_registry_mongo.yaml} (94%)
 create mode 100644 comps/prompt_registry/deployment/kubernetes/README.md
 rename comps/prompt_registry/{mongo => src}/Dockerfile (83%)
 rename comps/prompt_registry/{mongo => src}/README.md (81%)
 create mode 100644 comps/prompt_registry/src/__init__.py
 create mode 100644 comps/prompt_registry/src/integrations/__init__.py
 create mode 100644 comps/prompt_registry/src/integrations/mongo/__init__.py
 rename comps/prompt_registry/{ => src/integrations}/mongo/config.py (100%)
 rename comps/prompt_registry/{ => src/integrations}/mongo/mongo_conn.py (90%)
 rename comps/prompt_registry/{mongo/prompt.py => src/opea_prompt_microservice.py} (93%)
 rename comps/prompt_registry/{mongo/mongo_store.py => src/prompt_store.py} (94%)
 rename comps/prompt_registry/{mongo => src}/requirements.txt (100%)

diff --git a/.github/workflows/docker/compose/prompt_registry-compose.yaml b/.github/workflows/docker/compose/prompt_registry-compose.yaml
index 34d8973df5..f50daa42c4 100644
--- a/.github/workflows/docker/compose/prompt_registry-compose.yaml
+++ b/.github/workflows/docker/compose/prompt_registry-compose.yaml
@@ -3,7 +3,7 @@
 
 # this file should be run in the root of the repo
 services:
-  promptregistry-mongo-server:
+  promptregistry-server:
     build:
-      dockerfile: comps/prompt_registry/mongo/Dockerfile
-    image: ${REGISTRY:-opea}/promptregistry-mongo-server:${TAG:-latest}
+      dockerfile: comps/prompt_registry/src/Dockerfile
+    image: ${REGISTRY:-opea}/promptregistry-server:${TAG:-latest}
diff --git a/comps/prompt_registry/README.md b/comps/prompt_registry/README.md
index 6332a1a13a..ec29ad7417 100644
--- a/comps/prompt_registry/README.md
+++ b/comps/prompt_registry/README.md
@@ -18,4 +18,4 @@ The Prompt Registry microservice able to support various database backends for s
 
 ### Prompt Registry with MongoDB
 
-For more detail, please refer to this [README](./mongo/README.md)
+For more detail, please refer to this [README](./src/README.md)
diff --git a/comps/prompt_registry/mongo/docker-compose-prompt-registry-mongo.yaml b/comps/prompt_registry/deployment/docker_compose/compose_prompt_registry_mongo.yaml
similarity index 94%
rename from comps/prompt_registry/mongo/docker-compose-prompt-registry-mongo.yaml
rename to comps/prompt_registry/deployment/docker_compose/compose_prompt_registry_mongo.yaml
index 466d06093c..e956320bdc 100644
--- a/comps/prompt_registry/mongo/docker-compose-prompt-registry-mongo.yaml
+++ b/comps/prompt_registry/deployment/docker_compose/compose_prompt_registry_mongo.yaml
@@ -16,7 +16,7 @@ services:
 
   promptregistry-mongo:
     image: opea/promptregistry-mongo:latest
-    container_name: promptregistry-mongo-server
+    container_name: promptregistry-server
     ports:
       - "6018:6018"
     ipc: host
diff --git a/comps/prompt_registry/deployment/kubernetes/README.md b/comps/prompt_registry/deployment/kubernetes/README.md
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/comps/prompt_registry/mongo/Dockerfile b/comps/prompt_registry/src/Dockerfile
similarity index 83%
rename from comps/prompt_registry/mongo/Dockerfile
rename to comps/prompt_registry/src/Dockerfile
index 25e14cd33e..b00a8d56c7 100644
--- a/comps/prompt_registry/mongo/Dockerfile
+++ b/comps/prompt_registry/src/Dockerfile
@@ -20,11 +20,11 @@ COPY comps /home/user/comps
 COPY requirements.txt /home/user/
 
 RUN pip install --no-cache-dir --upgrade pip setuptools && \
-    pip install --no-cache-dir -r /home/user/comps/prompt_registry/mongo/requirements.txt && \
+    pip install --no-cache-dir -r /home/user/comps/prompt_registry/src/requirements.txt && \
     pip install --no-cache-dir -r /home/user/requirements.txt
 
 ENV PYTHONPATH=$PYTHONPATH:/home/user
 
-WORKDIR /home/user/comps/prompt_registry/mongo
+WORKDIR /home/user/comps/prompt_registry/src/
 
-ENTRYPOINT ["python", "prompt.py"]
+ENTRYPOINT ["python", "opea_prompt_microservice.py"]
diff --git a/comps/prompt_registry/mongo/README.md b/comps/prompt_registry/src/README.md
similarity index 81%
rename from comps/prompt_registry/mongo/README.md
rename to comps/prompt_registry/src/README.md
index 8db99e88a5..8142e23bd0 100644
--- a/comps/prompt_registry/mongo/README.md
+++ b/comps/prompt_registry/src/README.md
@@ -23,7 +23,7 @@ export COLLECTION_NAME=${COLLECTION_NAME}
 
 ```bash
 cd ~/GenAIComps
-docker build -t opea/promptregistry-mongo-server:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/prompt_registry/mongo/Dockerfile .
+docker build -t opea/promptregistry-server:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/prompt_registry/src/Dockerfile .
 ```
 
 ### Run Docker with CLI
@@ -37,7 +37,7 @@ docker build -t opea/promptregistry-mongo-server:latest --build-arg https_proxy=
 - Run Prompt Registry microservice
 
   ```bash
-  docker run -d --name="promptregistry-mongo-server" -p 6018:6018 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e no_proxy=$no_proxy -e MONGO_HOST=${MONGO_HOST} -e MONGO_PORT=${MONGO_PORT} -e DB_NAME=${DB_NAME} -e COLLECTION_NAME=${COLLECTION_NAME} opea/promptregistry-mongo-server:latest
+  docker run -d --name="promptregistry-server" -p 6018:6018 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e no_proxy=$no_proxy -e MONGO_HOST=${MONGO_HOST} -e MONGO_PORT=${MONGO_PORT} -e DB_NAME=${DB_NAME} -e COLLECTION_NAME=${COLLECTION_NAME} opea/promptregistry-server:latest
   ```
 
 ---
diff --git a/comps/prompt_registry/src/__init__.py b/comps/prompt_registry/src/__init__.py
new file mode 100644
index 0000000000..916f3a44b2
--- /dev/null
+++ b/comps/prompt_registry/src/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
diff --git a/comps/prompt_registry/src/integrations/__init__.py b/comps/prompt_registry/src/integrations/__init__.py
new file mode 100644
index 0000000000..916f3a44b2
--- /dev/null
+++ b/comps/prompt_registry/src/integrations/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
diff --git a/comps/prompt_registry/src/integrations/mongo/__init__.py b/comps/prompt_registry/src/integrations/mongo/__init__.py
new file mode 100644
index 0000000000..916f3a44b2
--- /dev/null
+++ b/comps/prompt_registry/src/integrations/mongo/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
diff --git a/comps/prompt_registry/mongo/config.py b/comps/prompt_registry/src/integrations/mongo/config.py
similarity index 100%
rename from comps/prompt_registry/mongo/config.py
rename to comps/prompt_registry/src/integrations/mongo/config.py
diff --git a/comps/prompt_registry/mongo/mongo_conn.py b/comps/prompt_registry/src/integrations/mongo/mongo_conn.py
similarity index 90%
rename from comps/prompt_registry/mongo/mongo_conn.py
rename to comps/prompt_registry/src/integrations/mongo/mongo_conn.py
index 52e4c8e130..48cd19a1af 100644
--- a/comps/prompt_registry/mongo/mongo_conn.py
+++ b/comps/prompt_registry/src/integrations/mongo/mongo_conn.py
@@ -4,7 +4,8 @@
 from typing import Any
 
 import motor.motor_asyncio as motor
-from config import DB_NAME, MONGO_HOST, MONGO_PORT
+
+from .config import DB_NAME, MONGO_HOST, MONGO_PORT
 
 
 class MongoClient:
diff --git a/comps/prompt_registry/mongo/prompt.py b/comps/prompt_registry/src/opea_prompt_microservice.py
similarity index 93%
rename from comps/prompt_registry/mongo/prompt.py
rename to comps/prompt_registry/src/opea_prompt_microservice.py
index 81047da1b0..036a6f66c5 100644
--- a/comps/prompt_registry/mongo/prompt.py
+++ b/comps/prompt_registry/src/opea_prompt_microservice.py
@@ -3,13 +3,13 @@
 import os
 from typing import Optional
 
-from mongo_store import PromptStore
+from prompt_store import PromptStore
 from pydantic import BaseModel
 
 from comps import CustomLogger
 from comps.cores.mega.micro_service import opea_microservices, register_microservice
 
-logger = CustomLogger("prompt_mongo")
+logger = CustomLogger("prompt_registry")
 logflag = os.getenv("LOGFLAG", False)
 
 
@@ -39,7 +39,7 @@ class PromptId(BaseModel):
 
 
 @register_microservice(
-    name="opea_service@prompt_mongo",
+    name="opea_service@prompt",
     endpoint="/v1/prompt/create",
     host="0.0.0.0",
     input_datatype=PromptCreate,
@@ -70,7 +70,7 @@ async def create_prompt(prompt: PromptCreate):
 
 
 @register_microservice(
-    name="opea_service@prompt_mongo",
+    name="opea_service@prompt",
     endpoint="/v1/prompt/get",
     host="0.0.0.0",
     input_datatype=PromptId,
@@ -106,7 +106,7 @@ async def get_prompt(prompt: PromptId):
 
 
 @register_microservice(
-    name="opea_service@prompt_mongo",
+    name="opea_service@prompt",
     endpoint="/v1/prompt/delete",
     host="0.0.0.0",
     input_datatype=PromptId,
@@ -140,4 +140,4 @@ async def delete_prompt(prompt: PromptId):
 
 
 if __name__ == "__main__":
-    opea_microservices["opea_service@prompt_mongo"].start()
+    opea_microservices["opea_service@prompt"].start()
diff --git a/comps/prompt_registry/mongo/mongo_store.py b/comps/prompt_registry/src/prompt_store.py
similarity index 94%
rename from comps/prompt_registry/mongo/mongo_store.py
rename to comps/prompt_registry/src/prompt_store.py
index 04b2b86694..7f363d9a68 100644
--- a/comps/prompt_registry/mongo/mongo_store.py
+++ b/comps/prompt_registry/src/prompt_store.py
@@ -4,8 +4,8 @@
 
 import bson.errors as BsonError
 from bson.objectid import ObjectId
-from config import COLLECTION_NAME
-from mongo_conn import MongoClient
+from integrations.mongo.config import COLLECTION_NAME
+from integrations.mongo.mongo_conn import MongoClient
 
 
 class PromptStore:
@@ -16,8 +16,9 @@ def __init__(
     ):
         self.user = user
 
-    def initialize_storage(self) -> None:
-        self.db_client = MongoClient.get_db_client()
+    def initialize_storage(self, db_type="mongo") -> None:
+        if db_type == "mongo":
+            self.db_client = MongoClient.get_db_client()
         self.collection = self.db_client[COLLECTION_NAME]
 
     async def save_prompt(self, prompt) -> str:
diff --git a/comps/prompt_registry/mongo/requirements.txt b/comps/prompt_registry/src/requirements.txt
similarity index 100%
rename from comps/prompt_registry/mongo/requirements.txt
rename to comps/prompt_registry/src/requirements.txt
diff --git a/tests/prompt_registry/test_prompt_registry_mongo.sh b/tests/prompt_registry/test_prompt_registry_mongo.sh
index 661a13ddac..47b84aa6af 100644
--- a/tests/prompt_registry/test_prompt_registry_mongo.sh
+++ b/tests/prompt_registry/test_prompt_registry_mongo.sh
@@ -17,18 +17,18 @@ function build_docker_images() {
     echo $(pwd)
     docker run -d -p 27017:27017 --name=test-comps-mongo mongo:latest
 
-    docker build --no-cache -t opea/promptregistry-mongo-server:comps --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/prompt_registry/mongo/Dockerfile .
+    docker build --no-cache -t opea/promptregistry-server:comps --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/prompt_registry/src/Dockerfile .
     if [ $? -ne 0 ]; then
-        echo "opea/promptregistry-mongo-server built fail"
+        echo "opea/promptregistry-server built fail"
         exit 1
     else
-        echo "opea/promptregistry-mongo-server built successful"
+        echo "opea/promptregistry-server built successful"
     fi
 }
 
 function start_service() {
 
-    docker run -d --name="test-comps-promptregistry-mongo-server" -p 6018:6018 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e no_proxy=$no_proxy -e MONGO_HOST=${MONGO_HOST} -e MONGO_PORT=${MONGO_PORT} -e DB_NAME=${DB_NAME} -e COLLECTION_NAME=${COLLECTION_NAME} opea/promptregistry-mongo-server:comps
+    docker run -d --name="test-comps-promptregistry-server" -p 6018:6018 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e no_proxy=$no_proxy -e MONGO_HOST=${MONGO_HOST} -e MONGO_PORT=${MONGO_PORT} -e DB_NAME=${DB_NAME} -e COLLECTION_NAME=${COLLECTION_NAME} opea/promptregistry-server:comps
 
     sleep 10s
 }
@@ -46,7 +46,7 @@ function validate_microservice() {
         echo "Correct result."
     else
         echo "Incorrect result."
-        docker logs test-comps-promptregistry-mongo-server
+        docker logs test-comps-promptregistry-server
         exit 1
     fi