Skip to content

Commit 1ee31fc

Browse files
authored
Merge pull request #1385 from asimurka/conversation_management
LCORE-1560: Output vector stores in user-facing format
2 parents 44afc1f + 6ff6acb commit 1ee31fc

5 files changed

Lines changed: 133 additions & 9 deletions

File tree

docs/responses.md

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ The following table maps LCORE query request fields to the OpenResponses request
119119
| `system_prompt` | `instructions` | Same meaning. Only change in attribute's name |
120120
| `attachments` | `input` items | Attachments can be passed as input messages with content of type `input_file` |
121121
| `no_tools` | `tool_choice` | `no_tools=true` mapped to `tool_choice="none"` |
122-
| `vector_store_ids` | `tools` + `tool_choice` | Vector stores can be explicitly specified and restricted by `file_search` tool type's `vector_store_ids` attribute |
122+
| `vector_store_ids` | `tools` + `tool_choice` | Restrict via `file_search.vector_store_ids` in **LCORE format**; translated to Llama Stack internally. |
123123
| `generate_topic_summary` | N/A | Exposed directly (LCORE-specific) |
124124
| `shield_ids` | N/A | Exposed directly (LCORE-specific) |
125125
| `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
332332

333333
**Tool types (each object has a required `type`):**
334334

335-
- `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`.
335+
- `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`.
336336
- `web_search`: Web search. `type` can be `"web_search"`, `"web_search_preview"`, or other variants. Optional: `search_context_size` (`"low"`, `"medium"`, `"high"`).
337337
- `function`: Call a named function. `name` (required). Optional: `description`, `parameters` (JSON schema), `strict`.
338338
- `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
515515

516516
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.
517517

518+
**Vector store IDs:** Accepts **LCORE format** in requests and also outputs it in responses; LCORE translates to/from Llama Stack format internally.
519+
518520
### LCORE-Specific Extensions
519521

520522
The API introduces extensions that are not part of the OpenResponses specification:

src/app/endpoints/responses.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@
6666
validate_model_provider_override,
6767
)
6868
from utils.quota import check_tokens_available, get_available_quotas
69+
from utils.tool_formatter import translate_vector_store_ids_to_user_facing
6970
from utils.responses import (
7071
build_tool_call_summary,
7172
build_turn_summary,
@@ -481,9 +482,16 @@ async def response_generator(
481482
chunk_dict["sequence_number"] = sequence_number
482483
sequence_number += 1
483484

484-
# Add conversation attribute to the response if chunk has it
485485
if "response" in chunk_dict:
486486
chunk_dict["response"]["conversation"] = normalized_conv_id
487+
tools = chunk_dict["response"].get("tools")
488+
if tools is not None:
489+
chunk_dict["response"]["tools"] = (
490+
translate_vector_store_ids_to_user_facing(
491+
tools,
492+
configuration.rag_id_mapping,
493+
)
494+
)
487495

488496
# Intermediate response - no quota consumption and text yet
489497
if event_type == "response.in_progress":
@@ -725,9 +733,16 @@ async def handle_non_streaming_response(
725733
skip_userid_check=skip_userid_check,
726734
topic_summary=topic_summary,
727735
)
736+
response_dict = api_response.model_dump(exclude_none=True)
737+
tools = response_dict.get("tools")
738+
if tools is not None:
739+
response_dict["tools"] = translate_vector_store_ids_to_user_facing(
740+
tools,
741+
configuration.rag_id_mapping,
742+
)
728743
response = ResponsesResponse.model_validate(
729744
{
730-
**api_response.model_dump(exclude_none=True),
745+
**response_dict,
731746
"available_quotas": available_quotas,
732747
"conversation": normalize_conversation_id(api_params.conversation),
733748
"completed_at": int(completed_at.timestamp()),

src/models/requests.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
)
2424
from pydantic import BaseModel, Field, field_validator, model_validator
2525

26+
from configuration import configuration
2627
from constants import (
2728
MCP_AUTH_CLIENT,
2829
MCP_AUTH_KUBERNETES,
@@ -32,6 +33,7 @@
3233
)
3334
from log import get_logger
3435
from utils import suid
36+
from utils.tool_formatter import translate_vector_store_ids_to_user_facing
3537
from utils.types import IncludeParameter, ResponseInput
3638

3739
logger = get_logger(__name__)
@@ -758,23 +760,24 @@ def check_previous_response_id(cls, value: Optional[str]) -> Optional[str]:
758760
return value
759761

760762
def echoed_params(self) -> dict[str, Any]:
761-
"""Dump attributes that are echoed back in the response.
762-
763-
The tools attribute is converted from input tool to output tool model.
763+
"""Build kwargs echoed into synthetic OpenAI-style responses (e.g. moderation blocks).
764764
765765
Returns:
766-
Dict of echoed attributes.
766+
dict[str, Any]: Field names and values to merge into the response object.
767767
"""
768768
data = self.model_dump(include=_ECHOED_FIELDS)
769769
if self.tools is not None:
770-
data["tools"] = [
770+
tool_dicts: list[dict[str, Any]] = [
771771
(
772772
OutputToolMCP.model_validate(t.model_dump()).model_dump()
773773
if t.type == "mcp"
774774
else t.model_dump()
775775
)
776776
for t in self.tools
777777
]
778+
data["tools"] = translate_vector_store_ids_to_user_facing(
779+
tool_dicts, configuration.rag_id_mapping
780+
)
778781

779782
return data
780783

src/utils/tool_formatter.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Utility functions for formatting and parsing MCP tool descriptions."""
22

3+
from collections.abc import Mapping
34
from typing import Any
45

56
from log import get_logger
@@ -130,3 +131,30 @@ def format_tools_list(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
130131
fields and cleaned descriptions.
131132
"""
132133
return [format_tool_response(tool) for tool in tools]
134+
135+
136+
def translate_vector_store_ids_to_user_facing(
137+
tools: list[dict[str, Any]],
138+
rag_id_mapping: Mapping[str, str],
139+
) -> list[dict[str, Any]]:
140+
"""
141+
Rewrite file_search tool dicts so vector_store_ids use user-facing RAG IDs.
142+
143+
Parameters:
144+
tools: Serialized tool dicts.
145+
rag_id_mapping: Llama Stack vector_db_id -> user-facing RAG id.
146+
147+
Returns:
148+
list[dict[str, Any]]: New list of tool dicts; file_search entries get
149+
updated vector_store_ids.
150+
"""
151+
if not rag_id_mapping:
152+
return list(tools)
153+
out: list[dict[str, Any]] = []
154+
for tool in tools:
155+
if tool["type"] == "file_search":
156+
mapped = [rag_id_mapping.get(vid, vid) for vid in tool["vector_store_ids"]]
157+
out.append({**tool, "vector_store_ids": mapped})
158+
else:
159+
out.append(tool)
160+
return out
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
"""Unit tests for tool_formatter utilities."""
2+
3+
from typing import Any
4+
5+
from utils.tool_formatter import translate_vector_store_ids_to_user_facing
6+
7+
8+
class TestTranslateVectorStoreIdsToUserFacing:
9+
"""Tests for translate_vector_store_ids_to_user_facing."""
10+
11+
def test_empty_mapping_returns_new_list_same_tool_objects(self) -> None:
12+
"""When mapping is empty, return a new list with the same tool dicts."""
13+
tools: list[dict[str, Any]] = [
14+
{"type": "file_search", "vector_store_ids": ["vs-1"]},
15+
]
16+
result = translate_vector_store_ids_to_user_facing(tools, {})
17+
assert result is not tools
18+
assert result == tools
19+
assert result[0] is tools[0]
20+
21+
def test_file_search_vector_store_ids_mapped(self) -> None:
22+
"""file_search tools get vector_store_ids rewritten via mapping."""
23+
tools: list[dict[str, Any]] = [
24+
{
25+
"type": "file_search",
26+
"vector_store_ids": ["llama-vs", "other-vs"],
27+
},
28+
]
29+
mapping = {"llama-vs": "user-rag-a", "other-vs": "user-rag-b"}
30+
result = translate_vector_store_ids_to_user_facing(tools, mapping)
31+
assert result[0]["vector_store_ids"] == ["user-rag-a", "user-rag-b"]
32+
assert result[0]["type"] == "file_search"
33+
34+
def test_unmapped_id_passthrough(self) -> None:
35+
"""IDs absent from mapping are left unchanged."""
36+
tools: list[dict[str, Any]] = [
37+
{"type": "file_search", "vector_store_ids": ["known", "unknown-id"]},
38+
]
39+
result = translate_vector_store_ids_to_user_facing(
40+
tools, {"known": "user-facing"}
41+
)
42+
assert result[0]["vector_store_ids"] == ["user-facing", "unknown-id"]
43+
44+
def test_non_file_search_tool_unchanged_identity(self) -> None:
45+
"""Non-file_search entries are appended as the same dict instance."""
46+
mcp_tool: dict[str, Any] = {"type": "mcp", "server_url": "http://x"}
47+
tools: list[dict[str, Any]] = [mcp_tool]
48+
result = translate_vector_store_ids_to_user_facing(tools, {"any": "mapping"})
49+
assert len(result) == 1
50+
assert result[0] is mcp_tool
51+
52+
def test_file_search_new_dict_instance(self) -> None:
53+
"""file_search entries are copied so original tool dict is not mutated."""
54+
original: dict[str, Any] = {
55+
"type": "file_search",
56+
"vector_store_ids": ["vs-1"],
57+
}
58+
tools: list[dict[str, Any]] = [original]
59+
result = translate_vector_store_ids_to_user_facing(tools, {"vs-1": "u-1"})
60+
assert result[0] is not original
61+
assert original["vector_store_ids"] == ["vs-1"]
62+
63+
def test_mixed_tools_order_preserved(self) -> None:
64+
"""Order and handling per type are stable across a mixed tool list."""
65+
tools: list[dict[str, Any]] = [
66+
{"type": "file_search", "vector_store_ids": ["a"]},
67+
{"type": "function", "name": "fn"},
68+
{"type": "file_search", "vector_store_ids": ["b", "c"]},
69+
]
70+
result = translate_vector_store_ids_to_user_facing(
71+
tools, {"a": "A", "b": "B", "c": "C"}
72+
)
73+
assert [t["type"] for t in result] == ["file_search", "function", "file_search"]
74+
assert result[0]["vector_store_ids"] == ["A"]
75+
assert result[1] is tools[1]
76+
assert result[2]["vector_store_ids"] == ["B", "C"]

0 commit comments

Comments
 (0)