Skip to content

Commit 62963c9

Browse files
authored
Merge pull request #407 from tisnik/update-documantation-for-query-endpoint-handler-module
LCORE-390: Update documentation for /query endpoint handler module
2 parents 8cf5462 + a0771bf commit 62963c9

File tree

1 file changed

+125
-9
lines changed

1 file changed

+125
-9
lines changed

src/app/endpoints/query.py

Lines changed: 125 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -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(
242256
def 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

305336
def _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

309349
def 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

314362
def 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
416499
def 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

446531
def 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-
513617
def 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

Comments
 (0)