lionagi 0.12.3__py3-none-any.whl → 0.12.5__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/libs/schema/load_pydantic_model_from_schema.py +259 -0
- lionagi/libs/token_transform/perplexity.py +2 -4
- 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/pile.py +94 -33
- lionagi/protocols/graph/node.py +25 -19
- lionagi/protocols/messages/assistant_response.py +20 -1
- 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/anthropic_models.py +159 -0
- lionagi/service/{providers/exa_/models.py → third_party/exa_models.py} +18 -13
- 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/version.py +1 -1
- {lionagi-0.12.3.dist-info → lionagi-0.12.5.dist-info}/METADATA +17 -19
- {lionagi-0.12.3.dist-info → lionagi-0.12.5.dist-info}/RECORD +43 -54
- lionagi/adapters/__init__.py +0 -1
- lionagi/adapters/adapter.py +0 -120
- lionagi/adapters/json_adapter.py +0 -181
- lionagi/adapters/pandas_/csv_adapter.py +0 -94
- lionagi/adapters/pandas_/excel_adapter.py +0 -94
- lionagi/adapters/pandas_/pd_dataframe_adapter.py +0 -81
- lionagi/adapters/pandas_/pd_series_adapter.py +0 -57
- lionagi/adapters/toml_adapter.py +0 -204
- lionagi/adapters/types.py +0 -21
- 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_/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 -144
- lionagi/service/providers/types.py +0 -17
- /lionagi/{adapters/pandas_/__init__.py → py.typed} +0 -0
- /lionagi/service/{providers/exa_ → third_party}/__init__.py +0 -0
- /lionagi/service/{endpoints/token_calculator.py → token_calculator.py} +0 -0
- {lionagi-0.12.3.dist-info → lionagi-0.12.5.dist-info}/WHEEL +0 -0
- {lionagi-0.12.3.dist-info → lionagi-0.12.5.dist-info}/licenses/LICENSE +0 -0
@@ -5,6 +5,7 @@
|
|
5
5
|
from __future__ import annotations
|
6
6
|
|
7
7
|
import asyncio
|
8
|
+
import json
|
8
9
|
import threading
|
9
10
|
from collections import deque
|
10
11
|
from collections.abc import (
|
@@ -19,20 +20,15 @@ from pathlib import Path
|
|
19
20
|
from typing import Any, ClassVar, Generic, TypeVar
|
20
21
|
|
21
22
|
import pandas as pd
|
22
|
-
from pydantic import Field
|
23
|
+
from pydantic import Field
|
23
24
|
from pydantic.fields import FieldInfo
|
25
|
+
from pydapter import Adapter, AdapterRegistry
|
26
|
+
from pydapter.adapters import CsvAdapter, JsonAdapter
|
27
|
+
from pydapter.extras.excel_ import ExcelAdapter
|
28
|
+
from pydapter.extras.pandas_ import DataFrameAdapter
|
24
29
|
from typing_extensions import Self, override
|
25
30
|
|
26
31
|
from lionagi._errors import ItemExistsError, ItemNotFoundError
|
27
|
-
from lionagi.adapters.types import (
|
28
|
-
Adapter,
|
29
|
-
AdapterRegistry,
|
30
|
-
CSVFileAdapter,
|
31
|
-
ExcelFileAdapter,
|
32
|
-
JsonAdapter,
|
33
|
-
JsonFileAdapter,
|
34
|
-
PandasDataFrameAdapter,
|
35
|
-
)
|
36
32
|
from lionagi.utils import UNDEFINED, is_same_dtype, to_list
|
37
33
|
|
38
34
|
from .._concepts import Observable
|
@@ -45,10 +41,9 @@ T = TypeVar("T", bound=E)
|
|
45
41
|
|
46
42
|
PILE_DEFAULT_ADAPTERS = (
|
47
43
|
JsonAdapter,
|
48
|
-
|
49
|
-
|
50
|
-
|
51
|
-
PandasDataFrameAdapter,
|
44
|
+
CsvAdapter,
|
45
|
+
ExcelAdapter,
|
46
|
+
DataFrameAdapter,
|
52
47
|
)
|
53
48
|
|
54
49
|
|
@@ -56,8 +51,9 @@ class PileAdapterRegistry(AdapterRegistry):
|
|
56
51
|
pass
|
57
52
|
|
58
53
|
|
54
|
+
pile_adapter_registry = PileAdapterRegistry()
|
59
55
|
for i in PILE_DEFAULT_ADAPTERS:
|
60
|
-
|
56
|
+
pile_adapter_registry.register(i)
|
61
57
|
|
62
58
|
|
63
59
|
__all__ = (
|
@@ -117,7 +113,7 @@ class Pile(Element, Collective[E], Generic[E]):
|
|
117
113
|
frozen=True,
|
118
114
|
)
|
119
115
|
|
120
|
-
_adapter_registry: ClassVar[AdapterRegistry] =
|
116
|
+
_adapter_registry: ClassVar[AdapterRegistry] = pile_adapter_registry
|
121
117
|
|
122
118
|
def __pydantic_extra__(self) -> dict[str, FieldInfo]:
|
123
119
|
return {
|
@@ -910,9 +906,23 @@ class Pile(Element, Collective[E], Generic[E]):
|
|
910
906
|
self.progression.insert(index, item_order)
|
911
907
|
self.collections.update(item_dict)
|
912
908
|
|
913
|
-
|
914
|
-
|
915
|
-
|
909
|
+
def to_dict(self) -> dict[str, Any]:
|
910
|
+
"""Convert pile to dictionary, properly handling collections."""
|
911
|
+
# Get base dict from parent class
|
912
|
+
dict_ = super().to_dict()
|
913
|
+
|
914
|
+
# Manually serialize collections
|
915
|
+
collections_list = []
|
916
|
+
for item in self.collections.values():
|
917
|
+
if hasattr(item, "to_dict"):
|
918
|
+
collections_list.append(item.to_dict())
|
919
|
+
elif hasattr(item, "model_dump"):
|
920
|
+
collections_list.append(item.model_dump())
|
921
|
+
else:
|
922
|
+
collections_list.append(str(item))
|
923
|
+
|
924
|
+
dict_["collections"] = collections_list
|
925
|
+
return dict_
|
916
926
|
|
917
927
|
class AsyncPileIterator:
|
918
928
|
def __init__(self, pile: Pile):
|
@@ -952,12 +962,25 @@ class Pile(Element, Collective[E], Generic[E]):
|
|
952
962
|
|
953
963
|
def adapt_to(self, obj_key: str, /, **kwargs: Any) -> Any:
|
954
964
|
"""Convert to another format."""
|
955
|
-
|
965
|
+
# For JSON adapter, we need to pass the dict representation
|
966
|
+
if obj_key in ["json", "csv", "toml"]:
|
967
|
+
data = self.to_dict()
|
956
968
|
|
957
|
-
|
958
|
-
|
959
|
-
|
960
|
-
|
969
|
+
# Create a simple object that has model_dump method
|
970
|
+
class _Wrapper:
|
971
|
+
def __init__(self, data):
|
972
|
+
self._data = data
|
973
|
+
|
974
|
+
def model_dump(self):
|
975
|
+
return self._data
|
976
|
+
|
977
|
+
wrapper = _Wrapper(data)
|
978
|
+
return self._get_adapter_registry().adapt_to(
|
979
|
+
wrapper, obj_key=obj_key, **kwargs
|
980
|
+
)
|
981
|
+
return self._get_adapter_registry().adapt_to(
|
982
|
+
self, obj_key=obj_key, **kwargs
|
983
|
+
)
|
961
984
|
|
962
985
|
@classmethod
|
963
986
|
def register_adapter(cls, adapter: type[Adapter]):
|
@@ -974,7 +997,7 @@ class Pile(Element, Collective[E], Generic[E]):
|
|
974
997
|
def adapt_from(cls, obj: Any, obj_key: str, /, **kwargs: Any):
|
975
998
|
"""Create from another format."""
|
976
999
|
dict_ = cls._get_adapter_registry().adapt_from(
|
977
|
-
cls, obj, obj_key, **kwargs
|
1000
|
+
cls, obj, obj_key=obj_key, **kwargs
|
978
1001
|
)
|
979
1002
|
if isinstance(dict_, list):
|
980
1003
|
dict_ = {"collections": dict_}
|
@@ -986,11 +1009,31 @@ class Pile(Element, Collective[E], Generic[E]):
|
|
986
1009
|
**kwargs: Any,
|
987
1010
|
) -> pd.DataFrame:
|
988
1011
|
"""Convert to DataFrame."""
|
989
|
-
|
1012
|
+
# For DataFrame, we need to pass a list of dicts
|
1013
|
+
data = [item.to_dict() for item in self.collections.values()]
|
1014
|
+
|
1015
|
+
# Create wrapper objects for each item
|
1016
|
+
class _ItemWrapper:
|
1017
|
+
def __init__(self, data):
|
1018
|
+
self._data = data
|
1019
|
+
|
1020
|
+
def model_dump(self):
|
1021
|
+
return self._data
|
1022
|
+
|
1023
|
+
wrappers = [_ItemWrapper(d) for d in data]
|
1024
|
+
df = self._get_adapter_registry().adapt_to(
|
1025
|
+
wrappers, obj_key="pd.DataFrame", many=True, **kwargs
|
1026
|
+
)
|
1027
|
+
|
1028
|
+
if columns:
|
1029
|
+
return df[columns]
|
1030
|
+
return df
|
990
1031
|
|
991
1032
|
def to_csv_file(self, fp: str | Path, **kwargs: Any) -> None:
|
992
1033
|
"""Save to CSV file."""
|
993
|
-
|
1034
|
+
# Convert to DataFrame first, then save as CSV
|
1035
|
+
df = self.to_df()
|
1036
|
+
df.to_csv(fp, index=False, **kwargs)
|
994
1037
|
|
995
1038
|
def to_json_file(
|
996
1039
|
self,
|
@@ -1011,8 +1054,14 @@ class Pile(Element, Collective[E], Generic[E]):
|
|
1011
1054
|
**kwargs: Additional arguments for json.dump() or DataFrame.to_json().
|
1012
1055
|
"""
|
1013
1056
|
if use_pd:
|
1014
|
-
return self.to_df().to_json(mode=mode, **kwargs)
|
1015
|
-
|
1057
|
+
return self.to_df().to_json(path_or_buf, mode=mode, **kwargs)
|
1058
|
+
|
1059
|
+
# Get JSON string from adapter
|
1060
|
+
json_str = self.adapt_to("json", many=many, **kwargs)
|
1061
|
+
|
1062
|
+
# Write to file
|
1063
|
+
with open(path_or_buf, mode, encoding="utf-8") as f:
|
1064
|
+
f.write(json_str)
|
1016
1065
|
|
1017
1066
|
|
1018
1067
|
def pile(
|
@@ -1038,16 +1087,28 @@ def pile(
|
|
1038
1087
|
"""
|
1039
1088
|
|
1040
1089
|
if df:
|
1041
|
-
return Pile.adapt_from(df, "
|
1090
|
+
return Pile.adapt_from(df, "pd.DataFrame", **kwargs)
|
1042
1091
|
|
1043
1092
|
if fp:
|
1044
1093
|
fp = Path(fp)
|
1045
1094
|
if fp.suffix == ".csv":
|
1046
|
-
|
1095
|
+
# Read CSV to DataFrame first
|
1096
|
+
df = pd.read_csv(fp, **kwargs)
|
1097
|
+
return Pile.adapt_from(df, "pd.DataFrame")
|
1047
1098
|
if fp.suffix == ".xlsx":
|
1048
|
-
|
1099
|
+
# Read Excel to DataFrame first
|
1100
|
+
df = pd.read_excel(fp, **kwargs)
|
1101
|
+
return Pile.adapt_from(df, "pd.DataFrame")
|
1049
1102
|
if fp.suffix in [".json", ".jsonl"]:
|
1050
|
-
|
1103
|
+
# Read JSON file
|
1104
|
+
with open(fp, encoding="utf-8") as f:
|
1105
|
+
data = json.load(f)
|
1106
|
+
if isinstance(data, dict):
|
1107
|
+
return Pile.from_dict(data)
|
1108
|
+
elif isinstance(data, list):
|
1109
|
+
return Pile.from_dict({"collections": data})
|
1110
|
+
else:
|
1111
|
+
raise ValueError(f"Invalid JSON data structure in {fp}")
|
1051
1112
|
|
1052
1113
|
return Pile(
|
1053
1114
|
collections,
|
lionagi/protocols/graph/node.py
CHANGED
@@ -6,26 +6,19 @@ import json
|
|
6
6
|
from typing import Any, ClassVar
|
7
7
|
|
8
8
|
from pydantic import field_validator
|
9
|
+
from pydapter import AdapterRegistry
|
10
|
+
from pydapter.adapters import JsonAdapter, TomlAdapter
|
11
|
+
from pydapter.extras.pandas_ import SeriesAdapter
|
9
12
|
|
10
13
|
from lionagi._class_registry import LION_CLASS_REGISTRY
|
11
|
-
from lionagi.adapters.types import (
|
12
|
-
AdapterRegistry,
|
13
|
-
JsonAdapter,
|
14
|
-
JsonFileAdapter,
|
15
|
-
PandasSeriesAdapter,
|
16
|
-
TomlAdapter,
|
17
|
-
TomlFileAdapter,
|
18
|
-
)
|
19
14
|
|
20
15
|
from .._concepts import Relational
|
21
16
|
from ..generic.element import Element
|
22
17
|
|
23
18
|
NODE_DEFAULT_ADAPTERS = (
|
24
19
|
JsonAdapter,
|
25
|
-
|
26
|
-
PandasSeriesAdapter,
|
20
|
+
SeriesAdapter,
|
27
21
|
TomlAdapter,
|
28
|
-
TomlFileAdapter,
|
29
22
|
)
|
30
23
|
|
31
24
|
|
@@ -33,8 +26,9 @@ class NodeAdapterRegistry(AdapterRegistry):
|
|
33
26
|
pass
|
34
27
|
|
35
28
|
|
29
|
+
node_adapter_registry = NodeAdapterRegistry()
|
36
30
|
for i in NODE_DEFAULT_ADAPTERS:
|
37
|
-
|
31
|
+
node_adapter_registry.register(i)
|
38
32
|
|
39
33
|
__all__ = ("Node",)
|
40
34
|
|
@@ -48,7 +42,7 @@ class Node(Element, Relational):
|
|
48
42
|
- Automatic subclass registration
|
49
43
|
"""
|
50
44
|
|
51
|
-
_adapter_registry: ClassVar[AdapterRegistry] =
|
45
|
+
_adapter_registry: ClassVar[AdapterRegistry] = node_adapter_registry
|
52
46
|
|
53
47
|
content: Any = None
|
54
48
|
embedding: list[float] | None = None
|
@@ -88,8 +82,24 @@ class Node(Element, Relational):
|
|
88
82
|
"""
|
89
83
|
Convert this Node to another format using a registered adapter.
|
90
84
|
"""
|
85
|
+
# For JSON/TOML adapters, we need to pass the dict representation
|
86
|
+
if obj_key in ["json", "toml"]:
|
87
|
+
data = self.to_dict()
|
88
|
+
|
89
|
+
# Create a simple object that has model_dump method
|
90
|
+
class _Wrapper:
|
91
|
+
def __init__(self, data):
|
92
|
+
self._data = data
|
93
|
+
|
94
|
+
def model_dump(self):
|
95
|
+
return self._data
|
96
|
+
|
97
|
+
wrapper = _Wrapper(data)
|
98
|
+
return self._get_adapter_registry().adapt_to(
|
99
|
+
wrapper, obj_key=obj_key, many=many, **kwargs
|
100
|
+
)
|
91
101
|
return self._get_adapter_registry().adapt_to(
|
92
|
-
self, obj_key, many=many, **kwargs
|
102
|
+
self, obj_key=obj_key, many=many, **kwargs
|
93
103
|
)
|
94
104
|
|
95
105
|
@classmethod
|
@@ -107,7 +117,7 @@ class Node(Element, Relational):
|
|
107
117
|
auto-delegate to the correct subclass via from_dict.
|
108
118
|
"""
|
109
119
|
result = cls._get_adapter_registry().adapt_from(
|
110
|
-
cls, obj, obj_key, many=many, **kwargs
|
120
|
+
cls, obj, obj_key=obj_key, many=many, **kwargs
|
111
121
|
)
|
112
122
|
# If adapter returned multiple items, choose the first or handle as needed.
|
113
123
|
if isinstance(result, list):
|
@@ -124,9 +134,5 @@ class Node(Element, Relational):
|
|
124
134
|
def register_adapter(cls, adapter: Any) -> None:
|
125
135
|
cls._get_adapter_registry().register(adapter)
|
126
136
|
|
127
|
-
@classmethod
|
128
|
-
def list_adapters(cls) -> list[str]:
|
129
|
-
return cls._get_adapter_registry().list_adapters()
|
130
|
-
|
131
137
|
|
132
138
|
# File: lionagi/protocols/graph/node.py
|
@@ -46,7 +46,7 @@ def prepare_assistant_response(
|
|
46
46
|
elif isinstance(j, str):
|
47
47
|
text_contents.append(j)
|
48
48
|
|
49
|
-
# openai standard
|
49
|
+
# openai chat completions standard
|
50
50
|
elif "choices" in i:
|
51
51
|
choices = i["choices"]
|
52
52
|
choices = (
|
@@ -58,6 +58,25 @@ def prepare_assistant_response(
|
|
58
58
|
elif "delta" in j:
|
59
59
|
text_contents.append(j["delta"]["content"] or "")
|
60
60
|
|
61
|
+
# openai responses API standard
|
62
|
+
elif "output" in i:
|
63
|
+
output = i["output"]
|
64
|
+
output = [output] if not isinstance(output, list) else output
|
65
|
+
for item in output:
|
66
|
+
if isinstance(item, dict):
|
67
|
+
if item.get("type") == "message":
|
68
|
+
# Extract content from message
|
69
|
+
content = item.get("content", [])
|
70
|
+
if isinstance(content, list):
|
71
|
+
for c in content:
|
72
|
+
if (
|
73
|
+
isinstance(c, dict)
|
74
|
+
and c.get("type") == "output_text"
|
75
|
+
):
|
76
|
+
text_contents.append(c.get("text", ""))
|
77
|
+
elif isinstance(c, str):
|
78
|
+
text_contents.append(c)
|
79
|
+
|
61
80
|
elif isinstance(i, str):
|
62
81
|
text_contents.append(i)
|
63
82
|
|
@@ -0,0 +1,15 @@
|
|
1
|
+
# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
|
2
|
+
#
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
4
|
+
|
5
|
+
from .endpoint import Endpoint
|
6
|
+
from .endpoint_config import EndpointConfig
|
7
|
+
from .header_factory import HeaderFactory
|
8
|
+
from .match_endpoint import match_endpoint
|
9
|
+
|
10
|
+
__all__ = (
|
11
|
+
"Endpoint",
|
12
|
+
"EndpointConfig",
|
13
|
+
"HeaderFactory",
|
14
|
+
"match_endpoint",
|
15
|
+
)
|
@@ -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
|