SID 1

1import os
2from openai import OpenAI
3
4client = OpenAI(
5    base_url="https://inference.baseten.co/v1",
6    api_key=os.environ["BASETEN_API_KEY"],
7    default_headers={
8        "Authorization": f"Api-Key {os.environ['BASETEN_API_KEY']}"
9    },
10)
11
12SYSTEM_PROMPT = """
13You 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.
14
15Steps:
161. Reflect on what information is needed to answer the question and use the search tools to find documents. Each document has an id.
172. 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. Most questions require at least 5-8 search requests. Many will need more.
183. Use the report_helpful_ids tool to report the most helpful document ids. List the most helpful document ids first (important!).
19
20The interaction ends once report_helpful_ids is called. Every assistant turn should contain a tool call: use search/read tools while gathering evidence, and use report_helpful_ids for the final ranked IDs. You will be scored based on whether you have found all the documents and whether you reported them in the correct order (NDCG)
21
22
23You have access to the following tools:
24
25- search: performs a semantic search with the query
26  - Arguments: query (required), limit (optional, default 5, max 15)
27- text_search: performs full-text search
28  - Arguments: query (required), limit (optional, default 5, max 15)
29- read: read the full document content for an ID returned by search or text_search
30  - Arguments: id (required)
31- report_helpful_ids: report helpful document IDs in order (most helpful first)
32  - Arguments: ids (required, list of strings)
33
34To use a tool, enclose it within <tool_call> tags with a Python dictionary containing "name" and "arguments". For example:
35
36<tool_call>
37{"name": "search", "arguments": {"query": "machine learning algorithms", "limit": 3}}
38</tool_call>
39
40The semantic search tool will match things that are conceptually related or use synonyms. This request above would also find texts that talk about linear regression, for example, although "linear regression" does not appear in the query directly. You can write long queries describing the document you want precisely with this tool.
41
42
43<tool_call>
44{"name": "text_search", "arguments": {"query": "data \"dimensionality reduction\" -PCA"}}
45</tool_call>
46
47For text_search queries, you can use \"\" (escaped double quotes) to find exact matches for a term. Since the query is inside a JSON string with double quotes, you need to escape the inner double quotes with backslashes (\"dimensionality reduction\"). The above will not return matches that only contain one of "reduction" or "dimensionality" -- it needs both.
48You can also use a - to exclude terms (like PCA in the example above). You don't need to use \"\" or - operators, but it can be helpful. If your text_search query has too many terms, there might not be a document that matches all the constraints and no data will be found.
49
50
51Both search tools return snippets (relevant excerpts) rather than full documents. Snippets are approximately 50 words long and show the most relevant portion of the document based on your query. If the document was truncated, you'll see "..." at the beginning or end.
52To read the full document content, use the read tool with the document ID from your search results. You can only read documents that were previously returned by search or text_search.
53
54<tool_call>
55{"name": "read", "arguments": {"id": "placeholder_1"}}
56</tool_call>
57
58After you've received the tool responses, you can report the helpful document IDs:
59
60<tool_call>
61{"name": "report_helpful_ids", "arguments": {"ids": ["placeholder_1", "placeholder_2", "placeholder_3"]}}
62</tool_call>
63"""
64
65# --- Connect these three functions to your retrieval backend ---
66
67def search(query: str, limit: int = 5) -> list[dict]:
68    """Semantic search. Return the documents your backend matched."""
69    raise NotImplementedError("Connect `search` to your retrieval backend.")
70
71
72def text_search(query: str, limit: int = 5) -> list[dict]:
73    """Full-text search. Return the documents your backend matched."""
74    raise NotImplementedError("Connect `text_search` to your retrieval backend.")
75
76
77def read(id: str) -> dict:
78    """Fetch one full document by id. Truncate the content to ~3000 tokens."""
79    raise NotImplementedError("Connect `read` to your retrieval backend.")
80
81
82def format_docs_xml(docs: list[dict]) -> str:
83    parts = []
84    for doc in docs:
85        title = doc.get("title", "")
86        doc_id = doc.get("id", "")
87        content = doc.get("content", "")
88        parts.append(f'<doc id="{doc_id}" title="{title}">\n{content}\n</doc>')
89    return "\n".join(parts)
90
91
92def turn_budget_notice(turns_left: int) -> str:
93    notice = f"\n\nYou have {turns_left} out of {MAX_TURNS} turns left."
94    if turns_left == 1:
95        notice += " You must call `report_helpful_ids` in the next turn."
96    return notice
97
98
99def execute_tool_call(tc) -> tuple[str, dict, str, bool]:
100    name = tc.function.name
101    arguments = json.loads(tc.function.arguments)
102    print(f"  Tool: {name}({arguments})")
103
104    if name == "search":
105        docs = search(arguments["query"], arguments.get("limit", 5))
106        return tc.id, arguments, format_docs_xml(docs), False
107
108    if name == "text_search":
109        docs = text_search(arguments["query"], arguments.get("limit", 5))
110        return tc.id, arguments, format_docs_xml(docs), False
111
112    if name == "read":
113        return tc.id, arguments, format_docs_xml([read(arguments["id"])]), False
114
115    if name == "report_helpful_ids":
116        return tc.id, arguments, json.dumps(arguments["ids"]), True
117
118    return tc.id, arguments, f"Unknown tool: {name}", False
119
120tools = [{"type": "function", "function": {"name": "dummy"}}]
121
122messages = [
123    {"role": "system", "content": SYSTEM_PROMPT},
124    {"role": "user", "content": "In which cities were the indoor and outdoor 4 minute mile records set?"},
125]
126response = client.chat.completions.create(model="sid/sid-1", messages=messages, tools=tools)
127
128print(response.to_json())