/ src / python / txtai / api / routers / llm.py
llm.py
 1  """
 2  Defines API paths for llm 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("/llm")
17  def llm(text: str, maxlength: Optional[int] = None, stream: Optional[bool] = False, stripthink: Optional[bool] = None):
18      """
19      Runs a LLM pipeline for the input text.
20  
21      Args:
22          text: input text
23          maxlength: optional response max length
24          stream: streams response if True
25          stripthink: strip thinking tags if True
26  
27      Returns:
28          response text
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("llm", text, **kwargs)
37  
38      # Handle both standard and streaming responses
39      return StreamingResponse(result) if stream else result
40  
41  
42  @router.post("/batchllm")
43  def batchllm(
44      texts: 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 LLM pipeline for the input texts.
51  
52      Args:
53          texts: input texts
54          maxlength: optional response max length
55          stream: streams response if True
56          stripthink: strip thinking tags if True
57  
58      Returns:
59          response texts
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("llm", texts, **kwargs)
68  
69      # Handle both standard and streaming responses
70      return StreamingResponse(result) if stream else result