Coverage for providers/vertex.py: 0%

1 import os

2 import json

3 import base64

4 from google import genai

5 from google.genai import types

6 from google.genai.types import HttpOptions

7 from qrclaw.providers.base import LLMProvider, LLMResponse, ToolCall

8 from qrclaw.config import OPENAI_MODEL, VERTEX_API_KEY

9 from qrclaw.logger import get_logger

10

11 logger = get_logger("qrclaw.providers.vertex")

12

13# thought_signature 在 tool_call 里的存储 key

14 _TS_KEY = "__thought_signature__"

15

16

17 def _build_vertex_tools(schemas: list[dict]) -> list[types.Tool] | None:

18 """把 OpenAI tool schema 转换成 Vertex AI 的 Tool 格式"""

19 if not schemas:

20 return None

21 declarations = []

22 for s in schemas:

23 fn = s["function"]

24 params = fn.get("parameters")

25 declarations.append(types.FunctionDeclaration(

26 name=fn["name"],

27 description=fn.get("description", ""),

28 parameters=params,

29 ))

30 return [types.Tool(function_declarations=declarations)]

31

32

33 class VertexProvider(LLMProvider):

34

35 def __init__(self):

36 os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"

37 self._client = genai.Client(

38 http_options=HttpOptions(api_version="v1"),

39 api_key=VERTEX_API_KEY,

40 )

41 logger.info("Vertex AI 渠道已初始化")

42

43 def chat(self, messages: list[dict], tools: list[dict] | None = None) -> LLMResponse:

44 contents = []

45 system_parts = []

46

47 i = 0

48 while i < len(messages):

49 msg = messages[i]

50 role = msg.get("role")

51 content = msg.get("content", "")

52

53 if role == "system":

54 if content:

55 system_parts.append(content)

56 i += 1

57

58 elif role == "user":

59 contents.append(types.Content(

60 role="user",

61 parts=[types.Part(text=content or "")]

62 ))

63 i += 1

64

65 elif role == "assistant":

66 parts = []

67 if content:

68 parts.append(types.Part(text=content))

69 # 处理 tool_calls，恢复 thought_signature

70 for tc in msg.get("tool_calls", []):

71 fn = tc["function"]

72 try:

73 args = json.loads(fn["arguments"])

74 except Exception:

75 args = {}

76 ts_b64 = tc.get(_TS_KEY)

77 ts_bytes = base64.b64decode(ts_b64) if ts_b64 else None

78

79 part = types.Part(

80 function_call=types.FunctionCall(

81 name=fn["name"],

82 args=args,

83 )

84 )

85 if ts_bytes:

86 part.thought_signature = ts_bytes

87 parts.append(part)

88 if parts:

89 contents.append(types.Content(role="model", parts=parts))

90 i += 1

91

92 elif role == "tool":

93 # 把连续的 tool 消息合并到同一个 Content，Vertex AI 要求 function_response 数量与 function_call 一致

94 tool_parts = []

95 while i < len(messages) and messages[i].get("role") == "tool":

96 m = messages[i]

97 tool_call_id = m.get("tool_call_id", "")

98 c = m.get("content", "")

99 try:

100 result = json.loads(c) if c else {}

101 if not isinstance(result, dict):

102 result = {"result": c}

103 except Exception:

104 result = {"result": c}

105 tool_parts.append(types.Part(

106 function_response=types.FunctionResponse(

107 name=tool_call_id,

108 response=result,

109 )

110 ))

111 i += 1

112 contents.append(types.Content(role="user", parts=tool_parts))

113

114 else:

115 i += 1

116

117 config_kwargs = {}

118 if system_parts:

119 config_kwargs["system_instruction"] = "\n".join(system_parts)

120 vertex_tools = _build_vertex_tools(tools)

121 if vertex_tools:

122 config_kwargs["tools"] = vertex_tools

123

124 config = types.GenerateContentConfig(**config_kwargs) if config_kwargs else None

125

126 response = self._client.models.generate_content(

127 model=OPENAI_MODEL,

128 contents=contents,

129 config=config,

130 )

131

132 # 解析响应，把 thought_signature 用 base64 存进 ToolCall

133 tool_calls = []

134 text_content = ""

135 finish_reason = "stop"

136

137 candidate = response.candidates[0]

138 for part in candidate.content.parts:

139 if part.text:

140 text_content += part.text

141 elif part.function_call:

142 fc = part.function_call

143 ts = getattr(part, "thought_signature", None)

144 ts_b64 = base64.b64encode(ts).decode() if ts else None

145 tool_calls.append(ToolCall(

146 id=fc.name,

147 name=fc.name,

148 arguments=json.dumps(dict(fc.args), ensure_ascii=False),

149 thought_signature=ts_b64,

150 ))

151

152 if tool_calls:

153 finish_reason = "tool_calls"

154

155 usage = response.usage_metadata

156 prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0

157 completion_tokens = getattr(usage, "candidates_token_count", 0) or 0

158

159 logger.info(f"Vertex AI 响应成功，tokens: {prompt_tokens + completion_tokens}")

160 return LLMResponse(

161 content=text_content,

162 tool_calls=tool_calls,

163 finish_reason=finish_reason,

164 prompt_tokens=prompt_tokens,

165 completion_tokens=completion_tokens,

166 total_tokens=prompt_tokens + completion_tokens,

167 raw=response,

168 )