rag.py
1 """ 2 Defines API paths for rag endpoints. 3 """ 4 5 from typing import List, Optional 6 7 from fastapi import APIRouter, Body 8 from fastapi.responses import StreamingResponse 9 10 from .. import application 11 from ..route import EncodingAPIRoute 12 13 router = APIRouter(route_class=EncodingAPIRoute) 14 15 16 @router.get("/rag") 17 def rag(query: str, maxlength: Optional[int] = None, stream: Optional[bool] = False, stripthink: Optional[bool] = None): 18 """ 19 Runs a RAG pipeline for the input query. 20 21 Args: 22 query: input RAG query 23 maxlength: optional response max length 24 stream: streams response if True 25 stripthink: strip thinking tags if True 26 27 Returns: 28 answer 29 """ 30 31 # Build keyword arguments 32 params = [("maxlength", maxlength), ("stream", stream), ("stripthink", stripthink)] 33 kwargs = {key: value for key, value in params if value} 34 35 # Run pipeline 36 result = application.get().pipeline("rag", query, **kwargs) 37 38 # Handle both standard and streaming responses 39 return StreamingResponse(result) if stream else result 40 41 42 @router.post("/batchrag") 43 def batchrag( 44 queries: List[str] = Body(...), 45 maxlength: Optional[int] = Body(default=None), 46 stream: Optional[bool] = Body(default=False), 47 stripthink: Optional[bool] = Body(default=None), 48 ): 49 """ 50 Runs a RAG pipeline for the input queries. 51 52 Args: 53 queries: input RAG queries 54 maxlength: optional response max length 55 stream: streams response if True 56 stripthink: strip thinking tags if True 57 58 Returns: 59 answers 60 """ 61 62 # Build keyword arguments 63 params = [("maxlength", maxlength), ("stream", stream), ("stripthink", stripthink)] 64 kwargs = {key: value for key, value in params if value} 65 66 # Run pipeline 67 result = application.get().pipeline("rag", queries, **kwargs) 68 69 # Handle both standard and streaming responses 70 return StreamingResponse(result) if stream else result