code-puppy 0.0.214__py3-none-any.whl → 0.0.366__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (231) hide show
  1. code_puppy/__init__.py +7 -1
  2. code_puppy/agents/__init__.py +2 -0
  3. code_puppy/agents/agent_c_reviewer.py +59 -6
  4. code_puppy/agents/agent_code_puppy.py +7 -1
  5. code_puppy/agents/agent_code_reviewer.py +12 -2
  6. code_puppy/agents/agent_cpp_reviewer.py +73 -6
  7. code_puppy/agents/agent_creator_agent.py +45 -4
  8. code_puppy/agents/agent_golang_reviewer.py +92 -3
  9. code_puppy/agents/agent_javascript_reviewer.py +101 -8
  10. code_puppy/agents/agent_manager.py +81 -4
  11. code_puppy/agents/agent_pack_leader.py +383 -0
  12. code_puppy/agents/agent_planning.py +163 -0
  13. code_puppy/agents/agent_python_programmer.py +165 -0
  14. code_puppy/agents/agent_python_reviewer.py +28 -6
  15. code_puppy/agents/agent_qa_expert.py +98 -6
  16. code_puppy/agents/agent_qa_kitten.py +12 -7
  17. code_puppy/agents/agent_security_auditor.py +113 -3
  18. code_puppy/agents/agent_terminal_qa.py +323 -0
  19. code_puppy/agents/agent_typescript_reviewer.py +106 -7
  20. code_puppy/agents/base_agent.py +802 -176
  21. code_puppy/agents/event_stream_handler.py +350 -0
  22. code_puppy/agents/pack/__init__.py +34 -0
  23. code_puppy/agents/pack/bloodhound.py +304 -0
  24. code_puppy/agents/pack/husky.py +321 -0
  25. code_puppy/agents/pack/retriever.py +393 -0
  26. code_puppy/agents/pack/shepherd.py +348 -0
  27. code_puppy/agents/pack/terrier.py +287 -0
  28. code_puppy/agents/pack/watchdog.py +367 -0
  29. code_puppy/agents/prompt_reviewer.py +145 -0
  30. code_puppy/agents/subagent_stream_handler.py +276 -0
  31. code_puppy/api/__init__.py +13 -0
  32. code_puppy/api/app.py +169 -0
  33. code_puppy/api/main.py +21 -0
  34. code_puppy/api/pty_manager.py +446 -0
  35. code_puppy/api/routers/__init__.py +12 -0
  36. code_puppy/api/routers/agents.py +36 -0
  37. code_puppy/api/routers/commands.py +217 -0
  38. code_puppy/api/routers/config.py +74 -0
  39. code_puppy/api/routers/sessions.py +232 -0
  40. code_puppy/api/templates/terminal.html +361 -0
  41. code_puppy/api/websocket.py +154 -0
  42. code_puppy/callbacks.py +142 -4
  43. code_puppy/chatgpt_codex_client.py +283 -0
  44. code_puppy/claude_cache_client.py +586 -0
  45. code_puppy/cli_runner.py +916 -0
  46. code_puppy/command_line/add_model_menu.py +1079 -0
  47. code_puppy/command_line/agent_menu.py +395 -0
  48. code_puppy/command_line/attachments.py +10 -5
  49. code_puppy/command_line/autosave_menu.py +605 -0
  50. code_puppy/command_line/clipboard.py +527 -0
  51. code_puppy/command_line/colors_menu.py +520 -0
  52. code_puppy/command_line/command_handler.py +176 -738
  53. code_puppy/command_line/command_registry.py +150 -0
  54. code_puppy/command_line/config_commands.py +715 -0
  55. code_puppy/command_line/core_commands.py +792 -0
  56. code_puppy/command_line/diff_menu.py +863 -0
  57. code_puppy/command_line/load_context_completion.py +15 -22
  58. code_puppy/command_line/mcp/base.py +0 -3
  59. code_puppy/command_line/mcp/catalog_server_installer.py +175 -0
  60. code_puppy/command_line/mcp/custom_server_form.py +688 -0
  61. code_puppy/command_line/mcp/custom_server_installer.py +195 -0
  62. code_puppy/command_line/mcp/edit_command.py +148 -0
  63. code_puppy/command_line/mcp/handler.py +9 -4
  64. code_puppy/command_line/mcp/help_command.py +6 -5
  65. code_puppy/command_line/mcp/install_command.py +15 -26
  66. code_puppy/command_line/mcp/install_menu.py +685 -0
  67. code_puppy/command_line/mcp/list_command.py +2 -2
  68. code_puppy/command_line/mcp/logs_command.py +174 -65
  69. code_puppy/command_line/mcp/remove_command.py +2 -2
  70. code_puppy/command_line/mcp/restart_command.py +12 -4
  71. code_puppy/command_line/mcp/search_command.py +16 -10
  72. code_puppy/command_line/mcp/start_all_command.py +18 -6
  73. code_puppy/command_line/mcp/start_command.py +47 -25
  74. code_puppy/command_line/mcp/status_command.py +4 -5
  75. code_puppy/command_line/mcp/stop_all_command.py +7 -1
  76. code_puppy/command_line/mcp/stop_command.py +8 -4
  77. code_puppy/command_line/mcp/test_command.py +2 -2
  78. code_puppy/command_line/mcp/wizard_utils.py +20 -16
  79. code_puppy/command_line/mcp_completion.py +174 -0
  80. code_puppy/command_line/model_picker_completion.py +75 -25
  81. code_puppy/command_line/model_settings_menu.py +884 -0
  82. code_puppy/command_line/motd.py +14 -8
  83. code_puppy/command_line/onboarding_slides.py +179 -0
  84. code_puppy/command_line/onboarding_wizard.py +340 -0
  85. code_puppy/command_line/pin_command_completion.py +329 -0
  86. code_puppy/command_line/prompt_toolkit_completion.py +463 -63
  87. code_puppy/command_line/session_commands.py +296 -0
  88. code_puppy/command_line/utils.py +54 -0
  89. code_puppy/config.py +898 -112
  90. code_puppy/error_logging.py +118 -0
  91. code_puppy/gemini_code_assist.py +385 -0
  92. code_puppy/gemini_model.py +602 -0
  93. code_puppy/http_utils.py +210 -148
  94. code_puppy/keymap.py +128 -0
  95. code_puppy/main.py +5 -698
  96. code_puppy/mcp_/__init__.py +17 -0
  97. code_puppy/mcp_/async_lifecycle.py +35 -4
  98. code_puppy/mcp_/blocking_startup.py +70 -43
  99. code_puppy/mcp_/captured_stdio_server.py +2 -2
  100. code_puppy/mcp_/config_wizard.py +4 -4
  101. code_puppy/mcp_/dashboard.py +15 -6
  102. code_puppy/mcp_/managed_server.py +65 -38
  103. code_puppy/mcp_/manager.py +146 -52
  104. code_puppy/mcp_/mcp_logs.py +224 -0
  105. code_puppy/mcp_/registry.py +6 -6
  106. code_puppy/mcp_/server_registry_catalog.py +24 -5
  107. code_puppy/messaging/__init__.py +199 -2
  108. code_puppy/messaging/bus.py +610 -0
  109. code_puppy/messaging/commands.py +167 -0
  110. code_puppy/messaging/markdown_patches.py +57 -0
  111. code_puppy/messaging/message_queue.py +17 -48
  112. code_puppy/messaging/messages.py +500 -0
  113. code_puppy/messaging/queue_console.py +1 -24
  114. code_puppy/messaging/renderers.py +43 -146
  115. code_puppy/messaging/rich_renderer.py +1027 -0
  116. code_puppy/messaging/spinner/__init__.py +21 -5
  117. code_puppy/messaging/spinner/console_spinner.py +86 -51
  118. code_puppy/messaging/subagent_console.py +461 -0
  119. code_puppy/model_factory.py +634 -83
  120. code_puppy/model_utils.py +167 -0
  121. code_puppy/models.json +66 -68
  122. code_puppy/models_dev_api.json +1 -0
  123. code_puppy/models_dev_parser.py +592 -0
  124. code_puppy/plugins/__init__.py +164 -10
  125. code_puppy/plugins/antigravity_oauth/__init__.py +10 -0
  126. code_puppy/plugins/antigravity_oauth/accounts.py +406 -0
  127. code_puppy/plugins/antigravity_oauth/antigravity_model.py +704 -0
  128. code_puppy/plugins/antigravity_oauth/config.py +42 -0
  129. code_puppy/plugins/antigravity_oauth/constants.py +136 -0
  130. code_puppy/plugins/antigravity_oauth/oauth.py +478 -0
  131. code_puppy/plugins/antigravity_oauth/register_callbacks.py +406 -0
  132. code_puppy/plugins/antigravity_oauth/storage.py +271 -0
  133. code_puppy/plugins/antigravity_oauth/test_plugin.py +319 -0
  134. code_puppy/plugins/antigravity_oauth/token.py +167 -0
  135. code_puppy/plugins/antigravity_oauth/transport.py +767 -0
  136. code_puppy/plugins/antigravity_oauth/utils.py +169 -0
  137. code_puppy/plugins/chatgpt_oauth/__init__.py +8 -0
  138. code_puppy/plugins/chatgpt_oauth/config.py +52 -0
  139. code_puppy/plugins/chatgpt_oauth/oauth_flow.py +328 -0
  140. code_puppy/plugins/chatgpt_oauth/register_callbacks.py +94 -0
  141. code_puppy/plugins/chatgpt_oauth/test_plugin.py +293 -0
  142. code_puppy/plugins/chatgpt_oauth/utils.py +489 -0
  143. code_puppy/plugins/claude_code_oauth/README.md +167 -0
  144. code_puppy/plugins/claude_code_oauth/SETUP.md +93 -0
  145. code_puppy/plugins/claude_code_oauth/__init__.py +6 -0
  146. code_puppy/plugins/claude_code_oauth/config.py +50 -0
  147. code_puppy/plugins/claude_code_oauth/register_callbacks.py +308 -0
  148. code_puppy/plugins/claude_code_oauth/test_plugin.py +283 -0
  149. code_puppy/plugins/claude_code_oauth/utils.py +518 -0
  150. code_puppy/plugins/customizable_commands/__init__.py +0 -0
  151. code_puppy/plugins/customizable_commands/register_callbacks.py +169 -0
  152. code_puppy/plugins/example_custom_command/README.md +280 -0
  153. code_puppy/plugins/example_custom_command/register_callbacks.py +2 -2
  154. code_puppy/plugins/file_permission_handler/__init__.py +4 -0
  155. code_puppy/plugins/file_permission_handler/register_callbacks.py +523 -0
  156. code_puppy/plugins/frontend_emitter/__init__.py +25 -0
  157. code_puppy/plugins/frontend_emitter/emitter.py +121 -0
  158. code_puppy/plugins/frontend_emitter/register_callbacks.py +261 -0
  159. code_puppy/plugins/oauth_puppy_html.py +228 -0
  160. code_puppy/plugins/shell_safety/__init__.py +6 -0
  161. code_puppy/plugins/shell_safety/agent_shell_safety.py +69 -0
  162. code_puppy/plugins/shell_safety/command_cache.py +156 -0
  163. code_puppy/plugins/shell_safety/register_callbacks.py +202 -0
  164. code_puppy/prompts/antigravity_system_prompt.md +1 -0
  165. code_puppy/prompts/codex_system_prompt.md +310 -0
  166. code_puppy/pydantic_patches.py +131 -0
  167. code_puppy/reopenable_async_client.py +8 -8
  168. code_puppy/round_robin_model.py +9 -12
  169. code_puppy/session_storage.py +2 -1
  170. code_puppy/status_display.py +21 -4
  171. code_puppy/summarization_agent.py +41 -13
  172. code_puppy/terminal_utils.py +418 -0
  173. code_puppy/tools/__init__.py +37 -1
  174. code_puppy/tools/agent_tools.py +536 -52
  175. code_puppy/tools/browser/__init__.py +37 -0
  176. code_puppy/tools/browser/browser_control.py +19 -23
  177. code_puppy/tools/browser/browser_interactions.py +41 -48
  178. code_puppy/tools/browser/browser_locators.py +36 -38
  179. code_puppy/tools/browser/browser_manager.py +316 -0
  180. code_puppy/tools/browser/browser_navigation.py +16 -16
  181. code_puppy/tools/browser/browser_screenshot.py +79 -143
  182. code_puppy/tools/browser/browser_scripts.py +32 -42
  183. code_puppy/tools/browser/browser_workflows.py +44 -27
  184. code_puppy/tools/browser/chromium_terminal_manager.py +259 -0
  185. code_puppy/tools/browser/terminal_command_tools.py +521 -0
  186. code_puppy/tools/browser/terminal_screenshot_tools.py +556 -0
  187. code_puppy/tools/browser/terminal_tools.py +525 -0
  188. code_puppy/tools/command_runner.py +930 -147
  189. code_puppy/tools/common.py +1113 -5
  190. code_puppy/tools/display.py +84 -0
  191. code_puppy/tools/file_modifications.py +288 -89
  192. code_puppy/tools/file_operations.py +226 -154
  193. code_puppy/tools/subagent_context.py +158 -0
  194. code_puppy/uvx_detection.py +242 -0
  195. code_puppy/version_checker.py +30 -11
  196. code_puppy-0.0.366.data/data/code_puppy/models.json +110 -0
  197. code_puppy-0.0.366.data/data/code_puppy/models_dev_api.json +1 -0
  198. {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/METADATA +149 -75
  199. code_puppy-0.0.366.dist-info/RECORD +217 -0
  200. {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/WHEEL +1 -1
  201. code_puppy/command_line/mcp/add_command.py +0 -183
  202. code_puppy/messaging/spinner/textual_spinner.py +0 -106
  203. code_puppy/tools/browser/camoufox_manager.py +0 -216
  204. code_puppy/tools/browser/vqa_agent.py +0 -70
  205. code_puppy/tui/__init__.py +0 -10
  206. code_puppy/tui/app.py +0 -1105
  207. code_puppy/tui/components/__init__.py +0 -21
  208. code_puppy/tui/components/chat_view.py +0 -551
  209. code_puppy/tui/components/command_history_modal.py +0 -218
  210. code_puppy/tui/components/copy_button.py +0 -139
  211. code_puppy/tui/components/custom_widgets.py +0 -63
  212. code_puppy/tui/components/human_input_modal.py +0 -175
  213. code_puppy/tui/components/input_area.py +0 -167
  214. code_puppy/tui/components/sidebar.py +0 -309
  215. code_puppy/tui/components/status_bar.py +0 -185
  216. code_puppy/tui/messages.py +0 -27
  217. code_puppy/tui/models/__init__.py +0 -8
  218. code_puppy/tui/models/chat_message.py +0 -25
  219. code_puppy/tui/models/command_history.py +0 -89
  220. code_puppy/tui/models/enums.py +0 -24
  221. code_puppy/tui/screens/__init__.py +0 -17
  222. code_puppy/tui/screens/autosave_picker.py +0 -175
  223. code_puppy/tui/screens/help.py +0 -130
  224. code_puppy/tui/screens/mcp_install_wizard.py +0 -803
  225. code_puppy/tui/screens/settings.py +0 -306
  226. code_puppy/tui/screens/tools.py +0 -74
  227. code_puppy/tui_state.py +0 -55
  228. code_puppy-0.0.214.data/data/code_puppy/models.json +0 -112
  229. code_puppy-0.0.214.dist-info/RECORD +0 -131
  230. {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/entry_points.txt +0 -0
  231. {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,118 @@
1
+ """Error logging utility for code_puppy.
2
+
3
+ Logs unexpected errors to XDG_STATE_HOME/code_puppy/logs/ for debugging purposes.
4
+ Per XDG spec, logs are "state data" (actions history), not configuration.
5
+ Because even good puppies make mistakes sometimes! 🐶
6
+ """
7
+
8
+ import os
9
+ import traceback
10
+ from datetime import datetime
11
+ from pathlib import Path
12
+ from typing import Optional
13
+
14
+ from code_puppy.config import STATE_DIR
15
+
16
+ # Logs directory within the state directory (per XDG spec, logs are state data)
17
+ LOGS_DIR = os.path.join(STATE_DIR, "logs")
18
+ ERROR_LOG_FILE = os.path.join(LOGS_DIR, "errors.log")
19
+
20
+
21
+ def _ensure_logs_dir() -> None:
22
+ """Create the logs directory if it doesn't exist (with 0700 perms per XDG spec)."""
23
+ Path(LOGS_DIR).mkdir(parents=True, exist_ok=True, mode=0o700)
24
+
25
+
26
+ def log_error(
27
+ error: Exception,
28
+ context: Optional[str] = None,
29
+ include_traceback: bool = True,
30
+ ) -> None:
31
+ """Log an error to the error log file.
32
+
33
+ Args:
34
+ error: The exception to log
35
+ context: Optional context string describing where the error occurred
36
+ include_traceback: Whether to include the full traceback (default True)
37
+ """
38
+ try:
39
+ _ensure_logs_dir()
40
+
41
+ timestamp = datetime.now().isoformat()
42
+ error_type = type(error).__name__
43
+ error_msg = str(error)
44
+
45
+ log_entry_parts = [
46
+ f"\n{'=' * 80}",
47
+ f"Timestamp: {timestamp}",
48
+ f"Error Type: {error_type}",
49
+ f"Error Message: {error_msg}",
50
+ ]
51
+
52
+ if context:
53
+ log_entry_parts.append(f"Context: {context}")
54
+
55
+ if include_traceback:
56
+ tb = traceback.format_exception(type(error), error, error.__traceback__)
57
+ log_entry_parts.append(f"Traceback:\n{''.join(tb)}")
58
+
59
+ if hasattr(error, "args") and error.args:
60
+ log_entry_parts.append(f"Args: {error.args}")
61
+
62
+ log_entry_parts.append(f"{'=' * 80}\n")
63
+
64
+ log_entry = "\n".join(log_entry_parts)
65
+
66
+ with open(ERROR_LOG_FILE, "a", encoding="utf-8") as f:
67
+ f.write(log_entry)
68
+
69
+ except Exception:
70
+ # If we can't log, we silently fail - don't want logging errors
71
+ # to cause more problems than they solve!
72
+ pass
73
+
74
+
75
+ def log_error_message(
76
+ message: str,
77
+ context: Optional[str] = None,
78
+ ) -> None:
79
+ """Log a simple error message without an exception object.
80
+
81
+ Args:
82
+ message: The error message to log
83
+ context: Optional context string describing where the error occurred
84
+ """
85
+ try:
86
+ _ensure_logs_dir()
87
+
88
+ timestamp = datetime.now().isoformat()
89
+
90
+ log_entry_parts = [
91
+ f"\n{'=' * 80}",
92
+ f"Timestamp: {timestamp}",
93
+ f"Message: {message}",
94
+ ]
95
+
96
+ if context:
97
+ log_entry_parts.append(f"Context: {context}")
98
+
99
+ log_entry_parts.append(f"{'=' * 80}\n")
100
+
101
+ log_entry = "\n".join(log_entry_parts)
102
+
103
+ with open(ERROR_LOG_FILE, "a", encoding="utf-8") as f:
104
+ f.write(log_entry)
105
+
106
+ except Exception:
107
+ # Silent fail - same reasoning as above
108
+ pass
109
+
110
+
111
+ def get_log_file_path() -> str:
112
+ """Return the path to the error log file."""
113
+ return ERROR_LOG_FILE
114
+
115
+
116
+ def get_logs_dir() -> str:
117
+ """Return the path to the logs directory."""
118
+ return LOGS_DIR
@@ -0,0 +1,385 @@
1
+ """Gemini Code Assist Model for pydantic_ai.
2
+
3
+ This module provides a custom Model implementation that uses Google's
4
+ Code Assist API (cloudcode-pa.googleapis.com) instead of the standard
5
+ Generative Language API. The Code Assist API supports OAuth authentication
6
+ and has a different request/response format.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import json
12
+ import logging
13
+ import uuid
14
+ from collections.abc import AsyncIterator
15
+ from contextlib import asynccontextmanager
16
+ from datetime import datetime, timezone
17
+ from typing import Any, Dict, Optional
18
+
19
+ import httpx
20
+ from pydantic_ai.messages import (
21
+ ModelMessage,
22
+ ModelRequest,
23
+ ModelResponse,
24
+ ModelResponsePart,
25
+ SystemPromptPart,
26
+ TextPart,
27
+ ToolCallPart,
28
+ ToolReturnPart,
29
+ UserPromptPart,
30
+ )
31
+ from pydantic_ai.models import Model, ModelRequestParameters
32
+ from pydantic_ai.settings import ModelSettings
33
+ from pydantic_ai.tools import ToolDefinition
34
+ from pydantic_ai.usage import RequestUsage
35
+
36
+ logger = logging.getLogger(__name__)
37
+
38
+
39
+ class GeminiCodeAssistModel(Model):
40
+ """Model implementation for Google's Code Assist API.
41
+
42
+ This uses the cloudcode-pa.googleapis.com endpoint which accepts OAuth
43
+ tokens and has a wrapped request/response format.
44
+ """
45
+
46
+ def __init__(
47
+ self,
48
+ model_name: str,
49
+ access_token: str,
50
+ project_id: str,
51
+ api_base_url: str = "https://cloudcode-pa.googleapis.com",
52
+ api_version: str = "v1internal",
53
+ ):
54
+ self._model_name = model_name
55
+ self.access_token = access_token
56
+ self.project_id = project_id
57
+ self.api_base_url = api_base_url
58
+ self.api_version = api_version
59
+
60
+ def model_name(self) -> str:
61
+ """Return the model name."""
62
+ return self._model_name
63
+
64
+ @property
65
+ def system(self) -> str:
66
+ return "google"
67
+
68
+ async def request(
69
+ self,
70
+ messages: list[ModelMessage],
71
+ model_settings: ModelSettings | None,
72
+ model_request_parameters: ModelRequestParameters,
73
+ ) -> ModelResponse:
74
+ """Make a non-streaming request to the Code Assist API."""
75
+ request_body = self._build_request(
76
+ messages, model_settings, model_request_parameters
77
+ )
78
+
79
+ url = f"{self.api_base_url}/{self.api_version}:generateContent"
80
+ headers = self._get_headers()
81
+
82
+ async with httpx.AsyncClient(timeout=180) as client:
83
+ response = await client.post(url, json=request_body, headers=headers)
84
+
85
+ if response.status_code != 200:
86
+ error_text = response.text
87
+ raise RuntimeError(
88
+ f"Code Assist API error {response.status_code}: {error_text}"
89
+ )
90
+
91
+ data = response.json()
92
+
93
+ return self._parse_response(data)
94
+
95
+ @asynccontextmanager
96
+ async def request_stream(
97
+ self,
98
+ messages: list[ModelMessage],
99
+ model_settings: ModelSettings | None,
100
+ model_request_parameters: ModelRequestParameters,
101
+ ) -> AsyncIterator[StreamedResponse]:
102
+ """Make a streaming request to the Code Assist API."""
103
+ request_body = self._build_request(
104
+ messages, model_settings, model_request_parameters
105
+ )
106
+
107
+ url = f"{self.api_base_url}/{self.api_version}:streamGenerateContent?alt=sse"
108
+ headers = self._get_headers()
109
+
110
+ async with httpx.AsyncClient(timeout=180) as client:
111
+ async with client.stream(
112
+ "POST", url, json=request_body, headers=headers
113
+ ) as response:
114
+ if response.status_code != 200:
115
+ error_text = await response.aread()
116
+ raise RuntimeError(
117
+ f"Code Assist API error {response.status_code}: {error_text.decode()}"
118
+ )
119
+
120
+ yield StreamedResponse(response, self._model_name)
121
+
122
+ def _get_headers(self) -> Dict[str, str]:
123
+ """Get HTTP headers for the request."""
124
+ return {
125
+ "Authorization": f"Bearer {self.access_token}",
126
+ "Content-Type": "application/json",
127
+ "Accept": "application/json",
128
+ }
129
+
130
+ def _build_request(
131
+ self,
132
+ messages: list[ModelMessage],
133
+ model_settings: ModelSettings | None,
134
+ model_request_parameters: ModelRequestParameters,
135
+ ) -> Dict[str, Any]:
136
+ """Build the Code Assist API request body."""
137
+ contents = []
138
+ system_instruction = None
139
+
140
+ for msg in messages:
141
+ if isinstance(msg, ModelRequest):
142
+ for part in msg.parts:
143
+ if isinstance(part, SystemPromptPart):
144
+ # Collect system prompt
145
+ if system_instruction is None:
146
+ system_instruction = {
147
+ "role": "user",
148
+ "parts": [{"text": part.content}],
149
+ }
150
+ else:
151
+ system_instruction["parts"].append({"text": part.content})
152
+ elif isinstance(part, UserPromptPart):
153
+ contents.append(
154
+ {
155
+ "role": "user",
156
+ "parts": [{"text": part.content}],
157
+ }
158
+ )
159
+ elif isinstance(part, ToolReturnPart):
160
+ # Serialize content to string if it's not already
161
+ content = part.content
162
+ if not isinstance(content, (str, int, float, bool, type(None))):
163
+ try:
164
+ content = json.dumps(content, default=str)
165
+ except (TypeError, ValueError):
166
+ content = str(content)
167
+ contents.append(
168
+ {
169
+ "role": "user",
170
+ "parts": [
171
+ {
172
+ "functionResponse": {
173
+ "name": part.tool_name,
174
+ "response": {"result": content},
175
+ }
176
+ }
177
+ ],
178
+ }
179
+ )
180
+ elif isinstance(msg, ModelResponse):
181
+ parts = []
182
+ first_func_call = True
183
+ for part in msg.parts:
184
+ if isinstance(part, TextPart):
185
+ parts.append({"text": part.content})
186
+ elif isinstance(part, ToolCallPart):
187
+ func_call_part = {
188
+ "functionCall": {
189
+ "name": part.tool_name,
190
+ "args": part.args_as_dict(),
191
+ }
192
+ }
193
+ # Code Assist API requires thoughtSignature on function calls
194
+ # Use synthetic signature to skip validation
195
+ if first_func_call:
196
+ func_call_part["thoughtSignature"] = (
197
+ "skip_thought_signature_validator"
198
+ )
199
+ first_func_call = False
200
+ parts.append(func_call_part)
201
+ if parts:
202
+ contents.append({"role": "model", "parts": parts})
203
+
204
+ # Build the inner request (Vertex-style format)
205
+ inner_request: Dict[str, Any] = {
206
+ "contents": contents,
207
+ }
208
+
209
+ if system_instruction:
210
+ inner_request["systemInstruction"] = system_instruction
211
+
212
+ # Add tools if available
213
+ if model_request_parameters.function_tools:
214
+ inner_request["tools"] = [
215
+ self._build_tools(model_request_parameters.function_tools)
216
+ ]
217
+
218
+ # Add generation config
219
+ generation_config = self._build_generation_config(model_settings)
220
+ if generation_config:
221
+ inner_request["generationConfig"] = generation_config
222
+
223
+ # Wrap in Code Assist format
224
+ return {
225
+ "model": self._model_name,
226
+ "project": self.project_id,
227
+ "user_prompt_id": str(uuid.uuid4()),
228
+ "request": inner_request,
229
+ }
230
+
231
+ def _build_tools(self, tools: list[ToolDefinition]) -> Dict[str, Any]:
232
+ """Build tool definitions for the API."""
233
+ function_declarations = []
234
+
235
+ for tool in tools:
236
+ func_decl: Dict[str, Any] = {
237
+ "name": tool.name,
238
+ "description": tool.description or "",
239
+ }
240
+
241
+ if tool.parameters_json_schema:
242
+ func_decl["parametersJsonSchema"] = tool.parameters_json_schema
243
+
244
+ function_declarations.append(func_decl)
245
+
246
+ return {"functionDeclarations": function_declarations}
247
+
248
+ def _build_generation_config(
249
+ self, model_settings: ModelSettings | None
250
+ ) -> Optional[Dict[str, Any]]:
251
+ """Build generation config from model settings."""
252
+ if not model_settings:
253
+ return None
254
+
255
+ config: Dict[str, Any] = {}
256
+
257
+ if (
258
+ hasattr(model_settings, "temperature")
259
+ and model_settings.temperature is not None
260
+ ):
261
+ config["temperature"] = model_settings.temperature
262
+
263
+ if hasattr(model_settings, "top_p") and model_settings.top_p is not None:
264
+ config["topP"] = model_settings.top_p
265
+
266
+ if (
267
+ hasattr(model_settings, "max_tokens")
268
+ and model_settings.max_tokens is not None
269
+ ):
270
+ config["maxOutputTokens"] = model_settings.max_tokens
271
+
272
+ return config if config else None
273
+
274
+ def _parse_response(self, data: Dict[str, Any]) -> ModelResponse:
275
+ """Parse the Code Assist API response."""
276
+ # Unwrap the Code Assist response format
277
+ inner_response = data.get("response", data)
278
+
279
+ candidates = inner_response.get("candidates", [])
280
+ if not candidates:
281
+ raise RuntimeError("No candidates in response")
282
+
283
+ candidate = candidates[0]
284
+ content = candidate.get("content", {})
285
+ parts = content.get("parts", [])
286
+
287
+ response_parts: list[ModelResponsePart] = []
288
+
289
+ for part in parts:
290
+ if "text" in part:
291
+ response_parts.append(TextPart(content=part["text"]))
292
+ elif "functionCall" in part:
293
+ func_call = part["functionCall"]
294
+ response_parts.append(
295
+ ToolCallPart(
296
+ tool_name=func_call["name"],
297
+ args=func_call.get("args", {}),
298
+ tool_call_id=str(uuid.uuid4()),
299
+ )
300
+ )
301
+
302
+ # Extract usage metadata
303
+ usage_meta = inner_response.get("usageMetadata", {})
304
+ usage = RequestUsage(
305
+ input_tokens=usage_meta.get("promptTokenCount", 0),
306
+ output_tokens=usage_meta.get("candidatesTokenCount", 0),
307
+ )
308
+
309
+ return ModelResponse(
310
+ parts=response_parts, model_name=self._model_name, usage=usage
311
+ )
312
+
313
+
314
+ class StreamedResponse:
315
+ """Handler for streaming responses from Code Assist API."""
316
+
317
+ def __init__(self, response: httpx.Response, model_name: str):
318
+ self._response = response
319
+ self._model_name = model_name
320
+ self._usage: Optional[RequestUsage] = None
321
+ self._timestamp = datetime.now(timezone.utc)
322
+
323
+ def __aiter__(self) -> AsyncIterator[str]:
324
+ return self._iter_chunks()
325
+
326
+ async def _iter_chunks(self) -> AsyncIterator[str]:
327
+ """Iterate over SSE chunks from the response."""
328
+ async for line in self._response.aiter_lines():
329
+ line = line.strip()
330
+
331
+ if line.startswith("data: "):
332
+ data_str = line[6:]
333
+ if data_str == "[DONE]":
334
+ break
335
+
336
+ try:
337
+ data = json.loads(data_str)
338
+ # Unwrap Code Assist format
339
+ inner = data.get("response", data)
340
+
341
+ # Extract usage if available
342
+ if "usageMetadata" in inner:
343
+ meta = inner["usageMetadata"]
344
+ self._usage = RequestUsage(
345
+ input_tokens=meta.get("promptTokenCount", 0),
346
+ output_tokens=meta.get("candidatesTokenCount", 0),
347
+ )
348
+
349
+ # Extract text from candidates
350
+ for candidate in inner.get("candidates", []):
351
+ content = candidate.get("content", {})
352
+ for part in content.get("parts", []):
353
+ if "text" in part:
354
+ yield part["text"]
355
+
356
+ except json.JSONDecodeError:
357
+ logger.warning("Failed to parse SSE data: %s", data_str)
358
+ continue
359
+
360
+ async def get_response_parts(self) -> list[ModelResponsePart]:
361
+ """Get all response parts after streaming is complete."""
362
+ text_content = ""
363
+ tool_calls = []
364
+
365
+ async for chunk in self:
366
+ text_content += chunk
367
+
368
+ parts: list[ModelResponsePart] = []
369
+ if text_content:
370
+ parts.append(TextPart(content=text_content))
371
+ parts.extend(tool_calls)
372
+
373
+ return parts
374
+
375
+ def usage(self) -> RequestUsage:
376
+ """Get usage statistics."""
377
+ return self._usage or RequestUsage()
378
+
379
+ def model_name(self) -> str:
380
+ """Get the model name."""
381
+ return self._model_name
382
+
383
+ def timestamp(self) -> datetime:
384
+ """Get the response timestamp."""
385
+ return self._timestamp