diff --git a/comps/reranks/README.md b/comps/reranks/README.md index f8176720bb..e06d1298cf 100644 --- a/comps/reranks/README.md +++ b/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) diff --git a/comps/reranks/predictionguard/Dockerfile b/comps/reranks/predictionguard/Dockerfile new file mode 100644 index 0000000000..f2761a475a --- /dev/null +++ b/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"] diff --git a/comps/reranks/predictionguard/README.md b/comps/reranks/predictionguard/README.md new file mode 100644 index 0000000000..6504792ea1 --- /dev/null +++ b/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..."}] + }' +``` diff --git a/comps/reranks/predictionguard/docker_compose.yaml b/comps/reranks/predictionguard/docker_compose.yaml new file mode 100644 index 0000000000..ecdb5d3b27 --- /dev/null +++ b/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 diff --git a/comps/reranks/predictionguard/entrypoint.sh b/comps/reranks/predictionguard/entrypoint.sh new file mode 100644 index 0000000000..9cbe393a39 --- /dev/null +++ b/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 diff --git a/comps/reranks/predictionguard/requirements.txt b/comps/reranks/predictionguard/requirements.txt new file mode 100644 index 0000000000..6c9f8340fd --- /dev/null +++ b/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 diff --git a/comps/reranks/predictionguard/src/__init__.py b/comps/reranks/predictionguard/src/__init__.py new file mode 100644 index 0000000000..a246c95e79 --- /dev/null +++ b/comps/reranks/predictionguard/src/__init__.py @@ -0,0 +1,2 @@ +# Copyright (C) 2024 Prediction Guard, Inc. +# SPDX-License-Identifier: Apache-2.0 diff --git a/comps/reranks/predictionguard/src/helpers.py b/comps/reranks/predictionguard/src/helpers.py new file mode 100644 index 0000000000..f948a99236 --- /dev/null +++ b/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 diff --git a/comps/reranks/predictionguard/src/reranks_predictionguard.py b/comps/reranks/predictionguard/src/reranks_predictionguard.py new file mode 100644 index 0000000000..8fd870e4cb --- /dev/null +++ b/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() diff --git a/comps/reranks/predictionguard/src/test_helpers.py b/comps/reranks/predictionguard/src/test_helpers.py new file mode 100644 index 0000000000..38badd1237 --- /dev/null +++ b/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() diff --git a/tests/reranks/test_reranks_predictionguard.sh b/tests/reranks/test_reranks_predictionguard.sh new file mode 100644 index 0000000000..9f6913898d --- /dev/null +++ b/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