jandres/GenAIExamples
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

move examples gateway (#992)

Browse Source 
Co-authored-by: root <root@idc708073.jf.intel.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sihan Chen <39623753+Spycsh@users.noreply.github.com>
This commit is contained in:
lkk
2024-12-06 14:40:25 +08:00
committed by GitHub
parent f5c08d4fbb
commit bde285dfce
17 changed files with 1236 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
DocIndexRetriever/retrieval_tool.py
View File

@@ -3,10 +3,14 @@
import asyncio
import os
from typing import Union
from comps import MicroService, RetrievalToolGateway, ServiceOrchestrator, ServiceType
from comps import Gateway, MegaServiceEndpoint, MicroService, ServiceOrchestrator, ServiceType
from comps.cores.proto.api_protocol import ChatCompletionRequest, EmbeddingRequest
from comps.cores.proto.docarray import LLMParamsDoc, RerankedDoc, RerankerParms, RetrieverParms, TextDoc
from fastapi import Request
from fastapi.responses import StreamingResponse
MEGA_SERVICE_HOST_IP = os.getenv("MEGA_SERVICE_HOST_IP", "0.0.0.0")
MEGA_SERVICE_PORT = os.getenv("MEGA_SERVICE_PORT", 8889)
EMBEDDING_SERVICE_HOST_IP = os.getenv("EMBEDDING_SERVICE_HOST_IP", "0.0.0.0")
EMBEDDING_SERVICE_PORT = os.getenv("EMBEDDING_SERVICE_PORT", 6000)
@@ -16,7 +20,7 @@ RERANK_SERVICE_HOST_IP = os.getenv("RERANK_SERVICE_HOST_IP", "0.0.0.0")
RERANK_SERVICE_PORT = os.getenv("RERANK_SERVICE_PORT", 8000)
class RetrievalToolService:
class RetrievalToolService(Gateway):
def __init__(self, host="0.0.0.0", port=8000):
self.host = host
self.port = port
@@ -51,9 +55,77 @@ class RetrievalToolService:
self.megaservice.add(embedding).add(retriever).add(rerank)
self.megaservice.flow_to(embedding, retriever)
self.megaservice.flow_to(retriever, rerank)
self.gateway = RetrievalToolGateway(megaservice=self.megaservice, host="0.0.0.0", port=self.port)
async def handle_request(self, request: Request):
def parser_input(data, TypeClass, key):
chat_request = None
try:
chat_request = TypeClass.parse_obj(data)
query = getattr(chat_request, key)
except:
query = None
return query, chat_request
data = await request.json()
query = None
for key, TypeClass in zip(["text", "input", "messages"], [TextDoc, EmbeddingRequest, ChatCompletionRequest]):
query, chat_request = parser_input(data, TypeClass, key)
if query is not None:
break
if query is None:
raise ValueError(f"Unknown request type: {data}")
if chat_request is None:
raise ValueError(f"Unknown request type: {data}")
if isinstance(chat_request, ChatCompletionRequest):
retriever_parameters = RetrieverParms(
search_type=chat_request.search_type if chat_request.search_type else "similarity",
k=chat_request.k if chat_request.k else 4,
distance_threshold=chat_request.distance_threshold if chat_request.distance_threshold else None,
fetch_k=chat_request.fetch_k if chat_request.fetch_k else 20,
lambda_mult=chat_request.lambda_mult if chat_request.lambda_mult else 0.5,
score_threshold=chat_request.score_threshold if chat_request.score_threshold else 0.2,
)
reranker_parameters = RerankerParms(
top_n=chat_request.top_n if chat_request.top_n else 1,
)
initial_inputs = {
"messages": query,
"input": query, # has to be input due to embedding expects either input or text
"search_type": chat_request.search_type if chat_request.search_type else "similarity",
"k": chat_request.k if chat_request.k else 4,
"distance_threshold": chat_request.distance_threshold if chat_request.distance_threshold else None,
"fetch_k": chat_request.fetch_k if chat_request.fetch_k else 20,
"lambda_mult": chat_request.lambda_mult if chat_request.lambda_mult else 0.5,
"score_threshold": chat_request.score_threshold if chat_request.score_threshold else 0.2,
"top_n": chat_request.top_n if chat_request.top_n else 1,
}
result_dict, runtime_graph = await self.megaservice.schedule(
initial_inputs=initial_inputs,
retriever_parameters=retriever_parameters,
reranker_parameters=reranker_parameters,
)
else:
result_dict, runtime_graph = await self.megaservice.schedule(initial_inputs={"text": query})
last_node = runtime_graph.all_leaves()[-1]
response = result_dict[last_node]
return response
def start(self):
super().__init__(
megaservice=self.megaservice,
host=self.host,
port=self.port,
endpoint=str(MegaServiceEndpoint.RETRIEVALTOOL),
input_datatype=Union[TextDoc, EmbeddingRequest, ChatCompletionRequest],
output_datatype=Union[RerankedDoc, LLMParamsDoc],
)
if __name__ == "__main__":
chatqna = RetrievalToolService(host=MEGA_SERVICE_HOST_IP, port=MEGA_SERVICE_PORT)
chatqna = RetrievalToolService(port=MEGA_SERVICE_PORT)
chatqna.add_remote_service()
chatqna.start()