Fix the bug of tools search error.

main.py

@@ -1,10 +1,10 @@
-import os
 import json
 import httpx
 import logging
 import yaml
 import secrets
 import traceback
+from fastapi.responses import JSONResponse
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, Request, HTTPException, Depends
@@ -21,7 +21,7 @@ from urllib.parse import urlparse
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # 启动时的代码
-    timeout = httpx.Timeout(connect=10.0, read=30.0, write=30.0, pool=30.0)
+    timeout = httpx.Timeout(connect=15.0, read=30.0, write=30.0, pool=30.0)
     app.state.client = httpx.AsyncClient(timeout=timeout)
     yield
     # 关闭时的代码
@@ -48,7 +48,7 @@ def load_config():
                 conf['providers'][index] = provider
             api_keys_db = conf['api_keys']
             api_list = [item["api"] for item in api_keys_db]
-            print(json.dumps(conf, indent=4, ensure_ascii=False))
+            # print(json.dumps(conf, indent=4, ensure_ascii=False))
             return conf, api_keys_db, api_list
     except FileNotFoundError:
         print("配置文件 'config.yaml' 未找到。请确保文件存在于正确的位置。")
@@ -59,6 +59,24 @@ def load_config():
 
 config, api_keys_db, api_list = load_config()
 
+async def error_handling_wrapper(generator, status_code=200):
+    try:
+        first_item = await generator.__anext__()
+        if isinstance(first_item, dict) and "error" in first_item:
+            # 如果第一个 yield 的项是错误信息，抛出 HTTPException
+            raise HTTPException(status_code=status_code, detail=first_item)
+
+        # 如果不是错误，创建一个新的生成器，首先yield第一个项，然后yield剩余的项
+        async def new_generator():
+            yield first_item
+            async for item in generator:
+                yield item
+
+        return new_generator()
+    except StopAsyncIteration:
+        # 处理生成器为空的情况
+        return []
+
 async def process_request(request: RequestModel, provider: Dict):
     print("provider: ", provider['provider'])
     url = provider['base_url']
@@ -84,7 +102,15 @@ async def process_request(request: RequestModel, provider: Dict):
 
     if request.stream:
         model = provider['model'][request.model]
-        return StreamingResponse(fetch_response_stream(app.state.client, url, headers, payload, engine, model), media_type="text/event-stream")
+        try:
+            generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
+            wrapped_generator = await error_handling_wrapper(generator, status_code=500)
+            return StreamingResponse(wrapped_generator, media_type="text/event-stream")
+        except HTTPException as e:
+            return JSONResponse(status_code=e.status_code, content={"error": str(e.detail)})
+        except Exception as e:
+            # 处理其他异常
+            return JSONResponse(status_code=500, content={"error": str(e)})
     else:
         return await fetch_response(app.state.client, url, headers, payload)
 

models.py

@@ -28,6 +28,7 @@ class ContentItem(BaseModel):
 class Message(BaseModel):
     role: str
     name: Optional[str] = None
+    arguments: Optional[str] = None
     content: Union[str, List[ContentItem]]
 
 class RequestModel(BaseModel):

request.py

@@ -1,3 +1,4 @@
+import json
 from models import RequestModel
 
 async def get_image_message(base64_image, engine = None):
@@ -191,6 +192,9 @@ async def get_claude_payload(request, engine, provider):
         else:
             content = msg.content
             name = msg.name
+            arguments = msg.arguments
+            if arguments:
+                arguments = json.loads(arguments)
         if name:
             # messages.append({"role": "assistant", "name": name, "content": content})
             messages.append(
@@ -201,7 +205,7 @@ async def get_claude_payload(request, engine, provider):
                             "type": "tool_use",
                             "id": "toolu_01RofFmKHUKsEaZvqESG5Hwz",
                             "name": name,
-                            "input": {"text": messages[-1]["content"][0]["text"]},
+                            "input": arguments,
                         }
                     ]
                 }
@@ -223,23 +227,30 @@ async def get_claude_payload(request, engine, provider):
         elif msg.role == "system":
             system_prompt = content
 
+    conversation_len = len(messages) - 1
+    message_index = 0
+    while message_index < conversation_len:
+        if messages[message_index]["role"] == messages[message_index + 1]["role"]:
+            if messages[message_index].get("content"):
+                if isinstance(messages[message_index]["content"], list):
+                    messages[message_index]["content"].extend(messages[message_index + 1]["content"])
+                elif isinstance(messages[message_index]["content"], str) and isinstance(messages[message_index + 1]["content"], list):
+                    content_list = [{"type": "text", "text": messages[message_index]["content"]}]
+                    content_list.extend(messages[message_index + 1]["content"])
+                    messages[message_index]["content"] = content_list
+                else:
+                    messages[message_index]["content"] += messages[message_index + 1]["content"]
+            messages.pop(message_index + 1)
+            conversation_len = conversation_len - 1
+        else:
+            message_index = message_index + 1
+
     model = provider['model'][request.model]
     payload = {
         "model": model,
         "messages": messages,
         "system": system_prompt,
     }
-    # json_post = {
-    #     "model": model or self.engine,
-    #     "messages": self.conversation[convo_id] if pass_history else [{
-    #         "role": "user",
-    #         "content": prompt
-    #     }],
-    #     "temperature": kwargs.get("temperature", self.temperature),
-    #     "top_p": kwargs.get("top_p", self.top_p),
-    #     "max_tokens": model_max_tokens,
-    #     "stream": True,
-    # }
 
     miss_fields = [
         'model',
@@ -258,7 +269,7 @@ async def get_claude_payload(request, engine, provider):
     if request.tools:
         tools = []
         for tool in request.tools:
-            print("tool", type(tool), tool)
+            # print("tool", type(tool), tool)
 
             json_tool = await gpt2claude_tools_json(tool.dict()["function"])
             tools.append(json_tool)
@@ -267,7 +278,6 @@ async def get_claude_payload(request, engine, provider):
             payload["tool_choice"] = {
                 "type": "auto"
             }
-    import json
     print("payload", json.dumps(payload, indent=2, ensure_ascii=False))
 
     return url, headers, payload

response.py

@@ -1,6 +1,7 @@
-from datetime import datetime
 import json
 import httpx
+from datetime import datetime
+
 
 async def generate_sse_response(timestamp, model, content=None, tools_id=None, function_call_name=None, function_call_content=None, role=None, tokens_use=None, total_tokens=None):
     sample_data = {
@@ -34,102 +35,108 @@ async def generate_sse_response(timestamp, model, content=None, tools_id=None, f
     return sse_response
 
 async def fetch_gemini_response_stream(client, url, headers, payload, model):
-    try:
-        timestamp = datetime.timestamp(datetime.now())
-        async with client.stream('POST', url, headers=headers, json=payload) as response:
-            buffer = ""
-            async for chunk in response.aiter_text():
-                buffer += chunk
-                while "\n" in buffer:
-                    line, buffer = buffer.split("\n", 1)
-                    print(line)
-                    if line and '\"text\": \"' in line:
-                        try:
-                            json_data = json.loads( "{" + line + "}")
-                            content = json_data.get('text', '')
-                            content = "\n".join(content.split("\\n"))
-                            sse_string = await generate_sse_response(timestamp, model, content)
-                            yield sse_string
-                        except json.JSONDecodeError:
-                            print(f"无法解析JSON: {line}")
-
-            # 处理缓冲区中剩余的内容
-            if buffer:
-                # print(buffer)
-                if '\"text\": \"' in buffer:
+    timestamp = datetime.timestamp(datetime.now())
+    async with client.stream('POST', url, headers=headers, json=payload) as response:
+        buffer = ""
+        async for chunk in response.aiter_text():
+            buffer += chunk
+            while "\n" in buffer:
+                line, buffer = buffer.split("\n", 1)
+                print(line)
+                if line and '\"text\": \"' in line:
                     try:
-                        json_data = json.loads(buffer)
+                        json_data = json.loads( "{" + line + "}")
                         content = json_data.get('text', '')
                         content = "\n".join(content.split("\\n"))
                         sse_string = await generate_sse_response(timestamp, model, content)
                         yield sse_string
                     except json.JSONDecodeError:
-                        print(f"无法解析JSON: {buffer}")
+                        print(f"无法解析JSON: {line}")
 
-            yield "data: [DONE]\n\n"
-    except httpx.ConnectError as e:
-        print(f"连接错误： {e}")
+        # 处理缓冲区中剩余的内容
+        if buffer:
+            # print(buffer)
+            if '\"text\": \"' in buffer:
+                try:
+                    json_data = json.loads(buffer)
+                    content = json_data.get('text', '')
+                    content = "\n".join(content.split("\\n"))
+                    sse_string = await generate_sse_response(timestamp, model, content)
+                    yield sse_string
+                except json.JSONDecodeError:
+                    print(f"无法解析JSON: {buffer}")
+
+        # yield "data: [DONE]\n\n"
 
 async def fetch_gpt_response_stream(client, url, headers, payload):
-    try:
-        async with client.stream('POST', url, headers=headers, json=payload) as response:
-            async for chunk in response.aiter_bytes():
-                print(chunk.decode('utf-8'))
-                yield chunk
-    except httpx.ConnectError as e:
-        print(f"连接错误： {e}")
+    async with client.stream('POST', url, headers=headers, json=payload) as response:
+        async for chunk in response.aiter_bytes():
+            print(chunk.decode('utf-8'))
+            yield chunk
 
 async def fetch_claude_response_stream(client, url, headers, payload, model):
-    try:
-        timestamp = datetime.timestamp(datetime.now())
-        async with client.stream('POST', url, headers=headers, json=payload) as response:
-            buffer = ""
-            async for chunk in response.aiter_bytes():
-                buffer += chunk.decode('utf-8')
-                while "\n" in buffer:
-                    line, buffer = buffer.split("\n", 1)
-                    print(line)
+    timestamp = datetime.timestamp(datetime.now())
+    async with client.stream('POST', url, headers=headers, json=payload) as response:
+        # response.raise_for_status()
+        if response.status_code == 200:
+            print("请求成功，状态码是200")
+        else:
+            print('\033[31m')
+            print(f"请求失败，状态码是{response.status_code}，错误信息：")
+            error_message = await response.aread()
+            error_str = error_message.decode('utf-8', errors='replace')
+            error_json = json.loads(error_str)
+            print(json.dumps(error_json, indent=4, ensure_ascii=False))
+            print('\033[0m')
+            yield {"error": f"HTTP Error {response.status_code}", "details": error_json}
+            # raise HTTPStatusError(f"HTTP Error {response.status_code}", request=response.request, response=response)
+            # raise HTTPException(status_code=response.status_code, detail=error_json)
+        buffer = ""
+        async for chunk in response.aiter_bytes():
+            buffer += chunk.decode('utf-8')
+            while "\n" in buffer:
+                line, buffer = buffer.split("\n", 1)
+                print(line)
 
-                    if line.startswith("data:"):
-                        print(line)
-                        line = line[6:]
-                        resp: dict = json.loads(line)
-                        message = resp.get("message")
-                        if message:
-                            tokens_use = resp.get("usage")
-                            role = message.get("role")
-                            if role:
-                                sse_string = await generate_sse_response(timestamp, model, None, None, None, None, role)
-                                yield sse_string
-                            if tokens_use:
-                                total_tokens = tokens_use["input_tokens"] + tokens_use["output_tokens"]
-                                # print("\n\rtotal_tokens", total_tokens)
-                        tool_use = resp.get("content_block")
-                        tools_id = None
-                        function_call_name = None
-                        if tool_use and "tool_use" == tool_use['type']:
-                            # print("tool_use", tool_use)
-                            tools_id = tool_use["id"]
-                            if "name" in tool_use:
-                                function_call_name = tool_use["name"]
-                                sse_string = await generate_sse_response(timestamp, model, None, tools_id, function_call_name, None)
-                                yield sse_string
-                        delta = resp.get("delta")
-                        # print("delta", delta)
-                        if not delta:
-                            continue
-                        if "text" in delta:
-                            content = delta["text"]
-                            sse_string = await generate_sse_response(timestamp, model, content, None, None)
+                if line.startswith("data:"):
+                    print(line)
+                    line = line[6:]
+                    resp: dict = json.loads(line)
+                    message = resp.get("message")
+                    if message:
+                        tokens_use = resp.get("usage")
+                        role = message.get("role")
+                        if role:
+                            sse_string = await generate_sse_response(timestamp, model, None, None, None, None, role)
                             yield sse_string
-                        if "partial_json" in delta:
-                            # {"type":"input_json_delta","partial_json":""}
-                            function_call_content = delta["partial_json"]
-                            sse_string = await generate_sse_response(timestamp, model, None, None, None, function_call_content)
+                        if tokens_use:
+                            total_tokens = tokens_use["input_tokens"] + tokens_use["output_tokens"]
+                            # print("\n\rtotal_tokens", total_tokens)
+                    tool_use = resp.get("content_block")
+                    tools_id = None
+                    function_call_name = None
+                    if tool_use and "tool_use" == tool_use['type']:
+                        # print("tool_use", tool_use)
+                        tools_id = tool_use["id"]
+                        if "name" in tool_use:
+                            function_call_name = tool_use["name"]
+                            sse_string = await generate_sse_response(timestamp, model, None, tools_id, function_call_name, None)
                             yield sse_string
-            yield "data: [DONE]\n\n"
-    except httpx.ConnectError as e:
-        print(f"连接错误： {e}")
+                    delta = resp.get("delta")
+                    # print("delta", delta)
+                    if not delta:
+                        continue
+                    if "text" in delta:
+                        content = delta["text"]
+                        sse_string = await generate_sse_response(timestamp, model, content, None, None)
+                        yield sse_string
+                    if "partial_json" in delta:
+                        # {"type":"input_json_delta","partial_json":""}
+                        function_call_content = delta["partial_json"]
+                        sse_string = await generate_sse_response(timestamp, model, None, None, None, function_call_content)
+                        yield sse_string
+
+        # yield "data: [DONE]\n\n"
 
 async def fetch_response(client, url, headers, payload):
     response = await client.post(url, headers=headers, json=payload)