Skip to content

Commit cc9b604

Browse files
harel-tussiharel-tussi
committed
wip
1 parent 6ed2649 commit cc9b604

8 files changed

Lines changed: 272 additions & 6 deletions

File tree

‎README.md‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
- [Installation](#Installation) 💿
2626
- [Usage - Chat Completions](#Usage)
2727
- [Maestro](#Maestro)
28+
- [Agents (Beta)](#Agents-Beta)
2829
- [Conversational RAG (Beta)](#Conversational-RAG-Beta)
2930
- [Older Models Support Usage](#Older-Models-Support-Usage)
3031
- [More Models](#More-Models)
@@ -277,6 +278,61 @@ For a more detailed example, see maestro [sync](examples/studio/maestro/run.py)
277278

278279
---
279280

281+
### Agents (Beta)
282+
283+
AI21 Agents provide a comprehensive way to create, manage, and run your Agents.
284+
285+
```python
286+
from ai21 import AI21Client
287+
from ai21.models.agents import BudgetLevel, AgentType
288+
289+
client = AI21Client()
290+
291+
# Run the agent
292+
run_response = client.beta.agents.runs.create_and_poll(
293+
agent_id=agent.id,
294+
input=[{"role": "user", "content": "What is 2+2?"}],
295+
poll_timeout_sec=120,
296+
)
297+
298+
print(f"Result: {run_response.result}")
299+
300+
```
301+
302+
#### Agent CRUD Operations
303+
304+
```python
305+
from ai21 import AI21Client
306+
from ai21.models.agents import BudgetLevel, AgentType
307+
308+
client = AI21Client()
309+
310+
# Create
311+
agent = client.beta.agents.create(
312+
name="Research Assistant",
313+
description="Specialized in research tasks",
314+
budget=BudgetLevel.HIGH,
315+
)
316+
317+
# Read
318+
retrieved_agent = client.beta.agents.get(agent.id)
319+
agents_list = client.beta.agents.list()
320+
321+
# Update
322+
modified_agent = client.beta.agents.modify(
323+
agent.id,
324+
name="Enhanced Research Assistant",
325+
description="Updated with enhanced capabilities",
326+
)
327+
328+
# Delete
329+
delete_response = client.beta.agents.delete(agent.id)
330+
```
331+
332+
For more detailed examples, see agent [CRUD operations](examples/studio/agents/agent_crud.py), [basic runs](examples/studio/agents/agent_run.py), and [async operations](examples/studio/agents/async_agent_run.py) examples.
333+
334+
---
335+
280336
### Conversational RAG (Beta)
281337

282338
Like chat, but with the ability to retrieve information from your Studio library.

‎ai21/clients/studio/resources/agents/agents.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
Visibility,
1717
)
1818
from ai21.models.agents.agent import ResponseLanguage
19-
from ai21.models.maestro.run import MaestroMessage, RunResponse
19+
from ai21.models.maestro.run import DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, MaestroMessage, RunResponse
2020
from ai21.types import NOT_GIVEN, NotGiven
2121

2222

@@ -57,6 +57,8 @@ def create_and_poll(
5757
agent_id: str,
5858
*,
5959
input: Union[str, List[MaestroMessage]],
60+
poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
61+
poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
6062
**kwargs,
6163
) -> RunResponse:
6264
"""Create and poll an agent run using maestro client"""
@@ -65,6 +67,8 @@ def create_and_poll(
6567
return self._maestro_runs.create_and_poll(
6668
input=input,
6769
**self.convert_agent_to_maestro_run_payload(agent),
70+
poll_interval_sec=poll_interval_sec,
71+
poll_timeout_sec=poll_timeout_sec,
6872
**kwargs,
6973
)
7074

@@ -84,7 +88,6 @@ def create(
8488
tool_resources: Union[Dict[str, Any], NotGiven] = NOT_GIVEN,
8589
requirements: Union[List[Requirement], NotGiven] = NOT_GIVEN,
8690
budget: Union[BudgetLevel, NotGiven] = NOT_GIVEN,
87-
agent_type: Union[AgentType, NotGiven] = NOT_GIVEN,
8891
response_language: Union[ResponseLanguage, NotGiven] = NOT_GIVEN,
8992
**kwargs,
9093
) -> Agent:
@@ -97,7 +100,6 @@ def create(
97100
tool_resources=tool_resources,
98101
requirements=requirements,
99102
budget=budget,
100-
agent_type=agent_type,
101103
response_language=response_language,
102104
**kwargs,
103105
)

‎ai21/clients/studio/resources/agents/async_agents.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
Requirement,
1717
Visibility,
1818
)
19-
from ai21.models.maestro.run import MaestroMessage, RunResponse
19+
from ai21.models.maestro.run import DEFAULT_RUN_POLL_INTERVAL, DEFAULT_RUN_POLL_TIMEOUT, MaestroMessage, RunResponse
2020
from ai21.types import NOT_GIVEN, NotGiven
2121

2222

@@ -55,6 +55,8 @@ async def create_and_poll(
5555
agent_id: str,
5656
*,
5757
input: List[MaestroMessage],
58+
poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
59+
poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
5860
**kwargs,
5961
) -> RunResponse:
6062
"""Create and poll an agent run using maestro client"""
@@ -63,6 +65,8 @@ async def create_and_poll(
6365
return await self._maestro_runs.create_and_poll(
6466
input=input,
6567
**self.convert_agent_to_maestro_run_payload(agent),
68+
poll_interval_sec=poll_interval_sec,
69+
poll_timeout_sec=poll_timeout_sec,
6670
**kwargs,
6771
)
6872

@@ -88,7 +92,6 @@ async def create(
8892
tool_resources: Dict[str, Any] | NotGiven = NOT_GIVEN,
8993
requirements: List[Requirement] | NotGiven = NOT_GIVEN,
9094
budget: Union[BudgetLevel, NotGiven] = NOT_GIVEN,
91-
agent_type: Union[AgentType, NotGiven] = NOT_GIVEN,
9295
**kwargs,
9396
) -> Agent:
9497
"""Create a new agent"""
@@ -102,7 +105,6 @@ async def create(
102105
tool_resources=tool_resources,
103106
requirements=requirements,
104107
budget=budget,
105-
agent_type=agent_type,
106108
**kwargs,
107109
)
108110

‎examples/studio/agents/__init__.py‎

Whitespace-only changes.

‎examples/studio/agents/agent_crud.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
from ai21 import AI21Client
2+
from ai21.models.agents import BudgetLevel, AgentType
3+
4+
client = AI21Client()
5+
6+
7+
def main():
8+
"""Example demonstrating CRUD operations for AI21 Agents"""
9+
10+
# Create an agent
11+
print("Creating an agent...")
12+
agent = client.beta.agents.create(
13+
name="My Assistant",
14+
description="A helpful AI assistant that can answer questions",
15+
budget=BudgetLevel.MEDIUM,
16+
)
17+
18+
print(f"Created agent: {agent.name} (ID: {agent.id})")
19+
print(f"Agent type: {agent.agent_type}")
20+
print(f"Budget: {agent.budget}")
21+
22+
agent_id = agent.id
23+
24+
try:
25+
# Get the agent
26+
print(f"\nRetrieving agent {agent_id}...")
27+
retrieved_agent = client.beta.agents.get(agent_id)
28+
print(f"Retrieved agent: {retrieved_agent.name}")
29+
30+
# List all agents
31+
print("\nListing all agents...")
32+
agents_list = client.beta.agents.list()
33+
print(f"Found {len(agents_list.results)} agents")
34+
35+
# Modify the agent
36+
print(f"\nModifying agent {agent_id}...")
37+
modified_agent = client.beta.agents.modify(
38+
agent_id,
39+
name="Updated Assistant",
40+
description="An updated AI assistant with enhanced capabilities",
41+
)
42+
print(f"Modified agent name: {modified_agent.name}")
43+
print(f"Modified agent description: {modified_agent.description}")
44+
45+
finally:
46+
# Delete the agent (cleanup)
47+
print(f"\nDeleting agent {agent_id}...")
48+
delete_response = client.beta.agents.delete(agent_id)
49+
print(f"Agent deleted: {delete_response.deleted}")
50+
51+
52+
if __name__ == "__main__":
53+
main()

‎examples/studio/agents/agent_run.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
from ai21 import AI21Client
2+
from ai21.models.agents import BudgetLevel, AgentType
3+
4+
client = AI21Client()
5+
6+
7+
def main():
8+
"""Example demonstrating how to create and run an AI21 Agent"""
9+
10+
# Create an agent
11+
print("Creating an agent...")
12+
agent = client.beta.agents.create(
13+
name="Math Assistant",
14+
description="An AI assistant specialized in solving math problems",
15+
budget=BudgetLevel.LOW,
16+
)
17+
18+
print(f"Created agent: {agent.name} (ID: {agent.id})")
19+
agent_id = agent.id
20+
21+
try:
22+
# Run the agent with a simple math question
23+
print("\nRunning agent with math question...")
24+
input_messages = [{"role": "user", "content": "What is 15 * 23? Please show your work."}]
25+
26+
run_response = client.beta.agents.runs.create_and_poll(
27+
agent_id=agent_id,
28+
input=input_messages,
29+
poll_timeout_sec=120, # 2 minutes timeout
30+
)
31+
32+
print(f"Run ID: {run_response.id}")
33+
print(f"Run status: {run_response.status}")
34+
35+
if run_response.status == "completed":
36+
print("Run completed successfully!")
37+
if run_response.result:
38+
print(f"Result: {run_response.result}")
39+
else:
40+
print(f"Run failed with status: {run_response.status}")
41+
42+
# Retrieve the run to show how to get run details
43+
print(f"\nRetrieving run details...")
44+
retrieved_run = client.beta.agents.runs.retrieve(str(run_response.id))
45+
print(f"Retrieved run status: {retrieved_run.status}")
46+
47+
except Exception as e:
48+
print(f"Error during agent run: {e}")
49+
finally:
50+
# Clean up - delete the agent
51+
print(f"\nCleaning up - deleting agent {agent_id}...")
52+
client.beta.agents.delete(agent_id)
53+
print("Agent deleted successfully")
54+
55+
56+
if __name__ == "__main__":
57+
main()

‎examples/studio/agents/async_agent_run.py‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import asyncio
2+
from ai21 import AsyncAI21Client
3+
from ai21.models.agents import BudgetLevel, AgentType
4+
5+
client = AsyncAI21Client()
6+
7+
8+
async def main():
9+
"""Example demonstrating async Agent operations with enhanced options"""
10+
11+
# Create an agent
12+
print("Creating an agent...")
13+
agent = await client.beta.agents.create(
14+
name="Research Assistant",
15+
description="An AI assistant that can help with research and analysis",
16+
budget=BudgetLevel.MEDIUM,
17+
)
18+
19+
print(f"Created agent: {agent.name} (ID: {agent.id})")
20+
agent_id = agent.id
21+
22+
try:
23+
# Run the agent with enhanced options
24+
print("\nRunning agent with research question...")
25+
input_messages = [{"role": "user", "content": "Explain the key benefits of renewable energy"}]
26+
27+
run_response = await client.beta.agents.runs.create_and_poll(
28+
agent_id=agent_id,
29+
input=input_messages,
30+
verbose=True,
31+
include=["data_sources", "requirements_result"],
32+
response_language="english",
33+
poll_timeout_sec=180, # 3 minutes timeout
34+
)
35+
36+
print(f"Run ID: {run_response.id}")
37+
print(f"Run status: {run_response.status}")
38+
39+
if run_response.status == "completed":
40+
print("Run completed successfully!")
41+
if run_response.result:
42+
print(f"Result: {run_response.result}")
43+
44+
# Show additional information if available
45+
if hasattr(run_response, "data_sources") and run_response.data_sources:
46+
print(f"Data sources used: {len(run_response.data_sources)}")
47+
48+
if hasattr(run_response, "requirements_result") and run_response.requirements_result:
49+
print(f"Requirements result: {run_response.requirements_result}")
50+
else:
51+
print(f"Run failed with status: {run_response.status}")
52+
53+
# Demonstrate multiple runs concurrently
54+
print("\nRunning multiple questions concurrently...")
55+
questions = [
56+
"What is photosynthesis?",
57+
"How do solar panels work?",
58+
"What are the main types of renewable energy?",
59+
]
60+
61+
tasks = []
62+
for i, question in enumerate(questions):
63+
input_msgs = [{"role": "user", "content": question}]
64+
task = client.beta.agents.runs.create_and_poll(
65+
agent_id=agent_id,
66+
input=input_msgs,
67+
poll_timeout_sec=120,
68+
)
69+
tasks.append(task)
70+
71+
# Wait for all runs to complete
72+
results = await asyncio.gather(*tasks, return_exceptions=True)
73+
74+
for i, result in enumerate(results):
75+
if isinstance(result, Exception):
76+
print(f"Question {i+1} failed: {result}")
77+
else:
78+
print(f"Question {i+1} status: {result.status}")
79+
80+
except Exception as e:
81+
print(f"Error during agent operations: {e}")
82+
finally:
83+
# Clean up - delete the agent
84+
print(f"\nCleaning up - deleting agent {agent_id}...")
85+
await client.beta.agents.delete(agent_id)
86+
print("Agent deleted successfully")
87+
88+
89+
if __name__ == "__main__":
90+
asyncio.run(main())

‎tests/integration_tests/clients/test_studio.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ def test_studio(test_file_name: str):
5454
("conversational_rag/async_conversational_rag.py",),
5555
("maestro/run.py",),
5656
("maestro/async_run.py",),
57+
("agents/agent_crud.py",),
58+
("agents/agent_run.py",),
59+
("agents/async_agent_run.py",),
5760
],
5861
ids=[
5962
"when_chat_completions__should_return_ok",
@@ -62,6 +65,9 @@ def test_studio(test_file_name: str):
6265
"when_async_conversational_rag__should_return_ok",
6366
"when_maestro_runs__should_return_ok",
6467
"when_maestro_async_runs__should_return_ok",
68+
"when_agent_crud__should_return_ok",
69+
"when_agent_run__should_return_ok",
70+
"when_async_agent_run__should_return_ok",
6571
],
6672
)
6773
async def test_async_studio(test_file_name: str):

0 commit comments

Comments
 (0)