Setting the file. One moment. File Cu Canary
helpers/knowledge_base_retrieve.py
Python·187 lines·8 KB
12
try
:
13 from ._common import (
14 SEARCH_AUDIENCE, HelperFailure, HttpResult, TokenProvider, Transport,
15 azure_cli_token, blocked_result, emit_result, http_request,
16 reject_secrets, require_allowed_fields, validate_search_endpoint,
17 )
18except ImportError:
19 from _common import ( # type: ignore[no-redef]
20 SEARCH_AUDIENCE, HelperFailure, HttpResult, TokenProvider, Transport,
21 azure_cli_token, blocked_result, emit_result, http_request,
22 reject_secrets, require_allowed_fields, validate_search_endpoint,
23 )
24
25
26API_VERSIONS = {"2026-04-01", "2026-08-01-preview"}
27
28
29def _invalid(message: str) -> HelperFailure:
30 return HelperFailure("input-schema-invalid", message, blocked_at="input-resolution")
31
32
33def _target(endpoint: Any, name: Any, api_version: Any) -> str:
34 if not isinstance(api_version, str) or api_version not in API_VERSIONS:
35 raise _invalid("Unsupported API version.")
36 if not isinstance(name, str) or not name.strip() or any(ord(c) < 32 for c in name):
37 raise _invalid("An exact nonempty resource name is required.")
38 try:
39 return validate_search_endpoint(endpoint)
40 except ValueError as exc:
41 raise _invalid("Invalid Search endpoint.") from exc
42
43
44def retrieval_request(
45 endpoint: str, name: str, api_version: str, effort: str, query: str,
46) -> tuple[str, str, dict[str, Any]]:
47 base = _target(endpoint, name, api_version) + "/knowledgebases"
48 if not isinstance(effort, str) or effort not in {"minimal", "low", "medium"}:
49 raise _invalid("Unsupported reasoning effort.")
50 if api_version == "2026-04-01" and effort != "minimal":
51 raise _invalid("GA requires minimal retrieval.")
52 if not isinstance(query, str) or not query.strip():
53 raise _invalid("A nonempty query is required.")
54 definition_url = base + "('" + quote(name.replace("'", "''"), safe="") + "')"
55 retrieve_url = base + "/" + quote(name, safe="") + "/retrieve"
56 suffix = "?api-version=" + api_version
57 body = (
58 {"intents": [{"type": "semantic", "search": query}]}
59 if effort == "minimal" else
60 {"messages": [{"role": "user", "content": [{"type": "text", "text": query}]}]}
61 )
62 return definition_url + suffix, retrieve_url + suffix, body
63
64
65def read_json(
66 method: str, url: str, body: dict[str, Any] | None = None,
67 forward_permissions: bool = False, *,
68 token_provider: TokenProvider = azure_cli_token,
69 transport: Transport = http_request,
70) -> HttpResult:
71 token = token_provider(SEARCH_AUDIENCE)
72 headers = {"Content-Type": "application/json"}
73 if forward_permissions:
74 headers["x-ms-query-source-authorization"] = token
75 result = transport(
76 method, url, token,
77 body=json.dumps(body).encode("utf-8") if body is not None else None,
78 headers=headers, follow_redirects=False, max_response_bytes=5 * 1024 * 1024,
79 )
80 if result.status != 200 or not isinstance(result.body, dict):
81 raise HelperFailure(
82 "retrieval-incomplete", "Expected complete JSON response.",
83 blocked_at="verification", status=result.status, request_id=result.request_id,
84 )
85 return result
86
87
88def _verify_signed_in_user() -> None:
89 executable = shutil.which("az")
90 if executable is None:
91 raise HelperFailure(
92 "azure-cli-unavailable", "Azure CLI is required for keyless authentication.",
93 blocked_at="execution",
94 )
95 try:
96 result = subprocess.run(
97 [executable, "account", "show", "--query", "user.type", "--output", "tsv"],
98 check=True, capture_output=True, text=True, timeout=60,
99 )
100 except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
101 raise HelperFailure(
102 "permission-forwarding-unavailable", "Could not verify the signed-in user.",
103 blocked_at="input-resolution",
104 ) from exc
105 if result.stdout.strip().casefold() != "user":
106 raise HelperFailure(
107 "permission-forwarding-unavailable", "Permission forwarding requires a signed-in user.",
108 blocked_at="input-resolution",
109 )
110
111
112def execute(
113 document: Any, *, token_provider: TokenProvider = azure_cli_token,
114 transport: Transport = http_request,
115) -> dict[str, Any]:
116 if not isinstance(document, dict):
117 raise _invalid("Input must be a JSON object.")
118 reject_secrets(document)
119 operation = document.get("operation")
120 if not isinstance(operation, str) or operation not in {
121 "get-knowledge-base", "get-knowledge-source", "retrieve",
122 }:
123 raise _invalid("Choose get-knowledge-base, get-knowledge-source, or retrieve.")
124 fields = {"operation", "endpoint", "name", "api_version"}
125 if operation == "retrieve":
126 fields |= {"effort", "query", "forward_permissions"}
127 require_allowed_fields(document, fields, label="read-only request")
128 endpoint = _target(document.get("endpoint"), document.get("name"), document.get("api_version"))
129 name, api_version = document["name"], document["api_version"]
130 body = None
131 forwarding = False
132 if operation == "retrieve":
133 forwarding = document.get("forward_permissions")
134 if type(forwarding) is not bool:
135 raise _invalid("Explicit forward_permissions true or false is required.")
136 _, url, body = retrieval_request(
137 endpoint, name, api_version, document.get("effort"), document.get("query"),
138 )
139 if forwarding:
140 _verify_signed_in_user()
141 method = "POST"
142 else:
143 collection = "knowledgebases" if operation == "get-knowledge-base" else "knowledgesources"
144 url = endpoint + "/" + collection + "('" + quote(name.replace("'", "''"), safe="") + "')"
145 url += "?api-version=" + api_version
146 method = "GET"
147 result = read_json(
148 method, url, body, forwarding, token_provider=token_provider, transport=transport,
149 )
150 return {
151 "status": "response-received", "operation": operation, "mutation": "none",
152 "writes_performed": [], "http_status": result.status, "request_id": result.request_id,
153 "response_body": result.body, "verification": "not-performed", "cleanup": "not-applicable",
154 }
155
156
157def main(argv: list[str] | None = None) -> int:
158 parser = argparse.ArgumentParser(description="Read Search KB/source definitions or retrieve from one KB.")
159 parser.add_argument("--input", required=True, type=Path)
160 args = parser.parse_args(argv)
161 try:
162 try:
163 document = json.loads(args.input.read_text(encoding="utf-8"))
164 except OSError as exc:
165 raise HelperFailure("input-unreadable", "Input file cannot be read.", blocked_at="input-resolution") from exc
166 except (UnicodeError, json.JSONDecodeError) as exc:
167 raise HelperFailure("input-invalid-json", "Input must be UTF-8 JSON.", blocked_at="input-resolution") from exc
168 result = execute(document)
169 except HelperFailure as failure:
170 # POST retrieve is read-only even when the generic transport marks it ambiguous.
171 readonly_failure = HelperFailure(
172 failure.code, failure.message, blocked_at=failure.blocked_at,
173 status=failure.http_status, request_id=failure.request_id,
174 )
175 result = blocked_result(readonly_failure, outcome="knowledge-base-retrieval", fingerprint=None)
176 result.update(
177 mutation="none", cleanup="not-applicable",
178 safe_next_decision="Resolve the first blocker before repeating this read-only request.",
179 )
180 emit_result(result)
181 return 2
182 emit_result(result)
183 return 0
184
185
186if __name__ == "__main__":
187 sys.exit(main())