cua-agent 0.4.34__tar.gz → 0.4.36__tar.gz
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.
Potentially problematic release.
This version of cua-agent might be problematic. Click here for more details.
- {cua_agent-0.4.34 → cua_agent-0.4.36}/PKG-INFO +18 -10
- {cua_agent-0.4.34 → cua_agent-0.4.36}/README.md +9 -8
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/__init__.py +4 -10
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/__main__.py +2 -1
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/huggingfacelocal_adapter.py +54 -61
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/human_adapter.py +116 -114
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/mlxvlm_adapter.py +110 -99
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/models/__init__.py +14 -6
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/models/generic.py +7 -4
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/models/internvl.py +66 -30
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/models/opencua.py +23 -8
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/models/qwen2_5_vl.py +7 -4
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/agent.py +184 -158
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/__init__.py +4 -4
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/base.py +45 -31
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/budget_manager.py +22 -10
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/image_retention.py +18 -13
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/logging.py +55 -42
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/operator_validator.py +3 -1
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/pii_anonymization.py +19 -16
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/telemetry.py +67 -61
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/trajectory_saver.py +90 -70
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/cli.py +115 -110
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/computers/__init__.py +13 -8
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/computers/base.py +32 -19
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/computers/cua.py +33 -25
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/computers/custom.py +78 -71
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/decorators.py +23 -14
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/human_tool/__init__.py +2 -7
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/human_tool/__main__.py +6 -2
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/human_tool/server.py +48 -37
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/human_tool/ui.py +235 -185
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/integrations/hud/__init__.py +15 -21
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/integrations/hud/agent.py +101 -83
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/integrations/hud/proxy.py +90 -57
- cua_agent-0.4.36/agent/loops/__init__.py +36 -0
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/anthropic.py +537 -483
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/base.py +13 -14
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/composed_grounded.py +135 -149
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/gemini.py +31 -12
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/glm45v.py +135 -133
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/gta1.py +47 -50
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/holo.py +4 -2
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/internvl.py +6 -11
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/moondream3.py +36 -12
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/omniparser.py +215 -210
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/openai.py +49 -50
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/opencua.py +29 -41
- cua_agent-0.4.36/agent/loops/qwen.py +510 -0
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/uitars.py +237 -202
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/proxy/examples.py +54 -50
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/proxy/handlers.py +27 -34
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/responses.py +330 -330
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/types.py +11 -5
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/ui/__init__.py +1 -1
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/ui/__main__.py +1 -1
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/ui/gradio/app.py +23 -18
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/ui/gradio/ui_components.py +310 -161
- {cua_agent-0.4.34 → cua_agent-0.4.36}/pyproject.toml +10 -2
- cua_agent-0.4.36/tests/conftest.py +84 -0
- cua_agent-0.4.36/tests/test_computer_agent.py +139 -0
- cua_agent-0.4.34/agent/loops/__init__.py +0 -32
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/adapters/__init__.py +0 -0
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/callbacks/prompt_instructions.py +0 -0
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/loops/model_types.csv +0 -0
- {cua_agent-0.4.34 → cua_agent-0.4.36}/agent/ui/gradio/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: cua-agent
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.36
|
|
4
4
|
Summary: CUA (Computer Use) Agent for AI-driven computer interaction
|
|
5
5
|
Author-Email: TryCua <gh@trycua.com>
|
|
6
6
|
Requires-Python: >=3.12
|
|
@@ -18,6 +18,10 @@ Requires-Dist: certifi>=2024.2.2
|
|
|
18
18
|
Requires-Dist: litellm>=1.74.12
|
|
19
19
|
Provides-Extra: openai
|
|
20
20
|
Provides-Extra: anthropic
|
|
21
|
+
Provides-Extra: qwen
|
|
22
|
+
Requires-Dist: qwen-vl-utils; extra == "qwen"
|
|
23
|
+
Requires-Dist: qwen-agent; extra == "qwen"
|
|
24
|
+
Requires-Dist: Pillow>=10.0.0; extra == "qwen"
|
|
21
25
|
Provides-Extra: omni
|
|
22
26
|
Requires-Dist: cua-som<0.2.0,>=0.1.0; extra == "omni"
|
|
23
27
|
Provides-Extra: uitars
|
|
@@ -34,7 +38,7 @@ Requires-Dist: transformers-v4.55.0-GLM-4.5V-preview; extra == "glm45v-hf"
|
|
|
34
38
|
Provides-Extra: opencua-hf
|
|
35
39
|
Requires-Dist: accelerate; extra == "opencua-hf"
|
|
36
40
|
Requires-Dist: torch; extra == "opencua-hf"
|
|
37
|
-
Requires-Dist: transformers
|
|
41
|
+
Requires-Dist: transformers>=4.53.0; extra == "opencua-hf"
|
|
38
42
|
Requires-Dist: tiktoken>=0.11.0; extra == "opencua-hf"
|
|
39
43
|
Requires-Dist: blobfile>=3.0.0; extra == "opencua-hf"
|
|
40
44
|
Provides-Extra: internvl-hf
|
|
@@ -70,6 +74,9 @@ Requires-Dist: python-dotenv>=1.0.1; extra == "all"
|
|
|
70
74
|
Requires-Dist: yaspin>=3.1.0; extra == "all"
|
|
71
75
|
Requires-Dist: hud-python==0.4.52; extra == "all"
|
|
72
76
|
Requires-Dist: google-genai>=1.41.0; extra == "all"
|
|
77
|
+
Requires-Dist: qwen-vl-utils; extra == "all"
|
|
78
|
+
Requires-Dist: qwen-agent; extra == "all"
|
|
79
|
+
Requires-Dist: Pillow>=10.0.0; extra == "all"
|
|
73
80
|
Description-Content-Type: text/markdown
|
|
74
81
|
|
|
75
82
|
<div align="center">
|
|
@@ -82,10 +89,11 @@ Description-Content-Type: text/markdown
|
|
|
82
89
|
</picture>
|
|
83
90
|
</div>
|
|
84
91
|
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
92
|
+
[](#)
|
|
93
|
+
[](#)
|
|
94
|
+
[](https://discord.com/invite/mVnXXpdE85)
|
|
95
|
+
[](https://pypi.org/project/cua-computer/)
|
|
96
|
+
|
|
89
97
|
</h1>
|
|
90
98
|
</div>
|
|
91
99
|
|
|
@@ -121,7 +129,7 @@ async def main():
|
|
|
121
129
|
name=os.getenv("CUA_CONTAINER_NAME"),
|
|
122
130
|
api_key=os.getenv("CUA_API_KEY")
|
|
123
131
|
) as computer:
|
|
124
|
-
|
|
132
|
+
|
|
125
133
|
# Create agent
|
|
126
134
|
agent = ComputerAgent(
|
|
127
135
|
model="anthropic/claude-3-5-sonnet-20241022",
|
|
@@ -130,10 +138,10 @@ async def main():
|
|
|
130
138
|
trajectory_dir="trajectories",
|
|
131
139
|
max_trajectory_budget=5.0 # $5 budget limit
|
|
132
140
|
)
|
|
133
|
-
|
|
141
|
+
|
|
134
142
|
# Run agent
|
|
135
143
|
messages = [{"role": "user", "content": "Take a screenshot and tell me what you see"}]
|
|
136
|
-
|
|
144
|
+
|
|
137
145
|
async for result in agent.run(messages):
|
|
138
146
|
for item in result["output"]:
|
|
139
147
|
if item["type"] == "message":
|
|
@@ -158,4 +166,4 @@ if __name__ == "__main__":
|
|
|
158
166
|
|
|
159
167
|
## License
|
|
160
168
|
|
|
161
|
-
MIT License - see LICENSE file for details.
|
|
169
|
+
MIT License - see LICENSE file for details.
|
|
@@ -8,10 +8,11 @@
|
|
|
8
8
|
</picture>
|
|
9
9
|
</div>
|
|
10
10
|
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
|
|
11
|
+
[](#)
|
|
12
|
+
[](#)
|
|
13
|
+
[](https://discord.com/invite/mVnXXpdE85)
|
|
14
|
+
[](https://pypi.org/project/cua-computer/)
|
|
15
|
+
|
|
15
16
|
</h1>
|
|
16
17
|
</div>
|
|
17
18
|
|
|
@@ -47,7 +48,7 @@ async def main():
|
|
|
47
48
|
name=os.getenv("CUA_CONTAINER_NAME"),
|
|
48
49
|
api_key=os.getenv("CUA_API_KEY")
|
|
49
50
|
) as computer:
|
|
50
|
-
|
|
51
|
+
|
|
51
52
|
# Create agent
|
|
52
53
|
agent = ComputerAgent(
|
|
53
54
|
model="anthropic/claude-3-5-sonnet-20241022",
|
|
@@ -56,10 +57,10 @@ async def main():
|
|
|
56
57
|
trajectory_dir="trajectories",
|
|
57
58
|
max_trajectory_budget=5.0 # $5 budget limit
|
|
58
59
|
)
|
|
59
|
-
|
|
60
|
+
|
|
60
61
|
# Run agent
|
|
61
62
|
messages = [{"role": "user", "content": "Take a screenshot and tell me what you see"}]
|
|
62
|
-
|
|
63
|
+
|
|
63
64
|
async for result in agent.run(messages):
|
|
64
65
|
for item in result["output"]:
|
|
65
66
|
if item["type"] == "message":
|
|
@@ -84,4 +85,4 @@ if __name__ == "__main__":
|
|
|
84
85
|
|
|
85
86
|
## License
|
|
86
87
|
|
|
87
|
-
MIT License - see LICENSE file for details.
|
|
88
|
+
MIT License - see LICENSE file for details.
|
|
@@ -5,19 +5,13 @@ agent - Decorator-based Computer Use Agent with liteLLM integration
|
|
|
5
5
|
import logging
|
|
6
6
|
import sys
|
|
7
7
|
|
|
8
|
-
from .decorators import register_agent
|
|
9
|
-
from .agent import ComputerAgent
|
|
10
|
-
from .types import Messages, AgentResponse
|
|
11
|
-
|
|
12
8
|
# Import loops to register them
|
|
13
9
|
from . import loops
|
|
10
|
+
from .agent import ComputerAgent
|
|
11
|
+
from .decorators import register_agent
|
|
12
|
+
from .types import AgentResponse, Messages
|
|
14
13
|
|
|
15
|
-
__all__ = [
|
|
16
|
-
"register_agent",
|
|
17
|
-
"ComputerAgent",
|
|
18
|
-
"Messages",
|
|
19
|
-
"AgentResponse"
|
|
20
|
-
]
|
|
14
|
+
__all__ = ["register_agent", "ComputerAgent", "Messages", "AgentResponse"]
|
|
21
15
|
|
|
22
16
|
__version__ = "0.4.0"
|
|
23
17
|
|
|
@@ -2,27 +2,30 @@ import asyncio
|
|
|
2
2
|
import functools
|
|
3
3
|
import warnings
|
|
4
4
|
from concurrent.futures import ThreadPoolExecutor
|
|
5
|
-
from typing import
|
|
6
|
-
|
|
5
|
+
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
|
|
6
|
+
|
|
7
|
+
from litellm import acompletion, completion
|
|
7
8
|
from litellm.llms.custom_llm import CustomLLM
|
|
8
|
-
from litellm import
|
|
9
|
+
from litellm.types.utils import GenericStreamingChunk, ModelResponse
|
|
9
10
|
|
|
10
11
|
# Try to import HuggingFace dependencies
|
|
11
12
|
try:
|
|
12
13
|
import torch
|
|
13
14
|
from transformers import AutoModelForImageTextToText, AutoProcessor
|
|
15
|
+
|
|
14
16
|
HF_AVAILABLE = True
|
|
15
17
|
except ImportError:
|
|
16
18
|
HF_AVAILABLE = False
|
|
17
19
|
|
|
18
20
|
from .models import load_model as load_model_handler
|
|
19
21
|
|
|
22
|
+
|
|
20
23
|
class HuggingFaceLocalAdapter(CustomLLM):
|
|
21
24
|
"""HuggingFace Local Adapter for running vision-language models locally."""
|
|
22
|
-
|
|
25
|
+
|
|
23
26
|
def __init__(self, device: str = "auto", trust_remote_code: bool = False, **kwargs):
|
|
24
27
|
"""Initialize the adapter.
|
|
25
|
-
|
|
28
|
+
|
|
26
29
|
Args:
|
|
27
30
|
device: Device to load model on ("auto", "cuda", "cpu", etc.)
|
|
28
31
|
trust_remote_code: Whether to trust remote code
|
|
@@ -34,129 +37,120 @@ class HuggingFaceLocalAdapter(CustomLLM):
|
|
|
34
37
|
# Cache for model handlers keyed by model_name
|
|
35
38
|
self._handlers: Dict[str, Any] = {}
|
|
36
39
|
self._executor = ThreadPoolExecutor(max_workers=1) # Single thread pool
|
|
37
|
-
|
|
40
|
+
|
|
38
41
|
def _get_handler(self, model_name: str):
|
|
39
42
|
"""Get or create a model handler for the given model name."""
|
|
40
43
|
if model_name not in self._handlers:
|
|
41
|
-
self._handlers[model_name] = load_model_handler(
|
|
44
|
+
self._handlers[model_name] = load_model_handler(
|
|
45
|
+
model_name=model_name, device=self.device, trust_remote_code=self.trust_remote_code
|
|
46
|
+
)
|
|
42
47
|
return self._handlers[model_name]
|
|
43
|
-
|
|
48
|
+
|
|
44
49
|
def _convert_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
45
50
|
"""Convert OpenAI format messages to HuggingFace format.
|
|
46
|
-
|
|
51
|
+
|
|
47
52
|
Args:
|
|
48
53
|
messages: Messages in OpenAI format
|
|
49
|
-
|
|
54
|
+
|
|
50
55
|
Returns:
|
|
51
56
|
Messages in HuggingFace format
|
|
52
57
|
"""
|
|
53
58
|
converted_messages = []
|
|
54
|
-
|
|
59
|
+
|
|
55
60
|
for message in messages:
|
|
56
|
-
converted_message = {
|
|
57
|
-
|
|
58
|
-
"content": []
|
|
59
|
-
}
|
|
60
|
-
|
|
61
|
+
converted_message = {"role": message["role"], "content": []}
|
|
62
|
+
|
|
61
63
|
content = message.get("content", [])
|
|
62
64
|
if isinstance(content, str):
|
|
63
65
|
# Simple text content
|
|
64
|
-
converted_message["content"].append({
|
|
65
|
-
"type": "text",
|
|
66
|
-
"text": content
|
|
67
|
-
})
|
|
66
|
+
converted_message["content"].append({"type": "text", "text": content})
|
|
68
67
|
elif isinstance(content, list):
|
|
69
68
|
# Multi-modal content
|
|
70
69
|
for item in content:
|
|
71
70
|
if item.get("type") == "text":
|
|
72
|
-
converted_message["content"].append(
|
|
73
|
-
"type": "text",
|
|
74
|
-
|
|
75
|
-
})
|
|
71
|
+
converted_message["content"].append(
|
|
72
|
+
{"type": "text", "text": item.get("text", "")}
|
|
73
|
+
)
|
|
76
74
|
elif item.get("type") == "image_url":
|
|
77
75
|
# Convert image_url format to image format
|
|
78
76
|
image_url = item.get("image_url", {}).get("url", "")
|
|
79
|
-
converted_message["content"].append({
|
|
80
|
-
|
|
81
|
-
"image": image_url
|
|
82
|
-
})
|
|
83
|
-
|
|
77
|
+
converted_message["content"].append({"type": "image", "image": image_url})
|
|
78
|
+
|
|
84
79
|
converted_messages.append(converted_message)
|
|
85
|
-
|
|
80
|
+
|
|
86
81
|
return converted_messages
|
|
87
|
-
|
|
82
|
+
|
|
88
83
|
def _generate(self, **kwargs) -> str:
|
|
89
84
|
"""Generate response using the local HuggingFace model.
|
|
90
|
-
|
|
85
|
+
|
|
91
86
|
Args:
|
|
92
87
|
**kwargs: Keyword arguments containing messages and model info
|
|
93
|
-
|
|
88
|
+
|
|
94
89
|
Returns:
|
|
95
90
|
Generated text response
|
|
96
91
|
"""
|
|
97
92
|
if not HF_AVAILABLE:
|
|
98
93
|
raise ImportError(
|
|
99
94
|
"HuggingFace transformers dependencies not found. "
|
|
100
|
-
|
|
95
|
+
'Please install with: pip install "cua-agent[uitars-hf]"'
|
|
101
96
|
)
|
|
102
|
-
|
|
97
|
+
|
|
103
98
|
# Extract messages and model from kwargs
|
|
104
|
-
messages = kwargs.get(
|
|
105
|
-
model_name = kwargs.get(
|
|
106
|
-
max_new_tokens = kwargs.get(
|
|
107
|
-
|
|
99
|
+
messages = kwargs.get("messages", [])
|
|
100
|
+
model_name = kwargs.get("model", "ByteDance-Seed/UI-TARS-1.5-7B")
|
|
101
|
+
max_new_tokens = kwargs.get("max_tokens", 128)
|
|
102
|
+
|
|
108
103
|
# Warn about ignored kwargs
|
|
109
|
-
ignored_kwargs = set(kwargs.keys()) - {
|
|
104
|
+
ignored_kwargs = set(kwargs.keys()) - {"messages", "model", "max_tokens"}
|
|
110
105
|
if ignored_kwargs:
|
|
111
106
|
warnings.warn(f"Ignoring unsupported kwargs: {ignored_kwargs}")
|
|
112
|
-
|
|
107
|
+
|
|
113
108
|
# Convert messages to HuggingFace format
|
|
114
109
|
hf_messages = self._convert_messages(messages)
|
|
115
|
-
|
|
110
|
+
|
|
116
111
|
# Delegate to model handler
|
|
117
112
|
handler = self._get_handler(model_name)
|
|
118
113
|
generated_text = handler.generate(hf_messages, max_new_tokens=max_new_tokens)
|
|
119
114
|
return generated_text
|
|
120
|
-
|
|
115
|
+
|
|
121
116
|
def completion(self, *args, **kwargs) -> ModelResponse:
|
|
122
117
|
"""Synchronous completion method.
|
|
123
|
-
|
|
118
|
+
|
|
124
119
|
Returns:
|
|
125
120
|
ModelResponse with generated text
|
|
126
121
|
"""
|
|
127
122
|
generated_text = self._generate(**kwargs)
|
|
128
|
-
|
|
123
|
+
|
|
129
124
|
return completion(
|
|
130
125
|
model=f"huggingface-local/{kwargs['model']}",
|
|
131
126
|
mock_response=generated_text,
|
|
132
127
|
)
|
|
133
|
-
|
|
128
|
+
|
|
134
129
|
async def acompletion(self, *args, **kwargs) -> ModelResponse:
|
|
135
130
|
"""Asynchronous completion method.
|
|
136
|
-
|
|
131
|
+
|
|
137
132
|
Returns:
|
|
138
133
|
ModelResponse with generated text
|
|
139
134
|
"""
|
|
140
135
|
# Run _generate in thread pool to avoid blocking
|
|
141
136
|
loop = asyncio.get_event_loop()
|
|
142
137
|
generated_text = await loop.run_in_executor(
|
|
143
|
-
self._executor,
|
|
144
|
-
functools.partial(self._generate, **kwargs)
|
|
138
|
+
self._executor, functools.partial(self._generate, **kwargs)
|
|
145
139
|
)
|
|
146
|
-
|
|
140
|
+
|
|
147
141
|
return await acompletion(
|
|
148
142
|
model=f"huggingface-local/{kwargs['model']}",
|
|
149
143
|
mock_response=generated_text,
|
|
150
144
|
)
|
|
151
|
-
|
|
145
|
+
|
|
152
146
|
def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
|
|
153
147
|
"""Synchronous streaming method.
|
|
154
|
-
|
|
148
|
+
|
|
155
149
|
Returns:
|
|
156
150
|
Iterator of GenericStreamingChunk
|
|
157
151
|
"""
|
|
158
152
|
generated_text = self._generate(**kwargs)
|
|
159
|
-
|
|
153
|
+
|
|
160
154
|
generic_streaming_chunk: GenericStreamingChunk = {
|
|
161
155
|
"finish_reason": "stop",
|
|
162
156
|
"index": 0,
|
|
@@ -165,22 +159,21 @@ class HuggingFaceLocalAdapter(CustomLLM):
|
|
|
165
159
|
"tool_use": None,
|
|
166
160
|
"usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
|
|
167
161
|
}
|
|
168
|
-
|
|
162
|
+
|
|
169
163
|
yield generic_streaming_chunk
|
|
170
|
-
|
|
164
|
+
|
|
171
165
|
async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
|
|
172
166
|
"""Asynchronous streaming method.
|
|
173
|
-
|
|
167
|
+
|
|
174
168
|
Returns:
|
|
175
169
|
AsyncIterator of GenericStreamingChunk
|
|
176
170
|
"""
|
|
177
171
|
# Run _generate in thread pool to avoid blocking
|
|
178
172
|
loop = asyncio.get_event_loop()
|
|
179
173
|
generated_text = await loop.run_in_executor(
|
|
180
|
-
self._executor,
|
|
181
|
-
functools.partial(self._generate, **kwargs)
|
|
174
|
+
self._executor, functools.partial(self._generate, **kwargs)
|
|
182
175
|
)
|
|
183
|
-
|
|
176
|
+
|
|
184
177
|
generic_streaming_chunk: GenericStreamingChunk = {
|
|
185
178
|
"finish_reason": "stop",
|
|
186
179
|
"index": 0,
|
|
@@ -189,5 +182,5 @@ class HuggingFaceLocalAdapter(CustomLLM):
|
|
|
189
182
|
"tool_use": None,
|
|
190
183
|
"usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
|
|
191
184
|
}
|
|
192
|
-
|
|
193
|
-
yield generic_streaming_chunk
|
|
185
|
+
|
|
186
|
+
yield generic_streaming_chunk
|