aient 1.1.35__py3-none-any.whl → 1.1.36__py3-none-any.whl

aient/core/models.py

@@ -105,6 +105,7 @@ class RequestModel(BaseRequest):
     response_format: Optional[ResponseFormat] = None
     thinking: Optional[Thinking] = None
     stream_options: Optional[StreamOptions] = None
+    chat_template_kwargs: Optional[Dict[str, Any]] = None
 
     def get_last_text_message(self) -> Optional[str]:
         for message in reversed(self.messages):

aient/core/request.py

@@ -341,7 +341,8 @@ async def get_vertex_gemini_payload(request, engine, provider, api_key=None):
     else:
         location = gemini1
 
-    if "gemini-2.5-pro-exp-03-25" == original_model or "gemini-2.5-pro-preview-06-05" == original_model:
+    if "gemini-2.5-flash-lite-preview-06-17" == original_model or \
+    "gemini-2.5-pro-preview-06-05" == original_model:
         location = gemini2_5_pro_exp
 
     if "google-vertex-ai" in provider.get("base_url", ""):
@@ -362,7 +363,8 @@ async def get_vertex_gemini_payload(request, engine, provider, api_key=None):
         else:
             url = f"https://aiplatform.googleapis.com/v1/publishers/google/models/{original_model}:{gemini_stream}?key={api_key}"
         headers.pop("Authorization", None)
-    elif "gemini-2.5-pro-exp-03-25" == original_model or "gemini-2.5-pro-preview-06-05" == original_model:
+    elif "gemini-2.5-flash-lite-preview-06-17" == original_model or \
+    "gemini-2.5-pro-preview-06-05" == original_model:
         url = "https://aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/{MODEL_ID}:{stream}".format(
             LOCATION=await location.next(),
             PROJECT_ID=project_id,
@@ -1195,6 +1197,126 @@ async def get_azure_payload(request, engine, provider, api_key=None):
 
     return url, headers, payload
 
+async def get_azure_databricks_payload(request, engine, provider, api_key=None):
+    api_key = base64.b64encode(f"token:{api_key}".encode()).decode()
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': f"Basic {api_key}",
+    }
+    model_dict = get_model_dict(provider)
+    original_model = model_dict[request.model]
+
+    base_url=provider['base_url']
+    url = urllib.parse.urljoin(base_url, f"/serving-endpoints/{original_model}/invocations")
+
+    messages = []
+    for msg in request.messages:
+        tool_calls = None
+        tool_call_id = None
+        if isinstance(msg.content, list):
+            content = []
+            for item in msg.content:
+                if item.type == "text":
+                    text_message = await get_text_message(item.text, engine)
+                    content.append(text_message)
+                elif item.type == "image_url" and provider.get("image", True) and "o1-mini" not in original_model:
+                    image_message = await get_image_message(item.image_url.url, engine)
+                    content.append(image_message)
+        else:
+            content = msg.content
+            tool_calls = msg.tool_calls
+            tool_call_id = msg.tool_call_id
+
+        if tool_calls:
+            tool_calls_list = []
+            for tool_call in tool_calls:
+                tool_calls_list.append({
+                    "id": tool_call.id,
+                    "type": tool_call.type,
+                    "function": {
+                        "name": tool_call.function.name,
+                        "arguments": tool_call.function.arguments
+                    }
+                })
+                if provider.get("tools"):
+                    messages.append({"role": msg.role, "tool_calls": tool_calls_list})
+        elif tool_call_id:
+            if provider.get("tools"):
+                messages.append({"role": msg.role, "tool_call_id": tool_call_id, "content": content})
+        else:
+            messages.append({"role": msg.role, "content": content})
+
+    if "claude-3-7-sonnet" in original_model:
+        max_tokens = 128000
+    elif "claude-3-5-sonnet" in original_model:
+        max_tokens = 8192
+    elif "claude-sonnet-4" in original_model:
+        max_tokens = 64000
+    elif "claude-opus-4" in original_model:
+        max_tokens = 32000
+    else:
+        max_tokens = 4096
+
+    payload = {
+        "model": original_model,
+        "messages": messages,
+        "max_tokens": max_tokens,
+    }
+
+    if request.max_tokens:
+        payload["max_tokens"] = int(request.max_tokens)
+
+    miss_fields = [
+        'model',
+        'messages',
+    ]
+
+    for field, value in request.model_dump(exclude_unset=True).items():
+        if field not in miss_fields and value is not None:
+            if field == "max_tokens" and "o1" in original_model:
+                payload["max_completion_tokens"] = value
+            else:
+                payload[field] = value
+
+    if provider.get("tools") == False or "o1" in original_model or "chatgpt-4o-latest" in original_model or "grok" in original_model:
+        payload.pop("tools", None)
+        payload.pop("tool_choice", None)
+
+    if "think" in request.model.lower():
+        payload["thinking"] = {
+            "budget_tokens": 4096,
+            "type": "enabled"
+        }
+        payload["temperature"] = 1
+        payload.pop("top_p", None)
+        payload.pop("top_k", None)
+        if request.model.split("-")[-1].isdigit():
+            think_tokens = int(request.model.split("-")[-1])
+            if think_tokens < max_tokens:
+                payload["thinking"] = {
+                    "budget_tokens": think_tokens,
+                    "type": "enabled"
+                }
+
+    if request.thinking:
+        payload["thinking"] = {
+            "budget_tokens": request.thinking.budget_tokens,
+            "type": request.thinking.type
+        }
+        payload["temperature"] = 1
+        payload.pop("top_p", None)
+        payload.pop("top_k", None)
+
+    if safe_get(provider, "preferences", "post_body_parameter_overrides", default=None):
+        for key, value in safe_get(provider, "preferences", "post_body_parameter_overrides", default={}).items():
+            if key == request.model:
+                for k, v in value.items():
+                    payload[k] = v
+            elif all(_model not in request.model.lower() for _model in ["gemini", "gpt", "claude"]):
+                payload[key] = value
+
+    return url, headers, payload
+
 async def get_openrouter_payload(request, engine, provider, api_key=None):
     headers = {
         'Content-Type': 'application/json'
@@ -1763,6 +1885,8 @@ async def get_payload(request: RequestModel, engine, provider, api_key=None):
         return await get_vertex_claude_payload(request, engine, provider, api_key)
     elif engine == "azure":
         return await get_azure_payload(request, engine, provider, api_key)
+    elif engine == "azure-databricks":
+        return await get_azure_databricks_payload(request, engine, provider, api_key)
     elif engine == "claude":
         return await get_claude_payload(request, engine, provider, api_key)
     elif engine == "gpt":

aient/core/response.py

@@ -49,6 +49,30 @@ async def fetch_gemini_response_stream(client, url, headers, payload, model):
             while "\n" in buffer:
                 line, buffer = buffer.split("\n", 1)
                 # line_index += 1
+                if line.startswith("data: "):
+                    json_line = line.lstrip("data: ").strip()
+                    response_json = json.loads(json_line)
+                    json_data = safe_get(response_json, "candidates", 0, "content", default=None)
+                    finishReason = safe_get(response_json, "candidates", 0 , "finishReason", default=None)
+                    if finishReason:
+                        promptTokenCount = safe_get(response_json, "usageMetadata", "promptTokenCount", default=0)
+                        candidatesTokenCount = safe_get(response_json, "usageMetadata", "candidatesTokenCount", default=0)
+                        totalTokenCount = safe_get(response_json, "usageMetadata", "totalTokenCount", default=0)
+
+                    content = safe_get(json_data, "parts", 0, "text", default="")
+                    b64_json = safe_get(json_data, "parts", 0, "inlineData", "data", default="")
+                    if b64_json:
+                        image_base64 = b64_json
+
+                    is_thinking = safe_get(json_data, "parts", 0, "thought", default=False)
+                    if is_thinking:
+                        sse_string = await generate_sse_response(timestamp, model, reasoning_content=content)
+                        yield sse_string
+                    elif not image_base64 and content:
+                        sse_string = await generate_sse_response(timestamp, model, content=content)
+                        yield sse_string
+
+                    continue
 
                 # https://ai.google.dev/api/generate-content?hl=zh-cn#FinishReason
                 if line and '\"finishReason\": \"' in line:
@@ -270,8 +294,15 @@ async def fetch_gpt_response_stream(client, url, headers, payload):
 
                     no_stream_content = safe_get(line, "choices", 0, "message", "content", default=None)
                     openrouter_reasoning = safe_get(line, "choices", 0, "delta", "reasoning", default="")
+                    azure_databricks_claude_summary_content = safe_get(line, "choices", 0, "delta", "content", 0, "summary", 0, "text", default="")
+                    azure_databricks_claude_signature_content = safe_get(line, "choices", 0, "delta", "content", 0, "summary", 0, "signature", default="")
                     # print("openrouter_reasoning", repr(openrouter_reasoning), openrouter_reasoning.endswith("\\\\"), openrouter_reasoning.endswith("\\"))
-                    if openrouter_reasoning:
+                    if azure_databricks_claude_signature_content:
+                        pass
+                    elif azure_databricks_claude_summary_content:
+                        sse_string = await generate_sse_response(timestamp, payload["model"], reasoning_content=azure_databricks_claude_summary_content)
+                        yield sse_string
+                    elif openrouter_reasoning:
                         if openrouter_reasoning.endswith("\\"):
                             enter_buffer += openrouter_reasoning
                             continue
@@ -640,15 +671,12 @@ async def fetch_response_stream(client, url, headers, payload, engine, model):
     elif engine == "aws":
         async for chunk in fetch_aws_response_stream(client, url, headers, payload, model):
             yield chunk
-    elif engine == "gpt":
+    elif engine == "gpt" or engine == "openrouter" or engine == "azure-databricks":
         async for chunk in fetch_gpt_response_stream(client, url, headers, payload):
             yield chunk
     elif engine == "azure":
         async for chunk in fetch_azure_response_stream(client, url, headers, payload):
             yield chunk
-    elif engine == "openrouter":
-        async for chunk in fetch_gpt_response_stream(client, url, headers, payload):
-            yield chunk
     elif engine == "cloudflare":
         async for chunk in fetch_cloudflare_response_stream(client, url, headers, payload, model):
             yield chunk

aient/core/utils.py

@@ -75,6 +75,8 @@ def get_engine(provider, endpoint=None, original_model=""):
         engine = "vertex"
     elif parsed_url.netloc.rstrip('/').endswith('azure.com'):
         engine = "azure"
+    elif parsed_url.netloc.rstrip('/').endswith('azuredatabricks.net'):
+        engine = "azure-databricks"
     elif parsed_url.netloc == 'api.cloudflare.com':
         engine = "cloudflare"
     elif parsed_url.netloc == 'api.anthropic.com' or parsed_url.path.endswith("v1/messages"):
@@ -482,7 +484,6 @@ async def generate_sse_response(timestamp, model, content=None, tools_id=None, f
     if role:
         sample_data["choices"][0]["delta"] = {"role": role, "content": ""}
     if total_tokens:
-        total_tokens = prompt_tokens + completion_tokens
         sample_data["usage"] = {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens}
         sample_data["choices"] = []
     if stop:
@@ -574,7 +575,6 @@ async def generate_no_stream_response(timestamp, model, content=None, tools_id=N
         }
 
     if total_tokens:
-        total_tokens = prompt_tokens + completion_tokens
         sample_data["usage"] = {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens}
 
     json_data = json.dumps(sample_data, ensure_ascii=False)
@@ -674,7 +674,7 @@ async def get_image_message(base64_image, engine = None):
         base64_image = f"data:image/png;base64,{png_base64}"
         image_type = "image/png"
 
-    if "gpt" == engine or "openrouter" == engine or "azure" == engine:
+    if "gpt" == engine or "openrouter" == engine or "azure" == engine or "azure-databricks" == engine:
         return {
             "type": "image_url",
             "image_url": {
@@ -702,7 +702,9 @@ async def get_image_message(base64_image, engine = None):
     raise ValueError("Unknown engine")
 
 async def get_text_message(message, engine = None):
-    if "gpt" == engine or "claude" == engine or "openrouter" == engine or "vertex-claude" == engine or "azure" == engine or "aws" == engine:
+    if "gpt" == engine or "claude" == engine or "openrouter" == engine or \
+    "vertex-claude" == engine or "azure" == engine or "aws" == engine or \
+    "azure-databricks" == engine:
         return {"type": "text", "text": message}
     if "gemini" == engine or "vertex-gemini" == engine:
         return {"text": message}
@@ -734,7 +736,7 @@ def parse_json_safely(json_str):
             return json.loads(json_str, strict=False)
         except json.JSONDecodeError as e:
             # 两种方法都失败，抛出异常
-            raise Exception(f"无法解析JSON字符串: {e}")
+            raise Exception(f"无法解析JSON字符串: {e}, {json_str}")
 
 if __name__ == "__main__":
     provider = {

{aient-1.1.35.dist-info → aient-1.1.36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aient
-Version: 1.1.35
+Version: 1.1.36
 Summary: Aient: The Awakening of Agent.
 Description-Content-Type: text/markdown
 License-File: LICENSE

{aient-1.1.35.dist-info → aient-1.1.36.dist-info}/RECORD

@@ -3,10 +3,10 @@ aient/core/.git,sha256=lrAcW1SxzRBUcUiuKL5tS9ykDmmTXxyLP3YYU-Y-Q-I,45
 aient/core/.gitignore,sha256=5JRRlYYsqt_yt6iFvvzhbqh2FTUQMqwo6WwIuFzlGR8,13
 aient/core/__init__.py,sha256=NxjebTlku35S4Dzr16rdSqSTWUvvwEeACe8KvHJnjPg,34
 aient/core/log_config.py,sha256=kz2_yJv1p-o3lUQOwA3qh-LSc3wMHv13iCQclw44W9c,274
-aient/core/models.py,sha256=oUGsytAVBuhE_MLco7PqIQGwWD8lEYkZxgb8HEuynmA,7444
-aient/core/request.py,sha256=Zp5Q6iZp3OWZt95cZLh0VmbwvF18jXQB88UfA_3AKt8,67279
-aient/core/response.py,sha256=WIJhvf3Th3-XW3MsDq8QdzxXcNe5q1mCKcCZx8pzm2A,33213
-aient/core/utils.py,sha256=zidsBUBd3KskzcxQcPB1y5x1RhtWcbZeWvmgb4LAadA,27318
+aient/core/models.py,sha256=d4MISNezTSe0ls0-fjuToI2SoT-sk5fWqAJuKVinIlo,7502
+aient/core/request.py,sha256=6Nwduj7kFuubFaZ0ZLkT_zd03XpT-bFhgrKVOZiGBOQ,71918
+aient/core/response.py,sha256=RYy70Ld_txixHHd61Dqtlo0tKHMU_OIXqxGWd6EfATI,35315
+aient/core/utils.py,sha256=fhI5wBxr01lVEp8nMfjG9dQ859AE-VdrWyb9suLzzqM,27400
 aient/core/test/test_base_api.py,sha256=pWnycRJbuPSXKKU9AQjWrMAX1wiLC_014Qc9hh5C2Pw,524
 aient/core/test/test_geminimask.py,sha256=HFX8jDbNg_FjjgPNxfYaR-0-roUrOO-ND-FVsuxSoiw,13254
 aient/core/test/test_image.py,sha256=_T4peNGdXKBHHxyQNx12u-NTyFE8TlYI6NvvagsG2LE,319
@@ -37,8 +37,8 @@ aient/plugins/write_file.py,sha256=7spYxloI_aUbeANEQK-oXrGPoBqSfsD7sdfMAWlNxhU,3
 aient/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aient/utils/prompt.py,sha256=UcSzKkFE4-h_1b6NofI6xgk3GoleqALRKY8VBaXLjmI,11311
 aient/utils/scripts.py,sha256=_43DEeoaiNVSA7ew1UUmp-gIV6XXe6rQPc2HTRuTzkw,40944
-aient-1.1.35.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
-aient-1.1.35.dist-info/METADATA,sha256=Md5vZ0BXPl9uQJCCnFmdJkppgneGrFrdHdVLplV4UN4,4968
-aient-1.1.35.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aient-1.1.35.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
-aient-1.1.35.dist-info/RECORD,,
+aient-1.1.36.dist-info/licenses/LICENSE,sha256=XNdbcWldt0yaNXXWB_Bakoqnxb3OVhUft4MgMA_71ds,1051
+aient-1.1.36.dist-info/METADATA,sha256=u-UNrKVDoYOocUU5VF-hi72Ej0bahPyP8SUKkj24LPU,4968
+aient-1.1.36.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aient-1.1.36.dist-info/top_level.txt,sha256=3oXzrP5sAVvyyqabpeq8A2_vfMtY554r4bVE-OHBrZk,6
+aient-1.1.36.dist-info/RECORD,,