@@ -152,7 +152,21 @@ async def query_endpoint_handler(
152152 auth : Annotated [AuthTuple , Depends (auth_dependency )],
153153 mcp_headers : dict [str , dict [str , str ]] = Depends (mcp_headers_dependency ),
154154) -> QueryResponse :
155- """Handle request to the /query endpoint."""
155+ """
156+ Handle request to the /query endpoint.
157+
158+ Processes a POST request to the /query endpoint, forwarding the
159+ user's query to a selected Llama Stack LLM or agent and
160+ returning the generated response.
161+
162+ Validates configuration and authentication, selects the appropriate model
163+ and provider, retrieves the LLM response, updates metrics, and optionally
164+ stores a transcript of the interaction. Handles connection errors to the
165+ Llama Stack service by returning an HTTP 500 error.
166+
167+ Returns:
168+ QueryResponse: Contains the conversation ID and the LLM-generated response.
169+ """
156170 check_configuration_loaded (configuration )
157171
158172 llama_stack_config = configuration .llama_stack_configuration
@@ -242,7 +256,24 @@ async def query_endpoint_handler(
242256def select_model_and_provider_id (
243257 models : ModelListResponse , model_id : str | None , provider_id : str | None
244258) -> tuple [str , str , str ]:
245- """Select the model ID and provider ID based on the request or available models."""
259+ """
260+ Select the model ID and provider ID based on the request or available models.
261+
262+ Determine and return the appropriate model and provider IDs for
263+ a query request.
264+
265+ If the request specifies both model and provider IDs, those are used.
266+ Otherwise, defaults from configuration are applied. If neither is
267+ available, selects the first available LLM model from the provided model
268+ list. Validates that the selected model exists among the available models.
269+
270+ Returns:
271+ A tuple containing the combined model ID (in the format
272+ "provider/model") and the provider ID.
273+
274+ Raises:
275+ HTTPException: If no suitable LLM model is found or the selected model is not available.
276+ """
246277 # If model_id and provider_id are provided in the request, use them
247278
248279 # If model_id is not provided in the request, check the configuration
@@ -303,16 +334,44 @@ def select_model_and_provider_id(
303334
304335
305336def _is_inout_shield (shield : Shield ) -> bool :
337+ """
338+ Determine if the shield identifier indicates an input/output shield.
339+
340+ Parameters:
341+ shield (Shield): The shield to check.
342+
343+ Returns:
344+ bool: True if the shield identifier starts with "inout_", otherwise False.
345+ """
306346 return shield .identifier .startswith ("inout_" )
307347
308348
309349def is_output_shield (shield : Shield ) -> bool :
310- """Determine if the shield is for monitoring output."""
350+ """
351+ Determine if the shield is for monitoring output.
352+
353+ Return True if the given shield is classified as an output or
354+ inout shield.
355+
356+ A shield is considered an output shield if its identifier
357+ starts with "output_" or "inout_".
358+ """
311359 return _is_inout_shield (shield ) or shield .identifier .startswith ("output_" )
312360
313361
314362def is_input_shield (shield : Shield ) -> bool :
315- """Determine if the shield is for monitoring input."""
363+ """
364+ Determine if the shield is for monitoring input.
365+
366+ Return True if the shield is classified as an input or inout
367+ shield.
368+
369+ Parameters:
370+ shield (Shield): The shield identifier to classify.
371+
372+ Returns:
373+ bool: True if the shield is for input or both input/output monitoring; False otherwise.
374+ """
316375 return _is_inout_shield (shield ) or not is_output_shield (shield )
317376
318377
@@ -323,7 +382,31 @@ async def retrieve_response( # pylint: disable=too-many-locals
323382 token : str ,
324383 mcp_headers : dict [str , dict [str , str ]] | None = None ,
325384) -> tuple [str , str ]:
326- """Retrieve response from LLMs and agents."""
385+ """
386+ Retrieve response from LLMs and agents.
387+
388+ Retrieves a response from the Llama Stack LLM or agent for a
389+ given query, handling shield configuration, tool usage, and
390+ attachment validation.
391+
392+ This function configures input/output shields, system prompts,
393+ and toolgroups (including RAG and MCP integration) as needed
394+ based on the query request and system configuration. It
395+ validates attachments, manages conversation and session
396+ context, and processes MCP headers for multi-component
397+ processing. Shield violations in the response are detected and
398+ corresponding metrics are updated.
399+
400+ Parameters:
401+ model_id (str): The identifier of the LLM model to use.
402+ query_request (QueryRequest): The user's query and associated metadata.
403+ token (str): The authentication token for authorization.
404+ mcp_headers (dict[str, dict[str, str]], optional): Headers for multi-component processing.
405+
406+ Returns:
407+ tuple[str, str]: A tuple containing the LLM or agent's response content
408+ and the conversation ID.
409+ """
327410 available_input_shields = [
328411 shield .identifier
329412 for shield in filter (is_input_shield , await client .shields .list ())
@@ -416,7 +499,9 @@ async def retrieve_response( # pylint: disable=too-many-locals
416499def validate_attachments_metadata (attachments : list [Attachment ]) -> None :
417500 """Validate the attachments metadata provided in the request.
418501
419- Raises HTTPException if any attachment has an improper type or content type.
502+ Raises:
503+ HTTPException: If any attachment has an invalid type or content type,
504+ an HTTP 422 error is raised.
420505 """
421506 for attachment in attachments :
422507 if attachment .attachment_type not in constants .ATTACHMENT_TYPES :
@@ -444,7 +529,19 @@ def validate_attachments_metadata(attachments: list[Attachment]) -> None:
444529
445530
446531def construct_transcripts_path (user_id : str , conversation_id : str ) -> Path :
447- """Construct path to transcripts."""
532+ """
533+ Construct path to transcripts.
534+
535+ Constructs a sanitized filesystem path for storing transcripts
536+ based on the user ID and conversation ID.
537+
538+ Parameters:
539+ user_id (str): The user identifier, which will be normalized and sanitized.
540+ conversation_id (str): The conversation identifier, which will be normalized and sanitized.
541+
542+ Returns:
543+ Path: The constructed path for storing transcripts for the specified user and conversation.
544+ """
448545 # these two normalizations are required by Snyk as it detects
449546 # this Path sanitization pattern
450547 uid = os .path .normpath ("/" + user_id ).lstrip ("/" )
@@ -468,7 +565,14 @@ def store_transcript( # pylint: disable=too-many-arguments,too-many-positional-
468565 truncated : bool ,
469566 attachments : list [Attachment ],
470567) -> None :
471- """Store transcript in the local filesystem.
568+ """
569+ Store transcript in the local filesystem.
570+
571+ Constructs a sanitized filesystem path for storing transcripts
572+ based on the user ID and conversation ID.
573+
574+ Returns:
575+ Path: The constructed path for storing transcripts for the specified user and conversation.
472576
473577 Args:
474578 user_id: The user ID (UUID).
@@ -513,7 +617,19 @@ def store_transcript( # pylint: disable=too-many-arguments,too-many-positional-
513617def get_rag_toolgroups (
514618 vector_db_ids : list [str ],
515619) -> list [Toolgroup ] | None :
516- """Return a list of RAG Tool groups if the given vector DB list is not empty."""
620+ """
621+ Return a list of RAG Tool groups if the given vector DB list is not empty.
622+
623+ Generate a list containing a RAG knowledge search toolgroup if
624+ vector database IDs are provided.
625+
626+ Parameters:
627+ vector_db_ids (list[str]): List of vector database identifiers to include in the toolgroup.
628+
629+ Returns:
630+ list[Toolgroup] | None: A list with a single RAG toolgroup if
631+ vector_db_ids is non-empty; otherwise, None.
632+ """
517633 return (
518634 [
519635 ToolgroupAgentToolGroupWithArgs (
0 commit comments