diff --git a/docs/responses.md b/docs/responses.md index 1eafb22ab..153c939f9 100644 --- a/docs/responses.md +++ b/docs/responses.md @@ -119,7 +119,7 @@ The following table maps LCORE query request fields to the OpenResponses request | `system_prompt` | `instructions` | Same meaning. Only change in attribute's name | | `attachments` | `input` items | Attachments can be passed as input messages with content of type `input_file` | | `no_tools` | `tool_choice` | `no_tools=true` mapped to `tool_choice="none"` | -| `vector_store_ids` | `tools` + `tool_choice` | Vector stores can be explicitly specified and restricted by `file_search` tool type's `vector_store_ids` attribute | +| `vector_store_ids` | `tools` + `tool_choice` | Restrict via `file_search.vector_store_ids` in **LCORE format**; translated to Llama Stack internally. | | `generate_topic_summary` | N/A | Exposed directly (LCORE-specific) | | `shield_ids` | N/A | Exposed directly (LCORE-specific) | | `solr` | N/A | Exposed directly (LCORE-specific) | @@ -332,7 +332,7 @@ Each item in `tools` declares one capability: search a set of vector stores (**f **Tool types (each object has a required `type`):** -- `file_search`: Search within given vector stores. `vector_store_ids` (required): list of vector store IDs. Optional: `max_num_results` (1–50, default 10), `filters`, `ranking_options`. +- `file_search`: Search within given vector stores. `vector_store_ids` (required): **LCORE format** IDs (mapped to Llama Stack internally). Optional: `max_num_results` (1–50, default 10), `filters`, `ranking_options`. - `web_search`: Web search. `type` can be `"web_search"`, `"web_search_preview"`, or other variants. Optional: `search_context_size` (`"low"`, `"medium"`, `"high"`). - `function`: Call a named function. `name` (required). Optional: `description`, `parameters` (JSON schema), `strict`. - `mcp`: Use tools from an MCP server. `server_label` (required), `server_url` (required). Optional: `headers`, `require_approval`, `allowed_tools`. @@ -515,6 +515,8 @@ Fields such as `media_type`, `tool_calls`, `tool_results`, `rag_chunks`, and `re Vector store IDs are configured within the `tools` as `file_search` tools rather than through separate parameters. MCP tools are configurable under `mcp` tool type. By default **all** tools that are configured in LCORE are used to support the response. The set of available tools can be maintained per-request by `tool_choice` or `tools` attributes. +**Vector store IDs:** Accepts **LCORE format** in requests and also outputs it in responses; LCORE translates to/from Llama Stack format internally. + ### LCORE-Specific Extensions The API introduces extensions that are not part of the OpenResponses specification: diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 8bbab146d..f05296238 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -66,6 +66,7 @@ validate_model_provider_override, ) from utils.quota import check_tokens_available, get_available_quotas +from utils.tool_formatter import translate_vector_store_ids_to_user_facing from utils.responses import ( build_tool_call_summary, build_turn_summary, @@ -480,9 +481,16 @@ async def response_generator( chunk_dict["sequence_number"] = sequence_number sequence_number += 1 - # Add conversation attribute to the response if chunk has it if "response" in chunk_dict: chunk_dict["response"]["conversation"] = normalized_conv_id + tools = chunk_dict["response"].get("tools") + if tools is not None: + chunk_dict["response"]["tools"] = ( + translate_vector_store_ids_to_user_facing( + tools, + configuration.rag_id_mapping, + ) + ) # Intermediate response - no quota consumption and text yet if event_type == "response.in_progress": @@ -724,9 +732,16 @@ async def handle_non_streaming_response( skip_userid_check=skip_userid_check, topic_summary=topic_summary, ) + response_dict = api_response.model_dump(exclude_none=True) + tools = response_dict.get("tools") + if tools is not None: + response_dict["tools"] = translate_vector_store_ids_to_user_facing( + tools, + configuration.rag_id_mapping, + ) response = ResponsesResponse.model_validate( { - **api_response.model_dump(exclude_none=True), + **response_dict, "available_quotas": available_quotas, "conversation": normalize_conversation_id(api_params.conversation), "completed_at": int(completed_at.timestamp()), diff --git a/src/models/requests.py b/src/models/requests.py index ecd53f485..7a8aba99c 100644 --- a/src/models/requests.py +++ b/src/models/requests.py @@ -23,6 +23,7 @@ ) from pydantic import BaseModel, Field, field_validator, model_validator +from configuration import configuration from constants import ( MCP_AUTH_CLIENT, MCP_AUTH_KUBERNETES, @@ -32,6 +33,7 @@ ) from log import get_logger from utils import suid +from utils.tool_formatter import translate_vector_store_ids_to_user_facing from utils.types import IncludeParameter, ResponseInput logger = get_logger(__name__) @@ -758,16 +760,14 @@ def check_previous_response_id(cls, value: Optional[str]) -> Optional[str]: return value def echoed_params(self) -> dict[str, Any]: - """Dump attributes that are echoed back in the response. - - The tools attribute is converted from input tool to output tool model. + """Build kwargs echoed into synthetic OpenAI-style responses (e.g. moderation blocks). Returns: - Dict of echoed attributes. + dict[str, Any]: Field names and values to merge into the response object. """ data = self.model_dump(include=_ECHOED_FIELDS) if self.tools is not None: - data["tools"] = [ + tool_dicts: list[dict[str, Any]] = [ ( OutputToolMCP.model_validate(t.model_dump()).model_dump() if t.type == "mcp" @@ -775,6 +775,9 @@ def echoed_params(self) -> dict[str, Any]: ) for t in self.tools ] + data["tools"] = translate_vector_store_ids_to_user_facing( + tool_dicts, configuration.rag_id_mapping + ) return data diff --git a/src/utils/tool_formatter.py b/src/utils/tool_formatter.py index 73b277777..e9e9bdc69 100644 --- a/src/utils/tool_formatter.py +++ b/src/utils/tool_formatter.py @@ -1,5 +1,6 @@ """Utility functions for formatting and parsing MCP tool descriptions.""" +from collections.abc import Mapping from typing import Any from log import get_logger @@ -130,3 +131,30 @@ def format_tools_list(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: fields and cleaned descriptions. """ return [format_tool_response(tool) for tool in tools] + + +def translate_vector_store_ids_to_user_facing( + tools: list[dict[str, Any]], + rag_id_mapping: Mapping[str, str], +) -> list[dict[str, Any]]: + """ + Rewrite file_search tool dicts so vector_store_ids use user-facing RAG IDs. + + Parameters: + tools: Serialized tool dicts. + rag_id_mapping: Llama Stack vector_db_id -> user-facing RAG id. + + Returns: + list[dict[str, Any]]: New list of tool dicts; file_search entries get + updated vector_store_ids. + """ + if not rag_id_mapping: + return list(tools) + out: list[dict[str, Any]] = [] + for tool in tools: + if tool["type"] == "file_search": + mapped = [rag_id_mapping.get(vid, vid) for vid in tool["vector_store_ids"]] + out.append({**tool, "vector_store_ids": mapped}) + else: + out.append(tool) + return out diff --git a/tests/unit/utils/test_tool_formatter.py b/tests/unit/utils/test_tool_formatter.py new file mode 100644 index 000000000..8da98b094 --- /dev/null +++ b/tests/unit/utils/test_tool_formatter.py @@ -0,0 +1,76 @@ +"""Unit tests for tool_formatter utilities.""" + +from typing import Any + +from utils.tool_formatter import translate_vector_store_ids_to_user_facing + + +class TestTranslateVectorStoreIdsToUserFacing: + """Tests for translate_vector_store_ids_to_user_facing.""" + + def test_empty_mapping_returns_new_list_same_tool_objects(self) -> None: + """When mapping is empty, return a new list with the same tool dicts.""" + tools: list[dict[str, Any]] = [ + {"type": "file_search", "vector_store_ids": ["vs-1"]}, + ] + result = translate_vector_store_ids_to_user_facing(tools, {}) + assert result is not tools + assert result == tools + assert result[0] is tools[0] + + def test_file_search_vector_store_ids_mapped(self) -> None: + """file_search tools get vector_store_ids rewritten via mapping.""" + tools: list[dict[str, Any]] = [ + { + "type": "file_search", + "vector_store_ids": ["llama-vs", "other-vs"], + }, + ] + mapping = {"llama-vs": "user-rag-a", "other-vs": "user-rag-b"} + result = translate_vector_store_ids_to_user_facing(tools, mapping) + assert result[0]["vector_store_ids"] == ["user-rag-a", "user-rag-b"] + assert result[0]["type"] == "file_search" + + def test_unmapped_id_passthrough(self) -> None: + """IDs absent from mapping are left unchanged.""" + tools: list[dict[str, Any]] = [ + {"type": "file_search", "vector_store_ids": ["known", "unknown-id"]}, + ] + result = translate_vector_store_ids_to_user_facing( + tools, {"known": "user-facing"} + ) + assert result[0]["vector_store_ids"] == ["user-facing", "unknown-id"] + + def test_non_file_search_tool_unchanged_identity(self) -> None: + """Non-file_search entries are appended as the same dict instance.""" + mcp_tool: dict[str, Any] = {"type": "mcp", "server_url": "http://x"} + tools: list[dict[str, Any]] = [mcp_tool] + result = translate_vector_store_ids_to_user_facing(tools, {"any": "mapping"}) + assert len(result) == 1 + assert result[0] is mcp_tool + + def test_file_search_new_dict_instance(self) -> None: + """file_search entries are copied so original tool dict is not mutated.""" + original: dict[str, Any] = { + "type": "file_search", + "vector_store_ids": ["vs-1"], + } + tools: list[dict[str, Any]] = [original] + result = translate_vector_store_ids_to_user_facing(tools, {"vs-1": "u-1"}) + assert result[0] is not original + assert original["vector_store_ids"] == ["vs-1"] + + def test_mixed_tools_order_preserved(self) -> None: + """Order and handling per type are stable across a mixed tool list.""" + tools: list[dict[str, Any]] = [ + {"type": "file_search", "vector_store_ids": ["a"]}, + {"type": "function", "name": "fn"}, + {"type": "file_search", "vector_store_ids": ["b", "c"]}, + ] + result = translate_vector_store_ids_to_user_facing( + tools, {"a": "A", "b": "B", "c": "C"} + ) + assert [t["type"] for t in result] == ["file_search", "function", "file_search"] + assert result[0]["vector_store_ids"] == ["A"] + assert result[1] is tools[1] + assert result[2]["vector_store_ids"] == ["B", "C"]