"""SID-2 preview.Before using this template: implement the vector_search and keyword_search backendfunctions, set ID_FIELD and SNIPPET_FIELD to your document's fields, andcustomize the INSTRUCTIONS, QUESTION, and TOOL_EXAMPLES to fit your data. You canalso define additional tools, or extend existing tools with more parameters (e.g. afilter). Just make sure to always update the tool's description and usage example.The report_helpful_ids tool is fixed and cannot be changed or modified.Run: SID_API_KEY=your-api-key uv run --with sid-sdk --with openai reference_implementation.py"""import jsonimport osimport uuidfrom concurrent.futures import ThreadPoolExecutorfrom openai import OpenAIfrom sid import DocumentCache, InvalidCharacterRange, parse_rendered_model_facing_idfrom sid.document_cache import SNIPPET_SIZE_DEFAULTclass ToolError(Exception): """A mistake the *model* made (bad id, malformed ref, invalid arguments), these errors are recoverable and passed back to the model as a message."""# ---------------------------------------------------------------------------# The episode. TODO: your question. The model is served by the first-party# SID API, under one SID_API_KEY.# ---------------------------------------------------------------------------QUESTION = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?"BASE_URL = os.environ.get("SID_BASE_URL", "https://api.sid-1.com/v1")MODEL = "sid-2"MAX_TURNS = 10REWARD = "f1" # the trained scoring prompt: "f1" (few, precise ids) | "ndcg" (many, best first)REASONING_EFFORT = "low" # "low" | "medium" | "high"DEFAULT_LIMIT = 5 # results per search when the model does not ask for a numberMAX_PARALLEL_FETCHES = 8 # backend queries in flight at once, across a turn's search calls# ---------------------------------------------------------------------------# TODO: your document backend. The tools below call these functions;# implement them against your database. A document is a dict with:# ID_FIELD : your document id (any stable string; never shown to the model)# SNIPPET_FIELD: the document's canonical text. Return the identical string# for the same id everywhere: the seen-ledger's character# offsets index it.# extras : (optional) any other fields, such as a "title", rendered# as display attributes. Every cached field is shown to the# model (ID_FIELD is popped before caching), so leave# anything private out of your documents.# ---------------------------------------------------------------------------ID_FIELD = "id"SNIPPET_FIELD = "content"def vector_search(query: str, limit: int) -> list[dict]: """TODO: your semantic (embedding) search: the ``limit`` most relevant documents for ``query``, best first.""" raise NotImplementedError("implement vector_search against your document store")def keyword_search(query: str, limit: int) -> list[dict]: """TODO: your full-text keyword (BM25) search: the ``limit`` best matching documents for ``query``.""" raise NotImplementedError("implement keyword_search against your document store")# ---------------------------------------------------------------------------# The system prompt: one sentence naming your subject matter, then the# search -> repeat -> report scaffolding. TODO: edit the first paragraph to# describe your data and goal; keep the steps.# ---------------------------------------------------------------------------INSTRUCTIONS = """You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called."""# ---------------------------------------------------------------------------# The tool schemas. Each tool should be demonstrated in TOOL_EXAMPLES.# ---------------------------------------------------------------------------SNIPPET_SIZE_DESCRIPTION = ( f"Words shown per result snippet (default {SNIPPET_SIZE_DEFAULT}). Must be a positive integer. " f"Pass a larger value to see more of each document.")SEARCH_TOOL = { "type": "function", "name": "search", "description": ( "Semantic (vector) search over the dataset. Best for conceptual or paraphrased " "queries: a long, descriptive query of what you want beats a few keywords; use " "text_search when you know the distinctive words a match would contain. Result " "ids may be ranged references 'id#start:end' naming the exact character span of " "the document shown; stretches you have already been shown are replaced by " "'[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the " "document's total character count." ), "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "The query to search the dataset with"}, "limit": {"type": "integer", "description": "The number of results to return"}, "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION}, }, "required": ["query"], },}TEXT_SEARCH_TOOL = { "type": "function", "name": "text_search", "description": ( "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you " "want matched as plain words separated by spaces: the distinctive content words a " "matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact " "character span of the document shown; stretches you have already been shown are " "replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length " "gives the document's total character count." ), "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": ( "Plain space-separated content words to match (BM25)." ), }, "limit": {"type": "integer", "description": "The number of results to return"}, "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION}, }, "required": ["query"], },}READ_TOOL = { "type": "function", "name": "read", "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.", "parameters": { "type": "object", "properties": { "id": { "type": "string", "description": ( "The id of the document to read. Also accepts a ranged reference " "'id#start:end' (character offsets, e.g. '123456#800:1600') to read only " "that part of the document's text; a range running past the end of the " "document is trimmed to the part that overlaps it." ), }, }, "required": ["id"], },}# The fixed terminal tool: do not change its name, description or parameters.REPORT_TOOL = { "type": "function", "name": "report_helpful_ids", "description": ( "Report the most helpful document ids from the search results. This is required in " "the last step. Each id may be reported at most once; repeated ids are penalized. " "Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a " "ranged reference counts as its whole document, so there is no need to strip the " "'#start:end' part." ), "parameters": { "type": "object", "properties": { "ids": { "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": {"type": "string"}, }, }, "required": ["ids"], },}TOOLS = [SEARCH_TOOL, TEXT_SEARCH_TOOL, READ_TOOL, REPORT_TOOL]TOOLS_BY_NAME = {tool["name"]: tool for tool in TOOLS}# Few-shot calls, sent on every request via the sid extension. TODO: write search# examples that fit your data; keep the read and report examples as they areTOOL_EXAMPLES = [ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": { "query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5, }, }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": {"query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5}, }, { "caption": "Example: reading one character range:", "name": "read", "arguments": {"id": "placeholder_1#800:1600"}, }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": {"ids": ["placeholder_1", "placeholder_2", "placeholder_3"]}, },]# ---------------------------------------------------------------------------# The tools, split in two: fetch_results is the slow backend call and runs in# parallel; everything below it touches the cache and runs serially, in call# order.# ---------------------------------------------------------------------------def fetch_results(name: str, arguments: dict) -> list[dict] | None: """The backend call behind one search tool call. `read` and the report have nothing to fetch (reads are served from the document cache), so they get ``None``.""" if name not in ("search", "text_search"): return None # Catch mistyped arguments as a ToolError that the model can recover from. properties = TOOLS_BY_NAME[name]["parameters"]["properties"] for key, value in arguments.items(): expected = str if properties[key]["type"] == "string" else int if not isinstance(value, expected) or isinstance(value, bool): raise ToolError( f"Invalid {name} arguments: {key} must be of type {properties[key]['type']}, " f"got {type(value).__name__}" ) query = arguments.get("query", "") # tool_arguments strips ""; treat missing as empty if not query: raise ToolError( "search requires a non-empty `query` to rank documents against. " "To fetch a specific document by id, use the `read` tool. " "To rank candidates against the question, call `search` with a meaningful query." ) snippet_size = arguments.get("snippet_size", SNIPPET_SIZE_DEFAULT) if snippet_size <= 0: raise ToolError( f"Invalid {name} arguments: snippet_size must be a positive integer, got {snippet_size}" ) limit = arguments.get("limit", DEFAULT_LIMIT) return vector_search(query, limit) if name == "search" else keyword_search(query, limit)def render_results(cache: DocumentCache, results: list[dict], query: str, snippet_size: int) -> str: """One <doc> per result, in backend order. Shared by both search tools; only their fetch differs.""" rendered = [] for result in results: if not isinstance(result, dict) or not isinstance(result.get(ID_FIELD), str) or not result[ID_FIELD]: raise ValueError(f"backend document must carry a non-empty string {ID_FIELD!r}: {result!r}") # Keep backend results intact when searches reuse the same dictionaries. document = dict(result) # Popped, not read: every cached field renders, and your database ids must not. data_id = document.pop(ID_FIELD) content = document.get(SNIPPET_FIELD) if not isinstance(content, str) or not content: # Not a model mistake: the seen-ledger's offsets index this string. raise ValueError(f"backend document {data_id!r} must carry non-empty {SNIPPET_FIELD!r} text") # Searches routinely return the same document, and the cache is add-once. if data_id not in cache: cache.add_document(data_id, document) # BM25 snippet, with already-shown stretches collapsed to [seen: ...]. view = cache.apply_snippet( data_id, SNIPPET_FIELD, query, snippet_size=snippet_size, ) rendered.append(view.render_xml()) cache.update_seen(view) # later views of this span collapse to [seen: ...] return "\n".join(rendered) if rendered else "No results found. Try a different query."def parse_reference(reference) -> tuple[str, tuple[int, int] | None]: """The model's ref in (``'R'`` or ``'R#start:end'``), split into id and optional char range. The SDK parser raises ``ValueError`` on malformed shapes with a descriptive message already written.""" try: return parse_rendered_model_facing_id(reference) except ValueError as error: raise ToolError(str(error)) from Nonedef read(cache: DocumentCache, reference: str | None) -> str: """Fetch one document (or a character range of it) by the id a search returned. Explicit reads are not masked: the model asked for this span, so it renders up to 10,000 characters. Model mistakes (bad ref grammar, an unknown or hallucinated id, an unusable range) raise ToolError.""" model_facing_id, char_range = parse_reference(reference) # ToolError on bad grammar if not cache.contains_model_facing_id(model_facing_id): raise ToolError( f"Invalid document id '{model_facing_id}'. " f"You can only read documents returned by search. " f"Use 'search' or 'text_search' first." ) data_id = cache.to_data_id(model_facing_id) # the one conversion back into your ids if char_range is not None: # Lenient mode trims partially overlapping ranges to the document. # Invalid ranges become recoverable tool errors. try: char_range = cache.resolve_char_range(data_id, SNIPPET_FIELD, char_range) except InvalidCharacterRange as error: raise ToolError( f"{error}. Re-read '{model_facing_id}' without a range to read from the " f"start, or pass a range inside 0:doc_length." ) from None # Cap document text per read; further ranged reads can continue from here. if char_range is not None: start, end = char_range char_range = (start, min(end, start + 10_000)) elif len(cache.get_document(data_id)[SNIPPET_FIELD]) > 10_000: char_range = (0, 10_000) view = cache.get_single_span_document_view( data_id, SNIPPET_FIELD, snippet_display_span=char_range, ) rendered = view.render_xml() cache.update_seen(view) return rendereddef report(cache: DocumentCache, ids) -> list[str]: """The terminal report: ids are doc-granular (a ranged ref denotes its document), deduplicated preserving order, and translated back to your data_ids. IDs outside the model-facing ID map are dropped.""" if not isinstance(ids, list) or not all(isinstance(reference, str) for reference in ids): raise ToolError("ids must be a list of strings") if not ids: raise ToolError("ids list cannot be empty") reported: list[str] = [] for reference in ids: model_facing_id, _ = parse_reference(reference) # ToolError on bad grammar if model_facing_id not in reported: reported.append(model_facing_id) return [ cache.to_data_id(model_facing_id) for model_facing_id in reported if cache.contains_model_facing_id(model_facing_id) ]# ---------------------------------------------------------------------------# The agent loop: one Responses API call per turn, tool outputs paired back# by call_id. The episode ends when the model reports.## Requests are stateless (store=False): nothing is retained server-side, so# the full conversation is resent on every call. The model's reasoning comes# back as an opaque ``encrypted_content`` blob on each reasoning item, and# carries over by echoing the output items back verbatim in the next input.## The prompt (instructions, tools, reasoning effort, tool_examples, reward) is# fixed for the life of the conversation. Build the request dict once and send# it unchanged on every turn; only `input` grows.# ---------------------------------------------------------------------------def main() -> None: client = OpenAI(base_url=BASE_URL, api_key=os.environ["SID_API_KEY"]) cache = DocumentCache(range_mode="lenient") # this episode's id map, document cache and seen-ledger session_id = uuid.uuid4().hex # KV cache routing information: one session per episode # Caps how many queries your backend sees at once (1 makes the turn serial). pool = ThreadPoolExecutor(max_workers=MAX_PARALLEL_FETCHES) # One request shape, only update the history at each turn. request = dict( model=MODEL, max_output_tokens=4096, store=False, include=["reasoning.encrypted_content"], instructions=INSTRUCTIONS, tools=TOOLS, # Reasoning effort, tool_examples, and the target reward are managed server-side reasoning={"effort": REASONING_EFFORT}, extra_body={"sid": {"tool_examples": TOOL_EXAMPLES, "reward": REWARD}}, extra_headers={"x-session-affinity": session_id}, ) # The whole conversation as Responses API input items, grown every turn. history: list[dict] = [{"role": "user", "content": QUESTION}] print(f"Q: {QUESTION}") response = client.responses.create(input=history, **request) for turn in range(1, MAX_TURNS + 1): print_model_output(turn, response) calls = [item for item in response.output if item.type == "function_call"] arguments = [tool_arguments(call) for call in calls] # Execute the model's tool calls in parallel fetches = [pool.submit(fetch_results, call.name, args) for call, args in zip(calls, arguments, strict=True)] # Compose observations serially, in call order: what earlier calls show, later calls mask. observations: list[str] = [] reported: list[str] | None = None for call, args, fetch in zip(calls, arguments, fetches, strict=True): try: if call.name in ("search", "text_search"): observation = render_results( cache, fetch.result(), args["query"], args.get("snippet_size", SNIPPET_SIZE_DEFAULT), ) elif call.name == "read": observation = read(cache, args.get("id")) elif call.name == "report_helpful_ids": if len(calls) > 1: raise ToolError("Combining report tool call with other tool calls is not allowed.") if turn == 1: raise ToolError("Cannot report in the first turn. Must make a search first.") reported = report(cache, args.get("ids")) observation = "reported ids submitted: " + ",".join(reported) else: observation = f"Tool '{call.name}' is not a valid tool. Available tools: {list(TOOLS_BY_NAME)}" except ToolError as e: # A recoverable mistake made by the model. observation = f"Error executing {call.name}: {e}" observations.append(observation) print(f" [{call.name}] -> {observation}") # verbatim: what you see is what the model receives # Termination conditions if reported is not None: print_report(cache, reported) return if turn == MAX_TURNS: print(f"\n[max turns reached: {MAX_TURNS}: episode ends without a report]") return if response.status == "incomplete": print("\n[generation hit max_output_tokens; ending the episode]") return # Non-terminal turn: append the model's output verbatim (the reasoning # items carry their encrypted_content), then feed back the tool outputs # (paired by call_id) plus the turn budget indicator as a user message. history += [item.model_dump(mode="json", exclude_none=True) for item in response.output] history += [ {"type": "function_call_output", "call_id": call.call_id, "output": observation} for call, observation in zip(calls, observations, strict=True) ] remaining = MAX_TURNS - turn messages = [f"You have {remaining} {'turn' if remaining == 1 else 'turns'} remaining out of {MAX_TURNS} total."] if turn == MAX_TURNS - 1: messages.append( f"This is your last turn: you must call {REPORT_TOOL['name']} now. " f"Anything you have found but not reported is lost when this turn ends." ) history += [{"role": "user", "content": message} for message in messages] for message in messages: print(f"[user] {message}") response = client.responses.create(input=history, **request)def tool_arguments(call) -> dict: """Parse a call's JSON arguments: keep only the parameters the tool's schema declares, and drop the ""/null "defaults" LLMs like to send. Anything else the model got wrong comes back as a ToolError from the fetch or the tool.""" try: arguments = json.loads(call.arguments) except json.JSONDecodeError: return {} schema = TOOLS_BY_NAME.get(call.name) if not isinstance(arguments, dict) or schema is None: return {} properties = schema["parameters"]["properties"] return {key: value for key, value in arguments.items() if key in properties and value not in ("", None)}def print_model_output(turn: int, response) -> None: """The model's side of one turn: its thinking, any prose, and its calls.""" print(f"\n--- turn {turn} (status={response.status}) " + "-" * 40) for item in response.output: if item.type == "reasoning": for part in getattr(item, "content", None) or []: print(f"[thinking] {part.text}") elif item.type == "message": print(f"[assistant] {response.output_text}") elif item.type == "function_call": print(f">>> {item.name}({item.arguments})")def print_report(cache: DocumentCache, reported_ids: list[str]) -> None: """The episode's answer: your database ids, most helpful first.""" print(f"\nReported for: {QUESTION}") for data_id in reported_ids: document = cache.get_document(data_id) print(f" {data_id} {document.get('title', '')}")if __name__ == "__main__": main()
/** SID-2 preview. * Implement vectorSearch and keywordSearch against your search backend. * Set ID_FIELD and SNIPPET_FIELD, then customize INSTRUCTIONS, QUESTION, * and TOOL_EXAMPLES to fit your data. Update tool descriptions and examples * together. The report_helpful_ids tool is fixed and cannot be modified. * * npm install @sid-ai/sid-sdk openai * npm install --save-dev tsx * SID_API_KEY=your-api-key npx tsx reference_implementation.ts */import { randomUUID } from "node:crypto";import { pathToFileURL } from "node:url";import OpenAI from "openai";import type { ResponseCreateParamsNonStreaming, ResponseInput, ResponseFunctionToolCall,} from "openai/resources/responses/responses";import { DocumentCache, InvalidCharacterRange, parseRenderedModelFacingId, SNIPPET_SIZE_DEFAULT,} from "@sid-ai/sid-sdk";class ToolError extends Error {}type BackendDocument = Record<string, unknown>;type Arguments = Record<string, unknown>;const BASE_URL = process.env.SID_BASE_URL ?? "https://api.sid-1.com/v1";const MODEL = "sid-2";const QUESTION = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?";const MAX_TURNS = 10;const REWARD = "f1";const REASONING_EFFORT = "low";const DEFAULT_LIMIT = 5;const MAX_PARALLEL_FETCHES = 8;const ID_FIELD = "id";const SNIPPET_FIELD = "content";const INSTRUCTIONS = `You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called.`;const SEARCH_TOOL = { "type": "function", "name": "search", "description": "Semantic (vector) search over the dataset. Best for conceptual or paraphrased queries: a long, descriptive query of what you want beats a few keywords; use text_search when you know the distinctive words a match would contain. Result ids may be ranged references 'id#start:end' naming the exact character span of the document shown; stretches you have already been shown are replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the document's total character count.", "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "The query to search the dataset with" }, "limit": { "type": "integer", "description": "The number of results to return" }, "snippet_size": { "type": "integer", "description": "Words shown per result snippet (default 50). Must be a positive integer. Pass a larger value to see more of each document." } }, "required": [ "query" ] }};const TEXT_SEARCH_TOOL = { "type": "function", "name": "text_search", "description": "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you want matched as plain words separated by spaces: the distinctive content words a matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact character span of the document shown; stretches you have already been shown are replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the document's total character count.", "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "Plain space-separated content words to match (BM25)." }, "limit": { "type": "integer", "description": "The number of results to return" }, "snippet_size": { "type": "integer", "description": "Words shown per result snippet (default 50). Must be a positive integer. Pass a larger value to see more of each document." } }, "required": [ "query" ] }};const READ_TOOL = { "type": "function", "name": "read", "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.", "parameters": { "type": "object", "properties": { "id": { "type": "string", "description": "The id of the document to read. Also accepts a ranged reference 'id#start:end' (character offsets, e.g. '123456#800:1600') to read only that part of the document's text; a range running past the end of the document is trimmed to the part that overlaps it." } }, "required": [ "id" ] }};const REPORT_TOOL = { "type": "function", "name": "report_helpful_ids", "description": "Report the most helpful document ids from the search results. This is required in the last step. Each id may be reported at most once; repeated ids are penalized. Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a ranged reference counts as its whole document, so there is no need to strip the '#start:end' part.", "parameters": { "type": "object", "properties": { "ids": { "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": { "type": "string" } } }, "required": [ "ids" ] }};const TOOL_EXAMPLES = [ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": { "query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5 } }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": { "query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5 } }, { "caption": "Example: reading one character range:", "name": "read", "arguments": { "id": "placeholder_1#800:1600" } }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": { "ids": [ "placeholder_1", "placeholder_2", "placeholder_3" ] } }];const TOOLS = [SEARCH_TOOL, TEXT_SEARCH_TOOL, READ_TOOL, REPORT_TOOL];const TOOLS_BY_NAME = new Map(TOOLS.map(tool => [tool.name, tool]));// Return full chunks with stable database IDs, best matches first.// Every field except ID_FIELD is shown to the model, so leave private fields out.async function vectorSearch(query: string, limit: number): Promise<BackendDocument[]> { throw new Error("Implement vectorSearch against your document store");}async function keywordSearch(query: string, limit: number): Promise<BackendDocument[]> { throw new Error("Implement keywordSearch against your document store");}async function fetchResults(name: string, args: Arguments): Promise<BackendDocument[] | null> { if (name !== "search" && name !== "text_search") return null; if (typeof args.query !== "string" || !args.query) { throw new ToolError("Search requires a non-empty query. Use read to fetch a specific document ID."); } for (const key of ["limit", "snippet_size"]) { if (args[key] !== undefined && !Number.isSafeInteger(args[key])) { throw new ToolError(`${key} must be an integer`); } } const snippetSize = (args.snippet_size ?? SNIPPET_SIZE_DEFAULT) as number; if (snippetSize <= 0) throw new ToolError("snippet_size must be a positive integer"); const limit = (args.limit ?? DEFAULT_LIMIT) as number; return name === "search" ? vectorSearch(args.query, limit) : keywordSearch(args.query, limit);}function renderResults(cache: DocumentCache, results: BackendDocument[], query: string, snippetSize: number): string { const rendered: string[] = []; for (const result of results) { const document = { ...result }; const dataId = document[ID_FIELD]; delete document[ID_FIELD]; if (typeof dataId !== "string" || !dataId) { throw new Error(`Backend document must carry a non-empty string ${ID_FIELD}`); } if (typeof document[SNIPPET_FIELD] !== "string" || !document[SNIPPET_FIELD]) { throw new Error(`Backend document ${dataId} must carry non-empty ${SNIPPET_FIELD} text`); } cache.addDocument(dataId, document); const view = cache.applySnippet(dataId, { snippetField: SNIPPET_FIELD, query, snippetSize }); rendered.push(view.renderXml()); cache.updateSeen(view); } return rendered.join("\n") || "No results found. Try a different query.";}function parseReference(reference: unknown) { if (typeof reference !== "string") throw new ToolError("Document reference must be a string"); try { return parseRenderedModelFacingId(reference); } catch (error) { if (error instanceof TypeError || error instanceof RangeError) throw new ToolError(error.message); throw error; }}function read(cache: DocumentCache, reference: unknown): string { const [modelId, range] = parseReference(reference); if (!cache.containsModelFacingId(modelId)) { throw new ToolError("Unknown document ID. Use search or text_search first."); } const dataId = cache.toDataId(modelId); let span = range ?? undefined; if (span) { try { span = cache.resolveCharRange(dataId, SNIPPET_FIELD, span); } catch (error) { if (error instanceof InvalidCharacterRange) throw new ToolError(error.message); throw error; } } // Count Unicode code points, matching the SDK's character offsets. if (span) { span = [span[0], Math.min(span[1], span[0] + 10_000)]; } else if (Array.from(cache.getDocument(dataId)[SNIPPET_FIELD] as string).length > 10_000) { span = [0, 10_000]; } const view = cache.getSingleSpanDocumentView(dataId, { snippetField: SNIPPET_FIELD, snippetDisplaySpan: span, }); const output = view.renderXml(); cache.updateSeen(view); return output;}function report(cache: DocumentCache, ids: unknown): string[] { if (!Array.isArray(ids) || !ids.every(id => typeof id === "string")) { throw new ToolError("ids must be a list of strings"); } if (!ids.length) throw new ToolError("ids list cannot be empty"); const modelIds = [...new Set(ids.map(id => parseReference(id)[0]))]; return modelIds .filter(id => cache.containsModelFacingId(id)) .map(id => cache.toDataId(id));}function toolArguments(call: ResponseFunctionToolCall): Arguments { let args: unknown; try { args = JSON.parse(call.arguments); } catch { return {}; } const tool = TOOLS_BY_NAME.get(call.name); if (!tool || !args || typeof args !== "object" || Array.isArray(args)) return {}; return Object.fromEntries(Object.entries(args).filter(([key, value]) => Object.hasOwn(tool.parameters.properties, key) && value !== "" && value !== null ));}// Limit backend concurrency, preserving the model's call order in the results.async function fetchInParallel(calls: ResponseFunctionToolCall[], args: Arguments[]) { const results: PromiseSettledResult<BackendDocument[] | null>[] = new Array(calls.length); let next = 0; async function worker() { while (next < calls.length) { const index = next++; try { results[index] = { status: "fulfilled", value: await fetchResults(calls[index].name, args[index]) }; } catch (reason) { results[index] = { status: "rejected", reason }; } } } await Promise.all(Array.from({ length: Math.min(MAX_PARALLEL_FETCHES, calls.length) }, worker)); return results;}async function main(): Promise<void> { const client = new OpenAI({ baseURL: BASE_URL, apiKey: process.env.SID_API_KEY }); const cache = new DocumentCache({ rangeMode: "lenient" }); const sessionId = randomUUID(); // One session per episode. const request: Omit<ResponseCreateParamsNonStreaming, "input"> & { sid: object } = { model: MODEL, max_output_tokens: 4096, store: false, include: ["reasoning.encrypted_content"], instructions: INSTRUCTIONS, tools: TOOLS as ResponseCreateParamsNonStreaming["tools"], reasoning: { effort: REASONING_EFFORT }, sid: { tool_examples: TOOL_EXAMPLES, reward: REWARD }, }; const options = { headers: { "x-session-affinity": sessionId } }; const history: ResponseInput = [{ role: "user", content: QUESTION }]; console.log(`Q: ${QUESTION}`); for (let turn = 1; turn <= MAX_TURNS; turn++) { const response = await client.responses.create({ ...request, input: history }, options); console.log(`Turn ${turn} (${response.status})`, response.output); const calls = response.output.filter((item): item is ResponseFunctionToolCall => item.type === "function_call"); const args = calls.map(toolArguments); const fetches = await fetchInParallel(calls, args); const observations: string[] = []; let reported: string[] | undefined; // Render serially: what earlier calls show, later calls mask. for (let i = 0; i < calls.length; i++) { const call = calls[i]; let observation: string; try { const fetch = fetches[i]; if (fetch.status === "rejected") throw fetch.reason; if (call.name === "search" || call.name === "text_search") { observation = renderResults(cache, fetch.value!, args[i].query as string, (args[i].snippet_size ?? SNIPPET_SIZE_DEFAULT) as number); } else if (call.name === "read") { observation = read(cache, args[i].id); } else if (call.name === "report_helpful_ids") { if (calls.length > 1) throw new ToolError("Combining report tool call with other tool calls is not allowed."); if (turn === 1) throw new ToolError("Cannot report in the first turn. Must make a search first."); reported = report(cache, args[i].ids); observation = `Reported IDs submitted: ${reported.join(",")}`; } else { observation = `Tool '${call.name}' is not valid. Available tools: ${[...TOOLS_BY_NAME.keys()].join(", ")}`; } } catch (error) { if (!(error instanceof ToolError)) throw error; observation = `Error executing ${call.name}: ${error.message}`; } observations.push(observation); console.log(`[${call.name}] ${observation}`); } if (reported !== undefined) { console.log(`Reported for: ${QUESTION}`); for (const id of reported) console.log(id, cache.getDocument(id).title ?? ""); return; } if (turn === MAX_TURNS || response.status === "incomplete") { console.log("Episode ends without a report."); return; } // Echo output items, including encrypted reasoning, without modifications. history.push(...response.output as ResponseInput); history.push(...calls.map((call, i) => ({ type: "function_call_output" as const, call_id: call.call_id, output: observations[i], }))); const remaining = MAX_TURNS - turn; history.push({ role: "user", content: `You have ${remaining} ${remaining === 1 ? "turn" : "turns"} remaining out of ${MAX_TURNS} total.` }); if (remaining === 1) { history.push({ role: "user", content: `This is your last turn: you must call ${REPORT_TOOL.name} now. Anything you have found but not reported is lost when this turn ends.` }); } }}if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) { main().catch(error => { console.error(error); process.exitCode = 1; });}
go mod init sid-examplego get github.com/sidhq/sid-go github.com/openai/openai-go/v3SID_API_KEY=your-api-key go run reference_implementation.go
reference_implementation.go
// SID-2 preview.//// Before using this template: implement the vectorSearch and keywordSearch// backend functions, set idField and snippetField to your document's fields,// and customize the instructions, question, and toolExamples to fit your data.// You can also define additional tools, or extend existing tools with more// parameters (e.g. a filter). Just make sure to always update the tool's// description and usage example. The report_helpful_ids tool is fixed and// cannot be changed or modified.//// Run://// go mod init sid-example && go get github.com/sidhq/sid-go github.com/openai/openai-go/v3// SID_API_KEY=your-api-key go run reference_implementation.gopackage mainimport ( "context" "crypto/rand" "encoding/hex" "encoding/json" "errors" "fmt" "log" "maps" "os" "slices" "strings" "sync" "unicode/utf8" "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" "github.com/openai/openai-go/v3/packages/param" "github.com/openai/openai-go/v3/responses" "github.com/openai/openai-go/v3/shared" sid "github.com/sidhq/sid-go")// ToolError is a mistake the *model* made (bad id, malformed ref, invalid// arguments). These errors are recoverable and passed back to the model as a// message.type ToolError struct{ message string }func (e *ToolError) Error() string { return e.message }func toolErrorf(format string, args ...any) error { return &ToolError{fmt.Sprintf(format, args...)}}// ---------------------------------------------------------------------------// The episode. TODO: your question. The model is served by the first-party// SID API, under one SID_API_KEY.// ---------------------------------------------------------------------------const question = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?"const ( model = "sid-2" maxTurns = 10 reward = "f1" // the trained scoring prompt: "f1" (few, precise ids) | "ndcg" (many, best first) reasoningEffort = "low" // "low" | "medium" | "high" defaultLimit = 5 // results per search when the model does not ask for a number maxReadChars = 10_000 // document text returned per read call maxParallelFetches = 8 // backend queries in flight at once, across a turn's search calls)func baseURL() string { if url := os.Getenv("SID_BASE_URL"); url != "" { return strings.TrimRight(url, "/") } return "https://api.sid-1.com/v1"}// ---------------------------------------------------------------------------// TODO: your document backend. The tools below call these functions;// implement them against your database. A document is a map with://// idField : your document id (any stable string; never shown to the model)// snippetField : the document's canonical text. Return the identical string// for the same id everywhere: the seen-ledger's character// offsets index it.// extras : (optional) any other fields, such as a "title", rendered as// display attributes. Every cached field is shown to the// model (idField is removed before caching), so leave// anything private out of your documents.// ---------------------------------------------------------------------------const ( idField = "id" snippetField = "content")// vectorSearch is your semantic (embedding) search: the limit most relevant// documents for query, best first.func vectorSearch(ctx context.Context, query string, limit int) ([]sid.Document, error) { return nil, errors.New("TODO: implement vectorSearch against your document store")}// keywordSearch is your full-text keyword (BM25) search: the limit best// matching documents for query.func keywordSearch(ctx context.Context, query string, limit int) ([]sid.Document, error) { return nil, errors.New("TODO: implement keywordSearch against your document store")}// ---------------------------------------------------------------------------// The system prompt: one sentence naming your subject matter, then the// search -> repeat -> report scaffolding. TODO: edit the first paragraph to// describe your data and goal; keep the steps.// ---------------------------------------------------------------------------const instructions = `You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called.`// ---------------------------------------------------------------------------// The tool schemas. Each tool should be demonstrated in toolExamples.// ---------------------------------------------------------------------------type object = map[string]anyvar snippetSizeDescription = fmt.Sprintf( "Words shown per result snippet (default %d). Must be a positive integer. "+ "Pass a larger value to see more of each document.", sid.SnippetSizeDefault,)var searchTool = object{ "type": "function", "name": "search", "description": "Semantic (vector) search over the dataset. Best for conceptual or paraphrased " + "queries: a long, descriptive query of what you want beats a few keywords; use " + "text_search when you know the distinctive words a match would contain. Result " + "ids may be ranged references 'id#start:end' naming the exact character span of " + "the document shown; stretches you have already been shown are replaced by " + "'[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the " + "document's total character count.", "parameters": object{ "type": "object", "properties": object{ "query": object{"type": "string", "description": "The query to search the dataset with"}, "limit": object{"type": "integer", "description": "The number of results to return"}, "snippet_size": object{"type": "integer", "description": snippetSizeDescription}, }, "required": []string{"query"}, },}var textSearchTool = object{ "type": "function", "name": "text_search", "description": "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you " + "want matched as plain words separated by spaces: the distinctive content words a " + "matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact " + "character span of the document shown; stretches you have already been shown are " + "replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length " + "gives the document's total character count.", "parameters": object{ "type": "object", "properties": object{ "query": object{"type": "string", "description": "Plain space-separated content words to match (BM25)."}, "limit": object{"type": "integer", "description": "The number of results to return"}, "snippet_size": object{"type": "integer", "description": snippetSizeDescription}, }, "required": []string{"query"}, },}var readTool = object{ "type": "function", "name": "read", "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. " + "A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.", "parameters": object{ "type": "object", "properties": object{ "id": object{ "type": "string", "description": "The id of the document to read. Also accepts a ranged reference " + "'id#start:end' (character offsets, e.g. '123456#800:1600') to read only " + "that part of the document's text; a range running past the end of the " + "document is trimmed to the part that overlaps it.", }, }, "required": []string{"id"}, },}// reportTool is the fixed terminal tool: do not change its name, description// or parameters.var reportTool = object{ "type": "function", "name": "report_helpful_ids", "description": "Report the most helpful document ids from the search results. This is required in " + "the last step. Each id may be reported at most once; repeated ids are penalized. " + "Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a " + "ranged reference counts as its whole document, so there is no need to strip the " + "'#start:end' part.", "parameters": object{ "type": "object", "properties": object{ "ids": object{ "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": object{"type": "string"}, }, }, "required": []string{"ids"}, },}var tools = []object{searchTool, textSearchTool, readTool, reportTool}func toolProperties(name string) object { for _, tool := range tools { if tool["name"] == name { return tool["parameters"].(object)["properties"].(object) } } return nil}func toolNames() []string { names := make([]string, len(tools)) for i, tool := range tools { names[i] = tool["name"].(string) } return names}// toolExamples are few-shot calls, sent on every request via the sid// extension. TODO: write search examples that fit your data; keep the read and// report examples as they are.var toolExamples = []object{ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": object{"query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5}, }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": object{"query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5}, }, { "caption": "Example: reading one character range:", "name": "read", "arguments": object{"id": "placeholder_1#800:1600"}, }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": object{"ids": []string{"placeholder_1", "placeholder_2", "placeholder_3"}}, },}// ---------------------------------------------------------------------------// The tools, split in two: fetchResults is the slow backend call and runs in// parallel; everything below it touches the cache and runs serially, in call// order.// ---------------------------------------------------------------------------// fetchResults is the backend call behind one search tool call. read and the// report have nothing to fetch (reads are served from the document cache), so// they get nil.func fetchResults(ctx context.Context, name string, arguments object) ([]sid.Document, error) { if name != "search" && name != "text_search" { return nil, nil } // Catch mistyped arguments as a ToolError that the model can recover from. properties := toolProperties(name) for key, value := range arguments { expected := properties[key].(object)["type"] _, isString := value.(string) if expected == "string" && !isString || expected == "integer" && !isInteger(value) { return nil, toolErrorf("Invalid %s arguments: %s must be of type %s, got %T", name, key, expected, value) } } query, _ := arguments["query"].(string) // toolArguments strips ""; treat missing as empty if query == "" { return nil, toolErrorf("search requires a non-empty `query` to rank documents against. " + "To fetch a specific document by id, use the `read` tool. " + "To rank candidates against the question, call `search` with a meaningful query.") } snippetSize := intArgument(arguments, "snippet_size", sid.SnippetSizeDefault) if snippetSize <= 0 { return nil, toolErrorf("Invalid %s arguments: snippet_size must be a positive integer, got %d", name, snippetSize) } limit := intArgument(arguments, "limit", defaultLimit) if name == "search" { return vectorSearch(ctx, query, limit) } return keywordSearch(ctx, query, limit)}// renderResults renders one <doc> per result, in backend order. Shared by// both search tools; only their fetch differs.func renderResults(cache *sid.DocumentCache, results []sid.Document, query string, snippetSize int) (string, error) { rendered := make([]string, 0, len(results)) for _, result := range results { dataID, ok := result[idField].(string) if !ok || dataID == "" { return "", fmt.Errorf("backend document must carry a non-empty string %q: %v", idField, result) } // Keep backend results intact when searches reuse the same maps. document := maps.Clone(result) // Removed, not read: every cached field renders, and your database ids must not. delete(document, idField) if content, ok := document[snippetField].(string); !ok || content == "" { // Not a model mistake: the seen-ledger's offsets index this string. return "", fmt.Errorf("backend document %q must carry non-empty %q text", dataID, snippetField) } // Searches routinely return the same document, and the cache is add-once. if _, err := cache.AddDocument(dataID, document); err != nil { return "", err } // BM25 snippet, with already-shown stretches collapsed to [seen: ...]. view, err := cache.ApplySnippet(dataID, sid.ApplySnippetOptions{ SnippetField: snippetField, Query: query, SnippetSize: snippetSize, }) if err != nil { return "", err } xml, err := view.RenderXML() if err != nil { return "", err } rendered = append(rendered, xml) if err := cache.UpdateSeen(view); err != nil { // later views of this span collapse to [seen: ...] return "", err } } if len(rendered) == 0 { return "No results found. Try a different query.", nil } return strings.Join(rendered, "\n"), nil}// parseReference splits the model's ref ('R' or 'R#start:end') into id and// optional char range. The SDK parser already writes a descriptive message.func parseReference(reference any) (string, *sid.CharacterRange, error) { value, ok := reference.(string) if !ok { return "", nil, toolErrorf("id must be a string, got %T", reference) } modelID, charRange, err := sid.ParseRenderedModelFacingID(value) if err != nil { return "", nil, &ToolError{err.Error()} } return modelID, charRange, nil}// read fetches one document (or a character range of it) by the id a search// returned. Explicit reads are not masked: the model asked for this span, so// it renders up to maxReadChars characters. Model mistakes (bad ref grammar,// an unknown or hallucinated id, an unusable range) return a ToolError.func read(cache *sid.DocumentCache, reference any) (string, error) { modelID, charRange, err := parseReference(reference) // ToolError on bad grammar if err != nil { return "", err } if !cache.ContainsModelFacingID(modelID) { return "", toolErrorf("Invalid document id '%s'. "+ "You can only read documents returned by search. "+ "Use 'search' or 'text_search' first.", modelID) } dataID, err := cache.ToDataID(modelID) // the one conversion back into your ids if err != nil { return "", err } if charRange != nil { // Lenient mode trims partially overlapping ranges to the document. // Invalid ranges become recoverable tool errors. resolved, err := cache.ResolveCharRange(dataID, snippetField, *charRange) var invalid *sid.InvalidCharacterRange if errors.As(err, &invalid) { return "", toolErrorf("%s. Re-read '%s' without a range to read from the "+ "start, or pass a range inside 0:doc_length.", invalid, modelID) } else if err != nil { return "", err } charRange = &resolved } // Cap document text per read; further ranged reads can continue from here. // Count Unicode code points, matching the SDK's character offsets. if charRange != nil { charRange = sid.RangeValue(charRange.Start(), min(charRange.End(), charRange.Start()+maxReadChars)) } else { document, err := cache.GetDocument(dataID) if err != nil { return "", err } if content, _ := document[snippetField].(string); utf8.RuneCountInString(content) > maxReadChars { charRange = sid.RangeValue(0, maxReadChars) } } view, err := cache.GetSingleSpanDocumentView(dataID, sid.SingleSpanOptions{ SnippetField: snippetField, SnippetDisplaySpan: charRange, }) if err != nil { return "", err } rendered, err := view.RenderXML() if err != nil { return "", err } return rendered, cache.UpdateSeen(view)}// report is the terminal report: ids are doc-granular (a ranged ref denotes// its document), deduplicated preserving order, and translated back to your// data ids. IDs outside the model-facing ID map are dropped.func report(cache *sid.DocumentCache, ids any) ([]string, error) { values, ok := ids.([]any) if !ok { return nil, toolErrorf("ids must be a list of strings") } if len(values) == 0 { return nil, toolErrorf("ids list cannot be empty") } var reported []string for _, reference := range values { if _, ok := reference.(string); !ok { return nil, toolErrorf("ids must be a list of strings") } modelID, _, err := parseReference(reference) // ToolError on bad grammar if err != nil { return nil, err } if !slices.Contains(reported, modelID) { reported = append(reported, modelID) } } dataIDs := make([]string, 0, len(reported)) for _, modelID := range reported { if dataID, err := cache.ToDataID(modelID); err == nil { dataIDs = append(dataIDs, dataID) } } return dataIDs, nil}// ---------------------------------------------------------------------------// The agent loop: one Responses API call per turn, tool outputs paired back// by call_id. The episode ends when the model reports.//// Requests are stateless (store=false): nothing is retained server-side, so// the full conversation is resent on every call. The model's reasoning comes// back as an opaque encrypted_content blob on each reasoning item, and// carries over by echoing the output items back verbatim in the next input.//// The prompt (instructions, tools, reasoning effort, tool_examples, reward) is// fixed for the life of the conversation. Build the request once and send it// unchanged on every turn; only input grows.// ---------------------------------------------------------------------------type fetched struct { results []sid.Document err error}func main() { ctx := context.Background() cache, err := sid.NewDocumentCache(sid.DocumentCacheOptions{RangeMode: sid.RangeModeLenient}) // this episode's id map, document cache and seen-ledger if err != nil { log.Fatal(err) } apiKey := os.Getenv("SID_API_KEY") if apiKey == "" { log.Fatal("SID_API_KEY is not set") } client := openai.NewClient(option.WithBaseURL(baseURL()), option.WithAPIKey(apiKey)) // One request shape, only update the history at each turn. request := responses.ResponseNewParams{ Model: model, MaxOutputTokens: openai.Int(4096), Store: openai.Bool(false), Include: []responses.ResponseIncludable{responses.ResponseIncludableReasoningEncryptedContent}, Instructions: openai.String(instructions), // Reasoning effort, tool_examples, and the target reward are managed server-side Reasoning: shared.ReasoningParam{Effort: reasoningEffort}, } options := []option.RequestOption{ // The tools and the sid extension are plain JSON, set on the request body. option.WithJSONSet("tools", tools), option.WithJSONSet("sid", object{"tool_examples": toolExamples, "reward": reward}), option.WithHeader("x-session-affinity", newSessionID()), // KV cache routing: one session per episode } // The whole conversation as Responses API input items, grown every turn. history := responses.ResponseInputParam{ responses.ResponseInputItemParamOfMessage(question, responses.EasyInputMessageRoleUser), } create := func() (*responses.Response, error) { request.Input = responses.ResponseNewParamsInputUnion{OfInputItemList: history} return client.Responses.New(ctx, request, options...) } fmt.Printf("Q: %s\n", question) resp, err := create() if err != nil { log.Fatal(err) } for turn := 1; turn <= maxTurns; turn++ { printModelOutput(turn, resp) var calls []responses.ResponseFunctionToolCall for _, item := range resp.Output { if item.Type == "function_call" { calls = append(calls, item.AsFunctionCall()) } } arguments := make([]object, len(calls)) for i, call := range calls { arguments[i] = toolArguments(call) } // Execute the model's tool calls in parallel; maxParallelFetches caps // how many queries your backend sees at once (1 makes the turn serial). fetches := make([]fetched, len(calls)) semaphore := make(chan struct{}, maxParallelFetches) var wg sync.WaitGroup for i, call := range calls { wg.Add(1) go func() { defer wg.Done() semaphore <- struct{}{} defer func() { <-semaphore }() results, err := fetchResults(ctx, call.Name, arguments[i]) fetches[i] = fetched{results, err} }() } wg.Wait() // Compose observations serially, in call order: what earlier calls show, later calls mask. observations := make([]string, len(calls)) var reported []string for i, call := range calls { observation, err := runTool(cache, turn, len(calls), call.Name, arguments[i], fetches[i], &reported) var toolError *ToolError if errors.As(err, &toolError) { // A recoverable mistake made by the model. observation = fmt.Sprintf("Error executing %s: %s", call.Name, toolError) } else if err != nil { log.Fatal(err) } observations[i] = observation fmt.Printf(" [%s] -> %s\n", call.Name, observation) // verbatim: what you see is what the model receives } // Termination conditions if reported != nil { printReport(cache, reported) return } if turn == maxTurns { fmt.Printf("\n[max turns reached: %d: episode ends without a report]\n", maxTurns) return } if resp.Status == responses.ResponseStatusIncomplete { fmt.Println("\n[generation hit max_output_tokens; ending the episode]") return } // Non-terminal turn: append the model's output verbatim (the reasoning // items carry their encrypted_content), then feed back the tool outputs // (paired by call_id) plus the turn budget indicator as a user message. for _, item := range resp.Output { history = append(history, param.Override[responses.ResponseInputItemUnionParam](json.RawMessage(item.RawJSON()))) } for i, call := range calls { output := responses.ResponseInputItemParamOfFunctionCallOutput(observations[i]) output.OfFunctionCallOutput.CallID = openai.String(call.CallID) history = append(history, output) } remaining := maxTurns - turn unit := "turns" if remaining == 1 { unit = "turn" } messages := []string{fmt.Sprintf("You have %d %s remaining out of %d total.", remaining, unit, maxTurns)} if turn == maxTurns-1 { messages = append(messages, fmt.Sprintf( "This is your last turn: you must call %s now. "+ "Anything you have found but not reported is lost when this turn ends.", reportTool["name"])) } for _, message := range messages { history = append(history, responses.ResponseInputItemParamOfMessage(message, responses.EasyInputMessageRoleUser)) fmt.Printf("[user] %s\n", message) } if resp, err = create(); err != nil { log.Fatal(err) } }}// runTool dispatches one call. A ToolError is a recoverable model mistake;// any other error is a bug or a backend failure and ends the episode.func runTool( cache *sid.DocumentCache, turn, callCount int, name string, arguments object, fetch fetched, reported *[]string,) (string, error) { if fetch.err != nil { return "", fetch.err } switch name { case "search", "text_search": return renderResults(cache, fetch.results, arguments["query"].(string), intArgument(arguments, "snippet_size", sid.SnippetSizeDefault)) case "read": return read(cache, arguments["id"]) case "report_helpful_ids": if callCount > 1 { return "", toolErrorf("Combining report tool call with other tool calls is not allowed.") } if turn == 1 { return "", toolErrorf("Cannot report in the first turn. Must make a search first.") } ids, err := report(cache, arguments["ids"]) if err != nil { return "", err } *reported = ids return "reported ids submitted: " + strings.Join(ids, ","), nil default: return fmt.Sprintf("Tool '%s' is not a valid tool. Available tools: %v", name, toolNames()), nil }}// toolArguments parses a call's JSON arguments: keep only the parameters the// tool's schema declares, and drop the ""/null "defaults" LLMs like to send.// Anything else the model got wrong comes back as a ToolError from the fetch// or the tool.func toolArguments(call responses.ResponseFunctionToolCall) object { var arguments object properties := toolProperties(call.Name) if json.Unmarshal([]byte(call.Arguments), &arguments) != nil || properties == nil { return object{} } for key, value := range arguments { if _, declared := properties[key]; !declared || value == nil || value == "" { delete(arguments, key) } } return arguments}// isInteger reports whether a decoded JSON value is a whole number.func isInteger(value any) bool { number, ok := value.(float64) return ok && number == float64(int(number))}func intArgument(arguments object, key string, fallback int) int { if value, ok := arguments[key].(float64); ok { return int(value) } return fallback}func newSessionID() string { var id [16]byte if _, err := rand.Read(id[:]); err != nil { log.Fatal(err) } return hex.EncodeToString(id[:])}// printModelOutput prints the model's side of one turn: its thinking, any// prose, and its calls.func printModelOutput(turn int, resp *responses.Response) { fmt.Printf("\n--- turn %d (status=%s) %s\n", turn, resp.Status, strings.Repeat("-", 40)) for _, item := range resp.Output { switch item.Type { case "reasoning": for _, part := range item.AsReasoning().Content { fmt.Printf("[thinking] %s\n", part.Text) } case "message": fmt.Printf("[assistant] %s\n", resp.OutputText()) case "function_call": fmt.Printf(">>> %s(%s)\n", item.Name, item.AsFunctionCall().Arguments) } }}// printReport prints the episode's answer: your database ids, most helpful// first.func printReport(cache *sid.DocumentCache, reported []string) { fmt.Printf("\nReported for: %s\n", question) for _, dataID := range reported { document, _ := cache.GetDocument(dataID) title, _ := document["title"].(string) fmt.Printf(" %s %s\n", dataID, title) }}
"""SID-2 preview on Baseten.Before using this template: implement the vector_search and keyword_search backendfunctions, set ID_FIELD and SNIPPET_FIELD to your document's fields, andcustomize the INSTRUCTIONS, QUESTION, and TOOL_EXAMPLES to fit your data. You canalso define additional tools, or extend existing tools with more parameters (e.g. afilter). Just make sure to always update the tool's description and usage example.The report_helpful_ids tool is fixed and cannot be changed or modified.Run: BASETEN_API_KEY=your-api-key uv run --with sid-sdk --with openai reference_implementation_baseten.py"""import jsonimport osimport uuidfrom concurrent.futures import ThreadPoolExecutorfrom openai import OpenAIfrom sid import DocumentCache, InvalidCharacterRange, parse_rendered_model_facing_idfrom sid.document_cache import SNIPPET_SIZE_DEFAULTclass ToolError(Exception): """A mistake the *model* made (bad id, malformed ref, invalid arguments), these errors are recoverable and passed back to the model as a message."""# ---------------------------------------------------------------------------# The episode. TODO: your question. The model is served by Baseten,# under one BASETEN_API_KEY.# ---------------------------------------------------------------------------QUESTION = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?"BASE_URL = os.environ.get("BASETEN_BASE_URL", "https://inference.baseten.co/v1")MODEL = "sid/sid-2-preview"MAX_TURNS = 10REWARD = "f1" # the trained scoring prompt: "f1" (few, precise ids) | "ndcg" (many, best first)REASONING_EFFORT = "low" # "low" | "medium" | "high"DEFAULT_LIMIT = 5 # results per search when the model does not ask for a numberMAX_PARALLEL_FETCHES = 8 # backend queries in flight at once, across a turn's search calls# ---------------------------------------------------------------------------# TODO: your document backend. The tools below call these functions;# implement them against your database. A document is a dict with:# ID_FIELD : your document id (any stable string; never shown to the model)# SNIPPET_FIELD: the document's canonical text. Return the identical string# for the same id everywhere: the seen-ledger's character# offsets index it.# extras : (optional) any other fields, such as a "title", rendered# as display attributes. Every cached field is shown to the# model (ID_FIELD is popped before caching), so leave# anything private out of your documents.# ---------------------------------------------------------------------------ID_FIELD = "id"SNIPPET_FIELD = "content"def vector_search(query: str, limit: int) -> list[dict]: """TODO: your semantic (embedding) search: the ``limit`` most relevant documents for ``query``, best first.""" raise NotImplementedError("implement vector_search against your document store")def keyword_search(query: str, limit: int) -> list[dict]: """TODO: your full-text keyword (BM25) search: the ``limit`` best matching documents for ``query``.""" raise NotImplementedError("implement keyword_search against your document store")# ---------------------------------------------------------------------------# The system prompt: one sentence naming your subject matter, then the# search -> repeat -> report scaffolding. TODO: edit the first paragraph to# describe your data and goal; keep the steps.# ---------------------------------------------------------------------------INSTRUCTIONS = """You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called."""# ---------------------------------------------------------------------------# The tool schemas. Each tool should be demonstrated in TOOL_EXAMPLES.# ---------------------------------------------------------------------------SNIPPET_SIZE_DESCRIPTION = ( f"Words shown per result snippet (default {SNIPPET_SIZE_DEFAULT}). Must be a positive integer. " f"Pass a larger value to see more of each document.")SEARCH_TOOL = { "type": "function", "name": "search", "description": ( "Semantic (vector) search over the dataset. Best for conceptual or paraphrased " "queries: a long, descriptive query of what you want beats a few keywords; use " "text_search when you know the distinctive words a match would contain. Result " "ids may be ranged references 'id#start:end' naming the exact character span of " "the document shown; stretches you have already been shown are replaced by " "'[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the " "document's total character count." ), "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "The query to search the dataset with"}, "limit": {"type": "integer", "description": "The number of results to return"}, "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION}, }, "required": ["query"], },}TEXT_SEARCH_TOOL = { "type": "function", "name": "text_search", "description": ( "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you " "want matched as plain words separated by spaces: the distinctive content words a " "matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact " "character span of the document shown; stretches you have already been shown are " "replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length " "gives the document's total character count." ), "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": ( "Plain space-separated content words to match (BM25)." ), }, "limit": {"type": "integer", "description": "The number of results to return"}, "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION}, }, "required": ["query"], },}READ_TOOL = { "type": "function", "name": "read", "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.", "parameters": { "type": "object", "properties": { "id": { "type": "string", "description": ( "The id of the document to read. Also accepts a ranged reference " "'id#start:end' (character offsets, e.g. '123456#800:1600') to read only " "that part of the document's text; a range running past the end of the " "document is trimmed to the part that overlaps it." ), }, }, "required": ["id"], },}# The fixed terminal tool: do not change its name, description or parameters.REPORT_TOOL = { "type": "function", "name": "report_helpful_ids", "description": ( "Report the most helpful document ids from the search results. This is required in " "the last step. Each id may be reported at most once; repeated ids are penalized. " "Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a " "ranged reference counts as its whole document, so there is no need to strip the " "'#start:end' part." ), "parameters": { "type": "object", "properties": { "ids": { "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": {"type": "string"}, }, }, "required": ["ids"], },}TOOLS = [SEARCH_TOOL, TEXT_SEARCH_TOOL, READ_TOOL, REPORT_TOOL]TOOLS_BY_NAME = {tool["name"]: tool for tool in TOOLS}# Few-shot calls, sent on every request via the sid extension. TODO: write search# examples that fit your data; keep the read and report examples as they areTOOL_EXAMPLES = [ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": { "query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5, }, }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": {"query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5}, }, { "caption": "Example: reading one character range:", "name": "read", "arguments": {"id": "placeholder_1#800:1600"}, }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": {"ids": ["placeholder_1", "placeholder_2", "placeholder_3"]}, },]# ---------------------------------------------------------------------------# The tools, split in two: fetch_results is the slow backend call and runs in# parallel; everything below it touches the cache and runs serially, in call# order.# ---------------------------------------------------------------------------def fetch_results(name: str, arguments: dict) -> list[dict] | None: """The backend call behind one search tool call. `read` and the report have nothing to fetch (reads are served from the document cache), so they get ``None``.""" if name not in ("search", "text_search"): return None # Catch mistyped arguments as a ToolError that the model can recover from. properties = TOOLS_BY_NAME[name]["parameters"]["properties"] for key, value in arguments.items(): expected = str if properties[key]["type"] == "string" else int if not isinstance(value, expected) or isinstance(value, bool): raise ToolError( f"Invalid {name} arguments: {key} must be of type {properties[key]['type']}, " f"got {type(value).__name__}" ) query = arguments.get("query", "") # tool_arguments strips ""; treat missing as empty if not query: raise ToolError( "search requires a non-empty `query` to rank documents against. " "To fetch a specific document by id, use the `read` tool. " "To rank candidates against the question, call `search` with a meaningful query." ) snippet_size = arguments.get("snippet_size", SNIPPET_SIZE_DEFAULT) if snippet_size <= 0: raise ToolError( f"Invalid {name} arguments: snippet_size must be a positive integer, got {snippet_size}" ) limit = arguments.get("limit", DEFAULT_LIMIT) return vector_search(query, limit) if name == "search" else keyword_search(query, limit)def render_results(cache: DocumentCache, results: list[dict], query: str, snippet_size: int) -> str: """One <doc> per result, in backend order. Shared by both search tools; only their fetch differs.""" rendered = [] for result in results: if not isinstance(result, dict) or not isinstance(result.get(ID_FIELD), str) or not result[ID_FIELD]: raise ValueError(f"backend document must carry a non-empty string {ID_FIELD!r}: {result!r}") # Keep backend results intact when searches reuse the same dictionaries. document = dict(result) # Popped, not read: every cached field renders, and your database ids must not. data_id = document.pop(ID_FIELD) content = document.get(SNIPPET_FIELD) if not isinstance(content, str) or not content: # Not a model mistake: the seen-ledger's offsets index this string. raise ValueError(f"backend document {data_id!r} must carry non-empty {SNIPPET_FIELD!r} text") # Searches routinely return the same document, and the cache is add-once. if data_id not in cache: cache.add_document(data_id, document) # BM25 snippet, with already-shown stretches collapsed to [seen: ...]. view = cache.apply_snippet( data_id, SNIPPET_FIELD, query, snippet_size=snippet_size, ) rendered.append(view.render_xml()) cache.update_seen(view) # later views of this span collapse to [seen: ...] return "\n".join(rendered) if rendered else "No results found. Try a different query."def parse_reference(reference) -> tuple[str, tuple[int, int] | None]: """The model's ref in (``'R'`` or ``'R#start:end'``), split into id and optional char range. The SDK parser raises ``ValueError`` on malformed shapes with a descriptive message already written.""" try: return parse_rendered_model_facing_id(reference) except ValueError as error: raise ToolError(str(error)) from Nonedef read(cache: DocumentCache, reference: str | None) -> str: """Fetch one document (or a character range of it) by the id a search returned. Explicit reads are not masked: the model asked for this span, so it renders up to 10,000 characters. Model mistakes (bad ref grammar, an unknown or hallucinated id, an unusable range) raise ToolError.""" model_facing_id, char_range = parse_reference(reference) # ToolError on bad grammar if not cache.contains_model_facing_id(model_facing_id): raise ToolError( f"Invalid document id '{model_facing_id}'. " f"You can only read documents returned by search. " f"Use 'search' or 'text_search' first." ) data_id = cache.to_data_id(model_facing_id) # the one conversion back into your ids if char_range is not None: # Lenient mode trims partially overlapping ranges to the document. # Invalid ranges become recoverable tool errors. try: char_range = cache.resolve_char_range(data_id, SNIPPET_FIELD, char_range) except InvalidCharacterRange as error: raise ToolError( f"{error}. Re-read '{model_facing_id}' without a range to read from the " f"start, or pass a range inside 0:doc_length." ) from None # Cap document text per read; further ranged reads can continue from here. if char_range is not None: start, end = char_range char_range = (start, min(end, start + 10_000)) elif len(cache.get_document(data_id)[SNIPPET_FIELD]) > 10_000: char_range = (0, 10_000) view = cache.get_single_span_document_view( data_id, SNIPPET_FIELD, snippet_display_span=char_range, ) rendered = view.render_xml() cache.update_seen(view) return rendereddef report(cache: DocumentCache, ids) -> list[str]: """The terminal report: ids are doc-granular (a ranged ref denotes its document), deduplicated preserving order, and translated back to your data_ids. IDs outside the model-facing ID map are dropped.""" if not isinstance(ids, list) or not all(isinstance(reference, str) for reference in ids): raise ToolError("ids must be a list of strings") if not ids: raise ToolError("ids list cannot be empty") reported: list[str] = [] for reference in ids: model_facing_id, _ = parse_reference(reference) # ToolError on bad grammar if model_facing_id not in reported: reported.append(model_facing_id) return [ cache.to_data_id(model_facing_id) for model_facing_id in reported if cache.contains_model_facing_id(model_facing_id) ]# ---------------------------------------------------------------------------# The agent loop: one Responses API call per turn, tool outputs paired back# by call_id. The episode ends when the model reports.## Requests are stateless (store=False): nothing is retained server-side, so# the full conversation is resent on every call. The model's reasoning comes# back as an opaque ``encrypted_content`` blob on each reasoning item, and# carries over by echoing the output items back verbatim in the next input.## The prompt (instructions, tools, reasoning effort, tool_examples, reward) is# fixed for the life of the conversation. Build the request dict once and send# it unchanged on every turn; only `input` grows.# ---------------------------------------------------------------------------def main() -> None: client = OpenAI(base_url=BASE_URL, api_key=os.environ["BASETEN_API_KEY"]) cache = DocumentCache(range_mode="lenient") # this episode's id map, document cache and seen-ledger session_id = uuid.uuid4().hex # KV cache routing information: one session per episode # Caps how many queries your backend sees at once (1 makes the turn serial). pool = ThreadPoolExecutor(max_workers=MAX_PARALLEL_FETCHES) # One request shape, only update the history at each turn. request = dict( model=MODEL, max_output_tokens=4096, store=False, include=["reasoning.encrypted_content"], instructions=INSTRUCTIONS, tools=TOOLS, # Reasoning effort, tool_examples, and the target reward are managed server-side reasoning={"effort": REASONING_EFFORT}, extra_body={"sid": {"tool_examples": TOOL_EXAMPLES, "reward": REWARD}}, extra_headers={"x-session-affinity": session_id}, ) # The whole conversation as Responses API input items, grown every turn. history: list[dict] = [{"role": "user", "content": QUESTION}] print(f"Q: {QUESTION}") response = client.responses.create(input=history, **request) for turn in range(1, MAX_TURNS + 1): print_model_output(turn, response) calls = [item for item in response.output if item.type == "function_call"] arguments = [tool_arguments(call) for call in calls] # Execute the model's tool calls in parallel fetches = [pool.submit(fetch_results, call.name, args) for call, args in zip(calls, arguments, strict=True)] # Compose observations serially, in call order: what earlier calls show, later calls mask. observations: list[str] = [] reported: list[str] | None = None for call, args, fetch in zip(calls, arguments, fetches, strict=True): try: if call.name in ("search", "text_search"): observation = render_results( cache, fetch.result(), args["query"], args.get("snippet_size", SNIPPET_SIZE_DEFAULT), ) elif call.name == "read": observation = read(cache, args.get("id")) elif call.name == "report_helpful_ids": if len(calls) > 1: raise ToolError("Combining report tool call with other tool calls is not allowed.") if turn == 1: raise ToolError("Cannot report in the first turn. Must make a search first.") reported = report(cache, args.get("ids")) observation = "reported ids submitted: " + ",".join(reported) else: observation = f"Tool '{call.name}' is not a valid tool. Available tools: {list(TOOLS_BY_NAME)}" except ToolError as e: # A recoverable mistake made by the model. observation = f"Error executing {call.name}: {e}" observations.append(observation) print(f" [{call.name}] -> {observation}") # verbatim: what you see is what the model receives # Termination conditions if reported is not None: print_report(cache, reported) return if turn == MAX_TURNS: print(f"\n[max turns reached: {MAX_TURNS}: episode ends without a report]") return if response.status == "incomplete": print("\n[generation hit max_output_tokens; ending the episode]") return # Non-terminal turn: append the model's output verbatim (the reasoning # items carry their encrypted_content), then feed back the tool outputs # (paired by call_id) plus the turn budget indicator as a user message. history += [item.model_dump(mode="json", exclude_none=True) for item in response.output] history += [ {"type": "function_call_output", "call_id": call.call_id, "output": observation} for call, observation in zip(calls, observations, strict=True) ] remaining = MAX_TURNS - turn messages = [f"You have {remaining} {'turn' if remaining == 1 else 'turns'} remaining out of {MAX_TURNS} total."] if turn == MAX_TURNS - 1: messages.append( f"This is your last turn: you must call {REPORT_TOOL['name']} now. " f"Anything you have found but not reported is lost when this turn ends." ) history += [{"role": "user", "content": message} for message in messages] for message in messages: print(f"[user] {message}") response = client.responses.create(input=history, **request)def tool_arguments(call) -> dict: """Parse a call's JSON arguments: keep only the parameters the tool's schema declares, and drop the ""/null "defaults" LLMs like to send. Anything else the model got wrong comes back as a ToolError from the fetch or the tool.""" try: arguments = json.loads(call.arguments) except json.JSONDecodeError: return {} schema = TOOLS_BY_NAME.get(call.name) if not isinstance(arguments, dict) or schema is None: return {} properties = schema["parameters"]["properties"] return {key: value for key, value in arguments.items() if key in properties and value not in ("", None)}def print_model_output(turn: int, response) -> None: """The model's side of one turn: its thinking, any prose, and its calls.""" print(f"\n--- turn {turn} (status={response.status}) " + "-" * 40) for item in response.output: if item.type == "reasoning": for part in getattr(item, "content", None) or []: print(f"[thinking] {part.text}") elif item.type == "message": print(f"[assistant] {response.output_text}") elif item.type == "function_call": print(f">>> {item.name}({item.arguments})")def print_report(cache: DocumentCache, reported_ids: list[str]) -> None: """The episode's answer: your database ids, most helpful first.""" print(f"\nReported for: {QUESTION}") for data_id in reported_ids: document = cache.get_document(data_id) print(f" {data_id} {document.get('title', '')}")if __name__ == "__main__": main()
/** SID-2 preview. * Implement vectorSearch and keywordSearch against your search backend. * Set ID_FIELD and SNIPPET_FIELD, then customize INSTRUCTIONS, QUESTION, * and TOOL_EXAMPLES to fit your data. Update tool descriptions and examples * together. The report_helpful_ids tool is fixed and cannot be modified. * * npm install @sid-ai/sid-sdk openai * npm install --save-dev tsx * BASETEN_API_KEY=your-api-key npx tsx reference_implementation_baseten.ts */import { randomUUID } from "node:crypto";import { pathToFileURL } from "node:url";import OpenAI from "openai";import type { ResponseCreateParamsNonStreaming, ResponseInput, ResponseFunctionToolCall,} from "openai/resources/responses/responses";import { DocumentCache, InvalidCharacterRange, parseRenderedModelFacingId, SNIPPET_SIZE_DEFAULT,} from "@sid-ai/sid-sdk";class ToolError extends Error {}type BackendDocument = Record<string, unknown>;type Arguments = Record<string, unknown>;const BASE_URL = process.env.BASETEN_BASE_URL ?? "https://inference.baseten.co/v1";const MODEL = "sid/sid-2-preview";const QUESTION = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?";const MAX_TURNS = 10;const REWARD = "f1";const REASONING_EFFORT = "low";const DEFAULT_LIMIT = 5;const MAX_PARALLEL_FETCHES = 8;const ID_FIELD = "id";const SNIPPET_FIELD = "content";const INSTRUCTIONS = `You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called.`;const SEARCH_TOOL = { "type": "function", "name": "search", "description": "Semantic (vector) search over the dataset. Best for conceptual or paraphrased queries: a long, descriptive query of what you want beats a few keywords; use text_search when you know the distinctive words a match would contain. Result ids may be ranged references 'id#start:end' naming the exact character span of the document shown; stretches you have already been shown are replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the document's total character count.", "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "The query to search the dataset with" }, "limit": { "type": "integer", "description": "The number of results to return" }, "snippet_size": { "type": "integer", "description": "Words shown per result snippet (default 50). Must be a positive integer. Pass a larger value to see more of each document." } }, "required": [ "query" ] }};const TEXT_SEARCH_TOOL = { "type": "function", "name": "text_search", "description": "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you want matched as plain words separated by spaces: the distinctive content words a matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact character span of the document shown; stretches you have already been shown are replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the document's total character count.", "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "Plain space-separated content words to match (BM25)." }, "limit": { "type": "integer", "description": "The number of results to return" }, "snippet_size": { "type": "integer", "description": "Words shown per result snippet (default 50). Must be a positive integer. Pass a larger value to see more of each document." } }, "required": [ "query" ] }};const READ_TOOL = { "type": "function", "name": "read", "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.", "parameters": { "type": "object", "properties": { "id": { "type": "string", "description": "The id of the document to read. Also accepts a ranged reference 'id#start:end' (character offsets, e.g. '123456#800:1600') to read only that part of the document's text; a range running past the end of the document is trimmed to the part that overlaps it." } }, "required": [ "id" ] }};const REPORT_TOOL = { "type": "function", "name": "report_helpful_ids", "description": "Report the most helpful document ids from the search results. This is required in the last step. Each id may be reported at most once; repeated ids are penalized. Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a ranged reference counts as its whole document, so there is no need to strip the '#start:end' part.", "parameters": { "type": "object", "properties": { "ids": { "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": { "type": "string" } } }, "required": [ "ids" ] }};const TOOL_EXAMPLES = [ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": { "query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5 } }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": { "query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5 } }, { "caption": "Example: reading one character range:", "name": "read", "arguments": { "id": "placeholder_1#800:1600" } }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": { "ids": [ "placeholder_1", "placeholder_2", "placeholder_3" ] } }];const TOOLS = [SEARCH_TOOL, TEXT_SEARCH_TOOL, READ_TOOL, REPORT_TOOL];const TOOLS_BY_NAME = new Map(TOOLS.map(tool => [tool.name, tool]));// Return full chunks with stable database IDs, best matches first.// Every field except ID_FIELD is shown to the model, so leave private fields out.async function vectorSearch(query: string, limit: number): Promise<BackendDocument[]> { throw new Error("Implement vectorSearch against your document store");}async function keywordSearch(query: string, limit: number): Promise<BackendDocument[]> { throw new Error("Implement keywordSearch against your document store");}async function fetchResults(name: string, args: Arguments): Promise<BackendDocument[] | null> { if (name !== "search" && name !== "text_search") return null; if (typeof args.query !== "string" || !args.query) { throw new ToolError("Search requires a non-empty query. Use read to fetch a specific document ID."); } for (const key of ["limit", "snippet_size"]) { if (args[key] !== undefined && !Number.isSafeInteger(args[key])) { throw new ToolError(`${key} must be an integer`); } } const snippetSize = (args.snippet_size ?? SNIPPET_SIZE_DEFAULT) as number; if (snippetSize <= 0) throw new ToolError("snippet_size must be a positive integer"); const limit = (args.limit ?? DEFAULT_LIMIT) as number; return name === "search" ? vectorSearch(args.query, limit) : keywordSearch(args.query, limit);}function renderResults(cache: DocumentCache, results: BackendDocument[], query: string, snippetSize: number): string { const rendered: string[] = []; for (const result of results) { const document = { ...result }; const dataId = document[ID_FIELD]; delete document[ID_FIELD]; if (typeof dataId !== "string" || !dataId) { throw new Error(`Backend document must carry a non-empty string ${ID_FIELD}`); } if (typeof document[SNIPPET_FIELD] !== "string" || !document[SNIPPET_FIELD]) { throw new Error(`Backend document ${dataId} must carry non-empty ${SNIPPET_FIELD} text`); } cache.addDocument(dataId, document); const view = cache.applySnippet(dataId, { snippetField: SNIPPET_FIELD, query, snippetSize }); rendered.push(view.renderXml()); cache.updateSeen(view); } return rendered.join("\n") || "No results found. Try a different query.";}function parseReference(reference: unknown) { if (typeof reference !== "string") throw new ToolError("Document reference must be a string"); try { return parseRenderedModelFacingId(reference); } catch (error) { if (error instanceof TypeError || error instanceof RangeError) throw new ToolError(error.message); throw error; }}function read(cache: DocumentCache, reference: unknown): string { const [modelId, range] = parseReference(reference); if (!cache.containsModelFacingId(modelId)) { throw new ToolError("Unknown document ID. Use search or text_search first."); } const dataId = cache.toDataId(modelId); let span = range ?? undefined; if (span) { try { span = cache.resolveCharRange(dataId, SNIPPET_FIELD, span); } catch (error) { if (error instanceof InvalidCharacterRange) throw new ToolError(error.message); throw error; } } // Count Unicode code points, matching the SDK's character offsets. if (span) { span = [span[0], Math.min(span[1], span[0] + 10_000)]; } else if (Array.from(cache.getDocument(dataId)[SNIPPET_FIELD] as string).length > 10_000) { span = [0, 10_000]; } const view = cache.getSingleSpanDocumentView(dataId, { snippetField: SNIPPET_FIELD, snippetDisplaySpan: span, }); const output = view.renderXml(); cache.updateSeen(view); return output;}function report(cache: DocumentCache, ids: unknown): string[] { if (!Array.isArray(ids) || !ids.every(id => typeof id === "string")) { throw new ToolError("ids must be a list of strings"); } if (!ids.length) throw new ToolError("ids list cannot be empty"); const modelIds = [...new Set(ids.map(id => parseReference(id)[0]))]; return modelIds .filter(id => cache.containsModelFacingId(id)) .map(id => cache.toDataId(id));}function toolArguments(call: ResponseFunctionToolCall): Arguments { let args: unknown; try { args = JSON.parse(call.arguments); } catch { return {}; } const tool = TOOLS_BY_NAME.get(call.name); if (!tool || !args || typeof args !== "object" || Array.isArray(args)) return {}; return Object.fromEntries(Object.entries(args).filter(([key, value]) => Object.hasOwn(tool.parameters.properties, key) && value !== "" && value !== null ));}// Limit backend concurrency, preserving the model's call order in the results.async function fetchInParallel(calls: ResponseFunctionToolCall[], args: Arguments[]) { const results: PromiseSettledResult<BackendDocument[] | null>[] = new Array(calls.length); let next = 0; async function worker() { while (next < calls.length) { const index = next++; try { results[index] = { status: "fulfilled", value: await fetchResults(calls[index].name, args[index]) }; } catch (reason) { results[index] = { status: "rejected", reason }; } } } await Promise.all(Array.from({ length: Math.min(MAX_PARALLEL_FETCHES, calls.length) }, worker)); return results;}async function main(): Promise<void> { const client = new OpenAI({ baseURL: BASE_URL, apiKey: process.env.BASETEN_API_KEY }); const cache = new DocumentCache({ rangeMode: "lenient" }); const sessionId = randomUUID(); // One session per episode. const request: Omit<ResponseCreateParamsNonStreaming, "input"> & { sid: object } = { model: MODEL, max_output_tokens: 4096, store: false, include: ["reasoning.encrypted_content"], instructions: INSTRUCTIONS, tools: TOOLS as ResponseCreateParamsNonStreaming["tools"], reasoning: { effort: REASONING_EFFORT }, sid: { tool_examples: TOOL_EXAMPLES, reward: REWARD }, }; const options = { headers: { "x-session-affinity": sessionId } }; const history: ResponseInput = [{ role: "user", content: QUESTION }]; console.log(`Q: ${QUESTION}`); for (let turn = 1; turn <= MAX_TURNS; turn++) { const response = await client.responses.create({ ...request, input: history }, options); console.log(`Turn ${turn} (${response.status})`, response.output); const calls = response.output.filter((item): item is ResponseFunctionToolCall => item.type === "function_call"); const args = calls.map(toolArguments); const fetches = await fetchInParallel(calls, args); const observations: string[] = []; let reported: string[] | undefined; // Render serially: what earlier calls show, later calls mask. for (let i = 0; i < calls.length; i++) { const call = calls[i]; let observation: string; try { const fetch = fetches[i]; if (fetch.status === "rejected") throw fetch.reason; if (call.name === "search" || call.name === "text_search") { observation = renderResults(cache, fetch.value!, args[i].query as string, (args[i].snippet_size ?? SNIPPET_SIZE_DEFAULT) as number); } else if (call.name === "read") { observation = read(cache, args[i].id); } else if (call.name === "report_helpful_ids") { if (calls.length > 1) throw new ToolError("Combining report tool call with other tool calls is not allowed."); if (turn === 1) throw new ToolError("Cannot report in the first turn. Must make a search first."); reported = report(cache, args[i].ids); observation = `Reported IDs submitted: ${reported.join(",")}`; } else { observation = `Tool '${call.name}' is not valid. Available tools: ${[...TOOLS_BY_NAME.keys()].join(", ")}`; } } catch (error) { if (!(error instanceof ToolError)) throw error; observation = `Error executing ${call.name}: ${error.message}`; } observations.push(observation); console.log(`[${call.name}] ${observation}`); } if (reported !== undefined) { console.log(`Reported for: ${QUESTION}`); for (const id of reported) console.log(id, cache.getDocument(id).title ?? ""); return; } if (turn === MAX_TURNS || response.status === "incomplete") { console.log("Episode ends without a report."); return; } // Echo output items, including encrypted reasoning, without modifications. history.push(...response.output as ResponseInput); history.push(...calls.map((call, i) => ({ type: "function_call_output" as const, call_id: call.call_id, output: observations[i], }))); const remaining = MAX_TURNS - turn; history.push({ role: "user", content: `You have ${remaining} ${remaining === 1 ? "turn" : "turns"} remaining out of ${MAX_TURNS} total.` }); if (remaining === 1) { history.push({ role: "user", content: `This is your last turn: you must call ${REPORT_TOOL.name} now. Anything you have found but not reported is lost when this turn ends.` }); } }}if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) { main().catch(error => { console.error(error); process.exitCode = 1; });}
go mod init sid-examplego get github.com/sidhq/sid-go github.com/openai/openai-go/v3BASETEN_API_KEY=your-api-key go run reference_implementation_baseten.go
reference_implementation_baseten.go
// SID-2 preview on Baseten.//// Before using this template: implement the vectorSearch and keywordSearch// backend functions, set idField and snippetField to your document's fields,// and customize the instructions, question, and toolExamples to fit your data.// You can also define additional tools, or extend existing tools with more// parameters (e.g. a filter). Just make sure to always update the tool's// description and usage example. The report_helpful_ids tool is fixed and// cannot be changed or modified.//// Run://// go mod init sid-example && go get github.com/sidhq/sid-go github.com/openai/openai-go/v3// BASETEN_API_KEY=your-api-key go run reference_implementation_baseten.gopackage mainimport ( "context" "crypto/rand" "encoding/hex" "encoding/json" "errors" "fmt" "log" "maps" "os" "slices" "strings" "sync" "unicode/utf8" "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" "github.com/openai/openai-go/v3/packages/param" "github.com/openai/openai-go/v3/responses" "github.com/openai/openai-go/v3/shared" sid "github.com/sidhq/sid-go")// ToolError is a mistake the *model* made (bad id, malformed ref, invalid// arguments). These errors are recoverable and passed back to the model as a// message.type ToolError struct{ message string }func (e *ToolError) Error() string { return e.message }func toolErrorf(format string, args ...any) error { return &ToolError{fmt.Sprintf(format, args...)}}// ---------------------------------------------------------------------------// The episode. TODO: your question. The model is served by Baseten,// under one BASETEN_API_KEY.// ---------------------------------------------------------------------------const question = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?"const ( model = "sid/sid-2-preview" maxTurns = 10 reward = "f1" // the trained scoring prompt: "f1" (few, precise ids) | "ndcg" (many, best first) reasoningEffort = "low" // "low" | "medium" | "high" defaultLimit = 5 // results per search when the model does not ask for a number maxReadChars = 10_000 // document text returned per read call maxParallelFetches = 8 // backend queries in flight at once, across a turn's search calls)func baseURL() string { if url := os.Getenv("BASETEN_BASE_URL"); url != "" { return strings.TrimRight(url, "/") } return "https://inference.baseten.co/v1"}// ---------------------------------------------------------------------------// TODO: your document backend. The tools below call these functions;// implement them against your database. A document is a map with://// idField : your document id (any stable string; never shown to the model)// snippetField : the document's canonical text. Return the identical string// for the same id everywhere: the seen-ledger's character// offsets index it.// extras : (optional) any other fields, such as a "title", rendered as// display attributes. Every cached field is shown to the// model (idField is removed before caching), so leave// anything private out of your documents.// ---------------------------------------------------------------------------const ( idField = "id" snippetField = "content")// vectorSearch is your semantic (embedding) search: the limit most relevant// documents for query, best first.func vectorSearch(ctx context.Context, query string, limit int) ([]sid.Document, error) { return nil, errors.New("TODO: implement vectorSearch against your document store")}// keywordSearch is your full-text keyword (BM25) search: the limit best// matching documents for query.func keywordSearch(ctx context.Context, query string, limit int) ([]sid.Document, error) { return nil, errors.New("TODO: implement keywordSearch against your document store")}// ---------------------------------------------------------------------------// The system prompt: one sentence naming your subject matter, then the// search -> repeat -> report scaffolding. TODO: edit the first paragraph to// describe your data and goal; keep the steps.// ---------------------------------------------------------------------------const instructions = `You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called.`// ---------------------------------------------------------------------------// The tool schemas. Each tool should be demonstrated in toolExamples.// ---------------------------------------------------------------------------type object = map[string]anyvar snippetSizeDescription = fmt.Sprintf( "Words shown per result snippet (default %d). Must be a positive integer. "+ "Pass a larger value to see more of each document.", sid.SnippetSizeDefault,)var searchTool = object{ "type": "function", "name": "search", "description": "Semantic (vector) search over the dataset. Best for conceptual or paraphrased " + "queries: a long, descriptive query of what you want beats a few keywords; use " + "text_search when you know the distinctive words a match would contain. Result " + "ids may be ranged references 'id#start:end' naming the exact character span of " + "the document shown; stretches you have already been shown are replaced by " + "'[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the " + "document's total character count.", "parameters": object{ "type": "object", "properties": object{ "query": object{"type": "string", "description": "The query to search the dataset with"}, "limit": object{"type": "integer", "description": "The number of results to return"}, "snippet_size": object{"type": "integer", "description": snippetSizeDescription}, }, "required": []string{"query"}, },}var textSearchTool = object{ "type": "function", "name": "text_search", "description": "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you " + "want matched as plain words separated by spaces: the distinctive content words a " + "matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact " + "character span of the document shown; stretches you have already been shown are " + "replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length " + "gives the document's total character count.", "parameters": object{ "type": "object", "properties": object{ "query": object{"type": "string", "description": "Plain space-separated content words to match (BM25)."}, "limit": object{"type": "integer", "description": "The number of results to return"}, "snippet_size": object{"type": "integer", "description": snippetSizeDescription}, }, "required": []string{"query"}, },}var readTool = object{ "type": "function", "name": "read", "description": "Search results show only short snippets. Read up to 10,000 characters of a document per call. " + "A bare ID starts at character 0; a ranged reference starts at its requested offset. Use further ranged reads to continue.", "parameters": object{ "type": "object", "properties": object{ "id": object{ "type": "string", "description": "The id of the document to read. Also accepts a ranged reference " + "'id#start:end' (character offsets, e.g. '123456#800:1600') to read only " + "that part of the document's text; a range running past the end of the " + "document is trimmed to the part that overlaps it.", }, }, "required": []string{"id"}, },}// reportTool is the fixed terminal tool: do not change its name, description// or parameters.var reportTool = object{ "type": "function", "name": "report_helpful_ids", "description": "Report the most helpful document ids from the search results. This is required in " + "the last step. Each id may be reported at most once; repeated ids are penalized. " + "Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a " + "ranged reference counts as its whole document, so there is no need to strip the " + "'#start:end' part.", "parameters": object{ "type": "object", "properties": object{ "ids": object{ "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": object{"type": "string"}, }, }, "required": []string{"ids"}, },}var tools = []object{searchTool, textSearchTool, readTool, reportTool}func toolProperties(name string) object { for _, tool := range tools { if tool["name"] == name { return tool["parameters"].(object)["properties"].(object) } } return nil}func toolNames() []string { names := make([]string, len(tools)) for i, tool := range tools { names[i] = tool["name"].(string) } return names}// toolExamples are few-shot calls, sent on every request via the sid// extension. TODO: write search examples that fit your data; keep the read and// report examples as they are.var toolExamples = []object{ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": object{"query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5}, }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": object{"query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5}, }, { "caption": "Example: reading one character range:", "name": "read", "arguments": object{"id": "placeholder_1#800:1600"}, }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": object{"ids": []string{"placeholder_1", "placeholder_2", "placeholder_3"}}, },}// ---------------------------------------------------------------------------// The tools, split in two: fetchResults is the slow backend call and runs in// parallel; everything below it touches the cache and runs serially, in call// order.// ---------------------------------------------------------------------------// fetchResults is the backend call behind one search tool call. read and the// report have nothing to fetch (reads are served from the document cache), so// they get nil.func fetchResults(ctx context.Context, name string, arguments object) ([]sid.Document, error) { if name != "search" && name != "text_search" { return nil, nil } // Catch mistyped arguments as a ToolError that the model can recover from. properties := toolProperties(name) for key, value := range arguments { expected := properties[key].(object)["type"] _, isString := value.(string) if expected == "string" && !isString || expected == "integer" && !isInteger(value) { return nil, toolErrorf("Invalid %s arguments: %s must be of type %s, got %T", name, key, expected, value) } } query, _ := arguments["query"].(string) // toolArguments strips ""; treat missing as empty if query == "" { return nil, toolErrorf("search requires a non-empty `query` to rank documents against. " + "To fetch a specific document by id, use the `read` tool. " + "To rank candidates against the question, call `search` with a meaningful query.") } snippetSize := intArgument(arguments, "snippet_size", sid.SnippetSizeDefault) if snippetSize <= 0 { return nil, toolErrorf("Invalid %s arguments: snippet_size must be a positive integer, got %d", name, snippetSize) } limit := intArgument(arguments, "limit", defaultLimit) if name == "search" { return vectorSearch(ctx, query, limit) } return keywordSearch(ctx, query, limit)}// renderResults renders one <doc> per result, in backend order. Shared by// both search tools; only their fetch differs.func renderResults(cache *sid.DocumentCache, results []sid.Document, query string, snippetSize int) (string, error) { rendered := make([]string, 0, len(results)) for _, result := range results { dataID, ok := result[idField].(string) if !ok || dataID == "" { return "", fmt.Errorf("backend document must carry a non-empty string %q: %v", idField, result) } // Keep backend results intact when searches reuse the same maps. document := maps.Clone(result) // Removed, not read: every cached field renders, and your database ids must not. delete(document, idField) if content, ok := document[snippetField].(string); !ok || content == "" { // Not a model mistake: the seen-ledger's offsets index this string. return "", fmt.Errorf("backend document %q must carry non-empty %q text", dataID, snippetField) } // Searches routinely return the same document, and the cache is add-once. if _, err := cache.AddDocument(dataID, document); err != nil { return "", err } // BM25 snippet, with already-shown stretches collapsed to [seen: ...]. view, err := cache.ApplySnippet(dataID, sid.ApplySnippetOptions{ SnippetField: snippetField, Query: query, SnippetSize: snippetSize, }) if err != nil { return "", err } xml, err := view.RenderXML() if err != nil { return "", err } rendered = append(rendered, xml) if err := cache.UpdateSeen(view); err != nil { // later views of this span collapse to [seen: ...] return "", err } } if len(rendered) == 0 { return "No results found. Try a different query.", nil } return strings.Join(rendered, "\n"), nil}// parseReference splits the model's ref ('R' or 'R#start:end') into id and// optional char range. The SDK parser already writes a descriptive message.func parseReference(reference any) (string, *sid.CharacterRange, error) { value, ok := reference.(string) if !ok { return "", nil, toolErrorf("id must be a string, got %T", reference) } modelID, charRange, err := sid.ParseRenderedModelFacingID(value) if err != nil { return "", nil, &ToolError{err.Error()} } return modelID, charRange, nil}// read fetches one document (or a character range of it) by the id a search// returned. Explicit reads are not masked: the model asked for this span, so// it renders up to maxReadChars characters. Model mistakes (bad ref grammar,// an unknown or hallucinated id, an unusable range) return a ToolError.func read(cache *sid.DocumentCache, reference any) (string, error) { modelID, charRange, err := parseReference(reference) // ToolError on bad grammar if err != nil { return "", err } if !cache.ContainsModelFacingID(modelID) { return "", toolErrorf("Invalid document id '%s'. "+ "You can only read documents returned by search. "+ "Use 'search' or 'text_search' first.", modelID) } dataID, err := cache.ToDataID(modelID) // the one conversion back into your ids if err != nil { return "", err } if charRange != nil { // Lenient mode trims partially overlapping ranges to the document. // Invalid ranges become recoverable tool errors. resolved, err := cache.ResolveCharRange(dataID, snippetField, *charRange) var invalid *sid.InvalidCharacterRange if errors.As(err, &invalid) { return "", toolErrorf("%s. Re-read '%s' without a range to read from the "+ "start, or pass a range inside 0:doc_length.", invalid, modelID) } else if err != nil { return "", err } charRange = &resolved } // Cap document text per read; further ranged reads can continue from here. // Count Unicode code points, matching the SDK's character offsets. if charRange != nil { charRange = sid.RangeValue(charRange.Start(), min(charRange.End(), charRange.Start()+maxReadChars)) } else { document, err := cache.GetDocument(dataID) if err != nil { return "", err } if content, _ := document[snippetField].(string); utf8.RuneCountInString(content) > maxReadChars { charRange = sid.RangeValue(0, maxReadChars) } } view, err := cache.GetSingleSpanDocumentView(dataID, sid.SingleSpanOptions{ SnippetField: snippetField, SnippetDisplaySpan: charRange, }) if err != nil { return "", err } rendered, err := view.RenderXML() if err != nil { return "", err } return rendered, cache.UpdateSeen(view)}// report is the terminal report: ids are doc-granular (a ranged ref denotes// its document), deduplicated preserving order, and translated back to your// data ids. IDs outside the model-facing ID map are dropped.func report(cache *sid.DocumentCache, ids any) ([]string, error) { values, ok := ids.([]any) if !ok { return nil, toolErrorf("ids must be a list of strings") } if len(values) == 0 { return nil, toolErrorf("ids list cannot be empty") } var reported []string for _, reference := range values { if _, ok := reference.(string); !ok { return nil, toolErrorf("ids must be a list of strings") } modelID, _, err := parseReference(reference) // ToolError on bad grammar if err != nil { return nil, err } if !slices.Contains(reported, modelID) { reported = append(reported, modelID) } } dataIDs := make([]string, 0, len(reported)) for _, modelID := range reported { if dataID, err := cache.ToDataID(modelID); err == nil { dataIDs = append(dataIDs, dataID) } } return dataIDs, nil}// ---------------------------------------------------------------------------// The agent loop: one Responses API call per turn, tool outputs paired back// by call_id. The episode ends when the model reports.//// Requests are stateless (store=false): nothing is retained server-side, so// the full conversation is resent on every call. The model's reasoning comes// back as an opaque encrypted_content blob on each reasoning item, and// carries over by echoing the output items back verbatim in the next input.//// The prompt (instructions, tools, reasoning effort, tool_examples, reward) is// fixed for the life of the conversation. Build the request once and send it// unchanged on every turn; only input grows.// ---------------------------------------------------------------------------type fetched struct { results []sid.Document err error}func main() { ctx := context.Background() cache, err := sid.NewDocumentCache(sid.DocumentCacheOptions{RangeMode: sid.RangeModeLenient}) // this episode's id map, document cache and seen-ledger if err != nil { log.Fatal(err) } apiKey := os.Getenv("BASETEN_API_KEY") if apiKey == "" { log.Fatal("BASETEN_API_KEY is not set") } client := openai.NewClient(option.WithBaseURL(baseURL()), option.WithAPIKey(apiKey)) // One request shape, only update the history at each turn. request := responses.ResponseNewParams{ Model: model, MaxOutputTokens: openai.Int(4096), Store: openai.Bool(false), Include: []responses.ResponseIncludable{responses.ResponseIncludableReasoningEncryptedContent}, Instructions: openai.String(instructions), // Reasoning effort, tool_examples, and the target reward are managed server-side Reasoning: shared.ReasoningParam{Effort: reasoningEffort}, } options := []option.RequestOption{ // The tools and the sid extension are plain JSON, set on the request body. option.WithJSONSet("tools", tools), option.WithJSONSet("sid", object{"tool_examples": toolExamples, "reward": reward}), option.WithHeader("x-session-affinity", newSessionID()), // KV cache routing: one session per episode } // The whole conversation as Responses API input items, grown every turn. history := responses.ResponseInputParam{ responses.ResponseInputItemParamOfMessage(question, responses.EasyInputMessageRoleUser), } create := func() (*responses.Response, error) { request.Input = responses.ResponseNewParamsInputUnion{OfInputItemList: history} return client.Responses.New(ctx, request, options...) } fmt.Printf("Q: %s\n", question) resp, err := create() if err != nil { log.Fatal(err) } for turn := 1; turn <= maxTurns; turn++ { printModelOutput(turn, resp) var calls []responses.ResponseFunctionToolCall for _, item := range resp.Output { if item.Type == "function_call" { calls = append(calls, item.AsFunctionCall()) } } arguments := make([]object, len(calls)) for i, call := range calls { arguments[i] = toolArguments(call) } // Execute the model's tool calls in parallel; maxParallelFetches caps // how many queries your backend sees at once (1 makes the turn serial). fetches := make([]fetched, len(calls)) semaphore := make(chan struct{}, maxParallelFetches) var wg sync.WaitGroup for i, call := range calls { wg.Add(1) go func() { defer wg.Done() semaphore <- struct{}{} defer func() { <-semaphore }() results, err := fetchResults(ctx, call.Name, arguments[i]) fetches[i] = fetched{results, err} }() } wg.Wait() // Compose observations serially, in call order: what earlier calls show, later calls mask. observations := make([]string, len(calls)) var reported []string for i, call := range calls { observation, err := runTool(cache, turn, len(calls), call.Name, arguments[i], fetches[i], &reported) var toolError *ToolError if errors.As(err, &toolError) { // A recoverable mistake made by the model. observation = fmt.Sprintf("Error executing %s: %s", call.Name, toolError) } else if err != nil { log.Fatal(err) } observations[i] = observation fmt.Printf(" [%s] -> %s\n", call.Name, observation) // verbatim: what you see is what the model receives } // Termination conditions if reported != nil { printReport(cache, reported) return } if turn == maxTurns { fmt.Printf("\n[max turns reached: %d: episode ends without a report]\n", maxTurns) return } if resp.Status == responses.ResponseStatusIncomplete { fmt.Println("\n[generation hit max_output_tokens; ending the episode]") return } // Non-terminal turn: append the model's output verbatim (the reasoning // items carry their encrypted_content), then feed back the tool outputs // (paired by call_id) plus the turn budget indicator as a user message. for _, item := range resp.Output { history = append(history, param.Override[responses.ResponseInputItemUnionParam](json.RawMessage(item.RawJSON()))) } for i, call := range calls { output := responses.ResponseInputItemParamOfFunctionCallOutput(observations[i]) output.OfFunctionCallOutput.CallID = openai.String(call.CallID) history = append(history, output) } remaining := maxTurns - turn unit := "turns" if remaining == 1 { unit = "turn" } messages := []string{fmt.Sprintf("You have %d %s remaining out of %d total.", remaining, unit, maxTurns)} if turn == maxTurns-1 { messages = append(messages, fmt.Sprintf( "This is your last turn: you must call %s now. "+ "Anything you have found but not reported is lost when this turn ends.", reportTool["name"])) } for _, message := range messages { history = append(history, responses.ResponseInputItemParamOfMessage(message, responses.EasyInputMessageRoleUser)) fmt.Printf("[user] %s\n", message) } if resp, err = create(); err != nil { log.Fatal(err) } }}// runTool dispatches one call. A ToolError is a recoverable model mistake;// any other error is a bug or a backend failure and ends the episode.func runTool( cache *sid.DocumentCache, turn, callCount int, name string, arguments object, fetch fetched, reported *[]string,) (string, error) { if fetch.err != nil { return "", fetch.err } switch name { case "search", "text_search": return renderResults(cache, fetch.results, arguments["query"].(string), intArgument(arguments, "snippet_size", sid.SnippetSizeDefault)) case "read": return read(cache, arguments["id"]) case "report_helpful_ids": if callCount > 1 { return "", toolErrorf("Combining report tool call with other tool calls is not allowed.") } if turn == 1 { return "", toolErrorf("Cannot report in the first turn. Must make a search first.") } ids, err := report(cache, arguments["ids"]) if err != nil { return "", err } *reported = ids return "reported ids submitted: " + strings.Join(ids, ","), nil default: return fmt.Sprintf("Tool '%s' is not a valid tool. Available tools: %v", name, toolNames()), nil }}// toolArguments parses a call's JSON arguments: keep only the parameters the// tool's schema declares, and drop the ""/null "defaults" LLMs like to send.// Anything else the model got wrong comes back as a ToolError from the fetch// or the tool.func toolArguments(call responses.ResponseFunctionToolCall) object { var arguments object properties := toolProperties(call.Name) if json.Unmarshal([]byte(call.Arguments), &arguments) != nil || properties == nil { return object{} } for key, value := range arguments { if _, declared := properties[key]; !declared || value == nil || value == "" { delete(arguments, key) } } return arguments}// isInteger reports whether a decoded JSON value is a whole number.func isInteger(value any) bool { number, ok := value.(float64) return ok && number == float64(int(number))}func intArgument(arguments object, key string, fallback int) int { if value, ok := arguments[key].(float64); ok { return int(value) } return fallback}func newSessionID() string { var id [16]byte if _, err := rand.Read(id[:]); err != nil { log.Fatal(err) } return hex.EncodeToString(id[:])}// printModelOutput prints the model's side of one turn: its thinking, any// prose, and its calls.func printModelOutput(turn int, resp *responses.Response) { fmt.Printf("\n--- turn %d (status=%s) %s\n", turn, resp.Status, strings.Repeat("-", 40)) for _, item := range resp.Output { switch item.Type { case "reasoning": for _, part := range item.AsReasoning().Content { fmt.Printf("[thinking] %s\n", part.Text) } case "message": fmt.Printf("[assistant] %s\n", resp.OutputText()) case "function_call": fmt.Printf(">>> %s(%s)\n", item.Name, item.AsFunctionCall().Arguments) } }}// printReport prints the episode's answer: your database ids, most helpful// first.func printReport(cache *sid.DocumentCache, reported []string) { fmt.Printf("\nReported for: %s\n", question) for _, dataID := range reported { document, _ := cache.GetDocument(dataID) title, _ := document["title"].(string) fmt.Printf(" %s %s\n", dataID, title) }}
Set SAGEMAKER_ENDPOINT_NAME to your endpoint's name and AWS_DEFAULT_REGION to its region. boto3 ignores AWS_REGION. The responses-sagemaker package signs requests with your AWS credentials, which need sagemaker:InvokeEndpoint on the endpoint.
"""SID-2 preview on Amazon SageMaker.Before using this template: implement the vector_search and keyword_search backendfunctions, set ID_FIELD and SNIPPET_FIELD to your document's fields, andcustomize the INSTRUCTIONS, QUESTION, and TOOL_EXAMPLES to fit your data. You canalso define additional tools, or extend existing tools with more parameters (e.g. afilter). Just make sure to always update the tool's description and usage example.The report_helpful_ids tool is fixed and cannot be changed or modified.Run: SAGEMAKER_ENDPOINT_NAME=your-endpoint AWS_DEFAULT_REGION=your-region uv run --with sid-sdk --with openai \ --with "responses-sagemaker @ git+https://github.com/sidhq/responses-sagemaker" reference_implementation_sagemaker.py"""import jsonimport osimport uuidfrom concurrent.futures import ThreadPoolExecutorimport boto3import httpxfrom openai import OpenAIfrom responses_sagemaker import SageMakerResponsesTransportfrom sid import DocumentCache, InvalidCharacterRange, parse_rendered_model_facing_idfrom sid.document_cache import SNIPPET_SIZE_DEFAULTclass ToolError(Exception): """A mistake the *model* made (bad id, malformed ref, invalid arguments), these errors are recoverable and passed back to the model as a message."""# ---------------------------------------------------------------------------# The episode. TODO: your question. The model is served by your SageMaker# endpoint, authenticated with your AWS credentials.# ---------------------------------------------------------------------------QUESTION = "What government position was held by the woman who portrayed Corliss Archer in the film Kiss and Tell?"ENDPOINT_NAME = os.environ["SAGEMAKER_ENDPOINT_NAME"]MODEL = "sid/sid-2-preview"MAX_TURNS = 10REWARD = "f1" # the trained scoring prompt: "f1" (few, precise ids) | "ndcg" (many, best first)REASONING_EFFORT = "low" # "low" | "medium" | "high"DEFAULT_LIMIT = 5 # results per search when the model does not ask for a numberMAX_PARALLEL_FETCHES = 8 # backend queries in flight at once, across a turn's search calls# ---------------------------------------------------------------------------# TODO: your document backend. The tools below call these functions;# implement them against your database. A document is a dict with:# ID_FIELD : your document id (any stable string; never shown to the model)# SNIPPET_FIELD: the document's canonical text. Return the identical string# for the same id everywhere: the seen-ledger's character# offsets index it.# extras : (optional) any other fields, such as a "title", rendered# as display attributes. Every cached field is shown to the# model (ID_FIELD is popped before caching), so leave# anything private out of your documents.# ---------------------------------------------------------------------------ID_FIELD = "id"SNIPPET_FIELD = "content"def vector_search(query: str, limit: int) -> list[dict]: """TODO: your semantic (embedding) search: the ``limit`` most relevant documents for ``query``, best first.""" raise NotImplementedError("implement vector_search against your document store")def keyword_search(query: str, limit: int) -> list[dict]: """TODO: your full-text keyword (BM25) search: the ``limit`` best matching documents for ``query``.""" raise NotImplementedError("implement keyword_search against your document store")# ---------------------------------------------------------------------------# The system prompt: one sentence naming your subject matter, then the# search -> repeat -> report scaffolding. TODO: edit the first paragraph to# describe your data and goal; keep the steps.# ---------------------------------------------------------------------------INSTRUCTIONS = """You are an expert research assistant that is given a question and must use the provided search tools to find all documents needed to answer the question.Steps:1. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.2. Repeat step 1 until all documents necessary and sufficient to answer the question have been found. Take as many turns and searches as needed – you can make multiple searches per turn! Most questions will require multiple turns. Many will need more.3. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).The episode ends once report_helpful_ids is called."""# ---------------------------------------------------------------------------# The tool schemas. Each tool should be demonstrated in TOOL_EXAMPLES.# ---------------------------------------------------------------------------SNIPPET_SIZE_DESCRIPTION = ( f"Words shown per result snippet (default {SNIPPET_SIZE_DEFAULT}). Must be a positive integer. " f"Pass a larger value to see more of each document.")SEARCH_TOOL = { "type": "function", "name": "search", "description": ( "Semantic (vector) search over the dataset. Best for conceptual or paraphrased " "queries: a long, descriptive query of what you want beats a few keywords; use " "text_search when you know the distinctive words a match would contain. Result " "ids may be ranged references 'id#start:end' naming the exact character span of " "the document shown; stretches you have already been shown are replaced by " "'[seen: \"#start:end\"]' markers instead of repeating, and doc_length gives the " "document's total character count." ), "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "The query to search the dataset with"}, "limit": {"type": "integer", "description": "The number of results to return"}, "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION}, }, "required": ["query"], },}TEXT_SEARCH_TOOL = { "type": "function", "name": "text_search", "description": ( "Full-text keyword search over the dataset, ranked by BM25. Provide the terms you " "want matched as plain words separated by spaces: the distinctive content words a " "matching document would contain. Result ids may be ranged references 'id#start:end' naming the exact " "character span of the document shown; stretches you have already been shown are " "replaced by '[seen: \"#start:end\"]' markers instead of repeating, and doc_length " "gives the document's total character count." ), "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": ( "Plain space-separated content words to match (BM25)." ), }, "limit": {"type": "integer", "description": "The number of results to return"}, "snippet_size": {"type": "integer", "description": SNIPPET_SIZE_DESCRIPTION}, }, "required": ["query"], },}READ_TOOL = { "type": "function", "name": "read", "description": "Search results show only short snippets. Read the full content of a specific document.", "parameters": { "type": "object", "properties": { "id": { "type": "string", "description": ( "The id of the document to read. Also accepts a ranged reference " "'id#start:end' (character offsets, e.g. '123456#800:1600') to read only " "that part of the document's text; a range running past the end of the " "document is trimmed to the part that overlaps it." ), }, }, "required": ["id"], },}# The fixed terminal tool: do not change its name, description or parameters.REPORT_TOOL = { "type": "function", "name": "report_helpful_ids", "description": ( "Report the most helpful document ids from the search results. This is required in " "the last step. Each id may be reported at most once; repeated ids are penalized. " "Ids may be whole documents ('123456') or ranged references ('123456#800:1600'); a " "ranged reference counts as its whole document, so there is no need to strip the " "'#start:end' part." ), "parameters": { "type": "object", "properties": { "ids": { "type": "array", "description": "The ids of the most helpful documents from the search results.", "items": {"type": "string"}, }, }, "required": ["ids"], },}TOOLS = [SEARCH_TOOL, TEXT_SEARCH_TOOL, READ_TOOL, REPORT_TOOL]TOOLS_BY_NAME = {tool["name"]: tool for tool in TOOLS}# Few-shot calls, sent on every request via the sid extension. TODO: write search# examples that fit your data; keep the read and report examples as they areTOOL_EXAMPLES = [ { "caption": "Example: a semantic search with a long, descriptive query:", "name": "search", "arguments": { "query": "the actress who played Corliss Archer in the 1945 film, and the government post she later held", "limit": 5, }, }, { "caption": "Example: a keyword search:", "name": "text_search", "arguments": {"query": "Kiss and Tell 1945 Shirley Temple Corliss Archer", "limit": 5}, }, { "caption": "Example: reading one character range:", "name": "read", "arguments": {"id": "placeholder_1#800:1600"}, }, { "caption": "Example: ending the episode with a report:", "name": "report_helpful_ids", "arguments": {"ids": ["placeholder_1", "placeholder_2", "placeholder_3"]}, },]# ---------------------------------------------------------------------------# The tools, split in two: fetch_results is the slow backend call and runs in# parallel; everything below it touches the cache and runs serially, in call# order.# ---------------------------------------------------------------------------def fetch_results(name: str, arguments: dict) -> list[dict] | None: """The backend call behind one search tool call. `read` and the report have nothing to fetch (reads are served from the document cache), so they get ``None``.""" if name not in ("search", "text_search"): return None # Catch mistyped arguments as a ToolError that the model can recover from. properties = TOOLS_BY_NAME[name]["parameters"]["properties"] for key, value in arguments.items(): expected = str if properties[key]["type"] == "string" else int if not isinstance(value, expected) or isinstance(value, bool): raise ToolError( f"Invalid {name} arguments: {key} must be of type {properties[key]['type']}, " f"got {type(value).__name__}" ) query = arguments.get("query", "") # tool_arguments strips ""; treat missing as empty if not query: raise ToolError( "search requires a non-empty `query` to rank documents against. " "To fetch a specific document by id, use the `read` tool. " "To rank candidates against the question, call `search` with a meaningful query." ) snippet_size = arguments.get("snippet_size", SNIPPET_SIZE_DEFAULT) if snippet_size <= 0: raise ToolError( f"Invalid {name} arguments: snippet_size must be a positive integer, got {snippet_size}" ) limit = arguments.get("limit", DEFAULT_LIMIT) return vector_search(query, limit) if name == "search" else keyword_search(query, limit)def render_results(cache: DocumentCache, results: list[dict], query: str, snippet_size: int) -> str: """One <doc> per result, in backend order. Shared by both search tools; only their fetch differs.""" rendered = [] for document in results: if not isinstance(document, dict) or not isinstance(document.get(ID_FIELD), str) or not document[ID_FIELD]: raise ValueError(f"backend document must carry a non-empty string {ID_FIELD!r}: {document!r}") # Popped, not read: every cached field renders, and your database ids must not. data_id = document.pop(ID_FIELD) content = document.get(SNIPPET_FIELD) if not isinstance(content, str) or not content: # Not a model mistake: the seen-ledger's offsets index this string. raise ValueError(f"backend document {data_id!r} must carry non-empty {SNIPPET_FIELD!r} text") # Searches routinely return the same document, and the cache is add-once. if data_id not in cache: cache.add_document(data_id, document) # BM25 snippet, with already-shown stretches collapsed to [seen: ...]. view = cache.apply_snippet( data_id, SNIPPET_FIELD, query, snippet_size=snippet_size, ) rendered.append(view.render_xml()) cache.update_seen(view) # later views of this span collapse to [seen: ...] return "\n".join(rendered) if rendered else "No results found. Try a different query."def parse_reference(reference) -> tuple[str, tuple[int, int] | None]: """The model's ref in (``'R'`` or ``'R#start:end'``), split into id and optional char range. The SDK parser raises ``ValueError`` on malformed shapes with a descriptive message already written.""" try: return parse_rendered_model_facing_id(reference) except ValueError as error: raise ToolError(str(error)) from Nonedef read(cache: DocumentCache, reference: str | None) -> str: """Fetch one document (or a character range of it) by the id a search returned. Explicit reads are not masked: the model asked for this span, so it renders in full. Model mistakes (bad ref grammar, an unknown or hallucinated id, an unusable range) raise ToolError.""" model_facing_id, char_range = parse_reference(reference) # ToolError on bad grammar if not cache.contains_model_facing_id(model_facing_id): raise ToolError( f"Invalid document id '{model_facing_id}'. " f"You can only read documents returned by search. " f"Use 'search' or 'text_search' first." ) data_id = cache.to_data_id(model_facing_id) # the one conversion back into your ids if char_range is not None: # Lenient mode trims partially overlapping ranges to the document. # Invalid ranges become recoverable tool errors. try: char_range = cache.resolve_char_range(data_id, SNIPPET_FIELD, char_range) except InvalidCharacterRange as error: raise ToolError( f"{error}. Re-read '{model_facing_id}' without a range to see the whole " f"document, or pass a range inside 0:doc_length." ) from None view = cache.get_single_span_document_view( data_id, SNIPPET_FIELD, snippet_display_span=char_range, ) rendered = view.render_xml() cache.update_seen(view) return rendereddef report(cache: DocumentCache, ids) -> list[str]: """The terminal report: ids are doc-granular (a ranged ref denotes its document), deduplicated preserving order, and translated back to your data_ids. IDs outside the model-facing ID map are dropped.""" if not isinstance(ids, list) or not all(isinstance(reference, str) for reference in ids): raise ToolError("ids must be a list of strings") if not ids: raise ToolError("ids list cannot be empty") reported: list[str] = [] for reference in ids: model_facing_id, _ = parse_reference(reference) # ToolError on bad grammar if model_facing_id not in reported: reported.append(model_facing_id) return [ cache.to_data_id(model_facing_id) for model_facing_id in reported if cache.contains_model_facing_id(model_facing_id) ]# ---------------------------------------------------------------------------# The agent loop: one Responses API call per turn, tool outputs paired back# by call_id. The episode ends when the model reports.## Requests are stateless (store=False): nothing is retained server-side, so# the full conversation is resent on every call. The model's reasoning comes# back as an opaque ``encrypted_content`` blob on each reasoning item, and# carries over by echoing the output items back verbatim in the next input.## The prompt (instructions, tools, reasoning effort, tool_examples, reward) is# fixed for the life of the conversation. Build the request dict once and send# it unchanged on every turn; only `input` grows.# ---------------------------------------------------------------------------def main() -> None: session = boto3.Session() # credentials and region, as boto3 resolves them client = OpenAI( base_url=f"https://runtime.sagemaker.{session.region_name}.amazonaws.com/endpoints/{ENDPOINT_NAME}", api_key="unused", # SageMaker authenticates with SigV4; the SDK requires a key http_client=httpx.Client( transport=SageMakerResponsesTransport(credentials=session.get_credentials(), region=session.region_name) ), ) cache = DocumentCache(range_mode="lenient") # this episode's id map, document cache and seen-ledger session_id = uuid.uuid4().hex # KV cache routing information: one session per episode # Caps how many queries your backend sees at once (1 makes the turn serial). pool = ThreadPoolExecutor(max_workers=MAX_PARALLEL_FETCHES) # One request shape, only update the history at each turn. request = dict( model=MODEL, max_output_tokens=4096, store=False, include=["reasoning.encrypted_content"], instructions=INSTRUCTIONS, tools=TOOLS, # Reasoning effort, tool_examples, and the target reward are managed server-side reasoning={"effort": REASONING_EFFORT}, extra_body={"sid": {"tool_examples": TOOL_EXAMPLES, "reward": REWARD}}, extra_headers={"x-session-affinity": session_id}, ) # The whole conversation as Responses API input items, grown every turn. history: list[dict] = [{"role": "user", "content": QUESTION}] print(f"Q: {QUESTION}") response = client.responses.create(input=history, **request) for turn in range(1, MAX_TURNS + 1): print_model_output(turn, response) calls = [item for item in response.output if item.type == "function_call"] arguments = [tool_arguments(call) for call in calls] # Execute the model's tool calls in parallel fetches = [pool.submit(fetch_results, call.name, args) for call, args in zip(calls, arguments, strict=True)] # Compose observations serially, in call order: what earlier calls show, later calls mask. observations: list[str] = [] reported: list[str] | None = None for call, args, fetch in zip(calls, arguments, fetches, strict=True): try: if call.name in ("search", "text_search"): observation = render_results( cache, fetch.result(), args["query"], args.get("snippet_size", SNIPPET_SIZE_DEFAULT), ) elif call.name == "read": observation = read(cache, args.get("id")) elif call.name == "report_helpful_ids": if len(calls) > 1: raise ToolError("Combining report tool call with other tool calls is not allowed.") if turn == 1: raise ToolError("Cannot report in the first turn. Must make a search first.") reported = report(cache, args.get("ids")) observation = "reported ids submitted: " + ",".join(reported) else: observation = f"Tool '{call.name}' is not a valid tool. Available tools: {list(TOOLS_BY_NAME)}" except ToolError as e: # A recoverable mistake made by the model. observation = f"Error executing {call.name}: {e}" observations.append(observation) print(f" [{call.name}] -> {observation}") # verbatim: what you see is what the model receives # Termination conditions if reported is not None: print_report(cache, reported) return if turn == MAX_TURNS: print(f"\n[max turns reached: {MAX_TURNS}: episode ends without a report]") return if response.status == "incomplete": print("\n[generation hit max_output_tokens; ending the episode]") return # Non-terminal turn: append the model's output verbatim (the reasoning # items carry their encrypted_content), then feed back the tool outputs # (paired by call_id) plus the turn budget indicator as a user message. history += [item.model_dump(mode="json", exclude_none=True) for item in response.output] history += [ {"type": "function_call_output", "call_id": call.call_id, "output": observation} for call, observation in zip(calls, observations, strict=True) ] remaining = MAX_TURNS - turn messages = [f"You have {remaining} {'turn' if remaining == 1 else 'turns'} remaining out of {MAX_TURNS} total."] if turn == MAX_TURNS - 1: messages.append( f"This is your last turn: you must call {REPORT_TOOL['name']} now. " f"Anything you have found but not reported is lost when this turn ends." ) history += [{"role": "user", "content": message} for message in messages] for message in messages: print(f"[user] {message}") response = client.responses.create(input=history, **request)def tool_arguments(call) -> dict: """Parse a call's JSON arguments: keep only the parameters the tool's schema declares, and drop the ""/null "defaults" LLMs like to send. Anything else the model got wrong comes back as a ToolError from the fetch or the tool.""" try: arguments = json.loads(call.arguments) except json.JSONDecodeError: return {} schema = TOOLS_BY_NAME.get(call.name) if not isinstance(arguments, dict) or schema is None: return {} properties = schema["parameters"]["properties"] return {key: value for key, value in arguments.items() if key in properties and value not in ("", None)}def print_model_output(turn: int, response) -> None: """The model's side of one turn: its thinking, any prose, and its calls.""" print(f"\n--- turn {turn} (status={response.status}) " + "-" * 40) for item in response.output: if item.type == "reasoning": for part in getattr(item, "content", None) or []: print(f"[thinking] {part.text}") elif item.type == "message": print(f"[assistant] {response.output_text}") elif item.type == "function_call": print(f">>> {item.name}({item.arguments})")def print_report(cache: DocumentCache, reported_ids: list[str]) -> None: """The episode's answer: your database ids, most helpful first.""" print(f"\nReported for: {QUESTION}") for data_id in reported_ids: document = cache.get_document(data_id) print(f" {data_id} {document.get('title', '')}")if __name__ == "__main__": main()