+7

main.py

··· 38 38 click.option("--model-api"), 39 39 click.option("--model-name"), 40 40 click.option("--model-api-key"), 41 41 + click.option("--model-endpoint"), 41 42 ] 42 43 43 44 ··· 59 60 model_api: Literal["anthropic", "openai", "openapi"] | None, 60 61 model_name: str | None, 61 62 model_api_key: str | None, 63 63 + model_endpoint: str | None, 62 64 ) -> tuple[Clickhouse, Osprey, ToolExecutor, Agent]: 63 65 http_client = httpx.AsyncClient() 64 66 ··· 94 96 model_api=model_api or CONFIG.model_api, 95 97 model_name=model_name or CONFIG.model_name, 96 98 model_api_key=model_api_key or CONFIG.model_api_key, 99 99 + model_endpoint=model_endpoint or CONFIG.model_endpoint or None, 97 100 tool_executor=executor, 98 101 ) 99 102 ··· 119 122 model_api: Literal["anthropic", "openai", "openapi"] | None, 120 123 model_name: str | None, 121 124 model_api_key: str | None, 125 125 + model_endpoint: str | None, 122 126 ): 123 127 clickhouse, osprey, executor, agent = build_services( 124 128 clickhouse_host=clickhouse_host, ··· 132 136 model_api=model_api, 133 137 model_name=model_name, 134 138 model_api_key=model_api_key, 139 139 + model_endpoint=model_endpoint, 135 140 ) 136 141 137 142 async def run(): ··· 161 166 model_api: Literal["anthropic", "openai", "openapi"] | None, 162 167 model_name: str | None, 163 168 model_api_key: str | None, 169 169 + model_endpoint: str | None, 164 170 ): 165 171 clickhouse, osprey, executor, agent = build_services( 166 172 clickhouse_host=clickhouse_host, ··· 174 180 model_api=model_api, 175 181 model_name=model_name, 176 182 model_api_key=model_api_key, 183 183 + model_endpoint=model_endpoint, 177 184 ) 178 185 179 186 async def run():

+224 -31

src/agent/agent.py

··· 1 1 from abc import ABC, abstractmethod 2 2 import asyncio 3 3 + import json 3 4 import logging 5 5 + from dataclasses import dataclass 4 6 from typing import Any, Literal 5 7 6 8 import anthropic 7 9 from anthropic.types import TextBlock, ToolUseBlock 8 8 - from pydantic import BaseModel 10 10 + import httpx 9 11 10 12 from src.agent.prompt import build_system_prompt 11 13 from src.tools.executor import ToolExecutor ··· 13 15 logger = logging.getLogger(__name__) 14 16 15 17 16 16 - class Message(BaseModel): 17 17 - role: Literal["user", "assistant"] 18 18 - content: str 18 18 + @dataclass 19 19 + class AgentTextBlock: 20 20 + text: str 21 21 + 22 22 + 23 23 + @dataclass 24 24 + class AgentToolUseBlock: 25 25 + id: str 26 26 + name: str 27 27 + input: dict[str, Any] 28 28 + 29 29 + 30 30 + @dataclass 31 31 + class AgentResponse: 32 32 + content: list[AgentTextBlock | AgentToolUseBlock] 33 33 + stop_reason: Literal["end_turn", "tool_use"] 34 34 + reasoning_content: str | None = None 19 35 20 36 21 37 class AgentClient(ABC): ··· 25 41 messages: list[dict[str, Any]], 26 42 system: str | None = None, 27 43 tools: list[dict[str, Any]] | None = None, 28 28 - ) -> anthropic.types.Message: 44 44 + ) -> AgentResponse: 29 45 pass 30 46 31 47 ··· 41 57 messages: list[dict[str, Any]], 42 58 system: str | None = None, 43 59 tools: list[dict[str, Any]] | None = None, 44 44 - ) -> anthropic.types.Message: 60 60 + ) -> AgentResponse: 45 61 system_text = system or build_system_prompt() 46 62 kwargs: dict[str, Any] = { 47 63 "model": self._model_name, ··· 57 73 } 58 74 59 75 if tools: 60 60 - tools = [dict(t) for t in tools] # shallow copy 76 76 + tools = [dict(t) for t in tools] 61 77 tools[-1]["cache_control"] = {"type": "ephemeral"} 62 78 kwargs["tools"] = tools 63 79 64 80 async with self._client.messages.stream(**kwargs) as stream: # type: ignore 65 65 - return await stream.get_final_message() 81 81 + msg = await stream.get_final_message() 82 82 + 83 83 + content: list[AgentTextBlock | AgentToolUseBlock] = [] 84 84 + for block in msg.content: 85 85 + if isinstance(block, TextBlock): 86 86 + content.append(AgentTextBlock(text=block.text)) 87 87 + elif isinstance(block, ToolUseBlock): 88 88 + content.append( 89 89 + AgentToolUseBlock( 90 90 + id=block.id, 91 91 + name=block.name, 92 92 + input=block.input, # type: ignore 93 93 + ) 94 94 + ) 95 95 + 96 96 + return AgentResponse( 97 97 + content=content, 98 98 + stop_reason=msg.stop_reason or "end_turn", # type: ignore TODO: fix this 99 99 + ) 100 100 + 101 101 + 102 102 + class OpenAICompatibleClient(AgentClient): 103 103 + """client for openapi compatible apis like openai, moonshot, etc""" 104 104 + 105 105 + def __init__(self, api_key: str, model_name: str, endpoint: str) -> None: 106 106 + self._api_key = api_key 107 107 + self._model_name = model_name 108 108 + self._endpoint = endpoint.rstrip("/") 109 109 + self._http = httpx.AsyncClient(timeout=300.0) 110 110 + 111 111 + async def complete( 112 112 + self, 113 113 + messages: list[dict[str, Any]], 114 114 + system: str | None = None, 115 115 + tools: list[dict[str, Any]] | None = None, 116 116 + ) -> AgentResponse: 117 117 + oai_messages = self._convert_messages(messages, system or build_system_prompt()) 118 118 + 119 119 + payload: dict[str, Any] = { 120 120 + "model": self._model_name, 121 121 + "messages": oai_messages, 122 122 + "max_tokens": 16_000, 123 123 + } 124 124 + 125 125 + if tools: 126 126 + payload["tools"] = self._convert_tools(tools) 127 127 + 128 128 + resp = await self._http.post( 129 129 + f"{self._endpoint}/chat/completions", 130 130 + headers={ 131 131 + "Authorization": f"Bearer {self._api_key}", 132 132 + "Content-Type": "application/json", 133 133 + }, 134 134 + json=payload, 135 135 + ) 136 136 + if not resp.is_success: 137 137 + logger.error( 138 138 + "API error %d: %s", resp.status_code, resp.text[:1000] 139 139 + ) 140 140 + resp.raise_for_status() 141 141 + data = resp.json() 142 142 + 143 143 + return self._parse_response(data) 144 144 + 145 145 + def _convert_messages( 146 146 + self, messages: list[dict[str, Any]], system: str 147 147 + ) -> list[dict[str, Any]]: 148 148 + """for anthropic chats, we'll convert the outputs into a similar format""" 149 149 + result: list[dict[str, Any]] = [{"role": "system", "content": system}] 150 150 + 151 151 + for msg in messages: 152 152 + role = msg["role"] 153 153 + content = msg["content"] 154 154 + 155 155 + if isinstance(content, str): 156 156 + result.append({"role": role, "content": content}) 157 157 + elif isinstance(content, list): 158 158 + if role == "assistant": 159 159 + text_parts = [] 160 160 + tool_calls = [] 161 161 + for block in content: 162 162 + if block.get("type") == "text": 163 163 + text_parts.append(block["text"]) 164 164 + elif block.get("type") == "tool_use": 165 165 + tool_calls.append( 166 166 + { 167 167 + "id": block["id"], 168 168 + "type": "function", 169 169 + "function": { 170 170 + "name": block["name"], 171 171 + "arguments": json.dumps(block["input"]), 172 172 + }, 173 173 + } 174 174 + ) 175 175 + oai_msg: dict[str, Any] = {"role": "assistant"} 176 176 + if msg.get("reasoning_content"): 177 177 + oai_msg["reasoning_content"] = msg["reasoning_content"] 178 178 + # some openai-compatible apis reject content: null on 179 179 + # assistant messages with tool_calls, so omit it when empty 180 180 + if text_parts: 181 181 + oai_msg["content"] = "\n".join(text_parts) 182 182 + else: 183 183 + oai_msg["content"] = "" 184 184 + if tool_calls: 185 185 + oai_msg["tool_calls"] = tool_calls 186 186 + result.append(oai_msg) 187 187 + elif role == "user": 188 188 + if content and content[0].get("type") == "tool_result": 189 189 + for block in content: 190 190 + result.append( 191 191 + { 192 192 + "role": "tool", 193 193 + "tool_call_id": block["tool_use_id"], 194 194 + "content": block.get("content", ""), 195 195 + } 196 196 + ) 197 197 + else: 198 198 + text = " ".join(b.get("text", str(b)) for b in content) 199 199 + result.append({"role": "user", "content": text}) 200 200 + 201 201 + return result 202 202 + 203 203 + def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]: 204 204 + """convert anthropic tool defs to oai function calling format""" 205 205 + result = [] 206 206 + for t in tools: 207 207 + func: dict[str, Any] = { 208 208 + "name": t["name"], 209 209 + "description": t.get("description", ""), 210 210 + } 211 211 + if "input_schema" in t: 212 212 + func["parameters"] = t["input_schema"] 213 213 + result.append({"type": "function", "function": func}) 214 214 + return result 215 215 + 216 216 + def _parse_response(self, data: dict[str, Any]) -> AgentResponse: 217 217 + """convert an oai chat completion resp to agentresponse""" 218 218 + choice = data["choices"][0] 219 219 + message = choice["message"] 220 220 + finish_reason = choice.get("finish_reason", "stop") 221 221 + 222 222 + content: list[AgentTextBlock | AgentToolUseBlock] = [] 223 223 + 224 224 + if message.get("content"): 225 225 + content.append(AgentTextBlock(text=message["content"])) 226 226 + 227 227 + if message.get("tool_calls"): 228 228 + for tc in message["tool_calls"]: 229 229 + try: 230 230 + args = json.loads(tc["function"]["arguments"]) 231 231 + except (json.JSONDecodeError, KeyError): 232 232 + args = {} 233 233 + content.append( 234 234 + AgentToolUseBlock( 235 235 + id=tc["id"], 236 236 + name=tc["function"]["name"], 237 237 + input=args, 238 238 + ) 239 239 + ) 240 240 + 241 241 + stop_reason = "tool_use" if finish_reason == "tool_calls" else "end_turn" 242 242 + reasoning_content = message.get("reasoning_content") 243 243 + return AgentResponse(content=content, stop_reason=stop_reason, reasoning_content=reasoning_content) 66 244 67 245 68 246 MAX_TOOL_RESULT_LENGTH = 10_000 ··· 74 252 model_api: Literal["anthropic", "openai", "openapi"], 75 253 model_name: str, 76 254 model_api_key: str | None, 255 255 + model_endpoint: str | None = None, 77 256 tool_executor: ToolExecutor | None = None, 78 257 ) -> None: 79 79 - if model_api != "anthropic": 80 80 - # TODO: implement other APIs 81 81 - raise NotImplementedError() 82 82 - 83 83 - if model_api == "anthropic": 84 84 - assert model_api_key 85 85 - self._client = AnthropicClient(api_key=model_api_key, model_name=model_name) 258 258 + match model_api: 259 259 + case "anthropic": 260 260 + assert model_api_key 261 261 + self._client: AgentClient = AnthropicClient( 262 262 + api_key=model_api_key, model_name=model_name 263 263 + ) 264 264 + case "openai": 265 265 + assert model_api_key 266 266 + self._client = OpenAICompatibleClient( 267 267 + api_key=model_api_key, 268 268 + model_name=model_name, 269 269 + endpoint="https://api.openai.com/v1", 270 270 + ) 271 271 + case "openapi": 272 272 + assert model_api_key 273 273 + assert model_endpoint, "model_endpoint is required for openapi" 274 274 + self._client = OpenAICompatibleClient( 275 275 + api_key=model_api_key, 276 276 + model_name=model_name, 277 277 + endpoint=model_endpoint, 278 278 + ) 86 279 87 280 self._tool_executor = tool_executor 88 281 self._conversation: list[dict[str, Any]] = [] ··· 94 287 return None 95 288 return [self._tool_executor.get_execute_code_tool_definition()] 96 289 97 97 - async def _handle_tool_call(self, tool_use: ToolUseBlock) -> dict[str, Any]: 290 290 + async def _handle_tool_call(self, tool_use: AgentToolUseBlock) -> dict[str, Any]: 98 291 """handle a tool call from the model""" 99 292 if tool_use.name == "execute_code" and self._tool_executor: 100 100 - code = tool_use.input.get("code", "") # type: ignore 101 101 - result = await self._tool_executor.execute_code(code) # type: ignore 293 293 + code = tool_use.input.get("code", "") 294 294 + result = await self._tool_executor.execute_code(code) 102 295 return result 103 296 else: 104 297 return {"error": f"Unknown tool: {tool_use.name}"} ··· 117 310 text_response = "" 118 311 119 312 for block in resp.content: 120 120 - if isinstance(block, TextBlock): 313 313 + if isinstance(block, AgentTextBlock): 121 314 assistant_content.append({"type": "text", "text": block.text}) 122 315 text_response += block.text 123 123 - elif isinstance(block, ToolUseBlock): 316 316 + elif isinstance(block, AgentToolUseBlock): # type: ignore TODO: for now this errors because there are no other types, but ignore for now 124 317 assistant_content.append( 125 318 { 126 319 "type": "tool_use", ··· 130 323 } 131 324 ) 132 325 133 133 - self._conversation.append( 134 134 - {"role": "assistant", "content": assistant_content} 135 135 - ) 326 326 + assistant_msg: dict[str, Any] = {"role": "assistant", "content": assistant_content} 327 327 + if resp.reasoning_content: 328 328 + assistant_msg["reasoning_content"] = resp.reasoning_content 329 329 + self._conversation.append(assistant_msg) 136 330 137 331 # find any tool calls that we need to handle 138 332 if resp.stop_reason == "tool_use": 139 333 tool_results: list[dict[str, Any]] = [] 140 334 for block in resp.content: 141 141 - if isinstance(block, ToolUseBlock): 142 142 - code = ( 143 143 - block.input.get("code", "") 144 144 - if isinstance(block.input, dict) # type: ignore 145 145 - else "" 146 146 - ) 335 335 + if isinstance(block, AgentToolUseBlock): 336 336 + code = block.input.get("code", "") 147 337 logger.info("Tool call: %s\n%s", block.name, code) 148 338 result = await self._handle_tool_call(block) 149 339 is_error = "error" in result ··· 155 345 ) 156 346 content_str = str(result) 157 347 if len(content_str) > MAX_TOOL_RESULT_LENGTH: 158 158 - content_str = content_str[:MAX_TOOL_RESULT_LENGTH] + "\n... (truncated)" 348 348 + content_str = ( 349 349 + content_str[:MAX_TOOL_RESULT_LENGTH] 350 350 + + "\n... (truncated)" 351 351 + ) 159 352 160 353 tool_results.append( 161 354 { ··· 167 360 168 361 self._conversation.append({"role": "user", "content": tool_results}) 169 362 else: 170 170 - # once there are no mroe tool calls, we proceed to the text response 363 363 + # once there are no more tool calls, we proceed to the text response 171 364 return text_response 172 365 173 366 async def run(self):

+46 -10

src/tools/executor.py

··· 130 130 tools_path = DENO_DIR / "tools.ts" 131 131 tools_path.write_text(tools_ts) 132 132 133 133 + @staticmethod 134 134 + def _kill_process(process: asyncio.subprocess.Process) -> None: 135 135 + """kill a subprocess, ignoring errors if it's already dead""" 136 136 + try: 137 137 + process.kill() 138 138 + except ProcessLookupError: 139 139 + pass 140 140 + 133 141 async def _run_deno(self, script_path: str) -> dict[str, Any]: 134 142 """run the input script in a deno subprocess""" 135 143 ··· 172 180 # calculate remaining time against the total execution deadline 173 181 remaining = deadline - asyncio.get_event_loop().time() 174 182 if remaining <= 0: 175 175 - process.kill() 183 183 + self._kill_process(process) 176 184 error = f"execution timed out after {MAX_EXECUTION_TIME:.0f} seconds (total)" 177 185 break 178 186 ··· 193 201 # track total output size to prevent stdout flooding 194 202 total_output_bytes += len(line) 195 203 if total_output_bytes > MAX_OUTPUT_SIZE: 196 196 - process.kill() 204 204 + self._kill_process(process) 197 205 error = f"output exceeded {MAX_OUTPUT_SIZE} bytes, killed" 198 206 break 199 207 ··· 208 216 if "__tool_call__" in message: 209 217 tool_call_count += 1 210 218 if tool_call_count > MAX_TOOL_CALLS: 211 211 - process.kill() 219 219 + self._kill_process(process) 212 220 error = f"exceeded maximum of {MAX_TOOL_CALLS} tool calls" 213 221 break 214 222 ··· 225 233 logger.exception(f"Tool error: {tool_name}") 226 234 response = json.dumps({"__tool_error__": str(e)}) 227 235 228 228 - process.stdin.write((response + "\n").encode()) 229 229 - await process.stdin.drain() 236 236 + try: 237 237 + process.stdin.write((response + "\n").encode()) 238 238 + await process.stdin.drain() 239 239 + except (ConnectionResetError, BrokenPipeError): 240 240 + error = f"deno process exited while sending tool result for {tool_name}" 241 241 + break 230 242 231 243 elif "__output__" in message: 232 244 outputs.append(message["__output__"]) ··· 239 251 240 252 # make sure that we kill deno subprocess if the execution times out 241 253 except asyncio.TimeoutError: 242 242 - process.kill() 254 254 + self._kill_process(process) 243 255 error = "execution timed out" 244 256 # also kill it for any other exceptions we encounter 245 257 except Exception as e: 246 246 - process.kill() 258 258 + self._kill_process(process) 247 259 error = str(e) 248 260 249 261 await process.wait() ··· 273 285 return result 274 286 275 287 def get_execute_code_tool_definition(self) -> dict[str, Any]: 276 276 - """get the anthropic tool definition for execute_code, including all the docs for available backend tools""" 288 288 + """get tool definition for execute_code, including all the docs for available backend tools""" 277 289 278 290 if self._tool_definition is not None: 279 291 return self._tool_definition ··· 290 302 {self._database_schema} 291 303 292 304 Use these exact column names when writing SQL queries. Do NOT guess column names. 305 305 + 306 306 + ## ClickHouse SQL Tips 307 307 + 308 308 + - **DateTime filtering**: The `__timestamp` column is `DateTime64(3)`. Do NOT use raw ISO strings. Use `parseDateTimeBestEffort()`: 309 309 + ```sql 310 310 + WHERE __timestamp >= parseDateTimeBestEffort('2026-02-06 04:30:00') 311 311 + ``` 312 312 + To compute a relative time in TypeScript, format it as `YYYY-MM-DD HH:MM:SS`: 313 313 + ```typescript 314 314 + const ts = new Date(Date.now() - 30 * 60 * 1000).toISOString().slice(0, 19).replace('T', ' '); 315 315 + ``` 316 316 + - **Array slicing**: ClickHouse does NOT support `array[1:5]` syntax. Use `arraySlice(array, offset, length)`: 317 317 + ```sql 318 318 + arraySlice(groupArray(DISTINCT UserId), 1, 5) as sample_accounts 319 319 + ``` 320 320 + - **Error handling**: When running multiple independent queries, use `Promise.allSettled()` instead of `Promise.all()` so one failure doesn't crash the rest. Check each result's `.status` field. 293 321 """ 294 322 295 323 osprey_section = "" ··· 329 357 330 358 Example: 331 359 ```typescript 332 332 - const result = await tools.clickhouse.query("SELECT count() FROM events"); 333 333 - output(result); 360 360 + // format a relative timestamp for ClickHouse DateTime64 columns 361 361 + const thirtyMinAgo = new Date(Date.now() - 30 * 60 * 1000).toISOString().slice(0, 19).replace('T', ' '); 362 362 + 363 363 + // run multiple independent queries safely with Promise.allSettled 364 364 + const results = await Promise.allSettled([ 365 365 + tools.clickhouse.query(`SELECT Count() as cnt FROM default.osprey_execution_results WHERE __timestamp >= parseDateTimeBestEffort('${{thirtyMinAgo}}') LIMIT 1`), 366 366 + tools.clickhouse.query(`SELECT UserId, Count() as n FROM default.osprey_execution_results WHERE __timestamp >= parseDateTimeBestEffort('${{thirtyMinAgo}}') GROUP BY UserId ORDER BY n DESC LIMIT 10`), 367 367 + ]); 368 368 + 369 369 + output(results.map(r => r.status === 'fulfilled' ? r.value : r.reason?.message)); 334 370 ``` 335 371 336 372 {tool_docs}{schema_section}{osprey_section}"""