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.
- code_puppy/__init__.py +7 -1
- code_puppy/agents/__init__.py +2 -0
- code_puppy/agents/agent_c_reviewer.py +59 -6
- code_puppy/agents/agent_code_puppy.py +7 -1
- code_puppy/agents/agent_code_reviewer.py +12 -2
- code_puppy/agents/agent_cpp_reviewer.py +73 -6
- code_puppy/agents/agent_creator_agent.py +45 -4
- code_puppy/agents/agent_golang_reviewer.py +92 -3
- code_puppy/agents/agent_javascript_reviewer.py +101 -8
- code_puppy/agents/agent_manager.py +81 -4
- code_puppy/agents/agent_pack_leader.py +383 -0
- code_puppy/agents/agent_planning.py +163 -0
- code_puppy/agents/agent_python_programmer.py +165 -0
- code_puppy/agents/agent_python_reviewer.py +28 -6
- code_puppy/agents/agent_qa_expert.py +98 -6
- code_puppy/agents/agent_qa_kitten.py +12 -7
- code_puppy/agents/agent_security_auditor.py +113 -3
- code_puppy/agents/agent_terminal_qa.py +323 -0
- code_puppy/agents/agent_typescript_reviewer.py +106 -7
- code_puppy/agents/base_agent.py +802 -176
- code_puppy/agents/event_stream_handler.py +350 -0
- code_puppy/agents/pack/__init__.py +34 -0
- code_puppy/agents/pack/bloodhound.py +304 -0
- code_puppy/agents/pack/husky.py +321 -0
- code_puppy/agents/pack/retriever.py +393 -0
- code_puppy/agents/pack/shepherd.py +348 -0
- code_puppy/agents/pack/terrier.py +287 -0
- code_puppy/agents/pack/watchdog.py +367 -0
- code_puppy/agents/prompt_reviewer.py +145 -0
- code_puppy/agents/subagent_stream_handler.py +276 -0
- code_puppy/api/__init__.py +13 -0
- code_puppy/api/app.py +169 -0
- code_puppy/api/main.py +21 -0
- code_puppy/api/pty_manager.py +446 -0
- code_puppy/api/routers/__init__.py +12 -0
- code_puppy/api/routers/agents.py +36 -0
- code_puppy/api/routers/commands.py +217 -0
- code_puppy/api/routers/config.py +74 -0
- code_puppy/api/routers/sessions.py +232 -0
- code_puppy/api/templates/terminal.html +361 -0
- code_puppy/api/websocket.py +154 -0
- code_puppy/callbacks.py +142 -4
- code_puppy/chatgpt_codex_client.py +283 -0
- code_puppy/claude_cache_client.py +586 -0
- code_puppy/cli_runner.py +916 -0
- code_puppy/command_line/add_model_menu.py +1079 -0
- code_puppy/command_line/agent_menu.py +395 -0
- code_puppy/command_line/attachments.py +10 -5
- code_puppy/command_line/autosave_menu.py +605 -0
- code_puppy/command_line/clipboard.py +527 -0
- code_puppy/command_line/colors_menu.py +520 -0
- code_puppy/command_line/command_handler.py +176 -738
- code_puppy/command_line/command_registry.py +150 -0
- code_puppy/command_line/config_commands.py +715 -0
- code_puppy/command_line/core_commands.py +792 -0
- code_puppy/command_line/diff_menu.py +863 -0
- code_puppy/command_line/load_context_completion.py +15 -22
- code_puppy/command_line/mcp/base.py +0 -3
- code_puppy/command_line/mcp/catalog_server_installer.py +175 -0
- code_puppy/command_line/mcp/custom_server_form.py +688 -0
- code_puppy/command_line/mcp/custom_server_installer.py +195 -0
- code_puppy/command_line/mcp/edit_command.py +148 -0
- code_puppy/command_line/mcp/handler.py +9 -4
- code_puppy/command_line/mcp/help_command.py +6 -5
- code_puppy/command_line/mcp/install_command.py +15 -26
- code_puppy/command_line/mcp/install_menu.py +685 -0
- code_puppy/command_line/mcp/list_command.py +2 -2
- code_puppy/command_line/mcp/logs_command.py +174 -65
- code_puppy/command_line/mcp/remove_command.py +2 -2
- code_puppy/command_line/mcp/restart_command.py +12 -4
- code_puppy/command_line/mcp/search_command.py +16 -10
- code_puppy/command_line/mcp/start_all_command.py +18 -6
- code_puppy/command_line/mcp/start_command.py +47 -25
- code_puppy/command_line/mcp/status_command.py +4 -5
- code_puppy/command_line/mcp/stop_all_command.py +7 -1
- code_puppy/command_line/mcp/stop_command.py +8 -4
- code_puppy/command_line/mcp/test_command.py +2 -2
- code_puppy/command_line/mcp/wizard_utils.py +20 -16
- code_puppy/command_line/mcp_completion.py +174 -0
- code_puppy/command_line/model_picker_completion.py +75 -25
- code_puppy/command_line/model_settings_menu.py +884 -0
- code_puppy/command_line/motd.py +14 -8
- code_puppy/command_line/onboarding_slides.py +179 -0
- code_puppy/command_line/onboarding_wizard.py +340 -0
- code_puppy/command_line/pin_command_completion.py +329 -0
- code_puppy/command_line/prompt_toolkit_completion.py +463 -63
- code_puppy/command_line/session_commands.py +296 -0
- code_puppy/command_line/utils.py +54 -0
- code_puppy/config.py +898 -112
- code_puppy/error_logging.py +118 -0
- code_puppy/gemini_code_assist.py +385 -0
- code_puppy/gemini_model.py +602 -0
- code_puppy/http_utils.py +210 -148
- code_puppy/keymap.py +128 -0
- code_puppy/main.py +5 -698
- code_puppy/mcp_/__init__.py +17 -0
- code_puppy/mcp_/async_lifecycle.py +35 -4
- code_puppy/mcp_/blocking_startup.py +70 -43
- code_puppy/mcp_/captured_stdio_server.py +2 -2
- code_puppy/mcp_/config_wizard.py +4 -4
- code_puppy/mcp_/dashboard.py +15 -6
- code_puppy/mcp_/managed_server.py +65 -38
- code_puppy/mcp_/manager.py +146 -52
- code_puppy/mcp_/mcp_logs.py +224 -0
- code_puppy/mcp_/registry.py +6 -6
- code_puppy/mcp_/server_registry_catalog.py +24 -5
- code_puppy/messaging/__init__.py +199 -2
- code_puppy/messaging/bus.py +610 -0
- code_puppy/messaging/commands.py +167 -0
- code_puppy/messaging/markdown_patches.py +57 -0
- code_puppy/messaging/message_queue.py +17 -48
- code_puppy/messaging/messages.py +500 -0
- code_puppy/messaging/queue_console.py +1 -24
- code_puppy/messaging/renderers.py +43 -146
- code_puppy/messaging/rich_renderer.py +1027 -0
- code_puppy/messaging/spinner/__init__.py +21 -5
- code_puppy/messaging/spinner/console_spinner.py +86 -51
- code_puppy/messaging/subagent_console.py +461 -0
- code_puppy/model_factory.py +634 -83
- code_puppy/model_utils.py +167 -0
- code_puppy/models.json +66 -68
- code_puppy/models_dev_api.json +1 -0
- code_puppy/models_dev_parser.py +592 -0
- code_puppy/plugins/__init__.py +164 -10
- code_puppy/plugins/antigravity_oauth/__init__.py +10 -0
- code_puppy/plugins/antigravity_oauth/accounts.py +406 -0
- code_puppy/plugins/antigravity_oauth/antigravity_model.py +704 -0
- code_puppy/plugins/antigravity_oauth/config.py +42 -0
- code_puppy/plugins/antigravity_oauth/constants.py +136 -0
- code_puppy/plugins/antigravity_oauth/oauth.py +478 -0
- code_puppy/plugins/antigravity_oauth/register_callbacks.py +406 -0
- code_puppy/plugins/antigravity_oauth/storage.py +271 -0
- code_puppy/plugins/antigravity_oauth/test_plugin.py +319 -0
- code_puppy/plugins/antigravity_oauth/token.py +167 -0
- code_puppy/plugins/antigravity_oauth/transport.py +767 -0
- code_puppy/plugins/antigravity_oauth/utils.py +169 -0
- code_puppy/plugins/chatgpt_oauth/__init__.py +8 -0
- code_puppy/plugins/chatgpt_oauth/config.py +52 -0
- code_puppy/plugins/chatgpt_oauth/oauth_flow.py +328 -0
- code_puppy/plugins/chatgpt_oauth/register_callbacks.py +94 -0
- code_puppy/plugins/chatgpt_oauth/test_plugin.py +293 -0
- code_puppy/plugins/chatgpt_oauth/utils.py +489 -0
- code_puppy/plugins/claude_code_oauth/README.md +167 -0
- code_puppy/plugins/claude_code_oauth/SETUP.md +93 -0
- code_puppy/plugins/claude_code_oauth/__init__.py +6 -0
- code_puppy/plugins/claude_code_oauth/config.py +50 -0
- code_puppy/plugins/claude_code_oauth/register_callbacks.py +308 -0
- code_puppy/plugins/claude_code_oauth/test_plugin.py +283 -0
- code_puppy/plugins/claude_code_oauth/utils.py +518 -0
- code_puppy/plugins/customizable_commands/__init__.py +0 -0
- code_puppy/plugins/customizable_commands/register_callbacks.py +169 -0
- code_puppy/plugins/example_custom_command/README.md +280 -0
- code_puppy/plugins/example_custom_command/register_callbacks.py +2 -2
- code_puppy/plugins/file_permission_handler/__init__.py +4 -0
- code_puppy/plugins/file_permission_handler/register_callbacks.py +523 -0
- code_puppy/plugins/frontend_emitter/__init__.py +25 -0
- code_puppy/plugins/frontend_emitter/emitter.py +121 -0
- code_puppy/plugins/frontend_emitter/register_callbacks.py +261 -0
- code_puppy/plugins/oauth_puppy_html.py +228 -0
- code_puppy/plugins/shell_safety/__init__.py +6 -0
- code_puppy/plugins/shell_safety/agent_shell_safety.py +69 -0
- code_puppy/plugins/shell_safety/command_cache.py +156 -0
- code_puppy/plugins/shell_safety/register_callbacks.py +202 -0
- code_puppy/prompts/antigravity_system_prompt.md +1 -0
- code_puppy/prompts/codex_system_prompt.md +310 -0
- code_puppy/pydantic_patches.py +131 -0
- code_puppy/reopenable_async_client.py +8 -8
- code_puppy/round_robin_model.py +9 -12
- code_puppy/session_storage.py +2 -1
- code_puppy/status_display.py +21 -4
- code_puppy/summarization_agent.py +41 -13
- code_puppy/terminal_utils.py +418 -0
- code_puppy/tools/__init__.py +37 -1
- code_puppy/tools/agent_tools.py +536 -52
- code_puppy/tools/browser/__init__.py +37 -0
- code_puppy/tools/browser/browser_control.py +19 -23
- code_puppy/tools/browser/browser_interactions.py +41 -48
- code_puppy/tools/browser/browser_locators.py +36 -38
- code_puppy/tools/browser/browser_manager.py +316 -0
- code_puppy/tools/browser/browser_navigation.py +16 -16
- code_puppy/tools/browser/browser_screenshot.py +79 -143
- code_puppy/tools/browser/browser_scripts.py +32 -42
- code_puppy/tools/browser/browser_workflows.py +44 -27
- code_puppy/tools/browser/chromium_terminal_manager.py +259 -0
- code_puppy/tools/browser/terminal_command_tools.py +521 -0
- code_puppy/tools/browser/terminal_screenshot_tools.py +556 -0
- code_puppy/tools/browser/terminal_tools.py +525 -0
- code_puppy/tools/command_runner.py +930 -147
- code_puppy/tools/common.py +1113 -5
- code_puppy/tools/display.py +84 -0
- code_puppy/tools/file_modifications.py +288 -89
- code_puppy/tools/file_operations.py +226 -154
- code_puppy/tools/subagent_context.py +158 -0
- code_puppy/uvx_detection.py +242 -0
- code_puppy/version_checker.py +30 -11
- code_puppy-0.0.366.data/data/code_puppy/models.json +110 -0
- code_puppy-0.0.366.data/data/code_puppy/models_dev_api.json +1 -0
- {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/METADATA +149 -75
- code_puppy-0.0.366.dist-info/RECORD +217 -0
- {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/WHEEL +1 -1
- code_puppy/command_line/mcp/add_command.py +0 -183
- code_puppy/messaging/spinner/textual_spinner.py +0 -106
- code_puppy/tools/browser/camoufox_manager.py +0 -216
- code_puppy/tools/browser/vqa_agent.py +0 -70
- code_puppy/tui/__init__.py +0 -10
- code_puppy/tui/app.py +0 -1105
- code_puppy/tui/components/__init__.py +0 -21
- code_puppy/tui/components/chat_view.py +0 -551
- code_puppy/tui/components/command_history_modal.py +0 -218
- code_puppy/tui/components/copy_button.py +0 -139
- code_puppy/tui/components/custom_widgets.py +0 -63
- code_puppy/tui/components/human_input_modal.py +0 -175
- code_puppy/tui/components/input_area.py +0 -167
- code_puppy/tui/components/sidebar.py +0 -309
- code_puppy/tui/components/status_bar.py +0 -185
- code_puppy/tui/messages.py +0 -27
- code_puppy/tui/models/__init__.py +0 -8
- code_puppy/tui/models/chat_message.py +0 -25
- code_puppy/tui/models/command_history.py +0 -89
- code_puppy/tui/models/enums.py +0 -24
- code_puppy/tui/screens/__init__.py +0 -17
- code_puppy/tui/screens/autosave_picker.py +0 -175
- code_puppy/tui/screens/help.py +0 -130
- code_puppy/tui/screens/mcp_install_wizard.py +0 -803
- code_puppy/tui/screens/settings.py +0 -306
- code_puppy/tui/screens/tools.py +0 -74
- code_puppy/tui_state.py +0 -55
- code_puppy-0.0.214.data/data/code_puppy/models.json +0 -112
- code_puppy-0.0.214.dist-info/RECORD +0 -131
- {code_puppy-0.0.214.dist-info → code_puppy-0.0.366.dist-info}/entry_points.txt +0 -0
- {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
|