diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index c02656eb8..749fa34f7 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -18,6 +18,7 @@ from llama_stack_client.types.agents.turn_create_params import ( Toolgroup, ToolgroupAgentToolGroupWithArgs, + Document, ) from llama_stack_client.types.model_list_response import ModelListResponse from llama_stack_client.types.shared.interleaved_content_item import TextContentItem @@ -692,10 +693,20 @@ async def retrieve_response( # pylint: disable=too-many-locals,too-many-branche if not toolgroups: toolgroups = None + # TODO: LCORE-881 - Remove if Llama Stack starts to support these mime types + documents: list[Document] = [ + ( + {"content": doc["content"], "mime_type": "text/plain"} + if doc["mime_type"].lower() in ("application/json", "application/xml") + else doc + ) + for doc in query_request.get_documents() + ] + response = await agent.create_turn( messages=[UserMessage(role="user", content=query_request.query)], session_id=session_id, - documents=query_request.get_documents(), + documents=documents, stream=False, toolgroups=toolgroups, ) diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 7f7caa79b..d4ad3088a 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -1,4 +1,4 @@ -"""Handler for REST API call to provide answer to streaming query.""" # pylint: disable=too-many-lines +"""Handler for REST API call to provide answer to streaming query.""" # pylint: disable=too-many-lines,too-many-locals,W0511 import ast import json @@ -21,6 +21,7 @@ ) from llama_stack_client.types.shared import ToolCall from llama_stack_client.types.shared.interleaved_content_item import TextContentItem +from llama_stack_client.types.agents.turn_create_params import Document from app.database import get_session from app.endpoints.query import ( @@ -62,6 +63,7 @@ from utils.transcripts import store_transcript from utils.types import TurnSummary + logger = logging.getLogger("app.endpoints.handlers") router = APIRouter(tags=["streaming_query"]) @@ -1039,10 +1041,20 @@ async def retrieve_response( if not toolgroups: toolgroups = None + # TODO: LCORE-881 - Remove if Llama Stack starts to support these mime types + documents: list[Document] = [ + ( + {"content": doc["content"], "mime_type": "text/plain"} + if doc["mime_type"].lower() in ("application/json", "application/xml") + else doc + ) + for doc in query_request.get_documents() + ] + response = await agent.create_turn( messages=[UserMessage(role="user", content=query_request.query)], session_id=session_id, - documents=query_request.get_documents(), + documents=documents, stream=True, toolgroups=toolgroups, ) diff --git a/tests/e2e/features/query.feature b/tests/e2e/features/query.feature index ed2175563..4505111b2 100644 --- a/tests/e2e/features/query.feature +++ b/tests/e2e/features/query.feature @@ -113,4 +113,30 @@ Scenario: Check if LLM responds for query request with error for missing query {"query": "Say hello"} """ Then The status code of the response is 500 - And The body of the response contains Unable to connect to Llama Stack \ No newline at end of file + And The body of the response contains Unable to connect to Llama Stack + + Scenario: Check if LLM responds properly when XML and JSON attachments are sent + Given The system is in default state + And I set the Authorization header to Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6Ikpva + When I use "query" to ask question with authorization header + """ + { + "query": "Say hello", + "attachments": [ + { + "attachment_type": "configuration", + "content": "UserSystemHello", + "content_type": "application/xml" + }, + { + "attachment_type": "configuration", + "content": "{\"foo\": \"bar\"}", + "content_type": "application/json" + } + ], + "model": "{MODEL}", + "provider": "{PROVIDER}", + "system_prompt": "You are a helpful assistant" + } + """ + Then The status code of the response is 200 diff --git a/tests/e2e/features/streaming_query.feature b/tests/e2e/features/streaming_query.feature index 17f9c9787..23d79ea0b 100644 --- a/tests/e2e/features/streaming_query.feature +++ b/tests/e2e/features/streaming_query.feature @@ -88,4 +88,30 @@ Feature: streaming_query endpoint API tests {"query": "Say hello", "model": "{MODEL}"} """ Then The status code of the response is 422 - And The body of the response contains Value error, Provider must be specified if model is specified \ No newline at end of file + And The body of the response contains Value error, Provider must be specified if model is specified + + Scenario: Check if LLM responds properly when XML and JSON attachments are sent + Given The system is in default state + And I set the Authorization header to Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6Ikpva + When I use "streaming_query" to ask question with authorization header + """ + { + "query": "Say hello", + "attachments": [ + { + "attachment_type": "configuration", + "content": "UserSystemHello", + "content_type": "application/xml" + }, + { + "attachment_type": "configuration", + "content": "{\"foo\": \"bar\"}", + "content_type": "application/json" + } + ], + "model": "{MODEL}", + "provider": "{PROVIDER}", + "system_prompt": "You are a helpful assistant" + } + """ + Then The status code of the response is 200