lionagi 0.12.2__py3-none-any.whl → 0.12.4__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.
- lionagi/config.py +123 -0
- lionagi/fields/file.py +1 -1
- lionagi/fields/reason.py +1 -1
- lionagi/libs/file/concat.py +1 -6
- lionagi/libs/file/concat_files.py +1 -5
- lionagi/libs/file/save.py +1 -1
- lionagi/libs/package/imports.py +8 -177
- lionagi/libs/parse.py +30 -0
- lionagi/libs/schema/load_pydantic_model_from_schema.py +259 -0
- lionagi/libs/token_transform/perplexity.py +2 -4
- lionagi/libs/token_transform/synthlang_/resources/frameworks/framework_options.json +46 -46
- lionagi/libs/token_transform/synthlang_/translate_to_synthlang.py +1 -1
- lionagi/operations/chat/chat.py +2 -2
- lionagi/operations/communicate/communicate.py +20 -5
- lionagi/operations/parse/parse.py +131 -43
- lionagi/protocols/generic/log.py +1 -2
- lionagi/protocols/generic/pile.py +18 -4
- lionagi/protocols/messages/assistant_response.py +20 -1
- lionagi/protocols/messages/templates/README.md +6 -10
- lionagi/service/connections/__init__.py +15 -0
- lionagi/service/connections/api_calling.py +230 -0
- lionagi/service/connections/endpoint.py +410 -0
- lionagi/service/connections/endpoint_config.py +137 -0
- lionagi/service/connections/header_factory.py +56 -0
- lionagi/service/connections/match_endpoint.py +49 -0
- lionagi/service/connections/providers/__init__.py +3 -0
- lionagi/service/connections/providers/anthropic_.py +87 -0
- lionagi/service/connections/providers/exa_.py +33 -0
- lionagi/service/connections/providers/oai_.py +166 -0
- lionagi/service/connections/providers/ollama_.py +122 -0
- lionagi/service/connections/providers/perplexity_.py +29 -0
- lionagi/service/imodel.py +36 -144
- lionagi/service/manager.py +1 -7
- lionagi/service/{endpoints/rate_limited_processor.py → rate_limited_processor.py} +4 -2
- lionagi/service/resilience.py +545 -0
- lionagi/service/third_party/README.md +71 -0
- lionagi/service/third_party/__init__.py +0 -0
- lionagi/service/third_party/anthropic_models.py +159 -0
- lionagi/service/third_party/exa_models.py +165 -0
- lionagi/service/third_party/openai_models.py +18241 -0
- lionagi/service/third_party/pplx_models.py +156 -0
- lionagi/service/types.py +5 -4
- lionagi/session/branch.py +12 -7
- lionagi/tools/file/reader.py +1 -1
- lionagi/tools/memory/tools.py +497 -0
- lionagi/utils.py +921 -123
- lionagi/version.py +1 -1
- {lionagi-0.12.2.dist-info → lionagi-0.12.4.dist-info}/METADATA +33 -16
- {lionagi-0.12.2.dist-info → lionagi-0.12.4.dist-info}/RECORD +53 -63
- lionagi/libs/file/create_path.py +0 -80
- lionagi/libs/file/file_util.py +0 -358
- lionagi/libs/parse/__init__.py +0 -3
- lionagi/libs/parse/fuzzy_parse_json.py +0 -117
- lionagi/libs/parse/to_dict.py +0 -336
- lionagi/libs/parse/to_json.py +0 -61
- lionagi/libs/parse/to_num.py +0 -378
- lionagi/libs/parse/to_xml.py +0 -57
- lionagi/libs/parse/xml_parser.py +0 -148
- lionagi/libs/schema/breakdown_pydantic_annotation.py +0 -48
- lionagi/service/endpoints/__init__.py +0 -3
- lionagi/service/endpoints/base.py +0 -706
- lionagi/service/endpoints/chat_completion.py +0 -116
- lionagi/service/endpoints/match_endpoint.py +0 -72
- lionagi/service/providers/__init__.py +0 -3
- lionagi/service/providers/anthropic_/__init__.py +0 -3
- lionagi/service/providers/anthropic_/messages.py +0 -99
- lionagi/service/providers/exa_/models.py +0 -3
- lionagi/service/providers/exa_/search.py +0 -80
- lionagi/service/providers/exa_/types.py +0 -7
- lionagi/service/providers/groq_/__init__.py +0 -3
- lionagi/service/providers/groq_/chat_completions.py +0 -56
- lionagi/service/providers/ollama_/__init__.py +0 -3
- lionagi/service/providers/ollama_/chat_completions.py +0 -134
- lionagi/service/providers/openai_/__init__.py +0 -3
- lionagi/service/providers/openai_/chat_completions.py +0 -101
- lionagi/service/providers/openai_/spec.py +0 -14
- lionagi/service/providers/openrouter_/__init__.py +0 -3
- lionagi/service/providers/openrouter_/chat_completions.py +0 -62
- lionagi/service/providers/perplexity_/__init__.py +0 -3
- lionagi/service/providers/perplexity_/chat_completions.py +0 -44
- lionagi/service/providers/perplexity_/models.py +0 -5
- lionagi/service/providers/types.py +0 -17
- /lionagi/{service/providers/exa_/__init__.py → py.typed} +0 -0
- /lionagi/service/{endpoints/token_calculator.py → token_calculator.py} +0 -0
- {lionagi-0.12.2.dist-info → lionagi-0.12.4.dist-info}/WHEEL +0 -0
- {lionagi-0.12.2.dist-info → lionagi-0.12.4.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,230 @@
|
|
1
|
+
# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
|
2
|
+
#
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
4
|
+
|
5
|
+
import asyncio
|
6
|
+
import logging
|
7
|
+
from typing import Any
|
8
|
+
|
9
|
+
from pydantic import BaseModel, Field, model_validator
|
10
|
+
from typing_extensions import Self
|
11
|
+
|
12
|
+
from lionagi.protocols.generic.event import Event, EventStatus
|
13
|
+
from lionagi.service.token_calculator import TokenCalculator
|
14
|
+
|
15
|
+
from .endpoint import Endpoint
|
16
|
+
|
17
|
+
logger = logging.getLogger(__name__)
|
18
|
+
|
19
|
+
|
20
|
+
class APICalling(Event):
|
21
|
+
"""Handles asynchronous API calls with automatic token usage tracking.
|
22
|
+
|
23
|
+
This class manages API calls through endpoints, handling both regular
|
24
|
+
and streaming responses with optional token usage tracking.
|
25
|
+
"""
|
26
|
+
|
27
|
+
endpoint: Endpoint = Field(
|
28
|
+
...,
|
29
|
+
description="Endpoint instance for making the API call",
|
30
|
+
exclude=True,
|
31
|
+
)
|
32
|
+
|
33
|
+
payload: dict = Field(
|
34
|
+
..., description="Request payload to send to the API"
|
35
|
+
)
|
36
|
+
|
37
|
+
headers: dict = Field(
|
38
|
+
default_factory=dict,
|
39
|
+
description="Additional headers for the request",
|
40
|
+
exclude=True,
|
41
|
+
)
|
42
|
+
|
43
|
+
cache_control: bool = Field(
|
44
|
+
default=False,
|
45
|
+
description="Whether to use cache control for this request",
|
46
|
+
exclude=True,
|
47
|
+
)
|
48
|
+
|
49
|
+
include_token_usage_to_model: bool = Field(
|
50
|
+
default=False,
|
51
|
+
description="Whether to include token usage information in messages",
|
52
|
+
exclude=True,
|
53
|
+
)
|
54
|
+
|
55
|
+
@model_validator(mode="after")
|
56
|
+
def _validate_streaming(self) -> Self:
|
57
|
+
"""Validate streaming configuration and add token usage if requested."""
|
58
|
+
if self.payload.get("stream") is True:
|
59
|
+
self.streaming = True
|
60
|
+
|
61
|
+
# Add token usage information to the last message if requested
|
62
|
+
if (
|
63
|
+
self.include_token_usage_to_model
|
64
|
+
and self.endpoint.config.requires_tokens
|
65
|
+
):
|
66
|
+
# Handle both messages format (chat completions) and input format (responses API)
|
67
|
+
if "messages" in self.payload and isinstance(
|
68
|
+
self.payload["messages"][-1], dict
|
69
|
+
):
|
70
|
+
required_tokens = self.required_tokens
|
71
|
+
content = self.payload["messages"][-1]["content"]
|
72
|
+
# Model token limit mapping
|
73
|
+
TOKEN_LIMITS = {
|
74
|
+
# OpenAI models
|
75
|
+
"gpt-4": 128_000,
|
76
|
+
"gpt-4-turbo": 128_000,
|
77
|
+
"o1-mini": 128_000,
|
78
|
+
"o1-preview": 128_000,
|
79
|
+
"o1": 200_000,
|
80
|
+
"o3": 200_000,
|
81
|
+
"gpt-4.1": 1_000_000,
|
82
|
+
# Anthropic models
|
83
|
+
"sonnet": 200_000,
|
84
|
+
"haiku": 200_000,
|
85
|
+
"opus": 200_000,
|
86
|
+
# Google models
|
87
|
+
"gemini": 1_000_000,
|
88
|
+
# Alibaba models
|
89
|
+
"qwen-turbo": 1_000_000,
|
90
|
+
}
|
91
|
+
|
92
|
+
token_msg = (
|
93
|
+
f"\n\nEstimated Current Token Usage: {required_tokens}"
|
94
|
+
)
|
95
|
+
|
96
|
+
# Find matching token limit
|
97
|
+
if "model" in self.payload:
|
98
|
+
model = self.payload["model"]
|
99
|
+
for model_prefix, limit in TOKEN_LIMITS.items():
|
100
|
+
if model_prefix in model.lower():
|
101
|
+
token_msg += f"/{limit:,}"
|
102
|
+
break
|
103
|
+
|
104
|
+
# Update content based on its type
|
105
|
+
if isinstance(content, str):
|
106
|
+
content += token_msg
|
107
|
+
elif isinstance(content, dict) and "text" in content:
|
108
|
+
content["text"] += token_msg
|
109
|
+
elif isinstance(content, list):
|
110
|
+
for item in reversed(content):
|
111
|
+
if isinstance(item, dict) and "text" in item:
|
112
|
+
item["text"] += token_msg
|
113
|
+
break
|
114
|
+
|
115
|
+
self.payload["messages"][-1]["content"] = content
|
116
|
+
|
117
|
+
return self
|
118
|
+
|
119
|
+
@property
|
120
|
+
def required_tokens(self) -> int | None:
|
121
|
+
"""Calculate the number of tokens required for this request."""
|
122
|
+
if not self.endpoint.config.requires_tokens:
|
123
|
+
return None
|
124
|
+
|
125
|
+
# Handle chat completions format
|
126
|
+
if "messages" in self.payload:
|
127
|
+
return TokenCalculator.calculate_message_tokens(
|
128
|
+
self.payload["messages"], **self.payload
|
129
|
+
)
|
130
|
+
# Handle responses API format
|
131
|
+
elif "input" in self.payload:
|
132
|
+
# Convert input to messages format for token calculation
|
133
|
+
input_val = self.payload["input"]
|
134
|
+
if isinstance(input_val, str):
|
135
|
+
messages = [{"role": "user", "content": input_val}]
|
136
|
+
elif isinstance(input_val, list):
|
137
|
+
# Handle array input format
|
138
|
+
messages = []
|
139
|
+
for item in input_val:
|
140
|
+
if isinstance(item, str):
|
141
|
+
messages.append({"role": "user", "content": item})
|
142
|
+
elif isinstance(item, dict) and "type" in item:
|
143
|
+
# Handle structured input items
|
144
|
+
if item["type"] == "message":
|
145
|
+
messages.append(item)
|
146
|
+
else:
|
147
|
+
return None
|
148
|
+
return TokenCalculator.calculate_message_tokens(
|
149
|
+
messages, **self.payload
|
150
|
+
)
|
151
|
+
# Handle embeddings endpoint
|
152
|
+
elif "embed" in self.endpoint.config.endpoint:
|
153
|
+
return TokenCalculator.calculate_embed_token(**self.payload)
|
154
|
+
|
155
|
+
return None
|
156
|
+
|
157
|
+
async def invoke(self) -> None:
|
158
|
+
"""Execute the API call through the endpoint.
|
159
|
+
|
160
|
+
Updates execution status and stores the response or error.
|
161
|
+
"""
|
162
|
+
start = asyncio.get_event_loop().time()
|
163
|
+
|
164
|
+
try:
|
165
|
+
self.execution.status = EventStatus.PROCESSING
|
166
|
+
|
167
|
+
# Make the API call
|
168
|
+
response = await self.endpoint.call(
|
169
|
+
request=self.payload,
|
170
|
+
cache_control=self.cache_control,
|
171
|
+
extra_headers=self.headers if self.headers else None,
|
172
|
+
)
|
173
|
+
|
174
|
+
self.execution.response = response
|
175
|
+
self.execution.status = EventStatus.COMPLETED
|
176
|
+
|
177
|
+
except asyncio.CancelledError:
|
178
|
+
self.execution.error = "API call cancelled"
|
179
|
+
self.execution.status = EventStatus.FAILED
|
180
|
+
raise
|
181
|
+
|
182
|
+
except Exception as e:
|
183
|
+
self.execution.error = str(e)
|
184
|
+
self.execution.status = EventStatus.FAILED
|
185
|
+
logger.error(f"API call failed: {e}")
|
186
|
+
|
187
|
+
finally:
|
188
|
+
self.execution.duration = asyncio.get_event_loop().time() - start
|
189
|
+
|
190
|
+
async def stream(self):
|
191
|
+
"""Stream the API response through the endpoint.
|
192
|
+
|
193
|
+
Yields:
|
194
|
+
Streaming chunks from the API.
|
195
|
+
"""
|
196
|
+
start = asyncio.get_event_loop().time()
|
197
|
+
response = []
|
198
|
+
|
199
|
+
try:
|
200
|
+
self.execution.status = EventStatus.PROCESSING
|
201
|
+
|
202
|
+
async for chunk in self.endpoint.stream(
|
203
|
+
request=self.payload,
|
204
|
+
extra_headers=self.headers if self.headers else None,
|
205
|
+
):
|
206
|
+
response.append(chunk)
|
207
|
+
yield chunk
|
208
|
+
|
209
|
+
self.execution.response = response
|
210
|
+
self.execution.status = EventStatus.COMPLETED
|
211
|
+
|
212
|
+
except Exception as e:
|
213
|
+
self.execution.error = str(e)
|
214
|
+
self.execution.status = EventStatus.FAILED
|
215
|
+
logger.error(f"Streaming failed: {e}")
|
216
|
+
|
217
|
+
finally:
|
218
|
+
self.execution.duration = asyncio.get_event_loop().time() - start
|
219
|
+
|
220
|
+
@property
|
221
|
+
def request(self) -> dict:
|
222
|
+
"""Get request information including token usage."""
|
223
|
+
return {
|
224
|
+
"required_tokens": self.required_tokens,
|
225
|
+
}
|
226
|
+
|
227
|
+
@property
|
228
|
+
def response(self):
|
229
|
+
"""Get the response from the execution."""
|
230
|
+
return self.execution.response if self.execution else None
|
@@ -0,0 +1,410 @@
|
|
1
|
+
# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
|
2
|
+
#
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
4
|
+
|
5
|
+
import asyncio
|
6
|
+
import logging
|
7
|
+
|
8
|
+
import aiohttp
|
9
|
+
import backoff
|
10
|
+
from aiocache import cached
|
11
|
+
from pydantic import BaseModel
|
12
|
+
|
13
|
+
from lionagi.config import settings
|
14
|
+
from lionagi.service.resilience import (
|
15
|
+
CircuitBreaker,
|
16
|
+
RetryConfig,
|
17
|
+
retry_with_backoff,
|
18
|
+
)
|
19
|
+
from lionagi.utils import to_dict
|
20
|
+
|
21
|
+
from .endpoint_config import EndpointConfig
|
22
|
+
from .header_factory import HeaderFactory
|
23
|
+
|
24
|
+
logger = logging.getLogger(__name__)
|
25
|
+
|
26
|
+
|
27
|
+
class Endpoint:
|
28
|
+
def __init__(
|
29
|
+
self,
|
30
|
+
config: dict | EndpointConfig,
|
31
|
+
circuit_breaker: CircuitBreaker | None = None,
|
32
|
+
retry_config: RetryConfig | None = None,
|
33
|
+
**kwargs,
|
34
|
+
):
|
35
|
+
"""
|
36
|
+
Initialize the endpoint.
|
37
|
+
|
38
|
+
This endpoint is designed to be stateless and thread-safe for parallel operations.
|
39
|
+
Each API call will create its own client session to avoid conflicts.
|
40
|
+
|
41
|
+
Args:
|
42
|
+
config: The endpoint configuration.
|
43
|
+
circuit_breaker: Optional circuit breaker for resilience.
|
44
|
+
retry_config: Optional retry configuration for resilience.
|
45
|
+
**kwargs: Additional keyword arguments to update the configuration.
|
46
|
+
"""
|
47
|
+
if isinstance(config, dict):
|
48
|
+
_config = EndpointConfig(**config, **kwargs)
|
49
|
+
elif isinstance(config, EndpointConfig):
|
50
|
+
_config = config.model_copy(
|
51
|
+
deep=True
|
52
|
+
) # Use deep copy to avoid sharing kwargs dict
|
53
|
+
_config.update(**kwargs)
|
54
|
+
else:
|
55
|
+
raise ValueError(
|
56
|
+
"Config must be a dict or EndpointConfig instance"
|
57
|
+
)
|
58
|
+
self.config = _config
|
59
|
+
self.circuit_breaker = circuit_breaker
|
60
|
+
self.retry_config = retry_config
|
61
|
+
|
62
|
+
logger.debug(
|
63
|
+
f"Initialized Endpoint with provider={self.config.provider}, "
|
64
|
+
f"endpoint={self.config.endpoint}, circuit_breaker={circuit_breaker is not None}, "
|
65
|
+
f"retry_config={retry_config is not None}"
|
66
|
+
)
|
67
|
+
|
68
|
+
def _create_http_session(self):
|
69
|
+
"""Create a new HTTP session (not thread-safe, create new for each request)."""
|
70
|
+
return aiohttp.ClientSession(
|
71
|
+
timeout=aiohttp.ClientTimeout(self.config.timeout),
|
72
|
+
**self.config.client_kwargs,
|
73
|
+
)
|
74
|
+
|
75
|
+
# Removed old context manager methods - endpoint is now stateless
|
76
|
+
|
77
|
+
@property
|
78
|
+
def request_options(self):
|
79
|
+
return self.config.request_options
|
80
|
+
|
81
|
+
@request_options.setter
|
82
|
+
def request_options(self, value):
|
83
|
+
self.config.request_options = EndpointConfig._validate_request_options(
|
84
|
+
value
|
85
|
+
)
|
86
|
+
|
87
|
+
def create_payload(
|
88
|
+
self,
|
89
|
+
request: dict | BaseModel,
|
90
|
+
extra_headers: dict | None = None,
|
91
|
+
**kwargs,
|
92
|
+
):
|
93
|
+
# First, create headers
|
94
|
+
headers = HeaderFactory.get_header(
|
95
|
+
auth_type=self.config.auth_type,
|
96
|
+
content_type=self.config.content_type,
|
97
|
+
api_key=self.config._api_key,
|
98
|
+
default_headers=self.config.default_headers,
|
99
|
+
)
|
100
|
+
if extra_headers:
|
101
|
+
headers.update(extra_headers)
|
102
|
+
|
103
|
+
# Convert request to dict if it's a BaseModel
|
104
|
+
request = (
|
105
|
+
request
|
106
|
+
if isinstance(request, dict)
|
107
|
+
else request.model_dump(exclude_none=True)
|
108
|
+
)
|
109
|
+
|
110
|
+
# Start with config defaults
|
111
|
+
payload = self.config.kwargs.copy()
|
112
|
+
|
113
|
+
# Update with request data
|
114
|
+
payload.update(request)
|
115
|
+
|
116
|
+
# Update with additional kwargs
|
117
|
+
if kwargs:
|
118
|
+
payload.update(kwargs)
|
119
|
+
|
120
|
+
# If we have request_options, use the model's fields to filter valid params
|
121
|
+
if self.config.request_options is not None:
|
122
|
+
# Get valid field names from the model
|
123
|
+
valid_fields = set(self.config.request_options.model_fields.keys())
|
124
|
+
|
125
|
+
# Filter payload to only include valid fields
|
126
|
+
filtered_payload = {
|
127
|
+
k: v for k, v in payload.items() if k in valid_fields
|
128
|
+
}
|
129
|
+
|
130
|
+
# Validate the filtered payload
|
131
|
+
payload = self.config.validate_payload(filtered_payload)
|
132
|
+
else:
|
133
|
+
# If no request_options, we still need to remove obvious non-API params
|
134
|
+
# These are parameters that are never part of any API payload
|
135
|
+
non_api_params = {
|
136
|
+
"task",
|
137
|
+
"provider",
|
138
|
+
"base_url",
|
139
|
+
"endpoint",
|
140
|
+
"endpoint_params",
|
141
|
+
"api_key",
|
142
|
+
"queue_capacity",
|
143
|
+
"capacity_refresh_time",
|
144
|
+
"interval",
|
145
|
+
"limit_requests",
|
146
|
+
"limit_tokens",
|
147
|
+
"invoke_with_endpoint",
|
148
|
+
"extra_headers",
|
149
|
+
"headers",
|
150
|
+
"cache_control",
|
151
|
+
"include_token_usage_to_model",
|
152
|
+
"chat_model",
|
153
|
+
"imodel",
|
154
|
+
"branch",
|
155
|
+
}
|
156
|
+
payload = {
|
157
|
+
k: v for k, v in payload.items() if k not in non_api_params
|
158
|
+
}
|
159
|
+
|
160
|
+
return (payload, headers)
|
161
|
+
|
162
|
+
async def call(
|
163
|
+
self, request: dict | BaseModel, cache_control: bool = False, **kwargs
|
164
|
+
):
|
165
|
+
"""
|
166
|
+
Make a call to the endpoint.
|
167
|
+
|
168
|
+
Args:
|
169
|
+
request: The request parameters or model.
|
170
|
+
cache_control: Whether to use cache control.
|
171
|
+
**kwargs: Additional keyword arguments for the request.
|
172
|
+
|
173
|
+
Returns:
|
174
|
+
The response from the endpoint.
|
175
|
+
"""
|
176
|
+
# Extract extra_headers before passing to create_payload
|
177
|
+
extra_headers = kwargs.pop("extra_headers", None)
|
178
|
+
payload, headers = self.create_payload(
|
179
|
+
request, extra_headers=extra_headers, **kwargs
|
180
|
+
)
|
181
|
+
|
182
|
+
async def _call(payload: dict, headers: dict, **kwargs):
|
183
|
+
# Direct call without context manager - each method handles its own resources
|
184
|
+
return await self._call_aiohttp(
|
185
|
+
payload=payload, headers=headers, **kwargs
|
186
|
+
)
|
187
|
+
|
188
|
+
# Apply resilience patterns if configured
|
189
|
+
call_func = _call
|
190
|
+
|
191
|
+
# Apply retry if configured
|
192
|
+
if self.retry_config:
|
193
|
+
|
194
|
+
async def call_func(p, h, **kw):
|
195
|
+
return await retry_with_backoff(
|
196
|
+
_call, p, h, **kw, **self.retry_config.as_kwargs()
|
197
|
+
)
|
198
|
+
|
199
|
+
# Apply circuit breaker if configured
|
200
|
+
if self.circuit_breaker:
|
201
|
+
if self.retry_config:
|
202
|
+
# If both are configured, apply circuit breaker to the retry-wrapped function
|
203
|
+
if not cache_control:
|
204
|
+
return await self.circuit_breaker.execute(
|
205
|
+
call_func, payload, headers, **kwargs
|
206
|
+
)
|
207
|
+
else:
|
208
|
+
# If only circuit breaker is configured, apply it directly
|
209
|
+
if not cache_control:
|
210
|
+
return await self.circuit_breaker.execute(
|
211
|
+
_call, payload, headers, **kwargs
|
212
|
+
)
|
213
|
+
|
214
|
+
# Handle caching if requested
|
215
|
+
if cache_control:
|
216
|
+
|
217
|
+
@cached(**settings.aiocache_config.as_kwargs())
|
218
|
+
async def _cached_call(payload: dict, headers: dict, **kwargs):
|
219
|
+
# Apply resilience patterns to cached call if configured
|
220
|
+
if self.circuit_breaker and self.retry_config:
|
221
|
+
return await self.circuit_breaker.execute(
|
222
|
+
call_func, payload, headers, **kwargs
|
223
|
+
)
|
224
|
+
if self.circuit_breaker:
|
225
|
+
return await self.circuit_breaker.execute(
|
226
|
+
_call, payload, headers, **kwargs
|
227
|
+
)
|
228
|
+
if self.retry_config:
|
229
|
+
return await call_func(payload, headers, **kwargs)
|
230
|
+
|
231
|
+
return await _call(payload, headers, **kwargs)
|
232
|
+
|
233
|
+
return await _cached_call(payload, headers, **kwargs)
|
234
|
+
|
235
|
+
# No caching, apply resilience patterns directly
|
236
|
+
if self.retry_config:
|
237
|
+
return await call_func(payload, headers, **kwargs)
|
238
|
+
|
239
|
+
return await _call(payload, headers, **kwargs)
|
240
|
+
|
241
|
+
async def _call_aiohttp(self, payload: dict, headers: dict, **kwargs):
|
242
|
+
"""
|
243
|
+
Make a call using aiohttp with a fresh session for each request.
|
244
|
+
|
245
|
+
Args:
|
246
|
+
payload: The request payload.
|
247
|
+
headers: The request headers.
|
248
|
+
**kwargs: Additional keyword arguments for the request.
|
249
|
+
|
250
|
+
Returns:
|
251
|
+
The response from the endpoint.
|
252
|
+
"""
|
253
|
+
|
254
|
+
async def _make_request_with_backoff():
|
255
|
+
# Create a new session for this request
|
256
|
+
async with self._create_http_session() as session:
|
257
|
+
response = None
|
258
|
+
try:
|
259
|
+
response = await session.request(
|
260
|
+
method=self.config.method,
|
261
|
+
url=self.config.full_url,
|
262
|
+
headers=headers,
|
263
|
+
json=payload,
|
264
|
+
**kwargs,
|
265
|
+
)
|
266
|
+
|
267
|
+
# Check for rate limit or server errors that should be retried
|
268
|
+
if response.status == 429 or response.status >= 500:
|
269
|
+
response.raise_for_status() # This will be caught by backoff
|
270
|
+
elif response.status != 200:
|
271
|
+
# Try to get error details from response body
|
272
|
+
try:
|
273
|
+
error_body = await response.json()
|
274
|
+
error_message = f"Request failed with status {response.status}: {error_body}"
|
275
|
+
except:
|
276
|
+
error_message = (
|
277
|
+
f"Request failed with status {response.status}"
|
278
|
+
)
|
279
|
+
|
280
|
+
raise aiohttp.ClientResponseError(
|
281
|
+
request_info=response.request_info,
|
282
|
+
history=response.history,
|
283
|
+
status=response.status,
|
284
|
+
message=error_message,
|
285
|
+
headers=response.headers,
|
286
|
+
)
|
287
|
+
|
288
|
+
# Extract and return the JSON response
|
289
|
+
return await response.json()
|
290
|
+
finally:
|
291
|
+
# Ensure response is properly released if coroutine is cancelled between retries
|
292
|
+
if response is not None and not response.closed:
|
293
|
+
await response.release()
|
294
|
+
|
295
|
+
# Define a giveup function for backoff
|
296
|
+
def giveup_on_client_error(e):
|
297
|
+
# Don't retry on 4xx errors except 429 (rate limit)
|
298
|
+
if isinstance(e, aiohttp.ClientResponseError):
|
299
|
+
return 400 <= e.status < 500 and e.status != 429
|
300
|
+
return False
|
301
|
+
|
302
|
+
# Use backoff for retries with exponential backoff and jitter
|
303
|
+
# Moved inside the method to reference runtime config
|
304
|
+
backoff_handler = backoff.on_exception(
|
305
|
+
backoff.expo,
|
306
|
+
(aiohttp.ClientError, asyncio.TimeoutError),
|
307
|
+
max_tries=self.config.max_retries,
|
308
|
+
giveup=giveup_on_client_error,
|
309
|
+
jitter=backoff.full_jitter,
|
310
|
+
)
|
311
|
+
|
312
|
+
# Apply the decorator at runtime
|
313
|
+
return await backoff_handler(_make_request_with_backoff)()
|
314
|
+
|
315
|
+
async def stream(
|
316
|
+
self,
|
317
|
+
request: dict | BaseModel,
|
318
|
+
extra_headers: dict | None = None,
|
319
|
+
**kwargs,
|
320
|
+
):
|
321
|
+
"""
|
322
|
+
Stream responses from the endpoint.
|
323
|
+
|
324
|
+
Args:
|
325
|
+
request: The request parameters or model.
|
326
|
+
extra_headers: Additional headers for the request.
|
327
|
+
**kwargs: Additional keyword arguments for the request.
|
328
|
+
|
329
|
+
Yields:
|
330
|
+
Streaming chunks from the API.
|
331
|
+
"""
|
332
|
+
payload, headers = self.create_payload(
|
333
|
+
request, extra_headers, **kwargs
|
334
|
+
)
|
335
|
+
|
336
|
+
# Direct streaming without context manager
|
337
|
+
async for chunk in self._stream_aiohttp(
|
338
|
+
payload=payload, headers=headers, **kwargs
|
339
|
+
):
|
340
|
+
yield chunk
|
341
|
+
|
342
|
+
async def _stream_aiohttp(self, payload: dict, headers: dict, **kwargs):
|
343
|
+
"""
|
344
|
+
Stream responses using aiohttp with a fresh session.
|
345
|
+
|
346
|
+
Args:
|
347
|
+
payload: The request payload.
|
348
|
+
headers: The request headers.
|
349
|
+
**kwargs: Additional keyword arguments for the request.
|
350
|
+
|
351
|
+
Yields:
|
352
|
+
Streaming chunks from the API.
|
353
|
+
"""
|
354
|
+
# Ensure stream is enabled
|
355
|
+
payload["stream"] = True
|
356
|
+
|
357
|
+
# Create a new session for streaming
|
358
|
+
async with self._create_http_session() as session:
|
359
|
+
async with session.request(
|
360
|
+
method=self.config.method,
|
361
|
+
url=self.config.full_url,
|
362
|
+
headers=headers,
|
363
|
+
json=payload,
|
364
|
+
**kwargs,
|
365
|
+
) as response:
|
366
|
+
if response.status != 200:
|
367
|
+
raise aiohttp.ClientResponseError(
|
368
|
+
request_info=response.request_info,
|
369
|
+
history=response.history,
|
370
|
+
status=response.status,
|
371
|
+
message=f"Request failed with status {response.status}",
|
372
|
+
headers=response.headers,
|
373
|
+
)
|
374
|
+
|
375
|
+
async for line in response.content:
|
376
|
+
if line:
|
377
|
+
yield line.decode("utf-8")
|
378
|
+
|
379
|
+
def to_dict(self):
|
380
|
+
return {
|
381
|
+
"retry_config": (
|
382
|
+
self.retry_config.to_dict() if self.retry_config else None
|
383
|
+
),
|
384
|
+
"circuit_breaker": (
|
385
|
+
self.circuit_breaker.to_dict()
|
386
|
+
if self.circuit_breaker
|
387
|
+
else None
|
388
|
+
),
|
389
|
+
"config": self.config.model_dump(exclude_none=True),
|
390
|
+
}
|
391
|
+
|
392
|
+
@classmethod
|
393
|
+
def from_dict(cls, data: dict):
|
394
|
+
data = to_dict(data, recursive=True)
|
395
|
+
retry_config = data.get("retry_config")
|
396
|
+
circuit_breaker = data.get("circuit_breaker")
|
397
|
+
config = data.get("config")
|
398
|
+
|
399
|
+
if retry_config:
|
400
|
+
retry_config = RetryConfig(**retry_config)
|
401
|
+
if circuit_breaker:
|
402
|
+
circuit_breaker = CircuitBreaker(**circuit_breaker)
|
403
|
+
if config:
|
404
|
+
config = EndpointConfig(**config)
|
405
|
+
|
406
|
+
return cls(
|
407
|
+
config=config,
|
408
|
+
circuit_breaker=circuit_breaker,
|
409
|
+
retry_config=retry_config,
|
410
|
+
)
|