Addition of predictionguard rerank genai component.

comps/reranks/README.md

@@ -22,6 +22,10 @@ However, a reranking model can further refine this process by rearranging potent
 
 For additional information, please refer to this [README](./fastrag/README.md)
 
+### Utilizing Reranking with Prediction Guard
+
+For additional information, please refer to this [README](./predictionguard/README.md)
+
 ### Utilizing Reranking with Mosec
 
 For additional information, please refer to this [README](./mosec/langchain/README.md)

comps/reranks/predictionguard/Dockerfile

@@ -0,0 +1,15 @@
+# Copyright (C) 2024 Prediction Guard, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+FROM python:3.11-slim
+
+COPY comps /home/comps
+
+RUN pip install --no-cache-dir --upgrade pip setuptools && \
+    pip install --no-cache-dir -r /home/comps/reranks/predictionguard/requirements.txt
+
+ENV PYTHONPATH=$PYTHONPATH:/home
+
+WORKDIR /home/comps/reranks/predictionguard
+
+ENTRYPOINT ["bash", "entrypoint.sh"]

comps/reranks/predictionguard/README.md

@@ -0,0 +1,31 @@
+# Prediction Guard Introduction
+
+[Prediction Guard](https://docs.predictionguard.com) allows you to utilize hosted open access LLMs, LVMs, and embedding functionality with seamlessly integrated safeguards. In addition to providing a scalable access to open models, Prediction Guard allows you to configure factual consistency checks, toxicity filters, PII filters, and prompt injection blocking. Join the [Prediction Guard Discord channel](https://discord.gg/TFHgnhAFKd) and request an API key to get started.
+
+## Get Started
+
+### Build Docker Image
+
+```bash
+cd ../../..
+docker build -t opea/reranking-predictionguard:latest -f comps/reranks/predictionguard/Dockerfile .
+```
+
+### Run the Predictionguard Microservice
+
+```bash
+docker run -d -p 9000:9000 -e PREDICTIONGUARD_API_KEY=$PREDICTIONGUARD_API_KEY  --name reranking-predictionguard opea/reranking-predictionguard:latest
+```
+
+## Consume the Prediction Guard Rerank Microservice
+
+See the [Prediction Guard docs](https://docs.predictionguard.com/options/reranker_models) for available model options.
+
+```bash
+curl -N -X POST http://localhost:9000/v1/reranking \
+    -H "Content-Type: application/json" \
+    -d '{
+        "initial_query": "What is Deep Learning?",
+        "retrieved_docs": [{"text":"Deep Learning is not..."}, {"text":"Deep learning is..."}]
+        }'
+```

comps/reranks/predictionguard/docker_compose.yaml

@@ -0,0 +1,20 @@
+# Copyright (C) 2024 Prediction Guard, Inc
+# SPDX-License-Identifier: Apache-2.0
+
+services:
+  reranking:
+    image: opea/reranks-predictionguard:latest
+    container_name: reranks-predictionguard
+    ports:
+      - "9000:9000"
+    ipc: host
+    environment:
+      no_proxy: ${no_proxy}
+      http_proxy: ${http_proxy}
+      https_proxy: ${https_proxy}
+      PREDICTIONGUARD_API_KEY: ${PREDICTIONGUARD_API_KEY}
+    restart: unless-stopped
+
+networks:
+  default:
+    driver: bridge

comps/reranks/predictionguard/entrypoint.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+# Copyright (C) 2024 Prediction Guard, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+#pip --no-cache-dir install -r requirements-runtime.txt
+
+python src/reranks_predictionguard.py

comps/reranks/predictionguard/requirements.txt

@@ -0,0 +1,12 @@
+aiohttp
+docarray
+fastapi
+opentelemetry-api
+opentelemetry-exporter-otlp
+opentelemetry-sdk
+Pillow
+predictionguard
+prometheus-fastapi-instrumentator
+shortuuid
+transformers
+uvicorn

comps/reranks/predictionguard/src/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Prediction Guard, Inc.
+# SPDX-License-Identifier: Apache-2.0

comps/reranks/predictionguard/src/helpers.py

@@ -0,0 +1,18 @@
+# Copyright (C) 2024 Prediction Guard, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List
+
+from comps.cores.proto.docarray import DocList
+
+
+def process_doc_list(docs: DocList) -> List[str]:
+    result = []
+    for doc in docs:
+        if doc.text is None:
+            continue
+        if isinstance(doc.text, list):
+            result.extend(doc.text)
+        else:
+            result.append(doc.text)
+    return result

comps/reranks/predictionguard/src/reranks_predictionguard.py

@@ -0,0 +1,66 @@
+# Copyright (C) 2024 Prediction Guard, Inc.
+# SPDX-License-Identified: Apache-2.0
+import logging
+import time
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from predictionguard import PredictionGuard
+
+from comps import (
+    GeneratedDoc,
+    LLMParamsDoc,
+    RerankedDoc,
+    SearchedDoc,
+    ServiceType,
+    TextDoc,
+    opea_microservices,
+    register_microservice,
+    register_statistics,
+    statistics_dict,
+)
+from comps.reranks.predictionguard.src.helpers import process_doc_list
+
+client = PredictionGuard()
+app = FastAPI()
+
+
+@register_microservice(
+    name="opea_service@reranks_predictionguard",
+    service_type=ServiceType.LLM,
+    endpoint="/v1/reranking",
+    host="0.0.0.0",
+    port=9000,
+    input_datatype=SearchedDoc,
+    output_datatype=RerankedDoc,
+)
+@register_statistics(names=["opea_service@reranks_predictionguard"])
+def reranks_generate(input: SearchedDoc) -> RerankedDoc:
+    start = time.time()
+    reranked_docs = []
+
+    if input.retrieved_docs:
+        docs = process_doc_list(input.retrieved_docs)
+
+        try:
+            rerank_result = client.rerank.create(
+                model="bge-reranker-v2-m3", query=input.initial_query, documents=docs, return_documents=True
+            )
+
+            # based on rerank_result, reorder the retrieved_docs to match the order of the retrieved_docs in the input
+            reranked_docs = [
+                TextDoc(id=input.retrieved_docs[doc["index"]].id, text=doc["text"]) for doc in rerank_result["results"]
+            ]
+
+        except ValueError as e:
+            logging.error(f"rerank failed with error: {e}. Inputs: query={input.initial_query}, documents={docs}")
+            raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+    else:
+        logging.info("reranking request input did not contain any documents")
+
+    statistics_dict["opea_service@reranks_predictionguard"].append_latency(time.time() - start, None)
+    return RerankedDoc(initial_query=input.initial_query, reranked_docs=reranked_docs)
+
+
+if __name__ == "__main__":
+    opea_microservices["opea_service@reranks_predictionguard"].start()

comps/reranks/predictionguard/src/test_helpers.py

@@ -0,0 +1,53 @@
+# Copyright (C) 2024 Prediction Guard, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+from comps.cores.proto.docarray import DocList, TextDoc
+from comps.reranks.predictionguard.src.helpers import process_doc_list
+
+
+class MyTestCase(unittest.TestCase):
+    def setUp(self):
+        # Test cases setup
+        self.simple_case = DocList(
+            [TextDoc(text="Hello world"), TextDoc(text="Another document"), TextDoc(text="Third document")]
+        )
+
+        self.list_case = DocList(
+            [TextDoc(text=["First sentence.", "Second sentence."]), TextDoc(text=["Another paragraph.", "More text."])]
+        )
+
+        self.mixed_case = DocList(
+            [
+                TextDoc(text="Single string document"),
+                TextDoc(text=["First part", "Second part"]),
+                TextDoc(text="Another single string"),
+            ]
+        )
+
+        self.none_case = DocList([TextDoc(text="Valid text"), TextDoc(text=[]), TextDoc(text="More text")])
+
+    def test_simple_case(self):
+        result = process_doc_list(self.simple_case)
+        expected = ["Hello world", "Another document", "Third document"]
+        self.assertEqual(result, expected)
+
+    def test_list_case(self):
+        result = process_doc_list(self.list_case)
+        expected = ["First sentence.", "Second sentence.", "Another paragraph.", "More text."]
+        self.assertEqual(result, expected)
+
+    def test_mixed_case(self):
+        result = process_doc_list(self.mixed_case)
+        expected = ["Single string document", "First part", "Second part", "Another single string"]
+        self.assertEqual(result, expected)
+
+    def test_none_case(self):
+        result = process_doc_list(self.none_case)
+        expected = ["Valid text", "More text"]
+        self.assertEqual(result, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/reranks/test_reranks_predictionguard.sh

@@ -0,0 +1,67 @@
+#!/bin/bash
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+set -x
+
+WORKPATH=$(dirname "$PWD")
+ip_address=$(hostname -I | awk '{print $1}')
+function build_docker_images() {
+    cd $WORKPATH
+    docker build --no-cache --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -t opea/reranking-predictionguard:comps -f comps/reranks/predictionguard/Dockerfile .
+    if [ $? -ne 0 ]; then
+        echo "opea/reranking-predictionguard built fail"
+        exit 1
+    else
+        echo "opea/reranking-predictionguard built successful"
+    fi
+}
+
+function start_service() {
+    predictionguard_service_port=9000
+    unset http_proxy
+
+    docker run -d --name="test-comps-reranking-predictionguard-server" \
+    -p ${predictionguard_service_port}:9000 \
+    --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy \
+    -e PREDICTIONGUARD_API_KEY=${PREDICTIONGUARD_API_KEY} \
+    opea/reranking-predictionguard:comps
+
+    sleep 1m
+}
+
+function validate_microservice() {
+    predictionguard_service_port=9000
+    result=$(http_proxy="" curl http://${ip_address}:${predictionguard_service_port}/v1/reranking\
+        -X POST \
+        -d '{"initial_query":"What is Deep Learning?", "retrieved_docs": [{"text":"Deep Learning is not..."}, {"text":"Deep learning is..."}]}' \
+        -H 'Content-Type: application/json')
+    if [[ $result == *"reranked_docs"* ]]; then
+        echo "Result correct."
+    else
+        echo "Result wrong. Received was $result"
+        docker logs test-comps-reranking-predictionguard-server
+        exit 1
+    fi
+}
+
+function stop_docker() {
+    cid=$(docker ps -aq --filter "name=test-comps-rerank*")
+    if [[ ! -z "$cid" ]]; then docker stop $cid && docker rm $cid && sleep 1s; fi
+}
+
+function main() {
+
+    stop_docker
+
+    build_docker_images
+    start_service
+
+    validate_microservice
+
+    stop_docker
+    echo y | docker system prune
+
+}
+
+main