/ src / python / txtai / api / routers / rag.py
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