diff --git a/gemini_web2api/server.py b/gemini_web2api/server.py index 3fc35b6..d56a0ae 100644 --- a/gemini_web2api/server.py +++ b/gemini_web2api/server.py @@ -13,6 +13,9 @@ from .multimodal import detect_image_mime, fetch_image_bytes, upload_image from . import __version__ +# Fence marker the model is instructed to wrap tool calls in (see tools.py). +TOOL_CALL_MARKER = "```tool_call" + def _usage(prompt: str, text: str) -> dict: p = len(prompt) // 4 @@ -192,6 +195,7 @@ def _handle_chat(self, body: bytes): return stream = req.get("stream", False) + log(f"Chat completions: stream={stream}, tools={len(tools) if tools else 0}, model={model_name}") cid = f"chatcmpl-{uuid.uuid4().hex[:12]}" try: file_refs = _upload_images(images) @@ -231,6 +235,80 @@ def _handle_chat(self, body: bytes): log(f"Stream error: {e}") return + if stream: + # Tools present + streaming: forward prose deltas in real time, + # but hold back ```tool_call fenced blocks (including a fence start + # straddling delta boundaries) so they can be parsed into OpenAI + # tool_calls at end of stream instead of leaking into chat text. + self._start_sse() + first_chunk = { + "id": cid, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model_name, + "choices": [{ + "index": 0, + "delta": {"role": "assistant"}, + "finish_reason": None, + }], + } + self.wfile.write(f"data: {json.dumps(first_chunk)}\n\n".encode()) + self.wfile.flush() + + def send_delta(content=None, tool_calls=None, finish=None): + delta = {} + if content is not None: + delta["content"] = content + if tool_calls: + delta["tool_calls"] = tool_calls + chunk = {"id": cid, "object": "chat.completion.chunk", "created": int(time.time()), + "model": model_name, "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]} + self.wfile.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()) + self.wfile.flush() + + full_text = "" + emitted = 0 + finish = "stop" + try: + for delta in generate_stream(prompt, model_id, think_mode, file_refs, extra_fields): + full_text += delta + marker_pos = full_text.find(TOOL_CALL_MARKER) + # Without a marker, hold back the last len(marker)-1 chars so a + # fence start split across deltas is not forwarded prematurely. + limit = marker_pos if marker_pos != -1 else len(full_text) - len(TOOL_CALL_MARKER) + 1 + if limit > emitted: + send_delta(content=full_text[emitted:limit]) + emitted = limit + clean, tool_calls = parse_tool_calls(full_text) + if tool_calls: + log(f"Chat tool-fenced streaming: parsed {len(tool_calls)} tool call(s)") + if len(clean) > emitted: + send_delta(content=clean[emitted:]) + for i, tc in enumerate(tool_calls): + send_delta(tool_calls=[{ + "index": i, + "id": tc["id"], + "type": "function", + "function": { + "name": tc["function"]["name"], + "arguments": tc["function"]["arguments"], + }, + }]) + finish = "tool_calls" + else: + # No (or malformed) tool call: forward whatever is still + # buffered, raw, so nothing disappears silently. + if len(full_text) > emitted: + send_delta(content=full_text[emitted:]) + send_delta(finish=finish) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + pass + except Exception as e: + log(f"Stream error: {e}") + return + try: text = generate(prompt, model_id, think_mode, file_refs, extra_fields) except Exception as e: diff --git a/tests/test_modular_sync.py b/tests/test_modular_sync.py index 7479e92..6a5c3dc 100644 --- a/tests/test_modular_sync.py +++ b/tests/test_modular_sync.py @@ -440,6 +440,98 @@ def test_responses_function_call_stream_has_complete_event_sequence( self.assertEqual(events[4][1]["arguments"], '{"city":"Shanghai"}') self.assertEqual(events[-1][1]["response"]["output"][0]["name"], "get_weather") + def _stream_chunks(self, body): + return [ + json.loads(line[len("data: "):]) + for line in body.splitlines() + if line.startswith("data: {") + ] + + @staticmethod + def _tools_payload(): + return [{ + "type": "function", + "function": {"name": "write_file", "description": "write a file", "parameters": {}}, + }] + + @mock.patch("gemini_web2api.server.generate_stream") + def test_chat_stream_with_tools_parses_tool_call_deltas(self, generate_stream): + generate_stream.return_value = iter([ + "Creating the file.\n", + "```tool_call\n{\"name\": \"write_file\", \"arguments\": {\"path\": \"a.txt\"}}\n```", + ]) + + status, _, body = self.post_json( + "/v1/chat/completions", + { + "model": "gemini-3.6-flash", + "messages": [{"role": "user", "content": "go"}], + "stream": True, + "tools": self._tools_payload(), + }, + ) + + self.assertEqual(status, 200) + chunks = self._stream_chunks(body) + contents = "".join(c["choices"][0]["delta"].get("content", "") for c in chunks) + self.assertEqual(contents, "Creating the file.\n") + self.assertNotIn("tool_call", contents) + tc_deltas = [c for c in chunks if c["choices"][0]["delta"].get("tool_calls")] + self.assertEqual(len(tc_deltas), 1) + call = tc_deltas[0]["choices"][0]["delta"]["tool_calls"][0] + self.assertTrue(call["id"].startswith("call_")) + self.assertEqual(call["type"], "function") + self.assertEqual(call["function"]["name"], "write_file") + self.assertEqual(json.loads(call["function"]["arguments"]), {"path": "a.txt"}) + finishes = [c["choices"][0]["finish_reason"] for c in chunks if c["choices"][0]["finish_reason"]] + self.assertEqual(finishes, ["tool_calls"]) + self.assertTrue(body.endswith("data: [DONE]\n\n")) + + @mock.patch("gemini_web2api.server.generate_stream") + def test_chat_stream_malformed_tool_call_falls_back_to_text(self, generate_stream): + raw = '```tool_call\n{"name": "write_file", "arguments": {"path":' + generate_stream.return_value = iter([raw]) + + status, _, body = self.post_json( + "/v1/chat/completions", + { + "model": "gemini-3.6-flash", + "messages": [{"role": "user", "content": "go"}], + "stream": True, + "tools": self._tools_payload(), + }, + ) + + self.assertEqual(status, 200) + chunks = self._stream_chunks(body) + contents = "".join(c["choices"][0]["delta"].get("content", "") for c in chunks) + self.assertEqual(contents, raw) + finishes = [c["choices"][0]["finish_reason"] for c in chunks if c["choices"][0]["finish_reason"]] + self.assertEqual(finishes, ["stop"]) + self.assertTrue(body.endswith("data: [DONE]\n\n")) + + @mock.patch("gemini_web2api.server.generate_stream") + def test_chat_stream_with_tools_plain_text_fully_streamed(self, generate_stream): + generate_stream.return_value = iter(["Hello", " there friend"]) + + status, _, body = self.post_json( + "/v1/chat/completions", + { + "model": "gemini-3.6-flash", + "messages": [{"role": "user", "content": "go"}], + "stream": True, + "tools": self._tools_payload(), + }, + ) + + self.assertEqual(status, 200) + chunks = self._stream_chunks(body) + contents = "".join(c["choices"][0]["delta"].get("content", "") for c in chunks) + self.assertEqual(contents, "Hello there friend") + finishes = [c["choices"][0]["finish_reason"] for c in chunks if c["choices"][0]["finish_reason"]] + self.assertEqual(finishes, ["stop"]) + self.assertTrue(body.endswith("data: [DONE]\n\n")) + if __name__ == "__main__": unittest.main()